diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index e918a71373..fc0f84a20d 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -33,10 +33,10 @@ jobs: poetry lock poetry install --with dev - - name: Run Black formatting + - name: Check Black formatting run: | cd litellm - poetry run black . + poetry run black --check --exclude '/enterprise/' . cd .. - name: Debug - Check file state diff --git a/deploy/charts/litellm-helm/templates/deployment.yaml b/deploy/charts/litellm-helm/templates/deployment.yaml index 51af22b7a4..3040fb45d8 100644 --- a/deploy/charts/litellm-helm/templates/deployment.yaml +++ b/deploy/charts/litellm-helm/templates/deployment.yaml @@ -20,6 +20,9 @@ spec: selector: matchLabels: {{- include "litellm.selectorLabels" . | nindent 6 }} + {{- if .Values.deploymentMinReadySeconds }} + minReadySeconds: {{ .Values.deploymentMinReadySeconds }} + {{- end }} template: metadata: annotations: diff --git a/deploy/charts/litellm-helm/tests/deployment_tests.yaml b/deploy/charts/litellm-helm/tests/deployment_tests.yaml index 2e9c48043d..0d278f2569 100644 --- a/deploy/charts/litellm-helm/tests/deployment_tests.yaml +++ b/deploy/charts/litellm-helm/tests/deployment_tests.yaml @@ -306,3 +306,16 @@ tests: - equal: path: spec.template.spec.containers[0].resources value: {} + - it: should be able to set minReadySeconds + template: deployment.yaml + set: + deploymentMinReadySeconds: 5 + asserts: + - equal: + path: spec.minReadySeconds + value: 5 + - it: should have minReadySeconds absent when deploymentMinReadySeconds is not set + template: deployment.yaml + asserts: + - notExists: + path: spec.minReadySeconds diff --git a/deploy/charts/litellm-helm/values.yaml b/deploy/charts/litellm-helm/values.yaml index f8944bddd5..690ca69e73 100644 --- a/deploy/charts/litellm-helm/values.yaml +++ b/deploy/charts/litellm-helm/values.yaml @@ -31,6 +31,8 @@ serviceAccount: # annotations for litellm deployment deploymentAnnotations: {} deploymentLabels: {} +deploymentMinReadySeconds: 0 + # annotations for litellm pods podAnnotations: {} podLabels: {} diff --git a/docs/my-website/docs/files_endpoints.md b/docs/my-website/docs/files_endpoints.md index 30677c748a..deb1793163 100644 --- a/docs/my-website/docs/files_endpoints.md +++ b/docs/my-website/docs/files_endpoints.md @@ -326,4 +326,10 @@ print("file content=", content.text) ### [Bedrock](./providers/bedrock_batches#4-retrieve-batch-results) +### [Anthropic](./providers/anthropic#files-api) + +:::note +Anthropic Files API has a different purpose than OpenAI's. It's **not** for Batches or Fine-tuning—it's for uploading files once and referencing them by `file_id` in multiple messages, avoiding re-uploads. File API operations are free — file content used in Messages requests is priced as input tokens. +::: + ## [Swagger API Reference](https://litellm-api.up.railway.app/#/files) diff --git a/docs/my-website/docs/providers/anthropic.md b/docs/my-website/docs/providers/anthropic.md index aa77ee7c26..50b964bd93 100644 --- a/docs/my-website/docs/providers/anthropic.md +++ b/docs/my-website/docs/providers/anthropic.md @@ -1965,6 +1965,98 @@ curl -L -X POST 'http://0.0.0.0:4000/v1/chat/completions' \ +## Files API + +Upload files once and reference them by `file_id` in multiple requests—no need to re-upload content each time. + +:::info +The `file_id` obtained from Anthropic only works with Anthropic Claude models. You cannot use it with other providers (OpenAI, Bedrock, etc.). +::: + +- **Max file size:** 500 MB | **Total storage:** 100 GB per org +- **Pricing:** File API operations are free. File content used in Messages requests is priced as input tokens. + +**Supported models by file type:** +- **Images:** All Claude 3+ models +- **PDFs:** All Claude 3.5+ models +- **Other file types** (for code execution): Claude 3.5 Haiku + all Claude 3.7+ models + +### Quick Start + +```python +import litellm +import os + +os.environ["ANTHROPIC_API_KEY"] = "sk-ant-..." + +# 1. Upload a file once +file = litellm.create_file( + file=open("document.pdf", "rb"), + purpose="messages", + custom_llm_provider="anthropic", +) + +# 2. Use file_id in messages (no re-upload needed) +response = litellm.completion( + model="anthropic/claude-sonnet-4-5-20250929", + messages=[{ + "role": "user", + "content": [ + {"type": "text", "text": "Summarize this document"}, + {"type": "file", "file": {"file_id": file.id, "format": "application/pdf"}} + ] + }] +) +``` + +### File Operations + +| Operation | Function | +|-----------|----------| +| Upload | `litellm.create_file(file, purpose="messages", custom_llm_provider="anthropic")` | +| List | `litellm.file_list(custom_llm_provider="anthropic")` | +| Retrieve | `litellm.file_retrieve(file_id, custom_llm_provider="anthropic")` | +| Delete | `litellm.file_delete(file_id, custom_llm_provider="anthropic")` | +| Download | `litellm.file_content(file_id, custom_llm_provider="anthropic")` | + +:::note +Download only works for files created by the [code execution tool](https://docs.anthropic.com/en/docs/agents-and-tools/tool-use/code-execution-tool), not uploaded files. +::: + +### Supported Formats + +| File Type | Format Value | +|-----------|-------------| +| PDF | `application/pdf` | +| Plain text | `text/plain` | +| JPEG | `image/jpeg` | +| PNG | `image/png` | +| GIF | `image/gif` | +| WebP | `image/webp` | + +### Using Images + +```python +# Upload image +image = litellm.create_file( + file=open("photo.jpg", "rb"), + purpose="messages", + custom_llm_provider="anthropic", +) + +# Use in message +response = litellm.completion( + model="anthropic/claude-sonnet-4-5-20250929", + messages=[{ + "role": "user", + "content": [ + {"type": "text", "text": "What's in this image?"}, + {"type": "file", "file": {"file_id": image.id, "format": "image/jpeg"}} + ] + }] +) +``` + ## Usage - passing 'user_id' to Anthropic LiteLLM translates the OpenAI `user` param to Anthropic's `metadata[user_id]` param. diff --git a/docs/my-website/docs/providers/openai.md b/docs/my-website/docs/providers/openai.md index 9d557303ef..80931ad821 100644 --- a/docs/my-website/docs/providers/openai.md +++ b/docs/my-website/docs/providers/openai.md @@ -638,7 +638,9 @@ This is useful when you want to use [Responses API](https://platform.openai.com/ :::tip gpt-5.4 + reasoning_effort + function tools -OpenAI does not support `reasoning_effort` with function tools for `gpt-5.4` in `/v1/chat/completions`. Use the responses bridge instead: +LiteLLM drops `reasoning_effort` from `gpt-5.4` requests to `litellm.completion()` that include tools, since that combination is supported in the Responses API. + +If you need reasoning **and** tools together, use the responses bridge instead: ```python response = litellm.completion( diff --git a/docs/my-website/docs/providers/vertex_partner.md b/docs/my-website/docs/providers/vertex_partner.md index 48a116eb7a..75ec3b9308 100644 --- a/docs/my-website/docs/providers/vertex_partner.md +++ b/docs/my-website/docs/providers/vertex_partner.md @@ -11,6 +11,7 @@ import TabItem from '@theme/TabItem'; |----------|---------------|---------------| | Anthropic (Claude) | `vertex_ai/claude-*` | [Vertex AI - Anthropic Models](https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/use-claude) | | DeepSeek | `vertex_ai/deepseek-ai/{MODEL}` | [Vertex AI - DeepSeek Models](https://cloud.google.com/vertex-ai/generative-ai/docs/maas/deepseek) | +| ZAI (GLM) | `vertex_ai/zai-org/{MODEL}` | [Vertex AI - GLM Models](https://cloud.google.com/vertex-ai/generative-ai/docs/maas/zaiorg/glm-47) | | Meta/Llama | `vertex_ai/meta/{MODEL}` | [Vertex AI - Meta Models](https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/llama) | | Mistral | `vertex_ai/mistral-*` | [Vertex AI - Mistral Models](https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/mistral) | | AI21 (Jamba) | `vertex_ai/jamba-*` | [Vertex AI - AI21 Models](https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/ai21) | @@ -226,6 +227,79 @@ ModelResponse( |------------------|------------------------------| | vertex_ai/deepseek-ai/deepseek-r1-0528-maas | `completion('vertex_ai/deepseek-ai/deepseek-r1-0528-maas', messages)` | +## VertexAI ZAI (GLM) + +| Property | Details | +|----------|---------| +| Provider Route | `vertex_ai/zai-org/{MODEL}` | +| Vertex Documentation | [Vertex AI - GLM Models](https://cloud.google.com/vertex-ai/generative-ai/docs/maas/zaiorg/glm-47) | + +**LiteLLM Supports all Vertex AI GLM Models.** Ensure you use the `vertex_ai/zai-org/` prefix for all Vertex AI GLM models. + +| Model Name | Usage | +|------------|-------| +| vertex_ai/zai-org/glm-4.7-maas | `completion('vertex_ai/zai-org/glm-4.7-maas', messages)` | + +#### Usage + + + + +```python +from litellm import completion +import os + +os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "" + +response = completion( + model="vertex_ai/zai-org/glm-4.7-maas", + messages=[{"role": "user", "content": "hi"}], + vertex_project="your-vertex-project", + # vertex_location routes to "global" +) +print("\nModel Response", response) +``` + + + +**1. Add to config** + +```yaml +model_list: + - model_name: glm-4.7 + litellm_params: + model: vertex_ai/zai-org/glm-4.7-maas + vertex_project: "my-project" + # vertex_location routes to "global" +``` + +**2. Start proxy** + +```bash +litellm --config /path/to/config.yaml + +# RUNNING at http://0.0.0.0:4000 +``` + +**3. Test it!** + +```bash +curl --location 'http://0.0.0.0:4000/chat/completions' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'Content-Type: application/json' \ + --data '{ + "model": "glm-4.7", + "messages": [ + { + "role": "user", + "content": "what llm are you" + } + ], + }' +``` + + + ## VertexAI Meta/Llama API diff --git a/docs/my-website/docs/reasoning_content.md b/docs/my-website/docs/reasoning_content.md index 5dd40122c7..8bf59f66a3 100644 --- a/docs/my-website/docs/reasoning_content.md +++ b/docs/my-website/docs/reasoning_content.md @@ -594,7 +594,9 @@ Expected Response :::tip gpt-5.4: reasoning_effort + function tools -OpenAI does not support `reasoning_effort` with function tools for `gpt-5.4` in `/v1/chat/completions`. Use `openai/responses/gpt-5.4` to route through the Responses API instead. See [Responses API Bridge](/docs/providers/openai#openai-chat-completion-to-responses-api-bridge) for details. +LiteLLM drops `reasoning_effort` from `gpt-5.4` requests to `litellm.completion()` that include tools, since that combination is supported in the Responses API. + +If you need reasoning **and** tools together, use `openai/responses/gpt-5.4` to route through the Responses API instead. See [Responses API Bridge](/docs/providers/openai#openai-chat-completion-to-responses-api-bridge) for details. ::: diff --git a/litellm-js/spend-logs/package-lock.json b/litellm-js/spend-logs/package-lock.json index 1a13a76820..b24ff0a494 100644 --- a/litellm-js/spend-logs/package-lock.json +++ b/litellm-js/spend-logs/package-lock.json @@ -6,7 +6,7 @@ "": { "dependencies": { "@hono/node-server": "^1.10.1", - "hono": "^4.10.3" + "hono": "^4.12.7" }, "devDependencies": { "@types/node": "^20.11.17", @@ -548,9 +548,9 @@ } }, "node_modules/hono": { - "version": "4.10.6", - "resolved": "https://registry.npmjs.org/hono/-/hono-4.10.6.tgz", - "integrity": "sha512-BIdolzGpDO9MQ4nu3AUuDwHZZ+KViNm+EZ75Ae55eMXMqLVhDFqEMXxtUe9Qh8hjL+pIna/frs2j6Y2yD5Ua/g==", + "version": "4.12.7", + "resolved": "https://registry.npmjs.org/hono/-/hono-4.12.7.tgz", + "integrity": "sha512-jq9l1DM0zVIvsm3lv9Nw9nlJnMNPOcAtsbsgiUhWcFzPE99Gvo6yRTlszSLLYacMeQ6quHD6hMfId8crVHvexw==", "license": "MIT", "engines": { "node": ">=16.9.0" diff --git a/litellm-js/spend-logs/package.json b/litellm-js/spend-logs/package.json index adfe49017d..a40b0fc2a8 100644 --- a/litellm-js/spend-logs/package.json +++ b/litellm-js/spend-logs/package.json @@ -4,7 +4,7 @@ }, "dependencies": { "@hono/node-server": "^1.10.1", - "hono": "^4.10.3" + "hono": "^4.12.7" }, "devDependencies": { "@types/node": "^20.11.17", diff --git a/litellm/__init__.py b/litellm/__init__.py index 0f7bac67c0..8b3723cb2b 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -81,6 +81,7 @@ from litellm.constants import ( DEFAULT_ALLOWED_FAILS, ) import httpx + # register_async_client_cleanup is lazy-loaded and called on first access litellm_mode = os.getenv("LITELLM_MODE", "DEV") # "PRODUCTION", "DEV" @@ -152,7 +153,9 @@ _known_custom_logger_compatible_callbacks: List = list( get_args(_custom_logger_compatible_callbacks_literal) ) callbacks: List[ - Union[Callable, _custom_logger_compatible_callbacks_literal, "CustomLogger"] # CustomLogger is lazy-loaded + Union[ + Callable, _custom_logger_compatible_callbacks_literal, "CustomLogger" + ] # CustomLogger is lazy-loaded ] = [] callback_settings: Dict[str, Dict[str, Any]] = {} initialized_langfuse_clients: int = 0 @@ -162,42 +165,50 @@ prometheus_initialize_budget_metrics: Optional[bool] = False require_auth_for_metrics_endpoint: Optional[bool] = False argilla_batch_size: Optional[int] = None datadog_use_v1: Optional[bool] = False # if you want to use v1 datadog logged payload. -gcs_pub_sub_use_v1: Optional[bool] = ( - False # if you want to use v1 gcs pubsub logged payload -) -generic_api_use_v1: Optional[bool] = ( - False # if you want to use v1 generic api logged payload -) +gcs_pub_sub_use_v1: Optional[ + bool +] = False # if you want to use v1 gcs pubsub logged payload +generic_api_use_v1: Optional[ + bool +] = False # if you want to use v1 generic api logged payload argilla_transformation_object: Optional[Dict[str, Any]] = None -_async_input_callback: List[Union[str, Callable, "CustomLogger"]] = ( # CustomLogger is lazy-loaded +_async_input_callback: List[ + Union[str, Callable, "CustomLogger"] +] = ( # CustomLogger is lazy-loaded [] ) # internal variable - async custom callbacks are routed here. -_async_success_callback: List[Union[str, Callable, "CustomLogger"]] = ( # CustomLogger is lazy-loaded +_async_success_callback: List[ + Union[str, Callable, "CustomLogger"] +] = ( # CustomLogger is lazy-loaded [] ) # internal variable - async custom callbacks are routed here. -_async_failure_callback: List[Union[str, Callable, "CustomLogger"]] = ( # CustomLogger is lazy-loaded +_async_failure_callback: List[ + Union[str, Callable, "CustomLogger"] +] = ( # CustomLogger is lazy-loaded [] ) # internal variable - async custom callbacks are routed here. pre_call_rules: List[Callable] = [] post_call_rules: List[Callable] = [] turn_off_message_logging: Optional[bool] = False -standard_logging_payload_excluded_fields: Optional[List[str]] = None # Fields to exclude from StandardLoggingPayload before callbacks receive it +standard_logging_payload_excluded_fields: Optional[ + List[str] +] = None # Fields to exclude from StandardLoggingPayload before callbacks receive it log_raw_request_response: bool = False redact_messages_in_exceptions: Optional[bool] = False redact_user_api_key_info: Optional[bool] = False filter_invalid_headers: Optional[bool] = False -add_user_information_to_llm_headers: Optional[bool] = ( - None # adds user_id, team_id, token hash (params from StandardLoggingMetadata) to request headers -) +add_user_information_to_llm_headers: Optional[ + bool +] = None # adds user_id, team_id, token hash (params from StandardLoggingMetadata) to request headers store_audit_logs = False # Enterprise feature, allow users to see audit logs ### end of callbacks ############# -email: Optional[str] = ( - None # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 -) -token: Optional[str] = ( - None # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 -) +email: Optional[ + str +] = None # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 +token: Optional[ + str +] = None # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 telemetry = True max_tokens: int = DEFAULT_MAX_TOKENS # OpenAI Defaults drop_params = bool(os.getenv("LITELLM_DROP_PARAMS", False)) @@ -259,9 +270,9 @@ use_client: bool = False ssl_verify: Union[str, bool] = True ssl_security_level: Optional[str] = None ssl_certificate: Optional[str] = None -ssl_ecdh_curve: Optional[str] = ( - None # Set to 'X25519' to disable PQC and improve performance -) +ssl_ecdh_curve: Optional[ + str +] = None # Set to 'X25519' to disable PQC and improve performance disable_streaming_logging: bool = False disable_token_counter: bool = False disable_add_transform_inline_image_block: bool = False @@ -314,24 +325,20 @@ enable_loadbalancing_on_batch_endpoints: Optional[bool] = None enable_caching_on_provider_specific_optional_params: bool = ( False # feature-flag for caching on optional params - e.g. 'top_k' ) -caching: bool = ( - False # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 -) -caching_with_models: bool = ( - False # # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 -) -cache: Optional["Cache"] = ( - None # cache object <- use this - https://docs.litellm.ai/docs/caching -) +caching: bool = False # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 +caching_with_models: bool = False # # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 +cache: Optional[ + "Cache" +] = None # cache object <- use this - https://docs.litellm.ai/docs/caching default_in_memory_ttl: Optional[float] = None default_redis_ttl: Optional[float] = None default_redis_batch_cache_expiry: Optional[float] = None model_alias_map: Dict[str, str] = {} model_group_settings: Optional["ModelGroupSettings"] = None max_budget: float = 0.0 # set the max budget across all providers -budget_duration: Optional[str] = ( - None # proxy only - resets budget after fixed duration. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). -) +budget_duration: Optional[ + str +] = None # proxy only - resets budget after fixed duration. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). default_soft_budget: float = ( DEFAULT_SOFT_BUDGET # by default all litellm proxy keys have a soft budget of 50.0 ) @@ -340,9 +347,7 @@ forward_traceparent_to_llm_provider: bool = False _current_cost = 0.0 # private variable, used if max budget is set error_logs: Dict = {} -add_function_to_prompt: bool = ( - False # if function calling not supported by api, append function call details to system prompt -) +add_function_to_prompt: bool = False # if function calling not supported by api, append function call details to system prompt client_session: Optional[httpx.Client] = None aclient_session: Optional[httpx.AsyncClient] = None model_fallbacks: Optional[List] = None # Deprecated for 'litellm.fallbacks' @@ -389,9 +394,7 @@ prometheus_emit_stream_label: bool = False disable_add_prefix_to_prompt: bool = ( False # used by anthropic, to disable adding prefix to prompt ) -disable_copilot_system_to_assistant: bool = ( - False # If false (default), converts all 'system' role messages to 'assistant' for GitHub Copilot compatibility. Set to true to disable this behavior. -) +disable_copilot_system_to_assistant: bool = False # If false (default), converts all 'system' role messages to 'assistant' for GitHub Copilot compatibility. Set to true to disable this behavior. public_mcp_servers: Optional[List[str]] = None public_model_groups: Optional[List[str]] = None public_agent_groups: Optional[List[str]] = None @@ -410,17 +413,13 @@ if TYPE_CHECKING: ######## Networking Settings ######## -use_aiohttp_transport: bool = ( - True # Older variable, aiohttp is now the default. use disable_aiohttp_transport instead. -) +use_aiohttp_transport: bool = True # Older variable, aiohttp is now the default. use disable_aiohttp_transport instead. aiohttp_trust_env: bool = False # set to true to use HTTP_ Proxy settings disable_aiohttp_transport: bool = False # Set this to true to use httpx instead disable_aiohttp_trust_env: bool = ( False # When False, aiohttp will respect HTTP(S)_PROXY env vars ) -force_ipv4: bool = ( - False # when True, litellm will force ipv4 for all LLM requests. Some users have seen httpx ConnectionError when using ipv6. -) +force_ipv4: bool = False # when True, litellm will force ipv4 for all LLM requests. Some users have seen httpx ConnectionError when using ipv6. network_mock: bool = False # When True, use mock transport — no real network calls ####### STOP SEQUENCE LIMIT ####### @@ -435,13 +434,13 @@ context_window_fallbacks: Optional[List] = None content_policy_fallbacks: Optional[List] = None allowed_fails: int = 3 allow_dynamic_callback_disabling: bool = True -num_retries_per_request: Optional[int] = ( - None # for the request overall (incl. fallbacks + model retries) -) +num_retries_per_request: Optional[ + int +] = None # for the request overall (incl. fallbacks + model retries) ####### SECRET MANAGERS ##################### -secret_manager_client: Optional[Any] = ( - None # list of instantiated key management clients - e.g. azure kv, infisical, etc. -) +secret_manager_client: Optional[ + Any +] = None # list of instantiated key management clients - e.g. azure kv, infisical, etc. _google_kms_resource_name: Optional[str] = None _key_management_system: Optional["KeyManagementSystem"] = None # Note: KeyManagementSettings must be eagerly imported because _key_management_settings @@ -454,12 +453,12 @@ output_parse_pii: bool = False from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map model_cost = get_model_cost_map(url=model_cost_map_url) -cost_discount_config: Dict[str, float] = ( - {} -) # Provider-specific cost discounts {"vertex_ai": 0.05} = 5% discount -cost_margin_config: Dict[str, Union[float, Dict[str, float]]] = ( - {} -) # Provider-specific or global cost margins. Examples: +cost_discount_config: Dict[ + str, float +] = {} # Provider-specific cost discounts {"vertex_ai": 0.05} = 5% discount +cost_margin_config: Dict[ + str, Union[float, Dict[str, float]] +] = {} # Provider-specific or global cost margins. Examples: # Percentage: {"openai": 0.10} = 10% margin # Fixed: {"openai": {"fixed_amount": 0.001}} = $0.001 per request # Global: {"global": 0.05} = 5% global margin on all providers @@ -1077,7 +1076,7 @@ models_by_provider: dict = { "aws_polly": aws_polly_models, "gigachat": gigachat_models, "llamagate": llamagate_models, - "bedrock_mantle": bedrock_mantle_models + "bedrock_mantle": bedrock_mantle_models, } # mapping for those models which have larger equivalents @@ -1128,10 +1127,12 @@ openai_video_generation_models = ["sora-2"] # Import KeyManagementSettings here (before utils import) because _key_management_settings # is accessed during import time in secret_managers/main.py (via dd_tracing -> datadog -> _service_logger -> utils) from litellm.types.secret_managers.main import KeyManagementSettings + _key_management_settings: KeyManagementSettings = KeyManagementSettings() # client must be imported immediately as it's used as a decorator at function definition time from .utils import client + # Note: Most other utils imports are lazy-loaded via __getattr__ to avoid loading utils.py # (which imports tiktoken) at import time @@ -1160,6 +1161,7 @@ from .llms.topaz.common_utils import TopazModelInfo # OpenAIOSeriesConfig is lazy loaded - openaiOSeriesConfig will be created on first access # OpenAIGPTConfig, OpenAIGPT5Config, etc. are lazy loaded - instances will be created on first access from .llms.xai.common_utils import XAIModelInfo + # PublicAI now uses JSON-based configuration (see litellm/llms/openai_like/providers.json) # All remaining configs are now lazy loaded - see _lazy_imports_registry.py @@ -1241,6 +1243,7 @@ from .batch_completion.main import * # type: ignore from .rerank_api.main import * from .llms.anthropic.experimental_pass_through.messages.handler import * from .responses.main import * + # Interactions API is available as litellm.interactions module # Usage: litellm.interactions.create(), litellm.interactions.get(), etc. from . import interactions @@ -1258,7 +1261,11 @@ from .containers.main import * from .ocr.main import * from .rag.main import * from .search.main import * -from .realtime_api.main import _arealtime, acreate_realtime_client_secret, arealtime_calls +from .realtime_api.main import ( + _arealtime, + acreate_realtime_client_secret, + arealtime_calls, +) from .responses.main import _aresponses_websocket from .fine_tuning.main import * from .files.main import * @@ -1300,12 +1307,12 @@ from . import rag from .types.llms.custom_llm import CustomLLMItem custom_provider_map: List[CustomLLMItem] = [] -_custom_providers: List[str] = ( - [] -) # internal helper util, used to track names of custom providers -disable_hf_tokenizer_download: Optional[bool] = ( - None # disable huggingface tokenizer download. Defaults to openai clk100 -) +_custom_providers: List[ + str +] = [] # internal helper util, used to track names of custom providers +disable_hf_tokenizer_download: Optional[ + bool +] = None # disable huggingface tokenizer download. Defaults to openai clk100 global_disable_no_log_param: bool = False ### CLI UTILITIES ### @@ -1344,131 +1351,327 @@ if TYPE_CHECKING: from litellm.caching.caching import Cache # Type stubs for lazy-loaded configs to help mypy - from .llms.bedrock.chat.converse_transformation import AmazonConverseConfig as AmazonConverseConfig - from .llms.openai_like.chat.handler import OpenAILikeChatConfig as OpenAILikeChatConfig - from .llms.galadriel.chat.transformation import GaladrielChatConfig as GaladrielChatConfig + from .llms.bedrock.chat.converse_transformation import ( + AmazonConverseConfig as AmazonConverseConfig, + ) + from .llms.openai_like.chat.handler import ( + OpenAILikeChatConfig as OpenAILikeChatConfig, + ) + from .llms.galadriel.chat.transformation import ( + GaladrielChatConfig as GaladrielChatConfig, + ) from .llms.github.chat.transformation import GithubChatConfig as GithubChatConfig - from .llms.azure_ai.anthropic.transformation import AzureAnthropicConfig as AzureAnthropicConfig + from .llms.azure_ai.anthropic.transformation import ( + AzureAnthropicConfig as AzureAnthropicConfig, + ) from .llms.bytez.chat.transformation import BytezChatConfig as BytezChatConfig - from .llms.compactifai.chat.transformation import CompactifAIChatConfig as CompactifAIChatConfig + from .llms.compactifai.chat.transformation import ( + CompactifAIChatConfig as CompactifAIChatConfig, + ) from .llms.empower.chat.transformation import EmpowerChatConfig as EmpowerChatConfig from .llms.minimax.chat.transformation import MinimaxChatConfig as MinimaxChatConfig - from .llms.aiohttp_openai.chat.transformation import AiohttpOpenAIChatConfig as AiohttpOpenAIChatConfig - from .llms.huggingface.chat.transformation import HuggingFaceChatConfig as HuggingFaceChatConfig - from .llms.huggingface.embedding.transformation import HuggingFaceEmbeddingConfig as HuggingFaceEmbeddingConfig + from .llms.aiohttp_openai.chat.transformation import ( + AiohttpOpenAIChatConfig as AiohttpOpenAIChatConfig, + ) + from .llms.huggingface.chat.transformation import ( + HuggingFaceChatConfig as HuggingFaceChatConfig, + ) + from .llms.huggingface.embedding.transformation import ( + HuggingFaceEmbeddingConfig as HuggingFaceEmbeddingConfig, + ) from .llms.oobabooga.chat.transformation import OobaboogaConfig as OobaboogaConfig from .llms.maritalk import MaritalkConfig as MaritalkConfig - from .llms.openrouter.chat.transformation import OpenrouterConfig as OpenrouterConfig + from .llms.openrouter.chat.transformation import ( + OpenrouterConfig as OpenrouterConfig, + ) from .llms.datarobot.chat.transformation import DataRobotConfig as DataRobotConfig from .llms.anthropic.chat.transformation import AnthropicConfig as AnthropicConfig - from .llms.anthropic.completion.transformation import AnthropicTextConfig as AnthropicTextConfig + from .llms.anthropic.completion.transformation import ( + AnthropicTextConfig as AnthropicTextConfig, + ) from .llms.groq.stt.transformation import GroqSTTConfig as GroqSTTConfig from .llms.triton.completion.transformation import TritonConfig as TritonConfig - from .llms.triton.completion.transformation import TritonGenerateConfig as TritonGenerateConfig - from .llms.triton.completion.transformation import TritonInferConfig as TritonInferConfig - from .llms.triton.embedding.transformation import TritonEmbeddingConfig as TritonEmbeddingConfig - from .llms.huggingface.rerank.transformation import HuggingFaceRerankConfig as HuggingFaceRerankConfig - from .llms.databricks.chat.transformation import DatabricksConfig as DatabricksConfig - from .llms.databricks.embed.transformation import DatabricksEmbeddingConfig as DatabricksEmbeddingConfig + from .llms.triton.completion.transformation import ( + TritonGenerateConfig as TritonGenerateConfig, + ) + from .llms.triton.completion.transformation import ( + TritonInferConfig as TritonInferConfig, + ) + from .llms.triton.embedding.transformation import ( + TritonEmbeddingConfig as TritonEmbeddingConfig, + ) + from .llms.huggingface.rerank.transformation import ( + HuggingFaceRerankConfig as HuggingFaceRerankConfig, + ) + from .llms.databricks.chat.transformation import ( + DatabricksConfig as DatabricksConfig, + ) + from .llms.databricks.embed.transformation import ( + DatabricksEmbeddingConfig as DatabricksEmbeddingConfig, + ) from .llms.predibase.chat.transformation import PredibaseConfig as PredibaseConfig from .llms.replicate.chat.transformation import ReplicateConfig as ReplicateConfig from .llms.snowflake.chat.transformation import SnowflakeConfig as SnowflakeConfig - from .llms.cohere.rerank.transformation import CohereRerankConfig as CohereRerankConfig - from .llms.cohere.rerank_v2.transformation import CohereRerankV2Config as CohereRerankV2Config - from .llms.azure_ai.rerank.transformation import AzureAIRerankConfig as AzureAIRerankConfig - from .llms.infinity.rerank.transformation import InfinityRerankConfig as InfinityRerankConfig - from .llms.jina_ai.rerank.transformation import JinaAIRerankConfig as JinaAIRerankConfig - from .llms.deepinfra.rerank.transformation import DeepinfraRerankConfig as DeepinfraRerankConfig - from .llms.hosted_vllm.rerank.transformation import HostedVLLMRerankConfig as HostedVLLMRerankConfig - from .llms.nvidia_nim.rerank.transformation import NvidiaNimRerankConfig as NvidiaNimRerankConfig - from .llms.nvidia_nim.rerank.ranking_transformation import NvidiaNimRankingConfig as NvidiaNimRankingConfig - from .llms.vertex_ai.rerank.transformation import VertexAIRerankConfig as VertexAIRerankConfig - from .llms.fireworks_ai.rerank.transformation import FireworksAIRerankConfig as FireworksAIRerankConfig - from .llms.voyage.rerank.transformation import VoyageRerankConfig as VoyageRerankConfig - from .llms.watsonx.rerank.transformation import IBMWatsonXRerankConfig as IBMWatsonXRerankConfig + from .llms.cohere.rerank.transformation import ( + CohereRerankConfig as CohereRerankConfig, + ) + from .llms.cohere.rerank_v2.transformation import ( + CohereRerankV2Config as CohereRerankV2Config, + ) + from .llms.azure_ai.rerank.transformation import ( + AzureAIRerankConfig as AzureAIRerankConfig, + ) + from .llms.infinity.rerank.transformation import ( + InfinityRerankConfig as InfinityRerankConfig, + ) + from .llms.jina_ai.rerank.transformation import ( + JinaAIRerankConfig as JinaAIRerankConfig, + ) + from .llms.deepinfra.rerank.transformation import ( + DeepinfraRerankConfig as DeepinfraRerankConfig, + ) + from .llms.hosted_vllm.rerank.transformation import ( + HostedVLLMRerankConfig as HostedVLLMRerankConfig, + ) + from .llms.nvidia_nim.rerank.transformation import ( + NvidiaNimRerankConfig as NvidiaNimRerankConfig, + ) + from .llms.nvidia_nim.rerank.ranking_transformation import ( + NvidiaNimRankingConfig as NvidiaNimRankingConfig, + ) + from .llms.vertex_ai.rerank.transformation import ( + VertexAIRerankConfig as VertexAIRerankConfig, + ) + from .llms.fireworks_ai.rerank.transformation import ( + FireworksAIRerankConfig as FireworksAIRerankConfig, + ) + from .llms.voyage.rerank.transformation import ( + VoyageRerankConfig as VoyageRerankConfig, + ) + from .llms.watsonx.rerank.transformation import ( + IBMWatsonXRerankConfig as IBMWatsonXRerankConfig, + ) from .llms.clarifai.chat.transformation import ClarifaiConfig as ClarifaiConfig from .llms.ai21.chat.transformation import AI21ChatConfig as AI21ChatConfig from .llms.meta_llama.chat.transformation import LlamaAPIConfig as LlamaAPIConfig - from .llms.together_ai.completion.transformation import TogetherAITextCompletionConfig as TogetherAITextCompletionConfig - from .llms.cloudflare.chat.transformation import CloudflareChatConfig as CloudflareChatConfig + from .llms.together_ai.completion.transformation import ( + TogetherAITextCompletionConfig as TogetherAITextCompletionConfig, + ) + from .llms.cloudflare.chat.transformation import ( + CloudflareChatConfig as CloudflareChatConfig, + ) from .llms.novita.chat.transformation import NovitaConfig as NovitaConfig from .llms.petals.completion.transformation import PetalsConfig as PetalsConfig from .llms.ollama.chat.transformation import OllamaChatConfig as OllamaChatConfig from .llms.ollama.completion.transformation import OllamaConfig as OllamaConfig - from .llms.sagemaker.completion.transformation import SagemakerConfig as SagemakerConfig - from .llms.sagemaker.chat.transformation import SagemakerChatConfig as SagemakerChatConfig + from .llms.sagemaker.completion.transformation import ( + SagemakerConfig as SagemakerConfig, + ) + from .llms.sagemaker.chat.transformation import ( + SagemakerChatConfig as SagemakerChatConfig, + ) from .llms.cohere.chat.transformation import CohereChatConfig as CohereChatConfig - from .llms.anthropic.experimental_pass_through.messages.transformation import AnthropicMessagesConfig as AnthropicMessagesConfig - from .llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import AmazonAnthropicClaudeMessagesConfig as AmazonAnthropicClaudeMessagesConfig + from .llms.anthropic.experimental_pass_through.messages.transformation import ( + AnthropicMessagesConfig as AnthropicMessagesConfig, + ) + from .llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeMessagesConfig as AmazonAnthropicClaudeMessagesConfig, + ) from .llms.together_ai.chat import TogetherAIConfig as TogetherAIConfig from .llms.nlp_cloud.chat.handler import NLPCloudConfig as NLPCloudConfig - from .llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig as VertexGeminiConfig - from .llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig as GoogleAIStudioGeminiConfig - from .llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import VertexAIAnthropicConfig as VertexAIAnthropicConfig - from .llms.vertex_ai.vertex_ai_partner_models.llama3.transformation import VertexAILlama3Config as VertexAILlama3Config - from .llms.vertex_ai.vertex_ai_partner_models.ai21.transformation import VertexAIAi21Config as VertexAIAi21Config - from .llms.bedrock.chat.invoke_handler import AmazonCohereChatConfig as AmazonCohereChatConfig - from .llms.bedrock.common_utils import AmazonBedrockGlobalConfig as AmazonBedrockGlobalConfig - from .llms.bedrock.chat.invoke_transformations.amazon_ai21_transformation import AmazonAI21Config as AmazonAI21Config - from .llms.bedrock.chat.invoke_transformations.amazon_nova_transformation import AmazonInvokeNovaConfig as AmazonInvokeNovaConfig - from .llms.bedrock.chat.invoke_transformations.amazon_qwen2_transformation import AmazonQwen2Config as AmazonQwen2Config - from .llms.bedrock.chat.invoke_transformations.amazon_qwen3_transformation import AmazonQwen3Config as AmazonQwen3Config - from .llms.bedrock.chat.invoke_transformations.anthropic_claude2_transformation import AmazonAnthropicConfig as AmazonAnthropicConfig - from .llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import AmazonAnthropicClaudeConfig as AmazonAnthropicClaudeConfig - from .llms.bedrock.chat.invoke_transformations.amazon_cohere_transformation import AmazonCohereConfig as AmazonCohereConfig - from .llms.bedrock.chat.invoke_transformations.amazon_llama_transformation import AmazonLlamaConfig as AmazonLlamaConfig - from .llms.bedrock.chat.invoke_transformations.amazon_deepseek_transformation import AmazonDeepSeekR1Config as AmazonDeepSeekR1Config - from .llms.bedrock.chat.invoke_transformations.amazon_mistral_transformation import AmazonMistralConfig as AmazonMistralConfig - from .llms.bedrock.chat.invoke_transformations.amazon_moonshot_transformation import AmazonMoonshotConfig as AmazonMoonshotConfig - from .llms.bedrock.chat.invoke_transformations.amazon_titan_transformation import AmazonTitanConfig as AmazonTitanConfig - from .llms.bedrock.chat.invoke_transformations.amazon_twelvelabs_pegasus_transformation import AmazonTwelveLabsPegasusConfig as AmazonTwelveLabsPegasusConfig - from .llms.bedrock.chat.invoke_transformations.base_invoke_transformation import AmazonInvokeConfig as AmazonInvokeConfig - from .llms.bedrock.chat.invoke_transformations.amazon_openai_transformation import AmazonBedrockOpenAIConfig as AmazonBedrockOpenAIConfig - from .llms.bedrock.image_generation.amazon_stability1_transformation import AmazonStabilityConfig as AmazonStabilityConfig - from .llms.bedrock.image_generation.amazon_stability3_transformation import AmazonStability3Config as AmazonStability3Config - from .llms.bedrock.image_generation.amazon_nova_canvas_transformation import AmazonNovaCanvasConfig as AmazonNovaCanvasConfig - from .llms.bedrock.embed.amazon_titan_g1_transformation import AmazonTitanG1Config as AmazonTitanG1Config - from .llms.bedrock.embed.amazon_titan_multimodal_transformation import AmazonTitanMultimodalEmbeddingG1Config as AmazonTitanMultimodalEmbeddingG1Config - from .llms.cohere.chat.v2_transformation import CohereV2ChatConfig as CohereV2ChatConfig - from .llms.bedrock.embed.cohere_transformation import BedrockCohereEmbeddingConfig as BedrockCohereEmbeddingConfig - from .llms.bedrock.embed.twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig as TwelveLabsMarengoEmbeddingConfig - from .llms.bedrock.embed.amazon_nova_transformation import AmazonNovaEmbeddingConfig as AmazonNovaEmbeddingConfig - from .llms.openai.openai import OpenAIConfig as OpenAIConfig, MistralEmbeddingConfig as MistralEmbeddingConfig - from .llms.openai.image_variations.transformation import OpenAIImageVariationConfig as OpenAIImageVariationConfig - from .llms.deepgram.audio_transcription.transformation import DeepgramAudioTranscriptionConfig as DeepgramAudioTranscriptionConfig - from .llms.topaz.image_variations.transformation import TopazImageVariationConfig as TopazImageVariationConfig - from litellm.llms.openai.completion.transformation import OpenAITextCompletionConfig as OpenAITextCompletionConfig + from .llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig as VertexGeminiConfig, + ) + from .llms.gemini.chat.transformation import ( + GoogleAIStudioGeminiConfig as GoogleAIStudioGeminiConfig, + ) + from .llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import ( + VertexAIAnthropicConfig as VertexAIAnthropicConfig, + ) + from .llms.vertex_ai.vertex_ai_partner_models.llama3.transformation import ( + VertexAILlama3Config as VertexAILlama3Config, + ) + from .llms.vertex_ai.vertex_ai_partner_models.ai21.transformation import ( + VertexAIAi21Config as VertexAIAi21Config, + ) + from .llms.bedrock.chat.invoke_handler import ( + AmazonCohereChatConfig as AmazonCohereChatConfig, + ) + from .llms.bedrock.common_utils import ( + AmazonBedrockGlobalConfig as AmazonBedrockGlobalConfig, + ) + from .llms.bedrock.chat.invoke_transformations.amazon_ai21_transformation import ( + AmazonAI21Config as AmazonAI21Config, + ) + from .llms.bedrock.chat.invoke_transformations.amazon_nova_transformation import ( + AmazonInvokeNovaConfig as AmazonInvokeNovaConfig, + ) + from .llms.bedrock.chat.invoke_transformations.amazon_qwen2_transformation import ( + AmazonQwen2Config as AmazonQwen2Config, + ) + from .llms.bedrock.chat.invoke_transformations.amazon_qwen3_transformation import ( + AmazonQwen3Config as AmazonQwen3Config, + ) + from .llms.bedrock.chat.invoke_transformations.anthropic_claude2_transformation import ( + AmazonAnthropicConfig as AmazonAnthropicConfig, + ) + from .llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeConfig as AmazonAnthropicClaudeConfig, + ) + from .llms.bedrock.chat.invoke_transformations.amazon_cohere_transformation import ( + AmazonCohereConfig as AmazonCohereConfig, + ) + from .llms.bedrock.chat.invoke_transformations.amazon_llama_transformation import ( + AmazonLlamaConfig as AmazonLlamaConfig, + ) + from .llms.bedrock.chat.invoke_transformations.amazon_deepseek_transformation import ( + AmazonDeepSeekR1Config as AmazonDeepSeekR1Config, + ) + from .llms.bedrock.chat.invoke_transformations.amazon_mistral_transformation import ( + AmazonMistralConfig as AmazonMistralConfig, + ) + from .llms.bedrock.chat.invoke_transformations.amazon_moonshot_transformation import ( + AmazonMoonshotConfig as AmazonMoonshotConfig, + ) + from .llms.bedrock.chat.invoke_transformations.amazon_titan_transformation import ( + AmazonTitanConfig as AmazonTitanConfig, + ) + from .llms.bedrock.chat.invoke_transformations.amazon_twelvelabs_pegasus_transformation import ( + AmazonTwelveLabsPegasusConfig as AmazonTwelveLabsPegasusConfig, + ) + from .llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( + AmazonInvokeConfig as AmazonInvokeConfig, + ) + from .llms.bedrock.chat.invoke_transformations.amazon_openai_transformation import ( + AmazonBedrockOpenAIConfig as AmazonBedrockOpenAIConfig, + ) + from .llms.bedrock.image_generation.amazon_stability1_transformation import ( + AmazonStabilityConfig as AmazonStabilityConfig, + ) + from .llms.bedrock.image_generation.amazon_stability3_transformation import ( + AmazonStability3Config as AmazonStability3Config, + ) + from .llms.bedrock.image_generation.amazon_nova_canvas_transformation import ( + AmazonNovaCanvasConfig as AmazonNovaCanvasConfig, + ) + from .llms.bedrock.embed.amazon_titan_g1_transformation import ( + AmazonTitanG1Config as AmazonTitanG1Config, + ) + from .llms.bedrock.embed.amazon_titan_multimodal_transformation import ( + AmazonTitanMultimodalEmbeddingG1Config as AmazonTitanMultimodalEmbeddingG1Config, + ) + from .llms.cohere.chat.v2_transformation import ( + CohereV2ChatConfig as CohereV2ChatConfig, + ) + from .llms.bedrock.embed.cohere_transformation import ( + BedrockCohereEmbeddingConfig as BedrockCohereEmbeddingConfig, + ) + from .llms.bedrock.embed.twelvelabs_marengo_transformation import ( + TwelveLabsMarengoEmbeddingConfig as TwelveLabsMarengoEmbeddingConfig, + ) + from .llms.bedrock.embed.amazon_nova_transformation import ( + AmazonNovaEmbeddingConfig as AmazonNovaEmbeddingConfig, + ) + from .llms.openai.openai import ( + OpenAIConfig as OpenAIConfig, + MistralEmbeddingConfig as MistralEmbeddingConfig, + ) + from .llms.openai.image_variations.transformation import ( + OpenAIImageVariationConfig as OpenAIImageVariationConfig, + ) + from .llms.deepgram.audio_transcription.transformation import ( + DeepgramAudioTranscriptionConfig as DeepgramAudioTranscriptionConfig, + ) + from .llms.topaz.image_variations.transformation import ( + TopazImageVariationConfig as TopazImageVariationConfig, + ) + from litellm.llms.openai.completion.transformation import ( + OpenAITextCompletionConfig as OpenAITextCompletionConfig, + ) from .llms.groq.chat.transformation import GroqChatConfig as GroqChatConfig - from .llms.bedrock_mantle.chat.transformation import BedrockMantleChatConfig as BedrockMantleChatConfig + from .llms.bedrock_mantle.chat.transformation import ( + BedrockMantleChatConfig as BedrockMantleChatConfig, + ) from .llms.a2a.chat.transformation import A2AConfig as A2AConfig - from .llms.voyage.embedding.transformation import VoyageEmbeddingConfig as VoyageEmbeddingConfig - from .llms.voyage.embedding.transformation_contextual import VoyageContextualEmbeddingConfig as VoyageContextualEmbeddingConfig - from .llms.infinity.embedding.transformation import InfinityEmbeddingConfig as InfinityEmbeddingConfig - from .llms.perplexity.embedding.transformation import PerplexityEmbeddingConfig as PerplexityEmbeddingConfig - from .llms.azure_ai.chat.transformation import AzureAIStudioConfig as AzureAIStudioConfig + from .llms.voyage.embedding.transformation import ( + VoyageEmbeddingConfig as VoyageEmbeddingConfig, + ) + from .llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig as VoyageContextualEmbeddingConfig, + ) + from .llms.infinity.embedding.transformation import ( + InfinityEmbeddingConfig as InfinityEmbeddingConfig, + ) + from .llms.perplexity.embedding.transformation import ( + PerplexityEmbeddingConfig as PerplexityEmbeddingConfig, + ) + from .llms.azure_ai.chat.transformation import ( + AzureAIStudioConfig as AzureAIStudioConfig, + ) from .llms.mistral.chat.transformation import MistralConfig as MistralConfig - from .llms.openai.responses.transformation import OpenAIResponsesAPIConfig as OpenAIResponsesAPIConfig - from .llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig as AzureOpenAIResponsesAPIConfig - from .llms.azure.responses.o_series_transformation import AzureOpenAIOSeriesResponsesAPIConfig as AzureOpenAIOSeriesResponsesAPIConfig - from .llms.xai.responses.transformation import XAIResponsesAPIConfig as XAIResponsesAPIConfig - from .llms.litellm_proxy.responses.transformation import LiteLLMProxyResponsesAPIConfig as LiteLLMProxyResponsesAPIConfig - from .llms.volcengine.responses.transformation import VolcEngineResponsesAPIConfig as VolcEngineResponsesAPIConfig - from .llms.manus.responses.transformation import ManusResponsesAPIConfig as ManusResponsesAPIConfig - from .llms.perplexity.responses.transformation import PerplexityResponsesConfig as PerplexityResponsesConfig - from .llms.databricks.responses.transformation import DatabricksResponsesAPIConfig as DatabricksResponsesAPIConfig - from .llms.openrouter.responses.transformation import OpenRouterResponsesAPIConfig as OpenRouterResponsesAPIConfig - from .llms.gemini.interactions.transformation import GoogleAIStudioInteractionsConfig as GoogleAIStudioInteractionsConfig - from .llms.openai.chat.o_series_transformation import OpenAIOSeriesConfig as OpenAIOSeriesConfig, OpenAIOSeriesConfig as OpenAIO1Config - from .llms.anthropic.skills.transformation import AnthropicSkillsConfig as AnthropicSkillsConfig - from .llms.base_llm.skills.transformation import BaseSkillsAPIConfig as BaseSkillsAPIConfig - from .llms.gradient_ai.chat.transformation import GradientAIConfig as GradientAIConfig + from .llms.openai.responses.transformation import ( + OpenAIResponsesAPIConfig as OpenAIResponsesAPIConfig, + ) + from .llms.azure.responses.transformation import ( + AzureOpenAIResponsesAPIConfig as AzureOpenAIResponsesAPIConfig, + ) + from .llms.azure.responses.o_series_transformation import ( + AzureOpenAIOSeriesResponsesAPIConfig as AzureOpenAIOSeriesResponsesAPIConfig, + ) + from .llms.xai.responses.transformation import ( + XAIResponsesAPIConfig as XAIResponsesAPIConfig, + ) + from .llms.litellm_proxy.responses.transformation import ( + LiteLLMProxyResponsesAPIConfig as LiteLLMProxyResponsesAPIConfig, + ) + from .llms.volcengine.responses.transformation import ( + VolcEngineResponsesAPIConfig as VolcEngineResponsesAPIConfig, + ) + from .llms.manus.responses.transformation import ( + ManusResponsesAPIConfig as ManusResponsesAPIConfig, + ) + from .llms.perplexity.responses.transformation import ( + PerplexityResponsesConfig as PerplexityResponsesConfig, + ) + from .llms.databricks.responses.transformation import ( + DatabricksResponsesAPIConfig as DatabricksResponsesAPIConfig, + ) + from .llms.openrouter.responses.transformation import ( + OpenRouterResponsesAPIConfig as OpenRouterResponsesAPIConfig, + ) + from .llms.gemini.interactions.transformation import ( + GoogleAIStudioInteractionsConfig as GoogleAIStudioInteractionsConfig, + ) + from .llms.openai.chat.o_series_transformation import ( + OpenAIOSeriesConfig as OpenAIOSeriesConfig, + OpenAIOSeriesConfig as OpenAIO1Config, + ) + from .llms.anthropic.skills.transformation import ( + AnthropicSkillsConfig as AnthropicSkillsConfig, + ) + from .llms.base_llm.skills.transformation import ( + BaseSkillsAPIConfig as BaseSkillsAPIConfig, + ) + from .llms.gradient_ai.chat.transformation import ( + GradientAIConfig as GradientAIConfig, + ) from .llms.openai.chat.gpt_transformation import OpenAIGPTConfig as OpenAIGPTConfig - from .llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config as OpenAIGPT5Config - from .llms.openai.transcriptions.whisper_transformation import OpenAIWhisperAudioTranscriptionConfig as OpenAIWhisperAudioTranscriptionConfig - from .llms.openai.transcriptions.gpt_transformation import OpenAIGPTAudioTranscriptionConfig as OpenAIGPTAudioTranscriptionConfig - from .llms.openai.chat.gpt_audio_transformation import OpenAIGPTAudioConfig as OpenAIGPTAudioConfig + from .llms.openai.chat.gpt_5_transformation import ( + OpenAIGPT5Config as OpenAIGPT5Config, + ) + from .llms.openai.transcriptions.whisper_transformation import ( + OpenAIWhisperAudioTranscriptionConfig as OpenAIWhisperAudioTranscriptionConfig, + ) + from .llms.openai.transcriptions.gpt_transformation import ( + OpenAIGPTAudioTranscriptionConfig as OpenAIGPTAudioTranscriptionConfig, + ) + from .llms.openai.chat.gpt_audio_transformation import ( + OpenAIGPTAudioConfig as OpenAIGPTAudioConfig, + ) from .llms.nvidia_nim.chat.transformation import NvidiaNimConfig as NvidiaNimConfig - from .llms.nvidia_nim.embed import NvidiaNimEmbeddingConfig as NvidiaNimEmbeddingConfig + from .llms.nvidia_nim.embed import ( + NvidiaNimEmbeddingConfig as NvidiaNimEmbeddingConfig, + ) # Type stubs for lazy-loaded config instances openaiOSeriesConfig: OpenAIOSeriesConfig @@ -1480,21 +1683,47 @@ if TYPE_CHECKING: # Import config classes that need type stubs (for mypy) - import with _ prefix to avoid circular reference from .llms.vllm.completion.transformation import VLLMConfig as _VLLMConfig - from .llms.deepseek.chat.transformation import DeepSeekChatConfig as _DeepSeekChatConfig - from .llms.sap.chat.transformation import GenAIHubOrchestrationConfig as _GenAIHubOrchestrationConfig - from .llms.sap.embed.transformation import GenAIHubEmbeddingConfig as _GenAIHubEmbeddingConfig - from .llms.azure.chat.o_series_transformation import AzureOpenAIO1Config as _AzureOpenAIO1Config - from .llms.perplexity.chat.transformation import PerplexityChatConfig as _PerplexityChatConfig + from .llms.deepseek.chat.transformation import ( + DeepSeekChatConfig as _DeepSeekChatConfig, + ) + from .llms.sap.chat.transformation import ( + GenAIHubOrchestrationConfig as _GenAIHubOrchestrationConfig, + ) + from .llms.sap.embed.transformation import ( + GenAIHubEmbeddingConfig as _GenAIHubEmbeddingConfig, + ) + from .llms.azure.chat.o_series_transformation import ( + AzureOpenAIO1Config as _AzureOpenAIO1Config, + ) + from .llms.perplexity.chat.transformation import ( + PerplexityChatConfig as _PerplexityChatConfig, + ) from .llms.nscale.chat.transformation import NscaleConfig as _NscaleConfig - from .llms.watsonx.chat.transformation import IBMWatsonXChatConfig as _IBMWatsonXChatConfig - from .llms.watsonx.completion.transformation import IBMWatsonXAIConfig as _IBMWatsonXAIConfig - from .llms.litellm_proxy.chat.transformation import LiteLLMProxyChatConfig as _LiteLLMProxyChatConfig + from .llms.watsonx.chat.transformation import ( + IBMWatsonXChatConfig as _IBMWatsonXChatConfig, + ) + from .llms.watsonx.completion.transformation import ( + IBMWatsonXAIConfig as _IBMWatsonXAIConfig, + ) + from .llms.litellm_proxy.chat.transformation import ( + LiteLLMProxyChatConfig as _LiteLLMProxyChatConfig, + ) from .llms.deepinfra.chat.transformation import DeepInfraConfig as _DeepInfraConfig - from .llms.llamafile.chat.transformation import LlamafileChatConfig as _LlamafileChatConfig - from .llms.lm_studio.chat.transformation import LMStudioChatConfig as _LMStudioChatConfig - from .llms.lm_studio.embed.transformation import LmStudioEmbeddingConfig as _LmStudioEmbeddingConfig - from .llms.watsonx.embed.transformation import IBMWatsonXEmbeddingConfig as _IBMWatsonXEmbeddingConfig - from .llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig as _VertexGeminiConfig + from .llms.llamafile.chat.transformation import ( + LlamafileChatConfig as _LlamafileChatConfig, + ) + from .llms.lm_studio.chat.transformation import ( + LMStudioChatConfig as _LMStudioChatConfig, + ) + from .llms.lm_studio.embed.transformation import ( + LmStudioEmbeddingConfig as _LmStudioEmbeddingConfig, + ) + from .llms.watsonx.embed.transformation import ( + IBMWatsonXEmbeddingConfig as _IBMWatsonXEmbeddingConfig, + ) + from .llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig as _VertexGeminiConfig, + ) # Type stubs for lazy-loaded config classes (to help mypy understand types) VLLMConfig: Type[_VLLMConfig] @@ -1514,56 +1743,125 @@ if TYPE_CHECKING: IBMWatsonXEmbeddingConfig: Type[_IBMWatsonXEmbeddingConfig] VertexAIConfig: Type[_VertexGeminiConfig] # Alias for VertexGeminiConfig - from .llms.featherless_ai.chat.transformation import FeatherlessAIConfig as FeatherlessAIConfig + from .llms.featherless_ai.chat.transformation import ( + FeatherlessAIConfig as FeatherlessAIConfig, + ) from .llms.cerebras.chat import CerebrasConfig as CerebrasConfig from .llms.baseten.chat import BasetenConfig as BasetenConfig from .llms.sambanova.chat import SambanovaConfig as SambanovaConfig - from .llms.sambanova.embedding.transformation import SambaNovaEmbeddingConfig as SambaNovaEmbeddingConfig - from .llms.fireworks_ai.chat.transformation import FireworksAIConfig as FireworksAIConfig - from .llms.fireworks_ai.completion.transformation import FireworksAITextCompletionConfig as FireworksAITextCompletionConfig - from .llms.fireworks_ai.audio_transcription.transformation import FireworksAIAudioTranscriptionConfig as FireworksAIAudioTranscriptionConfig - from .llms.fireworks_ai.embed.fireworks_ai_transformation import FireworksAIEmbeddingConfig as FireworksAIEmbeddingConfig - from .llms.friendliai.chat.transformation import FriendliaiChatConfig as FriendliaiChatConfig - from .llms.jina_ai.embedding.transformation import JinaAIEmbeddingConfig as JinaAIEmbeddingConfig + from .llms.sambanova.embedding.transformation import ( + SambaNovaEmbeddingConfig as SambaNovaEmbeddingConfig, + ) + from .llms.fireworks_ai.chat.transformation import ( + FireworksAIConfig as FireworksAIConfig, + ) + from .llms.fireworks_ai.completion.transformation import ( + FireworksAITextCompletionConfig as FireworksAITextCompletionConfig, + ) + from .llms.fireworks_ai.audio_transcription.transformation import ( + FireworksAIAudioTranscriptionConfig as FireworksAIAudioTranscriptionConfig, + ) + from .llms.fireworks_ai.embed.fireworks_ai_transformation import ( + FireworksAIEmbeddingConfig as FireworksAIEmbeddingConfig, + ) + from .llms.friendliai.chat.transformation import ( + FriendliaiChatConfig as FriendliaiChatConfig, + ) + from .llms.jina_ai.embedding.transformation import ( + JinaAIEmbeddingConfig as JinaAIEmbeddingConfig, + ) from .llms.xai.chat.transformation import XAIChatConfig as XAIChatConfig from .llms.zai.chat.transformation import ZAIChatConfig as ZAIChatConfig from .llms.aiml.chat.transformation import AIMLChatConfig as AIMLChatConfig - from .llms.volcengine.chat.transformation import VolcEngineChatConfig as VolcEngineChatConfig, VolcEngineChatConfig as VolcEngineConfig - from .llms.codestral.completion.transformation import CodestralTextCompletionConfig as CodestralTextCompletionConfig - from .llms.azure.azure import AzureOpenAIAssistantsAPIConfig as AzureOpenAIAssistantsAPIConfig + from .llms.volcengine.chat.transformation import ( + VolcEngineChatConfig as VolcEngineChatConfig, + VolcEngineChatConfig as VolcEngineConfig, + ) + from .llms.codestral.completion.transformation import ( + CodestralTextCompletionConfig as CodestralTextCompletionConfig, + ) + from .llms.azure.azure import ( + AzureOpenAIAssistantsAPIConfig as AzureOpenAIAssistantsAPIConfig, + ) from .llms.heroku.chat.transformation import HerokuChatConfig as HerokuChatConfig from .llms.cometapi.chat.transformation import CometAPIConfig as CometAPIConfig - from .llms.azure.chat.gpt_transformation import AzureOpenAIConfig as AzureOpenAIConfig - from .llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config as AzureOpenAIGPT5Config - from .llms.azure.completion.transformation import AzureOpenAITextConfig as AzureOpenAITextConfig - from .llms.hosted_vllm.chat.transformation import HostedVLLMChatConfig as HostedVLLMChatConfig - from .llms.hosted_vllm.embedding.transformation import HostedVLLMEmbeddingConfig as HostedVLLMEmbeddingConfig - from .llms.hosted_vllm.responses.transformation import HostedVLLMResponsesAPIConfig as HostedVLLMResponsesAPIConfig - from .llms.github_copilot.chat.transformation import GithubCopilotConfig as GithubCopilotConfig - from .llms.github_copilot.responses.transformation import GithubCopilotResponsesAPIConfig as GithubCopilotResponsesAPIConfig - from .llms.github_copilot.embedding.transformation import GithubCopilotEmbeddingConfig as GithubCopilotEmbeddingConfig + from .llms.azure.chat.gpt_transformation import ( + AzureOpenAIConfig as AzureOpenAIConfig, + ) + from .llms.azure.chat.gpt_5_transformation import ( + AzureOpenAIGPT5Config as AzureOpenAIGPT5Config, + ) + from .llms.azure.completion.transformation import ( + AzureOpenAITextConfig as AzureOpenAITextConfig, + ) + from .llms.hosted_vllm.chat.transformation import ( + HostedVLLMChatConfig as HostedVLLMChatConfig, + ) + from .llms.hosted_vllm.embedding.transformation import ( + HostedVLLMEmbeddingConfig as HostedVLLMEmbeddingConfig, + ) + from .llms.hosted_vllm.responses.transformation import ( + HostedVLLMResponsesAPIConfig as HostedVLLMResponsesAPIConfig, + ) + from .llms.github_copilot.chat.transformation import ( + GithubCopilotConfig as GithubCopilotConfig, + ) + from .llms.github_copilot.responses.transformation import ( + GithubCopilotResponsesAPIConfig as GithubCopilotResponsesAPIConfig, + ) + from .llms.github_copilot.embedding.transformation import ( + GithubCopilotEmbeddingConfig as GithubCopilotEmbeddingConfig, + ) from .llms.chatgpt.chat.transformation import ChatGPTConfig as ChatGPTConfig - from .llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig as ChatGPTResponsesAPIConfig + from .llms.chatgpt.responses.transformation import ( + ChatGPTResponsesAPIConfig as ChatGPTResponsesAPIConfig, + ) from .llms.gigachat.chat.transformation import GigaChatConfig as GigaChatConfig - from .llms.gigachat.embedding.transformation import GigaChatEmbeddingConfig as GigaChatEmbeddingConfig + from .llms.gigachat.embedding.transformation import ( + GigaChatEmbeddingConfig as GigaChatEmbeddingConfig, + ) from .llms.nebius.chat.transformation import NebiusConfig as NebiusConfig from .llms.wandb.chat.transformation import WandbConfig as WandbConfig - from .llms.dashscope.chat.transformation import DashScopeChatConfig as DashScopeChatConfig - from .llms.moonshot.chat.transformation import MoonshotChatConfig as MoonshotChatConfig - from .llms.docker_model_runner.chat.transformation import DockerModelRunnerChatConfig as DockerModelRunnerChatConfig + from .llms.dashscope.chat.transformation import ( + DashScopeChatConfig as DashScopeChatConfig, + ) + from .llms.moonshot.chat.transformation import ( + MoonshotChatConfig as MoonshotChatConfig, + ) + from .llms.docker_model_runner.chat.transformation import ( + DockerModelRunnerChatConfig as DockerModelRunnerChatConfig, + ) from .llms.v0.chat.transformation import V0ChatConfig as V0ChatConfig from .llms.oci.chat.transformation import OCIChatConfig as OCIChatConfig from .llms.morph.chat.transformation import MorphChatConfig as MorphChatConfig from .llms.ragflow.chat.transformation import RAGFlowConfig as RAGFlowConfig - from .llms.lambda_ai.chat.transformation import LambdaAIChatConfig as LambdaAIChatConfig - from .llms.hyperbolic.chat.transformation import HyperbolicChatConfig as HyperbolicChatConfig - from .llms.vercel_ai_gateway.chat.transformation import VercelAIGatewayConfig as VercelAIGatewayConfig - from .llms.ovhcloud.chat.transformation import OVHCloudChatConfig as OVHCloudChatConfig - from .llms.ovhcloud.embedding.transformation import OVHCloudEmbeddingConfig as OVHCloudEmbeddingConfig - from .llms.cometapi.embed.transformation import CometAPIEmbeddingConfig as CometAPIEmbeddingConfig - from .llms.lemonade.chat.transformation import LemonadeChatConfig as LemonadeChatConfig - from .llms.snowflake.embedding.transformation import SnowflakeEmbeddingConfig as SnowflakeEmbeddingConfig - from .llms.amazon_nova.chat.transformation import AmazonNovaChatConfig as AmazonNovaChatConfig + from .llms.lambda_ai.chat.transformation import ( + LambdaAIChatConfig as LambdaAIChatConfig, + ) + from .llms.hyperbolic.chat.transformation import ( + HyperbolicChatConfig as HyperbolicChatConfig, + ) + from .llms.vercel_ai_gateway.chat.transformation import ( + VercelAIGatewayConfig as VercelAIGatewayConfig, + ) + from .llms.ovhcloud.chat.transformation import ( + OVHCloudChatConfig as OVHCloudChatConfig, + ) + from .llms.ovhcloud.embedding.transformation import ( + OVHCloudEmbeddingConfig as OVHCloudEmbeddingConfig, + ) + from .llms.cometapi.embed.transformation import ( + CometAPIEmbeddingConfig as CometAPIEmbeddingConfig, + ) + from .llms.lemonade.chat.transformation import ( + LemonadeChatConfig as LemonadeChatConfig, + ) + from .llms.snowflake.embedding.transformation import ( + SnowflakeEmbeddingConfig as SnowflakeEmbeddingConfig, + ) + from .llms.amazon_nova.chat.transformation import ( + AmazonNovaChatConfig as AmazonNovaChatConfig, + ) from litellm.caching.llm_caching_handler import LLMClientCache from litellm.types.llms.bedrock import COHERE_EMBEDDING_INPUT_TYPES from litellm.types.utils import ( @@ -1624,6 +1922,7 @@ if TYPE_CHECKING: # Bedrock tool name mappings instance (lazy-loaded) from litellm.caching.caching import InMemoryCache + bedrock_tool_name_mappings: InMemoryCache # Azure exception class (lazy-loaded) @@ -1642,11 +1941,15 @@ if TYPE_CHECKING: from litellm.types.integrations.datadog_llm_obs import DatadogLLMObsInitParams # Logging callback manager class and instance (lazy-loaded) - from litellm.litellm_core_utils.logging_callback_manager import LoggingCallbackManager + from litellm.litellm_core_utils.logging_callback_manager import ( + LoggingCallbackManager, + ) + logging_callback_manager: LoggingCallbackManager # provider_list is lazy-loaded from litellm.types.utils import LlmProviders + provider_list: List[Union[LlmProviders, str]] # Note: AmazonConverseConfig and OpenAILikeChatConfig are imported above in TYPE_CHECKING block @@ -1671,7 +1974,10 @@ def __getattr__(name: str) -> Any: global _async_client_cleanup_registered # Register async client cleanup on first access (only once) if not _async_client_cleanup_registered: - from litellm.llms.custom_httpx.async_client_cleanup import register_async_client_cleanup + from litellm.llms.custom_httpx.async_client_cleanup import ( + register_async_client_cleanup, + ) + register_async_client_cleanup() _async_client_cleanup_registered = True @@ -1688,36 +1994,45 @@ def __getattr__(name: str) -> Any: # Lazy load encoding from main.py to avoid heavy tiktoken import if name == "encoding": from ._lazy_imports import _get_litellm_globals + _globals = _get_litellm_globals() # Check if already cached if "encoding" not in _globals: from .main import encoding as _encoding + _globals["encoding"] = _encoding return _globals["encoding"] # Lazy load bedrock_tool_name_mappings instance if name == "bedrock_tool_name_mappings": from ._lazy_imports import _get_litellm_globals + _globals = _get_litellm_globals() # Check if already cached if "bedrock_tool_name_mappings" not in _globals: - from .llms.bedrock.chat.invoke_handler import bedrock_tool_name_mappings as _bedrock_tool_name_mappings + from .llms.bedrock.chat.invoke_handler import ( + bedrock_tool_name_mappings as _bedrock_tool_name_mappings, + ) + _globals["bedrock_tool_name_mappings"] = _bedrock_tool_name_mappings return _globals["bedrock_tool_name_mappings"] # Lazy load AzureOpenAIError exception class if name == "AzureOpenAIError": from ._lazy_imports import _get_litellm_globals + _globals = _get_litellm_globals() # Check if already cached if "AzureOpenAIError" not in _globals: from .llms.azure.common_utils import AzureOpenAIError as _AzureOpenAIError + _globals["AzureOpenAIError"] = _AzureOpenAIError return _globals["AzureOpenAIError"] # Lazy load openaiOSeriesConfig instance if name == "openaiOSeriesConfig": from ._lazy_imports import _get_litellm_globals + _globals = _get_litellm_globals() if "openaiOSeriesConfig" not in _globals: # Import the config class and instantiate it @@ -1735,6 +2050,7 @@ def __getattr__(name: str) -> Any: } if name in _config_instances: from ._lazy_imports import _get_litellm_globals + _globals = _get_litellm_globals() if name not in _globals: # Import the config class and instantiate it @@ -1749,17 +2065,20 @@ def __getattr__(name: str) -> Any: # Lazy load provider_list if name == "provider_list": from ._lazy_imports import _get_litellm_globals + _globals = _get_litellm_globals() # Check if already cached if "provider_list" not in _globals: # LlmProviders is eagerly imported above, so we can import it directly from litellm.types.utils import LlmProviders + _globals["provider_list"] = list(LlmProviders) return _globals["provider_list"] # Lazy load priority_reservation_settings instance if name == "priority_reservation_settings": from ._lazy_imports import _get_litellm_globals + _globals = _get_litellm_globals() # Check if already cached if "priority_reservation_settings" not in _globals: @@ -1771,6 +2090,7 @@ def __getattr__(name: str) -> Any: # Lazy load logging_callback_manager instance if name == "logging_callback_manager": from ._lazy_imports import _get_litellm_globals + _globals = _get_litellm_globals() # Check if already cached if "logging_callback_manager" not in _globals: @@ -1782,19 +2102,41 @@ def __getattr__(name: str) -> Any: # Lazy load _service_logger module if name == "_service_logger": from ._lazy_imports import _get_litellm_globals + _globals = _get_litellm_globals() # Check if already cached if "_service_logger" not in _globals: # Import the module lazily import litellm._service_logger + _globals["_service_logger"] = litellm._service_logger return _globals["_service_logger"] # Lazy load evals module functions - if name in ["acreate_eval", "alist_evals", "aget_eval", "aupdate_eval", "adelete_eval", "acancel_eval", - "create_eval", "list_evals", "get_eval", "update_eval", "delete_eval", "cancel_eval", - "acreate_run", "alist_runs", "aget_run", "acancel_run", "adelete_run", - "create_run", "list_runs", "get_run", "cancel_run", "delete_run"]: + if name in [ + "acreate_eval", + "alist_evals", + "aget_eval", + "aupdate_eval", + "adelete_eval", + "acancel_eval", + "create_eval", + "list_evals", + "get_eval", + "update_eval", + "delete_eval", + "cancel_eval", + "acreate_run", + "alist_runs", + "aget_run", + "acancel_run", + "adelete_run", + "create_run", + "list_runs", + "get_run", + "cancel_run", + "delete_run", + ]: from litellm.evals.main import ( acreate_eval, alist_evals, @@ -1819,6 +2161,7 @@ def __getattr__(name: str) -> Any: cancel_run, delete_run, ) + return locals()[name] raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/litellm/_lazy_imports.py b/litellm/_lazy_imports.py index 3bfeba2e39..3604506d40 100644 --- a/litellm/_lazy_imports.py +++ b/litellm/_lazy_imports.py @@ -55,7 +55,7 @@ from ._lazy_imports_registry import ( def _get_litellm_globals() -> dict: """ Get the globals dictionary of the litellm module. - + This is where we cache imported attributes so we don't import them twice. When you do `litellm.some_function`, it gets stored in this dictionary. """ @@ -65,12 +65,13 @@ def _get_litellm_globals() -> dict: def _get_utils_globals() -> dict: """ Get the globals dictionary of the utils module. - + This is where we cache imported attributes so we don't import them twice. When you do `litellm.utils.some_function`, it gets stored in this dictionary. """ return sys.modules["litellm.utils"].__dict__ + # These are special lazy loaders for things that are used internally # They're separate from the main lazy import system because they have specific use cases @@ -81,10 +82,10 @@ _default_encoding: Optional[Any] = None def _get_default_encoding() -> Any: """ Lazily load and cache the default OpenAI encoding. - + This avoids importing `litellm.litellm_core_utils.default_encoding` (and thus tiktoken) at `litellm` import time. The encoding is cached after the first import. - + This is used internally by utils.py functions that need the encoding but shouldn't trigger its import during module load. """ @@ -103,10 +104,10 @@ _get_modified_max_tokens_func: Optional[Any] = None def _get_modified_max_tokens() -> Any: """ Lazily load and cache the get_modified_max_tokens function. - + This avoids importing `litellm.litellm_core_utils.token_counter` at `litellm` import time. The function is cached after the first import. - + This is used internally by utils.py functions that need the token counter but shouldn't trigger its import during module load. """ @@ -127,10 +128,10 @@ _token_counter_new_func: Optional[Any] = None def _get_token_counter_new() -> Any: """ Lazily load and cache the token_counter function (aliased as token_counter_new). - + This avoids importing `litellm.litellm_core_utils.token_counter` at `litellm` import time. The function is cached after the first import. - + This is used internally by utils.py functions that need the token counter but shouldn't trigger its import during module load. """ @@ -157,10 +158,10 @@ _LAZY_IMPORT_REGISTRY: Optional[dict[str, Callable[[str], Any]]] = None def _get_lazy_import_registry() -> dict[str, Callable[[str], Any]]: """ Build the registry that maps attribute names to their handler functions. - + This is called once, the first time someone accesses a lazy-loaded attribute. After that, we just look up the handler function in this dictionary. - + Returns: Dictionary like {"ModelResponse": _lazy_import_utils, ...} """ @@ -199,17 +200,19 @@ def _get_lazy_import_registry() -> dict[str, Callable[[str], Any]]: _LAZY_IMPORT_REGISTRY[name] = _lazy_import_llm_provider_logic for name in UTILS_MODULE_NAMES: _LAZY_IMPORT_REGISTRY[name] = _lazy_import_utils_module - + return _LAZY_IMPORT_REGISTRY -def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], category: str) -> Any: +def _generic_lazy_import( + name: str, import_map: dict[str, tuple[str, str]], category: str +) -> Any: """ Generic function that handles lazy importing for most attributes. - + This is the workhorse function - it does the actual importing and caching. Most handler functions just call this with their specific import map. - + Steps: 1. Check if the name exists in the import map (if not, raise error) 2. Check if we've already imported it (if yes, return cached value) @@ -218,7 +221,7 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate 5. Get the attribute from the module 6. Cache it in _globals so we don't import again 7. Return it - + Args: name: The attribute name someone is trying to access (e.g., "ModelResponse") import_map: Dictionary telling us where to find each attribute @@ -228,19 +231,19 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate # Step 1: Make sure this attribute exists in our map if name not in import_map: raise AttributeError(f"{category} lazy import: unknown attribute {name!r}") - + # Step 2: Get the cache (where we store imported things) _globals = _get_litellm_globals() - + # Step 3: If we've already imported it, just return the cached version if name in _globals: return _globals[name] - + # Step 4: Look up where to find this attribute # The map tells us: (module_path, attribute_name) # Example: (".utils", "ModelResponse") means "look in .utils module, get ModelResponse" module_path, attr_name = import_map[name] - + # Step 5: Import the module # Python automatically caches modules in sys.modules, so calling this twice is fast # If module_path starts with ".", it's a relative import (needs package="litellm") @@ -249,14 +252,14 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate module = importlib.import_module(module_path, package="litellm") else: module = importlib.import_module(module_path) - + # Step 6: Get the actual attribute from the module # Example: getattr(utils_module, "ModelResponse") returns the ModelResponse class value = getattr(module, attr_name) - + # Step 7: Cache it so we don't have to import again next time _globals[name] = value - + # Step 8: Return it return value @@ -268,6 +271,7 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate # Most of them just call _generic_lazy_import with their specific import map. # The registry (above) maps attribute names to these handler functions. + def _lazy_import_utils(name: str) -> Any: """Handler for utils module attributes (ModelResponse, token_counter, etc.)""" return _generic_lazy_import(name, _UTILS_IMPORT_MAP, "Utils") @@ -297,6 +301,7 @@ def _lazy_import_caching(name: str) -> Any: """Handler for caching classes (Cache, DualCache, RedisCache, etc.)""" return _generic_lazy_import(name, _CACHING_IMPORT_MAP, "Caching") + def _lazy_import_dotprompt(name: str) -> Any: """Handler for dotprompt integration globals""" return _generic_lazy_import(name, _DOTPROMPT_IMPORT_MAP, "Dotprompt") @@ -311,6 +316,7 @@ def _lazy_import_llm_configs(name: str) -> Any: """Handler for LLM config classes (AnthropicConfig, OpenAILikeChatConfig, etc.)""" return _generic_lazy_import(name, _LLM_CONFIGS_IMPORT_MAP, "LLM config") + def _lazy_import_litellm_logging(name: str) -> Any: """Handler for litellm_logging module (Logging, modify_integration)""" return _generic_lazy_import(name, _LITELLM_LOGGING_IMPORT_MAP, "Litellm logging") @@ -318,87 +324,91 @@ def _lazy_import_litellm_logging(name: str) -> Any: def _lazy_import_llm_provider_logic(name: str) -> Any: """Handler for LLM provider logic functions (get_llm_provider, etc.)""" - return _generic_lazy_import(name, _LLM_PROVIDER_LOGIC_IMPORT_MAP, "LLM provider logic") + return _generic_lazy_import( + name, _LLM_PROVIDER_LOGIC_IMPORT_MAP, "LLM provider logic" + ) def _lazy_import_utils_module(name: str) -> Any: """ Handler for utils module lazy imports. - + This uses a custom implementation because utils module needs to use _get_utils_globals() instead of _get_litellm_globals() for caching. """ # Check if this attribute exists in our map if name not in _UTILS_MODULE_IMPORT_MAP: raise AttributeError(f"Utils module lazy import: unknown attribute {name!r}") - + # Get the cache (where we store imported things) - use utils globals _globals = _get_utils_globals() - + # If we've already imported it, just return the cached version if name in _globals: return _globals[name] - + # Look up where to find this attribute module_path, attr_name = _UTILS_MODULE_IMPORT_MAP[name] - + # Import the module if module_path.startswith("."): module = importlib.import_module(module_path, package="litellm") else: module = importlib.import_module(module_path) - + # Get the actual attribute from the module value = getattr(module, attr_name) - + # Cache it so we don't have to import again next time _globals[name] = value - + # Return it return value + # ============================================================================ # SPECIAL HANDLERS # ============================================================================ # These handlers have custom logic that doesn't fit the generic pattern + def _lazy_import_llm_client_cache(name: str) -> Any: """ Handler for LLM client cache - has special logic for singleton instance. - + This one is different because: - "LLMClientCache" is the class itself - "in_memory_llm_clients_cache" is a singleton instance of that class So we need custom logic to handle both cases. """ _globals = _get_litellm_globals() - + # If already cached, return it if name in _globals: return _globals[name] - + # Import the class module = importlib.import_module("litellm.caching.llm_caching_handler") LLMClientCache = getattr(module, "LLMClientCache") - + # If they want the class itself, return it if name == "LLMClientCache": _globals["LLMClientCache"] = LLMClientCache return LLMClientCache - + # If they want the singleton instance, create it (only once) if name == "in_memory_llm_clients_cache": instance = LLMClientCache() _globals["in_memory_llm_clients_cache"] = instance return instance - + raise AttributeError(f"LLM client cache lazy import: unknown attribute {name!r}") def _lazy_import_http_handlers(name: str) -> Any: """ Handler for HTTP clients - has special logic for creating client instances. - + This one is different because: - These aren't just imports, they're actual client instances that need to be created - They need configuration (timeout, etc.) from the module globals @@ -413,14 +423,14 @@ def _lazy_import_http_handlers(name: str) -> Any: # Get timeout from module config (if set) timeout = _globals.get("request_timeout") params = {"timeout": timeout, "client_alias": "module level aclient"} - + # Create the client instance provider_id = cast(Any, "litellm_module_level_client") async_client = get_async_httpx_client( llm_provider=provider_id, params=params, ) - + # Cache it so we don't create it again _globals["module_level_aclient"] = async_client return async_client @@ -431,7 +441,7 @@ def _lazy_import_http_handlers(name: str) -> Any: timeout = _globals.get("request_timeout") sync_client = HTTPHandler(timeout=timeout) - + # Cache it _globals["module_level_client"] = sync_client return sync_client diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index 9e0453102d..f7f56d2b88 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -677,7 +677,10 @@ _LLM_CONFIGS_IMPORT_MAP = { "FireworksAIRerankConfig", ), "VoyageRerankConfig": (".llms.voyage.rerank.transformation", "VoyageRerankConfig"), - "IBMWatsonXRerankConfig": (".llms.watsonx.rerank.transformation", "IBMWatsonXRerankConfig"), + "IBMWatsonXRerankConfig": ( + ".llms.watsonx.rerank.transformation", + "IBMWatsonXRerankConfig", + ), "ClarifaiConfig": (".llms.clarifai.chat.transformation", "ClarifaiConfig"), "AI21ChatConfig": (".llms.ai21.chat.transformation", "AI21ChatConfig"), "LlamaAPIConfig": (".llms.meta_llama.chat.transformation", "LlamaAPIConfig"), @@ -859,7 +862,10 @@ _LLM_CONFIGS_IMPORT_MAP = { "OpenAITextCompletionConfig", ), "GroqChatConfig": (".llms.groq.chat.transformation", "GroqChatConfig"), - "BedrockMantleChatConfig": (".llms.bedrock_mantle.chat.transformation", "BedrockMantleChatConfig"), + "BedrockMantleChatConfig": ( + ".llms.bedrock_mantle.chat.transformation", + "BedrockMantleChatConfig", + ), "A2AConfig": (".llms.a2a.chat.transformation", "A2AConfig"), "GenAIHubOrchestrationConfig": ( ".llms.sap.chat.transformation", diff --git a/litellm/_redis.py b/litellm/_redis.py index c61582abd1..b754c1f433 100644 --- a/litellm/_redis.py +++ b/litellm/_redis.py @@ -34,7 +34,12 @@ def _get_redis_kwargs(): "retry", } - include_args = ["url", "redis_connect_func", "gcp_service_account", "gcp_ssl_ca_certs"] + include_args = [ + "url", + "redis_connect_func", + "gcp_service_account", + "gcp_ssl_ca_certs", + ] available_args = [x for x in arg_spec.args if x not in exclude_args] + include_args @@ -75,7 +80,9 @@ def _get_redis_cluster_kwargs(client=None): available_args.append("ssl_cert_reqs") available_args.append("ssl_check_hostname") available_args.append("ssl_ca_certs") - available_args.append("redis_connect_func") # Needed for sync clusters and IAM detection + available_args.append( + "redis_connect_func" + ) # Needed for sync clusters and IAM detection available_args.append("gcp_service_account") available_args.append("gcp_ssl_ca_certs") available_args.append("max_connections") @@ -103,10 +110,10 @@ def _redis_kwargs_from_environment(): def _generate_gcp_iam_access_token(service_account: str) -> str: """ Generate GCP IAM access token for Redis authentication. - + Args: service_account: GCP service account in format 'projects/-/serviceAccounts/name@project.iam.gserviceaccount.com' - + Returns: Access token string for GCP IAM authentication """ @@ -117,11 +124,11 @@ def _generate_gcp_iam_access_token(service_account: str) -> str: "google-cloud-iam is required for GCP IAM Redis authentication. " "Install it with: pip install google-cloud-iam" ) - + client = iam_credentials_v1.IAMCredentialsClient() request = iam_credentials_v1.GenerateAccessTokenRequest( name=service_account, - scope=['https://www.googleapis.com/auth/cloud-platform'], + scope=["https://www.googleapis.com/auth/cloud-platform"], ) response = client.generate_access_token(request=request) return str(response.access_token) @@ -133,14 +140,15 @@ def create_gcp_iam_redis_connect_func( ) -> Callable: """ Creates a custom Redis connection function for GCP IAM authentication. - + Args: service_account: GCP service account in format 'projects/-/serviceAccounts/name@project.iam.gserviceaccount.com' ssl_ca_certs: Path to SSL CA certificate file for secure connections - + Returns: A connection function that can be used with Redis clients """ + def iam_connect(self): """Initialize the connection and authenticate using GCP IAM""" from redis.exceptions import ( @@ -148,25 +156,25 @@ def create_gcp_iam_redis_connect_func( AuthenticationWrongNumberOfArgsError, ) from redis.utils import str_if_bytes - + self._parser.on_connect(self) - + auth_args = (_generate_gcp_iam_access_token(service_account),) self.send_command("AUTH", *auth_args, check_health=False) - + try: auth_response = self.read_response() except AuthenticationWrongNumberOfArgsError: # Fallback to password auth if IAM fails - if hasattr(self, 'password') and self.password: + if hasattr(self, "password") and self.password: self.send_command("AUTH", self.password, check_health=False) auth_response = self.read_response() else: raise - + if str_if_bytes(auth_response) != "OK": raise AuthenticationError("GCP IAM authentication failed") - + return iam_connect @@ -178,22 +186,20 @@ def get_redis_url_from_environment(): raise ValueError( "Either 'REDIS_URL' or both 'REDIS_HOST' and 'REDIS_PORT' must be specified for Redis." ) - + if "REDIS_SSL" in os.environ and os.environ["REDIS_SSL"].lower() == "true": redis_protocol = "rediss" else: redis_protocol = "redis" - + # Build authentication part of URL auth_part = "" if "REDIS_USERNAME" in os.environ and "REDIS_PASSWORD" in os.environ: auth_part = f"{os.environ['REDIS_USERNAME']}:{os.environ['REDIS_PASSWORD']}@" elif "REDIS_PASSWORD" in os.environ: auth_part = f"{os.environ['REDIS_PASSWORD']}@" - - return ( - f"{redis_protocol}://{auth_part}{os.environ['REDIS_HOST']}:{os.environ['REDIS_PORT']}" - ) + + return f"{redis_protocol}://{auth_part}{os.environ['REDIS_HOST']}:{os.environ['REDIS_PORT']}" def _get_redis_client_logic(**env_overrides): @@ -241,22 +247,27 @@ def _get_redis_client_logic(**env_overrides): redis_kwargs["service_name"] = _service_name # Handle GCP IAM authentication - _gcp_service_account = redis_kwargs.get("gcp_service_account") or get_secret_str("REDIS_GCP_SERVICE_ACCOUNT") - _gcp_ssl_ca_certs = redis_kwargs.get("gcp_ssl_ca_certs") or get_secret_str("REDIS_GCP_SSL_CA_CERTS") - + _gcp_service_account = redis_kwargs.get("gcp_service_account") or get_secret_str( + "REDIS_GCP_SERVICE_ACCOUNT" + ) + _gcp_ssl_ca_certs = redis_kwargs.get("gcp_ssl_ca_certs") or get_secret_str( + "REDIS_GCP_SSL_CA_CERTS" + ) + if _gcp_service_account is not None: - verbose_logger.debug("Setting up GCP IAM authentication for Redis with service account.") + verbose_logger.debug( + "Setting up GCP IAM authentication for Redis with service account." + ) redis_kwargs["redis_connect_func"] = create_gcp_iam_redis_connect_func( - service_account=_gcp_service_account, - ssl_ca_certs=_gcp_ssl_ca_certs + service_account=_gcp_service_account, ssl_ca_certs=_gcp_ssl_ca_certs ) # Store GCP service account in redis_connect_func for async cluster access redis_kwargs["redis_connect_func"]._gcp_service_account = _gcp_service_account - + # Remove GCP-specific kwargs that shouldn't be passed to Redis client redis_kwargs.pop("gcp_service_account", None) redis_kwargs.pop("gcp_ssl_ca_certs", None) - + # Only enable SSL if explicitly requested AND SSL CA certs are provided if _gcp_ssl_ca_certs and redis_kwargs.get("ssl", False): redis_kwargs["ssl_ca_certs"] = _gcp_ssl_ca_certs @@ -377,7 +388,8 @@ def get_redis_client(**env_overrides): def get_redis_async_client( - connection_pool: Optional[async_redis.BlockingConnectionPool] = None, **env_overrides, + connection_pool: Optional[async_redis.BlockingConnectionPool] = None, + **env_overrides, ) -> Union[async_redis.Redis, async_redis.RedisCluster]: redis_kwargs = _get_redis_client_logic(**env_overrides) if "url" in redis_kwargs and redis_kwargs["url"] is not None: @@ -411,39 +423,50 @@ def get_redis_async_client( # Get GCP service account - first try from redis_connect_func, then from environment gcp_service_account = None - if redis_connect_func and hasattr(redis_connect_func, '_gcp_service_account'): + if redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"): gcp_service_account = redis_connect_func._gcp_service_account else: - gcp_service_account = redis_kwargs.get("gcp_service_account") or get_secret_str("REDIS_GCP_SERVICE_ACCOUNT") - - verbose_logger.debug(f"DEBUG: Redis cluster kwargs: redis_connect_func={redis_connect_func is not None}, gcp_service_account_provided={gcp_service_account is not None}") - + gcp_service_account = redis_kwargs.get( + "gcp_service_account" + ) or get_secret_str("REDIS_GCP_SERVICE_ACCOUNT") + + verbose_logger.debug( + f"DEBUG: Redis cluster kwargs: redis_connect_func={redis_connect_func is not None}, gcp_service_account_provided={gcp_service_account is not None}" + ) + # If GCP IAM is configured (indicated by redis_connect_func), generate access token and use as password if redis_connect_func and gcp_service_account: - verbose_logger.debug("DEBUG: Generating IAM token for service account (value not logged for security reasons)") + verbose_logger.debug( + "DEBUG: Generating IAM token for service account (value not logged for security reasons)" + ) try: # Generate IAM access token using the helper function access_token = _generate_gcp_iam_access_token(gcp_service_account) cluster_kwargs["password"] = access_token - verbose_logger.debug("DEBUG: Successfully generated GCP IAM access token for async Redis cluster") + verbose_logger.debug( + "DEBUG: Successfully generated GCP IAM access token for async Redis cluster" + ) except Exception as e: verbose_logger.error(f"Failed to generate GCP IAM access token: {e}") from redis.exceptions import AuthenticationError + raise AuthenticationError("Failed to generate GCP IAM access token") else: - verbose_logger.debug(f"DEBUG: Not using GCP IAM auth - redis_connect_func={redis_connect_func is not None}, gcp_service_account_provided={gcp_service_account is not None}") - + verbose_logger.debug( + f"DEBUG: Not using GCP IAM auth - redis_connect_func={redis_connect_func is not None}, gcp_service_account_provided={gcp_service_account is not None}" + ) + new_startup_nodes: List[ClusterNode] = [] for item in redis_kwargs["startup_nodes"]: new_startup_nodes.append(ClusterNode(**item)) cluster_kwargs.pop("startup_nodes", None) - + # Create async RedisCluster with IAM token as password if available cluster_client = async_redis.RedisCluster( startup_nodes=new_startup_nodes, **cluster_kwargs # type: ignore ) - + return cluster_client # Check for Redis Sentinel @@ -463,7 +486,10 @@ def get_redis_connection_pool(**env_overrides): redis_kwargs = _get_redis_client_logic(**env_overrides) verbose_logger.debug("get_redis_connection_pool: redis_kwargs", redis_kwargs) if "url" in redis_kwargs and redis_kwargs["url"] is not None: - pool_kwargs = {"timeout": REDIS_CONNECTION_POOL_TIMEOUT, "url": redis_kwargs["url"]} + pool_kwargs = { + "timeout": REDIS_CONNECTION_POOL_TIMEOUT, + "url": redis_kwargs["url"], + } if "max_connections" in redis_kwargs: try: pool_kwargs["max_connections"] = int(redis_kwargs["max_connections"]) @@ -483,6 +509,7 @@ def get_redis_connection_pool(**env_overrides): timeout=REDIS_CONNECTION_POOL_TIMEOUT, **redis_kwargs ) + def _pretty_print_redis_config(redis_kwargs: dict) -> None: """Pretty print the Redis configuration using rich with sensitive data masking""" try: @@ -492,6 +519,7 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None: from rich.panel import Panel from rich.table import Table from rich.text import Text + if not verbose_logger.isEnabledFor(logging.DEBUG): return @@ -499,7 +527,7 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None: # Initialize the sensitive data masker masker = SensitiveDataMasker() - + # Mask sensitive data in redis_kwargs masked_redis_kwargs = masker.mask_dict(redis_kwargs) @@ -531,7 +559,7 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None: value_str = str(value) else: value_str = str(value) - + config_table.add_row(key, value_str) # Determine connection type @@ -568,4 +596,3 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None: verbose_logger.info(f"Redis configuration: {masked_redis_kwargs}") except Exception as e: verbose_logger.error(f"Error pretty printing Redis configuration: {e}") - diff --git a/litellm/_service_logger.py b/litellm/_service_logger.py index 8f9a3c5083..1a3be203fe 100644 --- a/litellm/_service_logger.py +++ b/litellm/_service_logger.py @@ -317,7 +317,7 @@ class ServiceLogging(CustomLogger): await self.async_service_success_hook( service=ServiceTypes.LITELLM, duration=_duration, - call_type=kwargs.get("call_type", "unknown") + call_type=kwargs.get("call_type", "unknown"), ) except Exception as e: raise e diff --git a/litellm/a2a_protocol/client.py b/litellm/a2a_protocol/client.py index 31f7c3b6a9..05e21284af 100644 --- a/litellm/a2a_protocol/client.py +++ b/litellm/a2a_protocol/client.py @@ -103,5 +103,7 @@ class A2AClient: from litellm.a2a_protocol.main import asend_message_streaming a2a_client = await self._get_client() - async for chunk in asend_message_streaming(a2a_client=a2a_client, request=request): + async for chunk in asend_message_streaming( + a2a_client=a2a_client, request=request + ): yield chunk diff --git a/litellm/a2a_protocol/cost_calculator.py b/litellm/a2a_protocol/cost_calculator.py index f3e84c5b84..f64174f8be 100644 --- a/litellm/a2a_protocol/cost_calculator.py +++ b/litellm/a2a_protocol/cost_calculator.py @@ -97,7 +97,11 @@ class A2ACostCalculator: completion_tokens = getattr(usage, "completion_tokens", 0) or 0 # Calculate costs - input_cost = prompt_tokens * (float(input_cost_per_token) if input_cost_per_token else 0.0) - output_cost = completion_tokens * (float(output_cost_per_token) if output_cost_per_token else 0.0) + input_cost = prompt_tokens * ( + float(input_cost_per_token) if input_cost_per_token else 0.0 + ) + output_cost = completion_tokens * ( + float(output_cost_per_token) if output_cost_per_token else 0.0 + ) return input_cost + output_cost diff --git a/litellm/a2a_protocol/litellm_completion_bridge/handler.py b/litellm/a2a_protocol/litellm_completion_bridge/handler.py index 1916b04454..c3d2e41523 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/handler.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/handler.py @@ -50,30 +50,28 @@ class A2ACompletionBridgeHandler: a2a_provider_config = A2AProviderConfigManager.get_provider_config( custom_llm_provider=custom_llm_provider ) - + # If provider config exists, use it if a2a_provider_config is not None: if api_base is None: raise ValueError(f"api_base is required for {custom_llm_provider}") - - verbose_logger.info( - f"A2A: Using provider config for {custom_llm_provider}" - ) - + + verbose_logger.info(f"A2A: Using provider config for {custom_llm_provider}") + response_data = await a2a_provider_config.handle_non_streaming( request_id=request_id, params=params, api_base=api_base, ) - + return response_data - + # Extract message from params message = params.get("message", {}) # Transform A2A message to OpenAI format - openai_messages = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages( - message + openai_messages = ( + A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message) ) # Get completion params @@ -100,7 +98,8 @@ class A2ACompletionBridgeHandler: } # Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.) litellm_params_to_add = { - k: v for k, v in litellm_params.items() + k: v + for k, v in litellm_params.items() if k not in ("model", "custom_llm_provider") } completion_params.update(litellm_params_to_add) @@ -109,9 +108,11 @@ class A2ACompletionBridgeHandler: response = await litellm.acompletion(**completion_params) # Transform response to A2A format - a2a_response = A2ACompletionBridgeTransformation.openai_response_to_a2a_response( - response=response, - request_id=request_id, + a2a_response = ( + A2ACompletionBridgeTransformation.openai_response_to_a2a_response( + response=response, + request_id=request_id, + ) ) verbose_logger.info(f"A2A completion bridge completed: request_id={request_id}") @@ -148,25 +149,25 @@ class A2ACompletionBridgeHandler: a2a_provider_config = A2AProviderConfigManager.get_provider_config( custom_llm_provider=custom_llm_provider ) - + # If provider config exists, use it if a2a_provider_config is not None: if api_base is None: raise ValueError(f"api_base is required for {custom_llm_provider}") - + verbose_logger.info( f"A2A: Using provider config for {custom_llm_provider} (streaming)" ) - + async for chunk in a2a_provider_config.handle_streaming( request_id=request_id, params=params, api_base=api_base, ): yield chunk - + return - + # Extract message from params message = params.get("message", {}) @@ -177,8 +178,8 @@ class A2ACompletionBridgeHandler: ) # Transform A2A message to OpenAI format - openai_messages = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages( - message + openai_messages = ( + A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message) ) # Get completion params @@ -205,7 +206,8 @@ class A2ACompletionBridgeHandler: } # Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.) litellm_params_to_add = { - k: v for k, v in litellm_params.items() + k: v + for k, v in litellm_params.items() if k not in ("model", "custom_llm_provider") } completion_params.update(litellm_params_to_add) @@ -244,9 +246,11 @@ class A2ACompletionBridgeHandler: # Emit artifact update with accumulated content if accumulated_text: - artifact_event = A2ACompletionBridgeTransformation.create_artifact_update_event( - ctx=ctx, - text=accumulated_text, + artifact_event = ( + A2ACompletionBridgeTransformation.create_artifact_update_event( + ctx=ctx, + text=accumulated_text, + ) ) yield artifact_event diff --git a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py index bbe7daa9fc..8a03569f68 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py @@ -124,9 +124,7 @@ class A2ACompletionBridgeTransformation: }, } - verbose_logger.debug( - f"OpenAI -> A2A transform: content_length={len(content)}" - ) + verbose_logger.debug(f"OpenAI -> A2A transform: content_length={len(content)}") return a2a_response diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index 8cf477ee5e..c86549da77 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -120,9 +120,9 @@ def _get_a2a_model_info(a2a_client: Any, kwargs: Dict[str, Any]) -> str: litellm_logging_obj.model = model litellm_logging_obj.custom_llm_provider = custom_llm_provider litellm_logging_obj.model_call_details["model"] = model - litellm_logging_obj.model_call_details["custom_llm_provider"] = ( - custom_llm_provider - ) + litellm_logging_obj.model_call_details[ + "custom_llm_provider" + ] = custom_llm_provider return agent_name @@ -664,9 +664,7 @@ async def create_a2a_client( if extra_headers: # Encode headers into a cache-key-only param so each unique header # set produces a distinct cache key. - _client_params["disable_aiohttp_transport"] = str( - sorted(extra_headers.items()) - ) + _client_params["disable_aiohttp_transport"] = str(sorted(extra_headers.items())) _async_handler = get_async_httpx_client( llm_provider=httpxSpecialProvider.A2AProvider, params=_client_params, diff --git a/litellm/a2a_protocol/providers/__init__.py b/litellm/a2a_protocol/providers/__init__.py index 873a5a8374..a21fa5f8f5 100644 --- a/litellm/a2a_protocol/providers/__init__.py +++ b/litellm/a2a_protocol/providers/__init__.py @@ -8,4 +8,3 @@ from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig from litellm.a2a_protocol.providers.config_manager import A2AProviderConfigManager __all__ = ["BaseA2AProviderConfig", "A2AProviderConfigManager"] - diff --git a/litellm/a2a_protocol/providers/base.py b/litellm/a2a_protocol/providers/base.py index 9931076a94..a2354b3495 100644 --- a/litellm/a2a_protocol/providers/base.py +++ b/litellm/a2a_protocol/providers/base.py @@ -9,7 +9,7 @@ from typing import Any, AsyncIterator, Dict class BaseA2AProviderConfig(ABC): """ Base configuration class for A2A protocol providers. - + Each provider should implement this interface to define how to handle A2A requests for their specific agent type. """ @@ -60,4 +60,3 @@ class BaseA2AProviderConfig(ABC): # The yield is here to make this a generator function if False: # pragma: no cover yield {} - diff --git a/litellm/a2a_protocol/providers/config_manager.py b/litellm/a2a_protocol/providers/config_manager.py index e0703ec466..a8b9566c17 100644 --- a/litellm/a2a_protocol/providers/config_manager.py +++ b/litellm/a2a_protocol/providers/config_manager.py @@ -12,7 +12,7 @@ from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig class A2AProviderConfigManager: """ Manager for A2A provider configurations. - + Similar to ProviderConfigManager in litellm.utils but specifically for A2A providers. """ @@ -31,7 +31,7 @@ class A2AProviderConfigManager: """ if custom_llm_provider is None: return None - + if custom_llm_provider == "pydantic_ai_agents": from litellm.a2a_protocol.providers.pydantic_ai_agents.config import ( PydanticAIProviderConfig, @@ -45,4 +45,3 @@ class A2AProviderConfigManager: # return AnotherProviderConfig() return None - diff --git a/litellm/a2a_protocol/providers/litellm_completion/__init__.py b/litellm/a2a_protocol/providers/litellm_completion/__init__.py index 3f2b88bfaa..fc2fc17f54 100644 --- a/litellm/a2a_protocol/providers/litellm_completion/__init__.py +++ b/litellm/a2a_protocol/providers/litellm_completion/__init__.py @@ -3,4 +3,3 @@ LiteLLM Completion bridge provider for A2A protocol. Routes A2A requests through litellm.acompletion based on custom_llm_provider. """ - diff --git a/litellm/a2a_protocol/providers/litellm_completion/handler.py b/litellm/a2a_protocol/providers/litellm_completion/handler.py index 57388a5d0e..730f8f6b36 100644 --- a/litellm/a2a_protocol/providers/litellm_completion/handler.py +++ b/litellm/a2a_protocol/providers/litellm_completion/handler.py @@ -52,26 +52,26 @@ class A2ACompletionBridgeHandler: if custom_llm_provider == "pydantic_ai_agents": if api_base is None: raise ValueError("api_base is required for Pydantic AI agents") - + verbose_logger.info( f"Pydantic AI: Routing to Pydantic AI agent at {api_base}" ) - + # Send request directly to Pydantic AI agent response_data = await PydanticAITransformation.send_non_streaming_request( api_base=api_base, request_id=request_id, params=params, ) - + return response_data - + # Extract message from params message = params.get("message", {}) # Transform A2A message to OpenAI format - openai_messages = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages( - message + openai_messages = ( + A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message) ) # Get completion params @@ -98,7 +98,8 @@ class A2ACompletionBridgeHandler: } # Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.) litellm_params_to_add = { - k: v for k, v in litellm_params.items() + k: v + for k, v in litellm_params.items() if k not in ("model", "custom_llm_provider") } completion_params.update(litellm_params_to_add) @@ -107,9 +108,11 @@ class A2ACompletionBridgeHandler: response = await litellm.acompletion(**completion_params) # Transform response to A2A format - a2a_response = A2ACompletionBridgeTransformation.openai_response_to_a2a_response( - response=response, - request_id=request_id, + a2a_response = ( + A2ACompletionBridgeTransformation.openai_response_to_a2a_response( + response=response, + request_id=request_id, + ) ) verbose_logger.info(f"A2A completion bridge completed: request_id={request_id}") @@ -146,27 +149,27 @@ class A2ACompletionBridgeHandler: if custom_llm_provider == "pydantic_ai_agents": if api_base is None: raise ValueError("api_base is required for Pydantic AI agents") - + verbose_logger.info( f"Pydantic AI: Faking streaming for Pydantic AI agent at {api_base}" ) - + # Get non-streaming response first response_data = await PydanticAITransformation.send_non_streaming_request( api_base=api_base, request_id=request_id, params=params, ) - + # Convert to fake streaming async for chunk in PydanticAITransformation.fake_streaming_from_response( response_data=response_data, request_id=request_id, ): yield chunk - + return - + # Extract message from params message = params.get("message", {}) @@ -177,8 +180,8 @@ class A2ACompletionBridgeHandler: ) # Transform A2A message to OpenAI format - openai_messages = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages( - message + openai_messages = ( + A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message) ) # Get completion params @@ -205,7 +208,8 @@ class A2ACompletionBridgeHandler: } # Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.) litellm_params_to_add = { - k: v for k, v in litellm_params.items() + k: v + for k, v in litellm_params.items() if k not in ("model", "custom_llm_provider") } completion_params.update(litellm_params_to_add) @@ -244,9 +248,11 @@ class A2ACompletionBridgeHandler: # Emit artifact update with accumulated content if accumulated_text: - artifact_event = A2ACompletionBridgeTransformation.create_artifact_update_event( - ctx=ctx, - text=accumulated_text, + artifact_event = ( + A2ACompletionBridgeTransformation.create_artifact_update_event( + ctx=ctx, + text=accumulated_text, + ) ) yield artifact_event diff --git a/litellm/a2a_protocol/providers/litellm_completion/transformation.py b/litellm/a2a_protocol/providers/litellm_completion/transformation.py index bbe7daa9fc..8a03569f68 100644 --- a/litellm/a2a_protocol/providers/litellm_completion/transformation.py +++ b/litellm/a2a_protocol/providers/litellm_completion/transformation.py @@ -124,9 +124,7 @@ class A2ACompletionBridgeTransformation: }, } - verbose_logger.debug( - f"OpenAI -> A2A transform: content_length={len(content)}" - ) + verbose_logger.debug(f"OpenAI -> A2A transform: content_length={len(content)}") return a2a_response diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/__init__.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/__init__.py index 2187400b2d..8e9cd6fc87 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/__init__.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/__init__.py @@ -14,4 +14,3 @@ from litellm.a2a_protocol.providers.pydantic_ai_agents.transformation import ( ) __all__ = ["PydanticAIHandler", "PydanticAITransformation", "PydanticAIProviderConfig"] - diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/config.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/config.py index acf09554e5..d4c5f6a298 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/config.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/config.py @@ -11,7 +11,7 @@ from litellm.a2a_protocol.providers.pydantic_ai_agents.handler import PydanticAI class PydanticAIProviderConfig(BaseA2AProviderConfig): """ Provider configuration for Pydantic AI agents. - + Pydantic AI agents follow A2A protocol but don't support streaming natively. This config provides fake streaming by converting non-streaming responses into streaming chunks. """ @@ -48,4 +48,3 @@ class PydanticAIProviderConfig(BaseA2AProviderConfig): delay_ms=kwargs.get("delay_ms", 10), ): yield chunk - diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py index 6680a9fe48..7d4167752f 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py @@ -16,7 +16,7 @@ from litellm.a2a_protocol.providers.pydantic_ai_agents.transformation import ( class PydanticAIHandler: """ Handler for Pydantic AI agent requests. - + Provides: - Direct non-streaming requests to Pydantic AI agents - Fake streaming by converting non-streaming responses into streaming chunks @@ -41,9 +41,7 @@ class PydanticAIHandler: Returns: A2A SendMessageResponse dict """ - verbose_logger.info( - f"Pydantic AI: Routing to Pydantic AI agent at {api_base}" - ) + verbose_logger.info(f"Pydantic AI: Routing to Pydantic AI agent at {api_base}") # Send request directly to Pydantic AI agent response_data = await PydanticAITransformation.send_non_streaming_request( @@ -102,5 +100,3 @@ class PydanticAIHandler: delay_ms=delay_ms, ): yield chunk - - diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py index 9352eab6c8..e73b17ac3c 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py @@ -10,13 +10,16 @@ from typing import Any, AsyncIterator, Dict, cast from uuid import uuid4 from litellm._logging import verbose_logger -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + get_async_httpx_client, +) class PydanticAITransformation: """ Transformation layer for Pydantic AI agents. - + Handles: - Direct A2A requests to Pydantic AI endpoints - Polling for task completion (since Pydantic AI doesn't support streaming) @@ -27,13 +30,13 @@ class PydanticAITransformation: def _remove_none_values(obj: Any) -> Any: """ Recursively remove None values from a dict/list structure. - + FastA2A/Pydantic AI servers don't accept None values for optional fields - they expect those fields to be omitted entirely. - + Args: obj: Dict, list, or other value to clean - + Returns: Cleaned object with None values removed """ @@ -56,10 +59,10 @@ class PydanticAITransformation: def _params_to_dict(params: Any) -> Dict[str, Any]: """ Convert params to a dict, handling Pydantic models. - + Args: params: Dict or Pydantic model - + Returns: Dict representation of params """ @@ -86,7 +89,7 @@ class PydanticAITransformation: ) -> Dict[str, Any]: """ Poll for task completion using tasks/get method. - + Args: client: HTTPX async client endpoint: API endpoint URL @@ -94,7 +97,7 @@ class PydanticAITransformation: request_id: JSON-RPC request ID max_attempts: Maximum polling attempts poll_interval: Seconds between poll attempts - + Returns: Completed task response """ @@ -105,7 +108,7 @@ class PydanticAITransformation: "method": "tasks/get", "params": {"id": task_id}, } - + response = await client.post( endpoint, json=poll_request, @@ -113,23 +116,25 @@ class PydanticAITransformation: ) response.raise_for_status() poll_data = response.json() - + result = poll_data.get("result", {}) status = result.get("status", {}) state = status.get("state", "") - + verbose_logger.debug( f"Pydantic AI: Poll attempt {attempt + 1}/{max_attempts}, state={state}" ) - + if state == "completed": return poll_data elif state in ("failed", "canceled"): raise Exception(f"Task {task_id} ended with state: {state}") - + await asyncio.sleep(poll_interval) - - raise TimeoutError(f"Task {task_id} did not complete within {max_attempts * poll_interval} seconds") + + raise TimeoutError( + f"Task {task_id} did not complete within {max_attempts * poll_interval} seconds" + ) @staticmethod async def _send_and_poll_raw( @@ -140,7 +145,7 @@ class PydanticAITransformation: ) -> Dict[str, Any]: """ Send a request to Pydantic AI agent and return the raw task response. - + This is an internal method used by both non-streaming and streaming handlers. Returns the raw Pydantic AI task format with history/artifacts. @@ -155,10 +160,10 @@ class PydanticAITransformation: """ # Convert params to dict if it's a Pydantic model params_dict = PydanticAITransformation._params_to_dict(params) - + # Remove None values - FastA2A doesn't accept null for optional fields params_dict = PydanticAITransformation._remove_none_values(params_dict) - + # Ensure the message has 'kind': 'message' as required by FastA2A/Pydantic AI if "message" in params_dict: params_dict["message"]["kind"] = "message" @@ -174,9 +179,7 @@ class PydanticAITransformation: # FastA2A uses root endpoint (/) not /messages endpoint = api_base.rstrip("/") - verbose_logger.info( - f"Pydantic AI: Sending non-streaming request to {endpoint}" - ) + verbose_logger.info(f"Pydantic AI: Sending non-streaming request to {endpoint}") # Send request to Pydantic AI agent using shared async HTTP client client = get_async_httpx_client( @@ -190,12 +193,12 @@ class PydanticAITransformation: ) response.raise_for_status() response_data = response.json() - + # Check if task is already completed result = response_data.get("result", {}) status = result.get("status", {}) state = status.get("state", "") - + if state != "completed": # Need to poll for completion task_id = result.get("id") @@ -210,7 +213,9 @@ class PydanticAITransformation: request_id=request_id, ) - verbose_logger.info(f"Pydantic AI: Received completed response for request_id={request_id}") + verbose_logger.info( + f"Pydantic AI: Received completed response for request_id={request_id}" + ) return response_data @@ -256,7 +261,7 @@ class PydanticAITransformation: ) -> Dict[str, Any]: """ Send a request to Pydantic AI agent and return the raw task response. - + Used by streaming handler to get raw response for fake streaming. Args: @@ -282,7 +287,7 @@ class PydanticAITransformation: ) -> Dict[str, Any]: """ Transform Pydantic AI task response to standard A2A non-streaming format. - + Pydantic AI returns a task with history/artifacts, but the standard A2A non-streaming format expects: { @@ -296,11 +301,11 @@ class PydanticAITransformation: } } } - + Args: response_data: Pydantic AI task response request_id: Original request ID - + Returns: Standard A2A non-streaming response format """ @@ -308,14 +313,14 @@ class PydanticAITransformation: full_text, message_id, parts = PydanticAITransformation._extract_response_text( response_data ) - + # Build standard A2A message a2a_message = { "role": "agent", "parts": parts if parts else [{"kind": "text", "text": full_text}], "messageId": message_id, } - + # Return standard A2A non-streaming format return { "jsonrpc": "2.0", @@ -329,19 +334,19 @@ class PydanticAITransformation: def _extract_response_text(response_data: Dict[str, Any]) -> tuple[str, str, list]: """ Extract response text from completed task response. - + Pydantic AI returns completed tasks with: - history: list of messages (user and agent) - artifacts: list of result artifacts - + Args: response_data: Completed task response - + Returns: Tuple of (full_text, message_id, parts) """ result = response_data.get("result", {}) - + # Try to extract from artifacts first (preferred for results) artifacts = result.get("artifacts", []) if artifacts: @@ -352,7 +357,7 @@ class PydanticAITransformation: text = part.get("text", "") if text: return text, str(uuid4()), parts - + # Fall back to history - get the last agent message history = result.get("history", []) for msg in reversed(history): @@ -365,7 +370,7 @@ class PydanticAITransformation: full_text += part.get("text", "") if full_text: return full_text, message_id, parts - + # Fall back to message field (original format) message = result.get("message", {}) if message: @@ -376,7 +381,7 @@ class PydanticAITransformation: if part.get("kind") == "text": full_text += part.get("text", "") return full_text, message_id, parts - + return "", str(uuid4()), [] @staticmethod @@ -408,7 +413,7 @@ class PydanticAITransformation: full_text, message_id, parts = PydanticAITransformation._extract_response_text( response_data ) - + # Extract input message from raw response for history result = response_data.get("result", {}) history = result.get("history", []) @@ -436,7 +441,9 @@ class PydanticAITransformation: "contextId": context_id, "kind": "message", "messageId": input_message_id, - "parts": input_message.get("parts", [{"kind": "text", "text": ""}]), + "parts": input_message.get( + "parts", [{"kind": "text", "text": ""}] + ), "role": "user", "taskId": task_id, } @@ -475,7 +482,7 @@ class PydanticAITransformation: if full_text: # Split text into chunks for i in range(0, len(full_text), chunk_size): - chunk_text = full_text[i:i + chunk_size] + chunk_text = full_text[i : i + chunk_size] is_last_chunk = (i + chunk_size) >= len(full_text) artifact_event = { @@ -521,5 +528,3 @@ class PydanticAITransformation: verbose_logger.info( f"Pydantic AI: Fake streaming completed for request_id={request_id}" ) - - diff --git a/litellm/a2a_protocol/streaming_iterator.py b/litellm/a2a_protocol/streaming_iterator.py index 921dc0e52e..98d45cf2ac 100644 --- a/litellm/a2a_protocol/streaming_iterator.py +++ b/litellm/a2a_protocol/streaming_iterator.py @@ -71,7 +71,11 @@ class A2AStreamingIterator: def _collect_text_from_chunk(self, chunk: Any) -> None: """Extract text from a streaming chunk and add to collected parts.""" try: - chunk_dict = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {} + chunk_dict = ( + chunk.model_dump(mode="json", exclude_none=True) + if hasattr(chunk, "model_dump") + else {} + ) text = A2ARequestUtils.extract_text_from_response(chunk_dict) if text: self.collected_text_parts.append(text) @@ -81,7 +85,11 @@ class A2AStreamingIterator: def _is_completed_chunk(self, chunk: Any) -> bool: """Check if chunk indicates stream completion.""" try: - chunk_dict = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {} + chunk_dict = ( + chunk.model_dump(mode="json", exclude_none=True) + if hasattr(chunk, "model_dump") + else {} + ) result = chunk_dict.get("result", {}) if isinstance(result, dict): status = result.get("status", {}) @@ -102,7 +110,9 @@ class A2AStreamingIterator: prompt_tokens = A2ARequestUtils.count_tokens(input_text) # Use the last (most complete) text from chunks - output_text = self.collected_text_parts[-1] if self.collected_text_parts else "" + output_text = ( + self.collected_text_parts[-1] if self.collected_text_parts else "" + ) completion_tokens = A2ARequestUtils.count_tokens(output_text) total_tokens = prompt_tokens + completion_tokens @@ -158,7 +168,9 @@ class A2AStreamingIterator: result: Dict[str, Any] = { "id": getattr(self.request, "id", "unknown"), "jsonrpc": "2.0", - "usage": usage.model_dump() if hasattr(usage, "model_dump") else dict(usage), + "usage": usage.model_dump() + if hasattr(usage, "model_dump") + else dict(usage), } # Add final chunk result if available @@ -170,4 +182,3 @@ class A2AStreamingIterator: pass return result - diff --git a/litellm/anthropic_beta_headers_manager.py b/litellm/anthropic_beta_headers_manager.py index 24df6296b9..efa57ca058 100644 --- a/litellm/anthropic_beta_headers_manager.py +++ b/litellm/anthropic_beta_headers_manager.py @@ -38,7 +38,7 @@ _BETA_HEADERS_CONFIG: Optional[Dict] = None class GetAnthropicBetaHeadersConfig: """ Handles fetching, validating, and loading the Anthropic beta headers configuration. - + Similar to GetModelCostMap, this class manages the lifecycle of the beta headers configuration with support for remote fetching and local fallback. """ @@ -62,7 +62,7 @@ class GetAnthropicBetaHeadersConfig: "bedrock": {}, "bedrock_converse": {}, "vertex_ai": {}, - "provider_aliases": {} + "provider_aliases": {}, } @staticmethod @@ -84,9 +84,15 @@ class GetAnthropicBetaHeadersConfig: return False # Check for at least one provider key - provider_keys = ["anthropic", "azure_ai", "bedrock", "bedrock_converse", "vertex_ai"] + provider_keys = [ + "anthropic", + "azure_ai", + "bedrock", + "bedrock_converse", + "vertex_ai", + ] has_provider = any(key in fetched_config for key in provider_keys) - + if not has_provider: verbose_logger.warning( "LiteLLM: Fetched beta headers config missing provider keys. " @@ -100,7 +106,7 @@ class GetAnthropicBetaHeadersConfig: def validate_beta_headers_config(cls, fetched_config: dict) -> bool: """ Validate the integrity of a fetched beta headers config. - + Returns True if all checks pass, False otherwise. """ return cls._check_is_valid_dict(fetched_config) @@ -109,7 +115,7 @@ class GetAnthropicBetaHeadersConfig: def fetch_remote_beta_headers_config(url: str, timeout: int = 5) -> dict: """ Fetch the beta headers config from a remote URL. - + Returns the parsed JSON dict. Raises on network/parse errors (caller is expected to handle). """ @@ -121,14 +127,14 @@ class GetAnthropicBetaHeadersConfig: def get_beta_headers_config(url: str) -> dict: """ Public entry point — returns the beta headers config dict. - + 1. If ``LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS`` is set, uses the local backup only. 2. Otherwise fetches from ``url``, validates integrity, and falls back to the local backup on any failure. - + Args: url: URL to fetch the remote beta headers configuration from - + Returns: Dict containing the beta headers configuration """ @@ -149,7 +155,9 @@ def get_beta_headers_config(url: str) -> dict: return GetAnthropicBetaHeadersConfig.load_local_beta_headers_config() # Validate the fetched config - if not GetAnthropicBetaHeadersConfig.validate_beta_headers_config(fetched_config=content): + if not GetAnthropicBetaHeadersConfig.validate_beta_headers_config( + fetched_config=content + ): verbose_logger.warning( "LiteLLM: Fetched beta headers config failed integrity check. " "Using local backup instead. url=%s", @@ -164,23 +172,23 @@ def _load_beta_headers_config() -> Dict: """ Load the beta headers configuration. Uses caching to avoid repeated fetches/file reads. - + This function is called by all public API functions and manages the global cache. - + Returns: Dict containing the beta headers configuration """ global _BETA_HEADERS_CONFIG - + if _BETA_HEADERS_CONFIG is not None: return _BETA_HEADERS_CONFIG - + # Get the URL from environment or use default from litellm import anthropic_beta_headers_url - + _BETA_HEADERS_CONFIG = get_beta_headers_config(url=anthropic_beta_headers_url) verbose_logger.debug("Loaded and cached beta headers config") - + return _BETA_HEADERS_CONFIG @@ -188,7 +196,7 @@ def reload_beta_headers_config() -> Dict: """ Force reload the beta headers configuration from source (remote or local). Clears the cache and fetches fresh configuration. - + Returns: Dict containing the newly loaded beta headers configuration """ @@ -201,10 +209,10 @@ def reload_beta_headers_config() -> Dict: def get_provider_name(provider: str) -> str: """ Resolve provider aliases to canonical provider names. - + Args: provider: Provider name (may be an alias) - + Returns: Canonical provider name """ @@ -219,53 +227,53 @@ def filter_and_transform_beta_headers( ) -> List[str]: """ Filter and transform beta headers based on provider's mapping configuration. - + This function: 1. Only allows headers that are present in the provider's mapping keys 2. Filters out headers with null values (unsupported) 3. Maps headers to provider-specific names (e.g., advanced-tool-use -> tool-search-tool) - + Args: beta_headers: List of Anthropic beta header values provider: Provider name (e.g., "anthropic", "bedrock", "vertex_ai") - + Returns: List of filtered and transformed beta headers for the provider """ if not beta_headers: return [] - + config = _load_beta_headers_config() provider = get_provider_name(provider) - + # Get the header mapping for this provider provider_mapping = config.get(provider, {}) - + filtered_headers: Set[str] = set() - + for header in beta_headers: header = header.strip() - + # Check if header is in the mapping if header not in provider_mapping: verbose_logger.debug( f"Dropping unknown beta header '{header}' for provider '{provider}' (not in mapping)" ) continue - + # Get the mapped header value mapped_header = provider_mapping[header] - + # Skip if header is unsupported (null value) if mapped_header is None: verbose_logger.debug( f"Dropping unsupported beta header '{header}' for provider '{provider}'" ) continue - + # Add the mapped header filtered_headers.add(mapped_header) - + return sorted(list(filtered_headers)) @@ -275,18 +283,18 @@ def is_beta_header_supported( ) -> bool: """ Check if a specific beta header is supported by a provider. - + Args: beta_header: The Anthropic beta header value provider: Provider name - + Returns: True if the header is in the mapping with a non-null value, False otherwise """ config = _load_beta_headers_config() provider = get_provider_name(provider) provider_mapping = config.get(provider, {}) - + # Header is supported if it's in the mapping and has a non-null value return beta_header in provider_mapping and provider_mapping[beta_header] is not None @@ -297,26 +305,26 @@ def get_provider_beta_header( ) -> Optional[str]: """ Get the provider-specific beta header name for a given Anthropic beta header. - + This function handles header transformations/mappings (e.g., advanced-tool-use -> tool-search-tool). - + Args: anthropic_beta_header: The Anthropic beta header value provider: Provider name - + Returns: The provider-specific header name if supported, or None if unsupported/unknown """ config = _load_beta_headers_config() provider = get_provider_name(provider) - + # Get the header mapping for this provider provider_mapping = config.get(provider, {}) - + # Check if header is in the mapping if anthropic_beta_header not in provider_mapping: return None - + # Return the mapped value (could be None if unsupported) return provider_mapping[anthropic_beta_header] @@ -328,50 +336,50 @@ def update_headers_with_filtered_beta( """ Update headers dict by filtering and transforming anthropic-beta header values. Modifies the headers dict in place and returns it. - + Args: headers: Request headers dict (will be modified in place) provider: Provider name - + Returns: Updated headers dict """ existing_beta = headers.get("anthropic-beta") if not existing_beta: return headers - + # Parse existing beta headers beta_values = [b.strip() for b in existing_beta.split(",") if b.strip()] - + # Filter and transform based on provider filtered_beta_values = filter_and_transform_beta_headers( beta_headers=beta_values, provider=provider, ) - + # Update or remove the header if filtered_beta_values: headers["anthropic-beta"] = ",".join(filtered_beta_values) else: # Remove the header if no values remain headers.pop("anthropic-beta", None) - + return headers def get_unsupported_headers(provider: str) -> List[str]: """ Get all beta headers that are unsupported by a provider (have null values in mapping). - + Args: provider: Provider name - + Returns: List of unsupported Anthropic beta header names """ config = _load_beta_headers_config() provider = get_provider_name(provider) provider_mapping = config.get(provider, {}) - + # Return headers with null values return [header for header, value in provider_mapping.items() if value is None] diff --git a/litellm/anthropic_interface/exceptions/exception_mapping_utils.py b/litellm/anthropic_interface/exceptions/exception_mapping_utils.py index b8a5079a4e..28020e763f 100644 --- a/litellm/anthropic_interface/exceptions/exception_mapping_utils.py +++ b/litellm/anthropic_interface/exceptions/exception_mapping_utils.py @@ -149,7 +149,9 @@ class AnthropicExceptionMapping: parsed = None # If parsed and already in Anthropic format - passthrough - if parsed is not None and AnthropicExceptionMapping._is_anthropic_error_dict(parsed): + if parsed is not None and AnthropicExceptionMapping._is_anthropic_error_dict( + parsed + ): # Optionally add request_id if provided and not present if request_id and "request_id" not in parsed: parsed["request_id"] = request_id @@ -157,7 +159,9 @@ class AnthropicExceptionMapping: # Extract message - use parsed dict if available, otherwise raw string if parsed is not None: - message = AnthropicExceptionMapping._extract_message_from_dict(parsed, raw_message) + message = AnthropicExceptionMapping._extract_message_from_dict( + parsed, raw_message + ) else: message = raw_message diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index c752e84b96..4b965d4e63 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -10,7 +10,9 @@ from litellm.utils import token_counter async def calculate_batch_cost_and_usage( file_content_dictionary: List[dict], - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], + custom_llm_provider: Literal[ + "openai", "azure", "vertex_ai", "hosted_vllm", "anthropic" + ], model_name: Optional[str] = None, model_info: Optional[ModelInfo] = None, ) -> Tuple[float, Usage, List[str]]: @@ -34,19 +36,23 @@ async def calculate_batch_cost_and_usage( custom_llm_provider=custom_llm_provider, model_name=model_name, ) - batch_models = _get_batch_models_from_file_content(file_content_dictionary, model_name) + batch_models = _get_batch_models_from_file_content( + file_content_dictionary, model_name + ) return batch_cost, batch_usage, batch_models async def _handle_completed_batch( batch: Batch, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], + custom_llm_provider: Literal[ + "openai", "azure", "vertex_ai", "hosted_vllm", "anthropic" + ], model_name: Optional[str] = None, litellm_params: Optional[dict] = None, ) -> Tuple[float, Usage, List[str]]: """Helper function to process a completed batch and handle logging - + Args: batch: The batch object custom_llm_provider: The LLM provider @@ -70,7 +76,9 @@ async def _handle_completed_batch( model_name=model_name, ) - batch_models = _get_batch_models_from_file_content(file_content_dictionary, model_name) + batch_models = _get_batch_models_from_file_content( + file_content_dictionary, model_name + ) return batch_cost, batch_usage, batch_models @@ -96,7 +104,9 @@ def _get_batch_models_from_file_content( def _batch_cost_calculator( file_content_dictionary: List[dict], - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai", + custom_llm_provider: Literal[ + "openai", "azure", "vertex_ai", "hosted_vllm", "anthropic" + ] = "openai", model_name: Optional[str] = None, model_info: Optional[ModelInfo] = None, ) -> float: @@ -105,10 +115,12 @@ def _batch_cost_calculator( """ # Handle Vertex AI with specialized method if custom_llm_provider == "vertex_ai" and model_name: - batch_cost, _ = calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name) + batch_cost, _ = calculate_vertex_ai_batch_cost_and_usage( + file_content_dictionary, model_name + ) verbose_logger.debug("vertex_ai_total_cost=%s", batch_cost) return batch_cost - + # For other providers, use the existing logic total_cost = _get_batch_job_cost_from_file_content( file_content_dictionary=file_content_dictionary, @@ -173,7 +185,10 @@ def calculate_vertex_ai_batch_cost_and_usage( verbose_logger.info( "vertex_ai batch cost: cost=%s, prompt=%d, completion=%d, total=%d", - total_cost, prompt_tokens, completion_tokens, total_tokens, + total_cost, + prompt_tokens, + completion_tokens, + total_tokens, ) return total_cost, Usage( @@ -185,12 +200,14 @@ def calculate_vertex_ai_batch_cost_and_usage( async def _get_batch_output_file_content_as_dictionary( batch: Batch, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai", + custom_llm_provider: Literal[ + "openai", "azure", "vertex_ai", "hosted_vllm", "anthropic" + ] = "openai", litellm_params: Optional[dict] = None, ) -> List[dict]: """ Get the batch output file content as a list of dictionaries - + Args: batch: The batch object custom_llm_provider: The LLM provider @@ -198,8 +215,9 @@ async def _get_batch_output_file_content_as_dictionary( Required for Azure and other providers that need authentication """ from litellm.files.main import afile_content - from litellm.proxy.openai_files_endpoints.common_utils import \ - _is_base64_encoded_unified_file_id + from litellm.proxy.openai_files_endpoints.common_utils import ( + _is_base64_encoded_unified_file_id, + ) if custom_llm_provider == "vertex_ai": raise ValueError("Vertex AI does not support file content retrieval") @@ -211,21 +229,27 @@ async def _get_batch_output_file_content_as_dictionary( is_base64_unified_file_id = _is_base64_encoded_unified_file_id(file_id) if is_base64_unified_file_id: try: - file_id = is_base64_unified_file_id.split("llm_output_file_id,")[1].split(";")[0] - verbose_logger.debug(f"Extracted LLM output file ID from unified file ID: {file_id}") + file_id = is_base64_unified_file_id.split("llm_output_file_id,")[1].split( + ";" + )[0] + verbose_logger.debug( + f"Extracted LLM output file ID from unified file ID: {file_id}" + ) except (IndexError, AttributeError) as e: - verbose_logger.error(f"Failed to extract LLM output file ID from unified file ID: {batch.output_file_id}, error: {e}") + verbose_logger.error( + f"Failed to extract LLM output file ID from unified file ID: {batch.output_file_id}, error: {e}" + ) # Build kwargs for afile_content with credentials from litellm_params file_content_kwargs = { "file_id": file_id, "custom_llm_provider": custom_llm_provider, } - + # Extract and add credentials for file access credentials = _extract_file_access_credentials(litellm_params) file_content_kwargs.update(credentials) - + _file_content = await afile_content(**file_content_kwargs) # type: ignore[reportArgumentType] return _get_file_content_as_dictionary(_file_content.content) @@ -233,30 +257,37 @@ async def _get_batch_output_file_content_as_dictionary( def _extract_file_access_credentials(litellm_params: Optional[dict]) -> dict: """ Extract credentials from litellm_params for file access operations. - + This method extracts relevant authentication and configuration parameters needed for accessing files across different providers (Azure, Vertex AI, etc.). - + Args: litellm_params: Dictionary containing litellm parameters with credentials - + Returns: Dictionary containing only the credentials needed for file access """ credentials = {} - + if litellm_params: # List of credential keys that should be passed to file operations credential_keys = [ - "api_key", "api_base", "api_version", "organization", - "azure_ad_token", "azure_ad_token_provider", - "vertex_project", "vertex_location", "vertex_credentials", - "timeout", "max_retries" + "api_key", + "api_base", + "api_version", + "organization", + "azure_ad_token", + "azure_ad_token_provider", + "vertex_project", + "vertex_location", + "vertex_credentials", + "timeout", + "max_retries", ] for key in credential_keys: if key in litellm_params: credentials[key] = litellm_params[key] - + return credentials @@ -279,7 +310,9 @@ def _get_file_content_as_dictionary(file_content: bytes) -> List[dict]: def _get_batch_job_cost_from_file_content( file_content_dictionary: List[dict], - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai", + custom_llm_provider: Literal[ + "openai", "azure", "vertex_ai", "hosted_vllm", "anthropic" + ] = "openai", model_info: Optional[ModelInfo] = None, ) -> float: """ @@ -321,7 +354,9 @@ def _get_batch_job_cost_from_file_content( def _get_batch_job_total_usage_from_file_content( file_content_dictionary: List[dict], - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai", + custom_llm_provider: Literal[ + "openai", "azure", "vertex_ai", "hosted_vllm", "anthropic" + ] = "openai", model_name: Optional[str] = None, ) -> Usage: """ @@ -329,9 +364,11 @@ def _get_batch_job_total_usage_from_file_content( """ # Handle Vertex AI with specialized method if custom_llm_provider == "vertex_ai" and model_name: - _, batch_usage = calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name) + _, batch_usage = calculate_vertex_ai_batch_cost_and_usage( + file_content_dictionary, model_name + ) return batch_usage - + # For other providers, use the existing logic total_tokens: int = 0 prompt_tokens: int = 0 @@ -349,6 +386,7 @@ def _get_batch_job_total_usage_from_file_content( completion_tokens=completion_tokens, ) + def _get_batch_job_input_file_usage( file_content_dictionary: List[dict], custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", @@ -358,25 +396,26 @@ def _get_batch_job_input_file_usage( Count the number of tokens in the input file Used for batch rate limiting to count the number of tokens in the input file - """ + """ prompt_tokens: int = 0 completion_tokens: int = 0 - + for _item in file_content_dictionary: body = _item.get("body", {}) model = body.get("model", model_name or "") messages = body.get("messages", []) - + if messages: item_tokens = token_counter(model=model, messages=messages) prompt_tokens += item_tokens - + return Usage( total_tokens=prompt_tokens + completion_tokens, prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, ) + def _get_batch_job_usage_from_response_body(response_body: dict) -> Usage: """ Get the tokens of a batch job from the response body @@ -400,4 +439,4 @@ def _batch_response_was_successful(batch_job_output_file: dict) -> bool: Check if the batch job response status == 200 """ _response: dict = batch_job_output_file.get("response", None) or {} - return _response.get("status_code", None) == 200 \ No newline at end of file + return _response.get("status_code", None) == 200 diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 723b59c6b4..1a03b172d3 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -109,7 +109,9 @@ async def acreate_batch( completion_window: Literal["24h"], endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"], input_file_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"] = "openai", + custom_llm_provider: Literal[ + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm" + ] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -159,7 +161,9 @@ def create_batch( # noqa: PLR0915 completion_window: Literal["24h"], endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"], input_file_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"] = "openai", + custom_llm_provider: Literal[ + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm" + ] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -220,7 +224,9 @@ def create_batch( # noqa: PLR0915 extra_body=extra_body, ) if output_expires_after is not None: - _create_batch_request["output_expires_after"] = cast(FileExpiresAfter, output_expires_after) + _create_batch_request["output_expires_after"] = cast( + FileExpiresAfter, output_expires_after + ) if model is not None: provider_config = ProviderConfigManager.get_provider_batches_config( model=model, @@ -364,7 +370,9 @@ def create_batch( # noqa: PLR0915 @client async def aretrieve_batch( batch_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"] = "openai", + custom_llm_provider: Literal[ + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic" + ] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -410,7 +418,9 @@ def _handle_retrieve_batch_providers_without_provider_config( litellm_params: dict, _retrieve_batch_request: RetrieveBatchRequest, _is_async: bool, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"] = "openai", + custom_llm_provider: Literal[ + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic" + ] = "openai", logging_obj: Optional[Any] = None, ): api_base: Optional[str] = None @@ -549,7 +559,9 @@ def _handle_retrieve_batch_providers_without_provider_config( @client def retrieve_batch( batch_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"] = "openai", + custom_llm_provider: Literal[ + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic" + ] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -929,7 +941,6 @@ def cancel_batch( LiteLLM Equivalent of POST https://api.openai.com/v1/batches/{batch_id}/cancel """ try: - try: if model is not None: _, custom_llm_provider, _, _ = get_llm_provider( @@ -1097,25 +1108,40 @@ def _handle_async_invoke_status( "inprogress": "in_progress", "in_progress": "in_progress", } - normalized_status: BatchJobStatus = status_mapping.get(aws_status_lower, "failed") # Default to "failed" if unknown status + normalized_status: BatchJobStatus = status_mapping.get( + aws_status_lower, "failed" + ) # Default to "failed" if unknown status # Get output S3 URI safely output_s3_uri = "" try: - output_s3_uri = status_response["outputDataConfig"]["s3OutputDataConfig"]["s3Uri"] + output_s3_uri = status_response["outputDataConfig"]["s3OutputDataConfig"][ + "s3Uri" + ] except (KeyError, TypeError): pass - + # Use BedrockBatchesConfig's timestamp parsing method (expects raw AWS status string) import time from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig - created_at, in_progress_at, completed_at, failed_at, _, _ = BedrockBatchesConfig()._parse_timestamps_and_status(status_response, aws_status_raw) + + ( + created_at, + in_progress_at, + completed_at, + failed_at, + _, + _, + ) = BedrockBatchesConfig()._parse_timestamps_and_status( + status_response, aws_status_raw + ) result = LiteLLMBatch( id=status_response["invocationArn"], object="batch", status=normalized_status, - created_at=created_at or int(time.time()), # Provide default timestamp if None + created_at=created_at + or int(time.time()), # Provide default timestamp if None in_progress_at=in_progress_at, completed_at=completed_at, failed_at=failed_at, diff --git a/litellm/caching/azure_blob_cache.py b/litellm/caching/azure_blob_cache.py index 45e551bdae..a2246640c3 100644 --- a/litellm/caching/azure_blob_cache.py +++ b/litellm/caching/azure_blob_cache.py @@ -22,7 +22,9 @@ class AzureBlobCache(BaseCache): from azure.storage.blob import BlobServiceClient from azure.core.exceptions import ResourceExistsError from azure.identity import DefaultAzureCredential - from azure.identity.aio import DefaultAzureCredential as AsyncDefaultAzureCredential + from azure.identity.aio import ( + DefaultAzureCredential as AsyncDefaultAzureCredential, + ) from azure.storage.blob.aio import BlobServiceClient as AsyncBlobServiceClient self.container_client = BlobServiceClient( @@ -50,14 +52,16 @@ class AzureBlobCache(BaseCache): print_verbose(f"LiteLLM SET Cache - Azure Blob. Key={key}. Value={value}") serialized_value = json.dumps(value) try: - await self.async_container_client.upload_blob(key, serialized_value, overwrite=True) + await self.async_container_client.upload_blob( + key, serialized_value, overwrite=True + ) except Exception as e: # NON blocking - notify users Azure Blob is throwing an exception print_verbose(f"LiteLLM set_cache() - Got exception from Azure Blob: {e}") def get_cache(self, key, **kwargs): from azure.core.exceptions import ResourceNotFoundError - + try: print_verbose(f"Get Azure Blob Cache: key: {key}") as_bytes = self.container_client.download_blob(key).readall() @@ -74,7 +78,7 @@ class AzureBlobCache(BaseCache): async def async_get_cache(self, key, **kwargs): from azure.core.exceptions import ResourceNotFoundError - + try: print_verbose(f"Get Azure Blob Cache: key: {key}") blob = await self.async_container_client.download_blob(key) diff --git a/litellm/caching/base_cache.py b/litellm/caching/base_cache.py index 8660e64efd..81f1d61bd0 100644 --- a/litellm/caching/base_cache.py +++ b/litellm/caching/base_cache.py @@ -53,12 +53,12 @@ class BaseCache(ABC): async def disconnect(self): raise NotImplementedError - + async def test_connection(self) -> dict: """ Test the cache connection. - + Returns: dict: {"status": "success" | "failed", "message": str, "error": Optional[str]} """ - raise NotImplementedError \ No newline at end of file + raise NotImplementedError diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 4e97197a9d..7cdbd3fc03 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -78,9 +78,7 @@ class CachingHandlerResponse(BaseModel): cached_result: Optional[Any] = None final_embedding_cached_response: Optional[EmbeddingResponse] = None - embedding_all_elements_cache_hit: bool = ( - False # this is set to True when all elements in the list have a cache hit in the embedding cache, if true return the final_embedding_cached_response no need to make an API call - ) + embedding_all_elements_cache_hit: bool = False # this is set to True when all elements in the list have a cache hit in the embedding cache, if true return the final_embedding_cached_response no need to make an API call in_memory_cache_obj = InMemoryCache() @@ -159,7 +157,7 @@ class LLMCachingHandler: ######################################################### parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) kwargs["parent_otel_span"] = parent_otel_span - + if litellm.cache is not None and self._is_call_type_supported_by_cache( original_function=original_function ): @@ -181,7 +179,9 @@ class LLMCachingHandler: api_base=kwargs.get("api_base", None), api_key=kwargs.get("api_key", None), ) - cache_duration_ms = (cache_check_end_time - cache_check_start_time) * 1000 + cache_duration_ms = ( + cache_check_end_time - cache_check_start_time + ) * 1000 self._update_litellm_logging_obj_environment( logging_obj=logging_obj, model=model, @@ -194,7 +194,6 @@ class LLMCachingHandler: call_type = original_function.__name__ - cached_result = self._convert_cached_result_to_model_response( cached_result=cached_result, call_type=call_type, @@ -244,7 +243,7 @@ class LLMCachingHandler: final_embedding_cached_response=final_embedding_cached_response, embedding_all_elements_cache_hit=embedding_all_elements_cache_hit, ) - + verbose_logger.debug(f"CACHE RESULT: {cached_result}") return CachingHandlerResponse( cached_result=cached_result, @@ -265,9 +264,8 @@ class LLMCachingHandler: ) -> CachingHandlerResponse: from litellm.utils import CustomStreamWrapper - cached_result: Optional[Any] = None - + # Check if caching should be performed BEFORE doing expensive kwargs copy if litellm.cache is not None and self._is_call_type_supported_by_cache( original_function=original_function @@ -325,7 +323,7 @@ class LLMCachingHandler: result=cached_result, start_time=start_time, end_time=end_time, - cache_hit=cache_hit + cache_hit=cache_hit, ) cache_key = litellm.cache.get_cache_key(**kwargs) if ( @@ -554,12 +552,18 @@ class LLMCachingHandler: GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( async_coroutine=logging_obj.async_success_handler( - result=cached_result, start_time=start_time, end_time=end_time, cache_hit=cache_hit + result=cached_result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, ) ) logging_obj.handle_sync_success_callbacks_for_async_calls( - result=cached_result, start_time=start_time, end_time=end_time, cache_hit=cache_hit + result=cached_result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, ) async def _retrieve_from_cache( @@ -728,10 +732,9 @@ class LLMCachingHandler: response_type="audio_transcription", hidden_params=hidden_params, ) - elif ( - call_type == "aresponses" - or call_type == "responses" - ) and isinstance(cached_result, dict): + elif (call_type == "aresponses" or call_type == "responses") and isinstance( + cached_result, dict + ): # Convert cached dict back to ResponsesAPIResponse object cached_result = ResponsesAPIResponse(**cached_result) @@ -741,7 +744,7 @@ class LLMCachingHandler: and isinstance(cached_result._hidden_params, dict) ): cached_result._hidden_params["cache_hit"] = True - + ######################################################### # Add final timing metrics to the cached result ######################################################### @@ -1011,9 +1014,9 @@ class LLMCachingHandler: } if litellm.cache is not None: - litellm_params["preset_cache_key"] = ( - litellm.cache._get_preset_cache_key_from_kwargs(**kwargs) - ) + litellm_params[ + "preset_cache_key" + ] = litellm.cache._get_preset_cache_key_from_kwargs(**kwargs) else: litellm_params["preset_cache_key"] = None diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index 6df570c72b..4020b8cc22 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -319,18 +319,20 @@ class DualCache(BaseCache): previous_access_times ) raise - + # Short-circuit if redis_result is None or contains only None values - if redis_result is None or all(v is None for v in redis_result.values()): + if redis_result is None or all( + v is None for v in redis_result.values() + ): return result # Pre-compute key-to-index mapping for O(1) lookup key_to_index = {key: i for i, key in enumerate(keys)} - + # Update both result and in-memory cache in a single loop for key, value in redis_result.items(): result[key_to_index[key]] = value - + if value is not None and self.in_memory_cache is not None: await self.in_memory_cache.async_set_cache( key, value, **kwargs @@ -346,6 +348,8 @@ class DualCache(BaseCache): ) try: if self.in_memory_cache is not None: + if "ttl" not in kwargs and self.default_in_memory_ttl is not None: + kwargs["ttl"] = self.default_in_memory_ttl await self.in_memory_cache.async_set_cache(key, value, **kwargs) if self.redis_cache is not None and local_only is False: @@ -367,6 +371,8 @@ class DualCache(BaseCache): ) try: if self.in_memory_cache is not None: + if "ttl" not in kwargs and self.default_in_memory_ttl is not None: + kwargs["ttl"] = self.default_in_memory_ttl await self.in_memory_cache.async_set_cache_pipeline( cache_list=cache_list, **kwargs ) diff --git a/litellm/caching/gcs_cache.py b/litellm/caching/gcs_cache.py index 88857ba0e7..a5bd092f15 100644 --- a/litellm/caching/gcs_cache.py +++ b/litellm/caching/gcs_cache.py @@ -16,13 +16,23 @@ from .base_cache import BaseCache class GCSCache(BaseCache): - def __init__(self, bucket_name: Optional[str] = None, path_service_account: Optional[str] = None, gcs_path: Optional[str] = None) -> None: + def __init__( + self, + bucket_name: Optional[str] = None, + path_service_account: Optional[str] = None, + gcs_path: Optional[str] = None, + ) -> None: super().__init__() self.bucket_name = bucket_name or GCSBucketBase(bucket_name=None).BUCKET_NAME - self.path_service_account = path_service_account or GCSBucketBase(bucket_name=None).path_service_account_json + self.path_service_account = ( + path_service_account + or GCSBucketBase(bucket_name=None).path_service_account_json + ) self.key_prefix = gcs_path.rstrip("/") + "/" if gcs_path else "" # create httpx clients - self.async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) + self.async_client = get_async_httpx_client( + llm_provider=httpxSpecialProvider.LoggingCallback + ) self.sync_client = _get_httpx_client() def _construct_headers(self) -> dict: @@ -52,7 +62,9 @@ class GCSCache(BaseCache): data = json.dumps(value) await self.async_client.post(url=url, data=data, headers=headers) except Exception as e: - print_verbose(f"GCS Caching: async_set_cache() - Got exception from GCS: {e}") + print_verbose( + f"GCS Caching: async_set_cache() - Got exception from GCS: {e}" + ) def get_cache(self, key, **kwargs): try: @@ -69,7 +81,9 @@ class GCSCache(BaseCache): return cached_response return None except Exception as e: - verbose_logger.error(f"GCS Caching: get_cache() - Got exception from GCS: {e}") + verbose_logger.error( + f"GCS Caching: get_cache() - Got exception from GCS: {e}" + ) async def async_get_cache(self, key, **kwargs): try: @@ -82,7 +96,9 @@ class GCSCache(BaseCache): return json.loads(response.text) return None except Exception as e: - verbose_logger.error(f"GCS Caching: async_get_cache() - Got exception from GCS: {e}") + verbose_logger.error( + f"GCS Caching: async_get_cache() - Got exception from GCS: {e}" + ) def flush_cache(self): pass diff --git a/litellm/caching/qdrant_semantic_cache.py b/litellm/caching/qdrant_semantic_cache.py index 181effa01d..5e3713e5a1 100644 --- a/litellm/caching/qdrant_semantic_cache.py +++ b/litellm/caching/qdrant_semantic_cache.py @@ -54,7 +54,9 @@ class QdrantSemanticCache(BaseCache): raise Exception("similarity_threshold must be provided, passed None") self.similarity_threshold = similarity_threshold self.embedding_model = embedding_model - self.vector_size = vector_size if vector_size is not None else QDRANT_VECTOR_SIZE + self.vector_size = ( + vector_size if vector_size is not None else QDRANT_VECTOR_SIZE + ) headers = {} # check if defined as os.environ/ variable diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index fa9b94bc2a..82794c116f 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -268,19 +268,19 @@ class RedisCache(BaseCache): def _parse_redis_major_version(self) -> int: """ Parse Redis version to extract the major version number. - + Handles multiple version formats: - Strings: "7.0.0", "6", "7.0.0-rc1", " 7.0.0 " - Floats: 7.0 (e.g., from AWS ElastiCache Valkey) - Integers: 7 - Malformed: "latest", "", "Unknown" (defaults to DEFAULT_REDIS_MAJOR_VERSION) - + Returns: int: The major version number (defaults to DEFAULT_REDIS_MAJOR_VERSION if unparseable) """ if self.redis_version == "Unknown": return DEFAULT_REDIS_MAJOR_VERSION - + try: version_str = str(self.redis_version).strip() # Handle cases where there's no dot (e.g., "7" or 7) @@ -1113,14 +1113,14 @@ class RedisCache(BaseCache): self.redis_client.close() except Exception as e: verbose_logger.debug("Error closing sync Redis client: %s", e) - + async def test_connection(self) -> dict: """ Test the Redis connection by creating a new client and pinging it. - + This creates a fresh connection without using cached clients or connection pools to ensure the credentials are actually valid. - + Returns: dict: {"status": "success" | "failed", "message": str, "error": Optional[str]} """ @@ -1129,29 +1129,26 @@ class RedisCache(BaseCache): # Create a fresh Redis client with current settings redis_client = redis_async.Redis(**self.redis_kwargs) - + # Test the connection ping_result = await redis_client.ping() # type: ignore[misc] # Close the connection await redis_client.aclose() # type: ignore[attr-defined] - + if ping_result: return { "status": "success", - "message": "Redis connection test successful" + "message": "Redis connection test successful", } else: - return { - "status": "failed", - "message": "Redis ping returned False" - } + return {"status": "failed", "message": "Redis ping returned False"} except Exception as e: verbose_logger.error(f"Redis connection test failed: {str(e)}") return { "status": "failed", "message": f"Redis connection failed: {str(e)}", - "error": str(e) + "error": str(e), } async def async_delete_cache(self, key: str): diff --git a/litellm/caching/redis_cluster_cache.py b/litellm/caching/redis_cluster_cache.py index 664578c870..b0f5754f58 100644 --- a/litellm/caching/redis_cluster_cache.py +++ b/litellm/caching/redis_cluster_cache.py @@ -57,11 +57,11 @@ class RedisClusterCache(RedisCache): """ async_redis_cluster_client = self.init_async_client() return await async_redis_cluster_client.mget_nonatomic(keys=keys) # type: ignore - + async def test_connection(self) -> dict: """ Test the Redis Cluster connection. - + Returns: dict: {"status": "success" | "failed", "message": str, "error": Optional[str]} """ @@ -72,37 +72,38 @@ class RedisClusterCache(RedisCache): # Create ClusterNode objects from startup_nodes cluster_kwargs = self.redis_kwargs.copy() startup_nodes = cluster_kwargs.pop("startup_nodes", []) - + new_startup_nodes: List[ClusterNode] = [] for item in startup_nodes: new_startup_nodes.append(ClusterNode(**item)) - + # Create a fresh Redis Cluster client with current settings redis_client = redis_async.RedisCluster( startup_nodes=new_startup_nodes, **cluster_kwargs # type: ignore ) - + # Test the connection ping_result = await redis_client.ping() # type: ignore[attr-defined, misc] # Close the connection await redis_client.aclose() # type: ignore[attr-defined] - + if ping_result: return { "status": "success", - "message": "Redis Cluster connection test successful" + "message": "Redis Cluster connection test successful", } else: return { "status": "failed", - "message": "Redis Cluster ping returned False" + "message": "Redis Cluster ping returned False", } except Exception as e: from litellm._logging import verbose_logger + verbose_logger.error(f"Redis Cluster connection test failed: {str(e)}") return { "status": "failed", "message": f"Redis Cluster connection failed: {str(e)}", - "error": str(e) - } \ No newline at end of file + "error": str(e), + } diff --git a/litellm/caching/s3_cache.py b/litellm/caching/s3_cache.py index 180964605f..e26fbe8981 100644 --- a/litellm/caching/s3_cache.py +++ b/litellm/caching/s3_cache.py @@ -110,7 +110,9 @@ class S3Cache(BaseCache): func = partial(self.set_cache, key, value, **kwargs) await loop.run_in_executor(None, func) except Exception as e: - verbose_logger.error(f"S3 Caching: async_set_cache() - Got exception from S3: {e}") + verbose_logger.error( + f"S3 Caching: async_set_cache() - Got exception from S3: {e}" + ) def get_cache(self, key, **kwargs): import botocore @@ -126,7 +128,7 @@ class S3Cache(BaseCache): if cached_response is not None: if "Expires" in cached_response: - expires_time = cached_response['Expires'] + expires_time = cached_response["Expires"] current_time = datetime.now(expires_time.tzinfo) if current_time > expires_time: diff --git a/litellm/completion_extras/litellm_responses_transformation/handler.py b/litellm/completion_extras/litellm_responses_transformation/handler.py index e9ac1d2ad7..2164a2c0f0 100644 --- a/litellm/completion_extras/litellm_responses_transformation/handler.py +++ b/litellm/completion_extras/litellm_responses_transformation/handler.py @@ -61,9 +61,7 @@ class ResponsesToCompletionBridgeHandler: existing.setdefault(key, value) return response - def _collect_response_from_stream( - self, stream_iter: Any - ) -> "ResponsesAPIResponse": + def _collect_response_from_stream(self, stream_iter: Any) -> "ResponsesAPIResponse": for _ in stream_iter: pass @@ -144,7 +142,9 @@ class ResponsesToCompletionBridgeHandler: custom_llm_provider=custom_llm_provider, ) - def completion(self, *args, **kwargs) -> Union[ + def completion( + self, *args, **kwargs + ) -> Union[ Coroutine[Any, Any, Union["ModelResponse", "CustomStreamWrapper"]], "ModelResponse", "CustomStreamWrapper", diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index babb575ee3..42359afef4 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -63,7 +63,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): def __init__(self): pass - def _handle_raw_dict_response_item(self, item: Dict[str, Any], index: int) -> Tuple[Optional[Any], int]: + def _handle_raw_dict_response_item( + self, item: Dict[str, Any], index: int + ) -> Tuple[Optional[Any], int]: """ Handle raw dict response items from Responses API (e.g., GPT-5 Codex format). @@ -106,9 +108,13 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): if item_type == "function_call": # Extract provider_specific_fields if present and pass through as-is provider_specific_fields = item.get("provider_specific_fields") - if provider_specific_fields and not isinstance(provider_specific_fields, dict): + if provider_specific_fields and not isinstance( + provider_specific_fields, dict + ): provider_specific_fields = ( - dict(provider_specific_fields) if hasattr(provider_specific_fields, "__dict__") else {} + dict(provider_specific_fields) + if hasattr(provider_specific_fields, "__dict__") + else {} ) tool_call_dict = { @@ -124,7 +130,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): if provider_specific_fields: tool_call_dict["provider_specific_fields"] = provider_specific_fields # Also add to function's provider_specific_fields for consistency - tool_call_dict["function"]["provider_specific_fields"] = provider_specific_fields + tool_call_dict["function"][ + "provider_specific_fields" + ] = provider_specific_fields msg = Message( content=None, @@ -232,10 +240,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): if key in ("max_tokens", "max_completion_tokens"): responses_api_request["max_output_tokens"] = value elif key == "tools" and value is not None: - responses_api_request["tools"] = ( - self._convert_tools_to_responses_format( - cast(List[Dict[str, Any]], value) - ) + responses_api_request[ + "tools" + ] = self._convert_tools_to_responses_format( + cast(List[Dict[str, Any]], value) ) elif key == "response_format": text_format = self._transform_response_format_to_text_format(value) @@ -250,9 +258,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): elif key == "web_search_options": self._add_web_search_tool(responses_api_request, value) - def _build_sanitized_litellm_params( - self, litellm_params: dict - ) -> Dict[str, Any]: + def _build_sanitized_litellm_params(self, litellm_params: dict) -> Dict[str, Any]: """Build sanitized litellm_params with merged metadata.""" responses_optional_param_keys = set( ResponsesAPIOptionalRequestParams.__annotations__.keys() @@ -337,7 +343,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): previous_response_id = optional_params.get("previous_response_id") if previous_response_id: # Use the existing session handler for responses API - verbose_logger.debug(f"Chat provider: Warning ignoring previous response ID: {previous_response_id}") + verbose_logger.debug( + f"Chat provider: Warning ignoring previous response ID: {previous_response_id}" + ) # Convert back to responses API format for the actual request @@ -347,9 +355,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): setattr(litellm_logging_obj, "call_type", CallTypes.responses.value) - sanitized_litellm_params = self._build_sanitized_litellm_params( - litellm_params - ) + sanitized_litellm_params = self._build_sanitized_litellm_params(litellm_params) request_data = { "model": api_model, @@ -359,7 +365,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): "client": client, } - verbose_logger.debug(f"Chat provider: Final request model={api_model}, input_items={len(input_items)}") + verbose_logger.debug( + f"Chat provider: Final request model={api_model}, input_items={len(input_items)}" + ) self._merge_responses_api_request_into_request_data( request_data, responses_api_request, instructions @@ -390,6 +398,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ResponseOutputMessage, ResponseReasoningItem, ) + from openai.types.responses.response_output_item import ( + ResponseApplyPatchToolCall, + ) from litellm.types.utils import Choices, Message @@ -439,11 +450,21 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): LiteLLMCompletionResponsesConfig, ) - tool_call_dict = ( - LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call( - tool_call_item=item, - index=tool_call_index, - ) + tool_call_dict = LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call( + tool_call_item=item, + index=tool_call_index, + ) + accumulated_tool_calls.append(tool_call_dict) + tool_call_index += 1 + + elif isinstance(item, ResponseApplyPatchToolCall): + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + tool_call_dict = LiteLLMCompletionResponsesConfig.convert_apply_patch_tool_call_to_chat_completion_tool_call( + tool_call_item=item, + index=tool_call_index, ) accumulated_tool_calls.append(tool_call_dict) tool_call_index += 1 @@ -463,7 +484,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): tool_calls=accumulated_tool_calls, reasoning_content=reasoning_content, ) - choices.append(Choices(message=msg, finish_reason="tool_calls", index=index)) + choices.append( + Choices(message=msg, finish_reason="tool_calls", index=index) + ) reasoning_content = None return choices @@ -499,10 +522,17 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ) if len(choices) == 0: - if raw_response.incomplete_details is not None and raw_response.incomplete_details.reason is not None: - raise ValueError(f"{model} unable to complete request: {raw_response.incomplete_details.reason}") + if ( + raw_response.incomplete_details is not None + and raw_response.incomplete_details.reason is not None + ): + raise ValueError( + f"{model} unable to complete request: {raw_response.incomplete_details.reason}" + ) else: - raise ValueError(f"Unknown items in responses API response: {raw_response.output}") + raise ValueError( + f"Unknown items in responses API response: {raw_response.output}" + ) setattr(model_response, "choices", choices) @@ -511,21 +541,28 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): setattr( model_response, "usage", - ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(raw_response.usage), + ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( + raw_response.usage + ), ) # Preserve hidden params from the ResponsesAPIResponse, especially the headers # which contain important provider information like x-request-id raw_response_hidden_params = getattr(raw_response, "_hidden_params", {}) if raw_response_hidden_params: - if not hasattr(model_response, "_hidden_params") or model_response._hidden_params is None: + if ( + not hasattr(model_response, "_hidden_params") + or model_response._hidden_params is None + ): model_response._hidden_params = {} # Merge the raw_response hidden params with model_response hidden params # Preserve existing keys in model_response but add/override with raw_response params for key, value in raw_response_hidden_params.items(): if key == "additional_headers" and key in model_response._hidden_params: # Merge additional_headers to preserve both sets - existing_additional_headers = model_response._hidden_params.get("additional_headers", {}) + existing_additional_headers = model_response._hidden_params.get( + "additional_headers", {} + ) merged_headers = {**value, **existing_additional_headers} model_response._hidden_params[key] = merged_headers else: @@ -535,13 +572,19 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): def get_model_response_iterator( self, - streaming_response: Union[Iterator[str], AsyncIterator[str], "ModelResponse", "BaseModel"], + streaming_response: Union[ + Iterator[str], AsyncIterator[str], "ModelResponse", "BaseModel" + ], sync_stream: bool, json_mode: Optional[bool] = False, ) -> BaseModelResponseIterator: - return OpenAiResponsesToChatCompletionStreamIterator(streaming_response, sync_stream, json_mode) + return OpenAiResponsesToChatCompletionStreamIterator( + streaming_response, sync_stream, json_mode + ) - def _convert_content_str_to_input_text(self, content: str, role: str) -> Dict[str, Any]: + def _convert_content_str_to_input_text( + self, content: str, role: str + ) -> Dict[str, Any]: if role == "user" or role == "system" or role == "tool": return {"type": "input_text", "text": content} else: @@ -568,7 +611,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): if actual_image_url is None: raise ValueError(f"Invalid image URL: {content_image_url}") - image_param = ResponseInputImageParam(image_url=actual_image_url, detail="auto", type="input_image") + image_param = ResponseInputImageParam( + image_url=actual_image_url, detail="auto", type="input_image" + ) if detail: image_param["detail"] = detail @@ -581,7 +626,13 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): Union[ str, List[Any], - Iterable[Union["OpenAIMessageContentListBlock", "ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock"]], + Iterable[ + Union[ + "OpenAIMessageContentListBlock", + "ChatCompletionThinkingBlock", + "ChatCompletionRedactedThinkingBlock", + ] + ], ] ], role: str, @@ -589,7 +640,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): """Convert chat completion content to responses API format""" from litellm.types.llms.openai import ChatCompletionImageObject - verbose_logger.debug(f"Chat provider: Converting content to responses format - input type: {type(content)}") + verbose_logger.debug( + f"Chat provider: Converting content to responses format - input type: {type(content)}" + ) if content is None: return [self._convert_content_str_to_input_text("", role)] @@ -600,7 +653,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): elif isinstance(content, list): result = [] for i, item in enumerate(content): - verbose_logger.debug(f"Chat provider: Processing content item {i}: {type(item)} = {item}") + verbose_logger.debug( + f"Chat provider: Processing content item {i}: {type(item)} = {item}" + ) if isinstance(item, str): converted = self._convert_content_str_to_input_text(item, role) result.append(converted) @@ -609,7 +664,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): # Handle multimodal content original_type = item.get("type") if original_type == "text": - converted = self._convert_content_str_to_input_text(item.get("text", ""), role) + converted = self._convert_content_str_to_input_text( + item.get("text", ""), role + ) result.append(converted) verbose_logger.debug(f"Chat provider: text -> {converted}") elif original_type == "image_url": @@ -621,14 +678,18 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ), ) result.append(converted) - verbose_logger.debug(f"Chat provider: image_url -> {converted}") + verbose_logger.debug( + f"Chat provider: image_url -> {converted}" + ) else: # Try to map other types to responses API format item_type = original_type or "input_text" if item_type == "image": converted = {"type": "input_image", **item} result.append(converted) - verbose_logger.debug(f"Chat provider: image -> {converted}") + verbose_logger.debug( + f"Chat provider: image -> {converted}" + ) elif item_type in [ "input_text", "input_image", @@ -640,12 +701,18 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ]: # Already in responses API format result.append(item) - verbose_logger.debug(f"Chat provider: passthrough -> {item}") + verbose_logger.debug( + f"Chat provider: passthrough -> {item}" + ) else: # Default to input_text for unknown types - converted = self._convert_content_str_to_input_text(str(item.get("text", item)), role) + converted = self._convert_content_str_to_input_text( + str(item.get("text", item)), role + ) result.append(converted) - verbose_logger.debug(f"Chat provider: unknown({original_type}) -> {converted}") + verbose_logger.debug( + f"Chat provider: unknown({original_type}) -> {converted}" + ) verbose_logger.debug(f"Chat provider: Final converted content: {result}") return result else: @@ -653,13 +720,17 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): verbose_logger.debug(f"Chat provider: Other content type -> {result}") return result - def _convert_tools_to_responses_format(self, tools: List[Dict[str, Any]]) -> List["ALL_RESPONSES_API_TOOL_PARAMS"]: + def _convert_tools_to_responses_format( + self, tools: List[Dict[str, Any]] + ) -> List["ALL_RESPONSES_API_TOOL_PARAMS"]: """Convert chat completion tools to responses API tools format""" responses_tools: List["ALL_RESPONSES_API_TOOL_PARAMS"] = [] for tool in tools: # convert function tool from chat completion to responses API format if tool.get("type") == "function": - function_tool = cast(ChatCompletionToolParamFunctionChunk, tool.get("function")) + function_tool = cast( + ChatCompletionToolParamFunctionChunk, tool.get("function") + ) responses_tools.append( FunctionToolParam( name=function_tool["name"], @@ -685,7 +756,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): if not extra_body: return optional_params - supported_responses_api_params = set(ResponsesAPIOptionalRequestParams.__annotations__.keys()) + supported_responses_api_params = set( + ResponsesAPIOptionalRequestParams.__annotations__.keys() + ) # Also include params we handle specially supported_responses_api_params.update( { @@ -703,7 +776,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): return optional_params - def _map_reasoning_effort(self, reasoning_effort: Union[str, Dict[str, Any]]) -> Optional[Reasoning]: + def _map_reasoning_effort( + self, reasoning_effort: Union[str, Dict[str, Any]] + ) -> Optional[Reasoning]: # If dict is passed, convert it directly to Reasoning object if isinstance(reasoning_effort, dict): return Reasoning(**reasoning_effort) # type: ignore[typeddict-item] @@ -711,25 +786,38 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): # Check if auto-summary is enabled via flag or environment variable # Priority: litellm.reasoning_auto_summary flag > LITELLM_REASONING_AUTO_SUMMARY env var auto_summary_enabled = ( - litellm.reasoning_auto_summary or os.getenv("LITELLM_REASONING_AUTO_SUMMARY", "false").lower() == "true" + litellm.reasoning_auto_summary + or os.getenv("LITELLM_REASONING_AUTO_SUMMARY", "false").lower() == "true" ) # If string is passed, map with optional summary based on flag/env var if reasoning_effort == "none": return Reasoning(effort="none", summary="detailed") if auto_summary_enabled else Reasoning(effort="none") # type: ignore elif reasoning_effort == "high": - return Reasoning(effort="high", summary="detailed") if auto_summary_enabled else Reasoning(effort="high") + return ( + Reasoning(effort="high", summary="detailed") + if auto_summary_enabled + else Reasoning(effort="high") + ) elif reasoning_effort == "xhigh": return Reasoning(effort="xhigh", summary="detailed") if auto_summary_enabled else Reasoning(effort="xhigh") # type: ignore[typeddict-item] elif reasoning_effort == "medium": return ( - Reasoning(effort="medium", summary="detailed") if auto_summary_enabled else Reasoning(effort="medium") + Reasoning(effort="medium", summary="detailed") + if auto_summary_enabled + else Reasoning(effort="medium") ) elif reasoning_effort == "low": - return Reasoning(effort="low", summary="detailed") if auto_summary_enabled else Reasoning(effort="low") + return ( + Reasoning(effort="low", summary="detailed") + if auto_summary_enabled + else Reasoning(effort="low") + ) elif reasoning_effort == "minimal": return ( - Reasoning(effort="minimal", summary="detailed") if auto_summary_enabled else Reasoning(effort="minimal") + Reasoning(effort="minimal", summary="detailed") + if auto_summary_enabled + else Reasoning(effort="minimal") ) return None @@ -745,7 +833,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): responses_api_request: The responses API request dict to modify web_search_options: Web search configuration (dict or other value) """ - if "tools" not in responses_api_request or responses_api_request["tools"] is None: + if ( + "tools" not in responses_api_request + or responses_api_request["tools"] is None + ): responses_api_request["tools"] = [] # Get the tools list with proper type narrowing @@ -835,13 +926,17 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): annotation_dict = annotation else: # Skip unsupported annotation types - verbose_logger.debug(f"Skipping unsupported annotation type: {type(annotation)}") + verbose_logger.debug( + f"Skipping unsupported annotation type: {type(annotation)}" + ) continue result.append(annotation_dict) # type: ignore except Exception as e: # Skip malformed annotations - verbose_logger.debug(f"Skipping malformed annotation: {annotation}, error: {e}") + verbose_logger.debug( + f"Skipping malformed annotation: {annotation}, error: {e}" + ) continue return result if result else None @@ -862,7 +957,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): - def __init__(self, streaming_response, sync_stream: bool, json_mode: Optional[bool] = False): + def __init__( + self, streaming_response, sync_stream: bool, json_mode: Optional[bool] = False + ): super().__init__(streaming_response, sync_stream, json_mode) def _handle_string_chunk( @@ -875,7 +972,9 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): if not str_line or str_line.startswith("event:"): # ignore. - return GenericStreamingChunk(text="", tool_use=None, is_finished=False, finish_reason="", usage=None) + return GenericStreamingChunk( + text="", tool_use=None, is_finished=False, finish_reason="", usage=None + ) index = str_line.find("data:") if index != -1: str_line = str_line[index + 5 :] @@ -938,9 +1037,13 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): if output_item.get("type") == "function_call": # Extract provider_specific_fields if present provider_specific_fields = output_item.get("provider_specific_fields") - if provider_specific_fields and not isinstance(provider_specific_fields, dict): + if provider_specific_fields and not isinstance( + provider_specific_fields, dict + ): provider_specific_fields = ( - dict(provider_specific_fields) if hasattr(provider_specific_fields, "__dict__") else {} + dict(provider_specific_fields) + if hasattr(provider_specific_fields, "__dict__") + else {} ) function_chunk = ChatCompletionToolCallFunctionChunk( @@ -949,7 +1052,9 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): ) if provider_specific_fields: - function_chunk["provider_specific_fields"] = provider_specific_fields + function_chunk[ + "provider_specific_fields" + ] = provider_specific_fields tool_call_index = parsed_chunk.get("output_index", 0) tool_call_chunk = ChatCompletionToolCallChunk( @@ -986,7 +1091,9 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): id=None, index=tool_call_index, type="function", - function=ChatCompletionToolCallFunctionChunk(name=None, arguments=content_part), + function=ChatCompletionToolCallFunctionChunk( + name=None, arguments=content_part + ), ) ] ), @@ -995,16 +1102,22 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): ] ) else: - raise ValueError(f"Chat provider: Invalid function argument delta {parsed_chunk}") + raise ValueError( + f"Chat provider: Invalid function argument delta {parsed_chunk}" + ) elif event_type == "response.output_item.done": # New output item added output_item = parsed_chunk.get("item", {}) if output_item.get("type") == "function_call": # Extract provider_specific_fields if present provider_specific_fields = output_item.get("provider_specific_fields") - if provider_specific_fields and not isinstance(provider_specific_fields, dict): + if provider_specific_fields and not isinstance( + provider_specific_fields, dict + ): provider_specific_fields = ( - dict(provider_specific_fields) if hasattr(provider_specific_fields, "__dict__") else {} + dict(provider_specific_fields) + if hasattr(provider_specific_fields, "__dict__") + else {} ) function_chunk = ChatCompletionToolCallFunctionChunk( @@ -1014,7 +1127,9 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): # Add provider_specific_fields to function if present if provider_specific_fields: - function_chunk["provider_specific_fields"] = provider_specific_fields + function_chunk[ + "provider_specific_fields" + ] = provider_specific_fields tool_call_index = parsed_chunk.get("output_index", 0) tool_call_chunk = ChatCompletionToolCallChunk( @@ -1090,11 +1205,22 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): output_items = response_data.get("output", []) if response_data else [] has_function_calls = any( - item.get("type") == "function_call" for item in output_items if isinstance(item, dict) + item.get("type") == "function_call" + for item in output_items + if isinstance(item, dict) ) finish_reason = "tool_calls" if has_function_calls else "stop" + usage = None + if response_data.get("usage"): + from litellm.responses.utils import ResponseAPILoggingUtils + + usage = ( + ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( + response_data.get("usage") + ) + ) return ModelResponseStream( choices=[ StreamingChoices( @@ -1102,12 +1228,15 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): delta=Delta(content=""), finish_reason=finish_reason, ) - ] + ], + usage=usage, ) else: pass # For any unhandled event types, create a minimal valid chunk or skip - verbose_logger.debug(f"Chat provider: Unhandled event type '{event_type}', creating empty chunk") + verbose_logger.debug( + f"Chat provider: Unhandled event type '{event_type}', creating empty chunk" + ) # Return a minimal valid chunk for unknown events return ModelResponseStream( @@ -1130,5 +1259,9 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): Returns: ModelResponseStream: OpenAI-formatted streaming chunk """ - verbose_logger.debug(f"Chat provider: transform_streaming_response called with chunk: {chunk}") - return OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(chunk) + verbose_logger.debug( + f"Chat provider: transform_streaming_response called with chunk: {chunk}" + ) + return OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream( + chunk + ) diff --git a/litellm/constants.py b/litellm/constants.py index 2486c223ec..34b6950a21 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -60,9 +60,7 @@ LITELLM_MAX_STREAMING_DURATION_SECONDS = ( # Maximum number of base64 characters to keep in logging payloads. # Data URIs exceeding this are replaced with a size placeholder. # Set to 0 to disable truncation. -MAX_BASE64_LENGTH_FOR_LOGGING = int( - os.getenv("MAX_BASE64_LENGTH_FOR_LOGGING", 64) -) +MAX_BASE64_LENGTH_FOR_LOGGING = int(os.getenv("MAX_BASE64_LENGTH_FOR_LOGGING", 64)) # When true, adds detailed per-phase timing breakdown headers to responses. # Headers: x-litellm-timing-{pre-processing,llm-api,post-processing,message-copy}-ms @@ -1421,9 +1419,7 @@ SPECIAL_LITELLM_AUTH_TOKEN = ["ui-token"] DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int( os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60) ) -DEFAULT_ACCESS_GROUP_CACHE_TTL = int( - os.getenv("DEFAULT_ACCESS_GROUP_CACHE_TTL", 600) -) +DEFAULT_ACCESS_GROUP_CACHE_TTL = int(os.getenv("DEFAULT_ACCESS_GROUP_CACHE_TTL", 600)) # Sentry Scrubbing Configuration SENTRY_DENYLIST = [ diff --git a/litellm/containers/__init__.py b/litellm/containers/__init__.py index e279cb429e..48ab5de418 100644 --- a/litellm/containers/__init__.py +++ b/litellm/containers/__init__.py @@ -42,4 +42,3 @@ __all__ = [ "retrieve_container_file", "retrieve_container_file_content", ] - diff --git a/litellm/containers/endpoint_factory.py b/litellm/containers/endpoint_factory.py index 0b73a19b92..22fd4226de 100644 --- a/litellm/containers/endpoint_factory.py +++ b/litellm/containers/endpoint_factory.py @@ -43,13 +43,13 @@ def _load_endpoints_config() -> Dict: def create_sync_endpoint_function(endpoint_config: Dict) -> Callable: """ Create a sync SDK function from endpoint config. - + Uses the generic container handler instead of individual handler methods. """ endpoint_name = endpoint_config["name"] response_type = RESPONSE_TYPES.get(endpoint_config["response_type"]) path_params = endpoint_config.get("path_params", []) - + @client def endpoint_func( timeout: int = 600, @@ -76,14 +76,16 @@ def create_sync_endpoint_function(endpoint_config: Dict) -> Callable: # Get provider config litellm_params = GenericLiteLLMParams(**kwargs) - container_provider_config: Optional[BaseContainerConfig] = ( - ProviderConfigManager.get_provider_container_config( - provider=litellm.LlmProviders(custom_llm_provider), - ) + container_provider_config: Optional[ + BaseContainerConfig + ] = ProviderConfigManager.get_provider_container_config( + provider=litellm.LlmProviders(custom_llm_provider), ) if container_provider_config is None: - raise ValueError(f"Container provider config not found for: {custom_llm_provider}") + raise ValueError( + f"Container provider config not found for: {custom_llm_provider}" + ) # Build optional params for logging optional_params = {k: kwargs.get(k) for k in path_params if k in kwargs} @@ -126,7 +128,7 @@ def create_async_endpoint_function( endpoint_config: Dict, ) -> Callable: """Create an async SDK function that wraps the sync function.""" - + @client async def async_endpoint_func( timeout: int = 600, @@ -176,21 +178,21 @@ def create_async_endpoint_function( def generate_container_endpoints() -> Dict[str, Callable]: """ Generate all container endpoint functions from the JSON config. - + Returns a dict mapping function names to their implementations. """ config = _load_endpoints_config() endpoints = {} - + for endpoint_config in config["endpoints"]: # Create sync function sync_func = create_sync_endpoint_function(endpoint_config) endpoints[endpoint_config["name"]] = sync_func - + # Create async function async_func = create_async_endpoint_function(sync_func, endpoint_config) endpoints[endpoint_config["async_name"]] = async_func - + return endpoints @@ -222,5 +224,9 @@ retrieve_container_file = _generated_endpoints.get("retrieve_container_file") aretrieve_container_file = _generated_endpoints.get("aretrieve_container_file") delete_container_file = _generated_endpoints.get("delete_container_file") adelete_container_file = _generated_endpoints.get("adelete_container_file") -retrieve_container_file_content = _generated_endpoints.get("retrieve_container_file_content") -aretrieve_container_file_content = _generated_endpoints.get("aretrieve_container_file_content") +retrieve_container_file_content = _generated_endpoints.get( + "retrieve_container_file_content" +) +aretrieve_container_file_content = _generated_endpoints.get( + "aretrieve_container_file_content" +) diff --git a/litellm/containers/main.py b/litellm/containers/main.py index 105e999ffe..88318ee039 100644 --- a/litellm/containers/main.py +++ b/litellm/containers/main.py @@ -39,6 +39,7 @@ __all__ = [ "upload_container_file", ] + ##### Container Create ####################### @client async def acreate_container( @@ -164,10 +165,7 @@ def create_container( extra_query: Optional[Dict[str, Any]] = None, extra_body: Optional[Dict[str, Any]] = None, **kwargs, -) -> Union[ - ContainerObject, - Coroutine[Any, Any, ContainerObject], -]: +) -> Union[ContainerObject, Coroutine[Any, Any, ContainerObject],]: """Create a container using the OpenAI Container API. Currently supports OpenAI @@ -175,7 +173,7 @@ def create_container( Example: ```python import litellm - + response = litellm.create_container( name="My Container", custom_llm_provider="openai", @@ -207,19 +205,23 @@ def create_container( **kwargs, ) # get provider config - container_provider_config: Optional[BaseContainerConfig] = ( - ProviderConfigManager.get_provider_container_config( - provider=litellm.LlmProviders(custom_llm_provider), - ) + container_provider_config: Optional[ + BaseContainerConfig + ] = ProviderConfigManager.get_provider_container_config( + provider=litellm.LlmProviders(custom_llm_provider), ) if container_provider_config is None: - raise ValueError(f"container operations are not supported for {custom_llm_provider}") + raise ValueError( + f"container operations are not supported for {custom_llm_provider}" + ) local_vars.update(kwargs) # Get ContainerCreateOptionalRequestParams with only valid parameters container_create_optional_params: ContainerCreateOptionalRequestParams = ( - ContainerRequestUtils.get_requested_container_create_optional_param(local_vars) + ContainerRequestUtils.get_requested_container_create_optional_param( + local_vars + ) ) # Get optional parameters for the container API @@ -388,10 +390,7 @@ def list_containers( extra_query: Optional[Dict[str, Any]] = None, extra_body: Optional[Dict[str, Any]] = None, **kwargs, -) -> Union[ - ContainerListResponse, - Coroutine[Any, Any, ContainerListResponse], -]: +) -> Union[ContainerListResponse, Coroutine[Any, Any, ContainerListResponse],]: """List containers using the OpenAI Container API. Currently supports OpenAI @@ -420,18 +419,22 @@ def list_containers( **kwargs, ) # get provider config - container_provider_config: Optional[BaseContainerConfig] = ( - ProviderConfigManager.get_provider_container_config( - provider=litellm.LlmProviders(custom_llm_provider), - ) + container_provider_config: Optional[ + BaseContainerConfig + ] = ProviderConfigManager.get_provider_container_config( + provider=litellm.LlmProviders(custom_llm_provider), ) if container_provider_config is None: - raise ValueError(f"Container provider config not found for provider: {custom_llm_provider}") + raise ValueError( + f"Container provider config not found for provider: {custom_llm_provider}" + ) # Get container list request parameters container_list_optional_params: ContainerListOptionalRequestParams = ( - ContainerRequestUtils.get_requested_container_list_optional_param(local_vars) + ContainerRequestUtils.get_requested_container_list_optional_param( + local_vars + ) ) # Pre Call logging @@ -582,10 +585,7 @@ def retrieve_container( extra_query: Optional[Dict[str, Any]] = None, extra_body: Optional[Dict[str, Any]] = None, **kwargs, -) -> Union[ - ContainerObject, - Coroutine[Any, Any, ContainerObject], -]: +) -> Union[ContainerObject, Coroutine[Any, Any, ContainerObject],]: """Retrieve a container using the OpenAI Container API. Currently supports OpenAI @@ -614,14 +614,16 @@ def retrieve_container( **kwargs, ) # get provider config - container_provider_config: Optional[BaseContainerConfig] = ( - ProviderConfigManager.get_provider_container_config( - provider=litellm.LlmProviders(custom_llm_provider), - ) + container_provider_config: Optional[ + BaseContainerConfig + ] = ProviderConfigManager.get_provider_container_config( + provider=litellm.LlmProviders(custom_llm_provider), ) if container_provider_config is None: - raise ValueError(f"Container provider config not found for provider: {custom_llm_provider}") + raise ValueError( + f"Container provider config not found for provider: {custom_llm_provider}" + ) # Pre Call logging litellm_logging_obj.update_environment_variables( @@ -768,10 +770,7 @@ def delete_container( extra_query: Optional[Dict[str, Any]] = None, extra_body: Optional[Dict[str, Any]] = None, **kwargs, -) -> Union[ - DeleteContainerResult, - Coroutine[Any, Any, DeleteContainerResult], -]: +) -> Union[DeleteContainerResult, Coroutine[Any, Any, DeleteContainerResult],]: """Delete a container using the OpenAI Container API. Currently supports OpenAI @@ -800,14 +799,16 @@ def delete_container( **kwargs, ) # get provider config - container_provider_config: Optional[BaseContainerConfig] = ( - ProviderConfigManager.get_provider_container_config( - provider=litellm.LlmProviders(custom_llm_provider), - ) + container_provider_config: Optional[ + BaseContainerConfig + ] = ProviderConfigManager.get_provider_container_config( + provider=litellm.LlmProviders(custom_llm_provider), ) if container_provider_config is None: - raise ValueError(f"Container provider config not found for provider: {custom_llm_provider}") + raise ValueError( + f"Container provider config not found for provider: {custom_llm_provider}" + ) # Pre Call logging litellm_logging_obj.update_environment_variables( @@ -968,10 +969,7 @@ def list_container_files( extra_query: Optional[Dict[str, Any]] = None, extra_body: Optional[Dict[str, Any]] = None, **kwargs, -) -> Union[ - ContainerFileListResponse, - Coroutine[Any, Any, ContainerFileListResponse], -]: +) -> Union[ContainerFileListResponse, Coroutine[Any, Any, ContainerFileListResponse],]: """List files in a container using the OpenAI Container API. Currently supports OpenAI @@ -1000,19 +998,26 @@ def list_container_files( **kwargs, ) # get provider config - container_provider_config: Optional[BaseContainerConfig] = ( - ProviderConfigManager.get_provider_container_config( - provider=litellm.LlmProviders(custom_llm_provider), - ) + container_provider_config: Optional[ + BaseContainerConfig + ] = ProviderConfigManager.get_provider_container_config( + provider=litellm.LlmProviders(custom_llm_provider), ) if container_provider_config is None: - raise ValueError(f"Container provider config not found for provider: {custom_llm_provider}") + raise ValueError( + f"Container provider config not found for provider: {custom_llm_provider}" + ) # Pre Call logging litellm_logging_obj.update_environment_variables( model="", - optional_params={"container_id": container_id, "after": after, "limit": limit, "order": order}, + optional_params={ + "container_id": container_id, + "after": after, + "limit": limit, + "order": order, + }, litellm_params={ "litellm_call_id": litellm_call_id, }, @@ -1180,10 +1185,7 @@ def upload_container_file( extra_query: Optional[Dict[str, Any]] = None, extra_body: Optional[Dict[str, Any]] = None, **kwargs, -) -> Union[ - ContainerFileObject, - Coroutine[Any, Any, ContainerFileObject], -]: +) -> Union[ContainerFileObject, Coroutine[Any, Any, ContainerFileObject],]: """Upload a file to a container using the OpenAI Container API. This endpoint allows uploading files directly to a container session, @@ -1241,14 +1243,16 @@ def upload_container_file( **kwargs, ) # get provider config - container_provider_config: Optional[BaseContainerConfig] = ( - ProviderConfigManager.get_provider_container_config( - provider=litellm.LlmProviders(custom_llm_provider), - ) + container_provider_config: Optional[ + BaseContainerConfig + ] = ProviderConfigManager.get_provider_container_config( + provider=litellm.LlmProviders(custom_llm_provider), ) if container_provider_config is None: - raise ValueError(f"Container provider config not found for provider: {custom_llm_provider}") + raise ValueError( + f"Container provider config not found for provider: {custom_llm_provider}" + ) # Pre Call logging litellm_logging_obj.update_environment_variables( diff --git a/litellm/containers/utils.py b/litellm/containers/utils.py index f30f1e154b..048f587fda 100644 --- a/litellm/containers/utils.py +++ b/litellm/containers/utils.py @@ -1,7 +1,10 @@ from typing import Dict from litellm.llms.base_llm.containers.transformation import BaseContainerConfig -from litellm.types.containers.main import ContainerCreateOptionalRequestParams, ContainerListOptionalRequestParams +from litellm.types.containers.main import ( + ContainerCreateOptionalRequestParams, + ContainerListOptionalRequestParams, +) class ContainerRequestUtils: diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 75d45af86e..c1daa109c7 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -120,37 +120,49 @@ else: LitellmLoggingObject = Any # Pre-resolved CallTypes enum values for fast membership checks -_A2A_CALL_TYPES = frozenset({ - CallTypes.asend_message.value, - CallTypes.send_message.value, -}) +_A2A_CALL_TYPES = frozenset( + { + CallTypes.asend_message.value, + CallTypes.send_message.value, + } +) -_VIDEO_CALL_TYPES = frozenset({ - CallTypes.create_video.value, - CallTypes.acreate_video.value, - CallTypes.video_remix.value, - CallTypes.avideo_remix.value, -}) +_VIDEO_CALL_TYPES = frozenset( + { + CallTypes.create_video.value, + CallTypes.acreate_video.value, + CallTypes.video_remix.value, + CallTypes.avideo_remix.value, + } +) -_SPEECH_CALL_TYPES = frozenset({ - CallTypes.speech.value, - CallTypes.aspeech.value, -}) +_SPEECH_CALL_TYPES = frozenset( + { + CallTypes.speech.value, + CallTypes.aspeech.value, + } +) -_TRANSCRIPTION_CALL_TYPES = frozenset({ - CallTypes.atranscription.value, - CallTypes.transcription.value, -}) +_TRANSCRIPTION_CALL_TYPES = frozenset( + { + CallTypes.atranscription.value, + CallTypes.transcription.value, + } +) -_RERANK_CALL_TYPES = frozenset({ - CallTypes.rerank.value, - CallTypes.arerank.value, -}) +_RERANK_CALL_TYPES = frozenset( + { + CallTypes.rerank.value, + CallTypes.arerank.value, + } +) -_SEARCH_CALL_TYPES = frozenset({ - CallTypes.search.value, - CallTypes.asearch.value, -}) +_SEARCH_CALL_TYPES = frozenset( + { + CallTypes.search.value, + CallTypes.asearch.value, + } +) _AREALTIME_CALL_TYPE = CallTypes.arealtime.value _MCP_CALL_TYPE = CallTypes.call_mcp_tool.value @@ -522,7 +534,10 @@ def cost_per_token( # noqa: PLR0915 return dashscope_cost_per_token(model=model, usage=usage_block) elif custom_llm_provider == "azure_ai": return azure_ai_cost_per_token( - model=model, usage=usage_block, response_time_ms=response_time_ms, request_model=request_model + model=model, + usage=usage_block, + response_time_ms=response_time_ms, + request_model=request_model, ) else: model_info = _cached_get_model_info_helper( @@ -1114,9 +1129,9 @@ def completion_cost( # noqa: PLR0915 or isinstance(completion_response, dict) ): # tts returns a custom class if isinstance(completion_response, dict): - usage_obj: Optional[Union[dict, Usage]] = ( - completion_response.get("usage", {}) - ) + usage_obj: Optional[ + Union[dict, Usage] + ] = completion_response.get("usage", {}) else: usage_obj = getattr(completion_response, "usage", {}) if isinstance(usage_obj, BaseModel) and not _is_known_usage_objects( @@ -1501,7 +1516,6 @@ def completion_cost( # noqa: PLR0915 else: additional_costs = None - _final_cost = ( prompt_tokens_cost_usd_dollar + completion_tokens_cost_usd_dollar ) @@ -1519,7 +1533,11 @@ def completion_cost( # noqa: PLR0915 # Apply discount from module-level config if configured original_cost = _final_cost if litellm.cost_discount_config: - _final_cost, discount_percent, discount_amount = _apply_cost_discount( + ( + _final_cost, + discount_percent, + discount_amount, + ) = _apply_cost_discount( base_cost=_final_cost, custom_llm_provider=custom_llm_provider, ) @@ -1976,9 +1994,7 @@ def default_video_cost_calculator( cost_info = litellm.model_cost[prefixed_model] if cost_info is None: - raise Exception( - f"Model not found in cost map for model={model}" - ) + raise Exception(f"Model not found in cost map for model={model}") # Check for video-specific cost per second first video_cost_per_second = cost_info.get("output_cost_per_video_per_second") @@ -2250,4 +2266,3 @@ def handle_realtime_stream_cost_calculation( total_cost = input_cost_per_token + output_cost_per_token return total_cost - diff --git a/litellm/evals/main.py b/litellm/evals/main.py index a39c283915..e57c75bd9b 100644 --- a/litellm/evals/main.py +++ b/litellm/evals/main.py @@ -152,16 +152,14 @@ def create_eval( custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ( - ProviderConfigManager.get_provider_evals_api_config( # type: ignore - provider=litellm.LlmProviders(custom_llm_provider), - ) + evals_api_provider_config: Optional[ + BaseEvalsAPIConfig + ] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + provider=litellm.LlmProviders(custom_llm_provider), ) if evals_api_provider_config is None: - raise ValueError( - f"CREATE eval is not supported for {custom_llm_provider}" - ) + raise ValueError(f"CREATE eval is not supported for {custom_llm_provider}") # Build create request create_request: CreateEvalRequest = { @@ -344,10 +342,10 @@ def list_evals( custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ( - ProviderConfigManager.get_provider_evals_api_config( # type: ignore - provider=litellm.LlmProviders(custom_llm_provider), - ) + evals_api_provider_config: Optional[ + BaseEvalsAPIConfig + ] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + provider=litellm.LlmProviders(custom_llm_provider), ) if evals_api_provider_config is None: @@ -513,10 +511,10 @@ def get_eval( custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ( - ProviderConfigManager.get_provider_evals_api_config( # type: ignore - provider=litellm.LlmProviders(custom_llm_provider), - ) + evals_api_provider_config: Optional[ + BaseEvalsAPIConfig + ] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + provider=litellm.LlmProviders(custom_llm_provider), ) if evals_api_provider_config is None: @@ -681,16 +679,14 @@ def update_eval( custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ( - ProviderConfigManager.get_provider_evals_api_config( # type: ignore - provider=litellm.LlmProviders(custom_llm_provider), - ) + evals_api_provider_config: Optional[ + BaseEvalsAPIConfig + ] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + provider=litellm.LlmProviders(custom_llm_provider), ) if evals_api_provider_config is None: - raise ValueError( - f"UPDATE eval is not supported for {custom_llm_provider}" - ) + raise ValueError(f"UPDATE eval is not supported for {custom_llm_provider}") # Build update request update_request: UpdateEvalRequest = {} @@ -701,20 +697,41 @@ def update_eval( if metadata is not None: # List of internal LiteLLM metadata keys that should NOT be sent to OpenAI internal_keys = { - "headers", "requester_metadata", "user_api_key_hash", "user_api_key_alias", - "user_api_key_spend", "user_api_key_max_budget", "user_api_key_team_id", - "user_api_key_user_id", "user_api_key_org_id", "user_api_key_team_alias", - "user_api_key_end_user_id", "user_api_key_user_email", "user_api_key_request_route", - "user_api_key_budget_reset_at", "user_api_key_auth_metadata", "user_api_key", - "user_api_end_user_max_budget", "user_api_key_auth", "litellm_api_version", - "global_max_parallel_requests", "user_api_key_team_max_budget", - "user_api_key_team_spend", "user_api_key_model_max_budget", - "user_api_key_user_spend", "user_api_key_user_max_budget", - "user_api_key_metadata", "endpoint", "litellm_parent_otel_span", - "requester_ip_address", "user_agent", + "headers", + "requester_metadata", + "user_api_key_hash", + "user_api_key_alias", + "user_api_key_spend", + "user_api_key_max_budget", + "user_api_key_team_id", + "user_api_key_user_id", + "user_api_key_org_id", + "user_api_key_team_alias", + "user_api_key_end_user_id", + "user_api_key_user_email", + "user_api_key_request_route", + "user_api_key_budget_reset_at", + "user_api_key_auth_metadata", + "user_api_key", + "user_api_end_user_max_budget", + "user_api_key_auth", + "litellm_api_version", + "global_max_parallel_requests", + "user_api_key_team_max_budget", + "user_api_key_team_spend", + "user_api_key_model_max_budget", + "user_api_key_user_spend", + "user_api_key_user_max_budget", + "user_api_key_metadata", + "endpoint", + "litellm_parent_otel_span", + "requester_ip_address", + "user_agent", } # Only include user-provided metadata keys - filtered_metadata = {k: v for k, v in metadata.items() if k not in internal_keys} + filtered_metadata = { + k: v for k, v in metadata.items() if k not in internal_keys + } if filtered_metadata: # Only add if there's user metadata update_request["metadata"] = filtered_metadata @@ -730,7 +747,11 @@ def update_eval( # Transform request api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE - url, headers, request_body = evals_api_provider_config.transform_update_eval_request( + ( + url, + headers, + request_body, + ) = evals_api_provider_config.transform_update_eval_request( eval_id=eval_id, update_request=update_request, api_base=api_base, @@ -868,10 +889,10 @@ def delete_eval( custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ( - ProviderConfigManager.get_provider_evals_api_config( # type: ignore - provider=litellm.LlmProviders(custom_llm_provider), - ) + evals_api_provider_config: Optional[ + BaseEvalsAPIConfig + ] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + provider=litellm.LlmProviders(custom_llm_provider), ) if evals_api_provider_config is None: @@ -1021,10 +1042,10 @@ def cancel_eval( custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ( - ProviderConfigManager.get_provider_evals_api_config( # type: ignore - provider=litellm.LlmProviders(custom_llm_provider), - ) + evals_api_provider_config: Optional[ + BaseEvalsAPIConfig + ] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + provider=litellm.LlmProviders(custom_llm_provider), ) if evals_api_provider_config is None: @@ -1038,7 +1059,11 @@ def cancel_eval( # Transform request api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE - url, headers, request_body = evals_api_provider_config.transform_cancel_eval_request( + ( + url, + headers, + request_body, + ) = evals_api_provider_config.transform_cancel_eval_request( eval_id=eval_id, api_base=api_base, litellm_params=litellm_params, @@ -1199,16 +1224,14 @@ def create_run( custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ( - ProviderConfigManager.get_provider_evals_api_config( # type: ignore - provider=litellm.LlmProviders(custom_llm_provider), - ) + evals_api_provider_config: Optional[ + BaseEvalsAPIConfig + ] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + provider=litellm.LlmProviders(custom_llm_provider), ) if evals_api_provider_config is None: - raise ValueError( - f"CREATE run is not supported for {custom_llm_provider}" - ) + raise ValueError(f"CREATE run is not supported for {custom_llm_provider}") # Build create request create_request: CreateRunRequest = { @@ -1388,10 +1411,10 @@ def list_runs( custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ( - ProviderConfigManager.get_provider_evals_api_config( # type: ignore - provider=litellm.LlmProviders(custom_llm_provider), - ) + evals_api_provider_config: Optional[ + BaseEvalsAPIConfig + ] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + provider=litellm.LlmProviders(custom_llm_provider), ) if evals_api_provider_config is None: @@ -1561,10 +1584,10 @@ def get_run( custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ( - ProviderConfigManager.get_provider_evals_api_config( # type: ignore - provider=litellm.LlmProviders(custom_llm_provider), - ) + evals_api_provider_config: Optional[ + BaseEvalsAPIConfig + ] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + provider=litellm.LlmProviders(custom_llm_provider), ) if evals_api_provider_config is None: @@ -1720,10 +1743,10 @@ def cancel_run( custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ( - ProviderConfigManager.get_provider_evals_api_config( # type: ignore - provider=litellm.LlmProviders(custom_llm_provider), - ) + evals_api_provider_config: Optional[ + BaseEvalsAPIConfig + ] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + provider=litellm.LlmProviders(custom_llm_provider), ) if evals_api_provider_config is None: @@ -1737,7 +1760,11 @@ def cancel_run( # Transform request api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE - url, headers, request_body = evals_api_provider_config.transform_cancel_run_request( + ( + url, + headers, + request_body, + ) = evals_api_provider_config.transform_cancel_run_request( eval_id=eval_id, run_id=run_id, api_base=api_base, @@ -1884,10 +1911,10 @@ def delete_run( custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ( - ProviderConfigManager.get_provider_evals_api_config( # type: ignore - provider=litellm.LlmProviders(custom_llm_provider), - ) + evals_api_provider_config: Optional[ + BaseEvalsAPIConfig + ] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + provider=litellm.LlmProviders(custom_llm_provider), ) if evals_api_provider_config is None: @@ -1901,7 +1928,11 @@ def delete_run( # Transform request api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE - url, headers, request_body = evals_api_provider_config.transform_delete_run_request( + ( + url, + headers, + request_body, + ) = evals_api_provider_config.transform_delete_run_request( eval_id=eval_id, run_id=run_id, api_base=api_base, diff --git a/litellm/exceptions.py b/litellm/exceptions.py index b36d4ef877..abdba09dd8 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -25,9 +25,7 @@ def _get_minimal_error_response() -> httpx.Response: if _MINIMAL_ERROR_RESPONSE is None: _MINIMAL_ERROR_RESPONSE = httpx.Response( status_code=400, - request=httpx.Request( - method="GET", url="https://litellm.ai" - ), + request=httpx.Request(method="GET", url="https://litellm.ai"), ) return _MINIMAL_ERROR_RESPONSE @@ -996,7 +994,7 @@ class MidStreamFallbackError(ServiceUnavailableError): # type: ignore max_retries=self.max_retries, num_retries=self.num_retries, ) - + # Restore the propagated status and original response/request objects self.status_code = int(original_status) if original_status is not None else 503 self.response = _saved_response diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 30a1ac20d0..a638a28aba 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -4,7 +4,18 @@ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers. import asyncio import base64 -from typing import Any, Awaitable, Callable, Dict, Generator, List, Optional, Tuple, TypeVar, Union +from typing import ( + Any, + Awaitable, + Callable, + Dict, + Generator, + List, + Optional, + Tuple, + TypeVar, + Union, +) import httpx from mcp import ClientSession, ReadResourceResult, Resource, StdioServerParameters @@ -14,7 +25,10 @@ from mcp.client.stdio import stdio_client streamable_http_client: Optional[Any] = None try: import mcp.client.streamable_http as streamable_http_module # type: ignore - streamable_http_client = getattr(streamable_http_module, "streamable_http_client", None) + + streamable_http_client = getattr( + streamable_http_module, "streamable_http_client", None + ) except ImportError: pass from mcp.types import CallToolRequestParams as MCPCallToolRequestParams @@ -188,12 +202,15 @@ class MCPClient: if self.transport_type == MCPTransport.sse: headers = self._get_auth_headers() httpx_client_factory = self._create_httpx_client_factory() - return sse_client( - url=self.server_url, - timeout=self.timeout, - headers=headers, - httpx_client_factory=httpx_client_factory, - ), None + return ( + sse_client( + url=self.server_url, + timeout=self.timeout, + headers=headers, + httpx_client_factory=httpx_client_factory, + ), + None, + ) # HTTP transport (default) if streamable_http_client is None: @@ -201,12 +218,10 @@ class MCPClient: "streamable_http_client is not available. " "Please install mcp with HTTP support." ) - + headers = self._get_auth_headers() httpx_client_factory = self._create_httpx_client_factory() - verbose_logger.debug( - "litellm headers for streamable_http_client: %s", headers - ) + verbose_logger.debug("litellm headers for streamable_http_client: %s", headers) http_client = httpx_client_factory( headers=headers, timeout=httpx.Timeout(self.timeout), @@ -392,7 +407,7 @@ class MCPClient: async def call_tool( self, call_tool_request_params: MCPCallToolRequestParams, - host_progress_callback: Optional[Callable] = None + host_progress_callback: Optional[Callable] = None, ) -> MCPCallToolResult: """ Call an MCP Tool. @@ -401,13 +416,15 @@ class MCPClient: f"MCP client calling tool '{call_tool_request_params.name}' with arguments: {call_tool_request_params.arguments}" ) - async def on_progress(progress: float, total: float | None, message: str | None): + async def on_progress( + progress: float, total: float | None, message: str | None + ): percentage = (progress / total * 100) if total else 0 verbose_logger.info( f"MCP Tool '{call_tool_request_params.name}' progress: " f"{progress}/{total} ({percentage:.0f}%) - {message or ''}" ) - + # Forward to Host if callback provided if host_progress_callback: try: @@ -421,8 +438,8 @@ class MCPClient: name=call_tool_request_params.name, arguments=call_tool_request_params.arguments, progress_callback=on_progress, - ) + try: tool_result = await self.run_with_session(_call_tool_operation) verbose_logger.info( diff --git a/litellm/experimental_mcp_client/tools.py b/litellm/experimental_mcp_client/tools.py index b716e3171e..bd42f7e711 100644 --- a/litellm/experimental_mcp_client/tools.py +++ b/litellm/experimental_mcp_client/tools.py @@ -18,7 +18,7 @@ from litellm.types.utils import ChatCompletionMessageToolCall def transform_mcp_tool_to_openai_tool(mcp_tool: MCPTool) -> ChatCompletionToolParam: """Convert an MCP tool to an OpenAI tool.""" normalized_parameters = _normalize_mcp_input_schema(mcp_tool.inputSchema) - + return ChatCompletionToolParam( type="function", function=FunctionDefinition( @@ -33,41 +33,39 @@ def transform_mcp_tool_to_openai_tool(mcp_tool: MCPTool) -> ChatCompletionToolPa def _normalize_mcp_input_schema(input_schema: dict) -> dict: """ Normalize MCP input schema to ensure it's valid for OpenAI function calling. - + OpenAI requires that function parameters have: - type: 'object' - properties: dict (can be empty) - additionalProperties: false (recommended) """ if not input_schema: - return { - "type": "object", - "properties": {}, - "additionalProperties": False - } - + return {"type": "object", "properties": {}, "additionalProperties": False} + # Make a copy to avoid modifying the original normalized_schema = dict(input_schema) - + # Ensure type is 'object' if "type" not in normalized_schema: normalized_schema["type"] = "object" - + # Ensure properties exists (can be empty) if "properties" not in normalized_schema: normalized_schema["properties"] = {} - + # Add additionalProperties if not present (recommended by OpenAI) if "additionalProperties" not in normalized_schema: normalized_schema["additionalProperties"] = False - + return normalized_schema -def transform_mcp_tool_to_openai_responses_api_tool(mcp_tool: MCPTool) -> FunctionToolParam: +def transform_mcp_tool_to_openai_responses_api_tool( + mcp_tool: MCPTool, +) -> FunctionToolParam: """Convert an MCP tool to an OpenAI Responses API tool.""" normalized_parameters = _normalize_mcp_input_schema(mcp_tool.inputSchema) - + return FunctionToolParam( name=mcp_tool.name, parameters=normalized_parameters, @@ -76,6 +74,7 @@ def transform_mcp_tool_to_openai_responses_api_tool(mcp_tool: MCPTool) -> Functi description=mcp_tool.description or "", ) + async def load_mcp_tools( session: ClientSession, format: Literal["mcp", "openai"] = "mcp" ) -> Union[List[MCPTool], List[ChatCompletionToolParam]]: diff --git a/litellm/files/main.py b/litellm/files/main.py index 2a10789e74..f7c89e0ba3 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -14,11 +14,30 @@ from typing import Any, Coroutine, Dict, Literal, Optional, Union, cast import httpx +# Type aliases for provider parameters +FileCreateProvider = Literal[ + "openai", + "azure", + "gemini", + "vertex_ai", + "bedrock", + "hosted_vllm", + "manus", + "anthropic", +] +FileRetrieveProvider = Literal[ + "openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "manus", "anthropic" +] +FileDeleteProvider = Literal["openai", "azure", "gemini", "manus", "anthropic"] +FileListProvider = Literal["openai", "azure", "manus", "anthropic"] +FileContentProvider = Literal[ + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic", "manus" +] + import litellm from litellm import get_secret_str from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.llms.anthropic.files.handler import AnthropicFilesHandler from litellm.llms.azure.common_utils import get_azure_credentials from litellm.llms.azure.files.handler import AzureOpenAIFilesAPI from litellm.llms.bedrock.files.handler import BedrockFilesHandler @@ -54,16 +73,15 @@ openai_files_instance = OpenAIFilesAPI() azure_files_instance = AzureOpenAIFilesAPI() vertex_ai_files_instance = VertexAIFilesHandler() bedrock_files_instance = BedrockFilesHandler() -anthropic_files_instance = AnthropicFilesHandler() ################################################# @client async def acreate_file( file: FileTypes, - purpose: Literal["assistants", "batch", "fine-tune"], + purpose: Literal["assistants", "batch", "fine-tune", "messages"], expires_after: Optional[FileExpiresAfter] = None, - custom_llm_provider: Literal["openai", "azure", "gemini", "vertex_ai", "bedrock", "hosted_vllm", "manus"] = "openai", + custom_llm_provider: FileCreateProvider = "openai", extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, @@ -106,9 +124,9 @@ async def acreate_file( @client def create_file( file: FileTypes, - purpose: Literal["assistants", "batch", "fine-tune"], + purpose: Literal["assistants", "batch", "fine-tune", "messages"], expires_after: Optional[FileExpiresAfter] = None, - custom_llm_provider: Optional[Literal["openai", "azure", "gemini", "vertex_ai", "bedrock", "hosted_vllm", "manus"]] = None, + custom_llm_provider: Optional[FileCreateProvider] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, @@ -218,7 +236,7 @@ def create_file( ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'create_file'. Only ['openai', 'azure', 'vertex_ai', 'manus'] are supported.".format( + message="LiteLLM doesn't support {} for 'create_file'. Only ['openai', 'azure', 'vertex_ai', 'manus', 'anthropic'] are supported.".format( custom_llm_provider ), model="n/a", @@ -237,7 +255,7 @@ def create_file( @client async def afile_retrieve( file_id: str, - custom_llm_provider: Literal["openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "manus"] = "openai", + custom_llm_provider: FileRetrieveProvider = "openai", extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, @@ -278,7 +296,7 @@ async def afile_retrieve( @client def file_retrieve( file_id: str, - custom_llm_provider: Literal["openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "manus"] = "openai", + custom_llm_provider: FileRetrieveProvider = "openai", extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, @@ -348,22 +366,25 @@ def file_retrieve( litellm_params_dict = get_litellm_params(**kwargs) litellm_params_dict["api_key"] = optional_params.api_key litellm_params_dict["api_base"] = optional_params.api_base - + logging_obj = kwargs.get("litellm_logging_obj") if logging_obj is None: from litellm.litellm_core_utils.litellm_logging import ( Logging as LiteLLMLoggingObj, ) + logging_obj = LiteLLMLoggingObj( model="", messages=[], stream=False, call_type="afile_retrieve" if _is_async else "file_retrieve", start_time=time.time(), - litellm_call_id=kwargs.get("litellm_call_id", str(uuid_module.uuid4())), + litellm_call_id=kwargs.get( + "litellm_call_id", str(uuid_module.uuid4()) + ), function_id=str(kwargs.get("id") or ""), ) - + client = kwargs.get("client") response = base_llm_http_handler.retrieve_file( file_id=file_id, @@ -382,7 +403,7 @@ def file_retrieve( ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'file_retrieve'. Only 'openai', 'azure', and 'manus' are supported.".format( + message="LiteLLM doesn't support {} for 'file_retrieve'. Only 'openai', 'azure', 'manus', and 'anthropic' are supported.".format( custom_llm_provider ), model="n/a", @@ -403,7 +424,7 @@ def file_retrieve( @client async def afile_delete( file_id: str, - custom_llm_provider: Literal["openai", "azure", "gemini", "manus"] = "openai", + custom_llm_provider: FileDeleteProvider = "openai", extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, @@ -447,7 +468,7 @@ async def afile_delete( def file_delete( file_id: str, model: Optional[str] = None, - custom_llm_provider: Union[Literal["openai", "azure", "gemini", "manus"], str] = "openai", + custom_llm_provider: Union[FileDeleteProvider, str] = "openai", extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, @@ -525,22 +546,25 @@ def file_delete( if provider_config is not None: litellm_params_dict["api_key"] = optional_params.api_key litellm_params_dict["api_base"] = optional_params.api_base - + logging_obj = kwargs.get("litellm_logging_obj") if logging_obj is None: from litellm.litellm_core_utils.litellm_logging import ( Logging as LiteLLMLoggingObj, ) + logging_obj = LiteLLMLoggingObj( model="", messages=[], stream=False, call_type="afile_delete" if _is_async else "file_delete", start_time=time.time(), - litellm_call_id=kwargs.get("litellm_call_id", str(uuid_module.uuid4())), + litellm_call_id=kwargs.get( + "litellm_call_id", str(uuid_module.uuid4()) + ), function_id=str(kwargs.get("id") or ""), ) - + response = base_llm_http_handler.delete_file( file_id=file_id, provider_config=provider_config, @@ -558,7 +582,7 @@ def file_delete( ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'file_delete'. Only 'openai', 'azure', 'gemini', and 'manus' are supported.".format( + message="LiteLLM doesn't support {} for 'file_delete'. Only 'openai', 'azure', 'gemini', 'manus', and 'anthropic' are supported.".format( custom_llm_provider ), model="n/a", @@ -577,7 +601,7 @@ def file_delete( # List files @client async def afile_list( - custom_llm_provider: Literal["openai", "azure", "manus"] = "openai", + custom_llm_provider: FileListProvider = "openai", purpose: Optional[str] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -618,7 +642,7 @@ async def afile_list( @client def file_list( - custom_llm_provider: Literal["openai", "azure", "manus"] = "openai", + custom_llm_provider: FileListProvider = "openai", purpose: Optional[str] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -648,7 +672,7 @@ def file_list( timeout = 600.0 _is_async = kwargs.pop("is_async", False) is True - + # Check if provider has a custom files config (e.g., Manus, Bedrock, Vertex AI) provider_config = ProviderConfigManager.get_provider_files_config( model="", @@ -658,22 +682,25 @@ def file_list( litellm_params_dict = get_litellm_params(**kwargs) litellm_params_dict["api_key"] = optional_params.api_key litellm_params_dict["api_base"] = optional_params.api_base - + logging_obj = kwargs.get("litellm_logging_obj") if logging_obj is None: from litellm.litellm_core_utils.litellm_logging import ( Logging as LiteLLMLoggingObj, ) + logging_obj = LiteLLMLoggingObj( model="", messages=[], stream=False, call_type="afile_list" if _is_async else "file_list", start_time=time.time(), - litellm_call_id=kwargs.get("litellm_call_id", str(uuid_module.uuid4())), + litellm_call_id=kwargs.get( + "litellm_call_id", str(uuid_module.uuid4()) + ), function_id=str(kwargs.get("id", "")), ) - + client = kwargs.get("client") response = base_llm_http_handler.list_files( purpose=purpose, @@ -723,7 +750,7 @@ def file_list( ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'file_list'. Only 'openai', 'azure', and 'manus' are supported.".format( + message="LiteLLM doesn't support {} for 'file_list'. Only 'openai', 'azure', 'manus', and 'anthropic' are supported.".format( custom_llm_provider ), model="n/a", @@ -742,7 +769,7 @@ def file_list( @client async def afile_content( file_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic", "manus"] = "openai", + custom_llm_provider: FileContentProvider = "openai", extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, @@ -786,9 +813,7 @@ async def afile_content( def file_content( file_id: str, model: Optional[str] = None, - custom_llm_provider: Optional[ - Union[Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic", "manus"], str] - ] = None, + custom_llm_provider: Optional[Union[FileContentProvider, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, @@ -834,15 +859,43 @@ def file_content( _is_async = kwargs.pop("afile_content", False) is True - # Check if this is an Anthropic batch results request - if custom_llm_provider == "anthropic": - response = anthropic_files_instance.file_content( - _is_async=_is_async, + # Check if provider has a custom files config (e.g., Anthropic, Manus) + provider_config = ProviderConfigManager.get_provider_files_config( + model="", + provider=LlmProviders(custom_llm_provider), + ) + if provider_config is not None: + litellm_params_dict["api_key"] = optional_params.api_key + litellm_params_dict["api_base"] = optional_params.api_base + + logging_obj = kwargs.get("litellm_logging_obj") + if logging_obj is None: + logging_obj = LiteLLMLoggingObj( + model="", + messages=[], + stream=False, + call_type="afile_content" if _is_async else "file_content", + start_time=time.time(), + litellm_call_id=kwargs.get( + "litellm_call_id", str(uuid_module.uuid4()) + ), + function_id=str(kwargs.get("id") or ""), + ) + + response = base_llm_http_handler.retrieve_file_content( file_content_request=_file_content_request, - api_base=optional_params.api_base, - api_key=optional_params.api_key, + provider_config=provider_config, + litellm_params=litellm_params_dict, + headers=extra_headers or {}, + logging_obj=logging_obj, + _is_async=_is_async, + client=( + client + if client is not None + and isinstance(client, (HTTPHandler, AsyncHTTPHandler)) + else None + ), timeout=timeout, - max_retries=optional_params.max_retries, ) return response @@ -915,7 +968,7 @@ def file_content( ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'file_content'. Supported providers are 'openai', 'azure', 'vertex_ai', 'bedrock', 'manus'.".format( + message="LiteLLM doesn't support {} for 'file_content'. Supported providers are 'openai', 'azure', 'vertex_ai', 'bedrock', 'manus', 'anthropic'.".format( custom_llm_provider ), model="n/a", diff --git a/litellm/files/utils.py b/litellm/files/utils.py index a56a29467d..a2b9a42c15 100644 --- a/litellm/files/utils.py +++ b/litellm/files/utils.py @@ -8,17 +8,22 @@ class FilesAPIUtils: """ Utils for files API interface on litellm """ + @staticmethod - def is_batch_jsonl_file(create_file_data: CreateFileRequest, extracted_file_data: ExtractedFileData) -> bool: + def is_batch_jsonl_file( + create_file_data: CreateFileRequest, extracted_file_data: ExtractedFileData + ) -> bool: """ Check if the file is a batch jsonl file """ return ( create_file_data.get("purpose") == "batch" - and FilesAPIUtils.valid_content_type(extracted_file_data.get("content_type")) + and FilesAPIUtils.valid_content_type( + extracted_file_data.get("content_type") + ) and extracted_file_data.get("content") is not None ) - + @staticmethod def valid_content_type(content_type: Optional[str]) -> bool: """ diff --git a/litellm/fine_tuning/main.py b/litellm/fine_tuning/main.py index 93fa56ff97..08373cda78 100644 --- a/litellm/fine_tuning/main.py +++ b/litellm/fine_tuning/main.py @@ -41,34 +41,34 @@ def _prepare_azure_extra_body( ) -> Dict[str, Any]: """ Prepare extra_body for Azure fine-tuning API by combining Azure-specific parameters. - + Azure fine-tuning API accepts additional parameters beyond the standard OpenAI spec: - trainingType: Type of training (e.g., 1 for supervised fine-tuning) - prompt_loss_weight: Weight for prompt loss in training - + These parameters must be passed in the extra_body field when calling the Azure OpenAI SDK. - + Args: extra_body: Optional existing extra_body dict kwargs: Request kwargs that may contain Azure-specific parameters azure_specific_hyperparams: Dict of Azure-specific hyperparameters already extracted - + Returns: Dict containing all Azure-specific parameters to be passed in extra_body """ if extra_body is None: extra_body = {} - + # Azure-specific root-level parameters azure_specific_params = ["trainingType"] for param in azure_specific_params: if param in kwargs: extra_body[param] = kwargs[param] - + # Add Azure-specific hyperparameters if azure_specific_hyperparams: extra_body.update(azure_specific_hyperparams) - + return extra_body @@ -126,7 +126,9 @@ async def acreate_fine_tuning_job( raise e -def _build_fine_tuning_job_data(model, training_file, hyperparameters, suffix, validation_file, integrations, seed): +def _build_fine_tuning_job_data( + model, training_file, hyperparameters, suffix, validation_file, integrations, seed +): return FineTuningJobCreate( model=model, training_file=training_file, @@ -177,7 +179,7 @@ def create_fine_tuning_job( # handle hyperparameters hyperparameters = hyperparameters or {} # original hyperparameters - + # For Azure, extract Azure-specific hyperparameters before creating OpenAI-spec hyperparameters azure_specific_hyperparams = {} if custom_llm_provider == "azure": @@ -185,7 +187,7 @@ def create_fine_tuning_job( for key in azure_hyperparameter_keys: if key in hyperparameters: azure_specific_hyperparams[key] = hyperparameters.pop(key) - + _oai_hyperparameters: Hyperparameters = Hyperparameters( **hyperparameters ) # Typed Hyperparameters for OpenAI Spec @@ -219,7 +221,13 @@ def create_fine_tuning_job( ) create_fine_tuning_job_data_dict = _build_fine_tuning_job_data( - model, training_file, _oai_hyperparameters, suffix, validation_file, integrations, seed, + model, + training_file, + _oai_hyperparameters, + suffix, + validation_file, + integrations, + seed, ).model_dump(exclude_none=True) response = openai_fine_tuning_apis_instance.create_fine_tuning_job( @@ -258,12 +266,20 @@ def create_fine_tuning_job( extra_body.pop("azure_ad_token", None) else: get_secret_str("AZURE_AD_TOKEN") # type: ignore - + # Prepare Azure-specific parameters for extra_body - extra_body = _prepare_azure_extra_body(extra_body, kwargs, azure_specific_hyperparams) - + extra_body = _prepare_azure_extra_body( + extra_body, kwargs, azure_specific_hyperparams + ) + create_fine_tuning_job_data_dict = _build_fine_tuning_job_data( - model, training_file, _oai_hyperparameters, suffix, validation_file, integrations, seed, + model, + training_file, + _oai_hyperparameters, + suffix, + validation_file, + integrations, + seed, ).model_dump(exclude_none=True) # Add extra_body if it has Azure-specific parameters @@ -298,7 +314,13 @@ def create_fine_tuning_job( response = vertex_fine_tuning_apis_instance.create_fine_tuning_job( _is_async=_is_async, create_fine_tuning_job_data=_build_fine_tuning_job_data( - model, training_file, _oai_hyperparameters, suffix, validation_file, integrations, seed, + model, + training_file, + _oai_hyperparameters, + suffix, + validation_file, + integrations, + seed, ), vertex_credentials=vertex_credentials, vertex_project=vertex_ai_project, diff --git a/litellm/google_genai/__init__.py b/litellm/google_genai/__init__.py index faeb1f227d..ca7b547c44 100644 --- a/litellm/google_genai/__init__.py +++ b/litellm/google_genai/__init__.py @@ -13,7 +13,7 @@ from .main import ( __all__ = [ "generate_content", - "agenerate_content", + "agenerate_content", "generate_content_stream", "agenerate_content_stream", -] \ No newline at end of file +] diff --git a/litellm/google_genai/adapters/__init__.py b/litellm/google_genai/adapters/__init__.py index 96ff777ebe..bfa9e71267 100644 --- a/litellm/google_genai/adapters/__init__.py +++ b/litellm/google_genai/adapters/__init__.py @@ -13,7 +13,7 @@ from .handler import GenerateContentToCompletionHandler from .transformation import GoogleGenAIAdapter, GoogleGenAIStreamWrapper __all__ = [ - "GoogleGenAIAdapter", + "GoogleGenAIAdapter", "GoogleGenAIStreamWrapper", - "GenerateContentToCompletionHandler" -] \ No newline at end of file + "GenerateContentToCompletionHandler", +] diff --git a/litellm/google_genai/main.py b/litellm/google_genai/main.py index 9ec56c3717..a937a35da2 100644 --- a/litellm/google_genai/main.py +++ b/litellm/google_genai/main.py @@ -168,7 +168,9 @@ class GenerateContentHelper: ) ) # Extract systemInstruction from kwargs to pass to transform - system_instruction = kwargs.get("systemInstruction") or kwargs.get("system_instruction") + system_instruction = kwargs.get("systemInstruction") or kwargs.get( + "system_instruction" + ) request_body = ( generate_content_provider_config.transform_generate_content_request( model=model, @@ -318,7 +320,9 @@ def generate_content( ) # Extract systemInstruction from kwargs to pass to handler - system_instruction = kwargs.get("systemInstruction") or kwargs.get("system_instruction") + system_instruction = kwargs.get("systemInstruction") or kwargs.get( + "system_instruction" + ) # Check if we should use the adapter (when provider config is None) if setup_result.generate_content_provider_config is None: @@ -407,7 +411,9 @@ async def agenerate_content_stream( ) # Extract systemInstruction from kwargs to pass to handler - system_instruction = kwargs.get("systemInstruction") or kwargs.get("system_instruction") + system_instruction = kwargs.get("systemInstruction") or kwargs.get( + "system_instruction" + ) # Check if we should use the adapter (when provider config is None) if setup_result.generate_content_provider_config is None: diff --git a/litellm/google_genai/streaming_iterator.py b/litellm/google_genai/streaming_iterator.py index d0fa5a0be6..8cb2ee0937 100644 --- a/litellm/google_genai/streaming_iterator.py +++ b/litellm/google_genai/streaming_iterator.py @@ -17,6 +17,7 @@ else: GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ = PassThroughEndpointLogging() + class BaseGoogleGenAIGenerateContentStreamingIterator: """ Base class for Google GenAI Generate Content streaming iterators that provides common logic @@ -42,6 +43,7 @@ class BaseGoogleGenAIGenerateContentStreamingIterator: from litellm.proxy.pass_through_endpoints.streaming_handler import ( PassThroughStreamingHandler, ) + end_time = datetime.now() asyncio.create_task( PassThroughStreamingHandler._route_streaming_logging_to_handler( @@ -58,7 +60,9 @@ class BaseGoogleGenAIGenerateContentStreamingIterator: ) -class GoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateContentStreamingIterator): +class GoogleGenAIGenerateContentStreamingIterator( + BaseGoogleGenAIGenerateContentStreamingIterator +): """ Streaming iterator specifically for Google GenAI generate content API. """ @@ -105,10 +109,14 @@ class GoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateContent async def __anext__(self): # This should not be used for sync responses # If you need async iteration, use AsyncGoogleGenAIGenerateContentStreamingIterator - raise NotImplementedError("Use AsyncGoogleGenAIGenerateContentStreamingIterator for async iteration") + raise NotImplementedError( + "Use AsyncGoogleGenAIGenerateContentStreamingIterator for async iteration" + ) -class AsyncGoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateContentStreamingIterator): +class AsyncGoogleGenAIGenerateContentStreamingIterator( + BaseGoogleGenAIGenerateContentStreamingIterator +): """ Async streaming iterator specifically for Google GenAI generate content API. """ @@ -148,4 +156,4 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateCo return chunk except StopAsyncIteration: await self._handle_async_streaming_logging() - raise StopAsyncIteration \ No newline at end of file + raise StopAsyncIteration diff --git a/litellm/images/main.py b/litellm/images/main.py index 553aa26da9..a3ae97b57d 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -86,7 +86,6 @@ def _get_ImageEditRequestUtils() -> "ImageEditRequestUtils": return _ImageEditRequestUtils_cache - ##### Image Generation ####################### @client async def aimage_generation(*args, **kwargs) -> ImageResponse: @@ -212,10 +211,7 @@ def image_generation( # noqa: PLR0915 api_version: Optional[str] = None, custom_llm_provider=None, **kwargs, -) -> Union[ - ImageResponse, - Coroutine[Any, Any, ImageResponse], -]: +) -> Union[ImageResponse, Coroutine[Any, Any, ImageResponse],]: """ Maps the https://api.openai.com/v1/images/generations endpoint. @@ -346,7 +342,7 @@ def image_generation( # noqa: PLR0915 azure_ad_token = optional_params.pop( "azure_ad_token", None ) or get_secret_str("AZURE_AD_TOKEN") - + # Create azure_ad_token_provider from tenant_id, client_id, client_secret if not already provided if azure_ad_token_provider is None: from litellm.llms.azure.common_utils import ( @@ -357,8 +353,11 @@ def image_generation( # noqa: PLR0915 tenant_id = litellm_params_dict.get("tenant_id") client_id = litellm_params_dict.get("client_id") client_secret = litellm_params_dict.get("client_secret") - azure_scope = litellm_params_dict.get("azure_scope") or "https://cognitiveservices.azure.com/.default" - + azure_scope = ( + litellm_params_dict.get("azure_scope") + or "https://cognitiveservices.azure.com/.default" + ) + # Create token provider if credentials are available if tenant_id and client_id and client_secret: azure_ad_token_provider = get_azure_ad_token_from_entra_id( @@ -375,7 +374,7 @@ def image_generation( # noqa: PLR0915 # Azure AD authentication will use Authorization header instead if api_key is not None: default_headers["api-key"] = api_key - + for k, v in default_headers.items(): if k not in headers: headers[k] = v @@ -462,7 +461,7 @@ def image_generation( # noqa: PLR0915 # Azure AD authentication will use Authorization header instead if api_key is not None: default_headers["api-key"] = api_key - + for k, v in default_headers.items(): if k not in headers: headers[k] = v @@ -738,7 +737,7 @@ def image_variation( @client def image_edit( # noqa: PLR0915 image: Optional[Union[FileTypes, List[FileTypes]]] = None, - prompt: Optional[str]= None, + prompt: Optional[str] = None, model: Optional[str] = None, mask: Optional[str] = None, n: Optional[int] = None, @@ -762,23 +761,23 @@ def image_edit( # noqa: PLR0915 local_vars = locals() try: openai_params = [ - "user", - "request_timeout", - "api_base", - "api_version", - "api_key", - "deployment_id", - "organization", - "base_url", - "default_headers", - "timeout", - "max_retries", - "n", - "quality", - "size", - "style", - "async_call", - ] + "user", + "request_timeout", + "api_base", + "api_version", + "api_key", + "deployment_id", + "organization", + "base_url", + "default_headers", + "timeout", + "max_retries", + "n", + "quality", + "size", + "style", + "async_call", + ] litellm_params_list = all_litellm_params default_params = openai_params + litellm_params_list non_default_params = { @@ -791,7 +790,9 @@ def image_edit( # noqa: PLR0915 _is_async = kwargs.pop("async_call", False) is True # add images / or return a single image - images = image if isinstance(image, list) else ([image] if image is not None else []) + images = ( + image if isinstance(image, list) else ([image] if image is not None else []) + ) headers_from_kwargs = kwargs.get("headers") merged_extra_headers: Dict[str, Any] = {} @@ -864,11 +865,11 @@ def image_edit( # noqa: PLR0915 ) # get provider config - image_edit_provider_config: Optional[BaseImageEditConfig] = ( - ProviderConfigManager.get_provider_image_edit_config( - model=model, - provider=litellm.LlmProviders(custom_llm_provider), - ) + image_edit_provider_config: Optional[ + BaseImageEditConfig + ] = ProviderConfigManager.get_provider_image_edit_config( + model=model, + provider=litellm.LlmProviders(custom_llm_provider), ) if image_edit_provider_config is None: @@ -877,7 +878,9 @@ def image_edit( # noqa: PLR0915 local_vars.update(kwargs) # Get ImageEditOptionalRequestParams with only valid parameters image_edit_optional_params: ImageEditOptionalRequestParams = ( - _get_ImageEditRequestUtils().get_requested_image_edit_optional_param(local_vars) + _get_ImageEditRequestUtils().get_requested_image_edit_optional_param( + local_vars + ) ) # Get optional parameters for the responses API image_edit_request_params: Dict = ( @@ -926,20 +929,20 @@ def image_edit( # noqa: PLR0915 elif custom_llm_provider == "stability": image_edit_request_params.update(non_default_params) return base_llm_http_handler.image_edit_handler( - model=model, - image=images, - prompt=prompt, - image_edit_provider_config=image_edit_provider_config, - image_edit_optional_request_params=image_edit_request_params, - custom_llm_provider=custom_llm_provider, - litellm_params=litellm_params, - logging_obj=litellm_logging_obj, - extra_headers=extra_headers, - extra_body=extra_body, - timeout=timeout or DEFAULT_REQUEST_TIMEOUT, - _is_async=_is_async, - client=kwargs.get("client"), - ) + model=model, + image=images, + prompt=prompt, + image_edit_provider_config=image_edit_provider_config, + image_edit_optional_request_params=image_edit_request_params, + custom_llm_provider=custom_llm_provider, + litellm_params=litellm_params, + logging_obj=litellm_logging_obj, + extra_headers=extra_headers, + extra_body=extra_body, + timeout=timeout or DEFAULT_REQUEST_TIMEOUT, + _is_async=_is_async, + client=kwargs.get("client"), + ) elif custom_llm_provider == "black_forest_labs": # Route to BFL-specific handler (polling required) if model is None: diff --git a/litellm/images/utils.py b/litellm/images/utils.py index fa271b61b6..8d3e96f143 100644 --- a/litellm/images/utils.py +++ b/litellm/images/utils.py @@ -40,9 +40,7 @@ class ImageEditRequestUtils: filtered_optional_params.pop(param, None) unsupported_params = [ - param - for param in filtered_optional_params - if param not in supported_params + param for param in filtered_optional_params if param not in supported_params ] if unsupported_params: diff --git a/litellm/integrations/SlackAlerting/hanging_request_check.py b/litellm/integrations/SlackAlerting/hanging_request_check.py index d2f70c9caf..b9c485dce8 100644 --- a/litellm/integrations/SlackAlerting/hanging_request_check.py +++ b/litellm/integrations/SlackAlerting/hanging_request_check.py @@ -102,10 +102,10 @@ class AlertingHangingRequestCheck: ) for request_id in hanging_requests: - hanging_request_data: Optional[HangingRequestData] = ( - await self.hanging_request_cache.async_get_cache( - key=request_id, - ) + hanging_request_data: Optional[ + HangingRequestData + ] = await self.hanging_request_cache.async_get_cache( + key=request_id, ) if hanging_request_data is None: diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 35634d5067..013cef7480 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -96,7 +96,9 @@ class SlackAlerting(CustomBatchLogger): self.alert_type_config: Dict[str, AlertTypeConfig] = {} if alert_type_config: for key, val in alert_type_config.items(): - self.alert_type_config[key] = AlertTypeConfig(**val) if isinstance(val, dict) else val + self.alert_type_config[key] = ( + AlertTypeConfig(**val) if isinstance(val, dict) else val + ) self.digest_buckets: Dict[str, DigestEntry] = {} self.digest_lock = asyncio.Lock() super().__init__(**kwargs, flush_lock=self.flush_lock) @@ -126,7 +128,9 @@ class SlackAlerting(CustomBatchLogger): self.periodic_started = True if alert_type_config is not None: for key, val in alert_type_config.items(): - self.alert_type_config[key] = AlertTypeConfig(**val) if isinstance(val, dict) else val + self.alert_type_config[key] = ( + AlertTypeConfig(**val) if isinstance(val, dict) else val + ) if alert_to_webhook_url is not None: # update the dict @@ -1367,7 +1371,7 @@ Model Info: return False - async def send_alert( # noqa: PLR0915 + async def send_alert( # noqa: PLR0915 self, message: str, level: Literal["Low", "Medium", "High"], @@ -1439,7 +1443,9 @@ Model Info: self.alert_to_webhook_url is not None and alert_type in self.alert_to_webhook_url ): - _digest_webhook: Optional[Union[str, List[str]]] = self.alert_to_webhook_url[alert_type] + _digest_webhook: Optional[ + Union[str, List[str]] + ] = self.alert_to_webhook_url[alert_type] elif self.default_webhook_url is not None: _digest_webhook = self.default_webhook_url else: @@ -1588,11 +1594,21 @@ Model Info: if isinstance(webhook_url, list): for url in webhook_url: self.log_queue.append( - {"url": url, "headers": headers, "payload": payload, "alert_type": alert_type_name} + { + "url": url, + "headers": headers, + "payload": payload, + "alert_type": alert_type_name, + } ) else: self.log_queue.append( - {"url": webhook_url, "headers": headers, "payload": payload, "alert_type": alert_type_name} + { + "url": webhook_url, + "headers": headers, + "payload": payload, + "alert_type": alert_type_name, + } ) flushed_keys.append(key) diff --git a/litellm/integrations/_types/open_inference.py b/litellm/integrations/_types/open_inference.py index 0fde1ff752..3404df7495 100644 --- a/litellm/integrations/_types/open_inference.py +++ b/litellm/integrations/_types/open_inference.py @@ -73,11 +73,15 @@ class SpanAttributes: """ Number of tokens in the prompt. """ - LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_WRITE = "llm.token_count.prompt_details.cache_write" + LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_WRITE = ( + "llm.token_count.prompt_details.cache_write" + ) """ Number of tokens in the prompt that were written to cache. """ - LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ = "llm.token_count.prompt_details.cache_read" + LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ = ( + "llm.token_count.prompt_details.cache_read" + ) """ Number of tokens in the prompt that were read from cache. """ @@ -89,11 +93,15 @@ class SpanAttributes: """ Number of tokens in the completion. """ - LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING = "llm.token_count.completion_details.reasoning" + LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING = ( + "llm.token_count.completion_details.reasoning" + ) """ Number of tokens used for reasoning steps in the completion. """ - LLM_TOKEN_COUNT_COMPLETION_DETAILS_AUDIO = "llm.token_count.completion_details.audio" + LLM_TOKEN_COUNT_COMPLETION_DETAILS_AUDIO = ( + "llm.token_count.completion_details.audio" + ) """ The number of audio input tokens generated by the model """ @@ -396,7 +404,7 @@ class OpenInferenceLLMProviderValues(Enum): class ErrorAttributes: """ Attributes for error information in spans. - + These attributes follow OpenTelemetry semantic conventions for exceptions and are used to record error information from StandardLoggingPayloadErrorInformation. """ diff --git a/litellm/integrations/agentops/__init__.py b/litellm/integrations/agentops/__init__.py index 6ad02ce0ba..003a12a611 100644 --- a/litellm/integrations/agentops/__init__.py +++ b/litellm/integrations/agentops/__init__.py @@ -1,3 +1,3 @@ from .agentops import AgentOps -__all__ = ["AgentOps"] \ No newline at end of file +__all__ = ["AgentOps"] diff --git a/litellm/integrations/agentops/agentops.py b/litellm/integrations/agentops/agentops.py index 11e76841e9..38b91c0658 100644 --- a/litellm/integrations/agentops/agentops.py +++ b/litellm/integrations/agentops/agentops.py @@ -7,6 +7,7 @@ from typing import Optional, Dict, Any from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig from litellm.llms.custom_httpx.http_handler import _get_httpx_client + @dataclass class AgentOpsConfig: endpoint: str = "https://otlp.agentops.cloud/v1/traces" @@ -22,9 +23,10 @@ class AgentOpsConfig: api_key=os.getenv("AGENTOPS_API_KEY"), service_name=os.getenv("AGENTOPS_SERVICE_NAME", "agentops"), deployment_environment=os.getenv("AGENTOPS_ENVIRONMENT", "production"), - auth_endpoint="https://api.agentops.ai/v3/auth/token" + auth_endpoint="https://api.agentops.ai/v3/auth/token", ) + class AgentOps(OpenTelemetry): """ AgentOps integration - built on top of OpenTelemetry @@ -32,7 +34,7 @@ class AgentOps(OpenTelemetry): Example usage: ```python import litellm - + litellm.success_callback = ["agentops"] response = litellm.completion( @@ -41,6 +43,7 @@ class AgentOps(OpenTelemetry): ) ``` """ + def __init__( self, config: Optional[AgentOpsConfig] = None, @@ -60,18 +63,13 @@ class AgentOps(OpenTelemetry): pass headers = f"Authorization=Bearer {jwt_token}" if jwt_token else None - + otel_config = OpenTelemetryConfig( - exporter="otlp_http", - endpoint=config.endpoint, - headers=headers + exporter="otlp_http", endpoint=config.endpoint, headers=headers ) # Initialize OpenTelemetry with our config - super().__init__( - config=otel_config, - callback_name="agentops" - ) + super().__init__(config=otel_config, callback_name="agentops") # Set AgentOps-specific resource attributes resource_attrs = { @@ -79,20 +77,20 @@ class AgentOps(OpenTelemetry): "deployment.environment": config.deployment_environment or "production", "telemetry.sdk.name": "agentops", } - + if project_id: resource_attrs["project.id"] = project_id - + self.resource_attributes = resource_attrs def _fetch_auth_token(self, api_key: str, auth_endpoint: str) -> Dict[str, Any]: """ Fetch JWT authentication token from AgentOps API - + Args: api_key: AgentOps API key auth_endpoint: Authentication endpoint - + Returns: Dict containing JWT token and project ID """ @@ -100,19 +98,19 @@ class AgentOps(OpenTelemetry): "Content-Type": "application/json", "Connection": "keep-alive", } - + client = _get_httpx_client() try: response = client.post( url=auth_endpoint, headers=headers, json={"api_key": api_key}, - timeout=10 + timeout=10, ) - + if response.status_code != 200: raise Exception(f"Failed to fetch auth token: {response.text}") - + return response.json() finally: - client.close() \ No newline at end of file + client.close() diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index 67b95c7694..8e4d40c460 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -99,10 +99,10 @@ class AnthropicCacheControlHook(CustomPromptManagement): targetted_index += len(messages) if 0 <= targetted_index < len(messages): - messages[targetted_index] = ( - AnthropicCacheControlHook._safe_insert_cache_control_in_message( - messages[targetted_index], control - ) + messages[ + targetted_index + ] = AnthropicCacheControlHook._safe_insert_cache_control_in_message( + messages[targetted_index], control ) else: verbose_logger.warning( diff --git a/litellm/integrations/arize/_utils.py b/litellm/integrations/arize/_utils.py index b75e296be4..8dfaa8b142 100644 --- a/litellm/integrations/arize/_utils.py +++ b/litellm/integrations/arize/_utils.py @@ -14,12 +14,12 @@ from litellm.types.utils import StandardLoggingPayload if TYPE_CHECKING: from opentelemetry.trace import Span from litellm.integrations._types.open_inference import ( - MessageAttributes, - ImageAttributes, - SpanAttributes, - AudioAttributes, - EmbeddingAttributes, - OpenInferenceSpanKindValues + MessageAttributes, + ImageAttributes, + SpanAttributes, + AudioAttributes, + EmbeddingAttributes, + OpenInferenceSpanKindValues, ) @@ -158,7 +158,9 @@ def _set_audio_outputs(span: "Span", response_obj, audio_attrs, span_attrs): audio_transcript = audio_item.get("transcript") if audio_transcript: - safe_set_attribute(span, f"{audio_attrs.AUDIO_TRANSCRIPT}.{i}", audio_transcript) + safe_set_attribute( + span, f"{audio_attrs.AUDIO_TRANSCRIPT}.{i}", audio_transcript + ) def _set_embedding_outputs(span: "Span", response_obj, embedding_attrs, span_attrs): @@ -212,7 +214,9 @@ def _set_structured_outputs(span: "Span", response_obj, msg_attrs, span_attrs): message_content = getattr(first_content, "text", "") message_role = getattr(item, "role", "assistant") safe_set_attribute(span, span_attrs.OUTPUT_VALUE, message_content) - safe_set_attribute(span, f"{prefix}.{msg_attrs.MESSAGE_CONTENT}", message_content) + safe_set_attribute( + span, f"{prefix}.{msg_attrs.MESSAGE_CONTENT}", message_content + ) safe_set_attribute(span, f"{prefix}.{msg_attrs.MESSAGE_ROLE}", message_role) @@ -221,16 +225,24 @@ def _set_usage_outputs(span: "Span", response_obj, span_attrs): if not usage: return - safe_set_attribute(span, span_attrs.LLM_TOKEN_COUNT_TOTAL, usage.get("total_tokens")) + safe_set_attribute( + span, span_attrs.LLM_TOKEN_COUNT_TOTAL, usage.get("total_tokens") + ) completion_tokens = usage.get("completion_tokens") or usage.get("output_tokens") if completion_tokens: - safe_set_attribute(span, span_attrs.LLM_TOKEN_COUNT_COMPLETION, completion_tokens) + safe_set_attribute( + span, span_attrs.LLM_TOKEN_COUNT_COMPLETION, completion_tokens + ) prompt_tokens = usage.get("prompt_tokens") or usage.get("input_tokens") if prompt_tokens: safe_set_attribute(span, span_attrs.LLM_TOKEN_COUNT_PROMPT, prompt_tokens) reasoning_tokens = usage.get("output_tokens_details", {}).get("reasoning_tokens") if reasoning_tokens: - safe_set_attribute(span, span_attrs.LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING, reasoning_tokens) + safe_set_attribute( + span, + span_attrs.LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING, + reasoning_tokens, + ) def _infer_open_inference_span_kind(call_type: Optional[str]) -> str: @@ -281,11 +293,15 @@ def _infer_open_inference_span_kind(call_type: Optional[str]) -> str: ): return OpenInferenceSpanKindValues.LLM.value - if any(keyword in lowered for keyword in ("file", "batch", "container", "fine_tuning_job")): + if any( + keyword in lowered + for keyword in ("file", "batch", "container", "fine_tuning_job") + ): return OpenInferenceSpanKindValues.CHAIN.value return OpenInferenceSpanKindValues.UNKNOWN.value + def _set_tool_attributes( span: "Span", optional_tools: Optional[list], metadata_tools: Optional[list] ): @@ -294,18 +310,30 @@ def _set_tool_attributes( for idx, tool in enumerate(optional_tools): if not isinstance(tool, dict): continue - function = tool.get("function") if isinstance(tool.get("function"), dict) else None + function = ( + tool.get("function") if isinstance(tool.get("function"), dict) else None + ) if not function: continue tool_name = function.get("name") if tool_name: - safe_set_attribute(span, f"{SpanAttributes.LLM_TOOLS}.{idx}.name", tool_name) + safe_set_attribute( + span, f"{SpanAttributes.LLM_TOOLS}.{idx}.name", tool_name + ) tool_description = function.get("description") if tool_description: - safe_set_attribute(span, f"{SpanAttributes.LLM_TOOLS}.{idx}.description", tool_description) + safe_set_attribute( + span, + f"{SpanAttributes.LLM_TOOLS}.{idx}.description", + tool_description, + ) params = function.get("parameters") if params is not None: - safe_set_attribute(span, f"{SpanAttributes.LLM_TOOLS}.{idx}.parameters", json.dumps(params)) + safe_set_attribute( + span, + f"{SpanAttributes.LLM_TOOLS}.{idx}.parameters", + json.dumps(params), + ) if metadata_tools and isinstance(metadata_tools, list): for idx, tool in enumerate(metadata_tools): @@ -343,7 +371,11 @@ def set_attributes( if standard_logging_payload is None: raise ValueError("standard_logging_object not found in kwargs") - metadata = standard_logging_payload.get("metadata") if standard_logging_payload else None + metadata = ( + standard_logging_payload.get("metadata") + if standard_logging_payload + else None + ) _set_metadata_attributes(span, metadata, SpanAttributes) metadata_tools = _extract_metadata_tools(metadata) @@ -362,13 +394,19 @@ def set_attributes( span_kind = _infer_open_inference_span_kind(call_type=call_type) _set_tool_attributes(span, optional_tools, metadata_tools) - if (optional_tools or metadata_tools) and span_kind != OpenInferenceSpanKindValues.TOOL.value: + if ( + optional_tools or metadata_tools + ) and span_kind != OpenInferenceSpanKindValues.TOOL.value: span_kind = OpenInferenceSpanKindValues.TOOL.value safe_set_attribute(span, SpanAttributes.OPENINFERENCE_SPAN_KIND, span_kind) attributes.set_messages(span, kwargs) - model_params = standard_logging_payload.get("model_parameters") if standard_logging_payload else None + model_params = ( + standard_logging_payload.get("model_parameters") + if standard_logging_payload + else None + ) _set_model_params(span, model_params, SpanAttributes) _set_response_attributes(span=span, response_obj=response_obj) @@ -418,17 +456,29 @@ def _set_request_attributes( if kwargs.get("model"): safe_set_attribute(span, span_attrs.LLM_MODEL_NAME, kwargs.get("model")) - safe_set_attribute(span, "llm.request.type", standard_logging_payload.get("call_type")) - safe_set_attribute(span, span_attrs.LLM_PROVIDER, litellm_params.get("custom_llm_provider", "Unknown")) + safe_set_attribute( + span, "llm.request.type", standard_logging_payload.get("call_type") + ) + safe_set_attribute( + span, + span_attrs.LLM_PROVIDER, + litellm_params.get("custom_llm_provider", "Unknown"), + ) if optional_params.get("max_tokens"): - safe_set_attribute(span, "llm.request.max_tokens", optional_params.get("max_tokens")) + safe_set_attribute( + span, "llm.request.max_tokens", optional_params.get("max_tokens") + ) if optional_params.get("temperature"): - safe_set_attribute(span, "llm.request.temperature", optional_params.get("temperature")) + safe_set_attribute( + span, "llm.request.temperature", optional_params.get("temperature") + ) if optional_params.get("top_p"): safe_set_attribute(span, "llm.request.top_p", optional_params.get("top_p")) - safe_set_attribute(span, "llm.is_streaming", str(optional_params.get("stream", False))) + safe_set_attribute( + span, "llm.is_streaming", str(optional_params.get("stream", False)) + ) if optional_params.get("user"): safe_set_attribute(span, "llm.user", optional_params.get("user")) @@ -443,7 +493,9 @@ def _set_model_params(span: "Span", model_params: Optional[dict], span_attrs) -> if not model_params: return - safe_set_attribute(span, span_attrs.LLM_INVOCATION_PARAMETERS, safe_dumps(model_params)) + safe_set_attribute( + span, span_attrs.LLM_INVOCATION_PARAMETERS, safe_dumps(model_params) + ) if model_params.get("user"): user_id = model_params.get("user") if user_id is not None: diff --git a/litellm/integrations/arize/arize_phoenix.py b/litellm/integrations/arize/arize_phoenix.py index 6720a93044..00bc24d418 100644 --- a/litellm/integrations/arize/arize_phoenix.py +++ b/litellm/integrations/arize/arize_phoenix.py @@ -12,7 +12,9 @@ if TYPE_CHECKING: from opentelemetry.trace import SpanKind from litellm.integrations.opentelemetry import OpenTelemetry as _OpenTelemetry - from litellm.integrations.opentelemetry import OpenTelemetryConfig as _OpenTelemetryConfig + from litellm.integrations.opentelemetry import ( + OpenTelemetryConfig as _OpenTelemetryConfig, + ) from litellm.types.integrations.arize import Protocol as _Protocol Protocol = _Protocol @@ -91,7 +93,9 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore @staticmethod def set_arize_phoenix_attributes(span: Span, kwargs, response_obj): - from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes import safe_set_attribute + from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes import ( + safe_set_attribute, + ) _utils.set_attributes(span, kwargs, response_obj, ArizeOTELAttributes) @@ -103,7 +107,9 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore # Fall back to static config from env var config = ArizePhoenixLogger.get_arize_phoenix_config() if config.project_name: - safe_set_attribute(span, "openinference.project.name", config.project_name) + safe_set_attribute( + span, "openinference.project.name", config.project_name + ) return @@ -172,7 +178,9 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore start_time_val = kwargs.get("start_time", kwargs.get("api_call_start_time")) parent_span = self.tracer.start_span( name="litellm_proxy_request", - start_time=self._to_ns(start_time_val) if start_time_val is not None else None, + start_time=self._to_ns(start_time_val) + if start_time_val is not None + else None, context=traceparent_ctx, kind=self.span_kind.SERVER, ) @@ -212,9 +220,7 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore # Raw-request sub-span (if enabled) — must be created before # ending the parent span so the hierarchy is valid. - self._maybe_log_raw_request( - kwargs, response_obj, start_time, end_time, span - ) + self._maybe_log_raw_request(kwargs, response_obj, start_time, end_time, span) span.end(end_time=self._to_ns(end_time)) # Guardrail span @@ -290,7 +296,9 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore if collector_endpoint: # Parse the endpoint to determine protocol - if collector_endpoint.startswith("grpc://") or (":4317" in collector_endpoint and "/v1/traces" not in collector_endpoint): + if collector_endpoint.startswith("grpc://") or ( + ":4317" in collector_endpoint and "/v1/traces" not in collector_endpoint + ): endpoint = collector_endpoint protocol = "otlp_grpc" else: @@ -334,11 +342,10 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore endpoint=endpoint, project_name=project_name, ) - + ## cannot suppress additional proxy server spans, removed previous methods. async def async_health_check(self): - config = self.get_arize_phoenix_config() if not config.otlp_auth_headers: @@ -350,4 +357,4 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore return { "status": "healthy", "message": "Arize-Phoenix credentials are configured properly", - } \ No newline at end of file + } diff --git a/litellm/integrations/azure_sentinel/__init__.py b/litellm/integrations/azure_sentinel/__init__.py index 46f2fed0a9..036711a80d 100644 --- a/litellm/integrations/azure_sentinel/__init__.py +++ b/litellm/integrations/azure_sentinel/__init__.py @@ -1,4 +1,3 @@ from litellm.integrations.azure_sentinel.azure_sentinel import AzureSentinelLogger __all__ = ["AzureSentinelLogger"] - diff --git a/litellm/integrations/azure_sentinel/azure_sentinel.py b/litellm/integrations/azure_sentinel/azure_sentinel.py index 875432de87..dd508e6c6c 100644 --- a/litellm/integrations/azure_sentinel/azure_sentinel.py +++ b/litellm/integrations/azure_sentinel/azure_sentinel.py @@ -62,18 +62,22 @@ class AzureSentinelLogger(CustomBatchLogger): llm_provider=httpxSpecialProvider.LoggingCallback ) - self.dcr_immutable_id = ( - dcr_immutable_id or os.getenv("AZURE_SENTINEL_DCR_IMMUTABLE_ID") + self.dcr_immutable_id = dcr_immutable_id or os.getenv( + "AZURE_SENTINEL_DCR_IMMUTABLE_ID" ) self.stream_name = stream_name or os.getenv( "AZURE_SENTINEL_STREAM_NAME", "Custom-LiteLLM" ) self.endpoint = endpoint or os.getenv("AZURE_SENTINEL_ENDPOINT") - self.tenant_id = tenant_id or os.getenv("AZURE_SENTINEL_TENANT_ID") or os.getenv( - "AZURE_TENANT_ID" + self.tenant_id = ( + tenant_id + or os.getenv("AZURE_SENTINEL_TENANT_ID") + or os.getenv("AZURE_TENANT_ID") ) - self.client_id = client_id or os.getenv("AZURE_SENTINEL_CLIENT_ID") or os.getenv( - "AZURE_CLIENT_ID" + self.client_id = ( + client_id + or os.getenv("AZURE_SENTINEL_CLIENT_ID") + or os.getenv("AZURE_CLIENT_ID") ) self.client_secret = ( client_secret @@ -103,9 +107,7 @@ class AzureSentinelLogger(CustomBatchLogger): ) # Build API endpoint: {Endpoint}/dataCollectionRules/{DCR Immutable ID}/streams/{Stream Name}?api-version=2023-01-01 - self.api_endpoint = ( - f"{self.endpoint.rstrip('/')}/dataCollectionRules/{self.dcr_immutable_id}/streams/{self.stream_name}?api-version=2023-01-01" - ) + self.api_endpoint = f"{self.endpoint.rstrip('/')}/dataCollectionRules/{self.dcr_immutable_id}/streams/{self.stream_name}?api-version=2023-01-01" # OAuth2 scope for Azure Monitor self.oauth_scope = "https://monitor.azure.com/.default" @@ -139,7 +141,9 @@ class AzureSentinelLogger(CustomBatchLogger): assert self.client_id is not None, "client_id is required" assert self.client_secret is not None, "client_secret is required" - token_url = f"https://login.microsoftonline.com/{self.tenant_id}/oauth2/v2.0/token" + token_url = ( + f"https://login.microsoftonline.com/{self.tenant_id}/oauth2/v2.0/token" + ) token_data = { "client_id": self.client_id, @@ -173,9 +177,7 @@ class AzureSentinelLogger(CustomBatchLogger): return self.oauth_token - async def async_log_success_event( - self, kwargs, response_obj, start_time, end_time - ): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """ Async Log success events to Azure Sentinel @@ -209,9 +211,7 @@ class AzureSentinelLogger(CustomBatchLogger): ) pass - async def async_log_failure_event( - self, kwargs, response_obj, start_time, end_time - ): + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): """ Async Log failure events to Azure Sentinel diff --git a/litellm/integrations/azure_storage/azure_storage.py b/litellm/integrations/azure_storage/azure_storage.py index 85f91199c1..6fc7b9c104 100644 --- a/litellm/integrations/azure_storage/azure_storage.py +++ b/litellm/integrations/azure_storage/azure_storage.py @@ -54,12 +54,12 @@ class AzureBlobStorageLogger(CustomBatchLogger): self._service_client_timeout: Optional[float] = None # Internal variables used for Token based authentication - self.azure_auth_token: Optional[str] = ( - None # the Azure AD token to use for Azure Storage API requests - ) - self.token_expiry: Optional[datetime] = ( - None # the expiry time of the currentAzure AD token - ) + self.azure_auth_token: Optional[ + str + ] = None # the Azure AD token to use for Azure Storage API requests + self.token_expiry: Optional[ + datetime + ] = None # the expiry time of the currentAzure AD token asyncio.create_task(self.periodic_flush()) self.flush_lock = asyncio.Lock() diff --git a/litellm/integrations/braintrust_logging.py b/litellm/integrations/braintrust_logging.py index 42e9680a7f..cb1b2bc553 100644 --- a/litellm/integrations/braintrust_logging.py +++ b/litellm/integrations/braintrust_logging.py @@ -41,7 +41,9 @@ class BraintrustLogger(CustomLogger): self.is_mock_mode = should_use_braintrust_mock() if self.is_mock_mode: create_mock_braintrust_client() - verbose_logger.info("[BRAINTRUST MOCK] Braintrust logger initialized in mock mode") + verbose_logger.info( + "[BRAINTRUST MOCK] Braintrust logger initialized in mock mode" + ) self.validate_environment(api_key=api_key) self.api_base = api_base or os.getenv("BRAINTRUST_API_BASE") or API_BASE self.default_project_id = None @@ -50,9 +52,9 @@ class BraintrustLogger(CustomLogger): "Authorization": "Bearer " + self.api_key, "Content-Type": "application/json", } - self._project_id_cache: Dict[str, str] = ( - {} - ) # Cache mapping project names to IDs + self._project_id_cache: Dict[ + str, str + ] = {} # Cache mapping project names to IDs self.global_braintrust_http_handler = get_async_httpx_client( llm_provider=httpxSpecialProvider.LoggingCallback ) @@ -214,7 +216,7 @@ class BraintrustLogger(CustomLogger): # Allow metadata override for span name span_name = dynamic_metadata.get("span_name", "Chat Completion") - + # Span parents is a special case span_parents = dynamic_metadata.get("span_parents") @@ -236,7 +238,7 @@ class BraintrustLogger(CustomLogger): "span_attributes": {"name": span_name, "type": "llm"}, } - # Braintrust cannot specify 'tags' for non-root spans + # Braintrust cannot specify 'tags' for non-root spans if dynamic_metadata.get("root_span_id") is None: request_data["tags"] = tags @@ -386,7 +388,7 @@ class BraintrustLogger(CustomLogger): "span_attributes": {"name": span_name, "type": "llm"}, } - # Braintrust cannot specify 'tags' for non-root spans + # Braintrust cannot specify 'tags' for non-root spans if dynamic_metadata.get("root_span_id") is None: request_data["tags"] = tags diff --git a/litellm/integrations/braintrust_mock_client.py b/litellm/integrations/braintrust_mock_client.py index 030aa62cd0..59e0988a10 100644 --- a/litellm/integrations/braintrust_mock_client.py +++ b/litellm/integrations/braintrust_mock_client.py @@ -13,7 +13,11 @@ import time from urllib.parse import urlparse from litellm._logging import verbose_logger -from litellm.integrations.mock_client_factory import MockClientConfig, MockResponse, create_mock_client_factory +from litellm.integrations.mock_client_factory import ( + MockClientConfig, + MockResponse, + create_mock_client_factory, +) # Use factory for should_use_mock and MockResponse # Braintrust uses both HTTPHandler (sync) and AsyncHTTPHandler (async) @@ -37,7 +41,10 @@ _config = MockClientConfig( # Get should_use_mock and create_mock_client from factory # We need to call the factory's create_mock_client to patch AsyncHTTPHandler.post -create_mock_braintrust_factory_client, should_use_braintrust_mock = create_mock_client_factory(_config) +( + create_mock_braintrust_factory_client, + should_use_braintrust_mock, +) = create_mock_client_factory(_config) # Store original HTTPHandler.post method (Braintrust-specific for sync calls with custom logic) _original_http_handler_post = None @@ -66,7 +73,19 @@ def _is_braintrust_url(url: str) -> bool: ) -def _mock_http_handler_post(self, url, data=None, json=None, params=None, headers=None, timeout=None, stream=False, files=None, content=None, logging_obj=None): +def _mock_http_handler_post( + self, + url, + data=None, + json=None, + params=None, + headers=None, + timeout=None, + stream=False, + files=None, + content=None, + logging_obj=None, +): """Monkey-patched HTTPHandler.post that intercepts Braintrust calls with endpoint-specific responses.""" # Only mock Braintrust API calls if isinstance(url, str) and _is_braintrust_url(url): @@ -86,46 +105,62 @@ def _mock_http_handler_post(self, url, data=None, json=None, params=None, header status_code=_config.default_status_code, json_data=mock_data, url=url, - elapsed_seconds=_MOCK_LATENCY_SECONDS + elapsed_seconds=_MOCK_LATENCY_SECONDS, ) if _original_http_handler_post is not None: - return _original_http_handler_post(self, url=url, data=data, json=json, params=params, headers=headers, timeout=timeout, stream=stream, files=files, content=content, logging_obj=logging_obj) + return _original_http_handler_post( + self, + url=url, + data=data, + json=json, + params=params, + headers=headers, + timeout=timeout, + stream=stream, + files=files, + content=content, + logging_obj=logging_obj, + ) raise RuntimeError("Original HTTPHandler.post not available") def create_mock_braintrust_client(): """ Monkey-patch HTTPHandler.post to intercept Braintrust sync calls. - + Braintrust uses HTTPHandler for sync calls and AsyncHTTPHandler for async calls. HTTPHandler.post uses self.client.send(), not self.client.post(), so we need custom patching for sync (similar to Helicone). AsyncHTTPHandler.post is patched by the factory. - + We use custom patching instead of factory's patch_http_handler because we need endpoint-specific responses (different for /project vs /project_logs). - + This function is idempotent - it only initializes mocks once, even if called multiple times. """ global _original_http_handler_post, _mocks_initialized - + if _mocks_initialized: return - + verbose_logger.debug("[BRAINTRUST MOCK] Initializing Braintrust mock client...") - + from litellm.llms.custom_httpx.http_handler import HTTPHandler - + if _original_http_handler_post is None: _original_http_handler_post = HTTPHandler.post HTTPHandler.post = _mock_http_handler_post # type: ignore verbose_logger.debug("[BRAINTRUST MOCK] Patched HTTPHandler.post") - + # CRITICAL: Call the factory's initialization function to patch AsyncHTTPHandler.post # This is required for async calls to be mocked create_mock_braintrust_factory_client() - - verbose_logger.debug(f"[BRAINTRUST MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS*1000:.0f}ms") - verbose_logger.debug("[BRAINTRUST MOCK] Braintrust mock client initialization complete") - + + verbose_logger.debug( + f"[BRAINTRUST MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS*1000:.0f}ms" + ) + verbose_logger.debug( + "[BRAINTRUST MOCK] Braintrust mock client initialization complete" + ) + _mocks_initialized = True diff --git a/litellm/integrations/cloudzero/cz_resource_names.py b/litellm/integrations/cloudzero/cz_resource_names.py index f1098d2038..20862c1c7e 100644 --- a/litellm/integrations/cloudzero/cz_resource_names.py +++ b/litellm/integrations/cloudzero/cz_resource_names.py @@ -30,7 +30,9 @@ class CZEntityType(str, Enum): class CZRNGenerator: """Generate CloudZero Resource Names (CZRNs) for LiteLLM resources.""" - CZRN_REGEX = re.compile(r'^czrn:([a-z0-9-]+):([a-zA-Z0-9-]+):([a-z0-9-]+):([a-z0-9-]+):([a-z0-9-]+):(.+)$') + CZRN_REGEX = re.compile( + r"^czrn:([a-z0-9-]+):([a-zA-Z0-9-]+):([a-z0-9-]+):([a-z0-9-]+):([a-z0-9-]+):(.+)$" + ) def __init__(self): """Initialize CZRN generator.""" @@ -38,9 +40,9 @@ class CZRNGenerator: def create_from_litellm_data(self, row: dict[str, Any]) -> str: """Create a CZRN from LiteLLM daily spend data. - + CZRN format: czrn:::::: - + For LiteLLM resources, we map: - service-type: 'litellm' (the service managing the LLM calls) - provider: The custom_llm_provider (e.g., 'openai', 'anthropic', 'azure') @@ -49,18 +51,18 @@ class CZRNGenerator: - resource-type: 'llm-usage' (represents LLM usage/inference) - cloud-local-id: model """ - service_type = 'litellm' - provider = self._normalize_provider(row.get('custom_llm_provider', 'unknown')) - region = 'cross-region' + service_type = "litellm" + provider = self._normalize_provider(row.get("custom_llm_provider", "unknown")) + region = "cross-region" # Use the actual entity_id (team_id or user_id) as the owner account - team_id = row.get('team_id', 'unknown') + team_id = row.get("team_id", "unknown") owner_account_id = self._normalize_component(team_id) - resource_type = 'llm-usage' + resource_type = "llm-usage" # Create a unique identifier with just the model (entity info already in owner_account_id) - model = row.get('model', 'unknown') + model = row.get("model", "unknown") cloud_local_id = model @@ -70,7 +72,7 @@ class CZRNGenerator: region=region, owner_account_id=owner_account_id, resource_type=resource_type, - cloud_local_id=cloud_local_id + cloud_local_id=cloud_local_id, ) def create_from_components( @@ -80,7 +82,7 @@ class CZRNGenerator: region: str, owner_account_id: str, resource_type: str, - cloud_local_id: str + cloud_local_id: str, ) -> str: """Create a CZRN from individual components.""" # Normalize components to ensure they meet CZRN requirements @@ -104,7 +106,7 @@ class CZRNGenerator: def extract_components(self, czrn: str) -> tuple[str, str, str, str, str, str]: """Extract all components from a CZRN. - + Returns: (service_type, provider, region, owner_account_id, resource_type, cloud_local_id) """ match = self.CZRN_REGEX.match(czrn) @@ -117,42 +119,43 @@ class CZRNGenerator: """Normalize provider names to standard CZRN format.""" # Map common provider names to CZRN standards provider_map = { - litellm.LlmProviders.AZURE.value: 'azure', - litellm.LlmProviders.AZURE_AI.value: 'azure', - litellm.LlmProviders.ANTHROPIC.value: 'anthropic', - litellm.LlmProviders.BEDROCK.value: 'aws', - litellm.LlmProviders.VERTEX_AI.value: 'gcp', - litellm.LlmProviders.GEMINI.value: 'google', - litellm.LlmProviders.COHERE.value: 'cohere', - litellm.LlmProviders.HUGGINGFACE.value: 'huggingface', - litellm.LlmProviders.REPLICATE.value: 'replicate', - litellm.LlmProviders.TOGETHER_AI.value: 'together-ai', + litellm.LlmProviders.AZURE.value: "azure", + litellm.LlmProviders.AZURE_AI.value: "azure", + litellm.LlmProviders.ANTHROPIC.value: "anthropic", + litellm.LlmProviders.BEDROCK.value: "aws", + litellm.LlmProviders.VERTEX_AI.value: "gcp", + litellm.LlmProviders.GEMINI.value: "google", + litellm.LlmProviders.COHERE.value: "cohere", + litellm.LlmProviders.HUGGINGFACE.value: "huggingface", + litellm.LlmProviders.REPLICATE.value: "replicate", + litellm.LlmProviders.TOGETHER_AI.value: "together-ai", } - normalized = provider.lower().replace('_', '-') + normalized = provider.lower().replace("_", "-") # use litellm custom llm provider if not in provider_map if normalized not in provider_map: return normalized return provider_map.get(normalized, normalized) - def _normalize_component(self, component: str, allow_uppercase: bool = False) -> str: + def _normalize_component( + self, component: str, allow_uppercase: bool = False + ) -> str: """Normalize a CZRN component to meet format requirements.""" if not component: - return 'unknown' + return "unknown" # Convert to lowercase unless uppercase is allowed if not allow_uppercase: component = component.lower() # Replace invalid characters with hyphens - component = re.sub(r'[^a-zA-Z0-9-]', '-', component) + component = re.sub(r"[^a-zA-Z0-9-]", "-", component) # Remove consecutive hyphens - component = re.sub(r'-+', '-', component) + component = re.sub(r"-+", "-", component) # Remove leading/trailing hyphens - component = component.strip('-') - - return component or 'unknown' + component = component.strip("-") + return component or "unknown" diff --git a/litellm/integrations/cloudzero/cz_stream_api.py b/litellm/integrations/cloudzero/cz_stream_api.py index 83b6e318ba..d673536e72 100644 --- a/litellm/integrations/cloudzero/cz_stream_api.py +++ b/litellm/integrations/cloudzero/cz_stream_api.py @@ -30,7 +30,9 @@ from rich.console import Console class CloudZeroStreamer: """Stream CBF data to CloudZero AnyCost API with proper batching and timezone handling.""" - def __init__(self, api_key: str, connection_id: str, user_timezone: Optional[str] = None): + def __init__( + self, api_key: str, connection_id: str, user_timezone: Optional[str] = None + ): """Initialize CloudZero streamer with credentials.""" self.api_key = api_key self.connection_id = connection_id @@ -43,12 +45,16 @@ class CloudZeroStreamer: try: self.user_timezone = zoneinfo.ZoneInfo(user_timezone) except zoneinfo.ZoneInfoNotFoundError: - self.console.print(f"[yellow]Warning: Unknown timezone '{user_timezone}', using UTC[/yellow]") + self.console.print( + f"[yellow]Warning: Unknown timezone '{user_timezone}', using UTC[/yellow]" + ) self.user_timezone = timezone.utc else: self.user_timezone = timezone.utc - def send_batched(self, data: pl.DataFrame, operation: str = "replace_hourly") -> None: + def send_batched( + self, data: pl.DataFrame, operation: str = "replace_hourly" + ) -> None: """Send CBF data in daily batches to CloudZero AnyCost API.""" if data.is_empty(): self.console.print("[yellow]No data to send to CloudZero[/yellow]") @@ -61,7 +67,9 @@ class CloudZeroStreamer: self.console.print("[yellow]No valid daily batches to send[/yellow]") return - self.console.print(f"[blue]Sending {len(daily_batches)} daily batch(es) with operation '{operation}'[/blue]") + self.console.print( + f"[blue]Sending {len(daily_batches)} daily batch(es) with operation '{operation}'[/blue]" + ) for batch_date, batch_data in daily_batches.items(): self._send_daily_batch(batch_date, batch_data, operation) @@ -71,21 +79,23 @@ class CloudZeroStreamer: daily_batches: dict[str, list[dict[str, Any]]] = {} # Ensure we have the required columns - if 'time/usage_start' not in data.columns: - self.console.print("[red]Error: Missing 'time/usage_start' column for date grouping[/red]") + if "time/usage_start" not in data.columns: + self.console.print( + "[red]Error: Missing 'time/usage_start' column for date grouping[/red]" + ) return {} - + timestamp_str: Optional[str] = None for row in data.iter_rows(named=True): try: # Parse the timestamp and convert to UTC - timestamp_str = row.get('time/usage_start') + timestamp_str = row.get("time/usage_start") if not timestamp_str: continue # Parse timestamp and handle timezone conversion dt = self._parse_and_convert_timestamp(timestamp_str) - batch_date = dt.strftime('%Y-%m-%d') + batch_date = dt.strftime("%Y-%m-%d") if batch_date not in daily_batches: daily_batches[batch_date] = [] @@ -93,25 +103,54 @@ class CloudZeroStreamer: daily_batches[batch_date].append(row) except Exception as e: - self.console.print(f"[yellow]Warning: Could not process timestamp '{timestamp_str}': {e}[/yellow]") + self.console.print( + f"[yellow]Warning: Could not process timestamp '{timestamp_str}': {e}[/yellow]" + ) continue # Convert lists back to DataFrames - return {date_key: pl.DataFrame(records) for date_key, records in daily_batches.items() if records} + return { + date_key: pl.DataFrame(records) + for date_key, records in daily_batches.items() + if records + } def _parse_and_convert_timestamp(self, timestamp_str: str) -> datetime: """Parse timestamp string and convert to UTC.""" # Try to parse the timestamp string try: # Handle various ISO 8601 formats - if timestamp_str.endswith('Z'): - dt = datetime.fromisoformat(timestamp_str.replace('Z', '+00:00')) - elif '+' in timestamp_str or timestamp_str.endswith(('-00:00', '-01:00', '-02:00', '-03:00', - '-04:00', '-05:00', '-06:00', '-07:00', - '-08:00', '-09:00', '-10:00', '-11:00', - '-12:00', '+01:00', '+02:00', '+03:00', - '+04:00', '+05:00', '+06:00', '+07:00', - '+08:00', '+09:00', '+10:00', '+11:00', '+12:00')): + if timestamp_str.endswith("Z"): + dt = datetime.fromisoformat(timestamp_str.replace("Z", "+00:00")) + elif "+" in timestamp_str or timestamp_str.endswith( + ( + "-00:00", + "-01:00", + "-02:00", + "-03:00", + "-04:00", + "-05:00", + "-06:00", + "-07:00", + "-08:00", + "-09:00", + "-10:00", + "-11:00", + "-12:00", + "+01:00", + "+02:00", + "+03:00", + "+04:00", + "+05:00", + "+06:00", + "+07:00", + "+08:00", + "+09:00", + "+10:00", + "+11:00", + "+12:00", + ) + ): dt = datetime.fromisoformat(timestamp_str) else: # Assume user timezone if no timezone info @@ -125,14 +164,16 @@ class CloudZeroStreamer: except ValueError as e: raise ValueError(f"Could not parse timestamp '{timestamp_str}': {e}") - def _send_daily_batch(self, batch_date: str, batch_data: pl.DataFrame, operation: str) -> None: + def _send_daily_batch( + self, batch_date: str, batch_data: pl.DataFrame, operation: str + ) -> None: """Send a single daily batch to CloudZero API.""" if batch_data.is_empty(): return headers = { - 'Authorization': f'Bearer {self.api_key}', - 'Content-Type': 'application/json' + "Authorization": f"Bearer {self.api_key}", + "Content-Type": "application/json", } # Use the correct API endpoint format from documentation @@ -143,29 +184,39 @@ class CloudZeroStreamer: try: with httpx.Client(timeout=30.0) as client: - self.console.print(f"[blue]Sending batch for {batch_date} ({len(batch_data)} records)[/blue]") + self.console.print( + f"[blue]Sending batch for {batch_date} ({len(batch_data)} records)[/blue]" + ) response = client.post(url, headers=headers, json=payload) response.raise_for_status() - self.console.print(f"[green]✓ Successfully sent batch for {batch_date} ({len(batch_data)} records)[/green]") + self.console.print( + f"[green]✓ Successfully sent batch for {batch_date} ({len(batch_data)} records)[/green]" + ) except httpx.RequestError as e: - self.console.print(f"[red]✗ Network error sending batch for {batch_date}: {e}[/red]") + self.console.print( + f"[red]✗ Network error sending batch for {batch_date}: {e}[/red]" + ) raise except httpx.HTTPStatusError as e: - self.console.print(f"[red]✗ HTTP error sending batch for {batch_date}: {e.response.status_code} {e.response.text}[/red]") + self.console.print( + f"[red]✗ HTTP error sending batch for {batch_date}: {e.response.status_code} {e.response.text}[/red]" + ) raise - def _prepare_batch_payload(self, batch_date: str, batch_data: pl.DataFrame, operation: str) -> dict[str, Any]: + def _prepare_batch_payload( + self, batch_date: str, batch_data: pl.DataFrame, operation: str + ) -> dict[str, Any]: """Prepare batch payload according to CloudZero AnyCost API format.""" # Convert batch_date to month for the API (YYYY-MM format) try: - date_obj = datetime.strptime(batch_date, '%Y-%m-%d') - month_str = date_obj.strftime('%Y-%m') + date_obj = datetime.strptime(batch_date, "%Y-%m-%d") + month_str = date_obj.strftime("%Y-%m") except ValueError: # Fallback to current month - month_str = datetime.now().strftime('%Y-%m') + month_str = datetime.now().strftime("%Y-%m") # Convert DataFrame rows to API format data_records = [] @@ -174,15 +225,13 @@ class CloudZeroStreamer: if record: data_records.append(record) - payload = { - 'month': month_str, - 'operation': operation, - 'data': data_records - } + payload = {"month": month_str, "operation": operation, "data": data_records} return payload - def _convert_cbf_to_api_format(self, row: dict[str, Any]) -> Optional[dict[str, Any]]: + def _convert_cbf_to_api_format( + self, row: dict[str, Any] + ) -> Optional[dict[str, Any]]: """Convert CBF row to CloudZero API format - keeping CBF field names as CloudZero expects them.""" try: # CloudZero expects CBF format field names directly, not converted names @@ -196,20 +245,24 @@ class CloudZeroStreamer: # Format floats to avoid scientific notation if isinstance(value, float): # Use a reasonable precision that avoids scientific notation - api_record[key] = f"{value:.10f}".rstrip('0').rstrip('.') + api_record[key] = f"{value:.10f}".rstrip("0").rstrip(".") else: api_record[key] = str(value) else: api_record[key] = value # Ensure timestamp is in UTC format - if 'time/usage_start' in api_record: - api_record['time/usage_start'] = self._ensure_utc_timestamp(api_record['time/usage_start']) + if "time/usage_start" in api_record: + api_record["time/usage_start"] = self._ensure_utc_timestamp( + api_record["time/usage_start"] + ) return api_record except Exception as e: - self.console.print(f"[yellow]Warning: Could not convert record to API format: {e}[/yellow]") + self.console.print( + f"[yellow]Warning: Could not convert record to API format: {e}[/yellow]" + ) return None def _ensure_utc_timestamp(self, timestamp_str: str) -> str: @@ -219,9 +272,7 @@ class CloudZeroStreamer: try: dt = self._parse_and_convert_timestamp(timestamp_str) - return dt.isoformat().replace('+00:00', 'Z') + return dt.isoformat().replace("+00:00", "Z") except Exception: # Fallback to current time in UTC - return datetime.now(timezone.utc).isoformat().replace('+00:00', 'Z') - - + return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z") diff --git a/litellm/integrations/cloudzero/transform.py b/litellm/integrations/cloudzero/transform.py index b40a71da1c..c1b0d5cf41 100644 --- a/litellm/integrations/cloudzero/transform.py +++ b/litellm/integrations/cloudzero/transform.py @@ -41,8 +41,8 @@ class CBFTransformer: # Filter out records with zero successful_requests first original_count = len(data) - if 'successful_requests' in data.columns: - filtered_data = data.filter(pl.col('successful_requests') > 0) + if "successful_requests" in data.columns: + filtered_data = data.filter(pl.col("successful_requests") > 0) zero_requests_dropped = original_count - len(filtered_data) else: filtered_data = data @@ -64,16 +64,23 @@ class CBFTransformer: # Print summary of dropped records if any from rich.console import Console + console = Console() if zero_requests_dropped > 0: - console.print(f"[yellow]⚠️ Dropped {zero_requests_dropped:,} of {original_count:,} records with zero successful_requests[/yellow]") + console.print( + f"[yellow]⚠️ Dropped {zero_requests_dropped:,} of {original_count:,} records with zero successful_requests[/yellow]" + ) if czrn_dropped_count > 0: - console.print(f"[yellow]⚠️ Dropped {czrn_dropped_count:,} of {filtered_count:,} filtered records due to invalid CZRNs[/yellow]") + console.print( + f"[yellow]⚠️ Dropped {czrn_dropped_count:,} of {filtered_count:,} filtered records due to invalid CZRNs[/yellow]" + ) if len(cbf_data) > 0: - console.print(f"[green]✓ Successfully transformed {len(cbf_data):,} records[/green]") + console.print( + f"[green]✓ Successfully transformed {len(cbf_data):,} records[/green]" + ) return pl.DataFrame(cbf_data) @@ -81,99 +88,116 @@ class CBFTransformer: """Create a single CBF record from LiteLLM daily spend row.""" # Parse date (daily spend tables use date strings like '2025-04-19') - usage_date = self._parse_date(row.get('date')) + usage_date = self._parse_date(row.get("date")) # Calculate total tokens - prompt_tokens = int(row.get('prompt_tokens', 0)) - completion_tokens = int(row.get('completion_tokens', 0)) + prompt_tokens = int(row.get("prompt_tokens", 0)) + completion_tokens = int(row.get("completion_tokens", 0)) total_tokens = prompt_tokens + completion_tokens # Create CloudZero Resource Name (CZRN) as resource_id resource_id = self.czrn_generator.create_from_litellm_data(row) # Build dimensions for CloudZero - model = str(row.get('model', '')) - api_key_hash = str(row.get('api_key', ''))[:8] # First 8 chars for identification - + model = str(row.get("model", "")) + api_key_hash = str(row.get("api_key", ""))[ + :8 + ] # First 8 chars for identification + # Handle team information with fallbacks - team_id = row.get('team_id') - team_alias = row.get('team_alias') - user_email = row.get('user_email') - + team_id = row.get("team_id") + team_alias = row.get("team_alias") + user_email = row.get("user_email") + # Use team_alias if available, otherwise team_id, otherwise fallback to 'unknown' - entity_id = str(team_alias) if team_alias else (str(team_id) if team_id else 'unknown') - + entity_id = ( + str(team_alias) if team_alias else (str(team_id) if team_id else "unknown") + ) + # Get alias fields if they exist - api_key_alias = row.get('api_key_alias') - organization_alias = row.get('organization_alias') - project_alias = row.get('project_alias') - user_alias = row.get('user_alias') + api_key_alias = row.get("api_key_alias") + organization_alias = row.get("organization_alias") + project_alias = row.get("project_alias") + user_alias = row.get("user_alias") dimensions = { - 'entity_type': CZEntityType.TEAM.value, - 'entity_id': entity_id, - 'team_alias': str(team_alias) if team_alias else 'unknown', - 'model': model, - 'model_group': str(row.get('model_group', '')), - 'provider': str(row.get('custom_llm_provider', '')), - 'api_key_prefix': api_key_hash, - 'api_key_alias': str(row.get('api_key_alias', '')), - 'user_email': str(user_email) if user_email else '', - 'api_requests': str(row.get('api_requests', 0)), - 'successful_requests': str(row.get('successful_requests', 0)), - 'failed_requests': str(row.get('failed_requests', 0)), - 'cache_creation_tokens': str(row.get('cache_creation_input_tokens', 0)), - 'cache_read_tokens': str(row.get('cache_read_input_tokens', 0)), - 'organization_alias': str(organization_alias) if organization_alias else '', - 'project_alias': str(project_alias) if project_alias else '', - 'user_alias': str(user_alias) if user_alias else '', + "entity_type": CZEntityType.TEAM.value, + "entity_id": entity_id, + "team_alias": str(team_alias) if team_alias else "unknown", + "model": model, + "model_group": str(row.get("model_group", "")), + "provider": str(row.get("custom_llm_provider", "")), + "api_key_prefix": api_key_hash, + "api_key_alias": str(row.get("api_key_alias", "")), + "user_email": str(user_email) if user_email else "", + "api_requests": str(row.get("api_requests", 0)), + "successful_requests": str(row.get("successful_requests", 0)), + "failed_requests": str(row.get("failed_requests", 0)), + "cache_creation_tokens": str(row.get("cache_creation_input_tokens", 0)), + "cache_read_tokens": str(row.get("cache_read_input_tokens", 0)), + "organization_alias": str(organization_alias) if organization_alias else "", + "project_alias": str(project_alias) if project_alias else "", + "user_alias": str(user_alias) if user_alias else "", } # Extract CZRN components to populate corresponding CBF columns czrn_components = self.czrn_generator.extract_components(resource_id) - service_type, provider, region, owner_account_id, resource_type, cloud_local_id = czrn_components + ( + service_type, + provider, + region, + owner_account_id, + resource_type, + cloud_local_id, + ) = czrn_components # Build resource/account as concat of api_key_alias and api_key_prefix - resource_account = f"{api_key_alias}|{api_key_hash}" if api_key_alias else api_key_hash + resource_account = ( + f"{api_key_alias}|{api_key_hash}" if api_key_alias else api_key_hash + ) # CloudZero CBF format with proper column names cbf_record = { # Required CBF fields - 'time/usage_start': usage_date.isoformat() if usage_date else None, # Required: ISO-formatted UTC datetime - 'cost/cost': float(row.get('spend', 0.0)), # Required: billed cost - 'resource/id': resource_id, # CZRN (CloudZero Resource Name) - + "time/usage_start": usage_date.isoformat() + if usage_date + else None, # Required: ISO-formatted UTC datetime + "cost/cost": float(row.get("spend", 0.0)), # Required: billed cost + "resource/id": resource_id, # CZRN (CloudZero Resource Name) # Usage metrics for token consumption - 'usage/amount': total_tokens, # Numeric value of tokens consumed - 'usage/units': 'tokens', # Description of token units - + "usage/amount": total_tokens, # Numeric value of tokens consumed + "usage/units": "tokens", # Description of token units # CBF fields - updated per LIT-1907 - 'resource/service': str(row.get('model_group', '')), # Send model_group - 'resource/account': resource_account, # Send api_key_alias|api_key_prefix - 'resource/region': region, # Maps to CZRN region (cross-region) - 'resource/usage_family': str(row.get('custom_llm_provider', '')), # Send provider - + "resource/service": str(row.get("model_group", "")), # Send model_group + "resource/account": resource_account, # Send api_key_alias|api_key_prefix + "resource/region": region, # Maps to CZRN region (cross-region) + "resource/usage_family": str( + row.get("custom_llm_provider", "") + ), # Send provider # Action field - 'action/operation': str(team_id) if team_id else '', # Send team_id - + "action/operation": str(team_id) if team_id else "", # Send team_id # Line item details - 'lineitem/type': 'Usage', # Standard usage line item + "lineitem/type": "Usage", # Standard usage line item } # Add CZRN components that don't have direct CBF column mappings as resource tags - cbf_record['resource/tag:provider'] = provider # CZRN provider component - cbf_record['resource/tag:model'] = cloud_local_id # CZRN cloud-local-id component (model) - + cbf_record["resource/tag:provider"] = provider # CZRN provider component + cbf_record[ + "resource/tag:model" + ] = cloud_local_id # CZRN cloud-local-id component (model) + # Add resource tags for all dimensions (using resource/tag: format) for key, value in dimensions.items(): - if value and value != 'N/A' and value != 'unknown': # Only add meaningful tags - cbf_record[f'resource/tag:{key}'] = str(value) + if ( + value and value != "N/A" and value != "unknown" + ): # Only add meaningful tags + cbf_record[f"resource/tag:{key}"] = str(value) # Add token breakdown as resource tags for analysis (excluding total_tokens per LIT-1907) if prompt_tokens > 0: - cbf_record['resource/tag:prompt_tokens'] = str(prompt_tokens) + cbf_record["resource/tag:prompt_tokens"] = str(prompt_tokens) if completion_tokens > 0: - cbf_record['resource/tag:completion_tokens'] = str(completion_tokens) + cbf_record["resource/tag:completion_tokens"] = str(completion_tokens) return CBFRecord(cbf_record) @@ -197,4 +221,3 @@ class CBFTransformer: return None return None - diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index b00584edbc..06ba9675ca 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -670,7 +670,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac return final_response """ pass - + async def async_should_run_chat_completion_agentic_loop( self, response: Any, @@ -869,9 +869,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac model_response_dict = model_response.model_dump() standard_logging_object_copy["response"] = model_response_dict - model_call_details_copy["standard_logging_object"] = ( - standard_logging_object_copy - ) + model_call_details_copy[ + "standard_logging_object" + ] = standard_logging_object_copy return model_call_details_copy async def get_proxy_server_request_from_cold_storage_with_object_key( diff --git a/litellm/integrations/custom_secret_manager.py b/litellm/integrations/custom_secret_manager.py index 2125aef220..45ffa2e08c 100644 --- a/litellm/integrations/custom_secret_manager.py +++ b/litellm/integrations/custom_secret_manager.py @@ -100,9 +100,7 @@ class CustomSecretManager(BaseSecretManager): """ super().__init__() self.secret_manager_name = secret_manager_name or "custom_secret_manager" - verbose_logger.info( - "Initialized custom secret manager" - ) + verbose_logger.info("Initialized custom secret manager") @abstractmethod async def async_read_secret( diff --git a/litellm/integrations/custom_sso_handler.py b/litellm/integrations/custom_sso_handler.py index bc80966f8c..7f60decabc 100644 --- a/litellm/integrations/custom_sso_handler.py +++ b/litellm/integrations/custom_sso_handler.py @@ -13,6 +13,7 @@ class CustomSSOLoginHandler(CustomLogger): Useful when you have an OAuth proxy in front of LiteLLM and you want to use the headers from the proxy to sign in the user """ + async def handle_custom_ui_sso_sign_in( self, request: Request, @@ -26,4 +27,4 @@ class CustomSSOLoginHandler(CustomLogger): display_name="Test", picture="https://test.com/test.png", provider="test", - ) \ No newline at end of file + ) diff --git a/litellm/integrations/datadog/datadog_llm_obs.py b/litellm/integrations/datadog/datadog_llm_obs.py index e5ce999749..de6cc02fa3 100644 --- a/litellm/integrations/datadog/datadog_llm_obs.py +++ b/litellm/integrations/datadog/datadog_llm_obs.py @@ -48,13 +48,15 @@ class DataDogLLMObsLogger(CustomBatchLogger): def __init__(self, **kwargs): try: verbose_logger.debug("DataDogLLMObs: Initializing logger") - + self.is_mock_mode = should_use_datadog_mock() - + if self.is_mock_mode: create_mock_datadog_client() - verbose_logger.debug("[DATADOG MOCK] DataDogLLMObs logger initialized in mock mode") - + verbose_logger.debug( + "[DATADOG MOCK] DataDogLLMObs logger initialized in mock mode" + ) + # Configure DataDog endpoint (Agent or Direct API) # Use LITELLM_DD_AGENT_HOST to avoid conflicts with ddtrace's DD_AGENT_HOST # Check for agent mode FIRST - agent mode doesn't require DD_API_KEY or DD_SITE @@ -189,9 +191,11 @@ class DataDogLLMObsLogger(CustomBatchLogger): verbose_logger.debug( f"DataDogLLMObs: Flushing {len(self.log_queue)} events" ) - + if self.is_mock_mode: - verbose_logger.debug("[DATADOG MOCK] Mock mode enabled - API calls will be intercepted") + verbose_logger.debug( + "[DATADOG MOCK] Mock mode enabled - API calls will be intercepted" + ) # Prepare the payload payload = { diff --git a/litellm/integrations/datadog/datadog_mock_client.py b/litellm/integrations/datadog/datadog_mock_client.py index a0a760deb0..7f9beab72c 100644 --- a/litellm/integrations/datadog/datadog_mock_client.py +++ b/litellm/integrations/datadog/datadog_mock_client.py @@ -8,7 +8,10 @@ Usage: Set DATADOG_MOCK=true in environment variables or config to enable mock mode. """ -from litellm.integrations.mock_client_factory import MockClientConfig, create_mock_client_factory +from litellm.integrations.mock_client_factory import ( + MockClientConfig, + create_mock_client_factory, +) # Create mock client using factory _config = MockClientConfig( @@ -25,4 +28,6 @@ _config = MockClientConfig( patch_sync_client=True, ) -create_mock_datadog_client, should_use_datadog_mock = create_mock_client_factory(_config) +create_mock_datadog_client, should_use_datadog_mock = create_mock_client_factory( + _config +) diff --git a/litellm/integrations/dotprompt/__init__.py b/litellm/integrations/dotprompt/__init__.py index 3847c8fa19..394929f4a2 100644 --- a/litellm/integrations/dotprompt/__init__.py +++ b/litellm/integrations/dotprompt/__init__.py @@ -25,6 +25,7 @@ def set_global_prompt_directory(directory: str) -> None: litellm.global_prompt_directory = directory # type: ignore + def _get_prompt_data_from_dotprompt_content(dotprompt_content: str) -> dict: """ Get the prompt data from the dotprompt content. @@ -36,12 +37,10 @@ def _get_prompt_data_from_dotprompt_content(dotprompt_content: str) -> dict: # Parse the dotprompt content to extract frontmatter and content temp_manager = PromptManager() metadata, content = temp_manager._parse_frontmatter(dotprompt_content) - + # Convert to prompt_data format - return { - "content": content.strip(), - "metadata": metadata - } + return {"content": content.strip(), "metadata": metadata} + def prompt_initializer( litellm_params: "PromptLiteLLMParams", prompt_spec: "PromptSpec" @@ -58,7 +57,7 @@ def prompt_initializer( ) prompt_file = getattr(litellm_params, "prompt_file", None) - + # Handle dotprompt_content from database dotprompt_content = getattr(litellm_params, "dotprompt_content", None) if dotprompt_content and not prompt_data and not prompt_file: @@ -74,7 +73,6 @@ def prompt_initializer( return dot_prompt_manager except Exception as e: - raise e diff --git a/litellm/integrations/dotprompt/dotprompt_manager.py b/litellm/integrations/dotprompt/dotprompt_manager.py index 9412ac3c84..37fdf7da69 100644 --- a/litellm/integrations/dotprompt/dotprompt_manager.py +++ b/litellm/integrations/dotprompt/dotprompt_manager.py @@ -128,7 +128,6 @@ class DotpromptManager(CustomPromptManagement): raise ValueError("prompt_id is required for dotprompt manager") try: - # Get the prompt template (versioned or base) template = self.prompt_manager.get_prompt( prompt_id=prompt_id, version=prompt_version @@ -205,7 +204,6 @@ class DotpromptManager(CustomPromptManagement): ignore_prompt_manager_model: Optional[bool] = False, ignore_prompt_manager_optional_params: Optional[bool] = False, ) -> Tuple[str, List[AllMessageValues], dict]: - from litellm.integrations.prompt_management_base import PromptManagementBase return PromptManagementBase.get_chat_completion_prompt( diff --git a/litellm/integrations/dotprompt/prompt_manager.py b/litellm/integrations/dotprompt/prompt_manager.py index fc5a325ffe..997a40d545 100644 --- a/litellm/integrations/dotprompt/prompt_manager.py +++ b/litellm/integrations/dotprompt/prompt_manager.py @@ -205,7 +205,7 @@ class PromptManager: """ # Get the template (versioned or base) template = self.get_prompt(prompt_id=prompt_id, version=version) - + if template is None: available_prompts = list(self.prompts.keys()) version_str = f" (version {version})" if version else "" @@ -266,11 +266,11 @@ class PromptManager: ) -> Optional[PromptTemplate]: """ Get a prompt template by ID and optional version. - + Args: prompt_id: The base prompt ID version: Optional version number. If provided, looks for {prompt_id}.v{version} - + Returns: The prompt template if found, None otherwise """ @@ -279,7 +279,7 @@ class PromptManager: versioned_id = f"{prompt_id}.v{version}" if versioned_id in self.prompts: return self.prompts[versioned_id] - + # Fall back to base prompt_id return self.prompts.get(prompt_id) diff --git a/litellm/integrations/email_templates/key_rotated_email.py b/litellm/integrations/email_templates/key_rotated_email.py index dab7172dc6..9e6dd41378 100644 --- a/litellm/integrations/email_templates/key_rotated_email.py +++ b/litellm/integrations/email_templates/key_rotated_email.py @@ -222,4 +222,3 @@ response = client.chat.completions.create(
""" - diff --git a/litellm/integrations/email_templates/templates.py b/litellm/integrations/email_templates/templates.py index 091351df2b..8df816dfec 100644 --- a/litellm/integrations/email_templates/templates.py +++ b/litellm/integrations/email_templates/templates.py @@ -131,4 +131,4 @@ MAX_BUDGET_ALERT_EMAIL_TEMPLATE = """ Best,
The LiteLLM team
-""" \ No newline at end of file +""" diff --git a/litellm/integrations/focus/focus_logger.py b/litellm/integrations/focus/focus_logger.py index ade1cf861b..f493d47e29 100644 --- a/litellm/integrations/focus/focus_logger.py +++ b/litellm/integrations/focus/focus_logger.py @@ -56,7 +56,9 @@ class FocusLogger(CustomLogger): self.interval_seconds = int(raw_interval) if raw_interval is not None else None env_prefix = os.getenv("FOCUS_PREFIX") self.prefix: str = ( - prefix if prefix is not None else (env_prefix if env_prefix else "focus_exports") + prefix + if prefix is not None + else (env_prefix if env_prefix else "focus_exports") ) self._destination_config = destination_config @@ -208,4 +210,5 @@ class FocusLogger(CustomLogger): frequency=self.frequency, ) + __all__ = ["FocusLogger"] diff --git a/litellm/integrations/gcs_bucket/gcs_bucket.py b/litellm/integrations/gcs_bucket/gcs_bucket.py index 0f1ba4a409..65296bafcf 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket.py @@ -34,7 +34,10 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): os.getenv("GCS_FLUSH_INTERVAL", GCS_DEFAULT_FLUSH_INTERVAL_SECONDS) ) self.use_batched_logging = ( - os.getenv("GCS_USE_BATCHED_LOGGING", str(GCS_DEFAULT_USE_BATCHED_LOGGING).lower()).lower() == "true" + os.getenv( + "GCS_USE_BATCHED_LOGGING", str(GCS_DEFAULT_USE_BATCHED_LOGGING).lower() + ).lower() + == "true" ) self.flush_lock = asyncio.Lock() super().__init__( @@ -112,9 +115,9 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): def _drain_queue_batch(self) -> List[GCSLogQueueItem]: """ Drain items from the queue (non-blocking), respecting batch_size limit. - + This prevents unbounded queue growth when processing is slower than log accumulation. - + Returns: List of items to process, up to batch_size items """ @@ -137,33 +140,45 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): """ Extract a synchronous grouping key from kwargs to group items by GCS config. This allows us to batch items with the same bucket/credentials together. - + Returns a string key that uniquely identifies the GCS config combination. This key may contain sensitive information (bucket names, paths) - use _sanitize_config_key() for logging purposes. """ - standard_callback_dynamic_params = kwargs.get("standard_callback_dynamic_params", None) or {} - - bucket_name = standard_callback_dynamic_params.get("gcs_bucket_name", None) or self.BUCKET_NAME or "default" - path_service_account = standard_callback_dynamic_params.get("gcs_path_service_account", None) or self.path_service_account_json or "default" - + standard_callback_dynamic_params = ( + kwargs.get("standard_callback_dynamic_params", None) or {} + ) + + bucket_name = ( + standard_callback_dynamic_params.get("gcs_bucket_name", None) + or self.BUCKET_NAME + or "default" + ) + path_service_account = ( + standard_callback_dynamic_params.get("gcs_path_service_account", None) + or self.path_service_account_json + or "default" + ) + return f"{bucket_name}|{path_service_account}" - + def _sanitize_config_key(self, config_key: str) -> str: """ Create a sanitized version of the config key for logging. Uses a hash to avoid exposing sensitive bucket names or service account paths. - + Returns a short hash prefix for safe logging. """ - hash_obj = hashlib.sha256(config_key.encode('utf-8')) + hash_obj = hashlib.sha256(config_key.encode("utf-8")) return f"config-{hash_obj.hexdigest()[:8]}" - - def _group_items_by_config(self, items: List[GCSLogQueueItem]) -> Dict[str, List[GCSLogQueueItem]]: + + def _group_items_by_config( + self, items: List[GCSLogQueueItem] + ) -> Dict[str, List[GCSLogQueueItem]]: """ Group items by their GCS config (bucket + credentials). This ensures items with different configs are processed separately. - + Returns a dict mapping config_key -> list of items with that config. """ grouped: Dict[str, List[GCSLogQueueItem]] = {} @@ -186,18 +201,20 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): lines.append(json_line) return "\n".join(lines) - async def _send_grouped_batch(self, items: List[GCSLogQueueItem], config_key: str) -> Tuple[int, int]: + async def _send_grouped_batch( + self, items: List[GCSLogQueueItem], config_key: str + ) -> Tuple[int, int]: """ Send a batch of items that share the same GCS config. - + Returns: (success_count, error_count) """ if not items: return (0, 0) - + first_kwargs = items[0]["kwargs"] - + try: gcs_logging_config: GCSLoggingConfig = await self.get_gcs_logging_config( first_kwargs @@ -208,23 +225,25 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): service_account_json=gcs_logging_config["path_service_account"], ) bucket_name = gcs_logging_config["bucket_name"] - - current_date = self._get_object_date_from_datetime(datetime.now(timezone.utc)) + + current_date = self._get_object_date_from_datetime( + datetime.now(timezone.utc) + ) batch_id = f"{int(time.time() * 1000)}-{uuid.uuid4().hex[:8]}" object_name = self._generate_batch_object_name(current_date, batch_id) combined_payload = self._combine_payloads_to_ndjson(items) - + await self._log_json_data_on_gcs( headers=headers, bucket_name=bucket_name, object_name=object_name, logging_payload=combined_payload, ) - + success_count = len(items) error_count = 0 return (success_count, error_count) - + except Exception as e: success_count = 0 error_count = len(items) @@ -255,13 +274,13 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): service_account_json=gcs_logging_config["path_service_account"], ) bucket_name = gcs_logging_config["bucket_name"] - + object_name = self._get_object_name( kwargs=item["kwargs"], logging_payload=item["payload"], response_obj=item["response_obj"], ) - + await self._log_json_data_on_gcs( headers=headers, bucket_name=bucket_name, @@ -289,7 +308,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): if self.use_batched_logging: grouped_items = self._group_items_by_config(items_to_process) - + for config_key, group_items in grouped_items.items(): await self._send_grouped_batch(group_items, config_key) else: diff --git a/litellm/integrations/gcs_bucket/gcs_bucket_base.py b/litellm/integrations/gcs_bucket/gcs_bucket_base.py index b1db9ec958..923f613291 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket_base.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket_base.py @@ -7,7 +7,7 @@ from litellm.integrations.gcs_bucket.gcs_bucket_mock_client import ( create_mock_gcs_client, mock_vertex_auth_methods, ) - + from litellm._logging import verbose_logger from litellm.integrations.custom_batch_logger import CustomBatchLogger @@ -28,11 +28,11 @@ IAM_AUTH_KEY = "IAM_AUTH" class GCSBucketBase(CustomBatchLogger): def __init__(self, bucket_name: Optional[str] = None, **kwargs) -> None: self.is_mock_mode = should_use_gcs_mock() - + if self.is_mock_mode: mock_vertex_auth_methods() create_mock_gcs_client() - + self.async_httpx_client = get_async_httpx_client( llm_provider=httpxSpecialProvider.LoggingCallback ) @@ -85,10 +85,10 @@ class GCSBucketBase(CustomBatchLogger): from litellm import vertex_chat_completion # Get project_id from environment if available, otherwise None - # This helps support use of this library to auth to pull secrets + # This helps support use of this library to auth to pull secrets # from Secret Manager. project_id = os.getenv("GOOGLE_SECRET_MANAGER_PROJECT_ID") - + _auth_header, vertex_project = vertex_chat_completion._ensure_access_token( credentials=self.path_service_account_json, project_id=project_id, diff --git a/litellm/integrations/gcs_bucket/gcs_bucket_mock_client.py b/litellm/integrations/gcs_bucket/gcs_bucket_mock_client.py index 2d14f5eb96..1761fe010c 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket_mock_client.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket_mock_client.py @@ -11,7 +11,11 @@ Usage: import asyncio from litellm._logging import verbose_logger -from litellm.integrations.mock_client_factory import MockClientConfig, create_mock_client_factory, MockResponse +from litellm.integrations.mock_client_factory import ( + MockClientConfig, + create_mock_client_factory, + MockResponse, +) # Use factory for POST handler _config = MockClientConfig( @@ -34,10 +38,14 @@ _mocks_initialized = False # Default mock latency in seconds (simulates network round-trip) # Typical GCS API calls take 100-300ms for uploads, 50-150ms for GET/DELETE -_MOCK_LATENCY_SECONDS = float(__import__("os").getenv("GCS_MOCK_LATENCY_MS", "150")) / 1000.0 +_MOCK_LATENCY_SECONDS = ( + float(__import__("os").getenv("GCS_MOCK_LATENCY_MS", "150")) / 1000.0 +) -async def _mock_async_handler_get(self, url, params=None, headers=None, follow_redirects=None): +async def _mock_async_handler_get( + self, url, params=None, headers=None, follow_redirects=None +): """Monkey-patched AsyncHTTPHandler.get that intercepts GCS calls.""" # Only mock GCS API calls if isinstance(url, str) and "storage.googleapis.com" in url: @@ -86,14 +94,30 @@ async def _mock_async_handler_get(self, url, params=None, headers=None, follow_r status_code=200, json_data=mock_payload, url=url, - elapsed_seconds=_MOCK_LATENCY_SECONDS + elapsed_seconds=_MOCK_LATENCY_SECONDS, ) if _original_async_handler_get is not None: - return await _original_async_handler_get(self, url=url, params=params, headers=headers, follow_redirects=follow_redirects) + return await _original_async_handler_get( + self, + url=url, + params=params, + headers=headers, + follow_redirects=follow_redirects, + ) raise RuntimeError("Original AsyncHTTPHandler.get not available") -async def _mock_async_handler_delete(self, url, data=None, json=None, params=None, headers=None, timeout=None, stream=False, content=None): +async def _mock_async_handler_delete( + self, + url, + data=None, + json=None, + params=None, + headers=None, + timeout=None, + stream=False, + content=None, +): """Monkey-patched AsyncHTTPHandler.delete that intercepts GCS calls.""" # Only mock GCS API calls if isinstance(url, str) and "storage.googleapis.com" in url: @@ -104,49 +128,61 @@ async def _mock_async_handler_delete(self, url, data=None, json=None, params=Non status_code=204, json_data=None, # Empty body for DELETE url=url, - elapsed_seconds=_MOCK_LATENCY_SECONDS + elapsed_seconds=_MOCK_LATENCY_SECONDS, ) if _original_async_handler_delete is not None: - return await _original_async_handler_delete(self, url=url, data=data, json=json, params=params, headers=headers, timeout=timeout, stream=stream, content=content) + return await _original_async_handler_delete( + self, + url=url, + data=data, + json=json, + params=params, + headers=headers, + timeout=timeout, + stream=stream, + content=content, + ) raise RuntimeError("Original AsyncHTTPHandler.delete not available") def create_mock_gcs_client(): """ Monkey-patch AsyncHTTPHandler methods to intercept GCS calls. - + AsyncHTTPHandler is used by LiteLLM's get_async_httpx_client() which is what GCSBucketBase uses for making API calls. - + This function is idempotent - it only initializes mocks once, even if called multiple times. """ global _original_async_handler_get, _original_async_handler_delete, _mocks_initialized - + # Use factory for POST handler _create_mock_gcs_post() - + # If already initialized, skip GET/DELETE patching if _mocks_initialized: return - + verbose_logger.debug("[GCS MOCK] Initializing GCS GET/DELETE handlers...") - + # Patch GET and DELETE handlers (GCS-specific) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler - + if _original_async_handler_get is None: _original_async_handler_get = AsyncHTTPHandler.get AsyncHTTPHandler.get = _mock_async_handler_get # type: ignore verbose_logger.debug("[GCS MOCK] Patched AsyncHTTPHandler.get") - + if _original_async_handler_delete is None: _original_async_handler_delete = AsyncHTTPHandler.delete AsyncHTTPHandler.delete = _mock_async_handler_delete # type: ignore verbose_logger.debug("[GCS MOCK] Patched AsyncHTTPHandler.delete") - - verbose_logger.debug(f"[GCS MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS*1000:.0f}ms") + + verbose_logger.debug( + f"[GCS MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS*1000:.0f}ms" + ) verbose_logger.debug("[GCS MOCK] GCS mock client initialization complete") - + _mocks_initialized = True @@ -154,38 +190,64 @@ def mock_vertex_auth_methods(): """ Monkey-patch Vertex AI auth methods to return fake tokens. This prevents auth failures when GCS_MOCK is enabled. - + This function is idempotent - it only patches once, even if called multiple times. """ from litellm.llms.vertex_ai.vertex_llm_base import VertexBase - + # Store original methods if not already stored - if not hasattr(VertexBase, '_original_ensure_access_token_async'): - setattr(VertexBase, '_original_ensure_access_token_async', VertexBase._ensure_access_token_async) - setattr(VertexBase, '_original_ensure_access_token', VertexBase._ensure_access_token) - setattr(VertexBase, '_original_get_token_and_url', VertexBase._get_token_and_url) - - async def _mock_ensure_access_token_async(self, credentials, project_id, custom_llm_provider): + if not hasattr(VertexBase, "_original_ensure_access_token_async"): + setattr( + VertexBase, + "_original_ensure_access_token_async", + VertexBase._ensure_access_token_async, + ) + setattr( + VertexBase, "_original_ensure_access_token", VertexBase._ensure_access_token + ) + setattr( + VertexBase, "_original_get_token_and_url", VertexBase._get_token_and_url + ) + + async def _mock_ensure_access_token_async( + self, credentials, project_id, custom_llm_provider + ): """Mock async auth method - returns fake token.""" - verbose_logger.debug("[GCS MOCK] Vertex AI auth: _ensure_access_token_async called") + verbose_logger.debug( + "[GCS MOCK] Vertex AI auth: _ensure_access_token_async called" + ) return ("mock-gcs-token", "mock-project-id") - - def _mock_ensure_access_token(self, credentials, project_id, custom_llm_provider): + + def _mock_ensure_access_token( + self, credentials, project_id, custom_llm_provider + ): """Mock sync auth method - returns fake token.""" - verbose_logger.debug("[GCS MOCK] Vertex AI auth: _ensure_access_token called") + verbose_logger.debug( + "[GCS MOCK] Vertex AI auth: _ensure_access_token called" + ) return ("mock-gcs-token", "mock-project-id") - - def _mock_get_token_and_url(self, model, auth_header, vertex_credentials, vertex_project, - vertex_location, gemini_api_key, stream, custom_llm_provider, api_base): + + def _mock_get_token_and_url( + self, + model, + auth_header, + vertex_credentials, + vertex_project, + vertex_location, + gemini_api_key, + stream, + custom_llm_provider, + api_base, + ): """Mock get_token_and_url - returns fake token.""" verbose_logger.debug("[GCS MOCK] Vertex AI auth: _get_token_and_url called") return ("mock-gcs-token", "https://storage.googleapis.com") - + # Patch the methods VertexBase._ensure_access_token_async = _mock_ensure_access_token_async # type: ignore VertexBase._ensure_access_token = _mock_ensure_access_token # type: ignore VertexBase._get_token_and_url = _mock_get_token_and_url # type: ignore - + verbose_logger.debug("[GCS MOCK] Patched Vertex AI auth methods") diff --git a/litellm/integrations/generic_api/generic_api_callback.py b/litellm/integrations/generic_api/generic_api_callback.py index 1c62ce9fcc..9a8060520d 100644 --- a/litellm/integrations/generic_api/generic_api_callback.py +++ b/litellm/integrations/generic_api/generic_api_callback.py @@ -164,7 +164,11 @@ class GenericAPILogger(CustomBatchLogger): self.callback_name: Optional[str] = callback_name # Validate and store log_format - if log_format is not None and log_format not in ["json_array", "ndjson", "single"]: + if log_format is not None and log_format not in [ + "json_array", + "ndjson", + "single", + ]: raise ValueError( f"Invalid log_format: {log_format}. Must be one of: 'json_array', 'ndjson', 'single'" ) diff --git a/litellm/integrations/generic_prompt_management/generic_prompt_manager.py b/litellm/integrations/generic_prompt_management/generic_prompt_manager.py index 9490d9fde1..858bfd458b 100644 --- a/litellm/integrations/generic_prompt_management/generic_prompt_manager.py +++ b/litellm/integrations/generic_prompt_management/generic_prompt_manager.py @@ -120,7 +120,6 @@ class GenericPromptManager(CustomPromptManagement): http_client = _get_httpx_client() try: - response = http_client.get( url, params=params, @@ -325,7 +324,6 @@ class GenericPromptManager(CustomPromptManagement): prompt_label: Optional[str] = None, prompt_version: Optional[int] = None, ) -> PromptManagementClient: - # Check cache first cached_prompt = self._common_caching_logic( prompt_id=prompt_id, diff --git a/litellm/integrations/gitlab/__init__.py b/litellm/integrations/gitlab/__init__.py index c73a23b687..24e7ddea9e 100644 --- a/litellm/integrations/gitlab/__init__.py +++ b/litellm/integrations/gitlab/__init__.py @@ -39,11 +39,8 @@ def prompt_initializer( gitlab_config = getattr(litellm_params, "gitlab_config", None) prompt_id = getattr(litellm_params, "prompt_id", None) - if not gitlab_config: - raise ValueError( - "gitlab_config is required for gitlab prompt integration" - ) + raise ValueError("gitlab_config is required for gitlab prompt integration") try: gitlab_prompt_manager = GitLabPromptManager( @@ -55,9 +52,10 @@ def prompt_initializer( except Exception as e: raise e + def _gitlab_prompt_initializer( - litellm_params: PromptLiteLLMParams, - prompt: PromptSpec, + litellm_params: PromptLiteLLMParams, + prompt: PromptSpec, ) -> CustomPromptManagement: """ Build a GitLab-backed prompt manager for this prompt. diff --git a/litellm/integrations/gitlab/gitlab_client.py b/litellm/integrations/gitlab/gitlab_client.py index ce03a35d48..60f7325618 100644 --- a/litellm/integrations/gitlab/gitlab_client.py +++ b/litellm/integrations/gitlab/gitlab_client.py @@ -45,7 +45,7 @@ class GitLabClient: self.auth_method = config.get("auth_method", "token") # 'token' or 'oauth' self.branch = config.get("branch", None) if not self.branch: - self.branch = 'main' + self.branch = "main" self.tag = config.get("tag") self.base_url = config.get("base_url", "https://gitlab.com/api/v4") @@ -86,7 +86,13 @@ class GitLabClient: ref_q = quote(ref or self.ref, safe="") return f"{self.base_url}/projects/{self._project_enc}/repository/files/{file_enc}?ref={ref_q}" - def _tree_url(self, directory_path: str = "", recursive: bool = False, *, ref: Optional[str] = None) -> str: + def _tree_url( + self, + directory_path: str = "", + recursive: bool = False, + *, + ref: Optional[str] = None, + ) -> str: path_q = f"&path={quote(directory_path, safe='')}" if directory_path else "" rec_q = "&recursive=true" if recursive else "" ref_q = quote(ref or self.ref, safe="") @@ -102,7 +108,9 @@ class GitLabClient: raise ValueError("ref must be a non-empty string") self.ref = ref - def get_file_content(self, file_path: str, *, ref: Optional[str] = None) -> Optional[str]: + def get_file_content( + self, file_path: str, *, ref: Optional[str] = None + ) -> Optional[str]: """ Fetch the content of a file from the GitLab repository at the given ref (tag, branch, or commit SHA). If `ref` is None, uses self.ref. @@ -124,7 +132,11 @@ class GitLabClient: resp.raise_for_status() ctype = (resp.headers.get("content-type") or "").lower() - if ctype.startswith("text/") or "charset=" in ctype or ctype.startswith("application/json"): + if ( + ctype.startswith("text/") + or "charset=" in ctype + or ctype.startswith("application/json") + ): return resp.text try: return resp.content.decode("utf-8") @@ -140,10 +152,14 @@ class GitLabClient: f"Access denied to file '{file_path}'. Check your GitLab permissions for project '{self.project}'." ) if status == 401: - raise Exception("Authentication failed. Check your GitLab token and auth_method.") + raise Exception( + "Authentication failed. Check your GitLab token and auth_method." + ) raise Exception(f"Failed to fetch file '{file_path}': {e}") - def _get_file_content_via_json(self, file_path: str, *, ref: Optional[str] = None) -> Optional[str]: + def _get_file_content_via_json( + self, file_path: str, *, ref: Optional[str] = None + ) -> Optional[str]: """ Fallback for get_file_content(): use the JSON file API which returns base64 content. """ @@ -171,16 +187,20 @@ class GitLabClient: f"Access denied to file '{file_path}'. Check your GitLab permissions for project '{self.project}'." ) if status == 401: - raise Exception("Authentication failed. Check your GitLab token and auth_method.") - raise Exception(f"Failed to fetch file '{file_path}' via JSON endpoint: {e}") + raise Exception( + "Authentication failed. Check your GitLab token and auth_method." + ) + raise Exception( + f"Failed to fetch file '{file_path}' via JSON endpoint: {e}" + ) def list_files( - self, - directory_path: str = "", - file_extension: str = ".prompt", - recursive: bool = False, - *, - ref: Optional[str] = None, + self, + directory_path: str = "", + file_extension: str = ".prompt", + recursive: bool = False, + *, + ref: Optional[str] = None, ) -> List[str]: """ List files in a directory with a specific extension using the repository tree API. @@ -220,7 +240,9 @@ class GitLabClient: f"Access denied to directory '{directory_path}'. Check your GitLab permissions for project '{self.project}'." ) if status == 401: - raise Exception("Authentication failed. Check your GitLab token and auth_method.") + raise Exception( + "Authentication failed. Check your GitLab token and auth_method." + ) raise Exception(f"Failed to list files in '{directory_path}': {e}") def get_repository_info(self) -> Dict[str, Any]: @@ -252,7 +274,9 @@ class GitLabClient: except Exception as e: raise Exception(f"Failed to get branches: {e}") - def get_file_metadata(self, file_path: str, *, ref: Optional[str] = None) -> Optional[Dict[str, Any]]: + def get_file_metadata( + self, file_path: str, *, ref: Optional[str] = None + ) -> Optional[Dict[str, Any]]: """ Get minimal metadata about a file via RAW endpoint headers at a given ref. diff --git a/litellm/integrations/helicone.py b/litellm/integrations/helicone.py index 51e6699c5f..376952033a 100644 --- a/litellm/integrations/helicone.py +++ b/litellm/integrations/helicone.py @@ -31,8 +31,10 @@ class HeliconeLogger: self.is_mock_mode = should_use_helicone_mock() if self.is_mock_mode: create_mock_helicone_client() - verbose_logger.info("[HELICONE MOCK] Helicone logger initialized in mock mode") - + verbose_logger.info( + "[HELICONE MOCK] Helicone logger initialized in mock mode" + ) + self.provider_url = "https://api.openai.com/v1" self.key = os.getenv("HELICONE_API_KEY") self.api_base = os.getenv("HELICONE_API_BASE") or "https://api.hconeai.com" @@ -111,7 +113,7 @@ class HeliconeLogger: for header_key in proxy_headers: if header_key.startswith("helicone_"): metadata[header_key] = proxy_headers.get(header_key) - + # Remove OpenTelemetry span from metadata as it's not JSON serializable # The span is used internally for tracing but shouldn't be logged to external services if "litellm_parent_otel_span" in metadata: @@ -134,14 +136,17 @@ class HeliconeLogger: metadata = self.add_metadata_from_header(litellm_params, metadata) # Check if model is a vertex_ai model - is_vertex_ai = custom_llm_provider == "vertex_ai" or model.startswith("vertex_ai/") + is_vertex_ai = custom_llm_provider == "vertex_ai" or model.startswith( + "vertex_ai/" + ) model = ( model if any( accepted_model in model for accepted_model in self.helicone_model_list - ) or is_vertex_ai + ) + or is_vertex_ai else "gpt-3.5-turbo" ) provider_request = {"model": model, "messages": messages} @@ -208,7 +213,9 @@ class HeliconeLogger: response = litellm.module_level_client.post(url, headers=headers, json=data) if response.status_code == 200: if self.is_mock_mode: - print_verbose("[HELICONE MOCK] Helicone Logging - Successfully mocked!") + print_verbose( + "[HELICONE MOCK] Helicone Logging - Successfully mocked!" + ) else: print_verbose("Helicone Logging - Success!") else: diff --git a/litellm/integrations/helicone_mock_client.py b/litellm/integrations/helicone_mock_client.py index 0f4670a1d2..c2d3dfdf5b 100644 --- a/litellm/integrations/helicone_mock_client.py +++ b/litellm/integrations/helicone_mock_client.py @@ -8,7 +8,10 @@ Usage: Set HELICONE_MOCK=true in environment variables or config to enable mock mode. """ -from litellm.integrations.mock_client_factory import MockClientConfig, create_mock_client_factory +from litellm.integrations.mock_client_factory import ( + MockClientConfig, + create_mock_client_factory, +) # Create mock client using factory # Helicone uses HTTPHandler which internally uses httpx.Client.send(), not httpx.Client.post() @@ -29,4 +32,6 @@ _config = MockClientConfig( patch_http_handler=True, # Patch HTTPHandler.post directly ) -create_mock_helicone_client, should_use_helicone_mock = create_mock_client_factory(_config) +create_mock_helicone_client, should_use_helicone_mock = create_mock_client_factory( + _config +) diff --git a/litellm/integrations/humanloop.py b/litellm/integrations/humanloop.py index 369df5ee0b..11414869a6 100644 --- a/litellm/integrations/humanloop.py +++ b/litellm/integrations/humanloop.py @@ -162,11 +162,7 @@ class HumanloopLogger(CustomLogger): prompt_version: Optional[int] = None, ignore_prompt_manager_model: Optional[bool] = False, ignore_prompt_manager_optional_params: Optional[bool] = False, - ) -> Tuple[ - str, - List[AllMessageValues], - dict, - ]: + ) -> Tuple[str, List[AllMessageValues], dict,]: humanloop_api_key = dynamic_callback_params.get( "humanloop_api_key" ) or get_secret_str("HUMANLOOP_API_KEY") diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 7bf97665fd..9d6ddd0f1e 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -123,7 +123,7 @@ class LangFuseLogger: self.langfuse_flush_interval = LangFuseLogger._get_langfuse_flush_interval( flush_interval ) - + if should_use_langfuse_mock(): self.langfuse_client = create_mock_langfuse_client() self.is_mock_mode = True @@ -607,9 +607,7 @@ class LangFuseLogger: # Use standard_logging_object.trace_id if available (when trace_id from metadata is None) # This allows standard trace_id to be used when provided in standard_logging_object if trace_id is None and standard_logging_object is not None: - trace_id = cast( - Optional[str], standard_logging_object.get("trace_id") - ) + trace_id = cast(Optional[str], standard_logging_object.get("trace_id")) # Fallback to litellm_call_id if no trace_id found if trace_id is None: trace_id = litellm_call_id diff --git a/litellm/integrations/langfuse/langfuse_mock_client.py b/litellm/integrations/langfuse/langfuse_mock_client.py index 8ed6cff8d4..b7862274f6 100644 --- a/litellm/integrations/langfuse/langfuse_mock_client.py +++ b/litellm/integrations/langfuse/langfuse_mock_client.py @@ -9,7 +9,10 @@ Usage: """ import httpx -from litellm.integrations.mock_client_factory import MockClientConfig, create_mock_client_factory +from litellm.integrations.mock_client_factory import ( + MockClientConfig, + create_mock_client_factory, +) # Create mock client using factory _config = MockClientConfig( @@ -26,7 +29,11 @@ _config = MockClientConfig( patch_sync_client=True, ) -_create_mock_langfuse_client_internal, should_use_langfuse_mock = create_mock_client_factory(_config) +( + _create_mock_langfuse_client_internal, + should_use_langfuse_mock, +) = create_mock_client_factory(_config) + # Langfuse needs to return an httpx.Client instance def create_mock_langfuse_client(): diff --git a/litellm/integrations/langfuse/langfuse_prompt_management.py b/litellm/integrations/langfuse/langfuse_prompt_management.py index 3986fc6a6e..f0e1e30b68 100644 --- a/litellm/integrations/langfuse/langfuse_prompt_management.py +++ b/litellm/integrations/langfuse/langfuse_prompt_management.py @@ -318,7 +318,7 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge ) except Exception as e: from litellm._logging import verbose_logger - + verbose_logger.exception( f"Langfuse Layer Error - Exception occurred while logging success event: {str(e)}" ) @@ -351,7 +351,7 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge ) except Exception as e: from litellm._logging import verbose_logger - + verbose_logger.exception( f"Langfuse Layer Error - Exception occurred while logging failure event: {str(e)}" ) diff --git a/litellm/integrations/langsmith.py b/litellm/integrations/langsmith.py index ebd005f880..03845af521 100644 --- a/litellm/integrations/langsmith.py +++ b/litellm/integrations/langsmith.py @@ -50,11 +50,13 @@ class LangsmithLogger(CustomBatchLogger): self.flush_lock = asyncio.Lock() super().__init__(**kwargs, flush_lock=self.flush_lock) self.is_mock_mode = should_use_langsmith_mock() - + if self.is_mock_mode: create_mock_langsmith_client() - verbose_logger.debug("[LANGSMITH MOCK] LangSmith logger initialized in mock mode") - + verbose_logger.debug( + "[LANGSMITH MOCK] LangSmith logger initialized in mock mode" + ) + self.default_credentials = self.get_credentials_from_env( langsmith_api_key=langsmith_api_key, langsmith_project=langsmith_project, @@ -399,7 +401,9 @@ class LangsmithLogger(CustomBatchLogger): "Sending batch of %s runs to Langsmith", len(elements_to_log) ) if self.is_mock_mode: - verbose_logger.debug("[LANGSMITH MOCK] Mock mode enabled - API calls will be intercepted") + verbose_logger.debug( + "[LANGSMITH MOCK] Mock mode enabled - API calls will be intercepted" + ) response = await self.async_httpx_client.post( url=url, json={"post": elements_to_log}, diff --git a/litellm/integrations/langsmith_mock_client.py b/litellm/integrations/langsmith_mock_client.py index ef60290823..0226bdecc2 100644 --- a/litellm/integrations/langsmith_mock_client.py +++ b/litellm/integrations/langsmith_mock_client.py @@ -8,7 +8,10 @@ Usage: Set LANGSMITH_MOCK=true in environment variables or config to enable mock mode. """ -from litellm.integrations.mock_client_factory import MockClientConfig, create_mock_client_factory +from litellm.integrations.mock_client_factory import ( + MockClientConfig, + create_mock_client_factory, +) # Create mock client using factory _config = MockClientConfig( @@ -26,4 +29,6 @@ _config = MockClientConfig( patch_sync_client=False, ) -create_mock_langsmith_client, should_use_langsmith_mock = create_mock_client_factory(_config) +create_mock_langsmith_client, should_use_langsmith_mock = create_mock_client_factory( + _config +) diff --git a/litellm/integrations/levo/levo.py b/litellm/integrations/levo/levo.py index 562f2fd906..4b08ce50f7 100644 --- a/litellm/integrations/levo/levo.py +++ b/litellm/integrations/levo/levo.py @@ -6,7 +6,9 @@ from litellm.integrations.opentelemetry import OpenTelemetry if TYPE_CHECKING: from opentelemetry.trace import Span as _Span - from litellm.integrations.opentelemetry import OpenTelemetryConfig as _OpenTelemetryConfig + from litellm.integrations.opentelemetry import ( + OpenTelemetryConfig as _OpenTelemetryConfig, + ) from litellm.types.integrations.arize import Protocol as _Protocol Protocol = _Protocol @@ -114,4 +116,3 @@ class LevoLogger(OpenTelemetry): "status": "unhealthy", "error_message": str(e), } - diff --git a/litellm/integrations/mock_client_factory.py b/litellm/integrations/mock_client_factory.py index 2f04fae9f7..3f2f0ae5b6 100644 --- a/litellm/integrations/mock_client_factory.py +++ b/litellm/integrations/mock_client_factory.py @@ -19,16 +19,21 @@ from litellm._logging import verbose_logger @dataclass class MockClientConfig: """Configuration for creating a mock client.""" + name: str # e.g., "GCS", "LANGFUSE", "LANGSMITH", "DATADOG" env_var: str # e.g., "GCS_MOCK", "LANGFUSE_MOCK" default_latency_ms: int = 100 # Default mock latency in milliseconds default_status_code: int = 200 # Default HTTP status code default_json_data: Optional[Dict] = None # Default JSON response data - url_matchers: Optional[List[str]] = None # List of strings to match in URLs (e.g., ["storage.googleapis.com"]) + url_matchers: Optional[ + List[str] + ] = None # List of strings to match in URLs (e.g., ["storage.googleapis.com"]) patch_async_handler: bool = True # Whether to patch AsyncHTTPHandler.post patch_sync_client: bool = False # Whether to patch httpx.Client.post - patch_http_handler: bool = False # Whether to patch HTTPHandler.post (for sync calls that use HTTPHandler) - + patch_http_handler: bool = ( + False # Whether to patch HTTPHandler.post (for sync calls that use HTTPHandler) + ) + def __post_init__(self): """Ensure url_matchers is a list.""" if self.url_matchers is None: @@ -37,8 +42,14 @@ class MockClientConfig: class MockResponse: """Generic mock httpx.Response that satisfies API requirements.""" - - def __init__(self, status_code: int = 200, json_data: Optional[Dict] = None, url: Optional[str] = None, elapsed_seconds: float = 0.0): + + def __init__( + self, + status_code: int = 200, + json_data: Optional[Dict] = None, + url: Optional[str] = None, + elapsed_seconds: float = 0.0, + ): self.status_code = status_code self._json_data = json_data or {"status": "success"} self.headers = httpx.Headers({}) @@ -49,25 +60,25 @@ class MockResponse: self.elapsed = timedelta(seconds=elapsed_seconds) self._text = json.dumps(self._json_data) if json_data else "" self._content = self._text.encode("utf-8") - + @property def text(self) -> str: """Return response text.""" return self._text - + @property def content(self) -> bytes: """Return response content.""" return self._content - + def json(self) -> Dict: """Return JSON response data.""" return self._json_data - + def read(self) -> bytes: """Read response content.""" return self._content - + def raise_for_status(self): """Raise exception for error status codes.""" if self.status_code >= 400: @@ -80,17 +91,17 @@ def _is_url_match(url, matchers: List[str]) -> bool: parsed_url = httpx.URL(url) if isinstance(url, str) else url url_str = str(parsed_url).lower() hostname = parsed_url.host or "" - + for matcher in matchers: if matcher.lower() in url_str or matcher.lower() in hostname.lower(): return True - + # Also check for localhost with matcher in path if hostname in ("localhost", "127.0.0.1"): for matcher in matchers: if matcher.lower() in url_str: return True - + return False except Exception: return False @@ -99,7 +110,7 @@ def _is_url_match(url, matchers: List[str]) -> bool: def create_mock_client_factory(config: MockClientConfig): # noqa: PLR0915 """ Factory function that creates mock client functions based on configuration. - + Returns: tuple: (create_mock_client_func, should_use_mock_func) """ @@ -108,19 +119,34 @@ def create_mock_client_factory(config: MockClientConfig): # noqa: PLR0915 _original_sync_client_post = None _original_http_handler_post = None _mocks_initialized = False - + # Calculate mock latency import os + latency_env = f"{config.name.upper()}_MOCK_LATENCY_MS" - _MOCK_LATENCY_SECONDS = float(os.getenv(latency_env, str(config.default_latency_ms))) / 1000.0 - + _MOCK_LATENCY_SECONDS = ( + float(os.getenv(latency_env, str(config.default_latency_ms))) / 1000.0 + ) + # Create URL matcher function def _is_mock_url(url) -> bool: # url_matchers is guaranteed to be a list after __post_init__ return _is_url_match(url, cast(List[str], config.url_matchers)) - + # Create async handler mock - async def _mock_async_handler_post(self, url, data=None, json=None, params=None, headers=None, timeout=None, stream=False, logging_obj=None, files=None, content=None): + async def _mock_async_handler_post( + self, + url, + data=None, + json=None, + params=None, + headers=None, + timeout=None, + stream=False, + logging_obj=None, + files=None, + content=None, + ): """Monkey-patched AsyncHTTPHandler.post that intercepts API calls.""" if isinstance(url, str) and _is_mock_url(url): verbose_logger.info(f"[{config.name} MOCK] POST to {url}") @@ -129,12 +155,24 @@ def create_mock_client_factory(config: MockClientConfig): # noqa: PLR0915 status_code=config.default_status_code, json_data=config.default_json_data, url=url, - elapsed_seconds=_MOCK_LATENCY_SECONDS + elapsed_seconds=_MOCK_LATENCY_SECONDS, ) if _original_async_handler_post is not None: - return await _original_async_handler_post(self, url=url, data=data, json=json, params=params, headers=headers, timeout=timeout, stream=stream, logging_obj=logging_obj, files=files, content=content) + return await _original_async_handler_post( + self, + url=url, + data=data, + json=json, + params=params, + headers=headers, + timeout=timeout, + stream=stream, + logging_obj=logging_obj, + files=files, + content=content, + ) raise RuntimeError("Original AsyncHTTPHandler.post not available") - + # Create sync client mock def _mock_sync_client_post(self, url, **kwargs): """Monkey-patched httpx.Client.post that intercepts API calls.""" @@ -144,73 +182,108 @@ def create_mock_client_factory(config: MockClientConfig): # noqa: PLR0915 status_code=config.default_status_code, json_data=config.default_json_data, url=url, - elapsed_seconds=_MOCK_LATENCY_SECONDS + elapsed_seconds=_MOCK_LATENCY_SECONDS, ) if _original_sync_client_post is not None: return _original_sync_client_post(self, url, **kwargs) - + # Create HTTPHandler mock (for sync calls that use HTTPHandler.post) - def _mock_http_handler_post(self, url, data=None, json=None, params=None, headers=None, timeout=None, stream=False, files=None, content=None, logging_obj=None): + def _mock_http_handler_post( + self, + url, + data=None, + json=None, + params=None, + headers=None, + timeout=None, + stream=False, + files=None, + content=None, + logging_obj=None, + ): """Monkey-patched HTTPHandler.post that intercepts API calls.""" if isinstance(url, str) and _is_mock_url(url): verbose_logger.info(f"[{config.name} MOCK] POST to {url}") import time + time.sleep(_MOCK_LATENCY_SECONDS) return MockResponse( status_code=config.default_status_code, json_data=config.default_json_data, url=url, - elapsed_seconds=_MOCK_LATENCY_SECONDS + elapsed_seconds=_MOCK_LATENCY_SECONDS, ) if _original_http_handler_post is not None: - return _original_http_handler_post(self, url=url, data=data, json=json, params=params, headers=headers, timeout=timeout, stream=stream, files=files, content=content, logging_obj=logging_obj) + return _original_http_handler_post( + self, + url=url, + data=data, + json=json, + params=params, + headers=headers, + timeout=timeout, + stream=stream, + files=files, + content=content, + logging_obj=logging_obj, + ) raise RuntimeError("Original HTTPHandler.post not available") - + # Create mock client initialization function def create_mock_client(): """Initialize the mock client by patching HTTP handlers.""" nonlocal _original_async_handler_post, _original_sync_client_post, _original_http_handler_post, _mocks_initialized - + if _mocks_initialized: return - - verbose_logger.debug(f"[{config.name} MOCK] Initializing {config.name} mock client...") - + + verbose_logger.debug( + f"[{config.name} MOCK] Initializing {config.name} mock client..." + ) + if config.patch_async_handler and _original_async_handler_post is None: from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + _original_async_handler_post = AsyncHTTPHandler.post AsyncHTTPHandler.post = _mock_async_handler_post # type: ignore verbose_logger.debug(f"[{config.name} MOCK] Patched AsyncHTTPHandler.post") - + if config.patch_sync_client and _original_sync_client_post is None: _original_sync_client_post = httpx.Client.post httpx.Client.post = _mock_sync_client_post # type: ignore verbose_logger.debug(f"[{config.name} MOCK] Patched httpx.Client.post") - + if config.patch_http_handler and _original_http_handler_post is None: from litellm.llms.custom_httpx.http_handler import HTTPHandler + _original_http_handler_post = HTTPHandler.post HTTPHandler.post = _mock_http_handler_post # type: ignore verbose_logger.debug(f"[{config.name} MOCK] Patched HTTPHandler.post") - - verbose_logger.debug(f"[{config.name} MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS*1000:.0f}ms") - verbose_logger.debug(f"[{config.name} MOCK] {config.name} mock client initialization complete") - + + verbose_logger.debug( + f"[{config.name} MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS*1000:.0f}ms" + ) + verbose_logger.debug( + f"[{config.name} MOCK] {config.name} mock client initialization complete" + ) + _mocks_initialized = True - + # Create should_use_mock function def should_use_mock() -> bool: """Determine if mock mode should be enabled.""" import os from litellm.secret_managers.main import str_to_bool - + mock_mode = os.getenv(config.env_var, "false") result = str_to_bool(mock_mode) result = bool(result) if result is not None else False - + if result: - verbose_logger.info(f"{config.name} Mock Mode: ENABLED - API calls will be mocked") - + verbose_logger.info( + f"{config.name} Mock Mode: ENABLED - API calls will be mocked" + ) + return result - + return create_mock_client, should_use_mock diff --git a/litellm/integrations/openmeter.py b/litellm/integrations/openmeter.py index b8fb64ec28..5a8ab4bcc9 100644 --- a/litellm/integrations/openmeter.py +++ b/litellm/integrations/openmeter.py @@ -66,19 +66,19 @@ class OpenMeterLogger(CustomLogger): } user_param = kwargs.get("user", None) # end-user passed in via 'user' param - + # If no user provided directly, try to get it from token user_id if user_param is None: # Check if user_id is available from the API key metadata litellm_params = kwargs.get("litellm_params", {}) metadata = litellm_params.get("metadata", {}) user_api_key_user_id = metadata.get("user_api_key_user_id", None) - + if user_api_key_user_id is not None: user_param = user_api_key_user_id else: raise Exception("OpenMeter: user is required") - + # Ensure subject is always a string for OpenMeter API subject = str(user_param) diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index a77a6f73b1..7689a6cc7e 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -1611,9 +1611,8 @@ class OpenTelemetry(CustomLogger): # the litellm call ID so every call type can be correlated # across LiteLLM UI, Phoenix traces, and provider logs (Issue #8). response_id = ( - (response_obj.get("id") if response_obj else None) - or standard_logging_payload.get("id") - ) + response_obj.get("id") if response_obj else None + ) or standard_logging_payload.get("id") if response_id: self.safe_set_attribute( span=span, diff --git a/litellm/integrations/opik/opik_payload_builder/api.py b/litellm/integrations/opik/opik_payload_builder/api.py index 99dbea165e..e3ffab80ae 100644 --- a/litellm/integrations/opik/opik_payload_builder/api.py +++ b/litellm/integrations/opik/opik_payload_builder/api.py @@ -97,11 +97,11 @@ def build_opik_payload( # Always create a span usage = utils.create_usage_object(response_obj["usage"]) - + # Extract provider and cost provider = extractors.normalize_provider_name(kwargs.get("custom_llm_provider")) cost = kwargs.get("response_cost") - + span_payload = payload_builders.build_span_payload( project_name=current_project_name, trace_id=trace_id, diff --git a/litellm/integrations/opik/opik_payload_builder/extractors.py b/litellm/integrations/opik/opik_payload_builder/extractors.py index e4ff021778..9779ccddac 100644 --- a/litellm/integrations/opik/opik_payload_builder/extractors.py +++ b/litellm/integrations/opik/opik_payload_builder/extractors.py @@ -9,16 +9,16 @@ from litellm import _logging def normalize_provider_name(provider: Optional[str]) -> Optional[str]: """ Normalize LiteLLM provider names to standardized string names. - + Args: provider: LiteLLM internal provider name - + Returns: Normalized provider name or the original if no mapping exists """ if provider is None: return None - + # Provider mapping to names used in Opik provider_mapping = { "openai": "openai", @@ -30,7 +30,7 @@ def normalize_provider_name(provider: Optional[str]) -> Optional[str]: "bedrock_converse": "bedrock", "groq": "groq", } - + return provider_mapping.get(provider, provider) diff --git a/litellm/integrations/posthog.py b/litellm/integrations/posthog.py index c4b6e843d6..17bb56b8f1 100644 --- a/litellm/integrations/posthog.py +++ b/litellm/integrations/posthog.py @@ -45,12 +45,14 @@ class PostHogLogger(CustomBatchLogger): """ try: verbose_logger.debug("PostHog: in init posthog logger") - + self.is_mock_mode = should_use_posthog_mock() if self.is_mock_mode: create_mock_posthog_client() - verbose_logger.debug("[POSTHOG MOCK] PostHog logger initialized in mock mode") - + verbose_logger.debug( + "[POSTHOG MOCK] PostHog logger initialized in mock mode" + ) + if os.getenv("POSTHOG_API_KEY", None) is None: raise Exception("POSTHOG_API_KEY is not set, set 'POSTHOG_API_KEY=<>'") @@ -58,10 +60,10 @@ class PostHogLogger(CustomBatchLogger): llm_provider=httpxSpecialProvider.LoggingCallback ) self.sync_client = _get_httpx_client() - + self.POSTHOG_API_KEY = os.getenv("POSTHOG_API_KEY") posthog_api_url = os.getenv("POSTHOG_API_URL", "https://us.i.posthog.com") - self.posthog_host = posthog_api_url.rstrip('/') + self.posthog_host = posthog_api_url.rstrip("/") self.capture_url = f"{self.posthog_host}/batch/" self._async_initialized = False @@ -141,17 +143,17 @@ class PostHogLogger(CustomBatchLogger): verbose_logger.exception(f"PostHog Layer Error - {str(e)}") pass - async def _log_async_event(self, kwargs, response_obj=None, start_time=0.0, end_time=0.0): + async def _log_async_event( + self, kwargs, response_obj=None, start_time=0.0, end_time=0.0 + ): # Note: response_obj, start_time, end_time not used - all data comes from kwargs api_key, api_url = self._get_credentials_for_request(kwargs) event_payload = self.create_posthog_event_payload(kwargs) # Store event with its credentials for batch sending - self.log_queue.append({ - "event": event_payload, - "api_key": api_key, - "api_url": api_url - }) + self.log_queue.append( + {"event": event_payload, "api_key": api_key, "api_url": api_url} + ) verbose_logger.debug( f"PostHog, event added to queue. Will flush in {self.flush_interval} seconds..." ) @@ -159,7 +161,9 @@ class PostHogLogger(CustomBatchLogger): if len(self.log_queue) >= self.batch_size: await self.flush_queue() - def create_posthog_event_payload(self, kwargs: Dict[str, Any]) -> PostHogEventPayload: + def create_posthog_event_payload( + self, kwargs: Dict[str, Any] + ) -> PostHogEventPayload: """ Helper function to create a PostHog event payload for logging @@ -203,7 +207,9 @@ class PostHogLogger(CustomBatchLogger): # Core model information properties["$ai_model"] = self._safe_get(standard_logging_object, "model", "") - properties["$ai_provider"] = self._safe_get(standard_logging_object, "custom_llm_provider", "") + properties["$ai_provider"] = self._safe_get( + standard_logging_object, "custom_llm_provider", "" + ) # Input/Output data messages = self._safe_get(standard_logging_object, "messages") @@ -216,16 +222,22 @@ class PostHogLogger(CustomBatchLogger): properties["$ai_output_choices"] = response # Token information - properties["$ai_input_tokens"] = self._safe_get(standard_logging_object, "prompt_tokens", 0) + properties["$ai_input_tokens"] = self._safe_get( + standard_logging_object, "prompt_tokens", 0 + ) if event_name == "$ai_generation": - properties["$ai_output_tokens"] = self._safe_get(standard_logging_object, "completion_tokens", 0) + properties["$ai_output_tokens"] = self._safe_get( + standard_logging_object, "completion_tokens", 0 + ) # Cost and performance response_cost = self._safe_get(standard_logging_object, "response_cost") if response_cost is not None: properties["$ai_total_cost_usd"] = response_cost - properties["$ai_latency"] = self._safe_get(standard_logging_object, "response_time", 0.0) + properties["$ai_latency"] = self._safe_get( + standard_logging_object, "response_time", 0.0 + ) # Error handling if self._safe_get(standard_logging_object, "status") == "failure": @@ -245,7 +257,9 @@ class PostHogLogger(CustomBatchLogger): def _add_trace_properties(self, properties: Dict[str, Any], kwargs: Dict[str, Any]): standard_logging_object = self._safe_get(kwargs, "standard_logging_object", {}) - trace_id = self._safe_get(standard_logging_object, "trace_id", self._safe_uuid()) + trace_id = self._safe_get( + standard_logging_object, "trace_id", self._safe_uuid() + ) properties["$ai_trace_id"] = trace_id span_id = self._safe_get(standard_logging_object, "id", self._safe_uuid()) @@ -256,22 +270,48 @@ class PostHogLogger(CustomBatchLogger): if parent_id: properties["$ai_parent_id"] = parent_id - def _add_custom_metadata_properties(self, properties: Dict[str, Any], kwargs: Dict[str, Any]): + def _add_custom_metadata_properties( + self, properties: Dict[str, Any], kwargs: Dict[str, Any] + ): """Add custom metadata fields to PostHog properties""" metadata = self._extract_metadata(kwargs) if not isinstance(metadata, dict): return litellm_internal_fields = { - "endpoint", "caching_groups", "user_api_key_hash", "user_api_key_alias", - "user_api_key_team_id", "user_api_key_user_id", "user_api_key_org_id", - "user_api_key_team_alias", "user_api_key_end_user_id", "user_api_key_user_email", - "user_api_key", "user_api_end_user_max_budget", "litellm_api_version", - "global_max_parallel_requests", "user_api_key_team_max_budget", "user_api_key_team_spend", - "user_api_key_spend", "user_api_key_max_budget", "user_api_key_model_max_budget", - "user_api_key_metadata", "headers", "litellm_parent_otel_span", "requester_ip_address", - "model_group", "model_group_size", "deployment", "model_info", "api_base", - "caching_groups", "hidden_params", "parent_run_id", "parent_id", "user_id" + "endpoint", + "caching_groups", + "user_api_key_hash", + "user_api_key_alias", + "user_api_key_team_id", + "user_api_key_user_id", + "user_api_key_org_id", + "user_api_key_team_alias", + "user_api_key_end_user_id", + "user_api_key_user_email", + "user_api_key", + "user_api_end_user_max_budget", + "litellm_api_version", + "global_max_parallel_requests", + "user_api_key_team_max_budget", + "user_api_key_team_spend", + "user_api_key_spend", + "user_api_key_max_budget", + "user_api_key_model_max_budget", + "user_api_key_metadata", + "headers", + "litellm_parent_otel_span", + "requester_ip_address", + "model_group", + "model_group_size", + "deployment", + "model_info", + "api_base", + "caching_groups", + "hidden_params", + "parent_run_id", + "parent_id", + "user_id", } for key, value in metadata.items(): @@ -294,7 +334,9 @@ class PostHogLogger(CustomBatchLogger): return self._safe_uuid() - def _get_credentials_for_request(self, kwargs: Dict[str, Any]) -> Tuple[Optional[str], Optional[str]]: + def _get_credentials_for_request( + self, kwargs: Dict[str, Any] + ) -> Tuple[Optional[str], Optional[str]]: """ Get PostHog credentials for this request. @@ -307,13 +349,19 @@ class PostHogLogger(CustomBatchLogger): Returns: tuple[str, str]: (api_key, api_url) """ - standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = ( - kwargs.get("standard_callback_dynamic_params", None) - ) + standard_callback_dynamic_params: Optional[ + StandardCallbackDynamicParams + ] = kwargs.get("standard_callback_dynamic_params", None) if standard_callback_dynamic_params is not None: - api_key = standard_callback_dynamic_params.get("posthog_api_key") or self.POSTHOG_API_KEY - api_url = standard_callback_dynamic_params.get("posthog_api_url") or self.posthog_host + api_key = ( + standard_callback_dynamic_params.get("posthog_api_key") + or self.POSTHOG_API_KEY + ) + api_url = ( + standard_callback_dynamic_params.get("posthog_api_url") + or self.posthog_host + ) else: api_key = self.POSTHOG_API_KEY api_url = self.posthog_host @@ -334,9 +382,11 @@ class PostHogLogger(CustomBatchLogger): verbose_logger.debug( f"PostHog: Sending batch of {len(self.log_queue)} events" ) - + if self.is_mock_mode: - verbose_logger.debug("[POSTHOG MOCK] Mock mode enabled - API calls will be intercepted") + verbose_logger.debug( + "[POSTHOG MOCK] Mock mode enabled - API calls will be intercepted" + ) # Group events by credentials for batch sending batches_by_credentials: Dict[tuple[str, str], list] = {} @@ -368,7 +418,9 @@ class PostHogLogger(CustomBatchLogger): ) if self.is_mock_mode: - verbose_logger.debug(f"[POSTHOG MOCK] Batch of {len(self.log_queue)} events successfully mocked") + verbose_logger.debug( + f"[POSTHOG MOCK] Batch of {len(self.log_queue)} events successfully mocked" + ) else: verbose_logger.debug( f"PostHog: Batch of {len(self.log_queue)} events successfully sent" @@ -384,7 +436,9 @@ class PostHogLogger(CustomBatchLogger): self._async_initialized = True verbose_logger.debug("PostHog: Async components initialized") except Exception as e: - verbose_logger.error(f"PostHog: Failed to initialize async components: {str(e)}") + verbose_logger.error( + f"PostHog: Failed to initialize async components: {str(e)}" + ) raise def _extract_metadata(self, kwargs: Dict[str, Any]) -> Dict[str, Any]: @@ -398,7 +452,7 @@ class PostHogLogger(CustomBatchLogger): return {"api_key": api_key, "batch": events} def _safe_get(self, obj: Any, key: str, default: Any = None) -> Any: - if obj is None or not hasattr(obj, 'get'): + if obj is None or not hasattr(obj, "get"): return default return obj.get(key, default) diff --git a/litellm/integrations/posthog_mock_client.py b/litellm/integrations/posthog_mock_client.py index b713587ed6..de085b855c 100644 --- a/litellm/integrations/posthog_mock_client.py +++ b/litellm/integrations/posthog_mock_client.py @@ -8,7 +8,10 @@ Usage: Set POSTHOG_MOCK=true in environment variables or config to enable mock mode. """ -from litellm.integrations.mock_client_factory import MockClientConfig, create_mock_client_factory +from litellm.integrations.mock_client_factory import ( + MockClientConfig, + create_mock_client_factory, +) # Create mock client using factory _config = MockClientConfig( @@ -27,4 +30,6 @@ _config = MockClientConfig( patch_sync_client=True, ) -create_mock_posthog_client, should_use_posthog_mock = create_mock_client_factory(_config) +create_mock_posthog_client, should_use_posthog_mock = create_mock_client_factory( + _config +) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 7a08432b9a..357e0229fc 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -1417,7 +1417,9 @@ class PrometheusLogger(CustomLogger): _sanitize_prometheus_label_value(user_api_team), _sanitize_prometheus_label_value(user_api_team_alias), _sanitize_prometheus_label_value(user_id), - _sanitize_prometheus_label_value(standard_logging_payload.get("model_id", "")), + _sanitize_prometheus_label_value( + standard_logging_payload.get("model_id", "") + ), ).inc() self.set_llm_deployment_failure_metrics(kwargs) except Exception as e: diff --git a/litellm/integrations/prompt_management_base.py b/litellm/integrations/prompt_management_base.py index b32f78c0de..71da650dc4 100644 --- a/litellm/integrations/prompt_management_base.py +++ b/litellm/integrations/prompt_management_base.py @@ -75,7 +75,6 @@ class PromptManagementBase(ABC): prompt_version: Optional[int] = None, prompt_spec: Optional[PromptSpec] = None, ) -> PromptManagementClient: - compiled_prompt_client = self._compile_prompt_helper( prompt_id=prompt_id, prompt_spec=prompt_spec, @@ -179,7 +178,6 @@ class PromptManagementBase(ABC): ignore_prompt_manager_model: Optional[bool] = False, ignore_prompt_manager_optional_params: Optional[bool] = False, ) -> Tuple[str, List[AllMessageValues], dict]: - if prompt_id is None: raise ValueError("prompt_id is required for Prompt Management Base class") if not self.should_run_prompt_management( diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index eddc80dbc1..c8db4be7ce 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -80,7 +80,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): s3_use_team_prefix=s3_use_team_prefix, s3_strip_base64_files=s3_strip_base64_files, s3_use_key_prefix=s3_use_key_prefix, - s3_use_virtual_hosted_style=s3_use_virtual_hosted_style + s3_use_virtual_hosted_style=s3_use_virtual_hosted_style, ) verbose_logger.debug(f"s3 logger using endpoint url {s3_endpoint_url}") @@ -91,7 +91,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): ) self.async_httpx_client = get_async_httpx_client( llm_provider=httpxSpecialProvider.LoggingCallback, - params={"ssl_verify": self.s3_verify} + params={"ssl_verify": self.s3_verify}, ) asyncio.create_task(self.periodic_flush()) @@ -158,10 +158,14 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): litellm.s3_callback_params.get("s3_api_version") or s3_api_version ) self.s3_use_ssl = ( - litellm.s3_callback_params.get("s3_use_ssl", True) if litellm.s3_callback_params.get("s3_use_ssl") is not None else s3_use_ssl + litellm.s3_callback_params.get("s3_use_ssl", True) + if litellm.s3_callback_params.get("s3_use_ssl") is not None + else s3_use_ssl ) self.s3_verify = ( - litellm.s3_callback_params.get("s3_verify") if litellm.s3_callback_params.get("s3_verify") is not None else s3_verify + litellm.s3_callback_params.get("s3_verify") + if litellm.s3_callback_params.get("s3_verify") is not None + else s3_verify ) self.s3_endpoint_url = ( litellm.s3_callback_params.get("s3_endpoint_url") or s3_endpoint_url @@ -211,8 +215,8 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): ) self.s3_use_key_prefix = ( - bool(litellm.s3_callback_params.get("s3_use_key_prefix", False)) - or s3_use_key_prefix + bool(litellm.s3_callback_params.get("s3_use_key_prefix", False)) + or s3_use_key_prefix ) self.s3_strip_base64_files = ( @@ -308,9 +312,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): verbose_logger.debug( f"s3_v2 logger - uploading data to s3 - {batch_logging_element.s3_object_key}" ) - verbose_logger.debug( - f"s3_v2 logger - s3_verify setting: {self.s3_verify}" - ) + verbose_logger.debug(f"s3_v2 logger - s3_verify setting: {self.s3_verify}") # Prepare the URL url = f"https://{self.s3_bucket_name}.s3.{self.s3_region_name}.amazonaws.com/{batch_logging_element.s3_object_key}" @@ -318,8 +320,14 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): if self.s3_endpoint_url and self.s3_bucket_name: if self.s3_use_virtual_hosted_style: # Virtual-hosted-style: bucket.endpoint/key - endpoint_host = self.s3_endpoint_url.replace("https://", "").replace("http://", "") - protocol = "https://" if self.s3_endpoint_url.startswith("https://") else "http://" + endpoint_host = self.s3_endpoint_url.replace( + "https://", "" + ).replace("http://", "") + protocol = ( + "https://" + if self.s3_endpoint_url.startswith("https://") + else "http://" + ) url = f"{protocol}{self.s3_bucket_name}.{endpoint_host}/{batch_logging_element.s3_object_key}" else: # Path-style: endpoint/bucket/key @@ -413,20 +421,25 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): return None if self.s3_strip_base64_files: - standard_logging_payload = self._strip_base64_from_messages_sync(standard_logging_payload) + standard_logging_payload = self._strip_base64_from_messages_sync( + standard_logging_payload + ) # Base prefix (default empty) prefix_components = [] if self.s3_use_team_prefix: - team_alias = standard_logging_payload.get("metadata", {}).get("user_api_key_team_alias", None) + team_alias = standard_logging_payload.get("metadata", {}).get( + "user_api_key_team_alias", None + ) if team_alias: prefix_components.append(team_alias) if self.s3_use_key_prefix: - user_api_key_alias = standard_logging_payload.get("metadata", {}).get("user_api_key_alias", None) + user_api_key_alias = standard_logging_payload.get("metadata", {}).get( + "user_api_key_alias", None + ) if user_api_key_alias: prefix_components.append(user_api_key_alias) - # Construct full prefix path prefix_path = "/".join(prefix_components) if prefix_path: @@ -435,7 +448,9 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): s3_file_name = ( litellm.utils.get_logging_id(start_time, standard_logging_payload) or "" ) - verbose_logger.debug(f"Creating s3 file with prefix_components={prefix_components},prefix_path={prefix_path} and {s3_file_name}") + verbose_logger.debug( + f"Creating s3 file with prefix_components={prefix_components},prefix_path={prefix_path} and {s3_file_name}" + ) s3_object_key = get_s3_object_key( s3_path=cast(Optional[str], self.s3_path) or "", prefix=prefix_path, @@ -479,8 +494,14 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): if self.s3_endpoint_url and self.s3_bucket_name: if self.s3_use_virtual_hosted_style: # Virtual-hosted-style: bucket.endpoint/key - endpoint_host = self.s3_endpoint_url.replace("https://", "").replace("http://", "") - protocol = "https://" if self.s3_endpoint_url.startswith("https://") else "http://" + endpoint_host = self.s3_endpoint_url.replace( + "https://", "" + ).replace("http://", "") + protocol = ( + "https://" + if self.s3_endpoint_url.startswith("https://") + else "http://" + ) url = f"{protocol}{self.s3_bucket_name}.{endpoint_host}/{batch_logging_element.s3_object_key}" else: # Path-style: endpoint/bucket/key @@ -525,7 +546,9 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): signed_headers = dict(aws_request.headers.items()) httpx_client = _get_httpx_client( - params={"ssl_verify": self.s3_verify} if self.s3_verify is not None else None + params={"ssl_verify": self.s3_verify} + if self.s3_verify is not None + else None ) # Make the request response = httpx_client.put(url, data=json_string, headers=signed_headers) @@ -580,8 +603,14 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): if self.s3_endpoint_url and self.s3_bucket_name: if self.s3_use_virtual_hosted_style: # Virtual-hosted-style: bucket.endpoint/key - endpoint_host = self.s3_endpoint_url.replace("https://", "").replace("http://", "") - protocol = "https://" if self.s3_endpoint_url.startswith("https://") else "http://" + endpoint_host = self.s3_endpoint_url.replace( + "https://", "" + ).replace("http://", "") + protocol = ( + "https://" + if self.s3_endpoint_url.startswith("https://") + else "http://" + ) url = f"{protocol}{self.s3_bucket_name}.{endpoint_host}/{s3_object_key}" else: # Path-style: endpoint/bucket/key @@ -653,4 +682,4 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): verbose_logger.exception( f"Error retrieving object {object_key} from cold storage: {str(e)}" ) - return None \ No newline at end of file + return None diff --git a/litellm/integrations/sqs.py b/litellm/integrations/sqs.py index 97a4c5723d..6cbd2c7974 100644 --- a/litellm/integrations/sqs.py +++ b/litellm/integrations/sqs.py @@ -42,31 +42,31 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): """Batching logger that writes logs to an AWS SQS queue, optionally encrypting the payload.""" def __init__( - self, - # --- Standard SQS params --- - sqs_queue_url: Optional[str] = None, - sqs_region_name: Optional[str] = None, - sqs_api_version: Optional[str] = None, - sqs_use_ssl: bool = True, - sqs_verify: Optional[bool] = None, - sqs_endpoint_url: Optional[str] = None, - sqs_aws_access_key_id: Optional[str] = None, - sqs_aws_secret_access_key: Optional[str] = None, - sqs_aws_session_token: Optional[str] = None, - sqs_aws_session_name: Optional[str] = None, - sqs_aws_profile_name: Optional[str] = None, - sqs_aws_role_name: Optional[str] = None, - sqs_aws_web_identity_token: Optional[str] = None, - sqs_aws_sts_endpoint: Optional[str] = None, - sqs_flush_interval: Optional[int] = DEFAULT_SQS_FLUSH_INTERVAL_SECONDS, - sqs_batch_size: Optional[int] = DEFAULT_SQS_BATCH_SIZE, - sqs_config=None, - sqs_strip_base64_files: bool = False, - # --- 🔐 Application-level encryption params --- - sqs_aws_use_application_level_encryption: bool = False, - sqs_app_encryption_key_b64: Optional[str] = None, - sqs_app_encryption_aad: Optional[str] = None, - **kwargs, + self, + # --- Standard SQS params --- + sqs_queue_url: Optional[str] = None, + sqs_region_name: Optional[str] = None, + sqs_api_version: Optional[str] = None, + sqs_use_ssl: bool = True, + sqs_verify: Optional[bool] = None, + sqs_endpoint_url: Optional[str] = None, + sqs_aws_access_key_id: Optional[str] = None, + sqs_aws_secret_access_key: Optional[str] = None, + sqs_aws_session_token: Optional[str] = None, + sqs_aws_session_name: Optional[str] = None, + sqs_aws_profile_name: Optional[str] = None, + sqs_aws_role_name: Optional[str] = None, + sqs_aws_web_identity_token: Optional[str] = None, + sqs_aws_sts_endpoint: Optional[str] = None, + sqs_flush_interval: Optional[int] = DEFAULT_SQS_FLUSH_INTERVAL_SECONDS, + sqs_batch_size: Optional[int] = DEFAULT_SQS_BATCH_SIZE, + sqs_config=None, + sqs_strip_base64_files: bool = False, + # --- 🔐 Application-level encryption params --- + sqs_aws_use_application_level_encryption: bool = False, + sqs_app_encryption_key_b64: Optional[str] = None, + sqs_app_encryption_aad: Optional[str] = None, + **kwargs, ) -> None: try: verbose_logger.debug( @@ -122,26 +122,26 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): raise e def _init_sqs_params( - self, - sqs_queue_url: Optional[str] = None, - sqs_region_name: Optional[str] = None, - sqs_api_version: Optional[str] = None, - sqs_use_ssl: bool = True, - sqs_verify: Optional[bool] = None, - sqs_endpoint_url: Optional[str] = None, - sqs_aws_access_key_id: Optional[str] = None, - sqs_aws_secret_access_key: Optional[str] = None, - sqs_aws_session_token: Optional[str] = None, - sqs_aws_session_name: Optional[str] = None, - sqs_aws_profile_name: Optional[str] = None, - sqs_aws_role_name: Optional[str] = None, - sqs_aws_web_identity_token: Optional[str] = None, - sqs_aws_sts_endpoint: Optional[str] = None, - sqs_strip_base64_files: bool = False, - sqs_aws_use_application_level_encryption: bool = False, - sqs_app_encryption_key_b64: Optional[str] = None, - sqs_app_encryption_aad: Optional[str] = None, - sqs_config=None, + self, + sqs_queue_url: Optional[str] = None, + sqs_region_name: Optional[str] = None, + sqs_api_version: Optional[str] = None, + sqs_use_ssl: bool = True, + sqs_verify: Optional[bool] = None, + sqs_endpoint_url: Optional[str] = None, + sqs_aws_access_key_id: Optional[str] = None, + sqs_aws_secret_access_key: Optional[str] = None, + sqs_aws_session_token: Optional[str] = None, + sqs_aws_session_name: Optional[str] = None, + sqs_aws_profile_name: Optional[str] = None, + sqs_aws_role_name: Optional[str] = None, + sqs_aws_web_identity_token: Optional[str] = None, + sqs_aws_sts_endpoint: Optional[str] = None, + sqs_strip_base64_files: bool = False, + sqs_aws_use_application_level_encryption: bool = False, + sqs_app_encryption_key_b64: Optional[str] = None, + sqs_app_encryption_aad: Optional[str] = None, + sqs_config=None, ) -> None: litellm.aws_sqs_callback_params = litellm.aws_sqs_callback_params or {} @@ -151,87 +151,98 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): litellm.aws_sqs_callback_params[key] = litellm.get_secret(value) self.sqs_queue_url = ( - litellm.aws_sqs_callback_params.get("sqs_queue_url") or sqs_queue_url + litellm.aws_sqs_callback_params.get("sqs_queue_url") or sqs_queue_url ) self.sqs_region_name = ( - litellm.aws_sqs_callback_params.get("sqs_region_name") or sqs_region_name + litellm.aws_sqs_callback_params.get("sqs_region_name") or sqs_region_name ) self.sqs_api_version = ( - litellm.aws_sqs_callback_params.get("sqs_api_version") or sqs_api_version + litellm.aws_sqs_callback_params.get("sqs_api_version") or sqs_api_version ) self.sqs_use_ssl = ( - litellm.aws_sqs_callback_params.get("sqs_use_ssl", True) or sqs_use_ssl + litellm.aws_sqs_callback_params.get("sqs_use_ssl", True) or sqs_use_ssl + ) + self.sqs_verify = ( + litellm.aws_sqs_callback_params.get("sqs_verify") or sqs_verify ) - self.sqs_verify = litellm.aws_sqs_callback_params.get("sqs_verify") or sqs_verify self.sqs_endpoint_url = ( - litellm.aws_sqs_callback_params.get("sqs_endpoint_url") or sqs_endpoint_url + litellm.aws_sqs_callback_params.get("sqs_endpoint_url") or sqs_endpoint_url ) self.sqs_aws_access_key_id = ( - litellm.aws_sqs_callback_params.get("sqs_aws_access_key_id") - or sqs_aws_access_key_id + litellm.aws_sqs_callback_params.get("sqs_aws_access_key_id") + or sqs_aws_access_key_id ) self.sqs_aws_secret_access_key = ( - litellm.aws_sqs_callback_params.get("sqs_aws_secret_access_key") - or sqs_aws_secret_access_key + litellm.aws_sqs_callback_params.get("sqs_aws_secret_access_key") + or sqs_aws_secret_access_key ) self.sqs_aws_session_token = ( - litellm.aws_sqs_callback_params.get("sqs_aws_session_token") - or sqs_aws_session_token + litellm.aws_sqs_callback_params.get("sqs_aws_session_token") + or sqs_aws_session_token ) self.sqs_aws_session_name = ( - litellm.aws_sqs_callback_params.get("sqs_aws_session_name") or sqs_aws_session_name + litellm.aws_sqs_callback_params.get("sqs_aws_session_name") + or sqs_aws_session_name ) self.sqs_aws_profile_name = ( - litellm.aws_sqs_callback_params.get("sqs_aws_profile_name") or sqs_aws_profile_name + litellm.aws_sqs_callback_params.get("sqs_aws_profile_name") + or sqs_aws_profile_name ) self.sqs_aws_role_name = ( - litellm.aws_sqs_callback_params.get("sqs_aws_role_name") or sqs_aws_role_name + litellm.aws_sqs_callback_params.get("sqs_aws_role_name") + or sqs_aws_role_name ) self.sqs_aws_web_identity_token = ( - litellm.aws_sqs_callback_params.get("sqs_aws_web_identity_token") - or sqs_aws_web_identity_token + litellm.aws_sqs_callback_params.get("sqs_aws_web_identity_token") + or sqs_aws_web_identity_token ) self.sqs_aws_sts_endpoint = ( - litellm.aws_sqs_callback_params.get("sqs_aws_sts_endpoint") or sqs_aws_sts_endpoint + litellm.aws_sqs_callback_params.get("sqs_aws_sts_endpoint") + or sqs_aws_sts_endpoint ) self.sqs_strip_base64_files = ( - litellm.aws_sqs_callback_params.get("sqs_strip_base64_files", False) - or sqs_strip_base64_files + litellm.aws_sqs_callback_params.get("sqs_strip_base64_files", False) + or sqs_strip_base64_files ) self.sqs_aws_use_application_level_encryption = ( - litellm.aws_sqs_callback_params.get("sqs_aws_use_application_level_encryption", False) - or sqs_aws_use_application_level_encryption + litellm.aws_sqs_callback_params.get( + "sqs_aws_use_application_level_encryption", False + ) + or sqs_aws_use_application_level_encryption ) self.sqs_app_encryption_key_b64 = ( - litellm.aws_sqs_callback_params.get("sqs_app_encryption_key_b64") - or sqs_app_encryption_key_b64 + litellm.aws_sqs_callback_params.get("sqs_app_encryption_key_b64") + or sqs_app_encryption_key_b64 ) self.sqs_app_encryption_aad = ( - litellm.aws_sqs_callback_params.get("sqs_app_encryption_aad") - or sqs_app_encryption_aad + litellm.aws_sqs_callback_params.get("sqs_app_encryption_aad") + or sqs_app_encryption_aad ) self.app_crypto: Optional["AppCrypto"] = None if self.sqs_aws_use_application_level_encryption: from litellm.litellm_core_utils.app_crypto import AppCrypto + if not self.sqs_app_encryption_key_b64: - raise ValueError("sqs_app_encryption_key_b64 is required when encryption is enabled.") + raise ValueError( + "sqs_app_encryption_key_b64 is required when encryption is enabled." + ) key = base64.b64decode(self.sqs_app_encryption_key_b64) self.app_crypto = AppCrypto(key) - verbose_logger.debug( - "SQSLogger: Application-level encryption enabled." - ) - self.sqs_config = litellm.aws_sqs_callback_params.get("sqs_config") or sqs_config + verbose_logger.debug("SQSLogger: Application-level encryption enabled.") + self.sqs_config = ( + litellm.aws_sqs_callback_params.get("sqs_config") or sqs_config + ) async def async_log_success_event( - self, kwargs, response_obj, start_time, end_time + self, kwargs, response_obj, start_time, end_time ) -> None: try: verbose_logger.debug( @@ -239,7 +250,9 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): ) standard_logging_payload = kwargs.get("standard_logging_object") if self.sqs_strip_base64_files: - standard_logging_payload = await self._strip_base64_from_messages(standard_logging_payload) + standard_logging_payload = await self._strip_base64_from_messages( + standard_logging_payload + ) if standard_logging_payload is None: raise ValueError("standard_logging_payload is None") @@ -258,7 +271,9 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): if standard_logging_payload is None: raise ValueError("standard_logging_payload is None") if self.sqs_strip_base64_files: - standard_logging_payload = await self._strip_base64_from_messages(standard_logging_payload) + standard_logging_payload = await self._strip_base64_from_messages( + standard_logging_payload + ) self.log_queue.append(standard_logging_payload) verbose_logger.debug( @@ -274,9 +289,7 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): pass async def async_send_batch(self) -> None: - verbose_logger.debug( - f"sqs logger - sending batch of {len(self.log_queue)}" - ) + verbose_logger.debug(f"sqs logger - sending batch of {len(self.log_queue)}") if not self.log_queue: return @@ -322,8 +335,8 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): json_string = safe_dumps(payload) body = ( - f"Action={SQS_SEND_MESSAGE_ACTION}&Version={SQS_API_VERSION}&MessageBody=" - + quote(json_string, safe="") + f"Action={SQS_SEND_MESSAGE_ACTION}&Version={SQS_API_VERSION}&MessageBody=" + + quote(json_string, safe="") ) headers = { @@ -341,9 +354,7 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): data=prepped.body, headers=prepped.headers, ) - SigV4Auth(credentials, "sqs", self.sqs_region_name).add_auth( - aws_request - ) + SigV4Auth(credentials, "sqs", self.sqs_region_name).add_auth(aws_request) signed_headers = dict(aws_request.headers.items()) @@ -364,10 +375,15 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): from litellm.litellm_core_utils.litellm_logging import ( create_dummy_standard_logging_payload, ) + # Create a minimal standard logging payload - standard_logging_object: StandardLoggingPayload = create_dummy_standard_logging_payload() + standard_logging_object: StandardLoggingPayload = ( + create_dummy_standard_logging_payload() + ) # Attempt to send a single message await self.async_send_message(standard_logging_object) return IntegrationHealthCheckStatus(status="healthy", error_message=None) except Exception as e: - return IntegrationHealthCheckStatus(status="unhealthy", error_message=str(e)) + return IntegrationHealthCheckStatus( + status="unhealthy", error_message=str(e) + ) diff --git a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py index c94b925ea2..50420fb713 100644 --- a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py +++ b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py @@ -82,17 +82,18 @@ class VectorStorePreCallHook(CustomLogger): prisma_client = None try: from litellm.proxy.proxy_server import prisma_client as _prisma_client + prisma_client = _prisma_client except ImportError: pass # Use database fallback to ensure synchronization across instances - vector_stores_to_run: List[LiteLLM_ManagedVectorStore] = ( - await litellm.vector_store_registry.pop_vector_stores_to_run_with_db_fallback( - non_default_params=non_default_params, - tools=tools, - prisma_client=prisma_client - ) + vector_stores_to_run: List[ + LiteLLM_ManagedVectorStore + ] = await litellm.vector_store_registry.pop_vector_stores_to_run_with_db_fallback( + non_default_params=non_default_params, + tools=tools, + prisma_client=prisma_client, ) if not vector_stores_to_run: @@ -111,7 +112,6 @@ class VectorStorePreCallHook(CustomLogger): all_search_results: List[VectorStoreSearchResponse] = [] for vector_store_to_run in vector_stores_to_run: - # Get vector store id from the vector store config vector_store_id = vector_store_to_run.get("vector_store_id", "") custom_llm_provider = vector_store_to_run.get("custom_llm_provider") @@ -147,9 +147,9 @@ class VectorStorePreCallHook(CustomLogger): # Store search results as-is (already in OpenAI-compatible format) if litellm_logging_obj and all_search_results: - litellm_logging_obj.model_call_details["search_results"] = ( - all_search_results - ) + litellm_logging_obj.model_call_details[ + "search_results" + ] = all_search_results return model, modified_messages, non_default_params @@ -208,9 +208,9 @@ class VectorStorePreCallHook(CustomLogger): Returns: Modified list of messages with context appended """ - search_response_data: Optional[List[VectorStoreSearchResult]] = ( - search_response.get("data") - ) + search_response_data: Optional[ + List[VectorStoreSearchResult] + ] = search_response.get("data") if not search_response_data: return messages @@ -268,9 +268,9 @@ class VectorStorePreCallHook(CustomLogger): ) # Get search results from model_call_details (already in OpenAI format) - search_results: Optional[List[VectorStoreSearchResponse]] = ( - litellm_logging_obj.model_call_details.get("search_results") - ) + search_results: Optional[ + List[VectorStoreSearchResponse] + ] = litellm_logging_obj.model_call_details.get("search_results") verbose_logger.debug(f"Search results found: {search_results is not None}") @@ -328,9 +328,9 @@ class VectorStorePreCallHook(CustomLogger): ) # Get search results from model_call_details (already in OpenAI format) - search_results: Optional[List[VectorStoreSearchResponse]] = ( - request_data.get("search_results") - ) + search_results: Optional[ + List[VectorStoreSearchResponse] + ] = request_data.get("search_results") verbose_logger.debug( f"Search results found for streaming chunk: {search_results is not None}" diff --git a/litellm/integrations/weave/weave_otel.py b/litellm/integrations/weave/weave_otel.py index 167deaf2cd..796a33a34d 100644 --- a/litellm/integrations/weave/weave_otel.py +++ b/litellm/integrations/weave/weave_otel.py @@ -9,7 +9,9 @@ from opentelemetry.trace import Status, StatusCode from typing_extensions import override from litellm._logging import verbose_logger -from litellm.integrations._types.open_inference import SpanAttributes as OpenInferenceSpanAttributes +from litellm.integrations._types.open_inference import ( + SpanAttributes as OpenInferenceSpanAttributes, +) from litellm.integrations.arize import _utils from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes import ( @@ -54,10 +56,14 @@ class WeaveLLMObsOTELAttributes(BaseLLMObsOTELAttributes): prompt["functions"] = functions if tools is not None: prompt["tools"] = tools - safe_set_attribute(span, OpenInferenceSpanAttributes.INPUT_VALUE, json.dumps(prompt)) + safe_set_attribute( + span, OpenInferenceSpanAttributes.INPUT_VALUE, json.dumps(prompt) + ) -def _set_weave_specific_attributes(span: Span, kwargs: dict[str, Any], response_obj: Any): +def _set_weave_specific_attributes( + span: Span, kwargs: dict[str, Any], response_obj: Any +): """ Sets Weave-specific metadata attributes onto the OTEL span. @@ -100,7 +106,9 @@ def _set_weave_specific_attributes(span: Span, kwargs: dict[str, Any], response_ output_dict = response_obj if output_dict: - safe_set_attribute(span, OpenInferenceSpanAttributes.OUTPUT_VALUE, safe_dumps(output_dict)) + safe_set_attribute( + span, OpenInferenceSpanAttributes.OUTPUT_VALUE, safe_dumps(output_dict) + ) def _get_weave_authorization_header(api_key: str) -> str: @@ -134,7 +142,9 @@ def get_weave_otel_config() -> WeaveOtelConfig: host = os.getenv("WANDB_HOST") if not api_key: - raise ValueError("WANDB_API_KEY must be set for Weave OpenTelemetry integration.") + raise ValueError( + "WANDB_API_KEY must be set for Weave OpenTelemetry integration." + ) if not project_id: raise ValueError( @@ -223,7 +233,9 @@ class WeaveOtelLogger(OpenTelemetry): super().__init__(config=config, callback_name=callback_name, **kwargs) - def _maybe_log_raw_request(self, kwargs, response_obj, start_time, end_time, parent_span): + def _maybe_log_raw_request( + self, kwargs, response_obj, start_time, end_time, parent_span + ): """ Override to skip creating the raw_gen_ai_request child span. @@ -281,7 +293,9 @@ class WeaveOtelLogger(OpenTelemetry): primary_span_parent = None # 1. Primary span - span = self._start_primary_span(kwargs, response_obj, start_time, end_time, ctx, primary_span_parent) + span = self._start_primary_span( + kwargs, response_obj, start_time, end_time, ctx, primary_span_parent + ) # 2. Raw-request sub-span (skipped for Weave via _maybe_log_raw_request override) self._maybe_log_raw_request(kwargs, response_obj, start_time, end_time, span) @@ -315,7 +329,9 @@ class WeaveOtelLogger(OpenTelemetry): dynamic_headers = {} dynamic_wandb_api_key = standard_callback_dynamic_params.get("wandb_api_key") - dynamic_weave_project_id = standard_callback_dynamic_params.get("weave_project_id") + dynamic_weave_project_id = standard_callback_dynamic_params.get( + "weave_project_id" + ) if dynamic_wandb_api_key: auth_header = _get_weave_authorization_header( diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index c31140d44d..2541a0bd7a 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -62,8 +62,7 @@ class WebSearchInterceptionLogger(CustomLogger): self.enabled_providers = [LlmProviders.BEDROCK.value] else: self.enabled_providers = [ - p.value if isinstance(p, LlmProviders) else p - for p in enabled_providers + p.value if isinstance(p, LlmProviders) else p for p in enabled_providers ] self.search_tool_name = search_tool_name self._request_has_websearch = False # Track if current request has web search @@ -80,10 +79,14 @@ class WebSearchInterceptionLogger(CustomLogger): """ # Check if this is for an enabled provider # Try top-level kwargs first, then nested litellm_params, then derive from model name - custom_llm_provider = kwargs.get("custom_llm_provider", "") or kwargs.get("litellm_params", {}).get("custom_llm_provider", "") + custom_llm_provider = kwargs.get("custom_llm_provider", "") or kwargs.get( + "litellm_params", {} + ).get("custom_llm_provider", "") if not custom_llm_provider: try: - _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=kwargs.get("model", "")) + _, custom_llm_provider, _, _ = litellm.get_llm_provider( + model=kwargs.get("model", "") + ) except Exception: custom_llm_provider = "" if custom_llm_provider not in self.enabled_providers: @@ -197,7 +200,10 @@ class WebSearchInterceptionLogger(CustomLogger): f" - enabled_providers={self.enabled_providers or 'ALL'}" ) - if self.enabled_providers is not None and custom_llm_provider not in self.enabled_providers: + if ( + self.enabled_providers is not None + and custom_llm_provider not in self.enabled_providers + ): verbose_logger.debug( f"WebSearchInterception: Skipping - provider {custom_llm_provider} not in {self.enabled_providers}" ) @@ -258,18 +264,23 @@ class WebSearchInterceptionLogger(CustomLogger): ) -> Tuple[bool, Dict]: """ Check if WebSearch tool interception is needed for Anthropic Messages API. - + This is the legacy method for Anthropic-style responses. For chat completions, use async_should_run_chat_completion_agentic_loop instead. """ - verbose_logger.debug(f"WebSearchInterception: Hook called! provider={custom_llm_provider}, stream={stream}") + verbose_logger.debug( + f"WebSearchInterception: Hook called! provider={custom_llm_provider}, stream={stream}" + ) verbose_logger.debug(f"WebSearchInterception: Response type: {type(response)}") # Check if provider should be intercepted # Note: custom_llm_provider is already normalized by get_llm_provider() # (e.g., "bedrock/invoke/..." -> "bedrock") - if self.enabled_providers is not None and custom_llm_provider not in self.enabled_providers: + if ( + self.enabled_providers is not None + and custom_llm_provider not in self.enabled_providers + ): verbose_logger.debug( f"WebSearchInterception: Skipping provider {custom_llm_provider} (not in enabled list: {self.enabled_providers})" ) @@ -278,9 +289,7 @@ class WebSearchInterceptionLogger(CustomLogger): # Check if tools include any web search tool (LiteLLM standard or native) has_websearch_tool = any(is_web_search_tool(t) for t in (tools or [])) if not has_websearch_tool: - verbose_logger.debug( - "WebSearchInterception: No web search tool in request" - ) + verbose_logger.debug("WebSearchInterception: No web search tool in request") return False, {} # Detect WebSearch tool_use in response (Anthropic format) @@ -324,16 +333,12 @@ class WebSearchInterceptionLogger(CustomLogger): # pattern in _detect_from_non_streaming_response thinking_block_dict: Dict = {"type": block_type} if block_type == "thinking": - thinking_block_dict["thinking"] = getattr( - block, "thinking", "" - ) + thinking_block_dict["thinking"] = getattr(block, "thinking", "") thinking_block_dict["signature"] = getattr( block, "signature", "" ) else: # redacted_thinking - thinking_block_dict["data"] = getattr( - block, "data", "" - ) + thinking_block_dict["data"] = getattr(block, "data", "") thinking_blocks.append(thinking_block_dict) if thinking_blocks: @@ -363,22 +368,29 @@ class WebSearchInterceptionLogger(CustomLogger): ) -> Tuple[bool, Dict]: """ Check if WebSearch tool interception is needed for Chat Completions API. - + Similar to async_should_run_agentic_loop but for OpenAI-style chat completions. """ - verbose_logger.debug(f"WebSearchInterception: Chat completion hook called! provider={custom_llm_provider}, stream={stream}") + verbose_logger.debug( + f"WebSearchInterception: Chat completion hook called! provider={custom_llm_provider}, stream={stream}" + ) verbose_logger.debug(f"WebSearchInterception: Response type: {type(response)}") # Check if provider should be intercepted - if self.enabled_providers is not None and custom_llm_provider not in self.enabled_providers: + if ( + self.enabled_providers is not None + and custom_llm_provider not in self.enabled_providers + ): verbose_logger.debug( f"WebSearchInterception: Skipping provider {custom_llm_provider} (not in enabled list: {self.enabled_providers})" ) return False, {} # Check if tools include any web search tool (strict check for chat completions) - has_websearch_tool = any(is_web_search_tool_chat_completion(t) for t in (tools or [])) + has_websearch_tool = any( + is_web_search_tool_chat_completion(t) for t in (tools or []) + ) if not has_websearch_tool: verbose_logger.debug( "WebSearchInterception: No litellm_web_search tool in request" @@ -425,7 +437,7 @@ class WebSearchInterceptionLogger(CustomLogger): ) -> Any: """ Execute agentic loop with WebSearch execution for Anthropic Messages API. - + This is the legacy method for Anthropic-style responses. """ @@ -460,7 +472,7 @@ class WebSearchInterceptionLogger(CustomLogger): ) -> Any: """ Execute agentic loop with WebSearch execution for Chat Completions API. - + Similar to async_run_agentic_loop but for OpenAI-style chat completions. """ @@ -510,7 +522,9 @@ class WebSearchInterceptionLogger(CustomLogger): verbose_logger.debug( "WebSearchInterception: max_tokens=%s <= thinking.budget_tokens=%s, " "adjusting to %s to satisfy Anthropic API constraint", - max_tokens, budget_tokens, adjusted, + max_tokens, + budget_tokens, + adjusted, ) max_tokens = adjusted return max_tokens @@ -526,10 +540,11 @@ class WebSearchInterceptionLogger(CustomLogger): call's spend from being recorded — the root cause of the SpendLog / AWS billing mismatch. """ - _internal_keys = {'litellm_logging_obj'} + _internal_keys = {"litellm_logging_obj"} return { - k: v for k, v in kwargs.items() - if not k.startswith('_websearch_interception') and k not in _internal_keys + k: v + for k, v in kwargs.items() + if not k.startswith("_websearch_interception") and k not in _internal_keys } async def _execute_agentic_loop( @@ -574,9 +589,7 @@ class WebSearchInterceptionLogger(CustomLogger): verbose_logger.error( f"WebSearchInterception: Search {i} failed with error: {str(result)}" ) - final_search_results.append( - f"Search failed: {str(result)}" - ) + final_search_results.append(f"Search failed: {str(result)}") elif isinstance(result, str): # Explicitly cast to str for type checker final_search_results.append(cast(str, result)) @@ -609,9 +622,8 @@ class WebSearchInterceptionLogger(CustomLogger): ) # Correlation context for structured logging - _call_id = ( - getattr(logging_obj, "litellm_call_id", None) - or kwargs.get("litellm_call_id", "unknown") + _call_id = getattr(logging_obj, "litellm_call_id", None) or kwargs.get( + "litellm_call_id", "unknown" ) full_model_name = model # safe default before try block @@ -628,8 +640,9 @@ class WebSearchInterceptionLogger(CustomLogger): # Create a copy of optional params without max_tokens (since we pass it explicitly) optional_params_without_max_tokens = { - k: v for k, v in anthropic_messages_optional_request_params.items() - if k != 'max_tokens' + k: v + for k, v in anthropic_messages_optional_request_params.items() + if k != "max_tokens" } kwargs_for_followup = self._prepare_followup_kwargs(kwargs) @@ -637,12 +650,14 @@ class WebSearchInterceptionLogger(CustomLogger): # Get model from logging_obj.model_call_details["agentic_loop_params"] # This preserves the full model name with provider prefix (e.g., "bedrock/invoke/...") if logging_obj is not None: - agentic_params = logging_obj.model_call_details.get("agentic_loop_params", {}) + agentic_params = logging_obj.model_call_details.get( + "agentic_loop_params", {} + ) full_model_name = agentic_params.get("model", model) verbose_logger.debug( f"WebSearchInterception: Using model name: {full_model_name}" ) - + final_response = await anthropic_messages.acreate( max_tokens=max_tokens, messages=follow_up_messages, @@ -661,8 +676,11 @@ class WebSearchInterceptionLogger(CustomLogger): verbose_logger.exception( "WebSearchInterception: Follow-up request failed " "[call_id=%s model=%s messages=%d searches=%d]: %s", - _call_id, full_model_name, len(follow_up_messages), - len(final_search_results), str(e), + _call_id, + full_model_name, + len(follow_up_messages), + len(final_search_results), + str(e), ) raise @@ -685,12 +703,15 @@ class WebSearchInterceptionLogger(CustomLogger): if self.search_tool_name: # Find specific search tool by name matching_tools = [ - tool for tool in llm_router.search_tools + tool + for tool in llm_router.search_tools if tool.get("search_tool_name") == self.search_tool_name ] if matching_tools: search_tool = matching_tools[0] - search_provider = search_tool.get("litellm_params", {}).get("search_provider") + search_provider = search_tool.get("litellm_params", {}).get( + "search_provider" + ) verbose_logger.debug( f"WebSearchInterception: Found search tool '{self.search_tool_name}' " f"with provider '{search_provider}'" @@ -704,7 +725,9 @@ class WebSearchInterceptionLogger(CustomLogger): # If no specific tool or not found, use first available if not search_provider and llm_router.search_tools: first_tool = llm_router.search_tools[0] - search_provider = first_tool.get("litellm_params", {}).get("search_provider") + search_provider = first_tool.get("litellm_params", {}).get( + "search_provider" + ) verbose_logger.debug( f"WebSearchInterception: Using first available search tool with provider '{search_provider}'" ) @@ -720,9 +743,7 @@ class WebSearchInterceptionLogger(CustomLogger): verbose_logger.debug( f"WebSearchInterception: Executing search for '{query}' using provider '{search_provider}'" ) - result = await litellm.asearch( - query=query, search_provider=search_provider - ) + result = await litellm.asearch(query=query, search_provider=search_provider) # Format using transformation function search_result_text = WebSearchTransformation.format_search_response(result) @@ -737,7 +758,7 @@ class WebSearchInterceptionLogger(CustomLogger): ) raise - async def _execute_chat_completion_agentic_loop( # noqa: PLR0915 + async def _execute_chat_completion_agentic_loop( # noqa: PLR0915 self, model: str, messages: List[Dict], @@ -763,7 +784,7 @@ class WebSearchInterceptionLogger(CustomLogger): args = func.get("arguments", {}) if isinstance(args, dict): query = args.get("query") - + if query: verbose_logger.debug( f"WebSearchInterception: Queuing search for query='{query}'" @@ -789,9 +810,7 @@ class WebSearchInterceptionLogger(CustomLogger): verbose_logger.error( f"WebSearchInterception: Search {i} failed with error: {str(result)}" ) - final_search_results.append( - f"Search failed: {str(result)}" - ) + final_search_results.append(f"Search failed: {str(result)}") elif isinstance(result, str): final_search_results.append(cast(str, result)) else: @@ -801,7 +820,10 @@ class WebSearchInterceptionLogger(CustomLogger): final_search_results.append(str(result)) # Build assistant and tool messages using transformation - assistant_message, tool_messages_or_user = WebSearchTransformation.transform_response( + ( + assistant_message, + tool_messages_or_user, + ) = WebSearchTransformation.transform_response( tool_calls=tool_calls, search_results=final_search_results, response_format=response_format, @@ -810,10 +832,15 @@ class WebSearchInterceptionLogger(CustomLogger): # Make follow-up request with search results # For OpenAI format, tool_messages_or_user is a list of tool messages if response_format == "openai": - follow_up_messages = messages + [assistant_message] + cast(List[Dict], tool_messages_or_user) + follow_up_messages = ( + messages + [assistant_message] + cast(List[Dict], tool_messages_or_user) + ) else: # For Anthropic format (shouldn't happen in this method, but handle it) - follow_up_messages = messages + [assistant_message, cast(Dict, tool_messages_or_user)] + follow_up_messages = messages + [ + assistant_message, + cast(Dict, tool_messages_or_user), + ] verbose_logger.debug( "WebSearchInterception: Making follow-up chat completion request with search results" @@ -826,17 +853,19 @@ class WebSearchInterceptionLogger(CustomLogger): try: # Remove internal parameters that shouldn't be passed to follow-up request internal_params = { - '_websearch_interception', - 'acompletion', - 'litellm_logging_obj', - 'custom_llm_provider', - 'model_alias_map', - 'stream_response', - 'custom_prompt_dict', + "_websearch_interception", + "acompletion", + "litellm_logging_obj", + "custom_llm_provider", + "model_alias_map", + "stream_response", + "custom_prompt_dict", } kwargs_for_followup = { - k: v for k, v in kwargs.items() - if not k.startswith('_websearch_interception') and k not in internal_params + k: v + for k, v in kwargs.items() + if not k.startswith("_websearch_interception") + and k not in internal_params } # Get full model name from kwargs @@ -848,21 +877,29 @@ class WebSearchInterceptionLogger(CustomLogger): # Check if model already has a provider prefix if "/" not in model: full_model_name = f"{custom_llm_provider}/{model}" - + verbose_logger.debug( f"WebSearchInterception: Using model name: {full_model_name}" ) # Prepare tools for follow-up request (same as original) tools_param = optional_params.get("tools") - + # Remove tools and extra_body from optional_params to avoid issues # extra_body often contains internal LiteLLM params that shouldn't be forwarded optional_params_clean = { - k: v for k, v in optional_params.items() - if k not in {"tools", "extra_body", "model_alias_map","stream_response", "custom_prompt_dict" } + k: v + for k, v in optional_params.items() + if k + not in { + "tools", + "extra_body", + "model_alias_map", + "stream_response", + "custom_prompt_dict", + } } - + final_response = await litellm.acompletion( model=full_model_name, messages=follow_up_messages, @@ -870,7 +907,7 @@ class WebSearchInterceptionLogger(CustomLogger): **optional_params_clean, **kwargs_for_followup, ) - + verbose_logger.debug( f"WebSearchInterception: Follow-up request completed, response type: {type(final_response)}" ) diff --git a/litellm/integrations/websearch_interception/tools.py b/litellm/integrations/websearch_interception/tools.py index 7ef2b35004..e373b64cdd 100644 --- a/litellm/integrations/websearch_interception/tools.py +++ b/litellm/integrations/websearch_interception/tools.py @@ -41,11 +41,11 @@ def get_litellm_web_search_tool() -> Dict[str, Any]: "properties": { "query": { "type": "string", - "description": "The search query to execute" + "description": "The search query to execute", } }, - "required": ["query"] - } + "required": ["query"], + }, } @@ -73,19 +73,19 @@ def get_litellm_web_search_tool_openai() -> Dict[str, Any]: "properties": { "query": { "type": "string", - "description": "The search query to execute" + "description": "The search query to execute", } }, - "required": ["query"] - } - } + "required": ["query"], + }, + }, } def is_web_search_tool_chat_completion(tool: Dict[str, Any]) -> bool: """ Check if a tool is a web search tool for Chat Completions API (strict check). - + This is a stricter version that ONLY checks for the exact LiteLLM web search tool name. Use this for Chat Completions API to avoid false positives with user-defined tools. @@ -111,7 +111,7 @@ def is_web_search_tool_chat_completion(tool: Dict[str, Any]) -> bool: """ tool_name = tool.get("name", "") tool_type = tool.get("type", "") - + # Check for OpenAI format: {"type": "function", "function": {"name": "litellm_web_search"}} if tool_type == "function" and "function" in tool: function_def = tool.get("function", {}) @@ -155,7 +155,7 @@ def is_web_search_tool(tool: Dict[str, Any]) -> bool: """ tool_name = tool.get("name", "") tool_type = tool.get("type", "") - + # Check for OpenAI format: {"type": "function", "function": {"name": "..."}} if tool_type == "function" and "function" in tool: function_def = tool.get("function", {}) diff --git a/litellm/integrations/websearch_interception/transformation.py b/litellm/integrations/websearch_interception/transformation.py index e016899e0c..f777a7d741 100644 --- a/litellm/integrations/websearch_interception/transformation.py +++ b/litellm/integrations/websearch_interception/transformation.py @@ -81,9 +81,7 @@ class WebSearchTransformation: content = response.content or [] if not content: - verbose_logger.debug( - "WebSearchInterception: Response has empty content" - ) + verbose_logger.debug("WebSearchInterception: Response has empty content") return False, [] # Find all WebSearch tool_use blocks @@ -104,7 +102,9 @@ class WebSearchTransformation: # Check for LiteLLM standard or legacy web search tools # Handles: litellm_web_search, WebSearch, web_search if block_type == "tool_use" and block_name in ( - LITELLM_WEB_SEARCH_TOOL_NAME, "WebSearch", "web_search" + LITELLM_WEB_SEARCH_TOOL_NAME, + "WebSearch", + "web_search", ): # Convert to dict for easier handling tool_call = { @@ -125,7 +125,7 @@ class WebSearchTransformation: response: Any, ) -> Tuple[bool, List[Dict]]: """Parse OpenAI-style response for WebSearch tool_calls""" - + # Handle both dict and ModelResponse objects if isinstance(response, dict): choices = response.get("choices", []) @@ -138,9 +138,7 @@ class WebSearchTransformation: choices = response.choices or [] if not choices: - verbose_logger.debug( - "WebSearchInterception: Response has empty choices" - ) + verbose_logger.debug("WebSearchInterception: Response has empty choices") return False, [] # Get first choice's message @@ -149,11 +147,9 @@ class WebSearchTransformation: message = first_choice.get("message", {}) else: message = getattr(first_choice, "message", None) - + if not message: - verbose_logger.debug( - "WebSearchInterception: First choice has no message" - ) + verbose_logger.debug("WebSearchInterception: First choice has no message") return False, [] # Get tool_calls from message @@ -163,9 +159,7 @@ class WebSearchTransformation: openai_tool_calls = getattr(message, "tool_calls", None) or [] if not openai_tool_calls: - verbose_logger.debug( - "WebSearchInterception: Message has no tool_calls" - ) + verbose_logger.debug("WebSearchInterception: Message has no tool_calls") return False, [] # Find all WebSearch tool calls @@ -176,18 +170,30 @@ class WebSearchTransformation: tool_id = tool_call.get("id") tool_type = tool_call.get("type") function = tool_call.get("function", {}) - function_name = function.get("name") if isinstance(function, dict) else getattr(function, "name", None) - function_arguments = function.get("arguments") if isinstance(function, dict) else getattr(function, "arguments", None) + function_name = ( + function.get("name") + if isinstance(function, dict) + else getattr(function, "name", None) + ) + function_arguments = ( + function.get("arguments") + if isinstance(function, dict) + else getattr(function, "arguments", None) + ) else: tool_id = getattr(tool_call, "id", None) tool_type = getattr(tool_call, "type", None) function = getattr(tool_call, "function", None) function_name = getattr(function, "name", None) if function else None - function_arguments = getattr(function, "arguments", None) if function else None + function_arguments = ( + getattr(function, "arguments", None) if function else None + ) # Check for LiteLLM standard or legacy web search tools if tool_type == "function" and function_name in ( - LITELLM_WEB_SEARCH_TOOL_NAME, "WebSearch", "web_search" + LITELLM_WEB_SEARCH_TOOL_NAME, + "WebSearch", + "web_search", ): # Parse arguments (might be JSON string) if isinstance(function_arguments, str): @@ -320,7 +326,9 @@ class WebSearchTransformation: "type": "function", "function": { "name": tc["name"], - "arguments": json.dumps(tc["input"]) if isinstance(tc["input"], dict) else str(tc["input"]), + "arguments": json.dumps(tc["input"]) + if isinstance(tc["input"], dict) + else str(tc["input"]), }, } for tc in tool_calls diff --git a/litellm/integrations/weights_biases.py b/litellm/integrations/weights_biases.py index 0d011e26ae..028b6e69a8 100644 --- a/litellm/integrations/weights_biases.py +++ b/litellm/integrations/weights_biases.py @@ -44,7 +44,9 @@ try: request, response, time_elapsed ) else: - logger.debug(f"Unknown OpenAI response object: {response['object']}") + logger.debug( + f"Unknown OpenAI response object: {response['object']}" + ) except Exception as e: logger.warning(f"Failed to resolve request/response: {e}") return None diff --git a/litellm/interactions/http_handler.py b/litellm/interactions/http_handler.py index 4b4ed9be4d..7fead07043 100644 --- a/litellm/interactions/http_handler.py +++ b/litellm/interactions/http_handler.py @@ -86,11 +86,17 @@ class InteractionsHTTPHandler: ) -> Union[ InteractionsAPIResponse, Iterator[InteractionsAPIStreamingResponse], - Coroutine[Any, Any, Union[InteractionsAPIResponse, AsyncIterator[InteractionsAPIStreamingResponse]]], + Coroutine[ + Any, + Any, + Union[ + InteractionsAPIResponse, AsyncIterator[InteractionsAPIStreamingResponse] + ], + ], ]: """ Create a new interaction (synchronous or async based on _is_async flag). - + Per Google's OpenAPI spec, the endpoint is POST /{api_version}/interactions """ if _is_async: @@ -199,7 +205,9 @@ class InteractionsHTTPHandler: timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[AsyncHTTPHandler] = None, stream: Optional[bool] = None, - ) -> Union[InteractionsAPIResponse, AsyncIterator[InteractionsAPIStreamingResponse]]: + ) -> Union[ + InteractionsAPIResponse, AsyncIterator[InteractionsAPIStreamingResponse] + ]: """ Create a new interaction (async version). """ @@ -287,7 +295,7 @@ class InteractionsHTTPHandler: interactions_api_config: BaseInteractionsAPIConfig, ) -> SyncInteractionsAPIStreamingIterator: """Create a synchronous streaming iterator. - + Google AI's streaming format uses SSE (Server-Sent Events). Returns a proper streaming iterator that yields chunks as they arrive. """ @@ -306,7 +314,7 @@ class InteractionsHTTPHandler: interactions_api_config: BaseInteractionsAPIConfig, ) -> InteractionsAPIStreamingIterator: """Create an asynchronous streaming iterator. - + Google AI's streaming format uses SSE (Server-Sent Events). Returns a proper streaming iterator that yields chunks as they arrive. """ @@ -687,4 +695,3 @@ class InteractionsHTTPHandler: # Initialize the HTTP handler singleton interactions_http_handler = InteractionsHTTPHandler() - diff --git a/litellm/interactions/litellm_responses_transformation/__init__.py b/litellm/interactions/litellm_responses_transformation/__init__.py index 2450a9f3d2..6f6b32503d 100644 --- a/litellm/interactions/litellm_responses_transformation/__init__.py +++ b/litellm/interactions/litellm_responses_transformation/__init__.py @@ -13,4 +13,3 @@ __all__ = [ "LiteLLMResponsesInteractionsHandler", "LiteLLMResponsesInteractionsConfig", # Transformation config class (not BaseInteractionsAPIConfig) ] - diff --git a/litellm/interactions/litellm_responses_transformation/handler.py b/litellm/interactions/litellm_responses_transformation/handler.py index c2df8f96ef..b121ee37de 100644 --- a/litellm/interactions/litellm_responses_transformation/handler.py +++ b/litellm/interactions/litellm_responses_transformation/handler.py @@ -56,7 +56,7 @@ class LiteLLMResponsesInteractionsHandler: ]: """ Handle Interactions API request by calling litellm.responses(). - + Args: model: The model to use input: The input content @@ -65,22 +65,20 @@ class LiteLLMResponsesInteractionsHandler: _is_async: Whether this is an async call stream: Whether to stream the response **kwargs: Additional parameters - + Returns: InteractionsAPIResponse or streaming iterator """ # Transform interactions request to responses request - responses_request = ( - LiteLLMResponsesInteractionsConfig.transform_interactions_request_to_responses_request( - model=model, - input=input, - optional_params=optional_params, - custom_llm_provider=custom_llm_provider, - stream=stream, - **kwargs, - ) + responses_request = LiteLLMResponsesInteractionsConfig.transform_interactions_request_to_responses_request( + model=model, + input=input, + optional_params=optional_params, + custom_llm_provider=custom_llm_provider, + stream=stream, + **kwargs, ) - + if _is_async: return self.async_interactions_api_handler( responses_request=responses_request, @@ -89,14 +87,14 @@ class LiteLLMResponsesInteractionsHandler: optional_params=optional_params, **kwargs, ) - + # Call litellm.responses() # Note: litellm.responses() returns Union[ResponsesAPIResponse, BaseResponsesAPIStreamingIterator] # but the type checker may see it as a coroutine in some contexts responses_response = litellm.responses( **responses_request, ) - + # Handle streaming response if isinstance(responses_response, BaseResponsesAPIStreamingIterator): return LiteLLMResponsesInteractionsStreamingIterator( @@ -107,11 +105,11 @@ class LiteLLMResponsesInteractionsHandler: custom_llm_provider=custom_llm_provider, litellm_metadata=kwargs.get("litellm_metadata", {}), ) - + # At this point, responses_response must be ResponsesAPIResponse (not streaming) # Cast to satisfy type checker since we've already checked it's not a streaming iterator responses_api_response = cast(ResponsesAPIResponse, responses_response) - + # Transform responses response to interactions response return LiteLLMResponsesInteractionsConfig.transform_responses_response_to_interactions_response( responses_response=responses_api_response, @@ -125,14 +123,16 @@ class LiteLLMResponsesInteractionsHandler: input: Optional[InteractionInput], optional_params: InteractionsAPIOptionalRequestParams, **kwargs, - ) -> Union[InteractionsAPIResponse, AsyncIterator[InteractionsAPIStreamingResponse]]: + ) -> Union[ + InteractionsAPIResponse, AsyncIterator[InteractionsAPIStreamingResponse] + ]: """Async handler for interactions API requests.""" # Call litellm.aresponses() # Note: litellm.aresponses() returns Union[ResponsesAPIResponse, BaseResponsesAPIStreamingIterator] responses_response = await litellm.aresponses( **responses_request, ) - + # Handle streaming response if isinstance(responses_response, BaseResponsesAPIStreamingIterator): return LiteLLMResponsesInteractionsStreamingIterator( @@ -143,14 +143,13 @@ class LiteLLMResponsesInteractionsHandler: custom_llm_provider=responses_request.get("custom_llm_provider"), litellm_metadata=kwargs.get("litellm_metadata", {}), ) - + # At this point, responses_response must be ResponsesAPIResponse (not streaming) # Cast to satisfy type checker since we've already checked it's not a streaming iterator responses_api_response = cast(ResponsesAPIResponse, responses_response) - + # Transform responses response to interactions response return LiteLLMResponsesInteractionsConfig.transform_responses_response_to_interactions_response( responses_response=responses_api_response, model=model, ) - diff --git a/litellm/interactions/litellm_responses_transformation/streaming_iterator.py b/litellm/interactions/litellm_responses_transformation/streaming_iterator.py index 511b69e83b..72a3afbc3c 100644 --- a/litellm/interactions/litellm_responses_transformation/streaming_iterator.py +++ b/litellm/interactions/litellm_responses_transformation/streaming_iterator.py @@ -26,7 +26,7 @@ from litellm.types.llms.openai import ( class LiteLLMResponsesInteractionsStreamingIterator: """ Iterator that wraps Responses API streaming and transforms chunks to Interactions API format. - + This class handles both sync and async iteration, transforming Responses API streaming events (output.text.delta, response.completed, etc.) to Interactions API streaming events (content.delta, interaction.complete, etc.). @@ -58,11 +58,11 @@ class LiteLLMResponsesInteractionsStreamingIterator: ) -> Optional[InteractionsAPIStreamingResponse]: """ Transform a Responses API streaming chunk to an Interactions API streaming chunk. - + Responses API events: - output.text.delta -> content.delta - response.completed -> interaction.complete - + Interactions API events: - interaction.start - content.start @@ -72,23 +72,26 @@ class LiteLLMResponsesInteractionsStreamingIterator: """ if not responses_chunk: return None - + # Handle OutputTextDeltaEvent -> content.delta if isinstance(responses_chunk, OutputTextDeltaEvent): - delta_text = responses_chunk.delta if isinstance(responses_chunk.delta, str) else "" + delta_text = ( + responses_chunk.delta if isinstance(responses_chunk.delta, str) else "" + ) self.collected_text += delta_text - + # Send interaction.start if not sent if not self.sent_interaction_start: self.sent_interaction_start = True return InteractionsAPIStreamingResponse( event_type="interaction.start", - id=getattr(responses_chunk, "item_id", None) or f"interaction_{id(self)}", + id=getattr(responses_chunk, "item_id", None) + or f"interaction_{id(self)}", object="interaction", status="in_progress", model=self.model, ) - + # Send content.start if not sent if not self.sent_content_start: self.sent_content_start = True @@ -98,7 +101,7 @@ class LiteLLMResponsesInteractionsStreamingIterator: object="content", delta={"type": "text", "text": ""}, ) - + # Send content.delta return InteractionsAPIStreamingResponse( event_type="content.delta", @@ -106,12 +109,16 @@ class LiteLLMResponsesInteractionsStreamingIterator: object="content", delta={"text": delta_text}, ) - + # Handle ResponseCreatedEvent or ResponseInProgressEvent -> interaction.start if isinstance(responses_chunk, (ResponseCreatedEvent, ResponseInProgressEvent)): if not self.sent_interaction_start: self.sent_interaction_start = True - response_id = getattr(responses_chunk.response, "id", None) if hasattr(responses_chunk, "response") else None + response_id = ( + getattr(responses_chunk.response, "id", None) + if hasattr(responses_chunk, "response") + else None + ) return InteractionsAPIStreamingResponse( event_type="interaction.start", id=response_id or f"interaction_{id(self)}", @@ -119,17 +126,17 @@ class LiteLLMResponsesInteractionsStreamingIterator: status="in_progress", model=self.model, ) - + # Handle ResponseCompletedEvent -> interaction.complete if isinstance(responses_chunk, ResponseCompletedEvent): self.finished = True response = responses_chunk.response - + # Send content.stop first if content was started if self.sent_content_start: # Note: We'll send this in the iterator, not here pass - + # Send interaction.complete return InteractionsAPIStreamingResponse( event_type="interaction.complete", @@ -144,7 +151,7 @@ class LiteLLMResponsesInteractionsStreamingIterator: } ], ) - + # For other event types, return None (skip) return None @@ -156,26 +163,36 @@ class LiteLLMResponsesInteractionsStreamingIterator: """Get next chunk in sync mode.""" if self.finished: raise StopIteration - + # Check if we have a pending interaction.complete to send if hasattr(self, "_pending_interaction_complete"): - pending: InteractionsAPIStreamingResponse = getattr(self, "_pending_interaction_complete") + pending: InteractionsAPIStreamingResponse = getattr( + self, "_pending_interaction_complete" + ) delattr(self, "_pending_interaction_complete") return pending - + # Use a loop instead of recursion to avoid stack overflow - sync_iterator = cast(SyncResponsesAPIStreamingIterator, self.responses_stream_iterator) + sync_iterator = cast( + SyncResponsesAPIStreamingIterator, self.responses_stream_iterator + ) while True: try: # Get next chunk from responses API stream chunk = next(sync_iterator) - + # Transform chunk (chunk is already a ResponsesAPIStreamingResponse) - transformed = self._transform_responses_chunk_to_interactions_chunk(chunk) - + transformed = self._transform_responses_chunk_to_interactions_chunk( + chunk + ) + if transformed: # If we finished and content was started, send content.stop before interaction.complete - if self.finished and self.sent_content_start and transformed.event_type == "interaction.complete": + if ( + self.finished + and self.sent_content_start + and transformed.event_type == "interaction.complete" + ): # Send content.stop first content_stop = InteractionsAPIStreamingResponse( event_type="content.stop", @@ -187,12 +204,12 @@ class LiteLLMResponsesInteractionsStreamingIterator: self._pending_interaction_complete = transformed return content_stop return transformed - + # If no transformation, continue to next chunk (loop continues) - + except StopIteration: self.finished = True - + # Send final events if needed if self.sent_content_start: return InteractionsAPIStreamingResponse( @@ -200,7 +217,7 @@ class LiteLLMResponsesInteractionsStreamingIterator: object="content", delta={"type": "text", "text": self.collected_text}, ) - + raise StopIteration def __aiter__(self) -> AsyncIterator[InteractionsAPIStreamingResponse]: @@ -211,26 +228,36 @@ class LiteLLMResponsesInteractionsStreamingIterator: """Get next chunk in async mode.""" if self.finished: raise StopAsyncIteration - + # Check if we have a pending interaction.complete to send if hasattr(self, "_pending_interaction_complete"): - pending: InteractionsAPIStreamingResponse = getattr(self, "_pending_interaction_complete") + pending: InteractionsAPIStreamingResponse = getattr( + self, "_pending_interaction_complete" + ) delattr(self, "_pending_interaction_complete") return pending - + # Use a loop instead of recursion to avoid stack overflow - async_iterator = cast(ResponsesAPIStreamingIterator, self.responses_stream_iterator) + async_iterator = cast( + ResponsesAPIStreamingIterator, self.responses_stream_iterator + ) while True: try: # Get next chunk from responses API stream chunk = await async_iterator.__anext__() - + # Transform chunk (chunk is already a ResponsesAPIStreamingResponse) - transformed = self._transform_responses_chunk_to_interactions_chunk(chunk) - + transformed = self._transform_responses_chunk_to_interactions_chunk( + chunk + ) + if transformed: # If we finished and content was started, send content.stop before interaction.complete - if self.finished and self.sent_content_start and transformed.event_type == "interaction.complete": + if ( + self.finished + and self.sent_content_start + and transformed.event_type == "interaction.complete" + ): # Send content.stop first content_stop = InteractionsAPIStreamingResponse( event_type="content.stop", @@ -242,12 +269,12 @@ class LiteLLMResponsesInteractionsStreamingIterator: self._pending_interaction_complete = transformed return content_stop return transformed - + # If no transformation, continue to next chunk (loop continues) - + except StopAsyncIteration: self.finished = True - + # Send final events if needed if self.sent_content_start: return InteractionsAPIStreamingResponse( @@ -255,6 +282,5 @@ class LiteLLMResponsesInteractionsStreamingIterator: object="content", delta={"type": "text", "text": self.collected_text}, ) - - raise StopAsyncIteration + raise StopAsyncIteration diff --git a/litellm/interactions/litellm_responses_transformation/transformation.py b/litellm/interactions/litellm_responses_transformation/transformation.py index 24b2c5dbde..b07e61c76d 100644 --- a/litellm/interactions/litellm_responses_transformation/transformation.py +++ b/litellm/interactions/litellm_responses_transformation/transformation.py @@ -32,7 +32,7 @@ class LiteLLMResponsesInteractionsConfig: ) -> Dict[str, Any]: """ Transform an Interactions API request to a Responses API request. - + Key transformations: - system_instruction -> instructions - input (string | Turn[]) -> input (ResponseInputParam) @@ -42,23 +42,23 @@ class LiteLLMResponsesInteractionsConfig: responses_request: Dict[str, Any] = { "model": model, } - + # Transform input if input is not None: - responses_request["input"] = ( - LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input( - input - ) + responses_request[ + "input" + ] = LiteLLMResponsesInteractionsConfig._transform_interactions_input_to_responses_input( + input ) - + # Transform system_instruction -> instructions if optional_params.get("system_instruction"): responses_request["instructions"] = optional_params["system_instruction"] - + # Transform tools (similar format, pass through for now) if optional_params.get("tools"): responses_request["tools"] = optional_params["tools"] - + # Transform generation_config to temperature, top_p, etc. generation_config = optional_params.get("generation_config") if generation_config: @@ -71,17 +71,19 @@ class LiteLLMResponsesInteractionsConfig: # Responses API doesn't have top_k, skip it pass if "max_output_tokens" in generation_config: - responses_request["max_output_tokens"] = generation_config["max_output_tokens"] - + responses_request["max_output_tokens"] = generation_config[ + "max_output_tokens" + ] + # Pass through other optional params that match passthrough_params = ["stream", "store", "metadata", "user"] for param in passthrough_params: if param in optional_params and optional_params[param] is not None: responses_request[param] = optional_params[param] - + # Add any extra kwargs responses_request.update(kwargs) - + return responses_request @staticmethod @@ -90,12 +92,12 @@ class LiteLLMResponsesInteractionsConfig: ) -> ResponseInputParam: """ Transform Interactions API input to Responses API input format. - + Interactions API input can be: - string: "Hello" - Turn[]: [{"role": "user", "content": [...]}] - Content object - + Responses API input is: - string: "Hello" - Message[]: [{"role": "user", "content": [...]}] @@ -103,7 +105,7 @@ class LiteLLMResponsesInteractionsConfig: if isinstance(input, str): # ResponseInputParam accepts str return cast(ResponseInputParam, input) - + if isinstance(input, list): # Turn[] format - convert to Responses API Message[] format messages = [] @@ -111,21 +113,25 @@ class LiteLLMResponsesInteractionsConfig: if isinstance(turn, dict): role = turn.get("role", "user") content = turn.get("content", []) - + # Transform content array transformed_content = ( - LiteLLMResponsesInteractionsConfig._transform_content_array(content) + LiteLLMResponsesInteractionsConfig._transform_content_array( + content + ) + ) + + messages.append( + { + "role": role, + "content": transformed_content, + } ) - - messages.append({ - "role": role, - "content": transformed_content, - }) elif isinstance(turn, Turn): # Pydantic model role = turn.role if hasattr(turn, "role") else "user" content = turn.content if hasattr(turn, "content") else [] - + # Ensure content is a list for _transform_content_array # Cast to List[Any] to handle various content types if isinstance(content, list): @@ -134,27 +140,38 @@ class LiteLLMResponsesInteractionsConfig: content_list = [content] else: content_list = [] - + transformed_content = ( - LiteLLMResponsesInteractionsConfig._transform_content_array(content_list) + LiteLLMResponsesInteractionsConfig._transform_content_array( + content_list + ) ) - - messages.append({ - "role": role, - "content": transformed_content, - }) - + + messages.append( + { + "role": role, + "content": transformed_content, + } + ) + return cast(ResponseInputParam, messages) - + # Single content object - wrap in message if isinstance(input, dict): - return cast(ResponseInputParam, [{ - "role": "user", - "content": LiteLLMResponsesInteractionsConfig._transform_content_array( - input.get("content", []) if isinstance(input.get("content"), list) else [input] - ), - }]) - + return cast( + ResponseInputParam, + [ + { + "role": "user", + "content": LiteLLMResponsesInteractionsConfig._transform_content_array( + input.get("content", []) + if isinstance(input.get("content"), list) + else [input] + ), + } + ], + ) + # Fallback: convert to string return cast(ResponseInputParam, str(input)) @@ -164,7 +181,7 @@ class LiteLLMResponsesInteractionsConfig: if not isinstance(content, list): # Single content item - wrap in array content = [content] - + transformed: List[Dict[str, Any]] = [] for item in content: if isinstance(item, dict): @@ -192,7 +209,7 @@ class LiteLLMResponsesInteractionsConfig: else: # Fallback: wrap in text format transformed.append({"type": "text", "text": str(item)}) - + return transformed @staticmethod @@ -202,7 +219,7 @@ class LiteLLMResponsesInteractionsConfig: ) -> InteractionsAPIResponse: """ Transform a Responses API response to an Interactions API response. - + Key transformations: - Extract text from output[].content[].text - Convert created_at (int) to created (ISO string) @@ -221,23 +238,29 @@ class LiteLLMResponsesInteractionsConfig: # Check if content_item has text attribute text = getattr(content_item, "text", None) if text is not None: - outputs.append({ - "type": "text", - "text": text, - }) - elif isinstance(content_item, dict) and content_item.get("type") == "text": + outputs.append( + { + "type": "text", + "text": text, + } + ) + elif ( + isinstance(content_item, dict) + and content_item.get("type") == "text" + ): outputs.append(content_item) - + # Convert created_at to ISO string created_at = getattr(responses_response, "created_at", None) if isinstance(created_at, int): from datetime import datetime + created = datetime.fromtimestamp(created_at).isoformat() elif created_at is not None and hasattr(created_at, "isoformat"): created = created_at.isoformat() else: created = None - + # Map status status = getattr(responses_response, "status", "completed") if status == "completed": @@ -246,7 +269,7 @@ class LiteLLMResponsesInteractionsConfig: interactions_status = "in_progress" else: interactions_status = status - + # Build interactions response interactions_response_dict: Dict[str, Any] = { "id": getattr(responses_response, "id", ""), @@ -256,7 +279,7 @@ class LiteLLMResponsesInteractionsConfig: "model": model or getattr(responses_response, "model", ""), "created": created, } - + # Add usage if available # Map Responses API usage (input_tokens, output_tokens) to Interactions API spec format # (total_input_tokens, total_output_tokens) @@ -266,12 +289,11 @@ class LiteLLMResponsesInteractionsConfig: "total_input_tokens": getattr(usage, "input_tokens", 0), "total_output_tokens": getattr(usage, "output_tokens", 0), } - + # Add role interactions_response_dict["role"] = "model" - + # Add updated (same as created for now) interactions_response_dict["updated"] = created - - return InteractionsAPIResponse(**interactions_response_dict) + return InteractionsAPIResponse(**interactions_response_dict) diff --git a/litellm/interactions/main.py b/litellm/interactions/main.py index fb811b25b2..2b1786ac3a 100644 --- a/litellm/interactions/main.py +++ b/litellm/interactions/main.py @@ -105,9 +105,9 @@ async def acreate( ) -> Union[InteractionsAPIResponse, AsyncIterator[InteractionsAPIStreamingResponse]]: """ Async: Create a new interaction using Google's Interactions API. - + Per OpenAPI spec, provide either `model` or `agent`. - + Args: model: The model to use (e.g., "gemini-2.5-flash") agent: The agent to use (e.g., "deep-research-pro-preview-12-2025") @@ -126,7 +126,7 @@ async def acreate( extra_body: Additional body parameters timeout: Request timeout custom_llm_provider: Override the LLM provider - + Returns: InteractionsAPIResponse or async iterator for streaming """ @@ -134,14 +134,14 @@ async def acreate( try: loop = asyncio.get_event_loop() kwargs["acreate_interaction"] = True - + if custom_llm_provider is None and model: _, custom_llm_provider, _, _ = litellm.get_llm_provider( model=model, api_base=kwargs.get("api_base", None) ) elif custom_llm_provider is None: custom_llm_provider = "gemini" - + func = partial( create, model=model, @@ -163,16 +163,16 @@ async def acreate( custom_llm_provider=custom_llm_provider, **kwargs, ) - + ctx = contextvars.copy_context() func_with_context = partial(ctx.run, func) init_response = await loop.run_in_executor(None, func_with_context) - + if asyncio.iscoroutine(init_response): response = await init_response else: response = init_response - + return response # type: ignore except Exception as e: raise litellm.exception_type( @@ -219,13 +219,17 @@ def create( ) -> Union[ InteractionsAPIResponse, Iterator[InteractionsAPIStreamingResponse], - Coroutine[Any, Any, Union[InteractionsAPIResponse, AsyncIterator[InteractionsAPIStreamingResponse]]], + Coroutine[ + Any, + Any, + Union[InteractionsAPIResponse, AsyncIterator[InteractionsAPIStreamingResponse]], + ], ]: """ Sync: Create a new interaction using Google's Interactions API. - + Per OpenAPI spec, provide either `model` or `agent`. - + Args: model: The model to use (e.g., "gemini-2.5-flash") agent: The agent to use (e.g., "deep-research-pro-preview-12-2025") @@ -244,47 +248,53 @@ def create( extra_body: Additional body parameters timeout: Request timeout custom_llm_provider: Override the LLM provider - + Returns: InteractionsAPIResponse or iterator for streaming """ local_vars = locals() - + try: litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) _is_async = kwargs.pop("acreate_interaction", False) is True - + litellm_params = GenericLiteLLMParams(**kwargs) - + if model: model, custom_llm_provider, _, _ = litellm.get_llm_provider( - model=model, - custom_llm_provider=custom_llm_provider, - api_base=litellm_params.api_base, - api_key=litellm_params.api_key, - ) + model=model, + custom_llm_provider=custom_llm_provider, + api_base=litellm_params.api_base, + api_key=litellm_params.api_key, + ) else: custom_llm_provider = custom_llm_provider or "gemini" - + interactions_api_config = get_provider_interactions_api_config( provider=custom_llm_provider, model=model, ) - + # Get optional params using utility (similar to responses API pattern) local_vars.update(kwargs) - optional_params = InteractionsAPIRequestUtils.get_requested_interactions_api_optional_params( - local_vars + optional_params = ( + InteractionsAPIRequestUtils.get_requested_interactions_api_optional_params( + local_vars + ) ) - + # Check if this is a bridge provider (litellm_responses) - similar to responses API # Either provider is explicitly "litellm_responses" or no config found (bridge to responses) - if custom_llm_provider == "litellm_responses" or interactions_api_config is None: + if ( + custom_llm_provider == "litellm_responses" + or interactions_api_config is None + ): # Bridge to litellm.responses() for non-native providers from litellm.interactions.litellm_responses_transformation.handler import ( LiteLLMResponsesInteractionsHandler, ) + handler = LiteLLMResponsesInteractionsHandler() return handler.interactions_api_handler( model=model or "", @@ -295,14 +305,14 @@ def create( stream=stream, **kwargs, ) - + litellm_logging_obj.update_environment_variables( model=model, optional_params=dict(optional_params), litellm_params={"litellm_call_id": litellm_call_id}, custom_llm_provider=custom_llm_provider, ) - + response = interactions_http_handler.create_interaction( model=model, agent=agent, @@ -318,7 +328,7 @@ def create( _is_async=_is_async, stream=stream, ) - + return response except Exception as e: raise litellm.exception_type( @@ -348,7 +358,7 @@ async def aget( try: loop = asyncio.get_event_loop() kwargs["aget_interaction"] = True - + func = partial( get, interaction_id=interaction_id, @@ -357,16 +367,16 @@ async def aget( custom_llm_provider=custom_llm_provider or "gemini", **kwargs, ) - + ctx = contextvars.copy_context() func_with_context = partial(ctx.run, func) init_response = await loop.run_in_executor(None, func_with_context) - + if asyncio.iscoroutine(init_response): response = await init_response else: response = init_response - + return response # type: ignore except Exception as e: raise litellm.exception_type( @@ -389,28 +399,30 @@ def get( """Sync: Get an interaction by its ID.""" local_vars = locals() custom_llm_provider = custom_llm_provider or "gemini" - + try: litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) _is_async = kwargs.pop("aget_interaction", False) is True - + litellm_params = GenericLiteLLMParams(**kwargs) - + interactions_api_config = get_provider_interactions_api_config( provider=custom_llm_provider, ) - + if interactions_api_config is None: - raise ValueError(f"Interactions API not supported for: {custom_llm_provider}") - + raise ValueError( + f"Interactions API not supported for: {custom_llm_provider}" + ) + litellm_logging_obj.update_environment_variables( model=None, optional_params={"interaction_id": interaction_id}, litellm_params={"litellm_call_id": litellm_call_id}, custom_llm_provider=custom_llm_provider, ) - + return interactions_http_handler.get_interaction( interaction_id=interaction_id, interactions_api_config=interactions_api_config, @@ -449,7 +461,7 @@ async def adelete( try: loop = asyncio.get_event_loop() kwargs["adelete_interaction"] = True - + func = partial( delete, interaction_id=interaction_id, @@ -458,16 +470,16 @@ async def adelete( custom_llm_provider=custom_llm_provider or "gemini", **kwargs, ) - + ctx = contextvars.copy_context() func_with_context = partial(ctx.run, func) init_response = await loop.run_in_executor(None, func_with_context) - + if asyncio.iscoroutine(init_response): response = await init_response else: response = init_response - + return response # type: ignore except Exception as e: raise litellm.exception_type( @@ -490,28 +502,30 @@ def delete( """Sync: Delete an interaction by its ID.""" local_vars = locals() custom_llm_provider = custom_llm_provider or "gemini" - + try: litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) _is_async = kwargs.pop("adelete_interaction", False) is True - + litellm_params = GenericLiteLLMParams(**kwargs) - + interactions_api_config = get_provider_interactions_api_config( provider=custom_llm_provider, ) - + if interactions_api_config is None: - raise ValueError(f"Interactions API not supported for: {custom_llm_provider}") - + raise ValueError( + f"Interactions API not supported for: {custom_llm_provider}" + ) + litellm_logging_obj.update_environment_variables( model=None, optional_params={"interaction_id": interaction_id}, litellm_params={"litellm_call_id": litellm_call_id}, custom_llm_provider=custom_llm_provider, ) - + return interactions_http_handler.delete_interaction( interaction_id=interaction_id, interactions_api_config=interactions_api_config, @@ -550,7 +564,7 @@ async def acancel( try: loop = asyncio.get_event_loop() kwargs["acancel_interaction"] = True - + func = partial( cancel, interaction_id=interaction_id, @@ -559,16 +573,16 @@ async def acancel( custom_llm_provider=custom_llm_provider or "gemini", **kwargs, ) - + ctx = contextvars.copy_context() func_with_context = partial(ctx.run, func) init_response = await loop.run_in_executor(None, func_with_context) - + if asyncio.iscoroutine(init_response): response = await init_response else: response = init_response - + return response # type: ignore except Exception as e: raise litellm.exception_type( @@ -591,28 +605,30 @@ def cancel( """Sync: Cancel an interaction by its ID.""" local_vars = locals() custom_llm_provider = custom_llm_provider or "gemini" - + try: litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) _is_async = kwargs.pop("acancel_interaction", False) is True - + litellm_params = GenericLiteLLMParams(**kwargs) - + interactions_api_config = get_provider_interactions_api_config( provider=custom_llm_provider, ) - + if interactions_api_config is None: - raise ValueError(f"Interactions API not supported for: {custom_llm_provider}") - + raise ValueError( + f"Interactions API not supported for: {custom_llm_provider}" + ) + litellm_logging_obj.update_environment_variables( model=None, optional_params={"interaction_id": interaction_id}, litellm_params={"litellm_call_id": litellm_call_id}, custom_llm_provider=custom_llm_provider, ) - + return interactions_http_handler.cancel_interaction( interaction_id=interaction_id, interactions_api_config=interactions_api_config, diff --git a/litellm/interactions/streaming_iterator.py b/litellm/interactions/streaming_iterator.py index f65d08d3ca..a5a7f9e06e 100644 --- a/litellm/interactions/streaming_iterator.py +++ b/litellm/interactions/streaming_iterator.py @@ -61,7 +61,9 @@ class BaseInteractionsAPIStreamingIterator: "litellm_params", {} ), ) - _model_info: Dict = litellm_metadata.get("model_info", {}) if litellm_metadata else {} + _model_info: Dict = ( + litellm_metadata.get("model_info", {}) if litellm_metadata else {} + ) self._hidden_params = { "model_id": _model_info.get("id", None), "api_base": _api_base, @@ -91,10 +93,12 @@ class BaseInteractionsAPIStreamingIterator: # Format as InteractionsAPIStreamingResponse if isinstance(parsed_chunk, dict): - streaming_response = self.interactions_api_config.transform_streaming_response( - model=self.model, - parsed_chunk=parsed_chunk, - logging_obj=self.logging_obj, + streaming_response = ( + self.interactions_api_config.transform_streaming_response( + model=self.model, + parsed_chunk=parsed_chunk, + logging_obj=self.logging_obj, + ) ) # Store the completed response (check for status=completed) @@ -110,7 +114,9 @@ class BaseInteractionsAPIStreamingIterator: return None except json.JSONDecodeError: # If we can't parse the chunk, continue - verbose_logger.debug(f"Failed to parse streaming chunk: {stripped_chunk[:200]}...") + verbose_logger.debug( + f"Failed to parse streaming chunk: {stripped_chunk[:200]}..." + ) return None def _handle_logging_completed_response(self): @@ -171,6 +177,7 @@ class InteractionsAPIStreamingIterator(BaseInteractionsAPIStreamingIterator): def _handle_logging_completed_response(self): """Handle logging for completed responses in async context.""" import copy + logging_response = copy.deepcopy(self.completed_response) asyncio.create_task( @@ -244,6 +251,7 @@ class SyncInteractionsAPIStreamingIterator(BaseInteractionsAPIStreamingIterator) def _handle_logging_completed_response(self): """Handle logging for completed responses in sync context.""" import copy + logging_response = copy.deepcopy(self.completed_response) run_async_function( @@ -261,4 +269,3 @@ class SyncInteractionsAPIStreamingIterator(BaseInteractionsAPIStreamingIterator) start_time=self.start_time, end_time=datetime.now(), ) - diff --git a/litellm/interactions/utils.py b/litellm/interactions/utils.py index 4fc40916e5..3a18ddf52f 100644 --- a/litellm/interactions/utils.py +++ b/litellm/interactions/utils.py @@ -29,22 +29,23 @@ def get_provider_interactions_api_config( ) -> Optional[BaseInteractionsAPIConfig]: """ Get the interactions API config for the given provider. - + Args: provider: The LLM provider name model: Optional model name - + Returns: The provider-specific interactions API config, or None if not supported """ from litellm.types.utils import LlmProviders - + if provider == LlmProviders.GEMINI.value or provider == "gemini": from litellm.llms.gemini.interactions.transformation import ( GoogleAIStudioInteractionsConfig, ) + return GoogleAIStudioInteractionsConfig() - + return None @@ -76,7 +77,9 @@ class InteractionsAPIRequestUtils: special_params=special_params, custom_llm_provider=custom_llm_provider, additional_drop_params=additional_drop_params, - default_param_values={k: None for k in INTERACTIONS_API_OPTIONAL_PARAMS}, + default_param_values={ + k: None for k in INTERACTIONS_API_OPTIONAL_PARAMS + }, additional_endpoint_specific_params=["input", "model", "agent"], ) ) diff --git a/litellm/litellm_core_utils/app_crypto.py b/litellm/litellm_core_utils/app_crypto.py index 5ce6d8d77f..e47962d6a3 100644 --- a/litellm/litellm_core_utils/app_crypto.py +++ b/litellm/litellm_core_utils/app_crypto.py @@ -30,4 +30,4 @@ class AppCrypto: ct = base64.b64decode(enc["ciphertext"]) tag = base64.b64decode(enc["tag"]) data = aes.decrypt(nonce, ct + tag, aad) - return json.loads(data.decode()) \ No newline at end of file + return json.loads(data.decode()) diff --git a/litellm/litellm_core_utils/audio_utils/utils.py b/litellm/litellm_core_utils/audio_utils/utils.py index a7d12841e5..2141df1873 100644 --- a/litellm/litellm_core_utils/audio_utils/utils.py +++ b/litellm/litellm_core_utils/audio_utils/utils.py @@ -135,7 +135,7 @@ def get_audio_file_content_hash(file_obj: FileTypes) -> str: """ file_content: Optional[bytes] = None fallback_filename: Optional[str] = None - + if isinstance(file_obj, tuple): if len(file_obj) < 2: fallback_filename = str(file_obj[0]) if len(file_obj) > 0 else None @@ -145,7 +145,7 @@ def get_audio_file_content_hash(file_obj: FileTypes) -> str: else: file_content_obj = file_obj fallback_filename = get_audio_file_name(file_obj) - + try: if isinstance(file_content_obj, (bytes, bytearray)): file_content = bytes(file_content_obj) @@ -160,7 +160,11 @@ def get_audio_file_content_hash(file_obj: FileTypes) -> str: file_content = None elif hasattr(file_content_obj, "read"): try: - current_position = file_content_obj.tell() if hasattr(file_content_obj, "tell") else None + current_position = ( + file_content_obj.tell() + if hasattr(file_content_obj, "tell") + else None + ) if hasattr(file_content_obj, "seek"): file_content_obj.seek(0) file_content = file_content_obj.read() # type: ignore @@ -172,20 +176,20 @@ def get_audio_file_content_hash(file_obj: FileTypes) -> str: file_content = None except Exception: file_content = None - + if file_content is not None and isinstance(file_content, bytes): try: hash_object = hashlib.sha256(file_content) return hash_object.hexdigest() except Exception: pass - + if fallback_filename: - hash_object = hashlib.sha256(fallback_filename.encode('utf-8')) + hash_object = hashlib.sha256(fallback_filename.encode("utf-8")) return hash_object.hexdigest() - + file_obj_str = str(file_obj) - hash_object = hashlib.sha256(file_obj_str.encode('utf-8')) + hash_object = hashlib.sha256(file_obj_str.encode("utf-8")) return hash_object.hexdigest() diff --git a/litellm/litellm_core_utils/cached_imports.py b/litellm/litellm_core_utils/cached_imports.py index c3ab292e9c..1a3943cc51 100644 --- a/litellm/litellm_core_utils/cached_imports.py +++ b/litellm/litellm_core_utils/cached_imports.py @@ -24,6 +24,7 @@ def get_litellm_logging_class() -> Type["Logging"]: if _LiteLLMLogging is not None: return _LiteLLMLogging from litellm.litellm_core_utils.litellm_logging import Logging + _LiteLLMLogging = Logging return _LiteLLMLogging @@ -34,6 +35,7 @@ def get_coroutine_checker() -> "CoroutineChecker": if _coroutine_checker is not None: return _coroutine_checker from litellm.litellm_core_utils.coroutine_checker import coroutine_checker + _coroutine_checker = coroutine_checker return _coroutine_checker @@ -44,6 +46,7 @@ def get_set_callbacks() -> Callable: if _set_callbacks is not None: return _set_callbacks from litellm.litellm_core_utils.litellm_logging import set_callbacks + _set_callbacks = set_callbacks return _set_callbacks diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index 2aedb1c19d..e2e304931a 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -23,9 +23,9 @@ def load_cli_token() -> Optional[dict]: token_file = get_cli_token_file_path() if not os.path.exists(token_file): return None - + try: - with open(token_file, 'r') as f: + with open(token_file, "r") as f: return json.load(f) except (json.JSONDecodeError, IOError): return None @@ -34,13 +34,13 @@ def load_cli_token() -> Optional[dict]: def get_litellm_gateway_api_key() -> Optional[str]: """ Get the stored CLI API key for use with LiteLLM SDK. - + This function reads the token file created by `litellm-proxy login` and returns the API key for use in Python scripts. - + Returns: str: The API key if found, None otherwise - + Example: >>> import litellm >>> api_key = litellm.get_litellm_gateway_api_key() @@ -53,6 +53,6 @@ def get_litellm_gateway_api_key() -> Optional[str]: >>> ) """ token_data = load_cli_token() - if token_data and 'key' in token_data: - return token_data['key'] + if token_data and "key" in token_data: + return token_data["key"] return None diff --git a/litellm/litellm_core_utils/coroutine_checker.py b/litellm/litellm_core_utils/coroutine_checker.py index 368aee62ed..bf065e5a15 100644 --- a/litellm/litellm_core_utils/coroutine_checker.py +++ b/litellm/litellm_core_utils/coroutine_checker.py @@ -10,14 +10,14 @@ from litellm.constants import ( class CoroutineChecker: """Utility class for checking coroutine status of functions and callables. - + Simple bounded cache using WeakKeyDictionary to avoid memory leaks. """ - + def __init__(self): self._cache = WeakKeyDictionary() self._max_size = COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY - + def is_async_callable(self, callback: Any) -> bool: """Fast, cached check for whether a callback is an async function. Falls back gracefully if the object cannot be weak-referenced or cached. @@ -52,12 +52,13 @@ class CoroutineChecker: # Simple size enforcement: clear cache if it gets too large if len(self._cache) >= self._max_size: self._cache.clear() - + self._cache[callback] = result except Exception: pass return result + # Global instance for backward compatibility and convenience coroutine_checker = CoroutineChecker() diff --git a/litellm/litellm_core_utils/dot_notation_indexing.py b/litellm/litellm_core_utils/dot_notation_indexing.py index 1e835004e9..65810e83c6 100644 --- a/litellm/litellm_core_utils/dot_notation_indexing.py +++ b/litellm/litellm_core_utils/dot_notation_indexing.py @@ -107,7 +107,7 @@ def _parse_path_segments(path: str) -> list: # Match field names OR bracket expressions # Pattern: field_name (anything except . or [) | [anything_in_brackets] - pattern = r'[^\.\[]+|\[[^\]]*\]' + pattern = r"[^\.\[]+|\[[^\]]*\]" segments = re.findall(pattern, path) return segments @@ -158,7 +158,9 @@ def _delete_nested_value_custom( # Only recurse if element is a dict or list (nested structure) element = data[index] if isinstance(element, (dict, list)): - _delete_nested_value_custom(element, segments, segment_index + 1) + _delete_nested_value_custom( + element, segments, segment_index + 1 + ) except (ValueError, IndexError): # Invalid index, skip pass @@ -172,15 +174,23 @@ def _delete_nested_value_custom( else: # Navigate deeper if segment in data: - next_segment = segments[segment_index + 1] if segment_index + 1 < len(segments) else None + next_segment = ( + segments[segment_index + 1] + if segment_index + 1 < len(segments) + else None + ) # If next segment is array notation, current field should be list if next_segment and (next_segment.startswith("[")): if isinstance(data[segment], list): - _delete_nested_value_custom(data[segment], segments, segment_index + 1) + _delete_nested_value_custom( + data[segment], segments, segment_index + 1 + ) # Otherwise navigate into dict elif isinstance(data[segment], dict): - _delete_nested_value_custom(data[segment], segments, segment_index + 1) + _delete_nested_value_custom( + data[segment], segments, segment_index + 1 + ) def delete_nested_value( diff --git a/litellm/litellm_core_utils/duration_parser.py b/litellm/litellm_core_utils/duration_parser.py index 70c28c4e06..6d2b4226ff 100644 --- a/litellm/litellm_core_utils/duration_parser.py +++ b/litellm/litellm_core_utils/duration_parser.py @@ -64,12 +64,10 @@ def duration_in_seconds(duration: str) -> int: now = time.time() current_time = datetime.fromtimestamp(now) - if current_time.month == 12: - target_year = current_time.year + 1 - target_month = 1 - else: - target_year = current_time.year - target_month = current_time.month + value + # Calculate target month and year, handling overflow past December + total_months = current_time.month - 1 + value # 0-indexed months + target_year = current_time.year + total_months // 12 + target_month = total_months % 12 + 1 # back to 1-indexed # Determine the day to set for next month target_day = current_time.day diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 951485130b..bc54786420 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -73,7 +73,10 @@ class ExceptionCheckers: # Exclude param validation errors (e.g. OpenAI "user" param max 64 chars) if "string_above_max_length" in _error_str_lowercase: return False - if "invalid 'user'" in _error_str_lowercase and "string too long" in _error_str_lowercase: + if ( + "invalid 'user'" in _error_str_lowercase + and "string too long" in _error_str_lowercase + ): return False known_exception_substrings = [ "exceed context limit", @@ -97,7 +100,7 @@ class ExceptionCheckers: return True return False - + @staticmethod def is_azure_content_policy_violation_error(error_str: str) -> bool: """ @@ -443,7 +446,10 @@ def exception_type( # type: ignore # noqa: PLR0915 response=getattr(original_exception, "response", None), litellm_debug_info=extra_information, ) - elif "invalid_encrypted_content" in error_str or "could not be verified" in error_str: + elif ( + "invalid_encrypted_content" in error_str + or "could not be verified" in error_str + ): exception_mapping_worked = True helpful_message = ( f"{exception_provider} - {message}\n\n" @@ -2093,13 +2099,18 @@ def exception_type( # type: ignore # noqa: PLR0915 # content policy violation even when the top-level # code is generic (e.g. "invalid_request_error"). if azure_error_code != "content_policy_violation": - _inner = ( - body_dict["error"].get("inner_error") # type: ignore[index] - or body_dict["error"].get("innererror") # type: ignore[index] - ) - if isinstance(_inner, dict) and _inner.get( - "code" - ) == "ResponsibleAIPolicyViolation": + _inner = body_dict["error"].get( + "inner_error" + ) or body_dict[ # type: ignore[index] + "error" + ].get( + "innererror" + ) # type: ignore[index] + if ( + isinstance(_inner, dict) + and _inner.get("code") + == "ResponsibleAIPolicyViolation" + ): azure_error_code = "content_policy_violation" else: azure_error_code = body_dict.get("code") @@ -2135,19 +2146,25 @@ def exception_type( # type: ignore # noqa: PLR0915 ) elif ( azure_error_code == "content_policy_violation" - or ExceptionCheckers.is_azure_content_policy_violation_error(error_str) + or ExceptionCheckers.is_azure_content_policy_violation_error( + error_str + ) ): exception_mapping_worked = True from litellm.llms.azure.exception_mapping import ( AzureOpenAIExceptionMapping, ) + raise AzureOpenAIExceptionMapping.create_content_policy_violation_error( message=message, model=model, extra_information=extra_information, original_exception=original_exception, ) - elif azure_error_code == "invalid_encrypted_content" or "could not be verified" in error_str: + elif ( + azure_error_code == "invalid_encrypted_content" + or "could not be verified" in error_str + ): exception_mapping_worked = True helpful_message = ( f"AzureException - {message}\n\n" diff --git a/litellm/litellm_core_utils/fallback_utils.py b/litellm/litellm_core_utils/fallback_utils.py index aa5bdd9271..52eb35663b 100644 --- a/litellm/litellm_core_utils/fallback_utils.py +++ b/litellm/litellm_core_utils/fallback_utils.py @@ -3,7 +3,10 @@ from typing import Optional import litellm from litellm._logging import verbose_logger -from litellm.litellm_core_utils.core_helpers import safe_deep_copy, filter_internal_params +from litellm.litellm_core_utils.core_helpers import ( + safe_deep_copy, + filter_internal_params, +) from .asyncify import run_async_function diff --git a/litellm/litellm_core_utils/get_blog_posts.py b/litellm/litellm_core_utils/get_blog_posts.py index 4f054c78ff..f54deb5929 100644 --- a/litellm/litellm_core_utils/get_blog_posts.py +++ b/litellm/litellm_core_utils/get_blog_posts.py @@ -51,9 +51,7 @@ class GetBlogPosts: def load_local_blog_posts() -> List[Dict[str, str]]: """Load the bundled local backup blog posts.""" content = json.loads( - files("litellm") - .joinpath("blog_posts.json") - .read_text(encoding="utf-8") + files("litellm").joinpath("blog_posts.json").read_text(encoding="utf-8") ) return content.get("posts", []) diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index c91e4b6de1..ad9538ac17 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -2,36 +2,38 @@ from typing import Optional # Pre-define optional kwargs keys as frozenset for O(1) lookups # These are extracted from kwargs only if present, avoiding unnecessary .get() calls -_OPTIONAL_KWARGS_KEYS = frozenset({ - "azure_ad_token", - "tenant_id", - "client_id", - "client_secret", - "azure_username", - "azure_password", - "azure_scope", - "timeout", - "bucket_name", - "vertex_credentials", - "vertex_project", - "vertex_location", - "vertex_ai_project", - "vertex_ai_location", - "vertex_ai_credentials", - "aws_region_name", - "aws_access_key_id", - "aws_secret_access_key", - "aws_session_token", - "aws_session_name", - "aws_profile_name", - "aws_role_name", - "aws_web_identity_token", - "aws_sts_endpoint", - "aws_external_id", - "aws_bedrock_runtime_endpoint", - "tpm", - "rpm", -}) +_OPTIONAL_KWARGS_KEYS = frozenset( + { + "azure_ad_token", + "tenant_id", + "client_id", + "client_secret", + "azure_username", + "azure_password", + "azure_scope", + "timeout", + "bucket_name", + "vertex_credentials", + "vertex_project", + "vertex_location", + "vertex_ai_project", + "vertex_ai_location", + "vertex_ai_credentials", + "aws_region_name", + "aws_access_key_id", + "aws_secret_access_key", + "aws_session_token", + "aws_session_name", + "aws_profile_name", + "aws_role_name", + "aws_web_identity_token", + "aws_sts_endpoint", + "aws_external_id", + "aws_bedrock_runtime_endpoint", + "tpm", + "rpm", + } +) def _get_base_model_from_litellm_call_metadata( diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index d1ee17fdd2..3621841737 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -279,10 +279,16 @@ def get_llm_provider( # noqa: PLR0915 elif endpoint == "api.moonshot.ai/v1": custom_llm_provider = "moonshot" dynamic_api_key = get_secret_str("MOONSHOT_API_KEY") - elif endpoint == "api.minimax.io/anthropic" or endpoint == "api.minimaxi.com/anthropic": + elif ( + endpoint == "api.minimax.io/anthropic" + or endpoint == "api.minimaxi.com/anthropic" + ): custom_llm_provider = "minimax" dynamic_api_key = get_secret_str("MINIMAX_API_KEY") - elif endpoint == "api.minimax.io/v1" or endpoint == "api.minimaxi.com/v1": + elif ( + endpoint == "api.minimax.io/v1" + or endpoint == "api.minimaxi.com/v1" + ): custom_llm_provider = "minimax" dynamic_api_key = get_secret_str("MINIMAX_API_KEY") elif endpoint == "platform.publicai.co/v1": @@ -586,7 +592,11 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915 if api_base is None: api_base = litellm.BasetenConfig.get_api_base_for_model(model) else: - api_base = api_base or get_secret_str("BASETEN_API_BASE") or "https://inference.baseten.co/v1" + api_base = ( + api_base + or get_secret_str("BASETEN_API_BASE") + or "https://inference.baseten.co/v1" + ) dynamic_api_key = api_key or get_secret_str("BASETEN_API_KEY") elif custom_llm_provider == "sambanova": api_base = ( @@ -611,9 +621,7 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915 dynamic_api_key = api_key or get_secret_str("NEBIUS_API_KEY") elif custom_llm_provider == "ollama": api_base = ( - api_base - or get_secret("OLLAMA_API_BASE") - or "http://localhost:11434" + api_base or get_secret("OLLAMA_API_BASE") or "http://localhost:11434" ) # type: ignore dynamic_api_key = api_key or get_secret_str("OLLAMA_API_KEY") elif (custom_llm_provider == "ai21_chat") or ( @@ -927,17 +935,13 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915 elif custom_llm_provider == "langgraph": # LangGraph is a custom provider, just need to set api_base api_base = ( - api_base - or get_secret_str("LANGGRAPH_API_BASE") - or "http://localhost:2024" + api_base or get_secret_str("LANGGRAPH_API_BASE") or "http://localhost:2024" ) dynamic_api_key = api_key or get_secret_str("LANGGRAPH_API_KEY") elif custom_llm_provider == "manus": # Manus is OpenAI compatible for responses API api_base = ( - api_base - or get_secret_str("MANUS_API_BASE") - or "https://api.manus.im" + api_base or get_secret_str("MANUS_API_BASE") or "https://api.manus.im" ) dynamic_api_key = api_key or get_secret_str("MANUS_API_KEY") diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 5673064a23..7679358bbc 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -92,7 +92,10 @@ class GetModelCostMap: ) return False - if backup_model_count > 0 and fetched_count < backup_model_count * max_shrink_ratio: + if ( + backup_model_count > 0 + and fetched_count < backup_model_count * max_shrink_ratio + ): verbose_logger.warning( "LiteLLM: Fetched model cost map shrank significantly " "(fetched=%d, backup=%d, threshold=%.0f%%). " @@ -286,7 +289,9 @@ def get_model_cost_map(url: str) -> dict: url, ) _cost_map_source_info.source = "local" - _cost_map_source_info.fallback_reason = "Remote data failed integrity validation" + _cost_map_source_info.fallback_reason = ( + "Remote data failed integrity validation" + ) return _expand_model_aliases(GetModelCostMap.load_local_model_cost_map()) _cost_map_source_info.source = "remote" diff --git a/litellm/litellm_core_utils/get_supported_openai_params.py b/litellm/litellm_core_utils/get_supported_openai_params.py index 07065aff32..b72d7abeae 100644 --- a/litellm/litellm_core_utils/get_supported_openai_params.py +++ b/litellm/litellm_core_utils/get_supported_openai_params.py @@ -89,7 +89,9 @@ def get_supported_openai_params( # noqa: PLR0915 elif custom_llm_provider == "groq": return litellm.GroqChatConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "bedrock_mantle": - return litellm.BedrockMantleChatConfig().get_supported_openai_params(model=model) + return litellm.BedrockMantleChatConfig().get_supported_openai_params( + model=model + ) elif custom_llm_provider == "hosted_vllm": return litellm.HostedVLLMChatConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "vllm": @@ -120,9 +122,13 @@ def get_supported_openai_params( # noqa: PLR0915 return litellm.OpenAIConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "sap": if request_type == "chat_completion": - return litellm.GenAIHubOrchestrationConfig().get_supported_openai_params(model=model) + return litellm.GenAIHubOrchestrationConfig().get_supported_openai_params( + model=model + ) elif request_type == "embeddings": - return litellm.GenAIHubEmbeddingConfig().get_supported_openai_params(model=model) + return litellm.GenAIHubEmbeddingConfig().get_supported_openai_params( + model=model + ) elif custom_llm_provider == "azure": if litellm.AzureOpenAIO1Config().is_o_series_model(model=model): return litellm.AzureOpenAIO1Config().get_supported_openai_params( diff --git a/litellm/litellm_core_utils/json_validation_rule.py b/litellm/litellm_core_utils/json_validation_rule.py index 315a90fe30..bbfd3e6de9 100644 --- a/litellm/litellm_core_utils/json_validation_rule.py +++ b/litellm/litellm_core_utils/json_validation_rule.py @@ -4,93 +4,105 @@ from typing import Any, Dict, List, Union from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH -def normalize_json_schema_types(schema: Union[Dict[str, Any], List[Any], Any], depth: int = 0, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH) -> Union[Dict[str, Any], List[Any], Any]: +def normalize_json_schema_types( + schema: Union[Dict[str, Any], List[Any], Any], + depth: int = 0, + max_depth: int = DEFAULT_MAX_RECURSE_DEPTH, +) -> Union[Dict[str, Any], List[Any], Any]: """ Normalize JSON schema types from uppercase to lowercase format. - + Some providers (like certain Google services) use uppercase types like 'BOOLEAN', 'STRING', 'ARRAY', 'OBJECT' but standard JSON Schema requires lowercase: 'boolean', 'string', 'array', 'object' - + This function recursively normalizes all type fields in a schema to lowercase. - + Args: schema: The schema to normalize (dict, list, or other) depth: Current recursion depth max_depth: Maximum recursion depth to prevent infinite loops - + Returns: The normalized schema with lowercase types """ # Prevent infinite recursion if depth >= max_depth: return schema - + if not isinstance(schema, (dict, list)): return schema - + # Type mapping from uppercase to lowercase type_mapping = { - 'BOOLEAN': 'boolean', - 'STRING': 'string', - 'ARRAY': 'array', - 'OBJECT': 'object', - 'NUMBER': 'number', - 'INTEGER': 'integer', - 'NULL': 'null' + "BOOLEAN": "boolean", + "STRING": "string", + "ARRAY": "array", + "OBJECT": "object", + "NUMBER": "number", + "INTEGER": "integer", + "NULL": "null", } - + if isinstance(schema, list): - return [normalize_json_schema_types(item, depth + 1, max_depth) for item in schema] - + return [ + normalize_json_schema_types(item, depth + 1, max_depth) for item in schema + ] + if isinstance(schema, dict): normalized_schema: Dict[str, Any] = {} - + for key, value in schema.items(): - if key == 'type' and isinstance(value, str) and value in type_mapping: + if key == "type" and isinstance(value, str) and value in type_mapping: normalized_schema[key] = type_mapping[value] - elif key == 'properties' and isinstance(value, dict): + elif key == "properties" and isinstance(value, dict): # Recursively normalize properties normalized_schema[key] = { - prop_key: normalize_json_schema_types(prop_value, depth + 1, max_depth) + prop_key: normalize_json_schema_types( + prop_value, depth + 1, max_depth + ) for prop_key, prop_value in value.items() } - elif key == 'items' and isinstance(value, (dict, list)): + elif key == "items" and isinstance(value, (dict, list)): # Recursively normalize array items - normalized_schema[key] = normalize_json_schema_types(value, depth + 1, max_depth) + normalized_schema[key] = normalize_json_schema_types( + value, depth + 1, max_depth + ) elif isinstance(value, (dict, list)): # Recursively normalize any nested dict or list - normalized_schema[key] = normalize_json_schema_types(value, depth + 1, max_depth) + normalized_schema[key] = normalize_json_schema_types( + value, depth + 1, max_depth + ) else: normalized_schema[key] = value - + return normalized_schema - + return schema def normalize_tool_schema(tool: Dict[str, Any]) -> Dict[str, Any]: """ Normalize a tool's parameter schema to use standard JSON Schema lowercase types. - + Args: tool: The tool definition containing function parameters - + Returns: The tool with normalized schema types """ if not isinstance(tool, dict): return tool - + normalized_tool = tool.copy() - + # Normalize function parameters if present - if 'function' in tool and isinstance(tool['function'], dict): - normalized_tool['function'] = tool['function'].copy() - if 'parameters' in tool['function']: - normalized_tool['function']['parameters'] = normalize_json_schema_types( - tool['function']['parameters'] + if "function" in tool and isinstance(tool["function"], dict): + normalized_tool["function"] = tool["function"].copy() + if "parameters" in tool["function"]: + normalized_tool["function"]["parameters"] = normalize_json_schema_types( + tool["function"]["parameters"] ) - + return normalized_tool diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index f9c5d74ee3..e22d057bb6 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -352,9 +352,9 @@ class Logging(LiteLLMLoggingBaseClass): ) self.function_id = function_id self.streaming_chunks: List[Any] = [] # for generating complete stream response - self.sync_streaming_chunks: List[Any] = ( - [] - ) # for generating complete stream response + self.sync_streaming_chunks: List[ + Any + ] = [] # for generating complete stream response self.log_raw_request_response = log_raw_request_response # Initialize dynamic callbacks @@ -746,9 +746,9 @@ class Logging(LiteLLMLoggingBaseClass): prompt_spec=prompt_spec, dynamic_callback_params=dynamic_callback_params, ): - self.model_call_details["prompt_integration"] = ( - logger.__class__.__name__ - ) + self.model_call_details[ + "prompt_integration" + ] = logger.__class__.__name__ return logger except Exception: # If check fails, continue to next logger @@ -816,9 +816,9 @@ class Logging(LiteLLMLoggingBaseClass): if anthropic_cache_control_logger := AnthropicCacheControlHook.get_custom_logger_for_anthropic_cache_control_hook( non_default_params ): - self.model_call_details["prompt_integration"] = ( - anthropic_cache_control_logger.__class__.__name__ - ) + self.model_call_details[ + "prompt_integration" + ] = anthropic_cache_control_logger.__class__.__name__ return anthropic_cache_control_logger ######################################################### @@ -830,9 +830,9 @@ class Logging(LiteLLMLoggingBaseClass): internal_usage_cache=None, llm_router=None, ) - self.model_call_details["prompt_integration"] = ( - vector_store_custom_logger.__class__.__name__ - ) + self.model_call_details[ + "prompt_integration" + ] = vector_store_custom_logger.__class__.__name__ # Add to global callbacks so post-call hooks are invoked if ( vector_store_custom_logger @@ -892,9 +892,9 @@ class Logging(LiteLLMLoggingBaseClass): model ): # if model name was changes pre-call, overwrite the initial model call name with the new one self.model_call_details["model"] = model - self.model_call_details["litellm_params"]["api_base"] = ( - self._get_masked_api_base(additional_args.get("api_base", "")) - ) + self.model_call_details["litellm_params"][ + "api_base" + ] = self._get_masked_api_base(additional_args.get("api_base", "")) def pre_call(self, input, api_key, model=None, additional_args={}): # noqa: PLR0915 # Log the exact input to the LLM API @@ -923,10 +923,10 @@ class Logging(LiteLLMLoggingBaseClass): try: # [Non-blocking Extra Debug Information in metadata] if turn_off_message_logging is True: - _metadata["raw_request"] = ( - "redacted by litellm. \ + _metadata[ + "raw_request" + ] = "redacted by litellm. \ 'litellm.turn_off_message_logging=True'" - ) else: curl_command = self._get_request_curl_command( api_base=additional_args.get("api_base", ""), @@ -937,34 +937,34 @@ class Logging(LiteLLMLoggingBaseClass): _metadata["raw_request"] = str(curl_command) # split up, so it's easier to parse in the UI - self.model_call_details["raw_request_typed_dict"] = ( - RawRequestTypedDict( - raw_request_api_base=str( - additional_args.get("api_base") or "" - ), - raw_request_body=self._get_raw_request_body( - additional_args.get("complete_input_dict", {}) - ), - # NOTE: setting ignore_sensitive_headers to True will cause - # the Authorization header to be leaked when calls to the health - # endpoint are made and fail. - raw_request_headers=self._get_masked_headers( - additional_args.get("headers", {}) or {}, - ), - error=None, - ) + self.model_call_details[ + "raw_request_typed_dict" + ] = RawRequestTypedDict( + raw_request_api_base=str( + additional_args.get("api_base") or "" + ), + raw_request_body=self._get_raw_request_body( + additional_args.get("complete_input_dict", {}) + ), + # NOTE: setting ignore_sensitive_headers to True will cause + # the Authorization header to be leaked when calls to the health + # endpoint are made and fail. + raw_request_headers=self._get_masked_headers( + additional_args.get("headers", {}) or {}, + ), + error=None, ) except Exception as e: - self.model_call_details["raw_request_typed_dict"] = ( - RawRequestTypedDict( - error=str(e), - ) + self.model_call_details[ + "raw_request_typed_dict" + ] = RawRequestTypedDict( + error=str(e), ) - _metadata["raw_request"] = ( - "Unable to Log \ + _metadata[ + "raw_request" + ] = "Unable to Log \ raw request: {}".format( - str(e) - ) + str(e) ) if getattr(self, "logger_fn", None) and callable(self.logger_fn): try: @@ -1265,13 +1265,13 @@ class Logging(LiteLLMLoggingBaseClass): for callback in callbacks: try: if isinstance(callback, CustomLogger): - response: Optional[MCPPostCallResponseObject] = ( - await callback.async_post_mcp_tool_call_hook( - kwargs=kwargs, - response_obj=post_mcp_tool_call_response_obj, - start_time=start_time, - end_time=end_time, - ) + response: Optional[ + MCPPostCallResponseObject + ] = await callback.async_post_mcp_tool_call_hook( + kwargs=kwargs, + response_obj=post_mcp_tool_call_response_obj, + start_time=start_time, + end_time=end_time, ) ###################################################################### # if any of the callbacks modify the response, use the modified response @@ -1466,9 +1466,9 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.debug( f"response_cost_failure_debug_information: {debug_info}" ) - self.model_call_details["response_cost_failure_debug_information"] = ( - debug_info - ) + self.model_call_details[ + "response_cost_failure_debug_information" + ] = debug_info return None try: @@ -1494,9 +1494,9 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.debug( f"response_cost_failure_debug_information: {debug_info}" ) - self.model_call_details["response_cost_failure_debug_information"] = ( - debug_info - ) + self.model_call_details[ + "response_cost_failure_debug_information" + ] = debug_info return None @@ -1652,9 +1652,9 @@ class Logging(LiteLLMLoggingBaseClass): result=logging_result ) - self.model_call_details["standard_logging_object"] = ( - self._build_standard_logging_payload(logging_result, start_time, end_time) - ) + self.model_call_details[ + "standard_logging_object" + ] = self._build_standard_logging_payload(logging_result, start_time, end_time) if ( standard_logging_payload := self.model_call_details.get( @@ -1732,9 +1732,9 @@ class Logging(LiteLLMLoggingBaseClass): end_time = datetime.datetime.now() if self.completion_start_time is None: self.completion_start_time = end_time - self.model_call_details["completion_start_time"] = ( - self.completion_start_time - ) + self.model_call_details[ + "completion_start_time" + ] = self.completion_start_time self.model_call_details["log_event_type"] = "successful_api_call" self.model_call_details["end_time"] = end_time @@ -1771,10 +1771,10 @@ class Logging(LiteLLMLoggingBaseClass): end_time=end_time, ) elif isinstance(result, dict) or isinstance(result, list): - self.model_call_details["standard_logging_object"] = ( - self._build_standard_logging_payload( - result, start_time, end_time - ) + self.model_call_details[ + "standard_logging_object" + ] = self._build_standard_logging_payload( + result, start_time, end_time ) if ( standard_logging_payload := self.model_call_details.get( @@ -1783,9 +1783,9 @@ class Logging(LiteLLMLoggingBaseClass): ) is not None: emit_standard_logging_payload(standard_logging_payload) elif standard_logging_object is not None: - self.model_call_details["standard_logging_object"] = ( - standard_logging_object - ) + self.model_call_details[ + "standard_logging_object" + ] = standard_logging_object else: self.model_call_details["response_cost"] = None @@ -1943,17 +1943,17 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.debug( "Logging Details LiteLLM-Success Call streaming complete" ) - self.model_call_details["complete_streaming_response"] = ( - complete_streaming_response - ) - self.model_call_details["response_cost"] = ( - self._response_cost_calculator(result=complete_streaming_response) - ) + self.model_call_details[ + "complete_streaming_response" + ] = complete_streaming_response + self.model_call_details[ + "response_cost" + ] = self._response_cost_calculator(result=complete_streaming_response) ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details["standard_logging_object"] = ( - self._build_standard_logging_payload( - complete_streaming_response, start_time, end_time - ) + self.model_call_details[ + "standard_logging_object" + ] = self._build_standard_logging_payload( + complete_streaming_response, start_time, end_time ) if ( standard_logging_payload := self.model_call_details.get( @@ -2287,10 +2287,10 @@ class Logging(LiteLLMLoggingBaseClass): ) else: if self.stream and complete_streaming_response: - self.model_call_details["complete_response"] = ( - self.model_call_details.get( - "complete_streaming_response", {} - ) + self.model_call_details[ + "complete_response" + ] = self.model_call_details.get( + "complete_streaming_response", {} ) result = self.model_call_details["complete_response"] openMeterLogger.log_success_event( @@ -2314,10 +2314,10 @@ class Logging(LiteLLMLoggingBaseClass): ) else: if self.stream and complete_streaming_response: - self.model_call_details["complete_response"] = ( - self.model_call_details.get( - "complete_streaming_response", {} - ) + self.model_call_details[ + "complete_response" + ] = self.model_call_details.get( + "complete_streaming_response", {} ) result = self.model_call_details["complete_response"] @@ -2456,9 +2456,9 @@ class Logging(LiteLLMLoggingBaseClass): if complete_streaming_response is not None: print_verbose("Async success callbacks: Got a complete streaming response") - self.model_call_details["async_complete_streaming_response"] = ( - complete_streaming_response - ) + self.model_call_details[ + "async_complete_streaming_response" + ] = complete_streaming_response try: if self.model_call_details.get("cache_hit", False) is True: @@ -2469,10 +2469,10 @@ class Logging(LiteLLMLoggingBaseClass): model_call_details=self.model_call_details ) # base_model defaults to None if not set on model_info - self.model_call_details["response_cost"] = ( - self._response_cost_calculator( - result=complete_streaming_response - ) + self.model_call_details[ + "response_cost" + ] = self._response_cost_calculator( + result=complete_streaming_response ) verbose_logger.debug( @@ -2485,10 +2485,10 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["response_cost"] = None ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details["standard_logging_object"] = ( - self._build_standard_logging_payload( - complete_streaming_response, start_time, end_time - ) + self.model_call_details[ + "standard_logging_object" + ] = self._build_standard_logging_payload( + complete_streaming_response, start_time, end_time ) # print standard logging payload @@ -2515,9 +2515,9 @@ class Logging(LiteLLMLoggingBaseClass): # _success_handler_helper_fn if self.model_call_details.get("standard_logging_object") is None: ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details["standard_logging_object"] = ( - self._build_standard_logging_payload(result, start_time, end_time) - ) + self.model_call_details[ + "standard_logging_object" + ] = self._build_standard_logging_payload(result, start_time, end_time) # print standard logging payload if ( @@ -2760,18 +2760,18 @@ class Logging(LiteLLMLoggingBaseClass): ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details["standard_logging_object"] = ( - get_standard_logging_object_payload( - kwargs=self.model_call_details, - init_response_obj={}, - start_time=start_time, - end_time=end_time, - logging_obj=self, - status="failure", - error_str=str(exception), - original_exception=exception, - standard_built_in_tools_params=self.standard_built_in_tools_params, - ) + self.model_call_details[ + "standard_logging_object" + ] = get_standard_logging_object_payload( + kwargs=self.model_call_details, + init_response_obj={}, + start_time=start_time, + end_time=end_time, + logging_obj=self, + status="failure", + error_str=str(exception), + original_exception=exception, + standard_built_in_tools_params=self.standard_built_in_tools_params, ) return start_time, end_time @@ -3735,9 +3735,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 service_name=arize_config.project_name, ) - os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = ( - f"space_id={arize_config.space_key or arize_config.space_id},api_key={arize_config.api_key}" - ) + os.environ[ + "OTEL_EXPORTER_OTLP_TRACES_HEADERS" + ] = f"space_id={arize_config.space_key or arize_config.space_id},api_key={arize_config.api_key}" for callback in _in_memory_loggers: if ( isinstance(callback, ArizeLogger) @@ -3763,13 +3763,13 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 existing_attrs = os.environ.get("OTEL_RESOURCE_ATTRIBUTES", "") # Add openinference.project.name attribute if existing_attrs: - os.environ["OTEL_RESOURCE_ATTRIBUTES"] = ( - f"{existing_attrs},openinference.project.name={arize_phoenix_config.project_name}" - ) + os.environ[ + "OTEL_RESOURCE_ATTRIBUTES" + ] = f"{existing_attrs},openinference.project.name={arize_phoenix_config.project_name}" else: - os.environ["OTEL_RESOURCE_ATTRIBUTES"] = ( - f"openinference.project.name={arize_phoenix_config.project_name}" - ) + os.environ[ + "OTEL_RESOURCE_ATTRIBUTES" + ] = f"openinference.project.name={arize_phoenix_config.project_name}" # Set Phoenix project name from environment variable phoenix_project_name = os.environ.get("PHOENIX_PROJECT_NAME", None) @@ -3777,19 +3777,19 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 existing_attrs = os.environ.get("OTEL_RESOURCE_ATTRIBUTES", "") # Add openinference.project.name attribute if existing_attrs: - os.environ["OTEL_RESOURCE_ATTRIBUTES"] = ( - f"{existing_attrs},openinference.project.name={phoenix_project_name}" - ) + os.environ[ + "OTEL_RESOURCE_ATTRIBUTES" + ] = f"{existing_attrs},openinference.project.name={phoenix_project_name}" else: - os.environ["OTEL_RESOURCE_ATTRIBUTES"] = ( - f"openinference.project.name={phoenix_project_name}" - ) + os.environ[ + "OTEL_RESOURCE_ATTRIBUTES" + ] = f"openinference.project.name={phoenix_project_name}" # auth can be disabled on local deployments of arize phoenix if arize_phoenix_config.otlp_auth_headers is not None: - os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = ( - arize_phoenix_config.otlp_auth_headers - ) + os.environ[ + "OTEL_EXPORTER_OTLP_TRACES_HEADERS" + ] = arize_phoenix_config.otlp_auth_headers for callback in _in_memory_loggers: if ( @@ -3965,9 +3965,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 exporter="otlp_http", endpoint="https://langtrace.ai/api/trace", ) - os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = ( - f"api_key={os.getenv('LANGTRACE_API_KEY')}" - ) + os.environ[ + "OTEL_EXPORTER_OTLP_TRACES_HEADERS" + ] = f"api_key={os.getenv('LANGTRACE_API_KEY')}" for callback in _in_memory_loggers: if ( isinstance(callback, OpenTelemetry) @@ -4881,10 +4881,10 @@ class StandardLoggingPayloadSetup: for key in StandardLoggingHiddenParams.__annotations__.keys(): if key in hidden_params: if key == "additional_headers": - clean_hidden_params["additional_headers"] = ( - StandardLoggingPayloadSetup.get_additional_headers( - hidden_params[key] - ) + clean_hidden_params[ + "additional_headers" + ] = StandardLoggingPayloadSetup.get_additional_headers( + hidden_params[key] ) else: clean_hidden_params[key] = hidden_params[key] # type: ignore @@ -5036,7 +5036,6 @@ class StandardLoggingPayloadSetup: dynamic_litellm_session_id = litellm_params.get("litellm_session_id") dynamic_litellm_trace_id = litellm_params.get("litellm_trace_id") - # Note: we recommend using `litellm_session_id` for session tracking # `litellm_trace_id` is an internal litellm param if dynamic_litellm_session_id: @@ -5355,7 +5354,10 @@ def get_standard_logging_object_payload( requested_model = kwargs.get("model") if ( isinstance(requested_model, str) - and ("model_router" in requested_model.lower() or "model-router" in requested_model.lower()) + and ( + "model_router" in requested_model.lower() + or "model-router" in requested_model.lower() + ) and isinstance(response_model_name, str) and response_model_name ): @@ -5521,9 +5523,9 @@ def scrub_sensitive_keys_in_metadata(litellm_params: Optional[dict]): ): for k, v in metadata["user_api_key_metadata"].items(): if k == "logging": # prevent logging user logging keys - cleaned_user_api_key_metadata[k] = ( - "scrubbed_by_litellm_for_sensitive_keys" - ) + cleaned_user_api_key_metadata[ + k + ] = "scrubbed_by_litellm_for_sensitive_keys" else: cleaned_user_api_key_metadata[k] = v diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py index 4a4a2508d2..4454fca3b0 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py @@ -233,9 +233,10 @@ class StandardBuiltInToolCostTracking: return 0.0 model_info_dict = dict(model_info) if model_info is not None else None - input_tokens, output_tokens = ( - StandardBuiltInToolCostTracking._extract_token_counts(computer_use_usage) - ) + ( + input_tokens, + output_tokens, + ) = StandardBuiltInToolCostTracking._extract_token_counts(computer_use_usage) return StandardBuiltInToolCostTracking.get_cost_for_computer_use( input_tokens=input_tokens, @@ -314,8 +315,10 @@ class StandardBuiltInToolCostTracking: if isinstance(response_object, ModelResponse): # chat completions only include url_citation annotations when a web search call is made - has_url_citations = StandardBuiltInToolCostTracking.response_includes_annotation_type( - response_object=response_object, annotation_type="url_citation" + has_url_citations = ( + StandardBuiltInToolCostTracking.response_includes_annotation_type( + response_object=response_object, annotation_type="url_citation" + ) ) if has_url_citations: return True @@ -325,7 +328,9 @@ class StandardBuiltInToolCostTracking: if ( hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details is not None - and isinstance(usage.prompt_tokens_details, PromptTokensDetailsWrapper) + and isinstance( + usage.prompt_tokens_details, PromptTokensDetailsWrapper + ) and hasattr(usage.prompt_tokens_details, "web_search_requests") and usage.prompt_tokens_details.web_search_requests is not None ): @@ -468,7 +473,9 @@ class StandardBuiltInToolCostTracking: """ if model_info is None: return 0.0 - search_context_raw: Any = model_info.get("search_context_cost_per_query", {}) or {} + search_context_raw: Any = ( + model_info.get("search_context_cost_per_query", {}) or {} + ) search_context_pricing: SearchContextCostPerQuery = ( SearchContextCostPerQuery(**search_context_raw) if search_context_raw @@ -603,21 +610,26 @@ class StandardBuiltInToolCostTracking: Get code interpreter cost per session from model cost map. """ import litellm - + try: container_model = f"{provider}/container" model_info = litellm.get_model_info( - model=container_model, - custom_llm_provider=provider + model=container_model, custom_llm_provider=provider ) - model_key = model_info.get("key") if isinstance(model_info, dict) else getattr(model_info, "key", None) - + model_key = ( + model_info.get("key") + if isinstance(model_info, dict) + else getattr(model_info, "key", None) + ) + if model_key and model_key in litellm.model_cost: - return litellm.model_cost[model_key].get("code_interpreter_cost_per_session") - + return litellm.model_cost[model_key].get( + "code_interpreter_cost_per_session" + ) + except Exception: pass - + return None @staticmethod @@ -646,7 +658,6 @@ class StandardBuiltInToolCostTracking: ) if cost_per_session is not None: return sessions * cost_per_session - return 0.0 diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index bf0b270936..191231f3e6 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -19,13 +19,15 @@ from litellm.types.utils import ( from litellm.utils import get_model_info # Pre-resolved CallTypes enum values for fast membership checks -_IMAGE_RESPONSE_CALL_TYPES = frozenset({ - CallTypes.image_generation.value, - CallTypes.aimage_generation.value, - PassthroughCallTypes.passthrough_image_generation.value, - CallTypes.image_edit.value, - CallTypes.aimage_edit.value, -}) +_IMAGE_RESPONSE_CALL_TYPES = frozenset( + { + CallTypes.image_generation.value, + CallTypes.aimage_generation.value, + PassthroughCallTypes.passthrough_image_generation.value, + CallTypes.image_edit.value, + CallTypes.aimage_edit.value, + } +) def _is_above_128k(tokens: float) -> bool: @@ -245,7 +247,10 @@ def _get_token_base_cost( else key ) prompt_base_cost = cast( - float, _get_cost_per_unit(model_info, tiered_input_key, prompt_base_cost) + float, + _get_cost_per_unit( + model_info, tiered_input_key, prompt_base_cost + ), ) tiered_output_key = ( _get_service_tier_cost_key( @@ -268,9 +273,7 @@ def _get_token_base_cost( cache_creation_tiered_key = ( f"cache_creation_input_token_cost_above_{threshold_str}_tokens" ) - cache_creation_1hr_tiered_key = ( - f"cache_creation_input_token_cost_above_1hr_above_{threshold_str}_tokens" - ) + cache_creation_1hr_tiered_key = f"cache_creation_input_token_cost_above_1hr_above_{threshold_str}_tokens" cache_read_tiered_key = ( f"cache_read_input_token_cost_above_{threshold_str}_tokens" ) @@ -576,7 +579,10 @@ def _calculate_input_cost( ) ### CACHE WRITING COST - Now uses tiered pricing - if prompt_tokens_details["cache_creation_tokens"] or prompt_tokens_details["cache_creation_token_details"] is not None: + if ( + prompt_tokens_details["cache_creation_tokens"] + or prompt_tokens_details["cache_creation_token_details"] is not None + ): prompt_cost += calculate_cache_writing_cost( cache_creation_tokens=prompt_tokens_details["cache_creation_tokens"], cache_creation_token_details=prompt_tokens_details[ @@ -589,7 +595,9 @@ def _calculate_input_cost( ### CHARACTER COST if prompt_tokens_details["character_count"]: prompt_cost += calculate_cost_component( - model_info, "input_cost_per_character", prompt_tokens_details["character_count"] + model_info, + "input_cost_per_character", + prompt_tokens_details["character_count"], ) ### IMAGE COUNT COST @@ -661,10 +669,14 @@ def generic_cost_per_token( # noqa: PLR0915 image_tokens = prompt_tokens_details["image_tokens"] # Check for double-counting: sum of details > prompt_tokens means overlap - total_details = text_tokens + cache_hit + audio_tokens + cache_creation + image_tokens + total_details = ( + text_tokens + cache_hit + audio_tokens + cache_creation + image_tokens + ) has_double_counting = cache_hit > 0 and total_details > usage.prompt_tokens - if (text_tokens == 0 and prompt_tokens_details["image_count"] == 0) or has_double_counting: + if ( + text_tokens == 0 and prompt_tokens_details["image_count"] == 0 + ) or has_double_counting: text_tokens = ( usage.prompt_tokens - cache_hit diff --git a/litellm/litellm_core_utils/llm_request_utils.py b/litellm/litellm_core_utils/llm_request_utils.py index 89f5728979..a2292d6e00 100644 --- a/litellm/litellm_core_utils/llm_request_utils.py +++ b/litellm/litellm_core_utils/llm_request_utils.py @@ -67,6 +67,7 @@ def pick_cheapest_chat_models_from_llm_provider(custom_llm_provider: str, n=1): # Return the top n cheapest models return [model for model, _ in model_costs[:n]] + def get_proxy_server_request_headers(litellm_params: Optional[dict]) -> dict: """ Get the `proxy_server_request` headers from the litellm_params.\ @@ -80,4 +81,4 @@ def get_proxy_server_request_headers(litellm_params: Optional[dict]) -> dict: litellm_params.get("proxy_server_request", {}).get("headers", {}) or {} ) - return proxy_request_headers \ No newline at end of file + return proxy_request_headers diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index 4bc9f0c835..20cc574666 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -471,7 +471,7 @@ def convert_to_model_response_object( # noqa: PLR0915 if hidden_params is None: hidden_params = {} - + # Preserve existing additional_headers if they contain important provider headers # For responses API, additional_headers may already be set with LLM provider headers existing_additional_headers = hidden_params.get("additional_headers", {}) @@ -482,7 +482,7 @@ def convert_to_model_response_object( # noqa: PLR0915 # Merge new headers with existing ones if existing_additional_headers: additional_headers.update(existing_additional_headers) - + hidden_params["additional_headers"] = additional_headers ### CHECK IF ERROR IN RESPONSE ### - openrouter returns these in the dictionary @@ -596,9 +596,9 @@ def convert_to_model_response_object( # noqa: PLR0915 provider_specific_fields["thinking_blocks"] = thinking_blocks if reasoning_content: - provider_specific_fields["reasoning_content"] = ( - reasoning_content - ) + provider_specific_fields[ + "reasoning_content" + ] = reasoning_content message = Message( content=content, @@ -654,7 +654,9 @@ def convert_to_model_response_object( # noqa: PLR0915 if "id" in response_object: # Preserve the auto-generated id from ModelResponse.__init__ # when the provider returns a falsy id (None, "") - model_response_object.id = response_object["id"] or model_response_object.id + model_response_object.id = ( + response_object["id"] or model_response_object.id + ) if "system_fingerprint" in response_object: model_response_object.system_fingerprint = response_object[ @@ -785,7 +787,9 @@ def convert_to_model_response_object( # noqa: PLR0915 # tracking without exposing it in the response body. Must be set # after hidden_params assignment to avoid being overwritten. if "_audio_transcription_duration" in response_object: - model_response_object._hidden_params["audio_transcription_duration"] = response_object["_audio_transcription_duration"] + model_response_object._hidden_params[ + "audio_transcription_duration" + ] = response_object["_audio_transcription_duration"] if _response_headers is not None: model_response_object._response_headers = _response_headers diff --git a/litellm/litellm_core_utils/logging_callback_manager.py b/litellm/litellm_core_utils/logging_callback_manager.py index 38da11e777..c5c150274c 100644 --- a/litellm/litellm_core_utils/logging_callback_manager.py +++ b/litellm/litellm_core_utils/logging_callback_manager.py @@ -150,11 +150,11 @@ class LoggingCallbackManager: def remove_callbacks_by_type(self, callback_list, callback_type): """ Remove all callbacks of a specific type from a callback list. - + Args: callback_list: The list to remove callbacks from (e.g., litellm.callbacks) callback_type: The class type to match (e.g., SemanticToolFilterHook) - + Example: litellm.logging_callback_manager.remove_callbacks_by_type( litellm.callbacks, SemanticToolFilterHook diff --git a/litellm/litellm_core_utils/logging_worker.py b/litellm/litellm_core_utils/logging_worker.py index d5eca9eeb5..7f00c47c1f 100644 --- a/litellm/litellm_core_utils/logging_worker.py +++ b/litellm/litellm_core_utils/logging_worker.py @@ -417,26 +417,30 @@ class LoggingWorker: """ # Check if logger has valid handlers before attempting to log # During shutdown, handlers may be closed, causing ValueError when writing - if not hasattr(verbose_logger, 'handlers') or not verbose_logger.handlers: + if not hasattr(verbose_logger, "handlers") or not verbose_logger.handlers: return - + # Check if any handler has a valid stream has_valid_handler = False for handler in verbose_logger.handlers: try: - if hasattr(handler, 'stream') and handler.stream and not handler.stream.closed: + if ( + hasattr(handler, "stream") + and handler.stream + and not handler.stream.closed + ): has_valid_handler = True break - elif not hasattr(handler, 'stream'): + elif not hasattr(handler, "stream"): # Non-stream handlers (like NullHandler) are always valid has_valid_handler = True break except (AttributeError, ValueError): continue - + if not has_valid_handler: return - + try: if level == "debug": verbose_logger.debug(message) diff --git a/litellm/litellm_core_utils/model_param_helper.py b/litellm/litellm_core_utils/model_param_helper.py index 4d45c47c22..66b174feac 100644 --- a/litellm/litellm_core_utils/model_param_helper.py +++ b/litellm/litellm_core_utils/model_param_helper.py @@ -93,9 +93,9 @@ class ModelParamHelper: streaming_params: Set[str] = set( getattr(CompletionCreateParamsStreaming, "__annotations__", {}).keys() ) - litellm_provider_specific_params: Set[str] = ( - ModelParamHelper.get_litellm_provider_specific_params_for_chat_params() - ) + litellm_provider_specific_params: Set[ + str + ] = ModelParamHelper.get_litellm_provider_specific_params_for_chat_params() all_chat_completion_kwargs: Set[str] = non_streaming_params.union( streaming_params ).union(litellm_provider_specific_params) diff --git a/litellm/litellm_core_utils/model_response_utils.py b/litellm/litellm_core_utils/model_response_utils.py index 00462221fe..6c290fa30c 100644 --- a/litellm/litellm_core_utils/model_response_utils.py +++ b/litellm/litellm_core_utils/model_response_utils.py @@ -114,23 +114,19 @@ def _is_choice_non_empty(choice: Any) -> bool: """ # Check finish_reason if hasattr(choice, "finish_reason") and choice.finish_reason is not None: - return True # Check logprobs if hasattr(choice, "logprobs") and choice.logprobs is not None: - return True # Check enhancements (if present) if hasattr(choice, "enhancements") and choice.enhancements is not None: - return True # Deep check delta object if hasattr(choice, "delta") and choice.delta is not None: if _is_delta_non_empty(choice.delta): - return True # Check model_extra for dynamically added fields on the choice @@ -138,19 +134,15 @@ def _is_choice_non_empty(choice: Any) -> bool: for extra_field_name, extra_field_value in choice.model_extra.items(): # Skip certain structural fields that are just default/None placeholders if extra_field_name == "index" and extra_field_value == 0: - continue if ( extra_field_name in {"finish_reason", "logprobs"} and extra_field_value is None ): - continue if extra_field_name == "delta": - continue if _has_meaningful_content(extra_field_value): - return True # Check for any other non-standard fields on the choice @@ -169,12 +161,10 @@ def _is_choice_non_empty(choice: Any) -> bool: "enhancements", } ): - continue attr_value = getattr(choice, attr_name, None) if _has_meaningful_content(attr_value): - return True return False @@ -195,7 +185,6 @@ def _is_delta_non_empty(delta: Delta) -> bool: for extra_field_name, extra_field_value in delta.model_extra.items(): # Even structural fields are meaningful if they have actual content if _has_meaningful_content(extra_field_value): - return True # Check all regular attributes of the delta object @@ -210,7 +199,6 @@ def _is_delta_non_empty(delta: Delta) -> bool: attr_value = getattr(delta, attr_name, None) if _has_meaningful_content(attr_value): - return True return False diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index d59b8d8871..a5d6bc936b 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -644,6 +644,10 @@ def extract_file_data(file_data: FileTypes) -> ExtractedFileData: content = f.read() elif isinstance(file_content, io.IOBase): # If it's a file-like object + # Try to get filename from file handle if not already set + if not filename and hasattr(file_content, "name"): + filename = Path(file_content.name).name + content = file_content.read() if isinstance(content, str): diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 610e3a368e..0df10ae9f9 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1037,8 +1037,7 @@ def convert_to_anthropic_tool_invoke_xml(tool_calls: list) -> str: ) if isinstance(parsed_args, dict): parameters = "".join( - f"<{param}>{val}\n" - for param, val in parsed_args.items() + f"<{param}>{val}\n" for param, val in parsed_args.items() ) else: parameters = f"{parsed_args}\n" @@ -1394,10 +1393,10 @@ def convert_to_gemini_tool_call_invoke( if tool_calls is not None: for idx, tool in enumerate(tool_calls): if "function" in tool: - gemini_function_call: Optional[VertexFunctionCall] = ( - _gemini_tool_call_invoke_helper( - function_call_params=tool["function"] - ) + gemini_function_call: Optional[ + VertexFunctionCall + ] = _gemini_tool_call_invoke_helper( + function_call_params=tool["function"] ) if gemini_function_call is not None: part_dict: VertexPartType = { @@ -1705,7 +1704,9 @@ def convert_to_anthropic_tool_result( anthropic_content_element=_anthropic_image_param, original_content_element=content, ) - anthropic_content_list.append(cast(AnthropicMessagesImageParam, _anthropic_image_param)) + anthropic_content_list.append( + cast(AnthropicMessagesImageParam, _anthropic_image_param) + ) anthropic_content = anthropic_content_list anthropic_tool_result: Optional[AnthropicMessagesToolResultParam] = None @@ -2036,7 +2037,7 @@ def _sanitize_empty_text_content( """ Case C: Sanitize empty text content - Replace empty or whitespace-only text content with a placeholder message. - + Returns: The message with sanitized content if needed, otherwise the original message """ @@ -2045,14 +2046,16 @@ def _sanitize_empty_text_content( if isinstance(content, str): if not content or not content.strip(): message = cast(AllMessageValues, dict(message)) # Make a copy - message["content"] = "[System: Empty message content sanitised to satisfy protocol]" + message[ + "content" + ] = "[System: Empty message content sanitised to satisfy protocol]" verbose_logger.debug( f"_sanitize_empty_text_content: Replaced empty text content in {message.get('role')} message" ) return message -def _add_missing_tool_results( # noqa: PLR0915 +def _add_missing_tool_results( # noqa: PLR0915 current_message: AllMessageValues, messages: List[AllMessageValues], current_index: int, @@ -2084,40 +2087,40 @@ def _add_missing_tool_results( # noqa: PLR0915 tool_call_id = getattr(tool_call, "id", None) if tool_call_id: expected_tool_call_ids.add(tool_call_id) - + # Collect actual tool result messages that follow this assistant message found_tool_call_ids = set() actual_tool_results: List[AllMessageValues] = [] j = current_index + 1 - + while j < len(messages): next_msg = messages[j] next_role = next_msg.get("role") - + if next_role == "assistant": break - + if next_role in ["tool", "function"]: tool_call_id = next_msg.get("tool_call_id") if tool_call_id and tool_call_id in expected_tool_call_ids: found_tool_call_ids.add(tool_call_id) actual_tool_results.append(next_msg) - + j += 1 - + # Find missing tool results missing_tool_call_ids = expected_tool_call_ids - found_tool_call_ids - + if missing_tool_call_ids: verbose_logger.debug( f"_add_missing_tool_results: Found {len(missing_tool_call_ids)} orphaned tool calls. Adding dummy tool results." ) - + result_messages.append(current_message) - + # Add existing tool results FIRST result_messages.extend(actual_tool_results) - + # Then add dummy tool results for missing ones for tool_call_id in missing_tool_call_ids: tool_name = "unknown_tool" @@ -2127,7 +2130,7 @@ def _add_missing_tool_results( # noqa: PLR0915 tc_id = tool_call.get("id") else: tc_id = getattr(tool_call, "id", None) - + if tc_id == tool_call_id: if isinstance(tool_call, dict): function = tool_call.get("function", {}) @@ -2140,17 +2143,17 @@ def _add_missing_tool_results( # noqa: PLR0915 if function: tool_name = getattr(function, "name", "unknown_tool") break - + dummy_tool_result: ChatCompletionToolMessage = { "role": "tool", "tool_call_id": tool_call_id, "content": f"[System: Tool execution skipped/interrupted by user. No result provided for tool '{tool_name}'.]", } result_messages.append(dummy_tool_result) - + # Return the messages and the number of original messages to skip return (result_messages, len(actual_tool_results)) - + return ([current_message], 0) @@ -2162,21 +2165,21 @@ def _is_orphaned_tool_result( Case B: Orphaned tool_result (unexpected result) - Check if a tool message references a tool_call_id that doesn't exist in the previous assistant message. - + Returns: True if this is an orphaned tool result that should be removed, False otherwise """ if current_message.get("role") not in ["tool", "function"]: return False - + tool_call_id = current_message.get("tool_call_id") - + if not tool_call_id: return False - + # Look back to find the most recent assistant message with tool_calls found_matching_tool_call = False - + for j in range(len(sanitized_messages) - 1, -1, -1): prev_msg = sanitized_messages[j] if prev_msg.get("role") == "assistant": @@ -2188,19 +2191,19 @@ def _is_orphaned_tool_result( tc_id = tool_call.get("id") else: tc_id = getattr(tool_call, "id", None) - + if tc_id == tool_call_id: found_matching_tool_call = True break - + break - + if not found_matching_tool_call: verbose_logger.debug( "_is_orphaned_tool_result: Found orphaned tool result with redacted tool_call_id" ) return True - + return False @@ -2209,58 +2212,60 @@ def sanitize_messages_for_tool_calling( ) -> List[AllMessageValues]: """ Sanitize messages for tool calling to handle common issues when modify_params=True: - + Case A: Missing tool_result for tool_use (orphaned tool calls) - If an assistant message has tool_calls but no corresponding tool result follows, add a dummy tool result message indicating the user did not provide the result. - + Case B: Orphaned tool_result (unexpected result) - If a tool message references a tool_call_id that doesn't exist in the previous assistant message, remove that tool message. - + Case C: Empty text content - Replace empty or whitespace-only text content with a placeholder message. - + Case D: Duplicate tool_result for same tool_use (duplicate results) - If multiple tool messages reference the same tool_call_id, keep only the last occurrence. Anthropic requires exactly one tool_result per tool_use and rejects with: "each tool_use must have a single result". - + This function operates on OpenAI format messages before they are converted to provider-specific formats. """ if not litellm.modify_params: return messages - + sanitized_messages: List[AllMessageValues] = [] i = 0 - + while i < len(messages): current_message = messages[i] - + # Case C: Sanitize empty text content current_message = _sanitize_empty_text_content(current_message) - + # Case A: Check if assistant message has tool_calls without following tool results if current_message.get("role") == "assistant": - result_messages, messages_consumed = _add_missing_tool_results(current_message, messages, i) - + result_messages, messages_consumed = _add_missing_tool_results( + current_message, messages, i + ) + # If dummy tool results were added, extend sanitized_messages and skip consumed messages if len(result_messages) > 1: sanitized_messages.extend(result_messages) # Skip the assistant message and any actual tool results that were included i += 1 + messages_consumed continue - + # Case B: Check for orphaned tool results if _is_orphaned_tool_result(current_message, sanitized_messages): i += 1 continue # Skip this orphaned tool result - + # Add the message to sanitized list sanitized_messages.append(current_message) i += 1 - + # Case D: Deduplicate tool results with the same tool_call_id. # Anthropic requires exactly one tool_result per tool_use. Session history # (e.g. from conversation resume) can contain duplicate tool_result messages @@ -2328,7 +2333,7 @@ def anthropic_messages_pt( # noqa: PLR0915 """ # Sanitize messages for tool calling issues when modify_params=True messages = sanitize_messages_for_tool_calling(messages) - + # add role=tool support to allow function call result/error submission user_message_types = {"user", "tool", "function"} # reformat messages to ensure user/assistant are alternating, if there's either 2 consecutive 'user' messages or 2 consecutive 'assistant' message, merge them. @@ -2383,9 +2388,9 @@ def anthropic_messages_pt( # noqa: PLR0915 # Convert ChatCompletionImageUrlObject to dict if needed image_url_value = m["image_url"] if isinstance(image_url_value, str): - image_url_input: Union[str, dict[str, Any]] = ( - image_url_value - ) + image_url_input: Union[ + str, dict[str, Any] + ] = image_url_value else: # ChatCompletionImageUrlObject or dict case - convert to dict image_url_input = { @@ -2412,9 +2417,9 @@ def anthropic_messages_pt( # noqa: PLR0915 ) if "cache_control" in _content_element: - _anthropic_content_element["cache_control"] = ( - _content_element["cache_control"] - ) + _anthropic_content_element[ + "cache_control" + ] = _content_element["cache_control"] user_content.append(_anthropic_content_element) elif m.get("type", "") == "text": m = cast(ChatCompletionTextObject, m) @@ -2452,9 +2457,9 @@ def anthropic_messages_pt( # noqa: PLR0915 ) if "cache_control" in _content_element: - _anthropic_content_text_element["cache_control"] = ( - _content_element["cache_control"] - ) + _anthropic_content_text_element[ + "cache_control" + ] = _content_element["cache_control"] user_content.append(_anthropic_content_text_element) @@ -2487,7 +2492,9 @@ def anthropic_messages_pt( # noqa: PLR0915 "provider_specific_fields" ) if isinstance(_provider_specific_fields_raw, dict): - _compaction_blocks = _provider_specific_fields_raw.get("compaction_blocks") + _compaction_blocks = _provider_specific_fields_raw.get( + "compaction_blocks" + ) if _compaction_blocks and isinstance(_compaction_blocks, list): # Add compaction blocks at the beginning of assistant content : https://platform.claude.com/docs/en/build-with-claude/compaction assistant_content.extend(_compaction_blocks) # type: ignore @@ -2507,7 +2514,11 @@ def anthropic_messages_pt( # noqa: PLR0915 if isinstance(_tc, dict) else getattr(_tc, "id", None) ) - if _tc_id and isinstance(_tc_id, str) and _tc_id.startswith("srvtoolu_"): + if ( + _tc_id + and isinstance(_tc_id, str) + and _tc_id.startswith("srvtoolu_") + ): _has_server_tool_calls = True break @@ -2583,9 +2594,9 @@ def anthropic_messages_pt( # noqa: PLR0915 original_content_element=dict(assistant_content_block), ) if "cache_control" in _content_element: - _anthropic_text_content_element["cache_control"] = ( - _content_element["cache_control"] - ) + _anthropic_text_content_element[ + "cache_control" + ] = _content_element["cache_control"] text_element = _anthropic_text_content_element # Interleave: each thinking block precedes its server tool group. @@ -2674,13 +2685,15 @@ def anthropic_messages_pt( # noqa: PLR0915 _list_has_thinking = False if _content_is_list: for _item in assistant_content_block["content"]: - if isinstance(_item, dict) and _item.get("type") in ("thinking", "redacted_thinking"): + if isinstance(_item, dict) and _item.get("type") in ( + "thinking", + "redacted_thinking", + ): _list_has_thinking = True break if ( - thinking_blocks is not None - and not _list_has_thinking + thinking_blocks is not None and not _list_has_thinking ): # IMPORTANT: ADD THIS FIRST, ELSE ANTHROPIC WILL RAISE AN ERROR assistant_content.extend(thinking_blocks) if _content_is_list: @@ -2738,9 +2751,9 @@ def anthropic_messages_pt( # noqa: PLR0915 ) if "cache_control" in _content_element: - _anthropic_text_content_element["cache_control"] = _content_element[ + _anthropic_text_content_element[ "cache_control" - ] + ] = _content_element["cache_control"] assistant_content.append(_anthropic_text_content_element) @@ -3795,16 +3808,12 @@ def _convert_to_bedrock_tool_call_invoke( # '{"cmd":"a"}{"cmd":"b"}{"cmd":"c"}' # Split them and emit one toolUse block per object. # Fixes: https://github.com/BerriAI/litellm/issues/20543 - parsed_objects = split_concatenated_json_objects( - arguments - ) + parsed_objects = split_concatenated_json_objects(arguments) if parsed_objects: # First object keeps the original tool id. for obj_idx, obj in enumerate(parsed_objects): block_id = ( - tool_id - if obj_idx == 0 - else f"{tool_id}_{obj_idx}" + tool_id if obj_idx == 0 else f"{tool_id}_{obj_idx}" ) bedrock_tool = BedrockToolUseBlock( input=obj, name=name, toolUseId=block_id @@ -3817,9 +3826,7 @@ def _convert_to_bedrock_tool_call_invoke( if tool.get("cache_control", None) is not None: _parts_list.append( BedrockContentBlock( - cachePoint=CachePointBlock( - type="default" - ) + cachePoint=CachePointBlock(type="default") ) ) continue @@ -4572,7 +4579,9 @@ class BedrockConverseMessagesProcessor: msg_i += 1 - assistant_content = _deduplicate_bedrock_content_blocks(assistant_content, "toolUse") + assistant_content = _deduplicate_bedrock_content_blocks( + assistant_content, "toolUse" + ) if assistant_content: contents.append( @@ -4888,7 +4897,9 @@ def _bedrock_converse_messages_pt( # noqa: PLR0915 # AWS Bedrock doesn't allow empty or whitespace-only text content # Skip completely empty strings to avoid blank content blocks if element.get("text", "").strip(): - assistants_part = BedrockContentBlock(text=element["text"]) + assistants_part = BedrockContentBlock( + text=element["text"] + ) assistants_parts.append(assistants_part) elif element["type"] == "image_url": if isinstance(element["image_url"], dict): @@ -4914,7 +4925,9 @@ def _bedrock_converse_messages_pt( # noqa: PLR0915 elif _assistant_content is not None and isinstance(_assistant_content, str): # Skip completely empty strings to avoid blank content blocks if _assistant_content.strip(): - assistant_content.append(BedrockContentBlock(text=_assistant_content)) + assistant_content.append( + BedrockContentBlock(text=_assistant_content) + ) # Add cache point block for assistant string content _cache_point_block = ( litellm.AmazonConverseConfig()._get_cache_point_block( @@ -4931,7 +4944,9 @@ def _bedrock_converse_messages_pt( # noqa: PLR0915 msg_i += 1 - assistant_content = _deduplicate_bedrock_content_blocks(assistant_content, "toolUse") + assistant_content = _deduplicate_bedrock_content_blocks( + assistant_content, "toolUse" + ) if assistant_content: contents.append( @@ -4995,18 +5010,18 @@ def add_cache_point_tool_block(tool: dict) -> Optional[BedrockToolBlock]: def _is_bedrock_tool_block(tool: dict) -> bool: """ Check if a tool is already a BedrockToolBlock. - + BedrockToolBlock has one of: systemTool, toolSpec, or cachePoint. This is used to detect tools that are already in Bedrock format (e.g., systemTool for Nova grounding) vs OpenAI-style function tools that need transformation. - + Args: tool: The tool dict to check - + Returns: True if the tool is already a BedrockToolBlock, False otherwise - + Examples: >>> _is_bedrock_tool_block({"systemTool": {"name": "nova_grounding"}}) True diff --git a/litellm/litellm_core_utils/prompt_templates/huggingface_template_handler.py b/litellm/litellm_core_utils/prompt_templates/huggingface_template_handler.py index 9305d5bbfc..fc8a0d2858 100644 --- a/litellm/litellm_core_utils/prompt_templates/huggingface_template_handler.py +++ b/litellm/litellm_core_utils/prompt_templates/huggingface_template_handler.py @@ -12,7 +12,7 @@ from litellm.types.llms.custom_http import httpxSpecialProvider def strftime_now(fmt: str) -> str: """ Custom function for templates that need current date/time formatting (e.g., gpt-oss) - + Args: fmt: Format string for datetime.now().strftime() @@ -25,10 +25,10 @@ def strftime_now(fmt: str) -> str: def _get_tokenizer_config(hf_model_name: str) -> Dict[str, Any]: """ Fetch tokenizer_config.json from HuggingFace (sync) - + Args: hf_model_name: HuggingFace model name (e.g., 'openai/gpt-oss-120b') - + Returns: Dict with 'status' and optionally 'tokenizer' keys """ @@ -48,10 +48,10 @@ def _get_tokenizer_config(hf_model_name: str) -> Dict[str, Any]: async def _aget_tokenizer_config(hf_model_name: str) -> Dict[str, Any]: """ Fetch tokenizer_config.json from HuggingFace (async) - + Args: hf_model_name: HuggingFace model name (e.g., 'openai/gpt-oss-120b') - + Returns: Dict with 'status' and optionally 'tokenizer' keys """ @@ -73,35 +73,38 @@ async def _aget_tokenizer_config(hf_model_name: str) -> Dict[str, Any]: def _get_chat_template_file(hf_model_name: str) -> Dict[str, Any]: """ Fetch chat template from separate .jinja file (sync) - + Args: hf_model_name: HuggingFace model name (e.g., 'openai/gpt-oss-120b') - + Returns: Dict with 'status' and optionally 'chat_template' keys """ template_filenames = ["chat_template.jinja", "chat_template.jinja2"] client = _get_httpx_client() - + for filename in template_filenames: try: url = f"https://huggingface.co/{hf_model_name}/raw/main/{filename}" response = client.get(url=url) if response.status_code == 200: - return {"status": "success", "chat_template": response.content.decode("utf-8")} + return { + "status": "success", + "chat_template": response.content.decode("utf-8"), + } except Exception: continue - + return {"status": "failure"} async def _aget_chat_template_file(hf_model_name: str) -> Dict[str, Any]: """ Fetch chat template from separate .jinja file (async) - + Args: hf_model_name: HuggingFace model name (e.g., 'openai/gpt-oss-120b') - + Returns: Dict with 'status' and optionally 'chat_template' keys """ @@ -109,26 +112,29 @@ async def _aget_chat_template_file(hf_model_name: str) -> Dict[str, Any]: client = get_async_httpx_client( llm_provider=httpxSpecialProvider.PromptFactory, ) - + for filename in template_filenames: try: url = f"https://huggingface.co/{hf_model_name}/raw/main/{filename}" response = await client.get(url=url) if response.status_code == 200: - return {"status": "success", "chat_template": response.content.decode("utf-8")} + return { + "status": "success", + "chat_template": response.content.decode("utf-8"), + } except Exception: continue - + return {"status": "failure"} def _extract_token_value(token_value: Union[None, str, Dict[str, Any]]) -> str: """ Extract token string from various formats (string, dict, etc.) - + Args: token_value: Token value in various formats (None, str, or dict with 'content' key) - + Returns: Extracted token string """ @@ -136,4 +142,4 @@ def _extract_token_value(token_value: Union[None, str, Dict[str, Any]]) -> str: return token_value or "" if isinstance(token_value, dict): return token_value.get("content", "") - return "" \ No newline at end of file + return "" diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 7137a4e422..eaf78b7bcf 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -35,7 +35,7 @@ def _process_image_response(response: Response, url: str) -> str: max_bytes = int(MAX_IMAGE_URL_DOWNLOAD_SIZE_MB * 1024 * 1024) image_bytes = bytearray() bytes_downloaded = 0 - + for chunk in response.iter_bytes(chunk_size=8192): bytes_downloaded += len(chunk) if bytes_downloaded > max_bytes: @@ -44,7 +44,7 @@ def _process_image_response(response: Response, url: str) -> str: f"Error: Image size ({size_mb:.2f}MB) exceeds maximum allowed size ({MAX_IMAGE_URL_DOWNLOAD_SIZE_MB}MB). url={url}" ) image_bytes.extend(chunk) - + base64_image = base64.b64encode(image_bytes).decode("utf-8") image_type = response.headers.get("Content-Type") diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 14a25e61d6..3723368071 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -111,9 +111,7 @@ class RealTimeStreaming: if self._should_store_message(message_obj): self.messages.append(message_obj) - def _collect_user_input_from_client_event( - self, message: Union[str, dict] - ) -> None: + def _collect_user_input_from_client_event(self, message: Union[str, dict]) -> None: """Extract user text content from client WebSocket events for spend logging.""" try: if isinstance(message, str): @@ -158,15 +156,10 @@ class RealTimeStreaming: """Extract user voice transcription from backend events for spend logging.""" try: event_type = event_obj.get("type", "") - if ( - event_type - == "conversation.item.input_audio_transcription.completed" - ): + if event_type == "conversation.item.input_audio_transcription.completed": transcript = cast(str, event_obj.get("transcript", "")) if transcript: - self.input_messages.append( - {"role": "user", "content": transcript} - ) + self.input_messages.append({"role": "user", "content": transcript}) except (AttributeError, TypeError): pass @@ -204,9 +197,7 @@ class RealTimeStreaming: """Log messages in list""" if self.logging_obj: if self.input_messages: - self.logging_obj.model_call_details["messages"] = ( - self.input_messages - ) + self.logging_obj.model_call_details["messages"] = self.input_messages if self.session_tools or self.tool_calls: self.logging_obj.model_call_details[ "realtime_tools" @@ -313,10 +304,13 @@ class RealTimeStreaming: except Exception as e: # Re-raise unexpected errors (no status_code/detail = programming bug, not a block). # HTTPException and guardrail-raised exceptions have a status_code or detail attr. - is_guardrail_block = hasattr(e, "status_code") or isinstance(e, ValueError) + is_guardrail_block = hasattr(e, "status_code") or isinstance( + e, ValueError + ) if not is_guardrail_block: verbose_logger.exception( - "[realtime guardrail] unexpected error in apply_guardrail: %s", e + "[realtime guardrail] unexpected error in apply_guardrail: %s", + e, ) raise # Extract the human-readable error from the detail dict (HTTPException) @@ -327,23 +321,30 @@ class RealTimeStreaming: elif detail is not None: safe_msg = str(detail) else: - safe_msg = str(e) or "I'm sorry, that request was blocked by the content filter." + safe_msg = ( + str(e) + or "I'm sorry, that request was blocked by the content filter." + ) # Use realtime_violation_message if configured; fall back to guardrail error text. - error_msg = getattr(callback, "realtime_violation_message", None) or safe_msg + error_msg = ( + getattr(callback, "realtime_violation_message", None) or safe_msg + ) # Cancel any in-progress LLM response (e.g. VAD auto-response). await self._send_to_backend(json.dumps({"type": "response.cancel"})) # Send the policy violation hint (shows as small gray status text in UI). await self.websocket.send_text( - json.dumps({ - "type": "error", - "error": { - "type": "guardrail_violation", - "message": error_msg, - "code": "content_policy_violation", - }, - }) + json.dumps( + { + "type": "error", + "error": { + "type": "guardrail_violation", + "message": error_msg, + "code": "content_policy_violation", + }, + } + ) ) # Ask the LLM to voice the exact guardrail message so the # user hears it as audio in voice sessions (not just text). @@ -351,23 +352,29 @@ class RealTimeStreaming: f"Say exactly the following message to the user, word for word, " f"do not add anything else: {error_msg}" ) - await self._send_to_backend(json.dumps({ - "type": "conversation.item.create", - "item": { - "type": "message", - "role": "user", - "content": [{"type": "input_text", "text": guardrail_prompt}], - }, - })) await self._send_to_backend( - json.dumps({"type": "response.create"}) + json.dumps( + { + "type": "conversation.item.create", + "item": { + "type": "message", + "role": "user", + "content": [ + {"type": "input_text", "text": guardrail_prompt} + ], + }, + } + ) ) + await self._send_to_backend(json.dumps({"type": "response.create"})) self._violation_count += 1 end_session_after: Optional[int] = getattr( callback, "end_session_after_n_fails", None ) - should_end = getattr(callback, "on_violation", None) == "end_session" or ( + should_end = getattr( + callback, "on_violation", None + ) == "end_session" or ( end_session_after is not None and self._violation_count >= end_session_after ) @@ -410,7 +417,9 @@ class RealTimeStreaming: self.current_conversation_id = returned_object["current_conversation_id"] self.current_item_chunks = returned_object["current_item_chunks"] self.current_delta_type = returned_object["current_delta_type"] - self.session_configuration_request = returned_object["session_configuration_request"] + self.session_configuration_request = returned_object[ + "session_configuration_request" + ] events = ( transformed_response if isinstance(transformed_response, list) @@ -446,12 +455,11 @@ class RealTimeStreaming: self.store_message(event_str) await self.websocket.send_text(event_str) blocked = await self.run_realtime_guardrails( - cast(str, transcript), item_id=cast(Optional[str], event.get("item_id")) + cast(str, transcript), + item_id=cast(Optional[str], event.get("item_id")), ) if not blocked: - await self._send_to_backend( - json.dumps({"type": "response.create"}) - ) + await self._send_to_backend(json.dumps({"type": "response.create"})) continue ## LOGGING self.store_message(event_str) @@ -502,9 +510,7 @@ class RealTimeStreaming: ) if not blocked: # Clean — trigger LLM response - await self._send_to_backend( - json.dumps({"type": "response.create"}) - ) + await self._send_to_backend(json.dumps({"type": "response.create"})) return True except (json.JSONDecodeError, AttributeError): pass @@ -579,7 +585,10 @@ class RealTimeStreaming: self._pending_guardrail_message = combined_text continue # don't forward the original blocked message - if msg_type == "response.create" and self._pending_guardrail_message: + if ( + msg_type == "response.create" + and self._pending_guardrail_message + ): # The guardrail already sent the synthetic AI bubble — drop this # response.create so OpenAI doesn't generate an additional response. self._pending_guardrail_message = None diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index ad68f3851a..dbeb411107 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -64,15 +64,64 @@ def _redact_responses_api_output(output_items): for content_part in output_item.content: if hasattr(content_part, "text"): content_part.text = "redacted-by-litellm" - + # Redact reasoning items in output array if hasattr(output_item, "type") and output_item.type == "reasoning": - if hasattr(output_item, "summary") and isinstance(output_item.summary, list): + if hasattr(output_item, "summary") and isinstance( + output_item.summary, list + ): for summary_item in output_item.summary: if hasattr(summary_item, "text"): summary_item.text = "redacted-by-litellm" +def _redact_standard_logging_object(model_call_details: dict): + """Redact messages and response inside standard_logging_object if present.""" + standard_logging_object = model_call_details.get("standard_logging_object") + if standard_logging_object is None: + return + + redacted_str = "redacted-by-litellm" + + if standard_logging_object.get("messages") is not None: + standard_logging_object["messages"] = [ + {"role": "user", "content": redacted_str} + ] + + response = standard_logging_object.get("response") + if response is not None: + if isinstance(response, dict) and "output" in response: + # ResponsesAPIResponse format - redact content in output items + if isinstance(response.get("output"), list): + for output_item in response["output"]: + if isinstance(output_item, dict) and "content" in output_item: + if isinstance(output_item["content"], list): + for content_item in output_item["content"]: + if ( + isinstance(content_item, dict) + and "text" in content_item + ): + content_item["text"] = redacted_str + elif isinstance(response, dict) and "choices" in response: + # ModelResponse dict format - redact content in choices + if isinstance(response.get("choices"), list): + for choice in response["choices"]: + if isinstance(choice, dict): + if "message" in choice and isinstance(choice["message"], dict): + choice["message"]["content"] = redacted_str + if "audio" in choice["message"]: + choice["message"]["audio"] = None + elif "delta" in choice and isinstance(choice["delta"], dict): + choice["delta"]["content"] = redacted_str + if "audio" in choice["delta"]: + choice["delta"]["audio"] = None + elif isinstance(response, str): + standard_logging_object["response"] = redacted_str + else: + # For other formats (empty dict, None, etc.), use simple text format + standard_logging_object["response"] = {"text": redacted_str} + + def perform_redaction(model_call_details: dict, result): """ Performs the actual redaction on the logging object and result. @@ -96,24 +145,56 @@ def perform_redaction(model_call_details: dict, result): elif hasattr(_streaming_response, "output"): _redact_responses_api_output(_streaming_response.output) # Redact reasoning field in ResponsesAPIResponse - if hasattr(_streaming_response, "reasoning") and _streaming_response.reasoning is not None: + if ( + hasattr(_streaming_response, "reasoning") + and _streaming_response.reasoning is not None + ): _streaming_response.reasoning = None # Redact result if result is not None: # Check if result is a coroutine, async generator, or other async object - these cannot be deepcopied - if (asyncio.iscoroutine(result) or - inspect.iscoroutinefunction(result) or - hasattr(result, '__aiter__') or # async generator - hasattr(result, '__anext__')): # async iterator + if ( + asyncio.iscoroutine(result) + or inspect.iscoroutinefunction(result) + or hasattr(result, "__aiter__") + or hasattr(result, "__anext__") # async generator + ): # async iterator # For async objects, return a simple redacted response without deepcopy return {"text": "redacted-by-litellm"} - + _result = copy.deepcopy(result) if isinstance(_result, litellm.ModelResponse): if hasattr(_result, "choices") and _result.choices is not None: for choice in _result.choices: _redact_choice_content(choice) + elif isinstance(_result, dict) and "choices" in _result: + # Handle dict representation of ModelResponse (e.g., from model_dump()) + if _result.get("choices") is not None: + for choice in _result["choices"]: + if isinstance(choice, dict): + if "message" in choice and isinstance(choice["message"], dict): + choice["message"]["content"] = "redacted-by-litellm" + if "reasoning_content" in choice["message"]: + choice["message"][ + "reasoning_content" + ] = "redacted-by-litellm" + if "thinking_blocks" in choice["message"]: + choice["message"]["thinking_blocks"] = None + if "audio" in choice["message"]: + choice["message"]["audio"] = None + elif "delta" in choice and isinstance(choice["delta"], dict): + choice["delta"]["content"] = "redacted-by-litellm" + if "reasoning_content" in choice["delta"]: + choice["delta"][ + "reasoning_content" + ] = "redacted-by-litellm" + if "thinking_blocks" in choice["delta"]: + choice["delta"]["thinking_blocks"] = None + if "audio" in choice["delta"]: + choice["delta"]["audio"] = None + else: + _redact_choice_content(choice) elif isinstance(_result, litellm.ResponsesAPIResponse): if hasattr(_result, "output"): _redact_responses_api_output(_result.output) @@ -131,14 +212,14 @@ def perform_redaction(model_call_details: dict, result): def should_redact_message_logging(model_call_details: dict) -> bool: """ Determine if message logging should be redacted. - + Priority order: 1. Dynamic parameter (turn_off_message_logging in request) 2. Headers (litellm-disable-message-redaction / litellm-enable-message-redaction) 3. Global setting (litellm.turn_off_message_logging) """ litellm_params = model_call_details.get("litellm_params", {}) - + metadata_field = get_metadata_variable_name_from_kwargs(litellm_params) metadata = litellm_params.get(metadata_field, {}) if not isinstance(metadata, dict): @@ -169,15 +250,17 @@ def should_redact_message_logging(model_call_details: dict) -> bool: break # Priority 1: Check dynamic parameter first (if explicitly set) - dynamic_turn_off = _get_turn_off_message_logging_from_dynamic_params(model_call_details) + dynamic_turn_off = _get_turn_off_message_logging_from_dynamic_params( + model_call_details + ) if dynamic_turn_off is not None: # Dynamic parameter is explicitly set, use it return dynamic_turn_off - + # Priority 2: Check if header explicitly enables redaction if is_redaction_enabled_via_header: return True - + # Priority 3: Fall back to global setting return litellm.turn_off_message_logging is True @@ -202,9 +285,9 @@ def _get_turn_off_message_logging_from_dynamic_params( handles boolean and string values of `turn_off_message_logging` """ - standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = ( - model_call_details.get("standard_callback_dynamic_params", None) - ) + standard_callback_dynamic_params: Optional[ + StandardCallbackDynamicParams + ] = model_call_details.get("standard_callback_dynamic_params", None) if standard_callback_dynamic_params: _turn_off_message_logging = standard_callback_dynamic_params.get( "turn_off_message_logging" diff --git a/litellm/litellm_core_utils/safe_json_loads.py b/litellm/litellm_core_utils/safe_json_loads.py index a7ab0d3e3b..bb4b72cfd9 100644 --- a/litellm/litellm_core_utils/safe_json_loads.py +++ b/litellm/litellm_core_utils/safe_json_loads.py @@ -4,6 +4,7 @@ Helper for safe JSON loading in LiteLLM. from typing import Any import json + def safe_json_loads(data: str, default: Any = None) -> Any: """ Safely parse a JSON string. If parsing fails, return the default value (None by default). @@ -11,4 +12,4 @@ def safe_json_loads(data: str, default: Any = None) -> Any: try: return json.loads(data) except Exception: - return default \ No newline at end of file + return default diff --git a/litellm/litellm_core_utils/sensitive_data_masker.py b/litellm/litellm_core_utils/sensitive_data_masker.py index 3ec34e6d9e..663c3fac80 100644 --- a/litellm/litellm_core_utils/sensitive_data_masker.py +++ b/litellm/litellm_core_utils/sensitive_data_masker.py @@ -98,7 +98,9 @@ class SensitiveDataMasker: masked_items.append(self._mask_value(item)) else: masked_items.append( - item if isinstance(item, (int, float, bool, str, list)) else str(item) + item + if isinstance(item, (int, float, bool, str, list)) + else str(item) ) return masked_items diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index ba35a2c7ca..1935372e5d 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -54,13 +54,16 @@ class ChunkProcessor: first_hidden_params = candidate if first_hidden_params.get("created_at"): + def _created_at(chunk: Any) -> Union[int, float]: if isinstance(chunk, dict): params = chunk.get("_hidden_params", {}) else: params = getattr(chunk, "_hidden_params", {}) if isinstance(params, dict): - return cast(Union[int, float], params.get("created_at", float("inf"))) + return cast( + Union[int, float], params.get("created_at", float("inf")) + ) return float("inf") return sorted(chunks, key=_created_at) @@ -88,7 +91,9 @@ class ChunkProcessor: return "" @staticmethod - def _get_model_from_chunks(chunks: List[Dict[str, Any]], first_chunk_model: str) -> str: + def _get_model_from_chunks( + chunks: List[Dict[str, Any]], first_chunk_model: str + ) -> str: """ Get the actual model from chunks, preferring a model that differs from the first chunk. @@ -151,13 +156,13 @@ class ChunkProcessor: ) return response - def get_combined_tool_content( # noqa: PLR0915 + def get_combined_tool_content( # noqa: PLR0915 self, tool_call_chunks: List[Dict[str, Any]] ) -> List[ChatCompletionMessageToolCall]: tool_calls_list: List[ChatCompletionMessageToolCall] = [] - tool_call_map: Dict[int, Dict[str, Any]] = ( - {} - ) # Map to store tool calls by index + tool_call_map: Dict[ + int, Dict[str, Any] + ] = {} # Map to store tool calls by index for chunk in tool_call_chunks: choices = chunk["choices"] @@ -169,14 +174,20 @@ class ChunkProcessor: # Handle both dict and object formats if not tool_call: continue - + # Check if tool_call has function (either as attribute or dict key) has_function = False if isinstance(tool_call, dict): - has_function = "function" in tool_call and tool_call["function"] is not None + has_function = ( + "function" in tool_call + and tool_call["function"] is not None + ) else: - has_function = hasattr(tool_call, "function") and tool_call.function is not None - + has_function = ( + hasattr(tool_call, "function") + and tool_call.function is not None + ) + if not has_function: continue @@ -185,7 +196,7 @@ class ChunkProcessor: index = tool_call.get("index", 0) else: index = getattr(tool_call, "index", 0) - + if index not in tool_call_map: tool_call_map[index] = { "id": None, @@ -201,19 +212,23 @@ class ChunkProcessor: tool_call_map[index]["id"] = tool_call["id"] if tool_call.get("type"): tool_call_map[index]["type"] = tool_call["type"] - + function = tool_call.get("function", {}) if isinstance(function, dict): if function.get("name"): tool_call_map[index]["name"] = function["name"] if function.get("arguments"): - tool_call_map[index]["arguments"].append(function["arguments"]) + tool_call_map[index]["arguments"].append( + function["arguments"] + ) else: # function is an object if hasattr(function, "name") and function.name: tool_call_map[index]["name"] = function.name if hasattr(function, "arguments") and function.arguments: - tool_call_map[index]["arguments"].append(function.arguments) + tool_call_map[index]["arguments"].append( + function.arguments + ) else: # tool_call is an object if hasattr(tool_call, "id") and tool_call.id: @@ -233,19 +248,32 @@ class ChunkProcessor: tool_call_map[index]["arguments"].append( tool_call.function.arguments ) - + # Preserve provider_specific_fields from streaming chunks provider_fields = None if isinstance(tool_call, dict): provider_fields = tool_call.get("provider_specific_fields") - if not provider_fields and isinstance(tool_call.get("function"), dict): - provider_fields = tool_call["function"].get("provider_specific_fields") + if not provider_fields and isinstance( + tool_call.get("function"), dict + ): + provider_fields = tool_call["function"].get( + "provider_specific_fields" + ) else: - if hasattr(tool_call, "provider_specific_fields") and tool_call.provider_specific_fields: + if ( + hasattr(tool_call, "provider_specific_fields") + and tool_call.provider_specific_fields + ): provider_fields = tool_call.provider_specific_fields - elif hasattr(tool_call, "function") and hasattr(tool_call.function, "provider_specific_fields") and tool_call.function.provider_specific_fields: - provider_fields = tool_call.function.provider_specific_fields - + elif ( + hasattr(tool_call, "function") + and hasattr(tool_call.function, "provider_specific_fields") + and tool_call.function.provider_specific_fields + ): + provider_fields = ( + tool_call.function.provider_specific_fields + ) + if provider_fields: # Merge provider_specific_fields if multiple chunks have them if tool_call_map[index]["provider_specific_fields"] is None: @@ -260,30 +288,31 @@ class ChunkProcessor: tool_call_data = tool_call_map[index] if tool_call_data["id"] and tool_call_data["name"]: combined_arguments = "".join(tool_call_data["arguments"]) or "{}" - + # Build function - provider_specific_fields should be on tool_call level, not function level function = Function( arguments=combined_arguments, name=tool_call_data["name"], ) - + # Prepare params for ChatCompletionMessageToolCall tool_call_params = { "id": tool_call_data["id"], "function": function, "type": tool_call_data["type"] or "function", } - + # Add provider_specific_fields if present (for thought signatures in Gemini 3) if tool_call_data.get("provider_specific_fields"): - tool_call_params["provider_specific_fields"] = tool_call_data["provider_specific_fields"] - + tool_call_params["provider_specific_fields"] = tool_call_data[ + "provider_specific_fields" + ] + tool_call = ChatCompletionMessageToolCall(**tool_call_params) tool_calls_list.append(tool_call) return tool_calls_list - def get_combined_function_call_content( self, function_call_chunks: List[Dict[str, Any]] ) -> FunctionCall: @@ -506,7 +535,7 @@ class ChunkProcessor: ## anthropic prompt caching information ## cache_creation_input_tokens: Optional[int] = None cache_read_input_tokens: Optional[int] = None - + server_tool_use: Optional[ServerToolUse] = None web_search_requests: Optional[int] = None completion_tokens_details: Optional[CompletionTokensDetails] = None @@ -551,7 +580,10 @@ class ChunkProcessor: completion_tokens_details = usage_chunk_dict[ "completion_tokens_details" ] - if hasattr(usage_chunk, 'server_tool_use') and usage_chunk.server_tool_use is not None: + if ( + hasattr(usage_chunk, "server_tool_use") + and usage_chunk.server_tool_use is not None + ): server_tool_use = usage_chunk.server_tool_use if ( usage_chunk_dict["prompt_tokens_details"] is not None @@ -611,12 +643,12 @@ class ChunkProcessor: web_search_requests: Optional[int] = calculated_usage_per_chunk[ "web_search_requests" ] - completion_tokens_details: Optional[CompletionTokensDetails] = ( - calculated_usage_per_chunk["completion_tokens_details"] - ) - prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = ( - calculated_usage_per_chunk["prompt_tokens_details"] - ) + completion_tokens_details: Optional[ + CompletionTokensDetails + ] = calculated_usage_per_chunk["completion_tokens_details"] + prompt_tokens_details: Optional[ + PromptTokensDetailsWrapper + ] = calculated_usage_per_chunk["prompt_tokens_details"] try: returned_usage.prompt_tokens = prompt_tokens or token_counter( @@ -650,8 +682,10 @@ class ChunkProcessor: ) # for anthropic if completion_tokens_details is not None: if isinstance(completion_tokens_details, CompletionTokensDetails): - returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper( - **completion_tokens_details.model_dump() + returned_usage.completion_tokens_details = ( + CompletionTokensDetailsWrapper( + **completion_tokens_details.model_dump() + ) ) else: returned_usage.completion_tokens_details = completion_tokens_details diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 317f103768..db2369d03d 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -485,7 +485,6 @@ class CustomStreamWrapper: def handle_openai_chat_completion_chunk(self, chunk): try: - str_line = chunk text = "" is_finished = False @@ -535,7 +534,6 @@ class CustomStreamWrapper: def handle_azure_text_completion_chunk(self, chunk): try: - text = "" is_finished = False finish_reason = None @@ -556,7 +554,6 @@ class CustomStreamWrapper: def handle_openai_text_completion_chunk(self, chunk): try: - text = "" is_finished = False finish_reason = None @@ -1100,8 +1097,7 @@ class CustomStreamWrapper: ): if self.received_finish_reason is not None: _chunk_has_content = isinstance(chunk, dict) and ( - bool(chunk.get("text", "")) - or chunk.get("tool_use") is not None + bool(chunk.get("text", "")) or chunk.get("tool_use") is not None ) if not _chunk_has_content and ( not isinstance(chunk, dict) @@ -1356,7 +1352,10 @@ class CustomStreamWrapper: if response_obj["is_finished"]: self.received_finish_reason = response_obj["finish_reason"] else: # openai / azure chat model - if self.custom_llm_provider in [LlmProviders.AZURE.value, LlmProviders.AZURE_AI.value]: + if self.custom_llm_provider in [ + LlmProviders.AZURE.value, + LlmProviders.AZURE_AI.value, + ]: if isinstance(chunk, BaseModel) and hasattr(chunk, "model"): # for azure, we need to pass the model from the original chunk self.model = getattr(chunk, "model", self.model) @@ -1612,10 +1611,12 @@ class CustomStreamWrapper: ) return chunk - def _add_mcp_list_tools_to_first_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: + def _add_mcp_list_tools_to_first_chunk( + self, chunk: ModelResponseStream + ) -> ModelResponseStream: """ Add mcp_list_tools from _hidden_params to the first chunk's delta.provider_specific_fields. - + This method checks if MCP metadata with mcp_list_tools is stored in _hidden_params and adds it to the first chunk's delta.provider_specific_fields. """ @@ -1623,43 +1624,53 @@ class CustomStreamWrapper: # Check if MCP metadata should be added to first chunk if not hasattr(self, "_hidden_params") or not self._hidden_params: return chunk - + mcp_metadata = self._hidden_params.get("mcp_metadata") if not mcp_metadata or not isinstance(mcp_metadata, dict): return chunk - + # Only add mcp_list_tools to first chunk (not tool_calls or tool_results) mcp_list_tools = mcp_metadata.get("mcp_list_tools") if not mcp_list_tools: return chunk - + # Add mcp_list_tools to delta.provider_specific_fields if hasattr(chunk, "choices") and chunk.choices: for choice in chunk.choices: - if isinstance(choice, StreamingChoices) and hasattr(choice, "delta") and choice.delta: + if ( + isinstance(choice, StreamingChoices) + and hasattr(choice, "delta") + and choice.delta + ): # Get existing provider_specific_fields or create new dict provider_fields = ( - getattr(choice.delta, "provider_specific_fields", None) or {} + getattr(choice.delta, "provider_specific_fields", None) + or {} ) - + # Add only mcp_list_tools to first chunk provider_fields["mcp_list_tools"] = mcp_list_tools - + # Set the provider_specific_fields - setattr(choice.delta, "provider_specific_fields", provider_fields) - + setattr( + choice.delta, "provider_specific_fields", provider_fields + ) + except Exception as e: from litellm._logging import verbose_logger + verbose_logger.exception( f"Error adding MCP list tools to first chunk: {str(e)}" ) - + return chunk - def _add_mcp_metadata_to_final_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: + def _add_mcp_metadata_to_final_chunk( + self, chunk: ModelResponseStream + ) -> ModelResponseStream: """ Add MCP metadata from _hidden_params to the final chunk's delta.provider_specific_fields. - + This method checks if MCP metadata is stored in _hidden_params and adds it to the chunk's delta.provider_specific_fields, similar to how RAG adds search results. """ @@ -1667,33 +1678,41 @@ class CustomStreamWrapper: # Check if MCP metadata should be added to final chunk if not hasattr(self, "_hidden_params") or not self._hidden_params: return chunk - + mcp_metadata = self._hidden_params.get("mcp_metadata") if not mcp_metadata: return chunk - + # Add MCP metadata to delta.provider_specific_fields if hasattr(chunk, "choices") and chunk.choices: for choice in chunk.choices: - if isinstance(choice, StreamingChoices) and hasattr(choice, "delta") and choice.delta: + if ( + isinstance(choice, StreamingChoices) + and hasattr(choice, "delta") + and choice.delta + ): # Get existing provider_specific_fields or create new dict provider_fields = ( - getattr(choice.delta, "provider_specific_fields", None) or {} + getattr(choice.delta, "provider_specific_fields", None) + or {} ) - + # Add MCP metadata if isinstance(mcp_metadata, dict): provider_fields.update(mcp_metadata) - + # Set the provider_specific_fields - setattr(choice.delta, "provider_specific_fields", provider_fields) - + setattr( + choice.delta, "provider_specific_fields", provider_fields + ) + except Exception as e: from litellm._logging import verbose_logger + verbose_logger.exception( f"Error adding MCP metadata to final chunk: {str(e)}" ) - + return chunk def cache_streaming_response(self, processed_chunk, cache_hit: bool): @@ -1813,12 +1832,12 @@ class CustomStreamWrapper: ) # HANDLE STREAM OPTIONS self.chunks.append(response) - + # Add mcp_list_tools to first chunk if present if not self.sent_first_chunk: response = self._add_mcp_list_tools_to_first_chunk(response) self.sent_first_chunk = True - + if hasattr( response, "usage" ): # remove usage from chunk, only send on final chunk @@ -1898,9 +1917,7 @@ class CustomStreamWrapper: and complete_streaming_response is not None and self._last_returned_hidden_params is not None ): - final_usage = getattr( - complete_streaming_response, "usage", None - ) + final_usage = getattr(complete_streaming_response, "usage", None) if final_usage is not None: self._last_returned_hidden_params["usage"] = final_usage @@ -1991,7 +2008,9 @@ class CustomStreamWrapper: ) # Add mcp_list_tools to first chunk if present if not self.sent_first_chunk: - processed_chunk = self._add_mcp_list_tools_to_first_chunk(processed_chunk) + processed_chunk = self._add_mcp_list_tools_to_first_chunk( + processed_chunk + ) self.sent_first_chunk = True _has_usage = ( @@ -2026,7 +2045,9 @@ class CustomStreamWrapper: if self.sent_last_chunk is True and self.stream_options is None: usage = calculate_total_usage(chunks=self.chunks) processed_chunk._hidden_params["usage"] = usage - self._last_returned_hidden_params = processed_chunk._hidden_params + self._last_returned_hidden_params = ( + processed_chunk._hidden_params + ) # Call post-call streaming deployment hook for final chunk if self.sent_last_chunk is True: @@ -2098,9 +2119,7 @@ class CustomStreamWrapper: and complete_streaming_response is not None and self._last_returned_hidden_params is not None ): - final_usage = getattr( - complete_streaming_response, "usage", None - ) + final_usage = getattr(complete_streaming_response, "usage", None) if final_usage is not None: self._last_returned_hidden_params["usage"] = final_usage @@ -2214,9 +2233,17 @@ class CustomStreamWrapper: # Raise non-retriable client errors directly (skip fallback). # Exception: 429 (rate-limit) IS retriable/transient — allow it # through so the Router can switch to a different model group. - if mapped_status_code is not None and 400 <= mapped_status_code < 500 and mapped_status_code != 429: + if ( + mapped_status_code is not None + and 400 <= mapped_status_code < 500 + and mapped_status_code != 429 + ): raise mapped_exception - if original_status_code is not None and 400 <= original_status_code < 500 and original_status_code != 429: + if ( + original_status_code is not None + and 400 <= original_status_code < 500 + and original_status_code != 429 + ): raise mapped_exception raise MidStreamFallbackError( diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index da357e51c2..09c62f2eb5 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -727,7 +727,9 @@ def _count_content_list( num_tokens += count_function(thinking_text) else: content_type = ( - c.get("type", type(c).__name__) if isinstance(c, dict) else type(c).__name__ + c.get("type", type(c).__name__) + if isinstance(c, dict) + else type(c).__name__ ) raise ValueError( f"Invalid content item type: {content_type}. " diff --git a/litellm/llms/a2a/chat/streaming_iterator.py b/litellm/llms/a2a/chat/streaming_iterator.py index 4b689414dd..72902f65f7 100644 --- a/litellm/llms/a2a/chat/streaming_iterator.py +++ b/litellm/llms/a2a/chat/streaming_iterator.py @@ -12,10 +12,10 @@ from ..common_utils import extract_text_from_a2a_response class A2AModelResponseIterator(BaseModelResponseIterator): """ Iterator for parsing A2A streaming responses. - + Converts A2A JSON-RPC streaming chunks to OpenAI-compatible format. """ - + def __init__( self, streaming_response, @@ -29,11 +29,13 @@ class A2AModelResponseIterator(BaseModelResponseIterator): json_mode=json_mode, ) self.model = model - - def chunk_parser(self, chunk: dict) -> Union[GenericStreamingChunk, ModelResponseStream]: + + def chunk_parser( + self, chunk: dict + ) -> Union[GenericStreamingChunk, ModelResponseStream]: """ Parse A2A streaming chunk to OpenAI format. - + A2A chunk format: { "jsonrpc": "2.0", @@ -44,7 +46,7 @@ class A2AModelResponseIterator(BaseModelResponseIterator): } } } - + Or for tasks: { "jsonrpc": "2.0", @@ -58,10 +60,10 @@ class A2AModelResponseIterator(BaseModelResponseIterator): try: # Extract text from A2A response text = extract_text_from_a2a_response(chunk) - + # Determine finish reason finish_reason = self._get_finish_reason(chunk) - + # Return generic streaming chunk return GenericStreamingChunk( text=text, @@ -81,11 +83,11 @@ class A2AModelResponseIterator(BaseModelResponseIterator): index=0, tool_use=None, ) - + def _get_finish_reason(self, chunk: dict) -> Optional[str]: """Extract finish reason from A2A chunk""" result = chunk.get("result", {}) - + # Check for task completion if isinstance(result, dict): status = result.get("status", {}) @@ -95,9 +97,9 @@ class A2AModelResponseIterator(BaseModelResponseIterator): return "stop" elif state == "failed": return "stop" # Map failed state to 'stop' (valid finish_reason) - + # Check for [DONE] marker if chunk.get("done") is True: return "stop" - + return None diff --git a/litellm/llms/a2a/chat/transformation.py b/litellm/llms/a2a/chat/transformation.py index 163cd5ab22..d088702863 100644 --- a/litellm/llms/a2a/chat/transformation.py +++ b/litellm/llms/a2a/chat/transformation.py @@ -22,10 +22,10 @@ from .streaming_iterator import A2AModelResponseIterator class A2AConfig(BaseConfig): """ Configuration for A2A (Agent-to-Agent) Protocol. - + Handles transformation between OpenAI and A2A JSON-RPC 2.0 formats. """ - + @staticmethod def resolve_agent_config_from_registry( model: str, @@ -36,58 +36,63 @@ class A2AConfig(BaseConfig): ) -> tuple[Optional[str], Optional[str], Optional[Dict[str, Any]]]: """ Resolve agent configuration from registry if model format is "a2a/". - + Extracts agent name from model string and looks up configuration in the agent registry (if available in proxy context). - + Args: model: Model string (e.g., "a2a/my-agent") api_base: Explicit api_base (takes precedence over registry) api_key: Explicit api_key (takes precedence over registry) headers: Explicit headers (takes precedence over registry) optional_params: Dict to merge additional litellm_params into - + Returns: Tuple of (api_base, api_key, headers) with registry values filled in """ # Extract agent name from model (e.g., "a2a/my-agent" -> "my-agent") agent_name = model.split("/", 1)[1] if "/" in model else None - + # Only lookup if agent name exists and some config is missing - if not agent_name or (api_base is not None and api_key is not None and headers is not None): + if not agent_name or ( + api_base is not None and api_key is not None and headers is not None + ): return api_base, api_key, headers - + # Try registry lookup (only available in proxy context) try: from litellm.proxy.agent_endpoints.agent_registry import ( global_agent_registry, ) - + agent = global_agent_registry.get_agent_by_name(agent_name) if agent: # Get api_base from agent card URL if api_base is None and agent.agent_card_params: api_base = agent.agent_card_params.get("url") - + # Get api_key, headers, and other params from litellm_params if agent.litellm_params: if api_key is None: api_key = agent.litellm_params.get("api_key") - + if headers is None: agent_headers = agent.litellm_params.get("headers") if agent_headers: headers = agent_headers - + # Merge other litellm_params (timeout, max_retries, etc.) for key, value in agent.litellm_params.items(): - if key not in ["api_key", "api_base", "headers", "model"] and key not in optional_params: + if ( + key not in ["api_key", "api_base", "headers", "model"] + and key not in optional_params + ): optional_params[key] = value except ImportError: pass # Registry not available (not running in proxy context) - + return api_base, api_key, headers - + def get_supported_openai_params(self, model: str) -> List[str]: """Return list of supported OpenAI parameters""" return [ @@ -96,7 +101,7 @@ class A2AConfig(BaseConfig): "max_tokens", "top_p", ] - + def map_openai_params( self, non_default_params: dict, @@ -106,7 +111,7 @@ class A2AConfig(BaseConfig): ) -> dict: """ Map OpenAI parameters to A2A parameters. - + For A2A protocol, we need to map the stream parameter so transform_request can determine which JSON-RPC method to use. """ @@ -114,9 +119,9 @@ class A2AConfig(BaseConfig): for param, value in non_default_params.items(): if param == "stream" and value is True: optional_params["stream"] = value - + return optional_params - + def validate_environment( self, headers: dict, @@ -129,7 +134,7 @@ class A2AConfig(BaseConfig): ) -> dict: """ Validate environment and set headers for A2A requests. - + Args: headers: Request headers dict model: Model name @@ -138,20 +143,20 @@ class A2AConfig(BaseConfig): litellm_params: LiteLLM parameters api_key: API key (optional for A2A) api_base: API base URL - + Returns: Updated headers dict """ # Ensure Content-Type is set to application/json for JSON-RPC 2.0 if "content-type" not in headers and "Content-Type" not in headers: headers["Content-Type"] = "application/json" - + # Add Authorization header if API key is provided if api_key is not None: headers["Authorization"] = f"Bearer {api_key}" - + return headers - + def get_complete_url( self, api_base: Optional[str], @@ -163,11 +168,11 @@ class A2AConfig(BaseConfig): ) -> str: """ Get the complete A2A agent endpoint URL. - + A2A agents use JSON-RPC 2.0 at the base URL, not specific paths. The method (message/send or message/stream) is specified in the JSON-RPC request body, not in the URL. - + Args: api_base: Base URL of the A2A agent (e.g., "http://0.0.0.0:9999") api_key: API key (not used for URL construction) @@ -175,17 +180,17 @@ class A2AConfig(BaseConfig): optional_params: Optional parameters litellm_params: LiteLLM parameters stream: Whether this is a streaming request (affects JSON-RPC method) - + Returns: Complete URL for the A2A endpoint (base URL) """ if api_base is None: raise ValueError("api_base is required for A2A provider") - + # A2A uses JSON-RPC 2.0 at the base URL # Remove trailing slash for consistency return api_base.rstrip("/") - + def transform_request( self, model: str, @@ -196,51 +201,49 @@ class A2AConfig(BaseConfig): ) -> dict: """ Transform OpenAI request to A2A JSON-RPC 2.0 format. - + Args: model: Model name messages: List of OpenAI messages optional_params: Optional parameters litellm_params: LiteLLM parameters headers: Request headers - + Returns: A2A JSON-RPC 2.0 request dict """ # Generate request ID request_id = str(uuid.uuid4()) - + if not messages: raise ValueError("At least one message is required for A2A completion") - + # Convert all messages to maintain conversation history # Use helper to format conversation with role prefixes full_context = convert_messages_to_prompt(messages) - + # Create single A2A message with full conversation context a2a_message = { "role": "user", "parts": [{"kind": "text", "text": full_context}], "messageId": str(uuid.uuid4()), } - + # Build JSON-RPC 2.0 request # For A2A protocol, the method is "message/send" for non-streaming # and "message/stream" for streaming stream = optional_params.get("stream", False) method = "message/stream" if stream else "message/send" - + request_data = { "jsonrpc": "2.0", "id": request_id, "method": method, - "params": { - "message": a2a_message - } + "params": {"message": a2a_message}, } - + return request_data - + def transform_response( self, model: str, @@ -257,7 +260,7 @@ class A2AConfig(BaseConfig): ) -> ModelResponse: """ Transform A2A JSON-RPC 2.0 response to OpenAI format. - + Args: model: Model name raw_response: HTTP response from A2A agent @@ -270,7 +273,7 @@ class A2AConfig(BaseConfig): encoding: Encoding object api_key: API key json_mode: JSON mode flag - + Returns: Populated ModelResponse object """ @@ -282,7 +285,7 @@ class A2AConfig(BaseConfig): message=f"Failed to parse A2A response: {str(e)}", headers=dict(raw_response.headers), ) - + # Check for JSON-RPC error if "error" in response_json: error = response_json["error"] @@ -291,10 +294,10 @@ class A2AConfig(BaseConfig): message=f"A2A error: {error.get('message', 'Unknown error')}", headers=dict(raw_response.headers), ) - + # Extract text from A2A response text = extract_text_from_a2a_response(response_json) - + # Populate model response model_response.choices = [ Choices( @@ -306,15 +309,15 @@ class A2AConfig(BaseConfig): ), ) ] - + # Set model model_response.model = model - + # Set ID from response model_response.id = response_json.get("id", str(uuid.uuid4())) - + return model_response - + def get_model_response_iterator( self, streaming_response: Union[Iterator, Any], @@ -323,12 +326,12 @@ class A2AConfig(BaseConfig): ) -> BaseModelResponseIterator: """ Get streaming iterator for A2A responses. - + Args: streaming_response: Streaming response iterator sync_stream: Whether this is a sync stream json_mode: JSON mode flag - + Returns: A2A streaming iterator """ @@ -337,26 +340,26 @@ class A2AConfig(BaseConfig): sync_stream=sync_stream, json_mode=json_mode, ) - + def _openai_message_to_a2a_message(self, message: Dict[str, Any]) -> Dict[str, Any]: """ Convert OpenAI message to A2A message format. - + Args: message: OpenAI message dict - + Returns: A2A message dict """ content = message.get("content", "") role = message.get("role", "user") - + return { "role": role, "parts": [{"kind": "text", "text": str(content)}], "messageId": str(uuid.uuid4()), } - + def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] ) -> BaseLLMException: diff --git a/litellm/llms/a2a/common_utils.py b/litellm/llms/a2a/common_utils.py index 116e120540..aa817ce0fe 100644 --- a/litellm/llms/a2a/common_utils.py +++ b/litellm/llms/a2a/common_utils.py @@ -31,13 +31,13 @@ class A2AError(BaseLLMException): def convert_messages_to_prompt(messages: List[AllMessageValues]) -> str: """ Convert OpenAI messages to a single prompt string for A2A agent. - + Formats each message as "{role}: {content}" and joins with newlines to preserve conversation history. Handles both string and list content. - + Args: messages: List of OpenAI-format messages - + Returns: Formatted prompt string with full conversation context """ @@ -45,7 +45,7 @@ def convert_messages_to_prompt(messages: List[AllMessageValues]) -> str: for msg in messages: # Use LiteLLM's helper to extract text from content (handles both str and list) content_text = convert_content_list_to_str(message=msg) - + # Get role if isinstance(msg, BaseModel): role = msg.model_dump().get("role", "user") @@ -53,10 +53,10 @@ def convert_messages_to_prompt(messages: List[AllMessageValues]) -> str: role = msg.get("role", "user") else: role = dict(msg).get("role", "user") # type: ignore - + if content_text: conversation_parts.append(f"{role}: {content_text}") - + return "\n".join(conversation_parts) @@ -65,21 +65,21 @@ def extract_text_from_a2a_message( ) -> str: """ Extract text content from A2A message parts. - + Args: message: A2A message dict with 'parts' containing text parts depth: Current recursion depth (internal use) max_depth: Maximum recursion depth to prevent infinite loops - + Returns: Concatenated text from all text parts """ if message is None or depth >= max_depth: return "" - + parts = message.get("parts", []) text_parts: List[str] = [] - + for part in parts: if part.get("kind") == "text": text_parts.append(part.get("text", "")) @@ -88,7 +88,7 @@ def extract_text_from_a2a_message( nested_text = extract_text_from_a2a_message(part, depth + 1, max_depth) if nested_text: text_parts.append(nested_text) - + return " ".join(text_parts) @@ -97,41 +97,39 @@ def extract_text_from_a2a_response( ) -> str: """ Extract text content from A2A response result. - + Args: response_dict: A2A response dict with 'result' containing message max_depth: Maximum recursion depth to prevent infinite loops - + Returns: Text from response message parts """ result = response_dict.get("result", {}) if not isinstance(result, dict): return "" - + # A2A response can have different formats: # 1. Direct message: {"result": {"kind": "message", "parts": [...]}} # 2. Nested message: {"result": {"message": {"parts": [...]}}} # 3. Task with artifacts: {"result": {"kind": "task", "artifacts": [{"parts": [...]}]}} # 4. Task with status message: {"result": {"kind": "task", "status": {"message": {"parts": [...]}}}} # 5. Streaming artifact-update: {"result": {"kind": "artifact-update", "artifact": {"parts": [...]}}} - + # Check if result itself has parts (direct message) if "parts" in result: return extract_text_from_a2a_message(result, depth=0, max_depth=max_depth) - + # Check for nested message message = result.get("message") if message: return extract_text_from_a2a_message(message, depth=0, max_depth=max_depth) - + # Check for streaming artifact-update (singular artifact) artifact = result.get("artifact") if artifact and isinstance(artifact, dict): - return extract_text_from_a2a_message( - artifact, depth=0, max_depth=max_depth - ) - + return extract_text_from_a2a_message(artifact, depth=0, max_depth=max_depth) + # Check for task status message (common in Gemini A2A agents) status = result.get("status", {}) if isinstance(status, dict): @@ -140,7 +138,7 @@ def extract_text_from_a2a_response( return extract_text_from_a2a_message( status_message, depth=0, max_depth=max_depth ) - + # Handle task result with artifacts (plural, array) artifacts = result.get("artifacts", []) if artifacts and len(artifacts) > 0: @@ -148,5 +146,5 @@ def extract_text_from_a2a_response( return extract_text_from_a2a_message( first_artifact, depth=0, max_depth=max_depth ) - + return "" diff --git a/litellm/llms/aiml/chat/transformation.py b/litellm/llms/aiml/chat/transformation.py index 0f3e333343..72e30a0817 100644 --- a/litellm/llms/aiml/chat/transformation.py +++ b/litellm/llms/aiml/chat/transformation.py @@ -20,4 +20,5 @@ class AIMLChatConfig(OpenAIGPTConfig): ) # type: ignore dynamic_api_key = api_key or get_secret_str("AIML_API_KEY") return api_base, dynamic_api_key - pass \ No newline at end of file + + pass diff --git a/litellm/llms/aiml/image_generation/cost_calculator.py b/litellm/llms/aiml/image_generation/cost_calculator.py index 1fecfb6a9a..4442f57c55 100644 --- a/litellm/llms/aiml/image_generation/cost_calculator.py +++ b/litellm/llms/aiml/image_generation/cost_calculator.py @@ -22,4 +22,6 @@ def cost_calculator( num_images = len(image_response.data) return output_cost_per_image * num_images else: - raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}") + raise ValueError( + f"image_response must be of type ImageResponse got type={type(image_response)}" + ) diff --git a/litellm/llms/aiml/image_generation/transformation.py b/litellm/llms/aiml/image_generation/transformation.py index d8f3e23fe7..39b1cc742d 100644 --- a/litellm/llms/aiml/image_generation/transformation.py +++ b/litellm/llms/aiml/image_generation/transformation.py @@ -24,19 +24,15 @@ else: class AimlImageGenerationConfig(BaseImageGenerationConfig): DEFAULT_BASE_URL: str = "https://api.aimlapi.com" IMAGE_GENERATION_ENDPOINT: str = "v1/images/generations" - + def get_supported_openai_params( self, model: str ) -> List[OpenAIImageGenerationOptionalParams]: """ https://api.aimlapi.com/v1/images/generations """ - return [ - "n", - "response_format", - "size" - ] - + return ["n", "response_format", "size"] + def map_openai_params( self, non_default_params: dict, @@ -45,7 +41,7 @@ class AimlImageGenerationConfig(BaseImageGenerationConfig): drop_params: bool, ) -> dict: supported_params = self.get_supported_openai_params(model) - + for k in non_default_params.keys(): if k not in optional_params.keys(): if k in supported_params: @@ -53,7 +49,7 @@ class AimlImageGenerationConfig(BaseImageGenerationConfig): if k == "n": optional_params["num_images"] = non_default_params[k] elif k == "response_format": - optional_params["output_format"] = non_default_params[k] + optional_params["output_format"] = non_default_params[k] elif k == "size": # Map OpenAI size format to AI/ML image_size size_value = non_default_params[k] @@ -61,7 +57,10 @@ class AimlImageGenerationConfig(BaseImageGenerationConfig): # Handle standard OpenAI sizes like "1024x1024" if "x" in size_value: width, height = map(int, size_value.split("x")) - optional_params["image_size"] = {"width": width, "height": height} + optional_params["image_size"] = { + "width": width, + "height": height, + } else: # Pass through predefined sizes optional_params["image_size"] = size_value @@ -91,9 +90,7 @@ class AimlImageGenerationConfig(BaseImageGenerationConfig): Get the complete url for the request """ complete_url: str = ( - api_base - or get_secret_str("AIML_API_BASE") - or self.DEFAULT_BASE_URL + api_base or get_secret_str("AIML_API_BASE") or self.DEFAULT_BASE_URL ) complete_url = complete_url.rstrip("/") @@ -114,15 +111,15 @@ class AimlImageGenerationConfig(BaseImageGenerationConfig): api_base: Optional[str] = None, ) -> dict: final_api_key: Optional[str] = ( - api_key or - get_secret_str("AIML_API_KEY") or - get_secret_str("AIMLAPI_KEY") # Alternative name + api_key + or get_secret_str("AIML_API_KEY") + or get_secret_str("AIMLAPI_KEY") # Alternative name ) if not final_api_key: raise ValueError("AIML_API_KEY or AIMLAPI_KEY is not set") - + headers["Authorization"] = f"Bearer {final_api_key}" - headers["Content-Type"] = "application/json" + headers["Content-Type"] = "application/json" return headers def transform_image_generation_request( @@ -138,10 +135,12 @@ class AimlImageGenerationConfig(BaseImageGenerationConfig): https://api.aimlapi.com/v1/images/generations """ - aiml_image_generation_request_body: AimlImageGenerationRequestParams = AimlImageGenerationRequestParams( - prompt=prompt, - model=model, - **optional_params, + aiml_image_generation_request_body: AimlImageGenerationRequestParams = ( + AimlImageGenerationRequestParams( + prompt=prompt, + model=model, + **optional_params, + ) ) return dict(aiml_image_generation_request_body) @@ -171,53 +170,65 @@ class AimlImageGenerationConfig(BaseImageGenerationConfig): status_code=raw_response.status_code, headers=raw_response.headers, ) - + if not model_response.data: model_response.data = [] - + # AI/ML API can return images in multiple formats: # 1. Top-level data array with url (OpenAI-like format) # 2. output.choices array with image_base64 # 3. images array with url (and optional width, height, content_type) - + if "data" in response_data and isinstance(response_data["data"], list): # Handle OpenAI-like format: {"data": [{"url": "...", "width": 1024, "height": 768, "content_type": "image/jpeg"}]} for image in response_data["data"]: if "url" in image: - model_response.data.append(ImageObject( - b64_json=None, - url=image["url"], - revised_prompt=image.get("revised_prompt"), - )) + model_response.data.append( + ImageObject( + b64_json=None, + url=image["url"], + revised_prompt=image.get("revised_prompt"), + ) + ) elif "b64_json" in image or "image_base64" in image: - model_response.data.append(ImageObject( - b64_json=image.get("b64_json") or image.get("image_base64"), - url=None, - revised_prompt=image.get("revised_prompt"), - )) + model_response.data.append( + ImageObject( + b64_json=image.get("b64_json") or image.get("image_base64"), + url=None, + revised_prompt=image.get("revised_prompt"), + ) + ) elif "output" in response_data and "choices" in response_data["output"]: for choice in response_data["output"]["choices"]: if "image_base64" in choice: - model_response.data.append(ImageObject( - b64_json=choice["image_base64"], - url=None, - )) + model_response.data.append( + ImageObject( + b64_json=choice["image_base64"], + url=None, + ) + ) elif "url" in choice: - model_response.data.append(ImageObject( - b64_json=None, - url=choice["url"], - )) + model_response.data.append( + ImageObject( + b64_json=None, + url=choice["url"], + ) + ) elif "images" in response_data: # Handle alternative format: {"images": [{"url": "...", "width": 1024, "height": 768, "content_type": "image/jpeg"}]} for image in response_data["images"]: if "url" in image: - model_response.data.append(ImageObject( - b64_json=None, - url=image["url"], - )) + model_response.data.append( + ImageObject( + b64_json=None, + url=image["url"], + ) + ) elif "image_base64" in image: - model_response.data.append(ImageObject( - b64_json=image["image_base64"], - url=None, - )) + model_response.data.append( + ImageObject( + b64_json=image["image_base64"], + url=None, + ) + ) return model_response diff --git a/litellm/llms/amazon_nova/chat/transformation.py b/litellm/llms/amazon_nova/chat/transformation.py index 6d321e298b..0fd08e6287 100644 --- a/litellm/llms/amazon_nova/chat/transformation.py +++ b/litellm/llms/amazon_nova/chat/transformation.py @@ -56,7 +56,7 @@ class AmazonNovaChatConfig(OpenAILikeChatConfig): or get_secret_str("AMAZON_NOVA_API_BASE") or "https://api.nova.amazon.com/v1" ) # type: ignore - + # Get API key from multiple sources key = ( api_key @@ -65,7 +65,7 @@ class AmazonNovaChatConfig(OpenAILikeChatConfig): or litellm.api_key ) return api_base, key - + def get_supported_openai_params(self, model: str) -> List: return [ "top_p", @@ -78,7 +78,7 @@ class AmazonNovaChatConfig(OpenAILikeChatConfig): "stream_options", "tools", "tool_choice", - "reasoning_effort" + "reasoning_effort", ] def transform_response( @@ -112,4 +112,4 @@ class AmazonNovaChatConfig(OpenAILikeChatConfig): # Storing amazon_nova in the model response for easier cost calculation later setattr(model_response, "model", "amazon-nova/" + model) - return model_response \ No newline at end of file + return model_response diff --git a/litellm/llms/amazon_nova/cost_calculation.py b/litellm/llms/amazon_nova/cost_calculation.py index 9d9cedde87..857369b76e 100644 --- a/litellm/llms/amazon_nova/cost_calculation.py +++ b/litellm/llms/amazon_nova/cost_calculation.py @@ -18,4 +18,4 @@ def cost_per_token(model: str, usage: "Usage") -> Tuple[float, float]: """ return generic_cost_per_token( model=model, usage=usage, custom_llm_provider="amazon_nova" - ) \ No newline at end of file + ) diff --git a/litellm/llms/anthropic/batches/__init__.py b/litellm/llms/anthropic/batches/__init__.py index 66d1a8f77f..dd9ae5273b 100644 --- a/litellm/llms/anthropic/batches/__init__.py +++ b/litellm/llms/anthropic/batches/__init__.py @@ -2,4 +2,3 @@ from .handler import AnthropicBatchesHandler from .transformation import AnthropicBatchesConfig __all__ = ["AnthropicBatchesHandler", "AnthropicBatchesConfig"] - diff --git a/litellm/llms/anthropic/batches/handler.py b/litellm/llms/anthropic/batches/handler.py index fd303e60af..52bf29a551 100644 --- a/litellm/llms/anthropic/batches/handler.py +++ b/litellm/llms/anthropic/batches/handler.py @@ -24,7 +24,7 @@ from .transformation import AnthropicBatchesConfig class AnthropicBatchesHandler: """ Handler for Anthropic Message Batches API. - + Supports: - retrieve_batch() - Retrieve batch status and information """ @@ -44,7 +44,7 @@ class AnthropicBatchesHandler: ) -> LiteLLMBatch: """ Async: Retrieve a batch from Anthropic. - + Args: batch_id: The batch ID to retrieve api_base: Anthropic API base URL @@ -52,20 +52,23 @@ class AnthropicBatchesHandler: timeout: Request timeout max_retries: Max retry attempts (unused for now) logging_obj: Optional logging object - + Returns: LiteLLMBatch: Batch information in OpenAI format """ # Resolve API credentials api_base = api_base or self.anthropic_model_info.get_api_base(api_base) api_key = api_key or self.anthropic_model_info.get_api_key() - + if not api_key: raise ValueError("Missing Anthropic API Key") - + # Create a minimal logging object if not provided if logging_obj is None: - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObjClass + from litellm.litellm_core_utils.litellm_logging import ( + Logging as LiteLLMLoggingObjClass, + ) + logging_obj = LiteLLMLoggingObjClass( model="anthropic/unknown", messages=[], @@ -75,7 +78,7 @@ class AnthropicBatchesHandler: litellm_call_id=f"batch_retrieve_{batch_id}", function_id="batch_retrieve", ) - + # Get the complete URL for batch retrieval retrieve_url = self.provider_config.get_retrieve_batch_url( api_base=api_base, @@ -83,7 +86,7 @@ class AnthropicBatchesHandler: optional_params={}, litellm_params={}, ) - + # Validate environment and get headers headers = self.provider_config.validate_environment( headers={}, @@ -106,12 +109,9 @@ class AnthropicBatchesHandler: ) # Make the request async_client = get_async_httpx_client(llm_provider=LlmProviders.ANTHROPIC) - response = await async_client.get( - url=retrieve_url, - headers=headers - ) + response = await async_client.get(url=retrieve_url, headers=headers) response.raise_for_status() - + # Transform response to LiteLLM format return self.provider_config.transform_retrieve_batch_response( model=None, @@ -132,7 +132,7 @@ class AnthropicBatchesHandler: ) -> Union[LiteLLMBatch, Coroutine[Any, Any, LiteLLMBatch]]: """ Retrieve a batch from Anthropic. - + Args: _is_async: Whether to run asynchronously batch_id: The batch ID to retrieve @@ -141,7 +141,7 @@ class AnthropicBatchesHandler: timeout: Request timeout max_retries: Max retry attempts (unused for now) logging_obj: Optional logging object - + Returns: LiteLLMBatch or Coroutine: Batch information in OpenAI format """ @@ -165,4 +165,3 @@ class AnthropicBatchesHandler: logging_obj=logging_obj, ) ) - diff --git a/litellm/llms/anthropic/batches/transformation.py b/litellm/llms/anthropic/batches/transformation.py index 750dd002ff..699f133f0f 100644 --- a/litellm/llms/anthropic/batches/transformation.py +++ b/litellm/llms/anthropic/batches/transformation.py @@ -84,7 +84,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig): ) -> Union[bytes, str, Dict[str, Any]]: """ Transform the batch creation request to Anthropic format. - + Not currently implemented - placeholder to satisfy abstract base class. """ raise NotImplementedError("Batch creation not yet implemented for Anthropic") @@ -98,7 +98,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig): ) -> LiteLLMBatch: """ Transform Anthropic MessageBatch creation response to LiteLLM format. - + Not currently implemented - placeholder to satisfy abstract base class. """ raise NotImplementedError("Batch creation not yet implemented for Anthropic") @@ -112,13 +112,13 @@ class AnthropicBatchesConfig(BaseBatchesConfig): ) -> str: """ Get the complete URL for batch retrieval request. - + Args: api_base: Base API URL (optional, will use default if not provided) batch_id: Batch ID to retrieve optional_params: Optional parameters litellm_params: LiteLLM parameters - + Returns: Complete URL for Anthropic batch retrieval: {api_base}/v1/messages/batches/{batch_id} """ @@ -133,7 +133,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig): ) -> Union[bytes, str, Dict[str, Any]]: """ Transform batch retrieval request for Anthropic. - + For Anthropic, the URL is constructed by get_retrieve_batch_url(), so this method returns an empty dict (no additional request params needed). """ @@ -156,9 +156,21 @@ class AnthropicBatchesConfig(BaseBatchesConfig): # Map Anthropic MessageBatch to OpenAI Batch format batch_id = response_data.get("id", "") processing_status = response_data.get("processing_status", "in_progress") - + # Map Anthropic processing_status to OpenAI status - status_mapping: Dict[str, Literal["validating", "failed", "in_progress", "finalizing", "completed", "expired", "cancelling", "cancelled"]] = { + status_mapping: Dict[ + str, + Literal[ + "validating", + "failed", + "in_progress", + "finalizing", + "completed", + "expired", + "cancelling", + "cancelled", + ], + ] = { "in_progress": "in_progress", "canceling": "cancelling", "ended": "completed", @@ -171,7 +183,8 @@ class AnthropicBatchesConfig(BaseBatchesConfig): return None try: from datetime import datetime - dt = datetime.fromisoformat(ts_str.replace('Z', '+00:00')) + + dt = datetime.fromisoformat(ts_str.replace("Z", "+00:00")) return int(dt.timestamp()) except Exception: return None @@ -185,14 +198,17 @@ class AnthropicBatchesConfig(BaseBatchesConfig): # Extract request counts request_counts_data = response_data.get("request_counts", {}) from openai.types.batch import BatchRequestCounts + request_counts = BatchRequestCounts( - total=sum([ - request_counts_data.get("processing", 0), - request_counts_data.get("succeeded", 0), - request_counts_data.get("errored", 0), - request_counts_data.get("canceled", 0), - request_counts_data.get("expired", 0), - ]), + total=sum( + [ + request_counts_data.get("processing", 0), + request_counts_data.get("succeeded", 0), + request_counts_data.get("errored", 0), + request_counts_data.get("canceled", 0), + request_counts_data.get("expired", 0), + ] + ), completed=request_counts_data.get("succeeded", 0), failed=request_counts_data.get("errored", 0), ) @@ -214,8 +230,12 @@ class AnthropicBatchesConfig(BaseBatchesConfig): completed_at=ended_at if processing_status == "ended" else None, failed_at=None, expired_at=archived_at if archived_at else None, - cancelling_at=cancel_initiated_at if processing_status == "canceling" else None, - cancelled_at=ended_at if processing_status == "canceling" and ended_at else None, + cancelling_at=cancel_initiated_at + if processing_status == "canceling" + else None, + cancelled_at=ended_at + if processing_status == "canceling" and ended_at + else None, request_counts=request_counts, metadata={}, ) @@ -232,7 +252,9 @@ class AnthropicBatchesConfig(BaseBatchesConfig): else: headers_obj = headers if isinstance(headers, Headers) else None - return AnthropicError(status_code=status_code, message=error_message, headers=headers_obj) + return AnthropicError( + status_code=status_code, message=error_message, headers=headers_obj + ) def transform_response( self, diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index a6df346e8a..0bc0777e37 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -75,11 +75,12 @@ class AnthropicMessagesHandler(BaseTranslation): if messages is None: return data - chat_completion_compatible_request, _tool_name_mapping = ( - LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( - # Use a shallow copy to avoid mutating request data (pop on litellm_metadata). - anthropic_message_request=cast(AnthropicMessagesRequest, data.copy()) - ) + ( + chat_completion_compatible_request, + _tool_name_mapping, + ) = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( + # Use a shallow copy to avoid mutating request data (pop on litellm_metadata). + anthropic_message_request=cast(AnthropicMessagesRequest, data.copy()) ) structured_messages = chat_completion_compatible_request.get("messages", []) @@ -205,7 +206,7 @@ class AnthropicMessagesHandler(BaseTranslation): openai_tools = self.adapter.translate_anthropic_tools_to_openai( tools=cast(List[AllAnthropicToolsValues], tools) ) - tools_to_check.extend(openai_tools) # type: ignore + tools_to_check.extend(openai_tools) # type: ignore async def _apply_guardrail_responses_to_input( self, @@ -375,10 +376,12 @@ class AnthropicMessagesHandler(BaseTranslation): has_ended = self._check_streaming_has_ended(responses_so_far) if has_ended: # build the model response from the responses_so_far - built_response = AnthropicPassthroughLoggingHandler._build_complete_streaming_response( - all_chunks=responses_so_far, - litellm_logging_obj=cast("LiteLLMLoggingObj", litellm_logging_obj), - model="", + built_response = ( + AnthropicPassthroughLoggingHandler._build_complete_streaming_response( + all_chunks=responses_so_far, + litellm_logging_obj=cast("LiteLLMLoggingObj", litellm_logging_obj), + model="", + ) ) # Check if model_response is valid and has choices before accessing @@ -407,7 +410,9 @@ class AnthropicMessagesHandler(BaseTranslation): logging_obj=litellm_logging_obj, ) else: - verbose_proxy_logger.debug("Skipping output guardrail - model response has no choices") + verbose_proxy_logger.debug( + "Skipping output guardrail - model response has no choices" + ) return responses_so_far string_so_far = self.get_streaming_string_so_far(responses_so_far) diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index f51adf9610..72cc7ecd9c 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -165,7 +165,10 @@ def make_sync_call( ) completion_stream = ModelResponseIterator( - streaming_response=response.iter_lines(), sync_stream=True, json_mode=json_mode, speed=speed + streaming_response=response.iter_lines(), + sync_stream=True, + json_mode=json_mode, + speed=speed, ) # LOGGING @@ -497,7 +500,11 @@ class AnthropicChatCompletion(BaseLLM): class ModelResponseIterator: def __init__( - self, streaming_response, sync_stream: bool, json_mode: Optional[bool] = False, speed: Optional[str] = None + self, + streaming_response, + sync_stream: bool, + json_mode: Optional[bool] = False, + speed: Optional[str] = None, ): self.streaming_response = streaming_response self.response_iterator = self.streaming_response @@ -525,7 +532,7 @@ class ModelResponseIterator: # Accumulate web_search_tool_result blocks for multi-turn reconstruction # See: https://github.com/BerriAI/litellm/issues/17737 self.web_search_results: List[Dict[str, Any]] = [] - + # Accumulate compaction blocks for multi-turn reconstruction self.compaction_blocks: List[Dict[str, Any]] = [] @@ -554,10 +561,14 @@ class ModelResponseIterator: def _handle_usage(self, anthropic_usage_chunk: Union[dict, UsageDelta]) -> Usage: return AnthropicConfig().calculate_usage( - usage_object=cast(dict, anthropic_usage_chunk), reasoning_content=None, speed=self.speed + usage_object=cast(dict, anthropic_usage_chunk), + reasoning_content=None, + speed=self.speed, ) - def _content_block_delta_helper(self, chunk: dict) -> Tuple[ + def _content_block_delta_helper( + self, chunk: dict + ) -> Tuple[ str, Optional[ChatCompletionToolCallChunk], List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]], @@ -608,11 +619,14 @@ class ModelResponseIterator: ) ] provider_specific_fields["thinking_blocks"] = thinking_blocks - elif "content" in content_block["delta"] and content_block["delta"].get("type") == "compaction_delta": + elif ( + "content" in content_block["delta"] + and content_block["delta"].get("type") == "compaction_delta" + ): # Handle compaction delta provider_specific_fields["compaction_delta"] = { "type": "compaction_delta", - "content": content_block["delta"]["content"] + "content": content_block["delta"]["content"], } return text, tool_use, thinking_blocks, provider_specific_fields @@ -710,10 +724,15 @@ class ModelResponseIterator: content_block_start = self.get_content_block_start(chunk=chunk) self.content_blocks = [] # reset content blocks when new block starts # Track current content block type for filtering deltas - self.current_content_block_type = content_block_start["content_block"]["type"] + self.current_content_block_type = content_block_start["content_block"][ + "type" + ] if content_block_start["content_block"]["type"] == "text": text = content_block_start["content_block"]["text"] - elif content_block_start["content_block"]["type"] == "tool_use" or content_block_start["content_block"]["type"] == "server_tool_use": + elif ( + content_block_start["content_block"]["type"] == "tool_use" + or content_block_start["content_block"]["type"] == "server_tool_use" + ): self.tool_index += 1 # Use empty string for arguments in content_block_start - actual arguments # come in subsequent content_block_delta chunks and get accumulated. @@ -746,21 +765,23 @@ class ModelResponseIterator: elif content_block_start["content_block"]["type"] == "compaction": # Handle compaction blocks # The full content comes in content_block_start - self.compaction_blocks.append( - content_block_start["content_block"] - ) - provider_specific_fields["compaction_blocks"] = ( - self.compaction_blocks - ) + self.compaction_blocks.append(content_block_start["content_block"]) + provider_specific_fields[ + "compaction_blocks" + ] = self.compaction_blocks provider_specific_fields["compaction_start"] = { "type": "compaction", - "content": content_block_start["content_block"].get("content", "") + "content": content_block_start["content_block"].get( + "content", "" + ), } - elif content_block_start["content_block"]["type"].endswith("_tool_result"): + elif content_block_start["content_block"]["type"].endswith( + "_tool_result" + ): # Handle all tool result types (web_search, bash_code_execution, text_editor, etc.) content_type = content_block_start["content_block"]["type"] - + # Special handling for web_search_tool_result for backwards compatibility if content_type == "web_search_tool_result": # Capture web_search_tool_result for multi-turn reconstruction @@ -769,9 +790,9 @@ class ModelResponseIterator: self.web_search_results.append( content_block_start["content_block"] ) - provider_specific_fields["web_search_results"] = ( - self.web_search_results - ) + provider_specific_fields[ + "web_search_results" + ] = self.web_search_results elif content_type == "web_fetch_tool_result": # Capture web_fetch_tool_result for multi-turn reconstruction # The full content comes in content_block_start, not in deltas @@ -779,9 +800,9 @@ class ModelResponseIterator: self.web_search_results.append( content_block_start["content_block"] ) - provider_specific_fields["web_search_results"] = ( - self.web_search_results - ) + provider_specific_fields[ + "web_search_results" + ] = self.web_search_results elif content_type != "tool_search_tool_result": # Handle other tool results (code execution, etc.) # Skip tool_search_tool_result as it's internal metadata @@ -932,7 +953,9 @@ class ModelResponseIterator: return text, tool_use - def _handle_message_delta(self, chunk: dict) -> Tuple[str, Optional[Usage], Optional[Dict[str, Any]]]: + def _handle_message_delta( + self, chunk: dict + ) -> Tuple[str, Optional[Usage], Optional[Dict[str, Any]]]: """ Handle message_delta event for finish_reason, usage, and container. @@ -1052,7 +1075,9 @@ class ModelResponseIterator: except StopIteration: raise StopIteration except ValueError as e: - raise RuntimeError(f"Error parsing chunk: {e},\nReceived chunk: {chunk}") + raise RuntimeError( + f"Error parsing chunk: {e},\nReceived chunk: {chunk}" + ) # Async iterator def __aiter__(self): @@ -1101,7 +1126,9 @@ class ModelResponseIterator: except StopAsyncIteration: raise StopAsyncIteration except ValueError as e: - raise RuntimeError(f"Error parsing chunk: {e},\nReceived chunk: {chunk}") + raise RuntimeError( + f"Error parsing chunk: {e},\nReceived chunk: {chunk}" + ) def convert_str_chunk_to_generic_chunk(self, chunk: str) -> ModelResponseStream: """ diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index fd1859f7d1..1b912bfc2a 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -173,8 +173,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): """Check if the model is specifically Claude Opus 4.6.""" model_lower = model.lower() return any( - v in model_lower - for v in ("opus-4-6", "opus_4_6", "opus-4.6", "opus_4.6") + v in model_lower for v in ("opus-4-6", "opus_4_6", "opus-4.6", "opus_4.6") ) def get_supported_openai_params(self, model: str): @@ -957,11 +956,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if mcp_servers: optional_params["mcp_servers"] = mcp_servers elif param == "tool_choice" or param == "parallel_tool_calls": - _tool_choice: Optional[AnthropicMessagesToolChoice] = ( - self._map_tool_choice( - tool_choice=non_default_params.get("tool_choice"), - parallel_tool_use=non_default_params.get("parallel_tool_calls"), - ) + _tool_choice: Optional[ + AnthropicMessagesToolChoice + ] = self._map_tool_choice( + tool_choice=non_default_params.get("tool_choice"), + parallel_tool_use=non_default_params.get("parallel_tool_calls"), ) if _tool_choice is not None: @@ -1059,9 +1058,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): self.map_openai_context_management_to_anthropic(value) ) if anthropic_context_management is not None: - optional_params["context_management"] = ( - anthropic_context_management - ) + optional_params[ + "context_management" + ] = anthropic_context_management elif param == "speed" and isinstance(value, str): # Pass through Anthropic-specific speed parameter for fast mode optional_params["speed"] = value @@ -1135,9 +1134,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): text=system_message_block["content"], ) if "cache_control" in system_message_block: - anthropic_system_message_content["cache_control"] = ( - system_message_block["cache_control"] - ) + anthropic_system_message_content[ + "cache_control" + ] = system_message_block["cache_control"] anthropic_system_message_list.append( anthropic_system_message_content ) @@ -1161,9 +1160,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) ) if "cache_control" in _content: - anthropic_system_message_content["cache_control"] = ( - _content["cache_control"] - ) + anthropic_system_message_content[ + "cache_control" + ] = _content["cache_control"] anthropic_system_message_list.append( anthropic_system_message_content @@ -1460,7 +1459,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) return _message - def extract_response_content(self, completion_response: dict) -> Tuple[ + def extract_response_content( + self, completion_response: dict + ) -> Tuple[ str, Optional[List[Any]], Optional[ diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 8f196966dc..ac35246787 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -31,6 +31,7 @@ def is_anthropic_oauth_key(value: Optional[str]) -> bool: value = value[7:] return value.startswith(ANTHROPIC_OAUTH_TOKEN_PREFIX) + def _merge_beta_headers(existing: Optional[str], new_beta: str) -> str: """Merge a new beta value into an existing comma-separated anthropic-beta header.""" if not existing: @@ -244,8 +245,14 @@ class AnthropicModelInfo(BaseLLMModelInfo): return any( v in model_lower for v in ( - "opus-4-6", "opus_4_6", "opus-4.6", "opus_4.6", - "sonnet-4-6", "sonnet_4_6", "sonnet-4.6", "sonnet_4.6", + "opus-4-6", + "opus_4_6", + "opus-4.6", + "opus_4.6", + "sonnet-4-6", + "sonnet_4_6", + "sonnet-4.6", + "sonnet_4.6", ) ) diff --git a/litellm/llms/anthropic/completion/transformation.py b/litellm/llms/anthropic/completion/transformation.py index a8798cd5d0..576ddb57fb 100644 --- a/litellm/llms/anthropic/completion/transformation.py +++ b/litellm/llms/anthropic/completion/transformation.py @@ -55,9 +55,9 @@ class AnthropicTextConfig(BaseConfig): to pass metadata to anthropic, it's {"user_id": "any-relevant-information"} """ - max_tokens_to_sample: Optional[int] = ( - litellm.max_tokens - ) # anthropic requires a default + max_tokens_to_sample: Optional[ + int + ] = litellm.max_tokens # anthropic requires a default stop_sequences: Optional[list] = None temperature: Optional[int] = None top_p: Optional[int] = None diff --git a/litellm/llms/anthropic/cost_calculation.py b/litellm/llms/anthropic/cost_calculation.py index cf9b18c464..3882d8f978 100644 --- a/litellm/llms/anthropic/cost_calculation.py +++ b/litellm/llms/anthropic/cost_calculation.py @@ -29,9 +29,13 @@ def _compute_cache_only_cost(model_info: "ModelInfo", usage: "Usage") -> float: return 0.0 prompt_tokens_details = _parse_prompt_tokens_details(usage) - _, _, cache_creation_cost, cache_creation_cost_above_1hr, cache_read_cost = ( - _get_token_base_cost(model_info=model_info, usage=usage) - ) + ( + _, + _, + cache_creation_cost, + cache_creation_cost_above_1hr, + cache_read_cost, + ) = _get_token_base_cost(model_info=model_info, usage=usage) cache_cost = float(prompt_tokens_details["cache_hit_tokens"]) * cache_read_cost @@ -68,7 +72,9 @@ def cost_per_token(model: str, usage: "Usage") -> Tuple[float, float]: # Apply provider_specific_entry multipliers for geo/speed routing try: - model_info = litellm.get_model_info(model=model, custom_llm_provider="anthropic") + model_info = litellm.get_model_info( + model=model, custom_llm_provider="anthropic" + ) provider_specific_entry: dict = model_info.get("provider_specific_entry") or {} multiplier = 1.0 @@ -77,9 +83,7 @@ def cost_per_token(model: str, usage: "Usage") -> Tuple[float, float]: and usage.inference_geo and usage.inference_geo.lower() not in ["global", "not_available"] ): - multiplier *= provider_specific_entry.get( - usage.inference_geo.lower(), 1.0 - ) + multiplier *= provider_specific_entry.get(usage.inference_geo.lower(), 1.0) if hasattr(usage, "speed") and usage.speed == "fast": multiplier *= provider_specific_entry.get("fast", 1.0) diff --git a/litellm/llms/anthropic/count_tokens/handler.py b/litellm/llms/anthropic/count_tokens/handler.py index 07481917af..4d0af0b36c 100644 --- a/litellm/llms/anthropic/count_tokens/handler.py +++ b/litellm/llms/anthropic/count_tokens/handler.py @@ -82,7 +82,9 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig): ) # Use provided timeout or fall back to litellm.request_timeout - request_timeout = timeout if timeout is not None else litellm.request_timeout + request_timeout = ( + timeout if timeout is not None else litellm.request_timeout + ) response = await async_client.post( endpoint_url, diff --git a/litellm/llms/anthropic/count_tokens/transformation.py b/litellm/llms/anthropic/count_tokens/transformation.py index 2d3f5b1942..ad5bbbda25 100644 --- a/litellm/llms/anthropic/count_tokens/transformation.py +++ b/litellm/llms/anthropic/count_tokens/transformation.py @@ -73,14 +73,10 @@ class AnthropicCountTokensConfig: "anthropic-version": "2023-06-01", "anthropic-beta": ANTHROPIC_TOKEN_COUNTING_BETA_VERSION, } - headers, _ = optionally_handle_anthropic_oauth( - headers=headers, api_key=api_key - ) + headers, _ = optionally_handle_anthropic_oauth(headers=headers, api_key=api_key) return headers - def validate_request( - self, model: str, messages: List[Dict[str, Any]] - ) -> None: + def validate_request(self, model: str, messages: List[Dict[str, Any]]) -> None: """ Validate the incoming count tokens request. diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py index 73e74c228b..8b1b21a0f9 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py @@ -65,7 +65,9 @@ class LiteLLMMessagesToCompletionTransformationHandler: model = completion_kwargs.get("model") try: - model_info = get_model_info(model=cast(str, model), custom_llm_provider=custom_llm_provider) + model_info = get_model_info( + model=cast(str, model), custom_llm_provider=custom_llm_provider + ) if model_info and model_info.get("supports_reasoning") is False: # Model doesn't support reasoning/responses API, don't route return @@ -75,7 +77,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: if isinstance(model, str) and model and not model.startswith("responses/"): # Prefix model with "responses/" to route to OpenAI Responses API completion_kwargs["model"] = f"responses/{model}" - + reasoning_effort = completion_kwargs.get("reasoning_effort") if isinstance(reasoning_effort, str) and reasoning_effort: completion_kwargs["reasoning_effort"] = { @@ -148,7 +150,10 @@ class LiteLLMMessagesToCompletionTransformationHandler: if output_format: request_data["output_format"] = output_format - openai_request, tool_name_mapping = ANTHROPIC_ADAPTER.translate_completion_input_params_with_tool_mapping( + ( + openai_request, + tool_name_mapping, + ) = ANTHROPIC_ADAPTER.translate_completion_input_params_with_tool_mapping( request_data ) @@ -210,24 +215,25 @@ class LiteLLMMessagesToCompletionTransformationHandler: **kwargs, ) -> Union[AnthropicMessagesResponse, AsyncIterator]: """Handle non-Anthropic models asynchronously using the adapter""" - completion_kwargs, tool_name_mapping = ( - LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs( - max_tokens=max_tokens, - messages=messages, - model=model, - metadata=metadata, - stop_sequences=stop_sequences, - stream=stream, - system=system, - temperature=temperature, - thinking=thinking, - tool_choice=tool_choice, - tools=tools, - top_k=top_k, - top_p=top_p, - output_format=output_format, - extra_kwargs=kwargs, - ) + ( + completion_kwargs, + tool_name_mapping, + ) = LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs( + max_tokens=max_tokens, + messages=messages, + model=model, + metadata=metadata, + stop_sequences=stop_sequences, + stream=stream, + system=system, + temperature=temperature, + thinking=thinking, + tool_choice=tool_choice, + tools=tools, + top_k=top_k, + top_p=top_p, + output_format=output_format, + extra_kwargs=kwargs, ) completion_response = await litellm.acompletion(**completion_kwargs) @@ -244,11 +250,9 @@ class LiteLLMMessagesToCompletionTransformationHandler: return transformed_stream raise ValueError("Failed to transform streaming response") else: - anthropic_response = ( - ANTHROPIC_ADAPTER.translate_completion_output_params( - cast(ModelResponse, completion_response), - tool_name_mapping=tool_name_mapping, - ) + anthropic_response = ANTHROPIC_ADAPTER.translate_completion_output_params( + cast(ModelResponse, completion_response), + tool_name_mapping=tool_name_mapping, ) if anthropic_response is not None: return anthropic_response @@ -297,24 +301,25 @@ class LiteLLMMessagesToCompletionTransformationHandler: **kwargs, ) - completion_kwargs, tool_name_mapping = ( - LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs( - max_tokens=max_tokens, - messages=messages, - model=model, - metadata=metadata, - stop_sequences=stop_sequences, - stream=stream, - system=system, - temperature=temperature, - thinking=thinking, - tool_choice=tool_choice, - tools=tools, - top_k=top_k, - top_p=top_p, - output_format=output_format, - extra_kwargs=kwargs, - ) + ( + completion_kwargs, + tool_name_mapping, + ) = LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs( + max_tokens=max_tokens, + messages=messages, + model=model, + metadata=metadata, + stop_sequences=stop_sequences, + stream=stream, + system=system, + temperature=temperature, + thinking=thinking, + tool_choice=tool_choice, + tools=tools, + top_k=top_k, + top_p=top_p, + output_format=output_format, + extra_kwargs=kwargs, ) completion_response = litellm.completion(**completion_kwargs) @@ -331,11 +336,9 @@ class LiteLLMMessagesToCompletionTransformationHandler: return transformed_stream raise ValueError("Failed to transform streaming response") else: - anthropic_response = ( - ANTHROPIC_ADAPTER.translate_completion_output_params( - cast(ModelResponse, completion_response), - tool_name_mapping=tool_name_mapping, - ) + anthropic_response = ANTHROPIC_ADAPTER.translate_completion_output_params( + cast(ModelResponse, completion_response), + tool_name_mapping=tool_name_mapping, ) if anthropic_response is not None: return anthropic_response diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index 7f17526e75..6bddad09f2 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -261,19 +261,37 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): # Add usage to the held chunk uncached_input_tokens = chunk.usage.prompt_tokens or 0 - if hasattr(chunk.usage, "prompt_tokens_details") and chunk.usage.prompt_tokens_details: - cached_tokens = getattr(chunk.usage.prompt_tokens_details, "cached_tokens", 0) or 0 + if ( + hasattr(chunk.usage, "prompt_tokens_details") + and chunk.usage.prompt_tokens_details + ): + cached_tokens = ( + getattr( + chunk.usage.prompt_tokens_details, "cached_tokens", 0 + ) + or 0 + ) uncached_input_tokens -= cached_tokens - + usage_dict: UsageDelta = { "input_tokens": uncached_input_tokens, "output_tokens": chunk.usage.completion_tokens or 0, } # Add cache tokens if available (for prompt caching support) - if hasattr(chunk.usage, "_cache_creation_input_tokens") and chunk.usage._cache_creation_input_tokens > 0: - usage_dict["cache_creation_input_tokens"] = chunk.usage._cache_creation_input_tokens - if hasattr(chunk.usage, "_cache_read_input_tokens") and chunk.usage._cache_read_input_tokens > 0: - usage_dict["cache_read_input_tokens"] = chunk.usage._cache_read_input_tokens + if ( + hasattr(chunk.usage, "_cache_creation_input_tokens") + and chunk.usage._cache_creation_input_tokens > 0 + ): + usage_dict[ + "cache_creation_input_tokens" + ] = chunk.usage._cache_creation_input_tokens + if ( + hasattr(chunk.usage, "_cache_read_input_tokens") + and chunk.usage._cache_read_input_tokens > 0 + ): + usage_dict[ + "cache_read_input_tokens" + ] = chunk.usage._cache_read_input_tokens merged_chunk["usage"] = usage_dict # Queue the merged chunk and reset @@ -439,12 +457,14 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): from typing import cast from litellm.types.llms.anthropic import ToolUseBlock - + tool_block = cast(ToolUseBlock, content_block_start) - + if tool_block.get("name"): truncated_name = tool_block["name"] - original_name = self.tool_name_mapping.get(truncated_name, truncated_name) + original_name = self.tool_name_mapping.get( + truncated_name, truncated_name + ) tool_block["name"] = original_name if block_type != self.current_content_block_type: @@ -458,7 +478,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): from typing import cast from litellm.types.llms.anthropic import ToolUseBlock - + tool_block = cast(ToolUseBlock, content_block_start) if tool_block.get("name"): self.current_content_block_type = block_type diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index a7362a9431..43a6fa8045 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -61,6 +61,7 @@ def create_tool_name_mapping( mapping[truncated_name] = original_name return mapping + from openai.types.chat.chat_completion_chunk import Choice as OpenAIStreamingChoice from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -172,10 +173,11 @@ class AnthropicAdapter: model=model, messages=messages, **kwargs ) - translated_body, tool_name_mapping = ( - LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( - anthropic_message_request=request_body - ) + ( + translated_body, + tool_name_mapping, + ) = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( + anthropic_message_request=request_body ) return translated_body, tool_name_mapping @@ -283,7 +285,11 @@ class LiteLLMAnthropicMessagesAdapter: model: Model name to check if cache_control should be preserved """ # TypedDict objects are dicts at runtime, so .get() works - cache_control = source.get("cache_control") if isinstance(source, dict) else getattr(source, "cache_control", None) + cache_control = ( + source.get("cache_control") + if isinstance(source, dict) + else getattr(source, "cache_control", None) + ) if cache_control and model and self.is_anthropic_claude_model(model): # TypedDict objects support dict operations at runtime # Use type ignore consistent with codebase pattern (see anthropic/chat/transformation.py:432) @@ -297,7 +303,15 @@ class LiteLLMAnthropicMessagesAdapter: """ Which anthropic params, we need to translate to the openai format. """ - return ["messages", "metadata", "system", "tool_choice", "tools", "thinking", "output_format"] + return [ + "messages", + "metadata", + "system", + "tool_choice", + "tools", + "thinking", + "output_format", + ] def _is_web_search_tool(self, tool: Dict[str, Any]) -> bool: """ @@ -350,13 +364,17 @@ class LiteLLMAnthropicMessagesAdapter: text_obj = ChatCompletionTextObject( type="text", text=content.get("text", "") ) - self._add_cache_control_if_applicable(content, text_obj, model) + self._add_cache_control_if_applicable( + content, text_obj, model + ) new_user_content_list.append(text_obj) # type: ignore elif content.get("type") == "image": # Convert Anthropic image format to OpenAI format source = content.get("source", {}) openai_image_url = ( - self._translate_anthropic_image_to_openai(cast(dict, source)) + self._translate_anthropic_image_to_openai( + cast(dict, source) + ) ) if openai_image_url: @@ -366,13 +384,17 @@ class LiteLLMAnthropicMessagesAdapter: image_obj = ChatCompletionImageObject( type="image_url", image_url=image_url_obj ) - self._add_cache_control_if_applicable(content, image_obj, model) + self._add_cache_control_if_applicable( + content, image_obj, model + ) new_user_content_list.append(image_obj) # type: ignore elif content.get("type") == "document": # Convert Anthropic document format (PDF, etc.) to OpenAI format source = content.get("source", {}) openai_image_url = ( - self._translate_anthropic_image_to_openai(cast(dict, source)) + self._translate_anthropic_image_to_openai( + cast(dict, source) + ) ) if openai_image_url: @@ -382,7 +404,9 @@ class LiteLLMAnthropicMessagesAdapter: doc_obj = ChatCompletionImageObject( type="image_url", image_url=image_url_obj ) - self._add_cache_control_if_applicable(content, doc_obj, model) + self._add_cache_control_if_applicable( + content, doc_obj, model + ) new_user_content_list.append(doc_obj) # type: ignore elif content.get("type") == "tool_result": if "content" not in content: @@ -391,7 +415,9 @@ class LiteLLMAnthropicMessagesAdapter: tool_call_id=content.get("tool_use_id", ""), content="", ) - self._add_cache_control_if_applicable(content, tool_result, model) + self._add_cache_control_if_applicable( + content, tool_result, model + ) tool_message_list.append(tool_result) # type: ignore[arg-type] elif isinstance(content.get("content"), str): tool_result = ChatCompletionToolMessage( @@ -399,7 +425,9 @@ class LiteLLMAnthropicMessagesAdapter: tool_call_id=content.get("tool_use_id", ""), content=str(content.get("content", "")), ) - self._add_cache_control_if_applicable(content, tool_result, model) + self._add_cache_control_if_applicable( + content, tool_result, model + ) tool_message_list.append(tool_result) # type: ignore[arg-type] elif isinstance(content.get("content"), list): # Combine all content items into a single tool message @@ -416,7 +444,9 @@ class LiteLLMAnthropicMessagesAdapter: tool_call_id=content.get("tool_use_id", ""), content=c, ) - self._add_cache_control_if_applicable(content, tool_result, model) + self._add_cache_control_if_applicable( + content, tool_result, model + ) tool_message_list.append(tool_result) # type: ignore[arg-type] elif isinstance(c, dict): if c.get("type") == "text": @@ -427,7 +457,9 @@ class LiteLLMAnthropicMessagesAdapter: ), content=c.get("text", ""), ) - self._add_cache_control_if_applicable(content, tool_result, model) + self._add_cache_control_if_applicable( + content, tool_result, model + ) tool_message_list.append(tool_result) # type: ignore[arg-type] elif c.get("type") == "image": source = c.get("source", {}) @@ -444,7 +476,9 @@ class LiteLLMAnthropicMessagesAdapter: ), content=openai_image_url, ) - self._add_cache_control_if_applicable(content, tool_result, model) + self._add_cache_control_if_applicable( + content, tool_result, model + ) tool_message_list.append(tool_result) # type: ignore[arg-type] else: # For multiple content items, combine into a single tool message @@ -494,7 +528,9 @@ class LiteLLMAnthropicMessagesAdapter: tool_call_id=content.get("tool_use_id", ""), content=combined_content_parts, # type: ignore ) - self._add_cache_control_if_applicable(content, tool_result, model) + self._add_cache_control_if_applicable( + content, tool_result, model + ) tool_message_list.append(tool_result) # type: ignore[arg-type] if len(tool_message_list) > 0: @@ -508,7 +544,9 @@ class LiteLLMAnthropicMessagesAdapter: ## ASSISTANT MESSAGE ## assistant_message_str: Optional[str] = None - assistant_content_list: List[Dict[str, Any]] = [] # For content blocks with cache_control + assistant_content_list: List[ + Dict[str, Any] + ] = [] # For content blocks with cache_control has_cache_control_in_text = False tool_calls: List[ChatCompletionAssistantToolCall] = [] thinking_blocks: List[ @@ -527,7 +565,9 @@ class LiteLLMAnthropicMessagesAdapter: "type": "text", "text": content.get("text", ""), } - self._add_cache_control_if_applicable(content, text_block, model) + self._add_cache_control_if_applicable( + content, text_block, model + ) if "cache_control" in text_block: has_cache_control_in_text = True assistant_content_list.append(text_block) @@ -549,19 +589,21 @@ class LiteLLMAnthropicMessagesAdapter: function_chunk.get("provider_specific_fields") or {} ) - provider_specific_fields["thought_signature"] = ( - signature - ) - function_chunk["provider_specific_fields"] = ( - provider_specific_fields - ) + provider_specific_fields[ + "thought_signature" + ] = signature + function_chunk[ + "provider_specific_fields" + ] = provider_specific_fields tool_call = ChatCompletionAssistantToolCall( id=content.get("id", ""), type="function", function=function_chunk, ) - self._add_cache_control_if_applicable(content, tool_call, model) + self._add_cache_control_if_applicable( + content, tool_call, model + ) tool_calls.append(tool_call) elif content.get("type") == "thinking": thinking_block = ChatCompletionThinkingBlock( @@ -660,10 +702,7 @@ class LiteLLMAnthropicMessagesAdapter: - vertex_ai/*claude* models """ model_lower = model.lower() - return ( - "anthropic" in model_lower - or "claude" in model_lower - ) + return "anthropic" in model_lower or "claude" in model_lower @staticmethod def translate_thinking_for_model( @@ -751,7 +790,9 @@ class LiteLLMAnthropicMessagesAdapter: for k, v in tool.items(): if k not in mapped_tool_params: # pass additional computer kwargs function_chunk.setdefault("parameters", {}).update({k: v}) - tool_param = ChatCompletionToolParam(type="function", function=function_chunk) + tool_param = ChatCompletionToolParam( + type="function", function=function_chunk + ) self._add_cache_control_if_applicable(tool, tool_param, model) new_tools.append(tool_param) # type: ignore[arg-type] @@ -883,10 +924,10 @@ class LiteLLMAnthropicMessagesAdapter: if "tool_choice" in anthropic_message_request: tool_choice = anthropic_message_request["tool_choice"] if tool_choice: - new_kwargs["tool_choice"] = ( - self.translate_anthropic_tool_choice_to_openai( - tool_choice=cast(AnthropicMessagesToolChoice, tool_choice) - ) + new_kwargs[ + "tool_choice" + ] = self.translate_anthropic_tool_choice_to_openai( + tool_choice=cast(AnthropicMessagesToolChoice, tool_choice) ) ## CONVERT TOOLS if "tools" in anthropic_message_request: @@ -907,7 +948,10 @@ class LiteLLMAnthropicMessagesAdapter: # Only translate regular tools (non-web-search) if regular_tools: - new_kwargs["tools"], tool_name_mapping = self.translate_anthropic_tools_to_openai( + ( + new_kwargs["tools"], + tool_name_mapping, + ) = self.translate_anthropic_tools_to_openai( tools=cast(List[AllAnthropicToolsValues], regular_tools), model=new_kwargs.get("model"), ) @@ -920,8 +964,10 @@ class LiteLLMAnthropicMessagesAdapter: if self.is_anthropic_claude_model(model): new_kwargs["thinking"] = thinking # type: ignore else: - reasoning_effort = self.translate_anthropic_thinking_to_reasoning_effort( - cast(Dict[str, Any], thinking) + reasoning_effort = ( + self.translate_anthropic_thinking_to_reasoning_effort( + cast(Dict[str, Any], thinking) + ) ) if reasoning_effort: new_kwargs["reasoning_effort"] = reasoning_effort @@ -1108,15 +1154,22 @@ class LiteLLMAnthropicMessagesAdapter: uncached_input_tokens = usage.prompt_tokens or 0 cached_tokens = 0 if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details: - cached_tokens = getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0 + cached_tokens = ( + getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0 + ) uncached_input_tokens -= cached_tokens anthropic_usage = AnthropicUsage( input_tokens=uncached_input_tokens, output_tokens=usage.completion_tokens or 0, ) - if hasattr(usage, "_cache_creation_input_tokens") and usage._cache_creation_input_tokens > 0: - anthropic_usage["cache_creation_input_tokens"] = usage._cache_creation_input_tokens + if ( + hasattr(usage, "_cache_creation_input_tokens") + and usage._cache_creation_input_tokens > 0 + ): + anthropic_usage[ + "cache_creation_input_tokens" + ] = usage._cache_creation_input_tokens if cached_tokens > 0: anthropic_usage["cache_read_input_tokens"] = cached_tokens @@ -1191,7 +1244,6 @@ class LiteLLMAnthropicMessagesAdapter: ContentThinkingSignatureBlockDelta, ], ]: - text: str = "" reasoning_content: str = "" reasoning_signature: str = "" @@ -1272,16 +1324,31 @@ class LiteLLMAnthropicMessagesAdapter: if litellm_usage_chunk is not None: uncached_input_tokens = litellm_usage_chunk.prompt_tokens or 0 cached_tokens = 0 - if hasattr(litellm_usage_chunk, "prompt_tokens_details") and litellm_usage_chunk.prompt_tokens_details: - cached_tokens = getattr(litellm_usage_chunk.prompt_tokens_details, "cached_tokens", 0) or 0 + if ( + hasattr(litellm_usage_chunk, "prompt_tokens_details") + and litellm_usage_chunk.prompt_tokens_details + ): + cached_tokens = ( + getattr( + litellm_usage_chunk.prompt_tokens_details, + "cached_tokens", + 0, + ) + or 0 + ) uncached_input_tokens -= cached_tokens usage_delta = UsageDelta( input_tokens=uncached_input_tokens, output_tokens=litellm_usage_chunk.completion_tokens or 0, ) - if hasattr(litellm_usage_chunk, "_cache_creation_input_tokens") and litellm_usage_chunk._cache_creation_input_tokens > 0: - usage_delta["cache_creation_input_tokens"] = litellm_usage_chunk._cache_creation_input_tokens + if ( + hasattr(litellm_usage_chunk, "_cache_creation_input_tokens") + and litellm_usage_chunk._cache_creation_input_tokens > 0 + ): + usage_delta[ + "cache_creation_input_tokens" + ] = litellm_usage_chunk._cache_creation_input_tokens if cached_tokens > 0: usage_delta["cache_read_input_tokens"] = cached_tokens else: diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/fake_stream_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/fake_stream_iterator.py index 542ae20b60..80afea7850 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/fake_stream_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/fake_stream_iterator.py @@ -19,11 +19,11 @@ from litellm.types.llms.anthropic_messages.anthropic_response import ( class FakeAnthropicMessagesStreamIterator: """ Fake streaming iterator for Anthropic Messages responses. - + Used when we need to convert a non-streaming response to a streaming format, such as when WebSearch interception converts stream=True to stream=False but the LLM doesn't make a tool call. - + This creates a proper Anthropic-style streaming response with multiple events: - message_start - content_block_start (for each content block) @@ -32,19 +32,19 @@ class FakeAnthropicMessagesStreamIterator: - message_delta (for usage) - message_stop """ - + def __init__(self, response: AnthropicMessagesResponse): self.response = response self.chunks = self._create_streaming_chunks() self.current_index = 0 - + def _create_streaming_chunks(self) -> List[bytes]: """Convert the non-streaming response to streaming chunks""" chunks = [] - + # Cast response to dict for easier access response_dict = cast(Dict[str, Any], self.response) - + # 1. message_start event usage = response_dict.get("usage", {}) message_start = { @@ -59,12 +59,14 @@ class FakeAnthropicMessagesStreamIterator: "stop_sequence": None, "usage": { "input_tokens": usage.get("input_tokens", 0) if usage else 0, - "output_tokens": 0 - } - } + "output_tokens": 0, + }, + }, } - chunks.append(f"event: message_start\ndata: {json.dumps(message_start)}\n\n".encode()) - + chunks.append( + f"event: message_start\ndata: {json.dumps(message_start)}\n\n".encode() + ) + # 2-4. For each content block, send start/delta/stop events content_blocks = response_dict.get("content", []) if content_blocks: @@ -72,38 +74,35 @@ class FakeAnthropicMessagesStreamIterator: # Cast block to dict for easier access block_dict = cast(Dict[str, Any], block) block_type = block_dict.get("type") - + if block_type == "text": # content_block_start content_block_start = { "type": "content_block_start", "index": index, - "content_block": { - "type": "text", - "text": "" - } + "content_block": {"type": "text", "text": ""}, } - chunks.append(f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode()) - + chunks.append( + f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode() + ) + # content_block_delta (send full text as one delta for simplicity) text = block_dict.get("text", "") content_block_delta = { "type": "content_block_delta", "index": index, - "delta": { - "type": "text_delta", - "text": text - } + "delta": {"type": "text_delta", "text": text}, } - chunks.append(f"event: content_block_delta\ndata: {json.dumps(content_block_delta)}\n\n".encode()) - + chunks.append( + f"event: content_block_delta\ndata: {json.dumps(content_block_delta)}\n\n".encode() + ) + # content_block_stop - content_block_stop = { - "type": "content_block_stop", - "index": index - } - chunks.append(f"event: content_block_stop\ndata: {json.dumps(content_block_stop)}\n\n".encode()) - + content_block_stop = {"type": "content_block_stop", "index": index} + chunks.append( + f"event: content_block_stop\ndata: {json.dumps(content_block_stop)}\n\n".encode() + ) + elif block_type == "thinking": # content_block_start for thinking content_block_start = { @@ -112,11 +111,13 @@ class FakeAnthropicMessagesStreamIterator: "content_block": { "type": "thinking", "thinking": "", - "signature": "" - } + "signature": "", + }, } - chunks.append(f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode()) - + chunks.append( + f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode() + ) + # content_block_delta for thinking text thinking_text = block_dict.get("thinking", "") if thinking_text: @@ -125,11 +126,13 @@ class FakeAnthropicMessagesStreamIterator: "index": index, "delta": { "type": "thinking_delta", - "thinking": thinking_text - } + "thinking": thinking_text, + }, } - chunks.append(f"event: content_block_delta\ndata: {json.dumps(content_block_delta)}\n\n".encode()) - + chunks.append( + f"event: content_block_delta\ndata: {json.dumps(content_block_delta)}\n\n".encode() + ) + # content_block_delta for signature (if present) signature = block_dict.get("signature", "") if signature: @@ -138,36 +141,36 @@ class FakeAnthropicMessagesStreamIterator: "index": index, "delta": { "type": "signature_delta", - "signature": signature - } + "signature": signature, + }, } - chunks.append(f"event: content_block_delta\ndata: {json.dumps(signature_delta)}\n\n".encode()) - + chunks.append( + f"event: content_block_delta\ndata: {json.dumps(signature_delta)}\n\n".encode() + ) + # content_block_stop - content_block_stop = { - "type": "content_block_stop", - "index": index - } - chunks.append(f"event: content_block_stop\ndata: {json.dumps(content_block_stop)}\n\n".encode()) - + content_block_stop = {"type": "content_block_stop", "index": index} + chunks.append( + f"event: content_block_stop\ndata: {json.dumps(content_block_stop)}\n\n".encode() + ) + elif block_type == "redacted_thinking": # content_block_start for redacted_thinking content_block_start = { "type": "content_block_start", "index": index, - "content_block": { - "type": "redacted_thinking" - } + "content_block": {"type": "redacted_thinking"}, } - chunks.append(f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode()) - + chunks.append( + f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode() + ) + # content_block_stop (no delta for redacted thinking) - content_block_stop = { - "type": "content_block_stop", - "index": index - } - chunks.append(f"event: content_block_stop\ndata: {json.dumps(content_block_stop)}\n\n".encode()) - + content_block_stop = {"type": "content_block_stop", "index": index} + chunks.append( + f"event: content_block_stop\ndata: {json.dumps(content_block_stop)}\n\n".encode() + ) + elif block_type == "tool_use": # content_block_start content_block_start = { @@ -177,11 +180,13 @@ class FakeAnthropicMessagesStreamIterator: "type": "tool_use", "id": block_dict.get("id"), "name": block_dict.get("name"), - "input": {} - } + "input": {}, + }, } - chunks.append(f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode()) - + chunks.append( + f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode() + ) + # content_block_delta (send input as JSON delta) input_data = block_dict.get("input", {}) content_block_delta = { @@ -189,58 +194,58 @@ class FakeAnthropicMessagesStreamIterator: "index": index, "delta": { "type": "input_json_delta", - "partial_json": json.dumps(input_data) - } + "partial_json": json.dumps(input_data), + }, } - chunks.append(f"event: content_block_delta\ndata: {json.dumps(content_block_delta)}\n\n".encode()) - + chunks.append( + f"event: content_block_delta\ndata: {json.dumps(content_block_delta)}\n\n".encode() + ) + # content_block_stop - content_block_stop = { - "type": "content_block_stop", - "index": index - } - chunks.append(f"event: content_block_stop\ndata: {json.dumps(content_block_stop)}\n\n".encode()) - + content_block_stop = {"type": "content_block_stop", "index": index} + chunks.append( + f"event: content_block_stop\ndata: {json.dumps(content_block_stop)}\n\n".encode() + ) + # 5. message_delta event (with final usage and stop_reason) message_delta = { "type": "message_delta", "delta": { "stop_reason": response_dict.get("stop_reason"), - "stop_sequence": response_dict.get("stop_sequence") + "stop_sequence": response_dict.get("stop_sequence"), }, - "usage": { - "output_tokens": usage.get("output_tokens", 0) if usage else 0 - } + "usage": {"output_tokens": usage.get("output_tokens", 0) if usage else 0}, } - chunks.append(f"event: message_delta\ndata: {json.dumps(message_delta)}\n\n".encode()) - + chunks.append( + f"event: message_delta\ndata: {json.dumps(message_delta)}\n\n".encode() + ) + # 6. message_stop event - message_stop = { - "type": "message_stop", - "usage": usage if usage else {} - } - chunks.append(f"event: message_stop\ndata: {json.dumps(message_stop)}\n\n".encode()) - + message_stop = {"type": "message_stop", "usage": usage if usage else {}} + chunks.append( + f"event: message_stop\ndata: {json.dumps(message_stop)}\n\n".encode() + ) + return chunks - + def __aiter__(self): return self - + async def __anext__(self): if self.current_index >= len(self.chunks): raise StopAsyncIteration - + chunk = self.chunks[self.current_index] self.current_index += 1 return chunk - + def __iter__(self): return self - + def __next__(self): if self.current_index >= len(self.chunks): raise StopIteration - + chunk = self.chunks[self.current_index] self.current_index += 1 return chunk diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py index 5b215c1fe5..1b5f03ec72 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py @@ -43,6 +43,7 @@ def _should_route_to_responses_api(custom_llm_provider: Optional[str]) -> bool: return False return custom_llm_provider in _RESPONSES_API_PROVIDERS + ####### ENVIRONMENT VARIABLES ################### # Initialize any necessary instances or variables here base_llm_http_handler = BaseLLMHTTPHandler() @@ -229,7 +230,7 @@ def anthropic_messages_handler( ]: """ Makes Anthropic `/v1/messages` API calls In the Anthropic API Spec - + Args: container: Container config with skills for code execution """ @@ -263,7 +264,7 @@ def anthropic_messages_handler( api_base=litellm_params.api_base, api_key=litellm_params.api_key, ) - + # Store agentic loop params in logging object for agentic hooks # This provides original request context needed for follow-up calls if litellm_logging_obj is not None: @@ -271,14 +272,15 @@ def anthropic_messages_handler( "model": original_model, "custom_llm_provider": custom_llm_provider, } - + # Check if stream was converted for WebSearch interception # This is set in the async wrapper above when stream=True is converted to stream=False if kwargs.get("_websearch_interception_converted_stream", False): - litellm_logging_obj.model_call_details["websearch_interception_converted_stream"] = True + litellm_logging_obj.model_call_details[ + "websearch_interception_converted_stream" + ] = True if litellm_params.mock_response and isinstance(litellm_params.mock_response, str): - return mock_response( model=model, messages=messages, @@ -324,8 +326,10 @@ def anthropic_messages_handler( return LiteLLMMessagesToResponsesAPIHandler.anthropic_messages_handler( **_shared_kwargs ) - return LiteLLMMessagesToCompletionTransformationHandler.anthropic_messages_handler( - **_shared_kwargs + return ( + LiteLLMMessagesToCompletionTransformationHandler.anthropic_messages_handler( + **_shared_kwargs + ) ) if custom_llm_provider is None: diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index df106c0e69..6cab38932a 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -12,6 +12,7 @@ from litellm.types.utils import GenericStreamingChunk, ModelResponseStream GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ = PassThroughEndpointLogging() + class BaseAnthropicMessagesStreamingIterator: """ Base class for Anthropic Messages streaming iterators that provides common logic @@ -27,7 +28,6 @@ class BaseAnthropicMessagesStreamingIterator: self.request_body = request_body self.start_time = datetime.now() - async def _handle_streaming_logging(self, collected_chunks: List[bytes]): """Handle the logging after all chunks have been collected.""" from litellm.proxy.pass_through_endpoints.streaming_handler import ( @@ -47,7 +47,7 @@ class BaseAnthropicMessagesStreamingIterator: end_time=end_time, ) ) - + def get_async_streaming_response_iterator( self, httpx_response, @@ -73,7 +73,7 @@ class BaseAnthropicMessagesStreamingIterator: def _convert_chunk_to_sse_format(self, chunk: Union[dict, Any]) -> bytes: """ Convert a chunk to Server-Sent Events format. - + This method should be overridden by subclasses if they need custom chunk formatting logic. """ @@ -94,15 +94,15 @@ class BaseAnthropicMessagesStreamingIterator: """ Generic async SSE wrapper that converts streaming chunks to SSE format and handles logging. - + This method provides the common logic for both Anthropic and Bedrock implementations. """ collected_chunks = [] - + async for chunk in completion_stream: encoded_chunk = self._convert_chunk_to_sse_format(chunk) collected_chunks.append(encoded_chunk) yield encoded_chunk - + # Handle logging after all chunks are processed - await self._handle_streaming_logging(collected_chunks) \ No newline at end of file + await self._handle_streaming_logging(collected_chunks) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py index e8d7a0383f..4d0a2cd829 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py @@ -165,15 +165,21 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): anthropic_messages_optional_request_params.pop("system", None) # Transform context_management from OpenAI format to Anthropic format if needed - context_management_param = anthropic_messages_optional_request_params.get("context_management") + context_management_param = anthropic_messages_optional_request_params.get( + "context_management" + ) if context_management_param is not None: from litellm.llms.anthropic.chat.transformation import AnthropicConfig - - transformed_context_management = AnthropicConfig.map_openai_context_management_to_anthropic( - context_management_param + + transformed_context_management = ( + AnthropicConfig.map_openai_context_management_to_anthropic( + context_management_param + ) ) if transformed_context_management is not None: - anthropic_messages_optional_request_params["context_management"] = transformed_context_management + anthropic_messages_optional_request_params[ + "context_management" + ] = transformed_context_management ####### get required params for all anthropic messages requests ###### verbose_logger.debug(f"TRANSFORMATION DEBUG - Messages: {messages}") diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py index ebc7d136f6..198ebe1ff8 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py @@ -43,7 +43,11 @@ def _build_responses_kwargs( Build the kwargs dict to pass directly to litellm.responses() / litellm.aresponses(). """ # Build a typed AnthropicMessagesRequest for the adapter - request_data: Dict[str, Any] = {"model": model, "messages": messages, "max_tokens": max_tokens} + request_data: Dict[str, Any] = { + "model": model, + "messages": messages, + "max_tokens": max_tokens, + } if context_management: request_data["context_management"] = context_management if output_config: @@ -142,7 +146,9 @@ class LiteLLMMessagesToResponsesAPIHandler: result = await litellm.aresponses(**responses_kwargs) if stream: - wrapper = AnthropicResponsesStreamWrapper(responses_stream=result, model=model) + wrapper = AnthropicResponsesStreamWrapper( + responses_stream=result, model=model + ) return wrapper.async_anthropic_sse_wrapper() if not isinstance(result, ResponsesAPIResponse): @@ -176,24 +182,26 @@ class LiteLLMMessagesToResponsesAPIHandler: Coroutine[Any, Any, Union[AnthropicMessagesResponse, AsyncIterator[Any]]], ]: if _is_async: - return LiteLLMMessagesToResponsesAPIHandler.async_anthropic_messages_handler( - max_tokens=max_tokens, - messages=messages, - model=model, - context_management=context_management, - metadata=metadata, - output_config=output_config, - stop_sequences=stop_sequences, - stream=stream, - system=system, - temperature=temperature, - thinking=thinking, - tool_choice=tool_choice, - tools=tools, - top_k=top_k, - top_p=top_p, - output_format=output_format, - **kwargs, + return ( + LiteLLMMessagesToResponsesAPIHandler.async_anthropic_messages_handler( + max_tokens=max_tokens, + messages=messages, + model=model, + context_management=context_management, + metadata=metadata, + output_config=output_config, + stop_sequences=stop_sequences, + stream=stream, + system=system, + temperature=temperature, + thinking=thinking, + tool_choice=tool_choice, + tools=tools, + top_k=top_k, + top_p=top_p, + output_format=output_format, + **kwargs, + ) ) # Sync path @@ -220,7 +228,9 @@ class LiteLLMMessagesToResponsesAPIHandler: result = litellm.responses(**responses_kwargs) if stream: - wrapper = AnthropicResponsesStreamWrapper(responses_stream=result, model=model) + wrapper = AnthropicResponsesStreamWrapper( + responses_stream=result, model=model + ) return wrapper.async_anthropic_sse_wrapper() if not isinstance(result, ResponsesAPIResponse): diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py index 926719c4ab..aa0738a071 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py @@ -35,7 +35,9 @@ class AnthropicResponsesStreamWrapper: # Map item_id -> content_block_index so we can stop the right block later self._item_id_to_block_index: Dict[str, int] = {} # Track open function_call items by item_id so we can emit tool_use start - self._pending_tool_ids: Dict[str, str] = {} # item_id -> call_id / name accumulator + self._pending_tool_ids: Dict[ + str, str + ] = {} # item_id -> call_id / name accumulator self._sent_message_start = False self._sent_message_stop = False self._chunk_queue: deque = deque() @@ -81,99 +83,168 @@ class AnthropicResponsesStreamWrapper: # ---- content_block_start for a new output message item ---- if event_type == "response.output_item.added": - item = getattr(event, "item", None) or (event.get("item") if isinstance(event, dict) else None) + item = getattr(event, "item", None) or ( + event.get("item") if isinstance(event, dict) else None + ) if item is None: return - item_type = getattr(item, "type", None) or (item.get("type") if isinstance(item, dict) else None) - item_id = getattr(item, "id", None) or (item.get("id") if isinstance(item, dict) else None) + item_type = getattr(item, "type", None) or ( + item.get("type") if isinstance(item, dict) else None + ) + item_id = getattr(item, "id", None) or ( + item.get("id") if isinstance(item, dict) else None + ) if item_type == "message": block_idx = self._next_block_index() if item_id: self._item_id_to_block_index[item_id] = block_idx - self._chunk_queue.append({ - "type": "content_block_start", - "index": block_idx, - "content_block": {"type": "text", "text": ""}, - }) + self._chunk_queue.append( + { + "type": "content_block_start", + "index": block_idx, + "content_block": {"type": "text", "text": ""}, + } + ) elif item_type == "function_call": - call_id = getattr(item, "call_id", None) or (item.get("call_id") if isinstance(item, dict) else None) or "" - name = getattr(item, "name", None) or (item.get("name") if isinstance(item, dict) else None) or "" + call_id = ( + getattr(item, "call_id", None) + or (item.get("call_id") if isinstance(item, dict) else None) + or "" + ) + name = ( + getattr(item, "name", None) + or (item.get("name") if isinstance(item, dict) else None) + or "" + ) block_idx = self._next_block_index() if item_id: self._item_id_to_block_index[item_id] = block_idx self._pending_tool_ids[item_id] = call_id - self._chunk_queue.append({ - "type": "content_block_start", - "index": block_idx, - "content_block": { - "type": "tool_use", - "id": call_id, - "name": name, - "input": {}, - }, - }) + self._chunk_queue.append( + { + "type": "content_block_start", + "index": block_idx, + "content_block": { + "type": "tool_use", + "id": call_id, + "name": name, + "input": {}, + }, + } + ) elif item_type == "reasoning": block_idx = self._next_block_index() if item_id: self._item_id_to_block_index[item_id] = block_idx - self._chunk_queue.append({ - "type": "content_block_start", - "index": block_idx, - "content_block": {"type": "thinking", "thinking": ""}, - }) + self._chunk_queue.append( + { + "type": "content_block_start", + "index": block_idx, + "content_block": {"type": "thinking", "thinking": ""}, + } + ) return # ---- text delta ---- if event_type == "response.output_text.delta": - item_id = getattr(event, "item_id", None) or (event.get("item_id") if isinstance(event, dict) else None) - delta = getattr(event, "delta", "") or (event.get("delta", "") if isinstance(event, dict) else "") - block_idx = self._item_id_to_block_index.get(item_id, self._current_block_index) if item_id else self._current_block_index - self._chunk_queue.append({ - "type": "content_block_delta", - "index": block_idx, - "delta": {"type": "text_delta", "text": delta}, - }) + item_id = getattr(event, "item_id", None) or ( + event.get("item_id") if isinstance(event, dict) else None + ) + delta = getattr(event, "delta", "") or ( + event.get("delta", "") if isinstance(event, dict) else "" + ) + block_idx = ( + self._item_id_to_block_index.get(item_id, self._current_block_index) + if item_id + else self._current_block_index + ) + self._chunk_queue.append( + { + "type": "content_block_delta", + "index": block_idx, + "delta": {"type": "text_delta", "text": delta}, + } + ) return # ---- reasoning summary text delta ---- if event_type == "response.reasoning_summary_text.delta": - item_id = getattr(event, "item_id", None) or (event.get("item_id") if isinstance(event, dict) else None) - delta = getattr(event, "delta", "") or (event.get("delta", "") if isinstance(event, dict) else "") - block_idx = self._item_id_to_block_index.get(item_id, self._current_block_index) if item_id else self._current_block_index - self._chunk_queue.append({ - "type": "content_block_delta", - "index": block_idx, - "delta": {"type": "thinking_delta", "thinking": delta}, - }) + item_id = getattr(event, "item_id", None) or ( + event.get("item_id") if isinstance(event, dict) else None + ) + delta = getattr(event, "delta", "") or ( + event.get("delta", "") if isinstance(event, dict) else "" + ) + block_idx = ( + self._item_id_to_block_index.get(item_id, self._current_block_index) + if item_id + else self._current_block_index + ) + self._chunk_queue.append( + { + "type": "content_block_delta", + "index": block_idx, + "delta": {"type": "thinking_delta", "thinking": delta}, + } + ) return # ---- function call arguments delta ---- if event_type == "response.function_call_arguments.delta": - item_id = getattr(event, "item_id", None) or (event.get("item_id") if isinstance(event, dict) else None) - delta = getattr(event, "delta", "") or (event.get("delta", "") if isinstance(event, dict) else "") - block_idx = self._item_id_to_block_index.get(item_id, self._current_block_index) if item_id else self._current_block_index - self._chunk_queue.append({ - "type": "content_block_delta", - "index": block_idx, - "delta": {"type": "input_json_delta", "partial_json": delta}, - }) + item_id = getattr(event, "item_id", None) or ( + event.get("item_id") if isinstance(event, dict) else None + ) + delta = getattr(event, "delta", "") or ( + event.get("delta", "") if isinstance(event, dict) else "" + ) + block_idx = ( + self._item_id_to_block_index.get(item_id, self._current_block_index) + if item_id + else self._current_block_index + ) + self._chunk_queue.append( + { + "type": "content_block_delta", + "index": block_idx, + "delta": {"type": "input_json_delta", "partial_json": delta}, + } + ) return # ---- output item done -> content_block_stop ---- if event_type == "response.output_item.done": - item = getattr(event, "item", None) or (event.get("item") if isinstance(event, dict) else None) - item_id = getattr(item, "id", None) or (item.get("id") if isinstance(item, dict) else None) if item else None - block_idx = self._item_id_to_block_index.get(item_id, self._current_block_index) if item_id else self._current_block_index - self._chunk_queue.append({ - "type": "content_block_stop", - "index": block_idx, - }) + item = getattr(event, "item", None) or ( + event.get("item") if isinstance(event, dict) else None + ) + item_id = ( + getattr(item, "id", None) + or (item.get("id") if isinstance(item, dict) else None) + if item + else None + ) + block_idx = ( + self._item_id_to_block_index.get(item_id, self._current_block_index) + if item_id + else self._current_block_index + ) + self._chunk_queue.append( + { + "type": "content_block_stop", + "index": block_idx, + } + ) return # ---- response completed -> message_delta + message_stop ---- - if event_type in ("response.completed", "response.failed", "response.incomplete"): - response_obj = getattr(event, "response", None) or (event.get("response") if isinstance(event, dict) else None) + if event_type in ( + "response.completed", + "response.failed", + "response.incomplete", + ): + response_obj = getattr(event, "response", None) or ( + event.get("response") if isinstance(event, dict) else None + ) stop_reason = "end_turn" input_tokens = 0 output_tokens = 0 @@ -191,14 +262,20 @@ class AnthropicResponsesStreamWrapper: cache_creation_tokens = getattr(usage, "input_tokens_details", None) # type: ignore[assignment] cache_read_tokens = getattr(usage, "output_tokens_details", None) # type: ignore[assignment] # Prefer direct cache fields if present - cache_creation_tokens = int(getattr(usage, "cache_creation_input_tokens", 0) or 0) - cache_read_tokens = int(getattr(usage, "cache_read_input_tokens", 0) or 0) + cache_creation_tokens = int( + getattr(usage, "cache_creation_input_tokens", 0) or 0 + ) + cache_read_tokens = int( + getattr(usage, "cache_read_input_tokens", 0) or 0 + ) # Check if tool_use was in the output to override stop_reason if response_obj is not None: output = getattr(response_obj, "output", []) or [] for out_item in output: - out_type = getattr(out_item, "type", None) or (out_item.get("type") if isinstance(out_item, dict) else None) + out_type = getattr(out_item, "type", None) or ( + out_item.get("type") if isinstance(out_item, dict) else None + ) if out_type == "function_call": stop_reason = "tool_use" break @@ -212,11 +289,13 @@ class AnthropicResponsesStreamWrapper: if cache_read_tokens: usage_delta["cache_read_input_tokens"] = cache_read_tokens - self._chunk_queue.append({ - "type": "message_delta", - "delta": {"stop_reason": stop_reason, "stop_sequence": None}, - "usage": usage_delta, - }) + self._chunk_queue.append( + { + "type": "message_delta", + "delta": {"stop_reason": stop_reason, "stop_sequence": None}, + "usage": usage_delta, + } + ) self._chunk_queue.append({"type": "message_stop"}) self._sent_message_stop = True return diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index 935babe438..ddd514146d 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -75,11 +75,13 @@ class LiteLLMAnthropicToResponsesAPIAdapter: if role == "user": if isinstance(content, str): - input_items.append({ - "type": "message", - "role": "user", - "content": [{"type": "input_text", "text": content}], - }) + input_items.append( + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": content}], + } + ) elif isinstance(content, list): user_parts: List[Dict[str, Any]] = [] for block in content: @@ -87,11 +89,17 @@ class LiteLLMAnthropicToResponsesAPIAdapter: continue btype = block.get("type") if btype == "text": - user_parts.append({"type": "input_text", "text": block.get("text", "")}) + user_parts.append( + {"type": "input_text", "text": block.get("text", "")} + ) elif btype == "image": - url = self._translate_anthropic_image_source_to_url(block.get("source", {})) + url = self._translate_anthropic_image_source_to_url( + block.get("source", {}) + ) if url: - user_parts.append({"type": "input_image", "image_url": url}) + user_parts.append( + {"type": "input_image", "image_url": url} + ) elif btype == "tool_result": tool_use_id = block.get("tool_use_id", "") inner = block.get("content") @@ -109,25 +117,31 @@ class LiteLLMAnthropicToResponsesAPIAdapter: else: output_text = str(inner) # tool_result is a top-level item, not inside the message - input_items.append({ - "type": "function_call_output", - "call_id": tool_use_id, - "output": output_text, - }) + input_items.append( + { + "type": "function_call_output", + "call_id": tool_use_id, + "output": output_text, + } + ) if user_parts: - input_items.append({ - "type": "message", - "role": "user", - "content": user_parts, - }) + input_items.append( + { + "type": "message", + "role": "user", + "content": user_parts, + } + ) elif role == "assistant": if isinstance(content, str): - input_items.append({ - "type": "message", - "role": "assistant", - "content": [{"type": "output_text", "text": content}], - }) + input_items.append( + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": content}], + } + ) elif isinstance(content, list): asst_parts: List[Dict[str, Any]] = [] for block in content: @@ -135,25 +149,33 @@ class LiteLLMAnthropicToResponsesAPIAdapter: continue btype = block.get("type") if btype == "text": - asst_parts.append({"type": "output_text", "text": block.get("text", "")}) + asst_parts.append( + {"type": "output_text", "text": block.get("text", "")} + ) elif btype == "tool_use": # tool_use becomes a top-level function_call item - input_items.append({ - "type": "function_call", - "call_id": block.get("id", ""), - "name": block.get("name", ""), - "arguments": json.dumps(block.get("input", {})), - }) + input_items.append( + { + "type": "function_call", + "call_id": block.get("id", ""), + "name": block.get("name", ""), + "arguments": json.dumps(block.get("input", {})), + } + ) elif btype == "thinking": thinking_text = block.get("thinking", "") if thinking_text: - asst_parts.append({"type": "output_text", "text": thinking_text}) + asst_parts.append( + {"type": "output_text", "text": thinking_text} + ) if asst_parts: - input_items.append({ - "type": "message", - "role": "assistant", - "content": asst_parts, - }) + input_items.append( + { + "type": "message", + "role": "assistant", + "content": asst_parts, + } + ) return input_items @@ -168,7 +190,9 @@ class LiteLLMAnthropicToResponsesAPIAdapter: tool_type = tool_dict.get("type", "") tool_name = tool_dict.get("name", "") # web_search tool - if (isinstance(tool_type, str) and tool_type.startswith("web_search")) or tool_name == "web_search": + if ( + isinstance(tool_type, str) and tool_type.startswith("web_search") + ) or tool_name == "web_search": result.append({"type": "web_search_preview"}) continue func_tool: Dict[str, Any] = {"type": "function", "name": tool_name} @@ -223,7 +247,9 @@ class LiteLLMAnthropicToResponsesAPIAdapter: return result if result else None @staticmethod - def translate_thinking_to_reasoning(thinking: Dict[str, Any]) -> Optional[Dict[str, Any]]: + def translate_thinking_to_reasoning( + thinking: Dict[str, Any] + ) -> Optional[Dict[str, Any]]: """ Convert Anthropic thinking param to Responses API reasoning param. @@ -253,7 +279,12 @@ class LiteLLMAnthropicToResponsesAPIAdapter: """ model: str = anthropic_request["model"] messages_list = cast( - List[Union[AnthropicMessagesUserMessageParam, AnthopicMessagesAssistantMessageParam]], + List[ + Union[ + AnthropicMessagesUserMessageParam, + AnthopicMessagesAssistantMessageParam, + ] + ], anthropic_request["messages"], ) @@ -296,7 +327,9 @@ class LiteLLMAnthropicToResponsesAPIAdapter: # tool_choice tool_choice = anthropic_request.get("tool_choice") if tool_choice: - responses_kwargs["tool_choice"] = self.translate_tool_choice_to_responses_api( + responses_kwargs[ + "tool_choice" + ] = self.translate_tool_choice_to_responses_api( cast(AnthropicMessagesToolChoice, tool_choice) ) @@ -314,7 +347,10 @@ class LiteLLMAnthropicToResponsesAPIAdapter: output_config = anthropic_request.get("output_config") if not isinstance(output_format, dict) and isinstance(output_config, dict): output_format = output_config.get("format") # type: ignore[assignment] - if isinstance(output_format, dict) and output_format.get("type") == "json_schema": + if ( + isinstance(output_format, dict) + and output_format.get("type") == "json_schema" + ): schema = output_format.get("schema") if schema: responses_kwargs["text"] = { @@ -329,7 +365,9 @@ class LiteLLMAnthropicToResponsesAPIAdapter: # context_management: Anthropic dict -> OpenAI array context_management = anthropic_request.get("context_management") if isinstance(context_management, dict): - openai_cm = self.translate_context_management_to_responses_api(context_management) + openai_cm = self.translate_context_management_to_responses_api( + context_management + ) if openai_cm is not None: responses_kwargs["context_management"] = openai_cm diff --git a/litellm/llms/anthropic/files/__init__.py b/litellm/llms/anthropic/files/__init__.py index b8b538ffb6..78c9dc89f7 100644 --- a/litellm/llms/anthropic/files/__init__.py +++ b/litellm/llms/anthropic/files/__init__.py @@ -1,4 +1,4 @@ from .handler import AnthropicFilesHandler +from .transformation import AnthropicFilesConfig -__all__ = ["AnthropicFilesHandler"] - +__all__ = ["AnthropicFilesHandler", "AnthropicFilesConfig"] diff --git a/litellm/llms/anthropic/files/handler.py b/litellm/llms/anthropic/files/handler.py index d46fc40131..77cc8c2731 100644 --- a/litellm/llms/anthropic/files/handler.py +++ b/litellm/llms/anthropic/files/handler.py @@ -40,7 +40,7 @@ ANTHROPIC_ERROR_STATUS_CODE_MAP = { class AnthropicFilesHandler: """ Handles Anthropic Files API operations. - + Currently supports: - file_content() for retrieving Anthropic Message Batch results """ @@ -58,17 +58,17 @@ class AnthropicFilesHandler: ) -> HttpxBinaryResponseContent: """ Async: Retrieve file content from Anthropic. - + For batch results, the file_id should be the batch_id. This will call Anthropic's /v1/messages/batches/{batch_id}/results endpoint. - + Args: file_content_request: Contains file_id (batch_id for batch results) api_base: Anthropic API base URL api_key: Anthropic API key timeout: Request timeout max_retries: Max retry attempts (unused for now) - + Returns: HttpxBinaryResponseContent: Binary content wrapped in compatible response format """ @@ -102,10 +102,7 @@ class AnthropicFilesHandler: # Make the request to Anthropic async_client = get_async_httpx_client(llm_provider=LlmProviders.ANTHROPIC) - anthropic_response = await async_client.get( - url=results_url, - headers=headers - ) + anthropic_response = await async_client.get(url=results_url, headers=headers) anthropic_response.raise_for_status() # Transform Anthropic batch results to OpenAI format @@ -124,7 +121,6 @@ class AnthropicFilesHandler: # Return the transformed response content return HttpxBinaryResponseContent(response=transformed_response) - def file_content( self, _is_async: bool, @@ -138,10 +134,10 @@ class AnthropicFilesHandler: ]: """ Retrieve file content from Anthropic. - + For batch results, the file_id should be the batch_id. This will call Anthropic's /v1/messages/batches/{batch_id}/results endpoint. - + Args: _is_async: Whether to run asynchronously file_content_request: Contains file_id (batch_id for batch results) @@ -149,7 +145,7 @@ class AnthropicFilesHandler: api_key: Anthropic API key timeout: Request timeout max_retries: Max retry attempts (unused for now) - + Returns: HttpxBinaryResponseContent or Coroutine: Binary content wrapped in compatible response format """ @@ -176,7 +172,7 @@ class AnthropicFilesHandler: ) -> bytes: """ Transform Anthropic batch results JSONL to OpenAI batch results JSONL format. - + Anthropic format: { "custom_id": "...", @@ -185,7 +181,7 @@ class AnthropicFilesHandler: "message": { ... } // Anthropic message format } } - + OpenAI format: { "custom_id": "...", @@ -199,28 +195,30 @@ class AnthropicFilesHandler: try: anthropic_config = AnthropicConfig() transformed_lines = [] - + # Parse JSONL content content_str = anthropic_content.decode("utf-8") for line in content_str.strip().split("\n"): if not line.strip(): continue - + anthropic_result = json.loads(line) custom_id = anthropic_result.get("custom_id", "") result = anthropic_result.get("result", {}) result_type = result.get("type", "") - + # Transform based on result type if result_type == "succeeded": # Transform Anthropic message to OpenAI format anthropic_message = result.get("message", {}) if anthropic_message: - openai_response_body = self._transform_anthropic_message_to_openai_format( - anthropic_message=anthropic_message, - anthropic_config=anthropic_config, + openai_response_body = ( + self._transform_anthropic_message_to_openai_format( + anthropic_message=anthropic_message, + anthropic_config=anthropic_config, + ) ) - + # Create OpenAI batch result format openai_result: OpenAIBatchResult = { "custom_id": custom_id, @@ -237,9 +235,9 @@ class AnthropicFilesHandler: error_obj = error.get("error", {}) error_message = error_obj.get("message", "Unknown error") error_type = error_obj.get("type", "api_error") - + status_code = ANTHROPIC_ERROR_STATUS_CODE_MAP.get(error_type, 500) - + error_body_errored: OpenAIErrorBody = { "error": { "message": error_message, @@ -272,7 +270,7 @@ class AnthropicFilesHandler: }, } transformed_lines.append(json.dumps(openai_result_canceled)) - + # Join lines and encode back to bytes transformed_content = "\n".join(transformed_lines) if transformed_lines: @@ -297,7 +295,7 @@ class AnthropicFilesHandler: status_code=200, content=json.dumps(anthropic_message).encode("utf-8"), ) - + # Create a ModelResponse object model_response = ModelResponse() # Initialize with required fields - will be populated by transform_parsed_response @@ -308,7 +306,7 @@ class AnthropicFilesHandler: message=litellm.Message(content="", role="assistant"), ) ] # type: ignore - + # Create a logging object for transformation logging_obj = Logging( model=anthropic_message.get("model", "claude-3-5-sonnet-20241022"), @@ -322,7 +320,7 @@ class AnthropicFilesHandler: kwargs={"optional_params": {}}, ) logging_obj.optional_params = {} - + # Transform using AnthropicConfig transformed_response = anthropic_config.transform_parsed_response( completion_response=anthropic_message, @@ -331,14 +329,16 @@ class AnthropicFilesHandler: json_mode=False, prefix_prompt=None, ) - + # Convert ModelResponse to OpenAI format dict - it's already in OpenAI format - openai_body: OpenAIChatCompletionResponse = transformed_response.model_dump(exclude_none=True) - + openai_body: OpenAIChatCompletionResponse = transformed_response.model_dump( + exclude_none=True + ) + # Ensure id comes from anthropic_message if not set if not openai_body.get("id"): openai_body["id"] = anthropic_message.get("id", "") - + return openai_body except Exception as e: verbose_logger.error( @@ -364,4 +364,3 @@ class AnthropicFilesHandler: }, } return error_response - diff --git a/litellm/llms/anthropic/skills/__init__.py b/litellm/llms/anthropic/skills/__init__.py index 60e78c2406..d7b3589db8 100644 --- a/litellm/llms/anthropic/skills/__init__.py +++ b/litellm/llms/anthropic/skills/__init__.py @@ -3,4 +3,3 @@ from .transformation import AnthropicSkillsConfig __all__ = ["AnthropicSkillsConfig"] - diff --git a/litellm/llms/anthropic/skills/transformation.py b/litellm/llms/anthropic/skills/transformation.py index ad0eff4297..af9863534e 100644 --- a/litellm/llms/anthropic/skills/transformation.py +++ b/litellm/llms/anthropic/skills/transformation.py @@ -47,10 +47,10 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): # Add required headers headers["x-api-key"] = api_key headers["anthropic-version"] = "2023-06-01" - + # Add beta header for skills API from litellm.constants import ANTHROPIC_SKILLS_API_BETA_VERSION - + if "anthropic-beta" not in headers: headers["anthropic-beta"] = ANTHROPIC_SKILLS_API_BETA_VERSION elif isinstance(headers["anthropic-beta"], list): @@ -58,8 +58,11 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): headers["anthropic-beta"].append(ANTHROPIC_SKILLS_API_BETA_VERSION) elif isinstance(headers["anthropic-beta"], str): if ANTHROPIC_SKILLS_API_BETA_VERSION not in headers["anthropic-beta"]: - headers["anthropic-beta"] = [headers["anthropic-beta"], ANTHROPIC_SKILLS_API_BETA_VERSION] - + headers["anthropic-beta"] = [ + headers["anthropic-beta"], + ANTHROPIC_SKILLS_API_BETA_VERSION, + ] + headers["content-type"] = "application/json" return headers @@ -87,13 +90,11 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): headers: dict, ) -> Dict: """Transform create skill request for Anthropic""" - verbose_logger.debug( - "Transforming create skill request: %s", create_request - ) - + verbose_logger.debug("Transforming create skill request: %s", create_request) + # Anthropic expects the request body directly request_body = {k: v for k, v in create_request.items() if v is not None} - + return request_body def transform_create_skill_response( @@ -103,10 +104,8 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): ) -> Skill: """Transform Anthropic response to Skill object""" response_json = raw_response.json() - verbose_logger.debug( - "Transforming create skill response: %s", response_json - ) - + verbose_logger.debug("Transforming create skill response: %s", response_json) + return Skill(**response_json) def transform_list_skills_request( @@ -122,7 +121,7 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): litellm_params.api_base if litellm_params else None ) url = self.get_complete_url(api_base=api_base, endpoint="skills") - + # Build query parameters query_params: Dict[str, Any] = {} if "limit" in list_params and list_params["limit"]: @@ -131,11 +130,12 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): query_params["page"] = list_params["page"] if "source" in list_params and list_params["source"]: query_params["source"] = list_params["source"] - + verbose_logger.debug( - "List skills request made to Anthropic Skills endpoint with params: %s", query_params + "List skills request made to Anthropic Skills endpoint with params: %s", + query_params, ) - + return url, query_params def transform_list_skills_response( @@ -145,10 +145,8 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): ) -> ListSkillsResponse: """Transform Anthropic response to ListSkillsResponse""" response_json = raw_response.json() - verbose_logger.debug( - "Transforming list skills response: %s", response_json - ) - + verbose_logger.debug("Transforming list skills response: %s", response_json) + return ListSkillsResponse(**response_json) def transform_get_skill_request( @@ -162,9 +160,9 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): url = self.get_complete_url( api_base=api_base, endpoint="skills", skill_id=skill_id ) - + verbose_logger.debug("Get skill request - URL: %s", url) - + return url, headers def transform_get_skill_response( @@ -174,10 +172,8 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): ) -> Skill: """Transform Anthropic response to Skill object""" response_json = raw_response.json() - verbose_logger.debug( - "Transforming get skill response: %s", response_json - ) - + verbose_logger.debug("Transforming get skill response: %s", response_json) + return Skill(**response_json) def transform_delete_skill_request( @@ -191,9 +187,9 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): url = self.get_complete_url( api_base=api_base, endpoint="skills", skill_id=skill_id ) - + verbose_logger.debug("Delete skill request - URL: %s", url) - + return url, headers def transform_delete_skill_response( @@ -203,9 +199,6 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): ) -> DeleteSkillResponse: """Transform Anthropic response to DeleteSkillResponse""" response_json = raw_response.json() - verbose_logger.debug( - "Transforming delete skill response: %s", response_json - ) - - return DeleteSkillResponse(**response_json) + verbose_logger.debug("Transforming delete skill response: %s", response_json) + return DeleteSkillResponse(**response_json) diff --git a/litellm/llms/aws_polly/text_to_speech/transformation.py b/litellm/llms/aws_polly/text_to_speech/transformation.py index dc6c40000f..caf6577039 100644 --- a/litellm/llms/aws_polly/text_to_speech/transformation.py +++ b/litellm/llms/aws_polly/text_to_speech/transformation.py @@ -43,20 +43,20 @@ class AWSPollyTextToSpeechConfig(BaseTextToSpeechConfig, BaseAWSLLM): # Voice name mappings from OpenAI voices to Polly voices VOICE_MAPPINGS = { - "alloy": "Joanna", # US English female - "echo": "Matthew", # US English male - "fable": "Amy", # British English female - "onyx": "Brian", # British English male - "nova": "Ivy", # US English female (child) - "shimmer": "Kendra", # US English female + "alloy": "Joanna", # US English female + "echo": "Matthew", # US English male + "fable": "Amy", # British English female + "onyx": "Brian", # British English male + "nova": "Ivy", # US English female (child) + "shimmer": "Kendra", # US English female } # Response format mappings from OpenAI to Polly FORMAT_MAPPINGS = { "mp3": "mp3", "opus": "ogg_vorbis", - "aac": "mp3", # Polly doesn't support AAC, use MP3 - "flac": "mp3", # Polly doesn't support FLAC, use MP3 + "aac": "mp3", # Polly doesn't support AAC, use MP3 + "flac": "mp3", # Polly doesn't support FLAC, use MP3 "wav": "pcm", "pcm": "pcm", } @@ -92,9 +92,9 @@ class AWSPollyTextToSpeechConfig(BaseTextToSpeechConfig, BaseAWSLLM): base_llm_http_handler: The BaseLLMHTTPHandler instance from main.py """ # Get AWS region from kwargs or environment - aws_region_name = kwargs.get("aws_region_name") or self._get_aws_region_name_for_polly( - optional_params=optional_params - ) + aws_region_name = kwargs.get( + "aws_region_name" + ) or self._get_aws_region_name_for_polly(optional_params=optional_params) # Convert voice to string if it's a dict voice_str: Optional[str] = None @@ -263,7 +263,9 @@ class AWSPollyTextToSpeechConfig(BaseTextToSpeechConfig, BaseAWSLLM): from botocore.auth import SigV4Auth from botocore.awsrequest import AWSRequest except ImportError: - raise ImportError("Missing boto3 to call AWS Polly. Run 'pip install boto3'.") + raise ImportError( + "Missing boto3 to call AWS Polly. Run 'pip install boto3'." + ) # Get AWS region aws_region_name = litellm_params.get("aws_region_name", self.DEFAULT_REGION) @@ -388,4 +390,3 @@ class AWSPollyTextToSpeechConfig(BaseTextToSpeechConfig, BaseAWSLLM): from litellm.types.llms.openai import HttpxBinaryResponseContent return HttpxBinaryResponseContent(raw_response) - diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 51b98c4af5..61cfd54b56 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -413,7 +413,9 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): litellm_params=litellm_params, ) if not isinstance(azure_client, (AsyncAzureOpenAI, AsyncOpenAI)): - raise ValueError("Azure client is not an instance of AsyncAzureOpenAI or AsyncOpenAI") + raise ValueError( + "Azure client is not an instance of AsyncAzureOpenAI or AsyncOpenAI" + ) ## LOGGING logging_obj.pre_call( input=data["messages"], @@ -595,7 +597,9 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): litellm_params=litellm_params, ) if not isinstance(azure_client, (AsyncAzureOpenAI, AsyncOpenAI)): - raise ValueError("Azure client is not an instance of AsyncAzureOpenAI or AsyncOpenAI") + raise ValueError( + "Azure client is not an instance of AsyncAzureOpenAI or AsyncOpenAI" + ) ## LOGGING logging_obj.pre_call( @@ -674,13 +678,15 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): litellm_params=litellm_params, ) if not isinstance(openai_aclient, (AsyncAzureOpenAI, AsyncOpenAI)): - raise ValueError("Azure client is not an instance of AsyncAzureOpenAI or AsyncOpenAI") + raise ValueError( + "Azure client is not an instance of AsyncAzureOpenAI or AsyncOpenAI" + ) raw_response = await openai_aclient.embeddings.with_raw_response.create( **data, timeout=timeout ) headers = dict(raw_response.headers) - + # Convert json.JSONDecodeError to AzureOpenAIError for two critical reasons: # # 1. ROUTER BEHAVIOR: The router relies on exception.status_code to determine cooldown logic: @@ -698,7 +704,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): except json.JSONDecodeError as json_error: raise AzureOpenAIError( status_code=raw_response.status_code or 500, - message=f"Failed to parse raw Azure embedding response: {str(json_error)}" + message=f"Failed to parse raw Azure embedding response: {str(json_error)}", ) from json_error if isinstance(response, str): raise AzureOpenAIError( @@ -1107,7 +1113,6 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): timeout=None, model: Optional[str] = None, ) -> ImageResponse: - response: Optional[dict] = None try: # response = await azure_client.images.generate(**data, timeout=timeout) @@ -1119,7 +1124,8 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): api_version: str = azure_client_params.get("api_version", "") # Use the deployment name (model) for URL construction, not the base_model from data img_gen_api_base = self.create_azure_base_url( - azure_client_params=azure_client_params, model=model or data.get("model", "") + azure_client_params=azure_client_params, + model=model or data.get("model", ""), ) ## LOGGING @@ -1212,13 +1218,17 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): and litellm_params is not None and litellm_params.get("base_model", None) is not None ): - model_response._hidden_params["model"] = litellm_params.get("base_model", None) + model_response._hidden_params["model"] = litellm_params.get( + "base_model", None + ) # Azure image generation API doesn't support extra_body parameter extra_body = optional_params.pop("extra_body", {}) flattened_params = {**optional_params, **extra_body} - - base_model = litellm_params.get("base_model", None) if litellm_params else None + + base_model = ( + litellm_params.get("base_model", None) if litellm_params else None + ) data = {"model": base_model or model, "prompt": prompt, **flattened_params} max_retries = data.pop("max_retries", 2) if not isinstance(max_retries, int): diff --git a/litellm/llms/azure/batches/handler.py b/litellm/llms/azure/batches/handler.py index 0e474a468e..6da3670b34 100644 --- a/litellm/llms/azure/batches/handler.py +++ b/litellm/llms/azure/batches/handler.py @@ -47,7 +47,9 @@ class AzureBatchesAPI(BaseAzureLLM): api_version: Optional[str], timeout: Union[float, httpx.Timeout], max_retries: Optional[int], - client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, + client: Optional[ + Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] + ] = None, litellm_params: Optional[dict] = None, ) -> Union[LiteLLMBatch, Coroutine[Any, Any, LiteLLMBatch]]: azure_client: Optional[ @@ -93,7 +95,9 @@ class AzureBatchesAPI(BaseAzureLLM): api_version: Optional[str], timeout: Union[float, httpx.Timeout], max_retries: Optional[int], - client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, + client: Optional[ + Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] + ] = None, litellm_params: Optional[dict] = None, ): azure_client: Optional[ @@ -141,7 +145,9 @@ class AzureBatchesAPI(BaseAzureLLM): api_version: Optional[str], timeout: Union[float, httpx.Timeout], max_retries: Optional[int], - client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, + client: Optional[ + Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] + ] = None, litellm_params: Optional[dict] = None, ): azure_client: Optional[ @@ -158,7 +164,7 @@ class AzureBatchesAPI(BaseAzureLLM): raise ValueError( "OpenAI client is not initialized. Make sure api_key is passed or OPENAI_API_KEY is set in the environment." ) - + if _is_async is True: if not isinstance(azure_client, (AsyncAzureOpenAI, AsyncOpenAI)): raise ValueError( @@ -167,7 +173,7 @@ class AzureBatchesAPI(BaseAzureLLM): return self.acancel_batch( # type: ignore cancel_batch_data=cancel_batch_data, client=azure_client ) - + # At this point, azure_client is guaranteed to be a sync client if not isinstance(azure_client, (AzureOpenAI, OpenAI)): raise ValueError( @@ -195,7 +201,9 @@ class AzureBatchesAPI(BaseAzureLLM): max_retries: Optional[int], after: Optional[str] = None, limit: Optional[int] = None, - client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, + client: Optional[ + Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] + ] = None, litellm_params: Optional[dict] = None, ): azure_client: Optional[ diff --git a/litellm/llms/azure/chat/gpt_5_transformation.py b/litellm/llms/azure/chat/gpt_5_transformation.py index 78d6372d02..6310df9cec 100644 --- a/litellm/llms/azure/chat/gpt_5_transformation.py +++ b/litellm/llms/azure/chat/gpt_5_transformation.py @@ -4,7 +4,10 @@ from typing import List import litellm from litellm.exceptions import UnsupportedParamsError -from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config +from litellm.llms.openai.chat.gpt_5_transformation import ( + OpenAIGPT5Config, + _get_effort_level, +) from litellm.types.llms.openai import AllMessageValues from .gpt_transformation import AzureOpenAIConfig @@ -38,7 +41,9 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config): used for manual routing. """ # gpt-5-chat* is a chat model and shouldn't go through GPT-5 reasoning restrictions. - return ("gpt-5" in model and "gpt-5-chat" not in model) or "gpt5_series" in model + return ( + "gpt-5" in model and "gpt-5-chat" not in model + ) or "gpt5_series" in model def get_supported_openai_params(self, model: str) -> List[str]: """Get supported parameters for Azure OpenAI GPT-5 models. @@ -61,7 +66,9 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config): # Only gpt-5.2+ has been verified to support logprobs on Azure. # The base OpenAI class includes logprobs for gpt-5.1+, but Azure # hasn't verified support for gpt-5.1, so remove them unless gpt-5.2/5.4+. - if self._supports_reasoning_effort_level(model, "none") and not self.is_model_gpt_5_2_model(model): + if self._supports_reasoning_effort_level( + model, "none" + ) and not self.is_model_gpt_5_2_model(model): params = [p for p in params if p not in ["logprobs", "top_logprobs"]] elif self.is_model_gpt_5_2_model(model): azure_supported_params = ["logprobs", "top_logprobs"] @@ -77,24 +84,27 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config): drop_params: bool, api_version: str = "", ) -> dict: - reasoning_effort_value = ( - non_default_params.get("reasoning_effort") - or optional_params.get("reasoning_effort") - ) + reasoning_effort_value = non_default_params.get( + "reasoning_effort" + ) or optional_params.get("reasoning_effort") + effective_effort = _get_effort_level(reasoning_effort_value) # gpt-5.1/5.2/5.4 support reasoning_effort='none', but other gpt-5 models don't # See: https://learn.microsoft.com/en-us/azure/ai-foundry/openai/how-to/reasoning supports_none = self._supports_reasoning_effort_level(model, "none") - if reasoning_effort_value == "none" and not supports_none: + if effective_effort == "none" and not supports_none: if litellm.drop_params is True or ( drop_params is not None and drop_params is True ): non_default_params = non_default_params.copy() optional_params = optional_params.copy() - if non_default_params.get("reasoning_effort") == "none": + if ( + _get_effort_level(non_default_params.get("reasoning_effort")) + == "none" + ): non_default_params.pop("reasoning_effort") - if optional_params.get("reasoning_effort") == "none": + if _get_effort_level(optional_params.get("reasoning_effort")) == "none": optional_params.pop("reasoning_effort") else: raise UnsupportedParamsError( @@ -117,9 +127,19 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config): ) # Only drop reasoning_effort='none' for models that don't support it - if result.get("reasoning_effort") == "none" and not supports_none: + result_effort = _get_effort_level(result.get("reasoning_effort")) + if result_effort == "none" and not supports_none: result.pop("reasoning_effort") + # Azure Chat Completions: gpt-5.4+ does not support tools + reasoning together. + # Drop reasoning_effort when both are present (OpenAI routes to Responses API; Azure does not). + if self.is_model_gpt_5_4_plus_model(model): + has_tools = bool( + non_default_params.get("tools") or optional_params.get("tools") + ) + if has_tools and result_effort not in (None, "none"): + result.pop("reasoning_effort", None) + return result def transform_request( diff --git a/litellm/llms/azure/chat/o_series_transformation.py b/litellm/llms/azure/chat/o_series_transformation.py index 778ec5f6de..cae7513245 100644 --- a/litellm/llms/azure/chat/o_series_transformation.py +++ b/litellm/llms/azure/chat/o_series_transformation.py @@ -44,7 +44,7 @@ class AzureOpenAIO1Config(OpenAIOSeriesConfig): return [ param for param in all_openai_params if param not in non_supported_params ] - + def _get_o_series_only_params(self, model: str) -> list: """ Helper function to get the o-series only params for the model @@ -52,7 +52,6 @@ class AzureOpenAIO1Config(OpenAIOSeriesConfig): - reasoning_effort """ o_series_only_param = [] - ######################################################### # Case 1: If the model is recognized and in litellm model cost map @@ -63,12 +62,12 @@ class AzureOpenAIO1Config(OpenAIOSeriesConfig): o_series_only_param.append("reasoning_effort") ######################################################### # Case 2: If the model is not recognized, then we assume it supports reasoning - # This is critical because several users tend to use custom deployment names + # This is critical because several users tend to use custom deployment names # for azure o-series models. ######################################################### else: o_series_only_param.append("reasoning_effort") - + return o_series_only_param def should_fake_stream( diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 7ed4306e29..fcdb3eca23 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -301,7 +301,9 @@ def get_azure_ad_token( ) tenant_id = litellm_params.get("tenant_id") or os.getenv("AZURE_TENANT_ID") client_id = litellm_params.get("client_id") or os.getenv("AZURE_CLIENT_ID") - client_secret = litellm_params.get("client_secret") or os.getenv("AZURE_CLIENT_SECRET") + client_secret = litellm_params.get("client_secret") or os.getenv( + "AZURE_CLIENT_SECRET" + ) azure_username = litellm_params.get("azure_username") or os.getenv("AZURE_USERNAME") azure_password = litellm_params.get("azure_password") or os.getenv("AZURE_PASSWORD") scope = litellm_params.get("azure_scope") or os.getenv( @@ -439,12 +441,16 @@ class BaseAzureLLM(BaseOpenAILLM): api_key: Optional[str], api_base: Optional[str], api_version: Optional[str] = None, - client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, + client: Optional[ + Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] + ] = None, litellm_params: Optional[dict] = None, _is_async: bool = False, model: Optional[str] = None, ) -> Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]]: - openai_client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None + openai_client: Optional[ + Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] + ] = None client_initialization_params: dict = locals() client_initialization_params["is_async"] = _is_async if client is None: @@ -453,7 +459,9 @@ class BaseAzureLLM(BaseOpenAILLM): client_type="azure", ) if cached_client: - if isinstance(cached_client, (AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI)): + if isinstance( + cached_client, (AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI) + ): return cached_client azure_client_params = self.initialize_azure_sdk_client( @@ -481,7 +489,9 @@ class BaseAzureLLM(BaseOpenAILLM): if "http_client" in azure_client_params: v1_params["http_client"] = azure_client_params["http_client"] - verbose_logger.debug(f"Using Azure v1 API with base_url: {v1_params['base_url']}") + verbose_logger.debug( + f"Using Azure v1 API with base_url: {v1_params['base_url']}" + ) if _is_async is True: openai_client = AsyncOpenAI(**v1_params) # type: ignore @@ -495,9 +505,11 @@ class BaseAzureLLM(BaseOpenAILLM): openai_client = AzureOpenAI(**azure_client_params) # type: ignore else: openai_client = client - if api_version is not None and isinstance( - openai_client, (AzureOpenAI, AsyncAzureOpenAI) - ) and isinstance(openai_client._custom_query, dict): + if ( + api_version is not None + and isinstance(openai_client, (AzureOpenAI, AsyncAzureOpenAI)) + and isinstance(openai_client._custom_query, dict) + ): # set api_version to version passed by user openai_client._custom_query.setdefault("api-version", api_version) @@ -524,11 +536,21 @@ class BaseAzureLLM(BaseOpenAILLM): # litellm_params sometimes contains the key, but the value is None # We should respect environment variables in this case - tenant_id = self._resolve_env_var(litellm_params, "tenant_id", "AZURE_TENANT_ID") - client_id = self._resolve_env_var(litellm_params, "client_id", "AZURE_CLIENT_ID") - client_secret = self._resolve_env_var(litellm_params, "client_secret", "AZURE_CLIENT_SECRET") - azure_username = self._resolve_env_var(litellm_params, "azure_username", "AZURE_USERNAME") - azure_password = self._resolve_env_var(litellm_params, "azure_password", "AZURE_PASSWORD") + tenant_id = self._resolve_env_var( + litellm_params, "tenant_id", "AZURE_TENANT_ID" + ) + client_id = self._resolve_env_var( + litellm_params, "client_id", "AZURE_CLIENT_ID" + ) + client_secret = self._resolve_env_var( + litellm_params, "client_secret", "AZURE_CLIENT_SECRET" + ) + azure_username = self._resolve_env_var( + litellm_params, "azure_username", "AZURE_USERNAME" + ) + azure_password = self._resolve_env_var( + litellm_params, "azure_password", "AZURE_PASSWORD" + ) scope = self._resolve_env_var(litellm_params, "azure_scope", "AZURE_SCOPE") if scope is None: scope = "https://cognitiveservices.azure.com/.default" @@ -777,9 +799,11 @@ class BaseAzureLLM(BaseOpenAILLM): return False return api_version in {"preview", "latest", "v1"} - def _resolve_env_var(self, litellm_params: Dict[str, Any], param_key: str, env_var_key: str) -> Optional[str]: + def _resolve_env_var( + self, litellm_params: Dict[str, Any], param_key: str, env_var_key: str + ) -> Optional[str]: """Resolve the environment variable for a given parameter key. - + The logic here is different from `params.get(key, os.getenv(env_var))` because litellm_params may contain the key with a None value, in which case we want to fallback to the environment variable. @@ -802,15 +826,9 @@ def get_azure_credentials( api_version: Optional[str] = None, ) -> AzureCredentials: """Resolve Azure credentials from params, litellm globals, and env vars.""" - resolved_api_base = ( - api_base - or litellm.api_base - or get_secret_str("AZURE_API_BASE") - ) + resolved_api_base = api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") resolved_api_version = ( - api_version - or litellm.api_version - or get_secret_str("AZURE_API_VERSION") + api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") ) resolved_api_key = ( api_key @@ -824,4 +842,3 @@ def get_azure_credentials( api_key=resolved_api_key, api_version=resolved_api_version, ) - diff --git a/litellm/llms/azure/exception_mapping.py b/litellm/llms/azure/exception_mapping.py index bcccad9352..dec7e7e5c9 100644 --- a/litellm/llms/azure/exception_mapping.py +++ b/litellm/llms/azure/exception_mapping.py @@ -24,9 +24,7 @@ class AzureOpenAIExceptionMapping: # Prefer the provider message/type/code when present. provider_message = ( - azure_error.get("message") - if isinstance(azure_error, dict) - else None + azure_error.get("message") if isinstance(azure_error, dict) else None ) or message provider_type = ( azure_error.get("type") if isinstance(azure_error, dict) else None diff --git a/litellm/llms/azure/files/handler.py b/litellm/llms/azure/files/handler.py index e53ced6b0e..72cbcba8a9 100644 --- a/litellm/llms/azure/files/handler.py +++ b/litellm/llms/azure/files/handler.py @@ -25,10 +25,12 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): super().__init__() @staticmethod - def _prepare_create_file_data(create_file_data: CreateFileRequest) -> dict[str, Any]: + def _prepare_create_file_data( + create_file_data: CreateFileRequest, + ) -> dict[str, Any]: """ Prepare create_file_data for OpenAI SDK. - + Removes expires_after if None to match SDK's Omit pattern. SDK expects file_create_params.ExpiresAfter | Omit, but FileExpiresAfter works at runtime. """ @@ -56,7 +58,9 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): api_version: Optional[str], timeout: Union[float, httpx.Timeout], max_retries: Optional[int], - client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, + client: Optional[ + Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] + ] = None, litellm_params: Optional[dict] = None, ) -> Union[OpenAIFileObject, Coroutine[Any, Any, OpenAIFileObject]]: openai_client: Optional[ @@ -102,7 +106,9 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): timeout: Union[float, httpx.Timeout], max_retries: Optional[int], api_version: Optional[str] = None, - client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, + client: Optional[ + Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] + ] = None, litellm_params: Optional[dict] = None, ) -> Union[ HttpxBinaryResponseContent, Coroutine[Any, Any, HttpxBinaryResponseContent] @@ -154,7 +160,9 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): timeout: Union[float, httpx.Timeout], max_retries: Optional[int], api_version: Optional[str] = None, - client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, + client: Optional[ + Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] + ] = None, litellm_params: Optional[dict] = None, ): openai_client: Optional[ @@ -206,7 +214,9 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): max_retries: Optional[int], organization: Optional[str] = None, api_version: Optional[str] = None, - client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, + client: Optional[ + Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] + ] = None, litellm_params: Optional[dict] = None, ): openai_client: Optional[ @@ -260,7 +270,9 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): max_retries: Optional[int], purpose: Optional[str] = None, api_version: Optional[str] = None, - client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, + client: Optional[ + Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] + ] = None, litellm_params: Optional[dict] = None, ): openai_client: Optional[ diff --git a/litellm/llms/azure/realtime/handler.py b/litellm/llms/azure/realtime/handler.py index 0ad6fb5735..6d00ecd51c 100644 --- a/litellm/llms/azure/realtime/handler.py +++ b/litellm/llms/azure/realtime/handler.py @@ -57,9 +57,12 @@ class AzureOpenAIRealtime(AzureChatCompletion): api_base = api_base.replace("https://", "wss://") # Determine path based on realtime_protocol (case-insensitive) - _is_ga = realtime_protocol is not None and realtime_protocol.upper() in ("GA", "V1") + _is_ga = realtime_protocol is not None and realtime_protocol.upper() in ( + "GA", + "V1", + ) if _is_ga: - path = "/openai/v1/realtime" + path = "/openai/v1/realtime" return f"{api_base}{path}?model={model}" else: # Default to beta path for backwards compatibility @@ -86,7 +89,9 @@ class AzureOpenAIRealtime(AzureChatCompletion): if api_base is None: raise ValueError("api_base is required for Azure OpenAI calls") - if api_version is None and (realtime_protocol is None or realtime_protocol.upper() not in ("GA", "V1")): + if api_version is None and ( + realtime_protocol is None or realtime_protocol.upper() not in ("GA", "V1") + ): raise ValueError("api_version is required for Azure OpenAI calls") url = self._construct_url( @@ -115,5 +120,7 @@ class AzureOpenAIRealtime(AzureChatCompletion): except websockets.exceptions.InvalidStatusCode as e: # type: ignore await websocket.close(code=e.status_code, reason=str(e)) except Exception: - verbose_proxy_logger.exception("Error in AzureOpenAIRealtime.async_realtime") + verbose_proxy_logger.exception( + "Error in AzureOpenAIRealtime.async_realtime" + ) pass diff --git a/litellm/llms/azure/realtime/http_transformation.py b/litellm/llms/azure/realtime/http_transformation.py index ef9a2d92d4..df1e2707af 100644 --- a/litellm/llms/azure/realtime/http_transformation.py +++ b/litellm/llms/azure/realtime/http_transformation.py @@ -9,22 +9,14 @@ from litellm.secret_managers.main import get_secret_str class AzureRealtimeHTTPConfig(BaseRealtimeHTTPConfig): def get_api_base(self, api_base: Optional[str], **kwargs) -> str: - return ( - api_base - or litellm.api_base - or get_secret_str("AZURE_API_BASE") - or "" - ) + return api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") or "" def get_api_key(self, api_key: Optional[str], **kwargs) -> str: - return ( - api_key - or litellm.api_key - or get_secret_str("AZURE_API_KEY") - or "" - ) + return api_key or litellm.api_key or get_secret_str("AZURE_API_KEY") or "" - def get_complete_url(self, api_base: Optional[str], model: str, api_version: Optional[str] = None) -> str: + def get_complete_url( + self, api_base: Optional[str], model: str, api_version: Optional[str] = None + ) -> str: base = self.get_api_base(api_base).rstrip("/") version = api_version or get_secret_str("AZURE_API_VERSION") or "2024-12-17" return f"{base}/openai/realtime/client_secrets?api-version={version}" @@ -41,7 +33,9 @@ class AzureRealtimeHTTPConfig(BaseRealtimeHTTPConfig): "Content-Type": "application/json", } - def get_realtime_calls_url(self, api_base: Optional[str], model: str, api_version: Optional[str] = None) -> str: + def get_realtime_calls_url( + self, api_base: Optional[str], model: str, api_version: Optional[str] = None + ) -> str: base = self.get_api_base(api_base).rstrip("/") version = api_version or get_secret_str("AZURE_API_VERSION") or "2024-12-17" return f"{base}/openai/realtime/calls?api-version={version}" diff --git a/litellm/llms/azure/responses/o_series_transformation.py b/litellm/llms/azure/responses/o_series_transformation.py index a0b2ef1630..3a554e9e19 100644 --- a/litellm/llms/azure/responses/o_series_transformation.py +++ b/litellm/llms/azure/responses/o_series_transformation.py @@ -27,7 +27,7 @@ else: class AzureOpenAIOSeriesResponsesAPIConfig(AzureOpenAIResponsesAPIConfig): """ Configuration for Azure OpenAI O-series models in Responses API. - + O-series models (o1, o3, etc.) do not support the temperature parameter in the responses API, so we need to drop it when drop_params is enabled. """ @@ -35,21 +35,22 @@ class AzureOpenAIOSeriesResponsesAPIConfig(AzureOpenAIResponsesAPIConfig): def get_supported_openai_params(self, model: str) -> list: """ Get supported parameters for Azure OpenAI O-series Responses API. - + O-series models don't support temperature parameter in responses API. """ # Get the base Azure supported params base_supported_params = super().get_supported_openai_params(model) - + # O-series models don't support temperature parameter in responses API o_series_unsupported_params = ["temperature"] - + # Filter out unsupported parameters for O-series models o_series_supported_params = [ - param for param in base_supported_params + param + for param in base_supported_params if param not in o_series_unsupported_params ] - + return o_series_supported_params def map_openai_params( @@ -60,34 +61,34 @@ class AzureOpenAIOSeriesResponsesAPIConfig(AzureOpenAIResponsesAPIConfig): ) -> Dict: """ Map OpenAI parameters for Azure OpenAI O-series Responses API. - + Drops temperature parameter if drop_params is True since O-series models don't support temperature in the responses API. """ mapped_params = dict(response_api_optional_params) - + # If drop_params is enabled, remove temperature parameter for O-series models if drop_params and "temperature" in mapped_params: verbose_logger.debug( f"Dropping unsupported parameter 'temperature' for Azure OpenAI O-series responses API model {model}" ) mapped_params.pop("temperature", None) - + return mapped_params def is_o_series_model(self, model: str) -> bool: """ Check if the model is an O-series model. - + Args: model: The model name to check - + Returns: True if it's an O-series model, False otherwise """ # Check if model name contains o_series or if it's a known O-series model if "o_series" in model.lower(): return True - + # Check if the model supports reasoning (which is O-series specific) - return supports_reasoning(model) \ No newline at end of file + return supports_reasoning(model) diff --git a/litellm/llms/azure/responses/transformation.py b/litellm/llms/azure/responses/transformation.py index 78631d3800..76a6d485bc 100644 --- a/litellm/llms/azure/responses/transformation.py +++ b/litellm/llms/azure/responses/transformation.py @@ -21,7 +21,6 @@ else: class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): - # Parameters not supported by Azure Responses API AZURE_UNSUPPORTED_PARAMS = ["context_management"] diff --git a/litellm/llms/azure/text_to_speech/__init__.py b/litellm/llms/azure/text_to_speech/__init__.py index ee923f122b..24dfb4fb49 100644 --- a/litellm/llms/azure/text_to_speech/__init__.py +++ b/litellm/llms/azure/text_to_speech/__init__.py @@ -5,4 +5,3 @@ from .transformation import AzureAVATextToSpeechConfig __all__ = [ "AzureAVATextToSpeechConfig", ] - diff --git a/litellm/llms/azure/text_to_speech/transformation.py b/litellm/llms/azure/text_to_speech/transformation.py index df582c3c09..a5dec24314 100644 --- a/litellm/llms/azure/text_to_speech/transformation.py +++ b/litellm/llms/azure/text_to_speech/transformation.py @@ -27,7 +27,7 @@ else: class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): """ Configuration for Azure AVA (Cognitive Services) Text-to-Speech - + Reference: https://learn.microsoft.com/en-us/azure/ai-services/speech-service/rest-text-to-speech """ @@ -78,9 +78,9 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): ]: """ Dispatch method to handle Azure AVA TTS requests - + This method encapsulates Azure-specific credential resolution and parameter handling - + Args: base_llm_http_handler: The BaseLLMHTTPHandler instance from main.py """ @@ -91,7 +91,7 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): or litellm.api_base or get_secret_str("AZURE_API_BASE") ) - + # Resolve api_key from multiple sources (Azure-specific) api_key = ( api_key @@ -101,7 +101,7 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): or get_secret_str("AZURE_OPENAI_API_KEY") or get_secret_str("AZURE_API_KEY") ) - + # Convert voice to string if it's a dict (for Azure AVA, voice must be a string) voice_str: Optional[str] = None if isinstance(voice, str): @@ -109,11 +109,13 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): elif isinstance(voice, dict): # Extract voice name from dict if needed voice_str = voice.get("name") if voice else None - - litellm_params_dict.update({ - "api_key": api_key, - "api_base": api_base, - }) + + litellm_params_dict.update( + { + "api_key": api_key, + "api_base": api_base, + } + ) # Call the text_to_speech_handler response = base_llm_http_handler.text_to_speech_handler( model=model, @@ -129,13 +131,13 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): client=None, _is_async=aspeech, ) - + return response def get_supported_openai_params(self, model: str) -> list: """ Azure AVA TTS supports these OpenAI parameters - + Note: Azure also supports additional SSML-specific parameters (style, styledegree, role) which can be passed but are not part of the OpenAI spec """ @@ -144,13 +146,13 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): def _convert_speed_to_azure_rate(self, speed: float) -> str: """ Convert OpenAI speed value to Azure SSML prosody rate percentage - + Args: speed: OpenAI speed value (0.25-4.0, default 1.0) - + Returns: Azure rate string with percentage (e.g., "+50%", "-50%", "+0%") - + Examples: speed=1.0 -> "+0%" (default) speed=2.0 -> "+100%" @@ -158,7 +160,7 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): """ rate_percentage = int((speed - 1.0) * 100) return f"{rate_percentage:+d}%" - + def _build_express_as_element( self, content: str, @@ -168,19 +170,19 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): ) -> str: """ Build mstts:express-as element with optional style, styledegree, and role attributes - + Args: content: The inner content to wrap style: Speaking style (e.g., "cheerful", "sad", "angry") styledegree: Style intensity (0.01 to 2) role: Voice role (e.g., "Girl", "Boy", "SeniorFemale", "SeniorMale") - + Returns: Content wrapped in mstts:express-as if any attributes provided, otherwise raw content """ if not (style or styledegree or role): return content - + express_as_attrs = [] if style: express_as_attrs.append(f"style='{style}'") @@ -188,10 +190,10 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): express_as_attrs.append(f"styledegree='{styledegree}'") if role: express_as_attrs.append(f"role='{role}'") - + express_as_attrs_str = " ".join(express_as_attrs) return f"{content}" - + def _get_voice_language( self, voice_name: Optional[str], @@ -199,14 +201,14 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): ) -> Optional[str]: """ Get the language for the voice element's xml:lang attribute - + Args: voice_name: The Azure voice name (e.g., "en-US-AriaNeural") explicit_lang: Explicitly provided language code (takes precedence) - + Returns: Language code if available (e.g., "es-ES"), or None - + Examples: - explicit_lang="es-ES" → "es-ES" (explicit takes precedence) - voice_name="en-US-AriaNeural", explicit_lang=None → None (use default from voice) @@ -215,7 +217,7 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): # If explicit language is provided, use it (for multilingual voices) if explicit_lang: return explicit_lang - + # For non-multilingual voices, we don't need to set xml:lang on the voice element # The voice name already encodes the language (e.g., en-US-AriaNeural) # Only return a language if explicitly set @@ -245,7 +247,7 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): else: # Assume it's already an Azure voice name mapped_voice = voice - + # Map response format if "response_format" in optional_params: format_name = optional_params["response_format"] @@ -257,23 +259,23 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): else: # Default to MP3 mapped_params["output_format"] = "audio-24khz-48kbitrate-mono-mp3" - + # Map speed (OpenAI: 0.25-4.0, Azure: prosody rate) if "speed" in optional_params: speed = optional_params["speed"] if speed is not None: mapped_params["rate"] = self._convert_speed_to_azure_rate(speed=speed) - + # Pass through Azure-specific SSML parameters if "style" in kwargs: mapped_params["style"] = kwargs["style"] - + if "styledegree" in kwargs: mapped_params["styledegree"] = kwargs["styledegree"] - + if "role" in kwargs: mapped_params["role"] = kwargs["role"] - + if "lang" in kwargs: mapped_params["lang"] = kwargs["lang"] return mapped_voice, mapped_params @@ -289,24 +291,24 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): Validate Azure environment and set up authentication headers """ validated_headers = headers.copy() - + # Azure AVA TTS requires either: # 1. Ocp-Apim-Subscription-Key header, or # 2. Authorization: Bearer header - + # We'll use the token-based auth via our token handler # The token will be added later in the handler - + if api_key: # If subscription key is provided, use it directly validated_headers["Ocp-Apim-Subscription-Key"] = api_key - + # Content-Type for SSML validated_headers["Content-Type"] = "application/ssml+xml" - + # User-Agent validated_headers["User-Agent"] = "litellm" - + return validated_headers def get_complete_url( @@ -317,7 +319,7 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): ) -> str: """ Get the complete URL for Azure AVA TTS request - + Azure TTS endpoint format: https://{region}.tts.speech.microsoft.com/cognitiveservices/v1 """ @@ -327,53 +329,50 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): f"Format: https://{{region}}.{self.COGNITIVE_SERVICES_DOMAIN} or " f"https://{{region}}.{self.TTS_SPEECH_DOMAIN}" ) - + # Remove trailing slash and parse URL api_base = api_base.rstrip("/") parsed_url = urlparse(api_base) hostname = parsed_url.hostname or "" - + # Check if it's a Cognitive Services endpoint (convert to TTS endpoint) if self._is_cognitive_services_endpoint(hostname=hostname): region = self._extract_region_from_hostname( - hostname=hostname, - domain=self.COGNITIVE_SERVICES_DOMAIN + hostname=hostname, domain=self.COGNITIVE_SERVICES_DOMAIN ) return self._build_tts_url(region=region) - + # Check if it's already a TTS endpoint if self._is_tts_endpoint(hostname=hostname): if not api_base.endswith(self.TTS_ENDPOINT_PATH): return f"{api_base}{self.TTS_ENDPOINT_PATH}" return api_base - + # Assume it's a custom endpoint, append the path return f"{api_base}{self.TTS_ENDPOINT_PATH}" def _is_cognitive_services_endpoint(self, hostname: str) -> bool: """Check if hostname is a Cognitive Services endpoint""" - return ( - hostname == self.COGNITIVE_SERVICES_DOMAIN - or hostname.endswith(f".{self.COGNITIVE_SERVICES_DOMAIN}") + return hostname == self.COGNITIVE_SERVICES_DOMAIN or hostname.endswith( + f".{self.COGNITIVE_SERVICES_DOMAIN}" ) def _is_tts_endpoint(self, hostname: str) -> bool: """Check if hostname is a TTS endpoint""" - return ( - hostname == self.TTS_SPEECH_DOMAIN - or hostname.endswith(f".{self.TTS_SPEECH_DOMAIN}") + return hostname == self.TTS_SPEECH_DOMAIN or hostname.endswith( + f".{self.TTS_SPEECH_DOMAIN}" ) def _extract_region_from_hostname(self, hostname: str, domain: str) -> str: """ Extract region from hostname - + Examples: eastus.api.cognitive.microsoft.com -> eastus api.cognitive.microsoft.com -> "" """ if hostname.endswith(f".{domain}"): - return hostname[:-len(f".{domain}")] + return hostname[: -len(f".{domain}")] return "" def _build_tts_url(self, region: str) -> str: @@ -382,7 +381,6 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): return f"https://{region}.{self.TTS_SPEECH_DOMAIN}{self.TTS_ENDPOINT_PATH}" return f"https://{self.TTS_SPEECH_DOMAIN}{self.TTS_ENDPOINT_PATH}" - def is_ssml_input(self, input: str) -> bool: """ Returns True if input is SSML, False otherwise @@ -402,30 +400,30 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): ) -> TextToSpeechRequestData: """ Transform OpenAI TTS request to Azure AVA TTS SSML format - + Note: optional_params should already be mapped via map_openai_params in main.py - + Supports Azure-specific SSML features: - style: Speaking style (e.g., "cheerful", "sad", "angry") - styledegree: Style intensity (0.01 to 2) - role: Voice role (e.g., "Girl", "Boy", "SeniorFemale", "SeniorMale") - lang: Language code for multilingual voices (e.g., "es-ES", "fr-FR") - + Auto-detects SSML: - If input contains , it's passed through as-is without transformation - + Returns: TextToSpeechRequestData: Contains SSML body and Azure-specific headers """ # Get voice (already mapped in main.py, or use default) azure_voice = voice or self.DEFAULT_VOICE - + # Get output format (already mapped in main.py) output_format = optional_params.get( "output_format", "audio-24khz-48kbitrate-mono-mp3" ) headers["X-Microsoft-OutputFormat"] = output_format - + # Auto-detect SSML: if input contains , pass it through as-is # Similar to Vertex AI behavior - check if input looks like SSML if self.is_ssml_input(input=input): @@ -433,14 +431,14 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): ssml_body=input, headers=headers, ) - + # Build SSML from plain text rate = optional_params.get("rate", "0%") style = optional_params.get("style") styledegree = optional_params.get("styledegree") role = optional_params.get("role") lang = optional_params.get("lang") - + # Escape XML special characters in input text escaped_input = ( input.replace("&", "&") @@ -449,19 +447,19 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): .replace('"', """) .replace("'", "'") ) - + # Determine if we need mstts namespace (for express-as element) use_mstts = style or role or styledegree - + # Build the xmlns attributes if use_mstts: xmlns = "xmlns='http://www.w3.org/2001/10/synthesis' xmlns:mstts='https://www.w3.org/2001/mstts'" else: xmlns = "xmlns='http://www.w3.org/2001/10/synthesis'" - + # Build the inner content with prosody prosody_content = f"{escaped_input}" - + # Wrap in mstts:express-as if style or role is specified voice_content = self._build_express_as_element( content=prosody_content, @@ -469,20 +467,20 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): styledegree=styledegree, role=role, ) - + # Build voice element with optional xml:lang attribute voice_lang = self._get_voice_language( voice_name=azure_voice, explicit_lang=lang, ) voice_lang_attr = f" xml:lang='{voice_lang}'" if voice_lang else "" - + ssml_body = f""" {voice_content} """ - + return { "ssml_body": ssml_body, "headers": headers, @@ -496,7 +494,7 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): ) -> "HttpxBinaryResponseContent": """ Transform Azure AVA TTS response to standard format - + Azure returns the audio data directly in the response body """ from litellm.types.llms.openai import HttpxBinaryResponseContent @@ -504,4 +502,3 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): # Azure returns audio data directly in the response body # Wrap it in HttpxBinaryResponseContent for consistent return type return HttpxBinaryResponseContent(raw_response) - diff --git a/litellm/llms/azure/vector_stores/transformation.py b/litellm/llms/azure/vector_stores/transformation.py index f1cd81b2bf..a98e7ae8cb 100644 --- a/litellm/llms/azure/vector_stores/transformation.py +++ b/litellm/llms/azure/vector_stores/transformation.py @@ -14,14 +14,12 @@ class AzureOpenAIVectorStoreConfig(OpenAIVectorStoreConfig): return BaseAzureLLM._get_base_azure_url( api_base=api_base, litellm_params=litellm_params, - route="/openai/vector_stores" + route="/openai/vector_stores", ) - def validate_environment( - self, headers: dict, litellm_params: Optional[GenericLiteLLMParams] + self, headers: dict, litellm_params: Optional[GenericLiteLLMParams] ) -> dict: return BaseAzureLLM._base_validate_azure_environment( - headers=headers, - litellm_params=litellm_params - ) \ No newline at end of file + headers=headers, litellm_params=litellm_params + ) diff --git a/litellm/llms/azure/videos/transformation.py b/litellm/llms/azure/videos/transformation.py index a6fbd8cef8..1ee0e95fb0 100644 --- a/litellm/llms/azure/videos/transformation.py +++ b/litellm/llms/azure/videos/transformation.py @@ -4,6 +4,7 @@ from litellm.types.videos.main import VideoCreateOptionalRequestParams from litellm.types.router import GenericLiteLLMParams from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.openai.videos.transformation import OpenAIVideoConfig + if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -64,16 +65,15 @@ class AzureVideoConfig(OpenAIVideoConfig): # If litellm_params is provided, use it; otherwise create a new one if litellm_params is None: litellm_params = GenericLiteLLMParams() - + if api_key and not litellm_params.api_key: litellm_params.api_key = api_key - + # Use the base Azure validation method which properly handles: # 1. Credentials from litellm_credential_name via litellm_params # 2. Sets the correct "api-key" header (not "Authorization: Bearer") return BaseAzureLLM._base_validate_azure_environment( - headers=headers, - litellm_params=litellm_params + headers=headers, litellm_params=litellm_params ) def get_complete_url( @@ -90,4 +90,4 @@ class AzureVideoConfig(OpenAIVideoConfig): litellm_params=litellm_params, route="/openai/v1/videos", default_api_version="", - ) \ No newline at end of file + ) diff --git a/litellm/llms/azure_ai/agents/handler.py b/litellm/llms/azure_ai/agents/handler.py index 379dc1e1c5..9eeec7f4e3 100644 --- a/litellm/llms/azure_ai/agents/handler.py +++ b/litellm/llms/azure_ai/agents/handler.py @@ -56,7 +56,7 @@ else: class AzureAIAgentsHandler: """ Handler for Azure AI Agent Service. - + Executes the complete agent flow which requires multiple API calls. """ @@ -72,16 +72,22 @@ class AzureAIAgentsHandler: def _build_thread_url(self, api_base: str, api_version: str) -> str: return f"{api_base}/threads?api-version={api_version}" - def _build_messages_url(self, api_base: str, thread_id: str, api_version: str) -> str: + def _build_messages_url( + self, api_base: str, thread_id: str, api_version: str + ) -> str: return f"{api_base}/threads/{thread_id}/messages?api-version={api_version}" def _build_runs_url(self, api_base: str, thread_id: str, api_version: str) -> str: return f"{api_base}/threads/{thread_id}/runs?api-version={api_version}" - def _build_run_status_url(self, api_base: str, thread_id: str, run_id: str, api_version: str) -> str: + def _build_run_status_url( + self, api_base: str, thread_id: str, run_id: str, api_version: str + ) -> str: return f"{api_base}/threads/{thread_id}/runs/{run_id}?api-version={api_version}" - def _build_list_messages_url(self, api_base: str, thread_id: str, api_version: str) -> str: + def _build_list_messages_url( + self, api_base: str, thread_id: str, api_version: str + ) -> str: return f"{api_base}/threads/{thread_id}/messages?api-version={api_version}" def _build_create_thread_and_run_url(self, api_base: str, api_version: str) -> str: @@ -112,12 +118,19 @@ class AzureAIAgentsHandler: from litellm.types.utils import Choices, Message, Usage model_response.choices = [ - Choices(finish_reason="stop", index=0, message=Message(content=content, role="assistant")) + Choices( + finish_reason="stop", + index=0, + message=Message(content=content, role="assistant"), + ) ] model_response.model = model # Store thread_id for conversation continuity - if not hasattr(model_response, "_hidden_params") or model_response._hidden_params is None: + if ( + not hasattr(model_response, "_hidden_params") + or model_response._hidden_params is None + ): model_response._hidden_params = {} model_response._hidden_params["thread_id"] = thread_id @@ -126,7 +139,9 @@ class AzureAIAgentsHandler: from litellm.utils import token_counter prompt_tokens = token_counter(model="gpt-3.5-turbo", messages=messages) - completion_tokens = token_counter(model="gpt-3.5-turbo", text=content, count_response_tokens=True) + completion_tokens = token_counter( + model="gpt-3.5-turbo", text=content, count_response_tokens=True + ) setattr( model_response, "usage", @@ -150,34 +165,43 @@ class AzureAIAgentsHandler: headers: Optional[dict], ) -> tuple: """Prepare common parameters for completion. - + Azure Foundry Agents API uses Bearer token authentication: - Authorization: Bearer (Azure AD token from 'az account get-access-token --resource https://ai.azure.com') - + See: https://learn.microsoft.com/en-us/azure/ai-foundry/agents/quickstart """ if headers is None: headers = {} headers["Content-Type"] = "application/json" - + # Azure Foundry Agents uses Bearer token authentication # The api_key here is expected to be an Azure AD token if api_key: headers["Authorization"] = f"Bearer {api_key}" - api_version = optional_params.get("api_version", self.config.DEFAULT_API_VERSION) + api_version = optional_params.get( + "api_version", self.config.DEFAULT_API_VERSION + ) agent_id = self.config._get_agent_id(model, optional_params) thread_id = optional_params.get("thread_id") api_base = api_base.rstrip("/") - verbose_logger.debug(f"Azure AI Agents completion - api_base: {api_base}, agent_id: {agent_id}") + verbose_logger.debug( + f"Azure AI Agents completion - api_base: {api_base}, agent_id: {agent_id}" + ) return headers, api_version, agent_id, thread_id, api_base - def _check_response(self, response: httpx.Response, expected_codes: List[int], error_msg: str): + def _check_response( + self, response: httpx.Response, expected_codes: List[int], error_msg: str + ): """Check response status and raise error if not expected.""" if response.status_code not in expected_codes: - raise AzureAIAgentsError(status_code=response.status_code, message=f"{error_msg}: {response.text}") + raise AzureAIAgentsError( + status_code=response.status_code, + message=f"{error_msg}: {response.text}", + ) # ------------------------------------------------------------------------- # Sync Completion @@ -200,16 +224,30 @@ class AzureAIAgentsHandler: from litellm.llms.custom_httpx.http_handler import _get_httpx_client if client is None: - client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)}) + client = _get_httpx_client( + params={"ssl_verify": litellm_params.get("ssl_verify", None)} + ) - headers, api_version, agent_id, thread_id, api_base = self._prepare_completion_params( + ( + headers, + api_version, + agent_id, + thread_id, + api_base, + ) = self._prepare_completion_params( model, api_base, api_key, optional_params, headers ) - def make_request(method: str, url: str, json_data: Optional[dict] = None) -> httpx.Response: + def make_request( + method: str, url: str, json_data: Optional[dict] = None + ) -> httpx.Response: if method == "GET": return client.get(url=url, headers=headers) - return client.post(url=url, headers=headers, data=json.dumps(json_data) if json_data else None) + return client.post( + url=url, + headers=headers, + data=json.dumps(json_data) if json_data else None, + ) # Execute the agent flow thread_id, content = self._execute_agent_flow_sync( @@ -222,7 +260,9 @@ class AzureAIAgentsHandler: optional_params=optional_params, ) - return self._build_model_response(model, content, model_response, thread_id, messages) + return self._build_model_response( + model, content, model_response, thread_id, messages + ) def _execute_agent_flow_sync( self, @@ -235,11 +275,15 @@ class AzureAIAgentsHandler: optional_params: dict, ) -> Tuple[str, str]: """Execute the agent flow synchronously. Returns (thread_id, content).""" - + # Step 1: Create thread if not provided if not thread_id: - verbose_logger.debug(f"Creating thread at: {self._build_thread_url(api_base, api_version)}") - response = make_request("POST", self._build_thread_url(api_base, api_version), {}) + verbose_logger.debug( + f"Creating thread at: {self._build_thread_url(api_base, api_version)}" + ) + response = make_request( + "POST", self._build_thread_url(api_base, api_version), {} + ) self._check_response(response, [200, 201], "Failed to create thread") thread_id = response.json()["id"] verbose_logger.debug(f"Created thread: {thread_id}") @@ -251,42 +295,58 @@ class AzureAIAgentsHandler: for msg in messages: if msg.get("role") in ["user", "system"]: url = self._build_messages_url(api_base, thread_id, api_version) - response = make_request("POST", url, {"role": "user", "content": msg.get("content", "")}) + response = make_request( + "POST", url, {"role": "user", "content": msg.get("content", "")} + ) self._check_response(response, [200, 201], "Failed to add message") # Step 3: Create run run_payload = {"assistant_id": agent_id} if "instructions" in optional_params: run_payload["instructions"] = optional_params["instructions"] - - response = make_request("POST", self._build_runs_url(api_base, thread_id, api_version), run_payload) + + response = make_request( + "POST", self._build_runs_url(api_base, thread_id, api_version), run_payload + ) self._check_response(response, [200, 201], "Failed to create run") run_id = response.json()["id"] verbose_logger.debug(f"Created run: {run_id}") # Step 4: Poll for completion - status_url = self._build_run_status_url(api_base, thread_id, run_id, api_version) + status_url = self._build_run_status_url( + api_base, thread_id, run_id, api_version + ) for _ in range(self.config.MAX_POLL_ATTEMPTS): response = make_request("GET", status_url) self._check_response(response, [200], "Failed to get run status") - + status = response.json().get("status") verbose_logger.debug(f"Run status: {status}") - + if status == "completed": break elif status in ["failed", "cancelled", "expired"]: - error_msg = response.json().get("last_error", {}).get("message", "Unknown error") - raise AzureAIAgentsError(status_code=500, message=f"Run {status}: {error_msg}") - + error_msg = ( + response.json() + .get("last_error", {}) + .get("message", "Unknown error") + ) + raise AzureAIAgentsError( + status_code=500, message=f"Run {status}: {error_msg}" + ) + time.sleep(self.config.POLL_INTERVAL_SECONDS) else: - raise AzureAIAgentsError(status_code=408, message="Run timed out waiting for completion") + raise AzureAIAgentsError( + status_code=408, message="Run timed out waiting for completion" + ) # Step 5: Get messages - response = make_request("GET", self._build_list_messages_url(api_base, thread_id, api_version)) + response = make_request( + "GET", self._build_list_messages_url(api_base, thread_id, api_version) + ) self._check_response(response, [200], "Failed to get messages") - + content = self._extract_content_from_messages(response.json()) return thread_id, content @@ -317,14 +377,26 @@ class AzureAIAgentsHandler: params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) - headers, api_version, agent_id, thread_id, api_base = self._prepare_completion_params( + ( + headers, + api_version, + agent_id, + thread_id, + api_base, + ) = self._prepare_completion_params( model, api_base, api_key, optional_params, headers ) - async def make_request(method: str, url: str, json_data: Optional[dict] = None) -> httpx.Response: + async def make_request( + method: str, url: str, json_data: Optional[dict] = None + ) -> httpx.Response: if method == "GET": return await client.get(url=url, headers=headers) - return await client.post(url=url, headers=headers, data=json.dumps(json_data) if json_data else None) + return await client.post( + url=url, + headers=headers, + data=json.dumps(json_data) if json_data else None, + ) # Execute the agent flow thread_id, content = await self._execute_agent_flow_async( @@ -337,7 +409,9 @@ class AzureAIAgentsHandler: optional_params=optional_params, ) - return self._build_model_response(model, content, model_response, thread_id, messages) + return self._build_model_response( + model, content, model_response, thread_id, messages + ) async def _execute_agent_flow_async( self, @@ -350,11 +424,15 @@ class AzureAIAgentsHandler: optional_params: dict, ) -> Tuple[str, str]: """Execute the agent flow asynchronously. Returns (thread_id, content).""" - + # Step 1: Create thread if not provided if not thread_id: - verbose_logger.debug(f"Creating thread at: {self._build_thread_url(api_base, api_version)}") - response = await make_request("POST", self._build_thread_url(api_base, api_version), {}) + verbose_logger.debug( + f"Creating thread at: {self._build_thread_url(api_base, api_version)}" + ) + response = await make_request( + "POST", self._build_thread_url(api_base, api_version), {} + ) self._check_response(response, [200, 201], "Failed to create thread") thread_id = response.json()["id"] verbose_logger.debug(f"Created thread: {thread_id}") @@ -366,42 +444,58 @@ class AzureAIAgentsHandler: for msg in messages: if msg.get("role") in ["user", "system"]: url = self._build_messages_url(api_base, thread_id, api_version) - response = await make_request("POST", url, {"role": "user", "content": msg.get("content", "")}) + response = await make_request( + "POST", url, {"role": "user", "content": msg.get("content", "")} + ) self._check_response(response, [200, 201], "Failed to add message") # Step 3: Create run run_payload = {"assistant_id": agent_id} if "instructions" in optional_params: run_payload["instructions"] = optional_params["instructions"] - - response = await make_request("POST", self._build_runs_url(api_base, thread_id, api_version), run_payload) + + response = await make_request( + "POST", self._build_runs_url(api_base, thread_id, api_version), run_payload + ) self._check_response(response, [200, 201], "Failed to create run") run_id = response.json()["id"] verbose_logger.debug(f"Created run: {run_id}") # Step 4: Poll for completion - status_url = self._build_run_status_url(api_base, thread_id, run_id, api_version) + status_url = self._build_run_status_url( + api_base, thread_id, run_id, api_version + ) for _ in range(self.config.MAX_POLL_ATTEMPTS): response = await make_request("GET", status_url) self._check_response(response, [200], "Failed to get run status") - + status = response.json().get("status") verbose_logger.debug(f"Run status: {status}") - + if status == "completed": break elif status in ["failed", "cancelled", "expired"]: - error_msg = response.json().get("last_error", {}).get("message", "Unknown error") - raise AzureAIAgentsError(status_code=500, message=f"Run {status}: {error_msg}") - + error_msg = ( + response.json() + .get("last_error", {}) + .get("message", "Unknown error") + ) + raise AzureAIAgentsError( + status_code=500, message=f"Run {status}: {error_msg}" + ) + await asyncio.sleep(self.config.POLL_INTERVAL_SECONDS) else: - raise AzureAIAgentsError(status_code=408, message="Run timed out waiting for completion") + raise AzureAIAgentsError( + status_code=408, message="Run timed out waiting for completion" + ) # Step 5: Get messages - response = await make_request("GET", self._build_list_messages_url(api_base, thread_id, api_version)) + response = await make_request( + "GET", self._build_list_messages_url(api_base, thread_id, api_version) + ) self._check_response(response, [200], "Failed to get messages") - + content = self._extract_content_from_messages(response.json()) return thread_id, content @@ -424,7 +518,13 @@ class AzureAIAgentsHandler: import litellm from litellm.llms.custom_httpx.http_handler import get_async_httpx_client - headers, api_version, agent_id, thread_id, api_base = self._prepare_completion_params( + ( + headers, + api_version, + agent_id, + thread_id, + api_base, + ) = self._prepare_completion_params( model, api_base, api_key, optional_params, headers ) @@ -432,20 +532,19 @@ class AzureAIAgentsHandler: thread_messages = [] for msg in messages: if msg.get("role") in ["user", "system"]: - thread_messages.append({ - "role": "user", - "content": msg.get("content", "") - }) + thread_messages.append( + {"role": "user", "content": msg.get("content", "")} + ) payload: Dict[str, Any] = { "assistant_id": agent_id, "stream": True, } - + # Add thread with messages if we don't have an existing thread if not thread_id: payload["thread"] = {"messages": thread_messages} - + if "instructions" in optional_params: payload["instructions"] = optional_params["instructions"] @@ -469,7 +568,7 @@ class AzureAIAgentsHandler: error_text = await response.aread() raise AzureAIAgentsError( status_code=response.status_code, - message=f"Streaming request failed: {error_text.decode()}" + message=f"Streaming request failed: {error_text.decode()}", ) async for chunk in self._process_sse_stream(response, model): @@ -482,23 +581,23 @@ class AzureAIAgentsHandler: ) -> AsyncIterator: """Process SSE stream and yield OpenAI-compatible streaming chunks.""" from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices - + response_id = f"chatcmpl-{uuid.uuid4().hex[:8]}" created = int(time.time()) thread_id = None - + current_event = None - + async for line in response.aiter_lines(): line = line.strip() - + if line.startswith("event:"): current_event = line[6:].strip() continue - + if line.startswith("data:"): data_str = line[5:].strip() - + if data_str == "[DONE]": # Send final chunk with finish_reason final_chunk = ModelResponseStream( @@ -518,17 +617,17 @@ class AzureAIAgentsHandler: final_chunk._hidden_params = {"thread_id": thread_id} yield final_chunk return - + try: data = json.loads(data_str) except json.JSONDecodeError: continue - + # Extract thread_id from thread.created event if current_event == "thread.created" and "id" in data: thread_id = data["id"] verbose_logger.debug(f"Stream created thread: {thread_id}") - + # Process message deltas - this is where the actual content comes if current_event == "thread.message.delta": delta_content = data.get("delta", {}).get("content", []) @@ -545,7 +644,9 @@ class AzureAIAgentsHandler: StreamingChoices( finish_reason=None, index=0, - delta=Delta(content=text_value, role="assistant"), + delta=Delta( + content=text_value, role="assistant" + ), ) ], ) diff --git a/litellm/llms/azure_ai/agents/transformation.py b/litellm/llms/azure_ai/agents/transformation.py index 01945aad32..777509fa82 100644 --- a/litellm/llms/azure_ai/agents/transformation.py +++ b/litellm/llms/azure_ai/agents/transformation.py @@ -56,9 +56,9 @@ class AzureAIAgentsConfig(BaseConfig): Azure AI Agent Service is a fully managed service for building AI agents that can understand natural language and perform tasks. - + Model format: azure_ai/agents/ - + The flow is: 1. Create a thread 2. Add user messages to the thread @@ -70,7 +70,7 @@ class AzureAIAgentsConfig(BaseConfig): # GA version: 2025-05-01, Preview: 2025-05-15-preview # See: https://learn.microsoft.com/en-us/azure/ai-foundry/agents/quickstart DEFAULT_API_VERSION = "2025-05-01" - + # Polling configuration MAX_POLL_ATTEMPTS = 60 POLL_INTERVAL_SECONDS = 1.0 @@ -82,7 +82,7 @@ class AzureAIAgentsConfig(BaseConfig): def is_azure_ai_agents_route(model: str) -> bool: """ Check if the model is an Azure AI Agents route. - + Model format: azure_ai/agents/ """ return "agents/" in model @@ -91,7 +91,7 @@ class AzureAIAgentsConfig(BaseConfig): def get_agent_id_from_model(model: str) -> str: """ Extract agent ID from the model string. - + Model format: azure_ai/agents/ -> or: agents/ -> """ @@ -153,12 +153,12 @@ class AzureAIAgentsConfig(BaseConfig): ) -> str: """ Get the base URL for Azure AI Agent Service. - + The actual endpoint will vary based on the operation: - /openai/threads for creating threads - /openai/threads/{thread_id}/messages for adding messages - /openai/threads/{thread_id}/runs for creating runs - + This returns the base URL that will be modified for each operation. """ if api_base is None: @@ -178,7 +178,9 @@ class AzureAIAgentsConfig(BaseConfig): model format: "azure_ai/agents/" or "agents/" or just "" """ - agent_id = optional_params.get("agent_id") or optional_params.get("assistant_id") + agent_id = optional_params.get("agent_id") or optional_params.get( + "assistant_id" + ) if agent_id: return agent_id @@ -195,7 +197,7 @@ class AzureAIAgentsConfig(BaseConfig): ) -> dict: """ Transform the request for Azure Agents. - + This stores the necessary data for the multi-step agent flow. The actual API calls happen in the custom handler. """ @@ -246,10 +248,10 @@ class AzureAIAgentsConfig(BaseConfig): ) -> dict: """ Validate and set up environment for Azure Foundry Agents requests. - + Azure Foundry Agents uses Bearer token authentication with Azure AD tokens. Get token via: az account get-access-token --resource 'https://ai.azure.com' - + See: https://learn.microsoft.com/en-us/azure/ai-foundry/agents/quickstart """ headers["Content-Type"] = "application/json" @@ -326,15 +328,15 @@ class AzureAIAgentsConfig(BaseConfig): ) -> Any: """ Dispatch method for Azure Foundry Agents completion. - + Routes to sync or async completion based on acompletion flag. Supports native streaming via SSE when stream=True and acompletion=True. - + Authentication: Uses Azure AD Bearer tokens. - Pass api_key directly as an Azure AD token - Or set up Azure AD credentials via environment variables for automatic token retrieval: - AZURE_TENANT_ID, AZURE_CLIENT_ID, AZURE_CLIENT_SECRET (Service Principal) - + See: https://learn.microsoft.com/en-us/azure/ai-foundry/agents/quickstart """ from litellm.llms.azure.common_utils import get_azure_ad_token @@ -349,7 +351,7 @@ class AzureAIAgentsConfig(BaseConfig): azure_auth_params = dict(litellm_params) if litellm_params else {} azure_auth_params["azure_scope"] = "https://ai.azure.com/.default" api_key = get_azure_ad_token(GenericLiteLLMParams(**azure_auth_params)) - + if api_key is None: raise ValueError( "api_key (Azure AD token) is required for Azure Foundry Agents. " diff --git a/litellm/llms/azure_ai/anthropic/__init__.py b/litellm/llms/azure_ai/anthropic/__init__.py index 233f22999f..931c71de3b 100644 --- a/litellm/llms/azure_ai/anthropic/__init__.py +++ b/litellm/llms/azure_ai/anthropic/__init__.py @@ -6,7 +6,11 @@ from .transformation import AzureAnthropicConfig try: from .messages_transformation import AzureAnthropicMessagesConfig - __all__ = ["AzureAnthropicChatCompletion", "AzureAnthropicConfig", "AzureAnthropicMessagesConfig"] + + __all__ = [ + "AzureAnthropicChatCompletion", + "AzureAnthropicConfig", + "AzureAnthropicMessagesConfig", + ] except ImportError: __all__ = ["AzureAnthropicChatCompletion", "AzureAnthropicConfig"] - diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index 2cba27925c..e24fc2097d 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -87,7 +87,9 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): ) # Use provided timeout or fall back to litellm.request_timeout - request_timeout = timeout if timeout is not None else litellm.request_timeout + request_timeout = ( + timeout if timeout is not None else litellm.request_timeout + ) response = await async_client.post( endpoint_url, diff --git a/litellm/llms/azure_ai/anthropic/handler.py b/litellm/llms/azure_ai/anthropic/handler.py index fe4524fd5b..a2263e72a1 100644 --- a/litellm/llms/azure_ai/anthropic/handler.py +++ b/litellm/llms/azure_ai/anthropic/handler.py @@ -64,7 +64,7 @@ class AzureAnthropicChatCompletion(AnthropicChatCompletion): # Use AzureAnthropicConfig for both azure_anthropic and azure_ai Claude models config = AzureAnthropicConfig() - + headers = config.validate_environment( api_key=api_key, headers=headers, @@ -224,4 +224,3 @@ class AzureAnthropicChatCompletion(AnthropicChatCompletion): encoding=encoding, json_mode=json_mode, ) - diff --git a/litellm/llms/azure_ai/anthropic/messages_transformation.py b/litellm/llms/azure_ai/anthropic/messages_transformation.py index 8e60e84391..59d8fb02c6 100644 --- a/litellm/llms/azure_ai/anthropic/messages_transformation.py +++ b/litellm/llms/azure_ai/anthropic/messages_transformation.py @@ -125,6 +125,7 @@ class AzureAnthropicMessagesConfig(AnthropicMessagesConfig): Processes both `system` and `messages` content blocks. """ + def _sanitize(cache_control: Any) -> None: if isinstance(cache_control, dict): cache_control.pop("scope", None) @@ -163,4 +164,3 @@ class AzureAnthropicMessagesConfig(AnthropicMessagesConfig): ) self._remove_scope_from_cache_control(anthropic_messages_request) return anthropic_messages_request - diff --git a/litellm/llms/azure_ai/anthropic/transformation.py b/litellm/llms/azure_ai/anthropic/transformation.py index c5510db68b..5d8f27b97d 100644 --- a/litellm/llms/azure_ai/anthropic/transformation.py +++ b/litellm/llms/azure_ai/anthropic/transformation.py @@ -49,7 +49,7 @@ class AzureAnthropicConfig(AnthropicConfig): # Set api_key if provided and not already set if api_key and not litellm_params_obj.api_key: litellm_params_obj.api_key = api_key - + # Use Azure authentication logic headers = BaseAzureLLM._base_validate_azure_environment( headers=headers, litellm_params=litellm_params_obj @@ -86,7 +86,6 @@ class AzureAnthropicConfig(AnthropicConfig): if "anthropic-version" not in headers: headers["anthropic-version"] = "2023-06-01" - return headers def transform_request( @@ -116,4 +115,3 @@ class AzureAnthropicConfig(AnthropicConfig): data.pop("stream_options", None) return data - diff --git a/litellm/llms/azure_ai/azure_model_router/transformation.py b/litellm/llms/azure_ai/azure_model_router/transformation.py index efda85f37b..57acb14706 100644 --- a/litellm/llms/azure_ai/azure_model_router/transformation.py +++ b/litellm/llms/azure_ai/azure_model_router/transformation.py @@ -17,7 +17,7 @@ from litellm.types.utils import ModelResponse class AzureModelRouterConfig(AzureAIStudioConfig): """ Configuration for Azure AI Foundry Model Router. - + Handles: - Stripping model_router prefix before sending to Azure API - Preserving full model path in responses for cost tracking @@ -34,7 +34,7 @@ class AzureModelRouterConfig(AzureAIStudioConfig): ) -> dict: """ Transform request for Model Router. - + Strips the model_router/ prefix so only the deployment name is sent to Azure. Example: model_router/azure-model-router -> azure-model-router """ @@ -42,7 +42,7 @@ class AzureModelRouterConfig(AzureAIStudioConfig): # Get base model name (strips routing prefixes like model_router/) base_model: str = AzureFoundryModelInfo.get_base_model(model) - + return super().transform_request( base_model, messages, optional_params, litellm_params, headers ) @@ -63,7 +63,7 @@ class AzureModelRouterConfig(AzureAIStudioConfig): ) -> ModelResponse: """ Transform response for Model Router. - + Extracts the actual model used from the Azure response (e.g., gpt-5-nano-2025-08-07) and returns it with the azure_ai/ prefix for proper display and cost tracking. """ @@ -71,8 +71,8 @@ class AzureModelRouterConfig(AzureAIStudioConfig): # Get base model for the parent call (strips routing prefixes for API compatibility) base_model: str = AzureFoundryModelInfo.get_base_model(model) - - # Call parent transform_response first - this will extract the actual model + + # Call parent transform_response first - this will extract the actual model # from the raw response (e.g., "gpt-5-nano-2025-08-07") model_response = super().transform_response( model=base_model, @@ -94,26 +94,26 @@ class AzureModelRouterConfig(AzureAIStudioConfig): ) -> Optional[dict]: """ Calculate additional costs for Azure Model Router. - + Adds a flat infrastructure cost of $0.14 per M input tokens for using the Model Router. - + Args: model: The model name (should be a model router model) prompt_tokens: Number of prompt tokens completion_tokens: Number of completion tokens - + Returns: Dictionary with additional costs, or None if not applicable. """ from litellm.llms.azure_ai.cost_calculator import ( calculate_azure_model_router_flat_cost, ) - + flat_cost = calculate_azure_model_router_flat_cost( model=model, prompt_tokens=prompt_tokens ) - + if flat_cost > 0: return {"Azure Model Router Flat Cost": flat_cost} - + return None diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 585efd3307..529ec71c53 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -90,7 +90,10 @@ class AzureAIStudioConfig(OpenAIConfig): """ parsed_url = urlparse(api_base) host = parsed_url.hostname - if host and (host.endswith(".services.ai.azure.com") or host.endswith(".openai.azure.com")): + if host and ( + host.endswith(".services.ai.azure.com") + or host.endswith(".openai.azure.com") + ): return True return False @@ -137,9 +140,13 @@ class AzureAIStudioConfig(OpenAIConfig): # Add the path to the base URL if "services.ai.azure.com" in api_base: - new_url = _add_path_to_api_base(api_base=api_base, ending_path="/models/chat/completions") + new_url = _add_path_to_api_base( + api_base=api_base, ending_path="/models/chat/completions" + ) else: - new_url = _add_path_to_api_base(api_base=api_base, ending_path="/chat/completions") + new_url = _add_path_to_api_base( + api_base=api_base, ending_path="/chat/completions" + ) # Use the new query_params dictionary final_url = httpx.URL(new_url).copy_with(params=query_params) @@ -209,7 +216,11 @@ class AzureAIStudioConfig(OpenAIConfig): dynamic_api_key = api_key or get_secret_str("AZURE_AI_API_KEY") if self._is_azure_openai_model(model=model, api_base=api_base): - verbose_logger.debug("Model={} is Azure OpenAI model. Setting custom_llm_provider='azure'.".format(model)) + verbose_logger.debug( + "Model={} is Azure OpenAI model. Setting custom_llm_provider='azure'.".format( + model + ) + ) custom_llm_provider = "azure" return api_base, dynamic_api_key, custom_llm_provider @@ -225,7 +236,9 @@ class AzureAIStudioConfig(OpenAIConfig): if extra_body and isinstance(extra_body, dict): optional_params.update(extra_body) optional_params.pop("max_retries", None) - return super().transform_request(model, messages, optional_params, litellm_params, headers) + return super().transform_request( + model, messages, optional_params, litellm_params, headers + ) def transform_response( self, @@ -264,30 +277,47 @@ class AzureAIStudioConfig(OpenAIConfig): if should_drop_params and "Extra inputs are not permitted" in error_text: return True - elif "unknown field: parameter index is not a valid field" in error_text: # remove index from tool calls + elif ( + "unknown field: parameter index is not a valid field" in error_text + ): # remove index from tool calls return True elif ( - AzureFoundryErrorStrings.SET_EXTRA_PARAMETERS_TO_PASS_THROUGH.value in error_text + AzureFoundryErrorStrings.SET_EXTRA_PARAMETERS_TO_PASS_THROUGH.value + in error_text ): # remove extra-parameters from tool calls return True - return super().should_retry_llm_api_inside_llm_translation_on_http_error(e=e, litellm_params=litellm_params) + return super().should_retry_llm_api_inside_llm_translation_on_http_error( + e=e, litellm_params=litellm_params + ) @property def max_retry_on_unprocessable_entity_error(self) -> int: return 2 - def transform_request_on_unprocessable_entity_error(self, e: httpx.HTTPStatusError, request_data: dict) -> dict: + def transform_request_on_unprocessable_entity_error( + self, e: httpx.HTTPStatusError, request_data: dict + ) -> dict: _messages = cast(Optional[List[AllMessageValues]], request_data.get("messages")) - if "unknown field: parameter index is not a valid field" in e.response.text and _messages is not None: + if ( + "unknown field: parameter index is not a valid field" in e.response.text + and _messages is not None + ): litellm.remove_index_from_tool_calls( messages=_messages, ) - elif AzureFoundryErrorStrings.SET_EXTRA_PARAMETERS_TO_PASS_THROUGH.value in e.response.text: - request_data = self._drop_extra_params_from_request_data(request_data, e.response.text) + elif ( + AzureFoundryErrorStrings.SET_EXTRA_PARAMETERS_TO_PASS_THROUGH.value + in e.response.text + ): + request_data = self._drop_extra_params_from_request_data( + request_data, e.response.text + ) data = drop_params_from_unprocessable_entity_error(e=e, data=request_data) return data - def _drop_extra_params_from_request_data(self, request_data: dict, error_text: str) -> dict: + def _drop_extra_params_from_request_data( + self, request_data: dict, error_text: str + ) -> dict: params_to_drop = self._extract_params_to_drop_from_error_text(error_text) if params_to_drop: for param in params_to_drop: @@ -295,7 +325,9 @@ class AzureAIStudioConfig(OpenAIConfig): request_data.pop(param, None) return request_data - def _extract_params_to_drop_from_error_text(self, error_text: str) -> Optional[List[str]]: + def _extract_params_to_drop_from_error_text( + self, error_text: str + ) -> Optional[List[str]]: """ Error text looks like this" "Extra parameters ['stream_options', 'extra-parameters'] are not allowed when extra-parameters is not set or set to be 'error'. diff --git a/litellm/llms/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py index 47d397d6e9..ecb36b2042 100644 --- a/litellm/llms/azure_ai/common_utils.py +++ b/litellm/llms/azure_ai/common_utils.py @@ -18,7 +18,7 @@ class AzureFoundryModelInfo(BaseLLMModelInfo): Get the Azure AI route for the given model. Similar to BedrockModelInfo.get_bedrock_route(). - + Supported routes: - agents: azure_ai/agents/ - model_router: azure_ai/model_router/ or models with "model-router"/"model_router" in name @@ -29,7 +29,7 @@ class AzureFoundryModelInfo(BaseLLMModelInfo): # Detect model router by prefix (model_router/) or by name containing "model-router"/"model_router" model_lower = model.lower() if ( - "model_router/" in model_lower + "model_router/" in model_lower or "model-router/" in model_lower or "model-router" in model_lower or "model_router" in model_lower @@ -78,7 +78,7 @@ class AzureFoundryModelInfo(BaseLLMModelInfo): ) -> List[str]: """ Returns a list of models supported by Azure AI. - + Azure AI doesn't have a standard model listing endpoint, so this returns an empty list. """ @@ -92,15 +92,15 @@ class AzureFoundryModelInfo(BaseLLMModelInfo): def strip_model_router_prefix(model: str) -> str: """ Strip the model_router prefix from model name. - + Examples: - "model_router/gpt-4o" -> "gpt-4o" - "model-router/gpt-4o" -> "gpt-4o" - "gpt-4o" -> "gpt-4o" - + Args: model: Model name potentially with model_router prefix - + Returns: Model name without the prefix """ @@ -109,15 +109,15 @@ class AzureFoundryModelInfo(BaseLLMModelInfo): if "model-router/" in model: return model.split("model-router/", 1)[1] return model - + @staticmethod def get_base_model(model: str) -> str: """ Get the base model name, stripping any Azure AI routing prefixes. - + Args: model: Model name potentially with routing prefixes - + Returns: Base model name """ @@ -129,32 +129,35 @@ class AzureFoundryModelInfo(BaseLLMModelInfo): def get_azure_ai_config_for_model(model: str): """ Get the appropriate Azure AI config class for the given model. - + Routes to specialized configs based on model type: - Model Router: AzureModelRouterConfig - - Claude models: AzureAnthropicConfig + - Claude models: AzureAnthropicConfig - Default: AzureAIStudioConfig - + Args: model: The model name - + Returns: The appropriate config instance """ azure_ai_route = AzureFoundryModelInfo.get_azure_ai_route(model) - + if azure_ai_route == "model_router": from litellm.llms.azure_ai.azure_model_router.transformation import ( AzureModelRouterConfig, ) + return AzureModelRouterConfig() elif "claude" in model.lower(): from litellm.llms.azure_ai.anthropic.transformation import ( AzureAnthropicConfig, ) + return AzureAnthropicConfig() else: from litellm.llms.azure_ai.chat.transformation import AzureAIStudioConfig + return AzureAIStudioConfig() def validate_environment( diff --git a/litellm/llms/azure_ai/cost_calculator.py b/litellm/llms/azure_ai/cost_calculator.py index 6fb2996267..3cca61b218 100644 --- a/litellm/llms/azure_ai/cost_calculator.py +++ b/litellm/llms/azure_ai/cost_calculator.py @@ -14,22 +14,22 @@ from litellm.utils import get_model_info def _is_azure_model_router(model: str) -> bool: """ Check if the model is Azure AI Foundry Model Router. - + Detects patterns like: - "azure-model-router" - - "model-router" + - "model-router" - "model_router/" - "model-router/" - + Args: model: The model name - + Returns: bool: True if this is a model router model """ model_lower = model.lower() return ( - "model-router" in model_lower + "model-router" in model_lower or "model_router" in model_lower or model_lower == "azure-model-router" ) @@ -38,50 +38,50 @@ def _is_azure_model_router(model: str) -> bool: def calculate_azure_model_router_flat_cost(model: str, prompt_tokens: int) -> float: """ Calculate the flat cost for Azure AI Foundry Model Router. - + Args: model: The model name (should be a model router model) prompt_tokens: Number of prompt tokens - + Returns: float: The flat cost in USD, or 0.0 if not applicable """ if not _is_azure_model_router(model): return 0.0 - + # Get the model router pricing from model_prices_and_context_window.json # Use "model_router" as the key (without actual model name suffix) model_info = get_model_info(model="model_router", custom_llm_provider="azure_ai") router_flat_cost_per_token = model_info.get("input_cost_per_token", 0) - + if router_flat_cost_per_token > 0: return prompt_tokens * router_flat_cost_per_token - + return 0.0 def cost_per_token( - model: str, - usage: Usage, + model: str, + usage: Usage, response_time_ms: Optional[float] = 0.0, request_model: Optional[str] = None, ) -> Tuple[float, float]: """ Calculate the cost per token for Azure AI models. - + For Azure AI Foundry Model Router: - Adds a flat cost of $0.14 per million input tokens (from model_prices_and_context_window.json) - Plus the cost of the actual model used (handled by generic_cost_per_token) - + Args: model: str, the model name without provider prefix (from response) usage: LiteLLM Usage block response_time_ms: Optional response time in milliseconds request_model: Optional[str], the original request model name (to detect router usage) - + Returns: Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd - + Raises: ValueError: If the model is not found in the cost map and cost cannot be calculated (except for Model Router models where we return just the routing flat cost) @@ -119,7 +119,9 @@ def cost_per_token( if is_router_request: # Use the request model for flat cost calculation if available, otherwise use response model router_model_for_calc = request_model if request_model else model - router_flat_cost = calculate_azure_model_router_flat_cost(router_model_for_calc, usage.prompt_tokens) + router_flat_cost = calculate_azure_model_router_flat_cost( + router_model_for_calc, usage.prompt_tokens + ) if router_flat_cost > 0: verbose_logger.debug( diff --git a/litellm/llms/azure_ai/image_edit/flux2_transformation.py b/litellm/llms/azure_ai/image_edit/flux2_transformation.py index 77d46ff917..0de163a771 100644 --- a/litellm/llms/azure_ai/image_edit/flux2_transformation.py +++ b/litellm/llms/azure_ai/image_edit/flux2_transformation.py @@ -101,10 +101,10 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): """ if prompt is None: raise ValueError("FLUX 2 image edit requires a prompt.") - + if image is None: raise ValueError("FLUX 2 image edit requires an image.") - + image_b64 = self._convert_image_to_base64(image) # Build request body with required params @@ -170,4 +170,3 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): model=model, api_version=api_version, ) - diff --git a/litellm/llms/azure_ai/image_generation/cost_calculator.py b/litellm/llms/azure_ai/image_generation/cost_calculator.py index 2fc7c554a3..b67de9cb70 100644 --- a/litellm/llms/azure_ai/image_generation/cost_calculator.py +++ b/litellm/llms/azure_ai/image_generation/cost_calculator.py @@ -22,4 +22,6 @@ def cost_calculator( num_images = len(image_response.data) return output_cost_per_image * num_images else: - raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}") + raise ValueError( + f"image_response must be of type ImageResponse got type={type(image_response)}" + ) diff --git a/litellm/llms/azure_ai/ocr/__init__.py b/litellm/llms/azure_ai/ocr/__init__.py index 7182a750b4..e49217a5ba 100644 --- a/litellm/llms/azure_ai/ocr/__init__.py +++ b/litellm/llms/azure_ai/ocr/__init__.py @@ -10,4 +10,3 @@ __all__ = [ "AzureDocumentIntelligenceOCRConfig", "get_azure_ai_ocr_config", ] - diff --git a/litellm/llms/azure_ai/ocr/common_utils.py b/litellm/llms/azure_ai/ocr/common_utils.py index ef470c7492..d736b89153 100644 --- a/litellm/llms/azure_ai/ocr/common_utils.py +++ b/litellm/llms/azure_ai/ocr/common_utils.py @@ -16,22 +16,22 @@ if TYPE_CHECKING: def get_azure_ai_ocr_config(model: str) -> Optional["BaseOCRConfig"]: """ Determine which Azure AI OCR configuration to use based on the model name. - + Azure AI supports multiple OCR services: - Azure Document Intelligence: azure_ai/doc-intelligence/ - Mistral OCR (via Azure AI): azure_ai/ - + Args: - model: The model name (e.g., "azure_ai/doc-intelligence/prebuilt-read", + model: The model name (e.g., "azure_ai/doc-intelligence/prebuilt-read", "azure_ai/pixtral-12b-2409") - + Returns: OCR configuration instance for the specified model - + Examples: >>> get_azure_ai_ocr_config("azure_ai/doc-intelligence/prebuilt-read") - + >>> get_azure_ai_ocr_config("azure_ai/pixtral-12b-2409") """ @@ -46,8 +46,7 @@ def get_azure_ai_ocr_config(model: str) -> Optional["BaseOCRConfig"]: f"Routing {model} to Azure Document Intelligence OCR config" ) return AzureDocumentIntelligenceOCRConfig() - + # Default to Mistral-based OCR for other azure_ai models verbose_logger.debug(f"Routing {model} to Azure AI (Mistral) OCR config") return AzureAIOCRConfig() - diff --git a/litellm/llms/azure_ai/ocr/document_intelligence/__init__.py b/litellm/llms/azure_ai/ocr/document_intelligence/__init__.py index 372a6a8d76..fb14fbbf0a 100644 --- a/litellm/llms/azure_ai/ocr/document_intelligence/__init__.py +++ b/litellm/llms/azure_ai/ocr/document_intelligence/__init__.py @@ -2,4 +2,3 @@ from .transformation import AzureDocumentIntelligenceOCRConfig __all__ = ["AzureDocumentIntelligenceOCRConfig"] - diff --git a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py index f6c6da2409..6ef309ca67 100644 --- a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py +++ b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py @@ -35,15 +35,15 @@ from litellm.secret_managers.main import get_secret_str class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): """ Azure Document Intelligence OCR transformation configuration. - + Supports Azure Document Intelligence v4.0 (2024-11-30) API. Model route: azure_ai/doc-intelligence/ - + Supported models: - prebuilt-layout: Extracts text with markdown, tables, and structure (closest to Mistral OCR) - prebuilt-read: Basic text extraction optimized for reading - prebuilt-document: General document analysis - + Reference: https://learn.microsoft.com/en-us/azure/ai-services/document-intelligence/ """ @@ -53,7 +53,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): def get_supported_ocr_params(self, model: str) -> list: """ Get supported OCR parameters for Azure Document Intelligence. - + Azure DI has minimal optional parameters compared to Mistral OCR. Most Mistral-specific params are ignored during transformation. """ @@ -70,7 +70,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): ) -> Dict: """ Validate environment and return headers for Azure Document Intelligence. - + Authentication uses Ocp-Apim-Subscription-Key header. """ # Get API key from environment if not provided @@ -109,16 +109,16 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): ) -> str: """ Get complete URL for Azure Document Intelligence endpoint. - + Format: {endpoint}/documentintelligence/documentModels/{modelId}:analyze?api-version=2024-11-30 - + Note: API version 2024-11-30 uses /documentintelligence/ path (not /formrecognizer/) - + Args: api_base: Azure Document Intelligence endpoint (e.g., https://your-resource.cognitiveservices.azure.com) model: Model ID (e.g., "prebuilt-layout", "prebuilt-read") optional_params: Optional parameters - + Returns: Complete URL for Azure DI analyze endpoint """ if api_base is None: @@ -146,10 +146,10 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): def _extract_base64_from_data_uri(self, data_uri: str) -> str: """ Extract base64 content from a data URI. - + Args: data_uri: Data URI like "data:application/pdf;base64,..." - + Returns: Base64 string without the data URI prefix """ @@ -169,7 +169,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): ) -> OCRRequestData: """ Transform OCR request to Azure Document Intelligence format. - + Mistral OCR format: { "document": { @@ -177,7 +177,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): "document_url": "https://example.com/doc.pdf" } } - + Azure DI format: { "urlSource": "https://example.com/doc.pdf" @@ -186,13 +186,13 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): { "base64Source": "base64_encoded_content" } - + Args: model: Model name document: Document dict from user (Mistral format) optional_params: Already mapped optional parameters headers: Request headers - + Returns: OCRRequestData with JSON data """ @@ -241,12 +241,12 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): def _extract_page_markdown(self, page_data: Dict[str, Any]) -> str: """ Extract text from Azure DI page and format as markdown. - + Azure DI provides text in 'lines' array. We concatenate them with newlines. - + Args: page_data: Azure DI page object - + Returns: Markdown-formatted text """ @@ -265,14 +265,14 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): ) -> OCRPageDimensions: """ Convert Azure DI dimensions to pixels. - + Azure DI provides dimensions in inches. We convert to pixels using configured DPI. - + Args: width: Width in specified unit height: Height in specified unit unit: Unit of measurement (e.g., "inch") - + Returns: OCRPageDimensions with pixel values """ @@ -292,11 +292,11 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): def _check_timeout(start_time: float, timeout_secs: int) -> None: """ Check if operation has timed out. - + Args: start_time: Start time of the operation timeout_secs: Timeout duration in seconds - + Raises: TimeoutError: If operation has exceeded timeout """ @@ -309,10 +309,10 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): def _get_retry_after(response: httpx.Response) -> int: """ Get retry-after duration from response headers. - + Args: response: HTTP response - + Returns: Retry-after duration in seconds (default: 2) """ @@ -324,13 +324,13 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): def _check_operation_status(response: httpx.Response) -> str: """ Check Azure DI operation status from response. - + Args: response: HTTP response from operation endpoint - + Returns: Operation status string - + Raises: ValueError: If operation failed or status is unknown """ @@ -366,15 +366,15 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): ) -> httpx.Response: """ Poll Azure Document Intelligence operation until completion (sync). - + Azure DI POST returns 202 with Operation-Location header. We need to poll that URL until status is "succeeded" or "failed". - + Args: operation_url: The Operation-Location URL to poll headers: Request headers (including auth) timeout_secs: Total timeout in seconds - + Returns: Final response with completed analysis """ @@ -409,12 +409,12 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): ) -> httpx.Response: """ Poll Azure Document Intelligence operation until completion (async). - + Args: operation_url: The Operation-Location URL to poll headers: Request headers (including auth) timeout_secs: Total timeout in seconds - + Returns: Final response with completed analysis """ @@ -451,10 +451,10 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): ) -> OCRResponse: """ Transform Azure Document Intelligence response to Mistral OCR format. - + Handles async operation polling: If response is 202 Accepted, polls Operation-Location until analysis completes. - + Azure DI response (after polling): { "status": "succeeded", @@ -471,7 +471,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): ] } } - + Mistral OCR format: { "pages": [ @@ -485,12 +485,12 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): "usage_info": {"pages_processed": 1}, "object": "ocr" } - + Args: model: Model name raw_response: Raw HTTP response from Azure DI (may be 202 Accepted) logging_obj: Logging object - + Returns: OCRResponse in Mistral format """ @@ -594,15 +594,15 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): ) -> OCRResponse: """ Async transform Azure Document Intelligence response to Mistral OCR format. - + Handles async operation polling: If response is 202 Accepted, polls Operation-Location until analysis completes using async polling. - + Args: model: Model name raw_response: Raw HTTP response from Azure DI (may be 202 Accepted) logging_obj: Logging object - + Returns: OCRResponse in Mistral format """ @@ -696,4 +696,3 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): f"Error parsing Azure Document Intelligence response (async): {e}" ) raise e - diff --git a/litellm/llms/azure_ai/ocr/transformation.py b/litellm/llms/azure_ai/ocr/transformation.py index 24fc9e8613..8f57bb3358 100644 --- a/litellm/llms/azure_ai/ocr/transformation.py +++ b/litellm/llms/azure_ai/ocr/transformation.py @@ -16,12 +16,12 @@ from litellm.secret_managers.main import get_secret_str class AzureAIOCRConfig(MistralOCRConfig): """ Azure AI OCR transformation configuration. - + Azure AI uses Mistral's OCR API but with a different endpoint format. Inherits transformation logic from MistralOCRConfig since they use the same format. - + Reference: Azure AI Foundry OCR documentation - + Important: Azure AI only supports base64 data URIs (data:image/..., data:application/pdf;base64,...). Regular URLs are not supported. """ @@ -40,7 +40,7 @@ class AzureAIOCRConfig(MistralOCRConfig): ) -> Dict: """ Validate environment and return headers for Azure AI OCR. - + Azure AI uses Bearer token authentication with AZURE_AI_API_KEY. """ # Get API key from environment if not provided @@ -55,7 +55,7 @@ class AzureAIOCRConfig(MistralOCRConfig): # Validate API base is provided if api_base is None: api_base = get_secret_str("AZURE_AI_API_BASE") - + if api_base is None: raise ValueError( "Missing Azure AI API Base - Set AZURE_AI_API_BASE environment variable or pass api_base parameter" @@ -79,14 +79,14 @@ class AzureAIOCRConfig(MistralOCRConfig): ) -> str: """ Get complete URL for Azure AI OCR endpoint. - + Azure AI endpoint format: https:///providers/mistral/azure/ocr - + Args: api_base: Azure AI API base URL model: Model name (not used in URL construction) optional_params: Optional parameters - + Returns: Complete URL for Azure AI OCR endpoint """ if api_base is None: @@ -96,54 +96,62 @@ class AzureAIOCRConfig(MistralOCRConfig): # Ensure no trailing slash api_base = api_base.rstrip("/") - + # Azure AI OCR endpoint format return f"{api_base}/providers/mistral/azure/ocr" def _convert_url_to_data_uri_sync(self, url: str) -> str: """ Synchronously convert a URL to a base64 data URI. - + Azure AI OCR doesn't have internet access, so we need to fetch URLs and convert them to base64 data URIs. - + Args: url: The URL to convert - + Returns: Base64 data URI string """ - verbose_logger.debug(f"Azure AI OCR: Converting URL to base64 data URI (sync): {url}") - + verbose_logger.debug( + f"Azure AI OCR: Converting URL to base64 data URI (sync): {url}" + ) + # Fetch and convert to base64 data URI # convert_url_to_base64 already returns a full data URI like "data:image/jpeg;base64,..." data_uri = convert_url_to_base64(url=url) - - verbose_logger.debug(f"Azure AI OCR: Converted URL to data URI (length: {len(data_uri)})") - + + verbose_logger.debug( + f"Azure AI OCR: Converted URL to data URI (length: {len(data_uri)})" + ) + return data_uri async def _convert_url_to_data_uri_async(self, url: str) -> str: """ Asynchronously convert a URL to a base64 data URI. - + Azure AI OCR doesn't have internet access, so we need to fetch URLs and convert them to base64 data URIs. - + Args: url: The URL to convert - + Returns: Base64 data URI string """ - verbose_logger.debug(f"Azure AI OCR: Converting URL to base64 data URI (async): {url}") - + verbose_logger.debug( + f"Azure AI OCR: Converting URL to base64 data URI (async): {url}" + ) + # Fetch and convert to base64 data URI asynchronously # async_convert_url_to_base64 already returns a full data URI like "data:image/jpeg;base64,..." data_uri = await async_convert_url_to_base64(url=url) - - verbose_logger.debug(f"Azure AI OCR: Converted URL to data URI (length: {len(data_uri)})") - + + verbose_logger.debug( + f"Azure AI OCR: Converted URL to data URI (length: {len(data_uri)})" + ) + return data_uri def transform_ocr_request( @@ -156,29 +164,31 @@ class AzureAIOCRConfig(MistralOCRConfig): ) -> OCRRequestData: """ Transform OCR request for Azure AI, converting URLs to base64 data URIs (sync). - + Azure AI OCR doesn't have internet access, so we automatically fetch any URLs and convert them to base64 data URIs synchronously. - + Args: model: Model name document: Document dict from user optional_params: Already mapped optional parameters headers: Request headers **kwargs: Additional arguments - + Returns: OCRRequestData with JSON data """ - verbose_logger.debug(f"Azure AI OCR transform_ocr_request (sync) - model: {model}") - + verbose_logger.debug( + f"Azure AI OCR transform_ocr_request (sync) - model: {model}" + ) + if not isinstance(document, dict): raise ValueError(f"Expected document dict, got {type(document)}") - + # Check if we need to convert URL to base64 doc_type = document.get("type") transformed_document = document.copy() - + if doc_type == "document_url": document_url = document.get("document_url", "") # If it's not already a data URI, convert it @@ -197,7 +207,7 @@ class AzureAIOCRConfig(MistralOCRConfig): ) data_uri = self._convert_url_to_data_uri_sync(url=image_url) transformed_document["image_url"] = data_uri - + # Call parent's transform to build the request return super().transform_ocr_request( model=model, @@ -217,29 +227,31 @@ class AzureAIOCRConfig(MistralOCRConfig): ) -> OCRRequestData: """ Transform OCR request for Azure AI, converting URLs to base64 data URIs (async). - + Azure AI OCR doesn't have internet access, so we automatically fetch any URLs and convert them to base64 data URIs asynchronously. - + Args: model: Model name document: Document dict from user optional_params: Already mapped optional parameters headers: Request headers **kwargs: Additional arguments - + Returns: OCRRequestData with JSON data """ - verbose_logger.debug(f"Azure AI OCR async_transform_ocr_request - model: {model}") - + verbose_logger.debug( + f"Azure AI OCR async_transform_ocr_request - model: {model}" + ) + if not isinstance(document, dict): raise ValueError(f"Expected document dict, got {type(document)}") - + # Check if we need to convert URL to base64 doc_type = document.get("type") transformed_document = document.copy() - + if doc_type == "document_url": document_url = document.get("document_url", "") # If it's not already a data URI, convert it @@ -258,7 +270,7 @@ class AzureAIOCRConfig(MistralOCRConfig): ) data_uri = await self._convert_url_to_data_uri_async(url=image_url) transformed_document["image_url"] = data_uri - + # Call parent's transform to build the request return super().transform_ocr_request( model=model, @@ -267,4 +279,3 @@ class AzureAIOCRConfig(MistralOCRConfig): headers=headers, **kwargs, ) - diff --git a/litellm/llms/azure_ai/rerank/transformation.py b/litellm/llms/azure_ai/rerank/transformation.py index f577a42ed5..b5993040ea 100644 --- a/litellm/llms/azure_ai/rerank/transformation.py +++ b/litellm/llms/azure_ai/rerank/transformation.py @@ -20,8 +20,8 @@ class AzureAIRerankConfig(CohereRerankConfig): """ def get_complete_url( - self, - api_base: Optional[str], + self, + api_base: Optional[str], model: str, optional_params: Optional[dict] = None, ) -> str: @@ -41,7 +41,9 @@ class AzureAIRerankConfig(CohereRerankConfig): # Allow callers to pass either full v1/v2 rerank endpoints: # - https://.services.ai.azure.com/v1/rerank # - https://.services.ai.azure.com/providers/cohere/v2/rerank - if normalized_path.endswith("/v1/rerank") or normalized_path.endswith("/v2/rerank"): + if normalized_path.endswith("/v1/rerank") or normalized_path.endswith( + "/v2/rerank" + ): return str(original_url.copy_with(path=normalized_path or "/")) # If callers pass just the version path (e.g. ".../v2" or ".../providers/cohere/v2"), append "/rerank" diff --git a/litellm/llms/azure_ai/vector_stores/__init__.py b/litellm/llms/azure_ai/vector_stores/__init__.py index 74ffe1afb1..d83363cbc5 100644 --- a/litellm/llms/azure_ai/vector_stores/__init__.py +++ b/litellm/llms/azure_ai/vector_stores/__init__.py @@ -1,4 +1,3 @@ from litellm.llms.azure_ai.vector_stores.transformation import AzureAIVectorStoreConfig __all__ = ["AzureAIVectorStoreConfig"] - diff --git a/litellm/llms/azure_ai/vector_stores/transformation.py b/litellm/llms/azure_ai/vector_stores/transformation.py index 96cea064ce..b62acb6516 100644 --- a/litellm/llms/azure_ai/vector_stores/transformation.py +++ b/litellm/llms/azure_ai/vector_stores/transformation.py @@ -58,7 +58,6 @@ class AzureAIVectorStoreConfig(BaseVectorStoreConfig, BaseAzureLLM): def validate_environment( self, headers: dict, litellm_params: Optional[GenericLiteLLMParams] ) -> dict: - basic_headers = self._base_validate_azure_environment(headers, litellm_params) basic_headers.update({"Content-Type": "application/json"}) return basic_headers diff --git a/litellm/llms/base_llm/base_model_iterator.py b/litellm/llms/base_llm/base_model_iterator.py index 6953b1c587..9bccf25c45 100644 --- a/litellm/llms/base_llm/base_model_iterator.py +++ b/litellm/llms/base_llm/base_model_iterator.py @@ -125,26 +125,34 @@ class BaseModelResponseIterator: ) def __next__(self): - try: - chunk = self.response_iterator.__next__() - except StopIteration: - raise StopIteration - except ValueError as e: - raise RuntimeError(f"Error receiving chunk from stream: {e}") + while True: + try: + chunk = self.response_iterator.__next__() + except StopIteration: + raise StopIteration + except ValueError as e: + raise RuntimeError(f"Error receiving chunk from stream: {e}") - try: - str_line = chunk - if isinstance(chunk, bytes): # Handle binary data - str_line = chunk.decode("utf-8") # Convert bytes to string - index = str_line.find("data:") - if index != -1: - str_line = str_line[index:] - # chunk is a str at this point - return self._handle_string_chunk(str_line=str_line) - except StopIteration: - raise StopIteration - except ValueError as e: - raise RuntimeError(f"Error parsing chunk: {e},\nReceived chunk: {chunk}") + try: + str_line = chunk + if isinstance(chunk, bytes): # Handle binary data + str_line = chunk.decode("utf-8") # Convert bytes to string + index = str_line.find("data:") + if index != -1: + str_line = str_line[index:] + + # Skip empty lines (common in SSE streams between events) + if not str_line or not str_line.strip(): + continue + + # chunk is a str at this point + return self._handle_string_chunk(str_line=str_line) + except StopIteration: + raise StopIteration + except ValueError as e: + raise RuntimeError( + f"Error parsing chunk: {e},\nReceived chunk: {chunk}" + ) # Async iterator def __aiter__(self): @@ -152,30 +160,37 @@ class BaseModelResponseIterator: return self async def __anext__(self): - try: - chunk = await self.async_response_iterator.__anext__() + while True: + try: + chunk = await self.async_response_iterator.__anext__() - except StopAsyncIteration: - raise StopAsyncIteration - except ValueError as e: - raise RuntimeError(f"Error receiving chunk from stream: {e}") + except StopAsyncIteration: + raise StopAsyncIteration + except ValueError as e: + raise RuntimeError(f"Error receiving chunk from stream: {e}") - try: - str_line = chunk - if isinstance(chunk, bytes): # Handle binary data - str_line = chunk.decode("utf-8") # Convert bytes to string - index = str_line.find("data:") - if index != -1: - str_line = str_line[index:] + try: + str_line = chunk + if isinstance(chunk, bytes): # Handle binary data + str_line = chunk.decode("utf-8") # Convert bytes to string + index = str_line.find("data:") + if index != -1: + str_line = str_line[index:] - # chunk is a str at this point - chunk = self._handle_string_chunk(str_line=str_line) + # Skip empty lines (common in SSE streams between events) + if not str_line or not str_line.strip(): + continue - return chunk - except StopAsyncIteration: - raise StopAsyncIteration - except ValueError as e: - raise RuntimeError(f"Error parsing chunk: {e},\nReceived chunk: {chunk}") + # chunk is a str at this point + chunk = self._handle_string_chunk(str_line=str_line) + + return chunk + except StopAsyncIteration: + raise StopAsyncIteration + except ValueError as e: + raise RuntimeError( + f"Error parsing chunk: {e},\nReceived chunk: {chunk}" + ) class MockResponseIterator: # for returning ai21 streaming responses diff --git a/litellm/llms/base_llm/base_utils.py b/litellm/llms/base_llm/base_utils.py index ecff9053dc..d2d3d5c0a9 100644 --- a/litellm/llms/base_llm/base_utils.py +++ b/litellm/llms/base_llm/base_utils.py @@ -98,7 +98,7 @@ class BaseLLMModelInfo(ABC): def get_token_counter(self) -> Optional[BaseTokenCounter]: """ Factory method to create a token counter for this provider. - + Returns: Optional TokenCounterInterface implementation for this provider, or None if token counting is not supported. diff --git a/litellm/llms/base_llm/batches/transformation.py b/litellm/llms/base_llm/batches/transformation.py index 9e67689fcd..aedaf0687c 100644 --- a/litellm/llms/base_llm/batches/transformation.py +++ b/litellm/llms/base_llm/batches/transformation.py @@ -26,7 +26,7 @@ else: class BaseBatchesConfig(ABC): """ Abstract base class for batch processing configurations across different LLM providers. - + This class defines the interface that all provider-specific batch configurations must implement to work with LiteLLM's unified batch processing system. """ @@ -73,7 +73,7 @@ class BaseBatchesConfig(ABC): ) -> dict: """ Validate and prepare environment-specific headers and parameters. - + Args: headers: HTTP headers dictionary model: Model name @@ -82,7 +82,7 @@ class BaseBatchesConfig(ABC): litellm_params: LiteLLM parameters api_key: API key api_base: API base URL - + Returns: Updated headers dictionary """ @@ -100,7 +100,7 @@ class BaseBatchesConfig(ABC): ) -> str: """ Get the complete URL for batch creation request. - + Args: api_base: Base API URL api_key: API key @@ -108,7 +108,7 @@ class BaseBatchesConfig(ABC): optional_params: Optional parameters litellm_params: LiteLLM parameters data: Batch creation request data - + Returns: Complete URL for the batch request """ @@ -124,13 +124,13 @@ class BaseBatchesConfig(ABC): ) -> Union[bytes, str, Dict[str, Any]]: """ Transform the batch creation request to provider-specific format. - + Args: model: Model name create_batch_data: Batch creation request data optional_params: Optional parameters litellm_params: LiteLLM parameters - + Returns: Transformed request data """ @@ -146,13 +146,13 @@ class BaseBatchesConfig(ABC): ) -> LiteLLMBatch: """ Transform provider-specific batch response to LiteLLM format. - + Args: model: Model name raw_response: Raw HTTP response logging_obj: Logging object litellm_params: LiteLLM parameters - + Returns: LiteLLM batch object """ @@ -167,12 +167,12 @@ class BaseBatchesConfig(ABC): ) -> Union[bytes, str, Dict[str, Any]]: """ Transform the batch retrieval request to provider-specific format. - + Args: batch_id: Batch ID to retrieve optional_params: Optional parameters litellm_params: LiteLLM parameters - + Returns: Transformed request data """ @@ -188,13 +188,13 @@ class BaseBatchesConfig(ABC): ) -> LiteLLMBatch: """ Transform provider-specific batch retrieval response to LiteLLM format. - + Args: model: Model name raw_response: Raw HTTP response logging_obj: Logging object litellm_params: LiteLLM parameters - + Returns: LiteLLM batch object """ @@ -206,12 +206,12 @@ class BaseBatchesConfig(ABC): ) -> "BaseLLMException": """ Get the appropriate error class for this provider. - + Args: error_message: Error message status_code: HTTP status code headers: Response headers - + Returns: Provider-specific exception class """ diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py index f22c8ee0d9..b71ae0fdde 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/litellm/llms/base_llm/chat/transformation.py @@ -135,7 +135,10 @@ class BaseConfig(ABC): if 'thinking' is enabled and 'max_tokens' or 'max_completion_tokens' is not specified, set 'max_tokens' to the thinking token budget + DEFAULT_MAX_TOKENS """ is_thinking_enabled = self.is_thinking_enabled(optional_params) - if is_thinking_enabled and ("max_tokens" not in non_default_params and "max_completion_tokens" not in non_default_params): + if is_thinking_enabled and ( + "max_tokens" not in non_default_params + and "max_completion_tokens" not in non_default_params + ): thinking_token_budget = cast(dict, optional_params["thinking"]).get( "budget_tokens", None ) @@ -447,14 +450,14 @@ class BaseConfig(ABC): ) -> Optional[dict]: """ Calculate any additional costs beyond standard token costs. - + This is used for provider-specific infrastructure costs, routing fees, etc. - + Args: model: The model name prompt_tokens: Number of prompt tokens completion_tokens: Number of completion tokens - + Returns: Optional dictionary with cost names and amounts, e.g.: {"Infrastructure Fee": 0.001, "Routing Cost": 0.0005} diff --git a/litellm/llms/base_llm/containers/transformation.py b/litellm/llms/base_llm/containers/transformation.py index 5ce374c773..dc75789156 100644 --- a/litellm/llms/base_llm/containers/transformation.py +++ b/litellm/llms/base_llm/containers/transformation.py @@ -89,7 +89,7 @@ class BaseContainerConfig(ABC): litellm_params: dict, ) -> str: """Get the complete url for the request. - + OPTIONAL - Some providers need `model` in `api_base`. """ if api_base is None: @@ -106,7 +106,7 @@ class BaseContainerConfig(ABC): headers: dict, ) -> dict: """Transform the container creation request. - + Returns: dict: Request data for container creation. """ @@ -133,7 +133,7 @@ class BaseContainerConfig(ABC): extra_query: dict[str, Any] | None = None, ) -> tuple[str, dict]: """Transform the container list request into a URL and params. - + Returns: tuple[str, dict]: (url, params) for the container list request. """ @@ -157,7 +157,7 @@ class BaseContainerConfig(ABC): headers: dict, ) -> tuple[str, dict]: """Transform the container retrieve request into a URL and data/params. - + Returns: tuple[str, dict]: (url, params) for the container retrieve request. """ @@ -181,7 +181,7 @@ class BaseContainerConfig(ABC): headers: dict, ) -> tuple[str, dict]: """Transform the container delete request into a URL and data. - + Returns: tuple[str, dict]: (url, data) for the container delete request. """ @@ -209,7 +209,7 @@ class BaseContainerConfig(ABC): extra_query: dict[str, Any] | None = None, ) -> tuple[str, dict]: """Transform the container file list request into a URL and params. - + Returns: tuple[str, dict]: (url, params) for the container file list request. """ @@ -234,7 +234,7 @@ class BaseContainerConfig(ABC): headers: dict, ) -> tuple[str, dict]: """Transform the container file content request into a URL and params. - + Returns: tuple[str, dict]: (url, params) for the container file content request. """ @@ -247,16 +247,16 @@ class BaseContainerConfig(ABC): logging_obj: LiteLLMLoggingObj, ) -> bytes: """Transform the container file content response. - + Returns: bytes: The raw file content. """ ... def get_error_class( - self, - error_message: str, - status_code: int, + self, + error_message: str, + status_code: int, headers: dict | httpx.Headers, ) -> BaseLLMException: from ..chat.transformation import BaseLLMException @@ -266,4 +266,3 @@ class BaseContainerConfig(ABC): message=error_message, headers=headers, ) - diff --git a/litellm/llms/base_llm/files/azure_blob_storage_backend.py b/litellm/llms/base_llm/files/azure_blob_storage_backend.py index db3aa50d89..a2155df404 100644 --- a/litellm/llms/base_llm/files/azure_blob_storage_backend.py +++ b/litellm/llms/base_llm/files/azure_blob_storage_backend.py @@ -20,26 +20,26 @@ from litellm.integrations.azure_storage.azure_storage import AzureBlobStorageLog class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): """ Azure Blob Storage backend implementation. - + Inherits from AzureBlobStorageLogger to reuse: - Authentication (account key and Azure AD) - Service client management - Token management - All Azure Storage helper methods - + Reads configuration from the same environment variables as AzureBlobStorageLogger. """ def __init__(self, **kwargs): """ Initialize Azure Blob Storage backend. - + Inherits all functionality from AzureBlobStorageLogger which handles: - Reading environment variables - Authentication (account key and Azure AD) - Service client management - Token management - + Environment variables (same as AzureBlobStorageLogger): - AZURE_STORAGE_ACCOUNT_NAME (required) - AZURE_STORAGE_FILE_SYSTEM (required) @@ -47,12 +47,12 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): - AZURE_STORAGE_TENANT_ID (optional, if using Azure AD) - AZURE_STORAGE_CLIENT_ID (optional, if using Azure AD) - AZURE_STORAGE_CLIENT_SECRET (optional, if using Azure AD) - + Note: We skip periodic_flush since we're not using this as a logger. """ # Initialize AzureBlobStorageLogger (handles all auth and config) AzureBlobStorageLogger.__init__(self, **kwargs) - + # Disable logging functionality - we're only using this for file storage # The periodic_flush task will be created but will do nothing since we override it @@ -87,12 +87,16 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): return quote(original_filename, safe="") elif file_naming_strategy == "timestamp": # Use timestamp - extension = original_filename.split(".")[-1] if "." in original_filename else "" + extension = ( + original_filename.split(".")[-1] if "." in original_filename else "" + ) timestamp = int(time.time() * 1000) # milliseconds return f"{timestamp}.{extension}" if extension else str(timestamp) else: # default to "uuid" # Use UUID - extension = original_filename.split(".")[-1] if "." in original_filename else "" + extension = ( + original_filename.split(".")[-1] if "." in original_filename else "" + ) file_uuid = str(uuid.uuid4()) return f"{file_uuid}.{extension}" if extension else file_uuid @@ -106,13 +110,13 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): ) -> str: """ Upload a file to Azure Blob Storage. - + Returns the blob URL in format: https://{account}.blob.core.windows.net/{container}/{path} """ try: # Generate file name file_name = self._generate_file_name(filename, file_naming_strategy) - + # Build full path if path_prefix: # Remove leading/trailing slashes and normalize @@ -140,7 +144,9 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): return storage_url except Exception as e: - verbose_logger.exception(f"Error uploading file to Azure Blob Storage: {str(e)}") + verbose_logger.exception( + f"Error uploading file to Azure Blob Storage: {str(e)}" + ) raise async def _upload_file_with_account_key( @@ -156,20 +162,22 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): # Create filesystem (container) if it doesn't exist if not await file_system_client.exists(): await file_system_client.create_file_system() - verbose_logger.debug(f"Created filesystem: {self.azure_storage_file_system}") + verbose_logger.debug( + f"Created filesystem: {self.azure_storage_file_system}" + ) # Extract directory and filename (similar to logger's pattern) path_parts = full_path.split("/") if len(path_parts) > 1: directory_path = "/".join(path_parts[:-1]) file_name = path_parts[-1] - + # Create directory if needed (like logger does) directory_client = file_system_client.get_directory_client(directory_path) if not await directory_client.exists(): await directory_client.create_directory() verbose_logger.debug(f"Created directory: {directory_path}") - + # Get file client from directory (same pattern as logger) file_client = directory_client.get_file_client(file_name) else: @@ -178,7 +186,9 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): # Create, append, and flush (same pattern as logger's upload_to_azure_data_lake_with_azure_account_key) await file_client.create_file() - await file_client.append_data(data=file_content, offset=0, length=len(file_content)) + await file_client.append_data( + data=file_content, offset=0, length=len(file_content) + ) await file_client.flush_data(position=len(file_content), offset=0) # Return blob URL (not DFS URL) @@ -191,12 +201,12 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): """Upload file using REST API with Azure AD authentication.""" # Reuse the logger's token management await self.set_valid_azure_ad_token() - + from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) - + async_client = get_async_httpx_client( llm_provider=httpxSpecialProvider.LoggingCallback ) @@ -215,12 +225,10 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): blob_url = f"https://{self.azure_storage_account_name}.blob.core.windows.net/{self.azure_storage_file_system}/{full_path}" return blob_url - async def _append_data_bytes( - self, client, base_url: str, file_content: bytes - ): + async def _append_data_bytes(self, client, base_url: str, file_content: bytes): """Append binary data to file using REST API.""" from litellm.constants import AZURE_STORAGE_MSFT_VERSION - + headers = { "x-ms-version": AZURE_STORAGE_MSFT_VERSION, "Content-Type": "application/octet-stream", @@ -236,10 +244,10 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): async def download_file(self, storage_url: str) -> bytes: """ Download a file from Azure Blob Storage. - + Args: storage_url: Blob URL in format: https://{account}.blob.core.windows.net/{container}/{path} - + Returns: bytes: File content """ @@ -253,7 +261,9 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): container_and_path = storage_url.split(".blob.core.windows.net/", 1)[1] path_parts = container_and_path.split("/", 1) if len(path_parts) < 2: - raise ValueError(f"Invalid Azure Blob Storage URL format: {storage_url}") + raise ValueError( + f"Invalid Azure Blob Storage URL format: {storage_url}" + ) file_path = path_parts[1] # Path after container name if self.azure_storage_account_key: @@ -264,7 +274,9 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): return await self._download_file_with_azure_ad(file_path) except Exception as e: - verbose_logger.exception(f"Error downloading file from Azure Blob Storage: {str(e)}") + verbose_logger.exception( + f"Error downloading file from Azure Blob Storage: {str(e)}" + ) raise async def _download_file_with_account_key(self, file_path: str) -> bytes: @@ -276,7 +288,9 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): ) # Ensure filesystem exists (should already exist, but check for safety) if not await file_system_client.exists(): - raise ValueError(f"Filesystem {self.azure_storage_file_system} does not exist") + raise ValueError( + f"Filesystem {self.azure_storage_file_system} does not exist" + ) file_client = file_system_client.get_file_client(file_path) # Download file download_response = await file_client.download_file() @@ -287,7 +301,7 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): """Download file using REST API with Azure AD token.""" # Reuse the logger's token management await self.set_valid_azure_ad_token() - + from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -300,13 +314,12 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): # Use blob endpoint for download (simpler than DFS) blob_url = f"https://{self.azure_storage_account_name}.blob.core.windows.net/{self.azure_storage_file_system}/{file_path}" - + headers = { "x-ms-version": AZURE_STORAGE_MSFT_VERSION, "Authorization": f"Bearer {self.azure_auth_token}", } - + response = await async_client.get(blob_url, headers=headers) response.raise_for_status() return response.content - diff --git a/litellm/llms/base_llm/files/storage_backend.py b/litellm/llms/base_llm/files/storage_backend.py index d957045295..31e68a7002 100644 --- a/litellm/llms/base_llm/files/storage_backend.py +++ b/litellm/llms/base_llm/files/storage_backend.py @@ -12,7 +12,7 @@ from typing import Optional class BaseFileStorageBackend(ABC): """ Abstract base class for file storage backends. - + All storage backends (Azure Blob Storage, S3, GCS, etc.) must implement these methods to provide a consistent interface for file operations. """ @@ -28,17 +28,17 @@ class BaseFileStorageBackend(ABC): ) -> str: """ Upload a file to the storage backend. - + Args: file_content: The file content as bytes filename: Original filename (may be used for naming strategy) content_type: MIME type of the file path_prefix: Optional path prefix for organizing files file_naming_strategy: Strategy for naming files ("uuid", "timestamp", "original_filename") - + Returns: str: The storage URL where the file can be accessed/downloaded - + Raises: Exception: If upload fails """ @@ -48,13 +48,13 @@ class BaseFileStorageBackend(ABC): async def download_file(self, storage_url: str) -> bytes: """ Download a file from the storage backend. - + Args: storage_url: The storage URL returned from upload_file - + Returns: bytes: The file content - + Raises: Exception: If download fails """ @@ -63,17 +63,16 @@ class BaseFileStorageBackend(ABC): async def delete_file(self, storage_url: str) -> None: """ Delete a file from the storage backend. - + This is optional and can be overridden by backends that support deletion. Default implementation does nothing. - + Args: storage_url: The storage URL of the file to delete - + Raises: Exception: If deletion fails """ # Default implementation: no-op # Backends can override if they support deletion pass - diff --git a/litellm/llms/base_llm/files/storage_backend_factory.py b/litellm/llms/base_llm/files/storage_backend_factory.py index 1685f3fbd2..12047f1122 100644 --- a/litellm/llms/base_llm/files/storage_backend_factory.py +++ b/litellm/llms/base_llm/files/storage_backend_factory.py @@ -15,22 +15,22 @@ from .storage_backend import BaseFileStorageBackend def get_storage_backend(backend_type: str) -> BaseFileStorageBackend: """ Factory function to create a storage backend instance. - + Backends are configured using the same environment variables as their corresponding callbacks. For example, "azure_storage" uses the same env vars as AzureBlobStorageLogger. - + Args: backend_type: Backend type identifier (e.g., "azure_storage") - + Returns: BaseFileStorageBackend: Instance of the appropriate storage backend - + Raises: ValueError: If backend_type is not supported """ verbose_logger.debug(f"Creating storage backend: type={backend_type}") - + if backend_type == "azure_storage": return AzureBlobStorageBackend() else: @@ -38,4 +38,3 @@ def get_storage_backend(backend_type: str) -> BaseFileStorageBackend: f"Unsupported storage backend type: {backend_type}. " f"Supported types: azure_storage" ) - diff --git a/litellm/llms/base_llm/files/transformation.py b/litellm/llms/base_llm/files/transformation.py index 58df15f0c4..c3abfafc55 100644 --- a/litellm/llms/base_llm/files/transformation.py +++ b/litellm/llms/base_llm/files/transformation.py @@ -81,7 +81,7 @@ class BaseFilesConfig(BaseConfig): ) -> Union[dict, str, bytes, "TwoStepFileUploadConfig"]: """ Transform OpenAI-style file creation request into provider-specific format. - + Returns: - dict: For pre-signed single-step uploads (e.g., Bedrock S3) - str/bytes: For traditional file uploads diff --git a/litellm/llms/base_llm/google_genai/transformation.py b/litellm/llms/base_llm/google_genai/transformation.py index 0a85e127bd..e8b3bf1a57 100644 --- a/litellm/llms/base_llm/google_genai/transformation.py +++ b/litellm/llms/base_llm/google_genai/transformation.py @@ -18,7 +18,7 @@ else: GenerateContentResponse = Any LiteLLMLoggingObj = Any ToolConfigDict = Any - + from litellm.types.router import GenericLiteLLMParams @@ -58,8 +58,9 @@ class BaseGoogleGenAIGenerateContentConfig(ABC): Returns: List of supported parameter names """ - raise NotImplementedError("get_supported_generate_content_optional_params is not implemented") - + raise NotImplementedError( + "get_supported_generate_content_optional_params is not implemented" + ) @abstractmethod def map_generate_content_optional_params( @@ -77,15 +78,17 @@ class BaseGoogleGenAIGenerateContentConfig(ABC): Returns: Mapped parameters for the provider """ - raise NotImplementedError("map_generate_content_optional_params is not implemented") + raise NotImplementedError( + "map_generate_content_optional_params is not implemented" + ) @abstractmethod def validate_environment( - self, + self, api_key: Optional[str], headers: Optional[dict], model: str, - litellm_params: Optional[Union[GenericLiteLLMParams, dict]] + litellm_params: Optional[Union[GenericLiteLLMParams, dict]], ) -> dict: """ Validate the environment and return headers for the request. @@ -100,7 +103,7 @@ class BaseGoogleGenAIGenerateContentConfig(ABC): Updated headers """ raise NotImplementedError("validate_environment is not implemented") - + def sync_get_auth_token_and_url( self, api_base: Optional[str], @@ -121,7 +124,7 @@ class BaseGoogleGenAIGenerateContentConfig(ABC): Tuple of headers and API base """ raise NotImplementedError("sync_get_auth_token_and_url is not implemented") - + async def get_auth_token_and_url( self, api_base: Optional[str], diff --git a/litellm/llms/base_llm/image_generation/transformation.py b/litellm/llms/base_llm/image_generation/transformation.py index 151e2893d1..7f13e6f3b4 100644 --- a/litellm/llms/base_llm/image_generation/transformation.py +++ b/litellm/llms/base_llm/image_generation/transformation.py @@ -24,7 +24,7 @@ class BaseImageGenerationConfig(ABC): self, model: str ) -> List[OpenAIImageGenerationOptionalParams]: pass - + @abstractmethod def map_openai_params( self, @@ -35,7 +35,6 @@ class BaseImageGenerationConfig(ABC): ) -> dict: pass - def get_complete_url( self, api_base: Optional[str], diff --git a/litellm/llms/base_llm/interactions/transformation.py b/litellm/llms/base_llm/interactions/transformation.py index 4ceb3f5387..be400628fd 100644 --- a/litellm/llms/base_llm/interactions/transformation.py +++ b/litellm/llms/base_llm/interactions/transformation.py @@ -41,11 +41,11 @@ else: class BaseInteractionsAPIConfig(ABC): """ Base configuration class for Google Interactions API implementations. - + Per OpenAPI spec, the Interactions API supports two types of interactions: - Model interactions (with model parameter) - Agent interactions (with agent parameter) - + Implementations should override the abstract methods to provide provider-specific transformations for requests and responses. """ @@ -87,10 +87,7 @@ class BaseInteractionsAPIConfig(ABC): @abstractmethod def validate_environment( - self, - headers: dict, - model: str, - litellm_params: Optional[GenericLiteLLMParams] + self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams] ) -> dict: """ Validate and prepare environment settings including headers. @@ -108,16 +105,16 @@ class BaseInteractionsAPIConfig(ABC): ) -> str: """ Get the complete URL for the interaction request. - + Per OpenAPI spec: POST /{api_version}/interactions - + Args: api_base: Base URL for the API model: The model name (for model interactions) agent: The agent name (for agent interactions) litellm_params: LiteLLM parameters stream: Whether this is a streaming request - + Returns: The complete URL for the request """ @@ -137,11 +134,11 @@ class BaseInteractionsAPIConfig(ABC): ) -> Dict: """ Transform the input request into the provider's expected format. - + Per OpenAPI spec, the request body should be either: - CreateModelInteractionParams (with model) - CreateAgentInteractionParams (with agent) - + Args: model: The model name (for model interactions) agent: The agent name (for agent interactions) @@ -149,7 +146,7 @@ class BaseInteractionsAPIConfig(ABC): optional_params: Optional parameters for the request litellm_params: LiteLLM-specific parameters headers: Request headers - + Returns: The transformed request body as a dictionary """ @@ -164,7 +161,7 @@ class BaseInteractionsAPIConfig(ABC): ) -> InteractionsAPIResponse: """ Transform the raw HTTP response into an InteractionsAPIResponse. - + Per OpenAPI spec, the response is an Interaction object. """ pass @@ -178,7 +175,7 @@ class BaseInteractionsAPIConfig(ABC): ) -> InteractionsAPIStreamingResponse: """ Transform a parsed streaming response chunk into an InteractionsAPIStreamingResponse. - + Per OpenAPI spec, streaming uses SSE with various event types. """ pass @@ -186,7 +183,7 @@ class BaseInteractionsAPIConfig(ABC): # ========================================================= # GET INTERACTION TRANSFORMATION # ========================================================= - + @abstractmethod def transform_get_interaction_request( self, @@ -197,9 +194,9 @@ class BaseInteractionsAPIConfig(ABC): ) -> Tuple[str, Dict]: """ Transform the get interaction request into URL and query params. - + Per OpenAPI spec: GET /{api_version}/interactions/{interaction_id} - + Returns: Tuple of (URL, query_params) """ @@ -219,7 +216,7 @@ class BaseInteractionsAPIConfig(ABC): # ========================================================= # DELETE INTERACTION TRANSFORMATION # ========================================================= - + @abstractmethod def transform_delete_interaction_request( self, @@ -230,9 +227,9 @@ class BaseInteractionsAPIConfig(ABC): ) -> Tuple[str, Dict]: """ Transform the delete interaction request into URL and body. - + Per OpenAPI spec: DELETE /{api_version}/interactions/{interaction_id} - + Returns: Tuple of (URL, request_body) """ @@ -253,7 +250,7 @@ class BaseInteractionsAPIConfig(ABC): # ========================================================= # CANCEL INTERACTION TRANSFORMATION # ========================================================= - + @abstractmethod def transform_cancel_interaction_request( self, @@ -264,7 +261,7 @@ class BaseInteractionsAPIConfig(ABC): ) -> Tuple[str, Dict]: """ Transform the cancel interaction request into URL and body. - + Returns: Tuple of (URL, request_body) """ @@ -307,7 +304,7 @@ class BaseInteractionsAPIConfig(ABC): ) -> bool: """ Returns True if litellm should fake a stream for the given model. - + Override in subclasses if the provider doesn't support native streaming. """ return False diff --git a/litellm/llms/base_llm/managed_resources/base_managed_resource.py b/litellm/llms/base_llm/managed_resources/base_managed_resource.py index 3c8ce748ad..5422af7678 100644 --- a/litellm/llms/base_llm/managed_resources/base_managed_resource.py +++ b/litellm/llms/base_llm/managed_resources/base_managed_resource.py @@ -39,27 +39,27 @@ else: Router = Any # Generic type for resource objects -ResourceObjectType = TypeVar('ResourceObjectType') +ResourceObjectType = TypeVar("ResourceObjectType") class BaseManagedResource(ABC, Generic[ResourceObjectType]): """ Base class for managing resources with target_model_names support. - + This class provides common functionality for: - Storing unified resource IDs with model mappings - Retrieving resources by unified ID - Deleting resources across multiple models - Creating resources for multiple models - Filtering deployments based on model mappings - + Subclasses should implement: - resource_type: str property - table_name: str property - create_resource_for_model: method to create resource on a specific model - get_unified_resource_id_format: method to generate unified ID format """ - + def __init__( self, internal_usage_cache: InternalUsageCache, @@ -98,15 +98,15 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): ) -> str: """ Generate the format string for the unified resource ID. - + This should return a string that will be base64 encoded. Example for files: "litellm_proxy:application/json;unified_id,{uuid};target_model_names,{models};..." - + Args: resource_object: The resource object returned from the provider target_model_names_list: List of target model names - + Returns: Format string to be base64 encoded """ @@ -122,13 +122,13 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): ) -> ResourceObjectType: """ Create a resource for a specific model. - + Args: llm_router: LiteLLM router instance model: Model name to create resource for request_data: Request data for resource creation litellm_parent_otel_span: OpenTelemetry span for tracing - + Returns: Resource object from the provider """ @@ -149,7 +149,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): ) -> None: """ Store unified resource ID with model mappings in cache and database. - + Args: unified_resource_id: The unified resource ID (base64 encoded) resource_object: The resource object to store (can be None) @@ -161,7 +161,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): verbose_logger.info( f"Storing LiteLLM Managed {self.resource_type} with id={unified_resource_id} in cache" ) - + # Prepare cache data cache_data = { "unified_resource_id": unified_resource_id, @@ -171,11 +171,11 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): "created_by": user_api_key_dict.user_id, "updated_by": user_api_key_dict.user_id, } - + # Add additional fields if provided if additional_db_fields: cache_data.update(additional_db_fields) - + # Store in cache if resource_object is not None: await self.internal_usage_cache.async_set_cache( @@ -192,7 +192,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): "created_by": user_api_key_dict.user_id, "updated_by": user_api_key_dict.user_id, } - + # Add resource object if available if resource_object is not None: # Handle both dict and Pydantic models @@ -200,14 +200,14 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): db_data["resource_object"] = resource_object.model_dump_json() # type: ignore elif isinstance(resource_object, dict): db_data["resource_object"] = json.dumps(resource_object) - + # Extract storage metadata from hidden params if present hidden_params = getattr(resource_object, "_hidden_params", {}) or {} if "storage_backend" in hidden_params: db_data["storage_backend"] = hidden_params["storage_backend"] if "storage_url" in hidden_params: db_data["storage_url"] = hidden_params["storage_url"] - + # Add additional fields to database if additional_db_fields: db_data.update(additional_db_fields) @@ -215,7 +215,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): # Store in database table = getattr(self.prisma_client.db, self.table_name) result = await table.create(data=db_data) - + verbose_logger.debug( f"LiteLLM Managed {self.resource_type} with id={unified_resource_id} stored in db: {result}" ) @@ -227,11 +227,11 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): ) -> Optional[Dict[str, Any]]: """ Retrieve unified resource by ID from cache or database. - + Args: unified_resource_id: The unified resource ID to retrieve litellm_parent_otel_span: OpenTelemetry span for tracing - + Returns: Dictionary containing resource data or None if not found """ @@ -255,7 +255,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): if db_object: return db_object.model_dump() - + return None async def delete_unified_resource_id( @@ -265,11 +265,11 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): ) -> Optional[ResourceObjectType]: """ Delete unified resource from cache and database. - + Args: unified_resource_id: The unified resource ID to delete litellm_parent_otel_span: OpenTelemetry span for tracing - + Returns: The deleted resource object or None if not found """ @@ -278,22 +278,22 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): initial_value = await table.find_first( where={"unified_resource_id": unified_resource_id} ) - + if initial_value is None: raise Exception( f"LiteLLM Managed {self.resource_type} with id={unified_resource_id} not found" ) - + # Delete from cache await self.internal_usage_cache.async_set_cache( key=unified_resource_id, value=None, litellm_parent_otel_span=litellm_parent_otel_span, ) - + # Delete from database await table.delete(where={"unified_resource_id": unified_resource_id}) - + return initial_value.resource_object async def can_user_access_unified_resource_id( @@ -304,20 +304,20 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): ) -> bool: """ Check if user has access to the unified resource ID. - + Uses get_unified_resource_id() which checks cache first before hitting the database, avoiding direct DB queries in the critical request path. - + Args: unified_resource_id: The unified resource ID to check user_api_key_dict: User API key authentication details litellm_parent_otel_span: OpenTelemetry span for tracing - + Returns: True if user has access, False otherwise """ user_id = user_api_key_dict.user_id - + # Use cached method instead of direct DB query resource = await self.get_unified_resource_id( unified_resource_id, litellm_parent_otel_span @@ -325,7 +325,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): if resource: return resource.get("created_by") == user_id - + return False # ============================================================================ @@ -339,14 +339,14 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): ) -> Dict[str, Dict[str, str]]: """ Get model-specific resource IDs for a list of unified resource IDs. - + Args: resource_ids: List of unified resource IDs litellm_parent_otel_span: OpenTelemetry span for tracing - + Returns: Dictionary mapping unified_resource_id -> model_id -> provider_resource_id - + Example: { "unified_resource_id_1": { @@ -365,11 +365,11 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): if unified_resource_object: model_mappings = unified_resource_object.get("model_mappings", {}) - + # Handle both JSON string and dict if isinstance(model_mappings, str): model_mappings = json.loads(model_mappings) - + resource_id_mapping[resource_id] = model_mappings return resource_id_mapping @@ -387,19 +387,19 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): ) -> List[ResourceObjectType]: """ Create a resource for each model in the target list. - + Args: llm_router: LiteLLM router instance request_data: Request data for resource creation target_model_names_list: List of target model names litellm_parent_otel_span: OpenTelemetry span for tracing - + Returns: List of resource objects created for each model """ if llm_router is None: raise Exception("LLM Router not initialized. Ensure models added to proxy.") - + responses = [] for model in target_model_names_list: individual_response = await self.create_resource_for_model( @@ -418,11 +418,11 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): ) -> str: """ Generate a unified resource ID from multiple resource objects. - + Args: resource_objects: List of resource objects from different models target_model_names_list: List of target model names - + Returns: Base64 encoded unified resource ID """ @@ -431,12 +431,12 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): resource_object=resource_objects[0], target_model_names_list=target_model_names_list, ) - + # Convert to URL-safe base64 and strip padding base64_unified_id = ( base64.urlsafe_b64encode(unified_id_format.encode()).decode().rstrip("=") ) - + return base64_unified_id def extract_model_mappings_from_responses( @@ -445,10 +445,10 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): ) -> Dict[str, str]: """ Extract model mappings from resource objects. - + Args: resource_objects: List of resource objects from different models - + Returns: Dictionary mapping model_id -> provider_resource_id """ @@ -458,8 +458,10 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): # Get hidden params if available hidden_params = getattr(resource_object, "_hidden_params", {}) or {} model_resource_id_mapping = hidden_params.get("model_resource_id_mapping") - - if model_resource_id_mapping and isinstance(model_resource_id_mapping, dict): + + if model_resource_id_mapping and isinstance( + model_resource_id_mapping, dict + ): model_mappings.update(model_resource_id_mapping) return model_mappings @@ -478,17 +480,17 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): ) -> List[Dict]: """ Filter deployments based on model mappings for a resource. - + This is used by the router to select only deployments that have the resource available. - + Args: model: Model name healthy_deployments: List of healthy deployments request_kwargs: Request kwargs containing resource_id and mappings parent_otel_span: OpenTelemetry span for tracing resource_id_key: Key to use for resource ID in request_kwargs - + Returns: Filtered list of deployments """ @@ -500,7 +502,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): Optional[Dict[str, Dict[str, str]]], request_kwargs.get("model_resource_id_mapping"), ) - + allowed_model_ids = [] if resource_id and model_resource_id_mapping: model_id_dict = model_resource_id_mapping.get(resource_id, {}) @@ -522,7 +524,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): def get_unified_id_prefix(self) -> str: """ Get the prefix for unified IDs for this resource type. - + Returns: Prefix string (e.g., "litellm_proxy:") """ @@ -537,29 +539,29 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): ) -> Dict[str, Any]: """ List resources created by a user. - + Args: user_api_key_dict: User API key authentication details limit: Maximum number of resources to return after: Cursor for pagination additional_filters: Additional filters to apply - + Returns: Dictionary with list of resources and pagination info """ where_clause: Dict[str, Any] = {} - + # Filter by user who created the resource if user_api_key_dict.user_id: where_clause["created_by"] = user_api_key_dict.user_id - + if after: where_clause["id"] = {"gt": after} - + # Add additional filters if additional_filters: where_clause.update(additional_filters) - + # Fetch resources fetch_limit = limit or 20 table = getattr(self.prisma_client.db, self.table_name) @@ -568,7 +570,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): take=fetch_limit, order={"created_at": "desc"}, ) - + resource_objects: List[Any] = [] for resource in resources: try: @@ -580,13 +582,13 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): resource_data = resource.resource_object if isinstance(resource_data, str): resource_data = json.loads(resource_data) - + # Set unified ID if hasattr(resource_data, "id"): resource_data.id = resource.unified_resource_id elif isinstance(resource_data, dict): resource_data["id"] = resource.unified_resource_id - + resource_objects.append(resource_data) except Exception as e: @@ -595,7 +597,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): f"{resource.unified_resource_id}: {e}" ) continue - + return { "object": "list", "data": resource_objects, diff --git a/litellm/llms/base_llm/managed_resources/utils.py b/litellm/llms/base_llm/managed_resources/utils.py index 0d843b6d12..59f5ff0d84 100644 --- a/litellm/llms/base_llm/managed_resources/utils.py +++ b/litellm/llms/base_llm/managed_resources/utils.py @@ -16,21 +16,21 @@ def is_base64_encoded_unified_id( ) -> Union[str, Literal[False]]: """ Check if a resource ID is a base64 encoded unified ID. - + Args: resource_id: The resource ID to check prefix: The expected prefix for unified IDs - + Returns: Decoded string if valid unified ID, False otherwise """ # Ensure resource_id is a string if not isinstance(resource_id, str): return False - + # Add padding back if needed padded = resource_id + "=" * (-len(resource_id) % 4) - + # Decode from base64 try: decoded = base64.urlsafe_b64decode(padded).decode() @@ -47,13 +47,13 @@ def extract_target_model_names_from_unified_id( ) -> List[str]: """ Extract target model names from a unified resource ID. - + Args: unified_id: The unified resource ID (decoded or encoded) - + Returns: List of target model names - + Example: unified_id = "litellm_proxy:vector_store;unified_id,uuid;target_model_names,gpt-4,gemini-2.0" returns: ["gpt-4", "gemini-2.0"] @@ -62,18 +62,18 @@ def extract_target_model_names_from_unified_id( # Ensure unified_id is a string if not isinstance(unified_id, str): return [] - + # Decode if it's base64 encoded decoded_id = is_base64_encoded_unified_id(unified_id) if decoded_id: unified_id = decoded_id - + # Extract model names using regex match = re.search(r"target_model_names,([^;]+)", unified_id) if match: # Split on comma and strip whitespace from each model name return [model.strip() for model in match.group(1).split(",")] - + return [] except Exception: return [] @@ -84,13 +84,13 @@ def extract_resource_type_from_unified_id( ) -> Optional[str]: """ Extract resource type from a unified resource ID. - + Args: unified_id: The unified resource ID (decoded or encoded) - + Returns: Resource type string or None - + Example: unified_id = "litellm_proxy:vector_store;unified_id,uuid;..." returns: "vector_store" @@ -99,17 +99,17 @@ def extract_resource_type_from_unified_id( # Ensure unified_id is a string if not isinstance(unified_id, str): return None - + # Decode if it's base64 encoded decoded_id = is_base64_encoded_unified_id(unified_id) if decoded_id: unified_id = decoded_id - + # Extract resource type (comes after prefix and before first semicolon) match = re.search(r"litellm_proxy:([^;]+)", unified_id) if match: return match.group(1).strip() - + return None except Exception: return None @@ -120,13 +120,13 @@ def extract_unified_uuid_from_unified_id( ) -> Optional[str]: """ Extract the UUID from a unified resource ID. - + Args: unified_id: The unified resource ID (decoded or encoded) - + Returns: UUID string or None - + Example: unified_id = "litellm_proxy:vector_store;unified_id,abc-123;..." returns: "abc-123" @@ -135,17 +135,17 @@ def extract_unified_uuid_from_unified_id( # Ensure unified_id is a string if not isinstance(unified_id, str): return None - + # Decode if it's base64 encoded decoded_id = is_base64_encoded_unified_id(unified_id) if decoded_id: unified_id = decoded_id - + # Extract UUID match = re.search(r"unified_id,([^;]+)", unified_id) if match: return match.group(1).strip() - + return None except Exception: return None @@ -156,13 +156,13 @@ def extract_model_id_from_unified_id( ) -> Optional[str]: """ Extract model ID from a unified resource ID. - + Args: unified_id: The unified resource ID (decoded or encoded) - + Returns: Model ID string or None - + Example: unified_id = "litellm_proxy:vector_store;...;model_id,gpt-4-model-id;..." returns: "gpt-4-model-id" @@ -171,17 +171,17 @@ def extract_model_id_from_unified_id( # Ensure unified_id is a string if not isinstance(unified_id, str): return None - + # Decode if it's base64 encoded decoded_id = is_base64_encoded_unified_id(unified_id) if decoded_id: unified_id = decoded_id - + # Extract model ID match = re.search(r"model_id,([^;]+)", unified_id) if match: return match.group(1).strip() - + return None except Exception: return None @@ -192,13 +192,13 @@ def extract_provider_resource_id_from_unified_id( ) -> Optional[str]: """ Extract provider resource ID from a unified resource ID. - + Args: unified_id: The unified resource ID (decoded or encoded) - + Returns: Provider resource ID string or None - + Example: unified_id = "litellm_proxy:vector_store;...;resource_id,vs_abc123;..." returns: "vs_abc123" @@ -207,24 +207,24 @@ def extract_provider_resource_id_from_unified_id( # Ensure unified_id is a string if not isinstance(unified_id, str): return None - + # Decode if it's base64 encoded decoded_id = is_base64_encoded_unified_id(unified_id) if decoded_id: unified_id = decoded_id - + # Extract resource ID (try multiple patterns for different resource types) patterns = [ r"resource_id,([^;]+)", r"vector_store_id,([^;]+)", r"file_id,([^;]+)", ] - + for pattern in patterns: match = re.search(pattern, unified_id) if match: return match.group(1).strip() - + return None except Exception: return None @@ -240,7 +240,7 @@ def generate_unified_id_string( ) -> str: """ Generate a unified ID string (before base64 encoding). - + Args: resource_type: Type of resource (e.g., "vector_store", "file") unified_uuid: UUID for this unified resource @@ -248,10 +248,10 @@ def generate_unified_id_string( provider_resource_id: Resource ID from the provider model_id: Model ID from the router additional_fields: Additional fields to include in the ID - + Returns: Unified ID string (not yet base64 encoded) - + Example: generate_unified_id_string( resource_type="vector_store", @@ -270,53 +270,49 @@ def generate_unified_id_string( f"resource_id,{provider_resource_id}", f"model_id,{model_id}", ] - + # Add additional fields if provided if additional_fields: for key, value in additional_fields.items(): parts.append(f"{key},{value}") - + return ";".join(parts) def encode_unified_id(unified_id_string: str) -> str: """ Encode a unified ID string to base64. - + Args: unified_id_string: The unified ID string to encode - + Returns: Base64 encoded unified ID (URL-safe, padding stripped) """ - return ( - base64.urlsafe_b64encode(unified_id_string.encode()) - .decode() - .rstrip("=") - ) + return base64.urlsafe_b64encode(unified_id_string.encode()).decode().rstrip("=") def decode_unified_id(encoded_unified_id: str) -> Optional[str]: """ Decode a base64 encoded unified ID. - + Args: encoded_unified_id: The base64 encoded unified ID - + Returns: Decoded unified ID string or None if invalid """ try: # Add padding back if needed padded = encoded_unified_id + "=" * (-len(encoded_unified_id) % 4) - + # Decode from base64 decoded = base64.urlsafe_b64decode(padded).decode() - + # Verify it starts with the expected prefix if decoded.startswith("litellm_proxy:"): return decoded - + return None except Exception: return None @@ -327,13 +323,13 @@ def parse_unified_id( ) -> Optional[dict]: """ Parse a unified ID into its components. - + Args: unified_id: The unified ID (encoded or decoded) - + Returns: Dictionary with parsed components or None if invalid - + Example: { "resource_type": "vector_store", @@ -352,12 +348,16 @@ def parse_unified_id( decoded_id = unified_id else: return None - + return { "resource_type": extract_resource_type_from_unified_id(decoded_id), "unified_uuid": extract_unified_uuid_from_unified_id(decoded_id), - "target_model_names": extract_target_model_names_from_unified_id(decoded_id), - "provider_resource_id": extract_provider_resource_id_from_unified_id(decoded_id), + "target_model_names": extract_target_model_names_from_unified_id( + decoded_id + ), + "provider_resource_id": extract_provider_resource_id_from_unified_id( + decoded_id + ), "model_id": extract_model_id_from_unified_id(decoded_id), } except Exception: diff --git a/litellm/llms/base_llm/ocr/transformation.py b/litellm/llms/base_llm/ocr/transformation.py index 29929a2bf6..7d16c696db 100644 --- a/litellm/llms/base_llm/ocr/transformation.py +++ b/litellm/llms/base_llm/ocr/transformation.py @@ -23,6 +23,7 @@ DocumentType = Dict[str, str] class OCRPageDimensions(LiteLLMPydanticObjectBase): """Page dimensions from OCR response.""" + dpi: Optional[int] = None height: Optional[int] = None width: Optional[int] = None @@ -30,27 +31,30 @@ class OCRPageDimensions(LiteLLMPydanticObjectBase): class OCRPageImage(LiteLLMPydanticObjectBase): """Image extracted from OCR page.""" + image_base64: Optional[str] = None bbox: Optional[Dict[str, Any]] = None - + model_config = {"extra": "allow"} class OCRPage(LiteLLMPydanticObjectBase): """Single page from OCR response.""" + index: int markdown: str images: Optional[List[OCRPageImage]] = None dimensions: Optional[OCRPageDimensions] = None - + model_config = {"extra": "allow"} class OCRUsageInfo(LiteLLMPydanticObjectBase): """Usage information from OCR response.""" + pages_processed: Optional[int] = None doc_size_bytes: Optional[int] = None - + model_config = {"extra": "allow"} @@ -59,12 +63,13 @@ class OCRResponse(LiteLLMPydanticObjectBase): Standard OCR response format. Standardized to Mistral OCR format - other providers should transform to this format. """ + pages: List[OCRPage] model: str document_annotation: Optional[Any] = None usage_info: Optional[OCRUsageInfo] = None object: str = "ocr" - + model_config = {"extra": "allow"} # Define private attributes using PrivateAttr @@ -73,6 +78,7 @@ class OCRResponse(LiteLLMPydanticObjectBase): class OCRRequestData(LiteLLMPydanticObjectBase): """OCR request data structure.""" + data: Optional[Union[Dict, bytes]] = None files: Optional[Dict[str, Any]] = None @@ -142,21 +148,23 @@ class BaseOCRConfig: """ Transform OCR request to provider-specific format. Override in provider-specific implementations. - + Note: By the time this method is called, any file-type documents have already been converted to document_url/image_url format with base64 data URIs by the preprocessing in litellm/ocr/main.py. - + Args: model: Model name document: Document to process - always a dict with type="document_url" or type="image_url" optional_params: Optional parameters for the request headers: Request headers - + Returns: OCRRequestData with data and files fields """ - raise NotImplementedError("transform_ocr_request must be implemented by provider") + raise NotImplementedError( + "transform_ocr_request must be implemented by provider" + ) async def async_transform_ocr_request( self, @@ -170,15 +178,15 @@ class BaseOCRConfig: Async transform OCR request to provider-specific format. Optional method - providers can override if they need async transformations (e.g., Azure AI for URL-to-base64 conversion). - + Default implementation falls back to sync transform_ocr_request. - + Args: model: Model name document: Document to process (Mistral format dict, or file path, bytes, etc.) optional_params: Optional parameters for the request headers: Request headers - + Returns: OCRRequestData with data and files fields """ @@ -202,7 +210,9 @@ class BaseOCRConfig: Transform provider-specific OCR response to standard format. Override in provider-specific implementations. """ - raise NotImplementedError("transform_ocr_response must be implemented by provider") + raise NotImplementedError( + "transform_ocr_response must be implemented by provider" + ) async def async_transform_ocr_response( self, @@ -215,14 +225,14 @@ class BaseOCRConfig: Async transform provider-specific OCR response to standard format. Optional method - providers can override if they need async transformations (e.g., Azure Document Intelligence for async operation polling). - + Default implementation falls back to sync transform_ocr_response. - + Args: model: Model name raw_response: Raw HTTP response logging_obj: Logging object - + Returns: OCRResponse in standard format """ @@ -246,4 +256,3 @@ class BaseOCRConfig: message=error_message, headers=headers, ) - diff --git a/litellm/llms/base_llm/passthrough/transformation.py b/litellm/llms/base_llm/passthrough/transformation.py index f925e6819d..9d4396dce4 100644 --- a/litellm/llms/base_llm/passthrough/transformation.py +++ b/litellm/llms/base_llm/passthrough/transformation.py @@ -39,16 +39,14 @@ class BasePassthroughConfig(BaseLLMModelInfo): import httpx - base = base_target_url.rstrip('/') - endpoint = endpoint.lstrip('/') + base = base_target_url.rstrip("/") + endpoint = endpoint.lstrip("/") full_url = f"{base}/{endpoint}" url = httpx.URL(full_url) if request_query_params: - url = url.copy_with( - query=urlencode(request_query_params).encode("ascii") - ) + url = url.copy_with(query=urlencode(request_query_params).encode("ascii")) return url diff --git a/litellm/llms/base_llm/realtime/http_transformation.py b/litellm/llms/base_llm/realtime/http_transformation.py index 7aadd49ffd..712ec42380 100644 --- a/litellm/llms/base_llm/realtime/http_transformation.py +++ b/litellm/llms/base_llm/realtime/http_transformation.py @@ -54,7 +54,9 @@ class BaseRealtimeHTTPConfig(ABC): # ------------------------------------------------------------------ # @abstractmethod - def get_complete_url(self, api_base: Optional[str], model: str, api_version: Optional[str] = None) -> str: + def get_complete_url( + self, api_base: Optional[str], model: str, api_version: Optional[str] = None + ) -> str: """Return the full URL for POST /realtime/client_secrets.""" @abstractmethod diff --git a/litellm/llms/base_llm/rerank/transformation.py b/litellm/llms/base_llm/rerank/transformation.py index b22d85e82b..7874201f7f 100644 --- a/litellm/llms/base_llm/rerank/transformation.py +++ b/litellm/llms/base_llm/rerank/transformation.py @@ -52,8 +52,8 @@ class BaseRerankConfig(ABC): @abstractmethod def get_complete_url( - self, - api_base: Optional[str], + self, + api_base: Optional[str], model: str, optional_params: Optional[dict] = None, ) -> str: diff --git a/litellm/llms/base_llm/responses/transformation.py b/litellm/llms/base_llm/responses/transformation.py index 4cc3583ed8..f429930e00 100644 --- a/litellm/llms/base_llm/responses/transformation.py +++ b/litellm/llms/base_llm/responses/transformation.py @@ -221,11 +221,11 @@ class BaseResponsesAPIConfig(ABC): def supports_native_websocket(self) -> bool: """ Returns True if the provider has a native WebSocket endpoint for Responses API. - + Providers with native websocket support can connect directly to wss:// endpoints. Providers without native support will use the ManagedResponsesWebSocketHandler which makes HTTP streaming calls and forwards events over the websocket. - + Default: False (use managed websocket handler) """ return False diff --git a/litellm/llms/base_llm/search/__init__.py b/litellm/llms/base_llm/search/__init__.py index 5a46482ed4..f185b4e595 100644 --- a/litellm/llms/base_llm/search/__init__.py +++ b/litellm/llms/base_llm/search/__init__.py @@ -12,4 +12,3 @@ __all__ = [ "SearchResponse", "SearchResult", ] - diff --git a/litellm/llms/base_llm/search/transformation.py b/litellm/llms/base_llm/search/transformation.py index 14941911f1..1fbc5b670a 100644 --- a/litellm/llms/base_llm/search/transformation.py +++ b/litellm/llms/base_llm/search/transformation.py @@ -17,12 +17,13 @@ else: class SearchResult(LiteLLMPydanticObjectBase): """Single search result.""" + title: str url: str snippet: str date: Optional[str] = None last_updated: Optional[str] = None - + model_config = {"extra": "allow"} @@ -31,9 +32,10 @@ class SearchResponse(LiteLLMPydanticObjectBase): Standard Search response format. Standardized to Perplexity Search format - other providers should transform to this format. """ + results: List[SearchResult] object: str = "search" - + model_config = {"extra": "allow"} # Define private attributes using PrivateAttr @@ -48,7 +50,7 @@ class BaseSearchConfig: def __init__(self) -> None: pass - + @staticmethod def ui_friendly_name() -> str: """ @@ -56,12 +58,12 @@ class BaseSearchConfig: Override in provider-specific implementations. """ return "Unknown Search Provider" - + def get_http_method(self) -> Literal["GET", "POST"]: """ Get HTTP method for search requests. Override in provider-specific implementations if needed. - + Returns: HTTP method ('GET' or 'POST'). Default is 'POST'. """ @@ -72,7 +74,7 @@ class BaseSearchConfig: """ Get the set of Perplexity unified search parameters. These are the standard parameters that providers should transform from. - + Returns: Set of parameter names that are part of the unified spec """ @@ -105,7 +107,7 @@ class BaseSearchConfig: ) -> str: """ Get complete URL for Search endpoint. - + Args: api_base: Base URL for the API optional_params: Optional parameters for the request @@ -114,10 +116,10 @@ class BaseSearchConfig: the request body to construct query parameters in the URL. Can be a dict or list of dicts depending on provider. **kwargs: Additional keyword arguments - + Returns: Complete URL for the search endpoint - + Note: Override in provider-specific implementations. """ @@ -132,15 +134,17 @@ class BaseSearchConfig: """ Transform Search request to provider-specific format. Override in provider-specific implementations. - + Args: query: Search query (string or list of strings) optional_params: Optional parameters for the request - + Returns: Dict with request data """ - raise NotImplementedError("transform_search_request must be implemented by provider") + raise NotImplementedError( + "transform_search_request must be implemented by provider" + ) def transform_search_response( self, @@ -152,7 +156,9 @@ class BaseSearchConfig: Transform provider-specific Search response to standard format. Override in provider-specific implementations. """ - raise NotImplementedError("transform_search_response must be implemented by provider") + raise NotImplementedError( + "transform_search_response must be implemented by provider" + ) def get_error_class( self, @@ -166,4 +172,3 @@ class BaseSearchConfig: message=error_message, headers=headers, ) - diff --git a/litellm/llms/base_llm/skills/__init__.py b/litellm/llms/base_llm/skills/__init__.py index 3c523a0d12..e0b860ffb7 100644 --- a/litellm/llms/base_llm/skills/__init__.py +++ b/litellm/llms/base_llm/skills/__init__.py @@ -3,4 +3,3 @@ from .transformation import BaseSkillsAPIConfig __all__ = ["BaseSkillsAPIConfig"] - diff --git a/litellm/llms/base_llm/skills/transformation.py b/litellm/llms/base_llm/skills/transformation.py index 7c2ebc3529..017587c0b0 100644 --- a/litellm/llms/base_llm/skills/transformation.py +++ b/litellm/llms/base_llm/skills/transformation.py @@ -43,11 +43,11 @@ class BaseSkillsAPIConfig(ABC): ) -> dict: """ Validate and update headers with provider-specific requirements - + Args: headers: Base headers dictionary litellm_params: LiteLLM parameters - + Returns: Updated headers dictionary """ @@ -62,12 +62,12 @@ class BaseSkillsAPIConfig(ABC): ) -> str: """ Get the complete URL for the API request - + Args: api_base: Base API URL endpoint: API endpoint (e.g., 'skills', 'skills/{id}') skill_id: Optional skill ID for specific skill operations - + Returns: Complete URL """ @@ -84,12 +84,12 @@ class BaseSkillsAPIConfig(ABC): ) -> Dict: """ Transform create skill request to provider-specific format - + Args: create_request: Skill creation parameters litellm_params: LiteLLM parameters headers: Request headers - + Returns: Provider-specific request body """ @@ -103,11 +103,11 @@ class BaseSkillsAPIConfig(ABC): ) -> Skill: """ Transform provider response to Skill object - + Args: raw_response: Raw HTTP response logging_obj: Logging object - + Returns: Skill object """ @@ -122,12 +122,12 @@ class BaseSkillsAPIConfig(ABC): ) -> Tuple[str, Dict]: """ Transform list skills request parameters - + Args: list_params: List parameters (pagination, filters) litellm_params: LiteLLM parameters headers: Request headers - + Returns: Tuple of (url, query_params) """ @@ -141,11 +141,11 @@ class BaseSkillsAPIConfig(ABC): ) -> ListSkillsResponse: """ Transform provider response to ListSkillsResponse - + Args: raw_response: Raw HTTP response logging_obj: Logging object - + Returns: ListSkillsResponse object """ @@ -161,13 +161,13 @@ class BaseSkillsAPIConfig(ABC): ) -> Tuple[str, Dict]: """ Transform get skill request - + Args: skill_id: Skill ID api_base: Base API URL litellm_params: LiteLLM parameters headers: Request headers - + Returns: Tuple of (url, headers) """ @@ -181,11 +181,11 @@ class BaseSkillsAPIConfig(ABC): ) -> Skill: """ Transform provider response to Skill object - + Args: raw_response: Raw HTTP response logging_obj: Logging object - + Returns: Skill object """ @@ -201,13 +201,13 @@ class BaseSkillsAPIConfig(ABC): ) -> Tuple[str, Dict]: """ Transform delete skill request - + Args: skill_id: Skill ID api_base: Base API URL litellm_params: LiteLLM parameters headers: Request headers - + Returns: Tuple of (url, headers) """ @@ -221,11 +221,11 @@ class BaseSkillsAPIConfig(ABC): ) -> DeleteSkillResponse: """ Transform provider response to DeleteSkillResponse - + Args: raw_response: Raw HTTP response logging_obj: Logging object - + Returns: DeleteSkillResponse object """ @@ -243,4 +243,3 @@ class BaseSkillsAPIConfig(ABC): message=error_message, headers=headers, ) - diff --git a/litellm/llms/base_llm/text_to_speech/transformation.py b/litellm/llms/base_llm/text_to_speech/transformation.py index 31f581cec0..0e30ddae5f 100644 --- a/litellm/llms/base_llm/text_to_speech/transformation.py +++ b/litellm/llms/base_llm/text_to_speech/transformation.py @@ -24,10 +24,11 @@ else: class TextToSpeechRequestData(TypedDict, total=False): """ Structured return type for text-to-speech transformations. - + This ensures a consistent interface across all TTS providers. Providers should set ONE of: dict_body, ssml_body, or text_body. """ + dict_body: Dict[str, Any] # JSON request body (e.g., OpenAI TTS) ssml_body: str # SSML/XML string body (e.g., Azure AVA TTS) headers: Dict[str, str] # Provider-specific headers to merge with base headers @@ -116,7 +117,7 @@ class BaseTextToSpeechConfig(ABC): ) -> TextToSpeechRequestData: """ Transform request to provider-specific format. - + Returns: TextToSpeechRequestData: A structured dict containing: - body: The request body (JSON dict, XML string, or binary data) @@ -146,4 +147,3 @@ class BaseTextToSpeechConfig(ABC): message=error_message, headers=headers, ) - diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 935fd53c19..5fbf0a4b19 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -27,7 +27,6 @@ else: class BaseVectorStoreConfig: - def get_supported_openai_params( self, model: str ) -> List[VECTOR_STORE_OPENAI_PARAMS]: @@ -61,7 +60,6 @@ class BaseVectorStoreConfig: litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, ) -> Tuple[str, Dict]: - pass async def atransform_search_vector_store_request( diff --git a/litellm/llms/base_llm/vector_store_files/transformation.py b/litellm/llms/base_llm/vector_store_files/transformation.py index f751022faa..f13de56382 100644 --- a/litellm/llms/base_llm/vector_store_files/transformation.py +++ b/litellm/llms/base_llm/vector_store_files/transformation.py @@ -58,9 +58,9 @@ class BaseVectorStoreFilesConfig(ABC): ... @abstractmethod - def get_vector_store_file_endpoints_by_type(self) -> Dict[ - str, Tuple[Tuple[str, str], ...] - ]: + def get_vector_store_file_endpoints_by_type( + self, + ) -> Dict[str, Tuple[Tuple[str, str], ...]]: ... @abstractmethod diff --git a/litellm/llms/base_llm/videos/transformation.py b/litellm/llms/base_llm/videos/transformation.py index 1ad91a43df..2201a63363 100644 --- a/litellm/llms/base_llm/videos/transformation.py +++ b/litellm/llms/base_llm/videos/transformation.py @@ -145,13 +145,13 @@ class BaseVideoConfig(ABC): Async transform video content download response to bytes. Optional method - providers can override if they need async transformations (e.g., RunwayML for downloading video from CloudFront URL). - + Default implementation falls back to sync transform_video_content_response. - + Args: raw_response: Raw HTTP response logging_obj: Logging object - + Returns: Video content as bytes """ @@ -173,7 +173,7 @@ class BaseVideoConfig(ABC): ) -> Tuple[str, Dict]: """ Transform the video remix request into a URL and data - + Returns: Tuple[str, Dict]: (url, data) for the video remix request """ @@ -201,7 +201,7 @@ class BaseVideoConfig(ABC): ) -> Tuple[str, Dict]: """ Transform the video list request into a URL and params - + Returns: Tuple[str, Dict]: (url, params) for the video list request """ @@ -213,7 +213,7 @@ class BaseVideoConfig(ABC): raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, custom_llm_provider: Optional[str] = None, - ) -> Dict[str,str]: + ) -> Dict[str, str]: pass @abstractmethod @@ -226,7 +226,7 @@ class BaseVideoConfig(ABC): ) -> Tuple[str, Dict]: """ Transform the video delete request into a URL and data - + Returns: Tuple[str, Dict]: (url, data) for the video delete request """ @@ -250,7 +250,7 @@ class BaseVideoConfig(ABC): ) -> Tuple[str, Dict]: """ Transform the video retrieve request into a URL and data/params - + Returns: Tuple[str, Dict]: (url, params) for the video retrieve request """ diff --git a/litellm/llms/baseten/chat.py b/litellm/llms/baseten/chat.py index 05fc9961ac..1e49b34608 100644 --- a/litellm/llms/baseten/chat.py +++ b/litellm/llms/baseten/chat.py @@ -82,14 +82,16 @@ class BasetenConfig(OpenAIGPTConfig): optional_params[param] = value return optional_params - def _get_openai_compatible_provider_info(self, api_base: str, api_key: str) -> tuple: + def _get_openai_compatible_provider_info( + self, api_base: str, api_key: str + ) -> tuple: """ Get the OpenAI compatible provider info for Baseten """ # Default to Model API default_api_base = "https://inference.baseten.co/v1" default_api_key = api_key or "BASETEN_API_KEY" - + return default_api_base, default_api_key @staticmethod @@ -99,10 +101,11 @@ class BasetenConfig(OpenAIGPTConfig): """ # Remove 'baseten/' prefix if present model_id = model.replace("baseten/", "") - + # Check if it's an 8-digit alphanumeric code import re - return bool(re.match(r'^[a-zA-Z0-9]{8}$', model_id)) + + return bool(re.match(r"^[a-zA-Z0-9]{8}$", model_id)) @staticmethod def get_api_base_for_model(model: str) -> str: @@ -115,4 +118,4 @@ class BasetenConfig(OpenAIGPTConfig): return f"https://model-{model_id}.api.baseten.co/environments/production/sync/v1" else: # Use Model API - return "https://inference.baseten.co/v1" \ No newline at end of file + return "https://inference.baseten.co/v1" diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 5da118a8f5..697fccd268 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -747,7 +747,10 @@ class BaseAWSLLM: with open(web_identity_token_file, "r") as f: web_identity_token = f.read().strip() - irsa_sts_kwargs: dict = {"region_name": region, "verify": self._get_ssl_verify(ssl_verify)} + irsa_sts_kwargs: dict = { + "region_name": region, + "verify": self._get_ssl_verify(ssl_verify), + } if aws_sts_endpoint is not None: irsa_sts_kwargs["endpoint_url"] = aws_sts_endpoint @@ -814,7 +817,10 @@ class BaseAWSLLM: """Handle same-account role assumption for IRSA.""" import boto3 - irsa_sts_kwargs: dict = {"region_name": region, "verify": self._get_ssl_verify(ssl_verify)} + irsa_sts_kwargs: dict = { + "region_name": region, + "verify": self._get_ssl_verify(ssl_verify), + } if aws_sts_endpoint is not None: irsa_sts_kwargs["endpoint_url"] = aws_sts_endpoint @@ -889,7 +895,11 @@ class BaseAWSLLM: web_identity_token_file = os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE") irsa_role_arn = os.getenv("AWS_ROLE_ARN") - region = aws_region_name or os.getenv("AWS_REGION") or os.getenv("AWS_DEFAULT_REGION") + region = ( + aws_region_name + or os.getenv("AWS_REGION") + or os.getenv("AWS_DEFAULT_REGION") + ) # If we have IRSA environment variables and no explicit credentials, # we need to use the web identity token flow diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index 4a26bd4334..e0c7c08836 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -12,6 +12,7 @@ class BedrockBatchesHandler: E.g. Twelve Labs Embedding Async Invoke """ + @staticmethod def _handle_async_invoke_status( batch_id: str, aws_region_name: str, logging_obj=None, **kwargs diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index a9bc1b26c8..5d008038ca 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -29,7 +29,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): """ Config for Bedrock Batches - handles batch job creation and management for Bedrock """ - + def __init__(self): super().__init__() self.common_utils = CommonBatchFilesUtils() @@ -69,19 +69,15 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): Bedrock batch jobs are created via the model invocation job API. """ aws_region_name = self._get_aws_region_name(optional_params, model) - + # Bedrock model invocation job endpoint # Format: https://bedrock.{region}.amazonaws.com/model-invocation-job - bedrock_endpoint = f"https://bedrock.{aws_region_name}.amazonaws.com/model-invocation-job" - + bedrock_endpoint = ( + f"https://bedrock.{aws_region_name}.amazonaws.com/model-invocation-job" + ) + return bedrock_endpoint - - - - - - def transform_create_batch_request( self, model: str, @@ -91,7 +87,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): ) -> Dict[str, Any]: """ Transform the batch creation request to Bedrock format. - + Bedrock batch inference requires: - modelId: The Bedrock model ID - jobName: Unique name for the batch job @@ -103,19 +99,21 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): input_file_id = create_batch_data.get("input_file_id") if not input_file_id: raise ValueError("input_file_id is required for Bedrock batch creation") - + # Extract S3 information from file ID using common utility input_bucket, input_key = self.common_utils.parse_s3_uri(input_file_id) - + # Get output S3 configuration - output_bucket = litellm_params.get("s3_output_bucket_name") or os.getenv("AWS_S3_OUTPUT_BUCKET_NAME") + output_bucket = litellm_params.get("s3_output_bucket_name") or os.getenv( + "AWS_S3_OUTPUT_BUCKET_NAME" + ) if not output_bucket: # Use same bucket as input if no output bucket specified output_bucket = input_bucket - + # Get IAM role ARN role_arn = ( - litellm_params.get("aws_batch_role_arn") + litellm_params.get("aws_batch_role_arn") or optional_params.get("aws_batch_role_arn") or os.getenv("AWS_BATCH_ROLE_ARN") ) @@ -125,47 +123,47 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): "Set 'aws_batch_role_arn' in litellm_params or AWS_BATCH_ROLE_ARN env var" ) - if not model: - raise ValueError("Could not determine Bedrock model ID. Please pass `model` in your request body.") - + raise ValueError( + "Could not determine Bedrock model ID. Please pass `model` in your request body." + ) + # Generate job name with the correct model ID using common utility job_name = self.common_utils.generate_unique_job_name(model, prefix="litellm") output_key = f"litellm-batch-outputs/{job_name}/" - + # Build input data config input_data_config: BedrockInputDataConfig = { "s3InputDataConfig": BedrockS3InputDataConfig( s3Uri=f"s3://{input_bucket}/{input_key}" ) } - + # Build output data config s3_output_config: BedrockS3OutputDataConfig = BedrockS3OutputDataConfig( s3Uri=f"s3://{output_bucket}/{output_key}" ) - + # Add optional KMS encryption key ID if provided - s3_encryption_key_id = ( - litellm_params.get("s3_encryption_key_id") - or get_secret_str("AWS_S3_ENCRYPTION_KEY_ID") - ) + s3_encryption_key_id = litellm_params.get( + "s3_encryption_key_id" + ) or get_secret_str("AWS_S3_ENCRYPTION_KEY_ID") if s3_encryption_key_id: s3_output_config["s3EncryptionKeyId"] = s3_encryption_key_id - + output_data_config: BedrockOutputDataConfig = { "s3OutputDataConfig": s3_output_config } - + # Create Bedrock batch request with proper typing bedrock_request: BedrockCreateBatchRequest = { "modelId": model, "jobName": job_name, "inputDataConfig": input_data_config, "outputDataConfig": output_data_config, - "roleArn": role_arn + "roleArn": role_arn, } - + # Add optional parameters if provided completion_window = create_batch_data.get("completion_window") if completion_window: @@ -182,15 +180,15 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): data=bedrock_request, endpoint_url=endpoint_url, optional_params=optional_params, - method="POST" + method="POST", ) - + # Return a pre-signed request format that the HTTP handler can use return { "method": "POST", "url": endpoint_url, "headers": signed_headers, - "data": signed_data.decode('utf-8') + "data": signed_data.decode("utf-8"), } def transform_create_batch_response( @@ -207,17 +205,17 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): response_data: BedrockCreateBatchResponse = raw_response.json() except Exception as e: raise ValueError(f"Failed to parse Bedrock batch response: {e}") - + # Extract information from typed Bedrock response job_arn = response_data.get("jobArn", "") status_str: str = str(response_data.get("status", "Submitted")) - + # Map Bedrock status to OpenAI-compatible status status_mapping: Dict[str, str] = { "Submitted": "validating", "Validating": "validating", "Scheduled": "in_progress", - "InProgress": "in_progress", + "InProgress": "in_progress", "PartiallyCompleted": "completed", "Completed": "completed", "Failed": "failed", @@ -225,12 +223,24 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): "Stopped": "cancelled", "Expired": "expired", } - - openai_status = cast(Literal["validating", "failed", "in_progress", "finalizing", "completed", "expired", "cancelling", "cancelled"], status_mapping.get(status_str, "validating")) - + + openai_status = cast( + Literal[ + "validating", + "failed", + "in_progress", + "finalizing", + "completed", + "expired", + "cancelling", + "cancelled", + ], + status_mapping.get(status_str, "validating"), + ) + # Get original request data from litellm_params if available original_request = litellm_params.get("original_batch_request", {}) - + # Create LiteLLM batch object return LiteLLMBatch( id=job_arn, # Use ARN as the batch ID @@ -263,12 +273,12 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): ) -> Dict[str, Any]: """ Transform batch retrieval request for Bedrock. - + Args: batch_id: Bedrock job ARN optional_params: Optional parameters litellm_params: LiteLLM parameters - + Returns: Transformed request data for Bedrock GetModelInvocationJob API """ @@ -276,66 +286,113 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): # The GetModelInvocationJob API expects the full ARN as the identifier if not batch_id.startswith("arn:aws:bedrock:"): raise ValueError(f"Invalid batch_id format. Expected ARN, got: {batch_id}") - + # Extract the job identifier from the ARN - use the full ARN path part # ARN format: arn:aws:bedrock:region:account:model-invocation-job/job-name arn_parts = batch_id.split(":") if len(arn_parts) < 6: raise ValueError(f"Invalid ARN format: {batch_id}") - + region = arn_parts[3] # arn_parts[5] contains "model-invocation-job/{jobId}" - + # Build the endpoint URL for GetModelInvocationJob # AWS API format: GET /model-invocation-job/{jobIdentifier} # Use the FULL ARN as jobIdentifier and URL-encode it (includes ':' and '/') import urllib.parse as _ul + encoded_arn = _ul.quote(batch_id, safe="") - endpoint_url = f"https://bedrock.{region}.amazonaws.com/model-invocation-job/{encoded_arn}" - + endpoint_url = ( + f"https://bedrock.{region}.amazonaws.com/model-invocation-job/{encoded_arn}" + ) + # Use common utility for AWS signing signed_headers, _ = self.common_utils.sign_aws_request( service_name="bedrock", data={}, # GET request has no body endpoint_url=endpoint_url, optional_params=optional_params, - method="GET" + method="GET", ) - + # Return pre-signed request format return { "method": "GET", "url": endpoint_url, "headers": signed_headers, - "data": None + "data": None, } def _parse_timestamps_and_status(self, response_data, status_str: str): """Helper to parse timestamps based on status.""" import datetime + def parse_timestamp(ts_str: Optional[str]) -> Optional[int]: if not ts_str: return None try: - dt = datetime.datetime.fromisoformat(ts_str.replace('Z', '+00:00')) + dt = datetime.datetime.fromisoformat(ts_str.replace("Z", "+00:00")) return int(dt.timestamp()) except Exception: return None - - created_at = parse_timestamp(str(response_data.get("submitTime")) if response_data.get("submitTime") is not None else None) + + created_at = parse_timestamp( + str(response_data.get("submitTime")) + if response_data.get("submitTime") is not None + else None + ) in_progress_states = {"InProgress", "Validating", "Scheduled"} in_progress_at = ( - parse_timestamp(str(response_data.get("lastModifiedTime")) if response_data.get("lastModifiedTime") is not None else None) + parse_timestamp( + str(response_data.get("lastModifiedTime")) + if response_data.get("lastModifiedTime") is not None + else None + ) if status_str in in_progress_states else None ) - completed_at = parse_timestamp(str(response_data.get("endTime")) if response_data.get("endTime") is not None else None) if status_str in {"Completed", "PartiallyCompleted"} else None - failed_at = parse_timestamp(str(response_data.get("endTime")) if response_data.get("endTime") is not None else None) if status_str == "Failed" else None - cancelled_at = parse_timestamp(str(response_data.get("endTime")) if response_data.get("endTime") is not None else None) if status_str == "Stopped" else None - expires_at = parse_timestamp(str(response_data.get("jobExpirationTime")) if response_data.get("jobExpirationTime") is not None else None) - - return created_at, in_progress_at, completed_at, failed_at, cancelled_at, expires_at - + completed_at = ( + parse_timestamp( + str(response_data.get("endTime")) + if response_data.get("endTime") is not None + else None + ) + if status_str in {"Completed", "PartiallyCompleted"} + else None + ) + failed_at = ( + parse_timestamp( + str(response_data.get("endTime")) + if response_data.get("endTime") is not None + else None + ) + if status_str == "Failed" + else None + ) + cancelled_at = ( + parse_timestamp( + str(response_data.get("endTime")) + if response_data.get("endTime") is not None + else None + ) + if status_str == "Stopped" + else None + ) + expires_at = parse_timestamp( + str(response_data.get("jobExpirationTime")) + if response_data.get("jobExpirationTime") is not None + else None + ) + + return ( + created_at, + in_progress_at, + completed_at, + failed_at, + cancelled_at, + expires_at, + ) + def _extract_file_configs(self, response_data): """Helper to extract input and output file configurations.""" # Extract input file ID @@ -345,7 +402,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): s3_input_config = input_data_config.get("s3InputDataConfig", {}) if isinstance(s3_input_config, dict): input_file_id = s3_input_config.get("s3Uri", "") - + # Extract output file ID output_file_id = None output_data_config = response_data.get("outputDataConfig", {}) @@ -353,9 +410,9 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): s3_output_config = output_data_config.get("s3OutputDataConfig", {}) if isinstance(s3_output_config, dict): output_file_id = s3_output_config.get("s3Uri", "") - + return input_file_id, output_file_id - + def _extract_errors_and_metadata(self, response_data, raw_response): """Helper to extract errors and enriched metadata.""" # Extract errors @@ -364,11 +421,12 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): if message: from openai.types.batch import Errors from openai.types.batch_error import BatchError + errors = Errors( data=[BatchError(message=message, code=str(raw_response.status_code))], - object="list" + object="list", ) - + # Enrich metadata with useful Bedrock fields enriched_metadata_raw: Dict[str, Any] = { "jobName": response_data.get("jobName"), @@ -379,6 +437,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): "vpcConfig": response_data.get("vpcConfig"), } import json as _json + enriched_metadata: Dict[str, str] = {} for _k, _v in enriched_metadata_raw.items(): if _v is None: @@ -390,7 +449,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): enriched_metadata[_k] = str(_v) else: enriched_metadata[_k] = str(_v) - + return errors, enriched_metadata def transform_retrieve_batch_response( @@ -404,31 +463,60 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): Transform Bedrock batch retrieval response to LiteLLM format. """ from litellm.types.llms.bedrock import BedrockGetBatchResponse + try: response_data: BedrockGetBatchResponse = raw_response.json() except Exception as e: raise ValueError(f"Failed to parse Bedrock batch response: {e}") - + job_arn = response_data.get("jobArn", "") status_str: str = str(response_data.get("status", "Submitted")) - + # Map Bedrock status to OpenAI-compatible status status_mapping: Dict[str, str] = { - "Submitted": "validating", "Validating": "validating", "Scheduled": "in_progress", - "InProgress": "in_progress", "PartiallyCompleted": "completed", "Completed": "completed", - "Failed": "failed", "Stopping": "cancelling", "Stopped": "cancelled", "Expired": "expired" + "Submitted": "validating", + "Validating": "validating", + "Scheduled": "in_progress", + "InProgress": "in_progress", + "PartiallyCompleted": "completed", + "Completed": "completed", + "Failed": "failed", + "Stopping": "cancelling", + "Stopped": "cancelled", + "Expired": "expired", } - openai_status = cast(Literal["validating", "failed", "in_progress", "finalizing", "completed", "expired", "cancelling", "cancelled"], status_mapping.get(status_str, "validating")) - + openai_status = cast( + Literal[ + "validating", + "failed", + "in_progress", + "finalizing", + "completed", + "expired", + "cancelling", + "cancelled", + ], + status_mapping.get(status_str, "validating"), + ) + # Parse timestamps - created_at, in_progress_at, completed_at, failed_at, cancelled_at, expires_at = self._parse_timestamps_and_status(response_data, status_str) - + ( + created_at, + in_progress_at, + completed_at, + failed_at, + cancelled_at, + expires_at, + ) = self._parse_timestamps_and_status(response_data, status_str) + # Extract file configurations input_file_id, output_file_id = self._extract_file_configs(response_data) - + # Extract errors and metadata - errors, enriched_metadata = self._extract_errors_and_metadata(response_data, raw_response) - + errors, enriched_metadata = self._extract_errors_and_metadata( + response_data, raw_response + ) + return LiteLLMBatch( id=job_arn, object="batch", @@ -459,5 +547,3 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): Get Bedrock-specific error class using common utility. """ return self.common_utils.get_error_class(error_message, status_code, headers) - - diff --git a/litellm/llms/bedrock/chat/agentcore/__init__.py b/litellm/llms/bedrock/chat/agentcore/__init__.py index a2f1387620..2c83261fc9 100644 --- a/litellm/llms/bedrock/chat/agentcore/__init__.py +++ b/litellm/llms/bedrock/chat/agentcore/__init__.py @@ -1,4 +1,3 @@ from .transformation import AmazonAgentCoreConfig __all__ = ["AmazonAgentCoreConfig"] - diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py index 560fadad7c..d6eb5a734c 100644 --- a/litellm/llms/bedrock/chat/agentcore/transformation.py +++ b/litellm/llms/bedrock/chat/agentcore/transformation.py @@ -26,7 +26,15 @@ from litellm.types.llms.bedrock_agentcore import ( AgentCoreUsage, ) from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import Choices, Delta, Message, ModelResponse, ModelResponseStream, StreamingChoices, Usage +from litellm.types.utils import ( + Choices, + Delta, + Message, + ModelResponse, + ModelResponseStream, + StreamingChoices, + Usage, +) if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -364,9 +372,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): ) # Strategy 2: {"response": [{"text": "..."}]} - Strands agent content blocks - if "response" in response_json and isinstance( - response_json["response"], list - ): + if "response" in response_json and isinstance(response_json["response"], list): content = self._extract_content_from_message( {"content": response_json["response"]} # type: ignore ) @@ -498,11 +504,11 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): buffer += text_chunk # Process complete lines - while '\n' in buffer: - line, buffer = buffer.split('\n', 1) + while "\n" in buffer: + line, buffer = buffer.split("\n", 1) line = line.strip() - if not line or not line.startswith('data:'): + if not line or not line.startswith("data:"): continue json_str = line[5:].strip() @@ -556,11 +562,15 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): ) ] usage_data: AgentCoreUsage = metadata["usage"] # type: ignore - setattr(chunk, "usage", Usage( - prompt_tokens=usage_data.get("inputTokens", 0), - completion_tokens=usage_data.get("outputTokens", 0), - total_tokens=usage_data.get("totalTokens", 0), - )) + setattr( + chunk, + "usage", + Usage( + prompt_tokens=usage_data.get("inputTokens", 0), + completion_tokens=usage_data.get("outputTokens", 0), + total_tokens=usage_data.get("totalTokens", 0), + ), + ) yield chunk # Process final message @@ -710,11 +720,11 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): buffer += text_chunk # Process complete lines - while '\n' in buffer: - line, buffer = buffer.split('\n', 1) + while "\n" in buffer: + line, buffer = buffer.split("\n", 1) line = line.strip() - if not line or not line.startswith('data:'): + if not line or not line.startswith("data:"): continue json_str = line[5:].strip() @@ -768,11 +778,15 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): ) ] usage_data: AgentCoreUsage = metadata["usage"] # type: ignore - setattr(chunk, "usage", Usage( - prompt_tokens=usage_data.get("inputTokens", 0), - completion_tokens=usage_data.get("outputTokens", 0), - total_tokens=usage_data.get("totalTokens", 0), - )) + setattr( + chunk, + "usage", + Usage( + prompt_tokens=usage_data.get("inputTokens", 0), + completion_tokens=usage_data.get("outputTokens", 0), + total_tokens=usage_data.get("totalTokens", 0), + ), + ) yield chunk # Process final message @@ -863,7 +877,9 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): ) parsed = self._parse_json_response(response_json) - async def _json_as_async_stream() -> AsyncGenerator[ModelResponseStream, None]: + async def _json_as_async_stream() -> AsyncGenerator[ + ModelResponseStream, None + ]: # Content chunk content_chunk = ModelResponseStream( id=f"chatcmpl-{uuid.uuid4()}", diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 26986aab58..ef46ae5c18 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -70,7 +70,9 @@ def make_sync_call( ) else: decoder = AWSEventStreamDecoder(model=model, json_mode=json_mode) - completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=stream_chunk_size)) + completion_stream = decoder.iter_bytes( + response.iter_bytes(chunk_size=stream_chunk_size) + ) # LOGGING logging_obj.post_call( @@ -124,7 +126,7 @@ class BedrockConverseLLM(BaseAWSLLM): endpoint_url=api_base, data=data, headers=headers, - api_key=api_key + api_key=api_key, ) ## LOGGING @@ -184,7 +186,7 @@ class BedrockConverseLLM(BaseAWSLLM): headers=headers, ) data = json.dumps(request_data) - + prepped = self.get_request_headers( credentials=credentials, aws_region_name=litellm_params.get("aws_region_name") or "us-west-2", @@ -192,7 +194,7 @@ class BedrockConverseLLM(BaseAWSLLM): endpoint_url=api_base, data=data, headers=headers, - api_key=api_key + api_key=api_key, ) ## LOGGING @@ -278,7 +280,7 @@ class BedrockConverseLLM(BaseAWSLLM): _stripped = _model_for_id for rp in ["bedrock/converse/", "bedrock/", "converse/"]: if _stripped.startswith(rp): - _stripped = _stripped[len(rp):] + _stripped = _stripped[len(rp) :] break # Strip embedded region prefix (e.g. "bedrock/us-east-1/model" -> "model") # and capture it so it can be used as aws_region_name below. @@ -294,7 +296,10 @@ class BedrockConverseLLM(BaseAWSLLM): break modelId = self.encode_model_id(model_id=_model_for_id) # Inject region extracted from model path so _get_aws_region_name picks it up - if _region_from_model is not None and "aws_region_name" not in optional_params: + if ( + _region_from_model is not None + and "aws_region_name" not in optional_params + ): optional_params["aws_region_name"] = _region_from_model fake_stream = litellm.AmazonConverseConfig().should_fake_stream( @@ -304,7 +309,6 @@ class BedrockConverseLLM(BaseAWSLLM): custom_llm_provider="bedrock", ) - ### SET REGION NAME ### aws_region_name = self._get_aws_region_name( optional_params=optional_params, @@ -362,7 +366,7 @@ class BedrockConverseLLM(BaseAWSLLM): headers = {"Content-Type": "application/json"} if extra_headers is not None: headers = {"Content-Type": "application/json", **extra_headers} - + # Filter beta headers in HTTP headers before making the request headers = update_headers_with_filtered_beta( headers=headers, provider="bedrock_converse" @@ -408,7 +412,7 @@ class BedrockConverseLLM(BaseAWSLLM): timeout=timeout, client=client, credentials=credentials, - api_key=api_key + api_key=api_key, ) # type: ignore ## TRANSFORMATION ## @@ -421,7 +425,7 @@ class BedrockConverseLLM(BaseAWSLLM): headers=extra_headers, ) data = json.dumps(_data) - + prepped = self.get_request_headers( credentials=credentials, aws_region_name=aws_region_name, @@ -429,7 +433,7 @@ class BedrockConverseLLM(BaseAWSLLM): endpoint_url=proxy_endpoint_url, data=data, headers=headers, - api_key=api_key + api_key=api_key, ) ## LOGGING diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 7dd32b99bc..229457a73b 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -51,6 +51,7 @@ from litellm.types.llms.openai import ( ) from litellm.types.utils import ( ChatCompletionMessageToolCall, + CompletionTokensDetailsWrapper, Function, Message, ModelResponse, @@ -63,6 +64,7 @@ from litellm.utils import ( has_tool_call_blocks, last_assistant_with_tool_calls_has_no_thinking_blocks, supports_reasoning, + token_counter, ) from ..common_utils import ( @@ -348,7 +350,9 @@ class AmazonConverseConfig(BaseConfig): # Check if the model is a Nova 2 model (matches nova-2-lite, nova-2-pro, etc.) # Also check for nova-2/ spec prefix for imported models - return model_without_region.startswith("amazon.nova-2-") or model_without_region.startswith("nova-2/") + return model_without_region.startswith( + "amazon.nova-2-" + ) or model_without_region.startswith("nova-2/") def _map_web_search_options( self, web_search_options: dict, model: str @@ -762,8 +766,7 @@ class AmazonConverseConfig(BaseConfig): def _supports_native_structured_outputs(model: str) -> bool: """Check if the Bedrock model supports native structured outputs (outputConfig.textFormat).""" return any( - substring in model - for substring in BEDROCK_NATIVE_STRUCTURED_OUTPUT_MODELS + substring in model for substring in BEDROCK_NATIVE_STRUCTURED_OUTPUT_MODELS ) @staticmethod @@ -917,9 +920,7 @@ class AmazonConverseConfig(BaseConfig): if param == "parallel_tool_calls": disable_parallel = not value optional_params["_parallel_tool_use_config"] = { - "tool_choice": { - "disable_parallel_tool_use": disable_parallel - } + "tool_choice": {"disable_parallel_tool_use": disable_parallel} } if param == "thinking": optional_params["thinking"] = value @@ -1208,7 +1209,12 @@ class AmazonConverseConfig(BaseConfig): if request_metadata is not None: self._validate_request_metadata(request_metadata) - output_config: Optional[OutputConfigBlock] = inference_params.pop("outputConfig", None) + output_config: Optional[OutputConfigBlock] = inference_params.pop( + "outputConfig", None + ) + inference_params.pop( + "output_config", None + ) # Bedrock Converse doesn't support it # keep supported params in 'inference_params', and set all model-specific params in 'additional_request_params' additional_request_params = { @@ -1219,10 +1225,16 @@ class AmazonConverseConfig(BaseConfig): } # Handle parallel_tool_calls configuration - parallel_tool_use_config = additional_request_params.pop("_parallel_tool_use_config", None) + parallel_tool_use_config = additional_request_params.pop( + "_parallel_tool_use_config", None + ) if parallel_tool_use_config is not None and is_claude_4_5_on_bedrock(model): for key, value in parallel_tool_use_config.items(): - if key in additional_request_params and isinstance(additional_request_params[key], dict) and isinstance(value, dict): + if ( + key in additional_request_params + and isinstance(additional_request_params[key], dict) + and isinstance(value, dict) + ): additional_request_params[key].update(value) else: additional_request_params[key] = value @@ -1304,7 +1316,16 @@ class AmazonConverseConfig(BaseConfig): # "computer-use-2025-01-24" for Claude Sonnet 4.5, Haiku 4.5, Opus 4.1, Sonnet 4, Opus 4, and Sonnet 3.7 # "computer-use-2024-10-22" for older models model_lower = model.lower() - if "opus-4.6" in model_lower or "opus_4.6" in model_lower or "opus-4-6" in model_lower or "opus_4_6" in model_lower or "sonnet-4.6" in model_lower or "sonnet_4.6" in model_lower or "sonnet-4-6" in model_lower or "sonnet_4_6" in model_lower: + if ( + "opus-4.6" in model_lower + or "opus_4.6" in model_lower + or "opus-4-6" in model_lower + or "opus_4_6" in model_lower + or "sonnet-4.6" in model_lower + or "sonnet_4.6" in model_lower + or "sonnet-4-6" in model_lower + or "sonnet_4_6" in model_lower + ): computer_use_header = "computer-use-2025-11-24" elif ( "opus-4.5" in model_lower @@ -1623,7 +1644,11 @@ class AmazonConverseConfig(BaseConfig): thinking_blocks_list.append(_redacted_block) return thinking_blocks_list - def _transform_usage(self, usage: ConverseTokenUsageBlock) -> Usage: + def _transform_usage( + self, + usage: ConverseTokenUsageBlock, + reasoning_content: Optional[str] = None, + ) -> Usage: input_tokens = usage["inputTokens"] output_tokens = usage["outputTokens"] total_tokens = usage["totalTokens"] @@ -1640,6 +1665,19 @@ class AmazonConverseConfig(BaseConfig): prompt_tokens_details = PromptTokensDetailsWrapper( cached_tokens=cache_read_input_tokens ) + reasoning_tokens = ( + token_counter(text=reasoning_content, count_response_tokens=True) + if reasoning_content + else 0 + ) + completion_tokens_details = CompletionTokensDetailsWrapper( + reasoning_tokens=reasoning_tokens, + text_tokens=( + output_tokens - reasoning_tokens + if reasoning_tokens > 0 + else output_tokens + ), + ) openai_usage = Usage( prompt_tokens=input_tokens, completion_tokens=output_tokens, @@ -1647,6 +1685,7 @@ class AmazonConverseConfig(BaseConfig): prompt_tokens_details=prompt_tokens_details, cache_creation_input_tokens=cache_creation_input_tokens, cache_read_input_tokens=cache_read_input_tokens, + completion_tokens_details=completion_tokens_details, ) return openai_usage @@ -1709,7 +1748,9 @@ class AmazonConverseConfig(BaseConfig): return message, returned_finish_reason - def _translate_message_content(self, content_blocks: List[ContentBlock]) -> Tuple[ + def _translate_message_content( + self, content_blocks: List[ContentBlock] + ) -> Tuple[ str, List[ChatCompletionToolCallChunk], Optional[List[BedrockConverseReasoningContentBlock]], @@ -1726,9 +1767,9 @@ class AmazonConverseConfig(BaseConfig): """ content_str = "" tools: List[ChatCompletionToolCallChunk] = [] - reasoningContentBlocks: Optional[List[BedrockConverseReasoningContentBlock]] = ( - None - ) + reasoningContentBlocks: Optional[ + List[BedrockConverseReasoningContentBlock] + ] = None citationsContentBlocks: Optional[List[CitationsContentBlock]] = None for idx, content in enumerate(content_blocks): """ @@ -1841,9 +1882,7 @@ class AmazonConverseConfig(BaseConfig): verbose_logger.debug( "Processing JSON tool call response for response_format" ) - json_mode_content_str: Optional[str] = tools[0]["function"].get( - "arguments" - ) + json_mode_content_str: Optional[str] = tools[0]["function"].get("arguments") if json_mode_content_str is not None: json_mode_content_str = AmazonConverseConfig._unwrap_bedrock_properties( json_mode_content_str @@ -1941,9 +1980,9 @@ class AmazonConverseConfig(BaseConfig): chat_completion_message: ChatCompletionResponseMessage = {"role": "assistant"} content_str = "" tools: List[ChatCompletionToolCallChunk] = [] - reasoningContentBlocks: Optional[List[BedrockConverseReasoningContentBlock]] = ( - None - ) + reasoningContentBlocks: Optional[ + List[BedrockConverseReasoningContentBlock] + ] = None citationsContentBlocks: Optional[List[CitationsContentBlock]] = None if message is not None: @@ -1962,17 +2001,17 @@ class AmazonConverseConfig(BaseConfig): provider_specific_fields["citationsContent"] = citationsContentBlocks if provider_specific_fields: - chat_completion_message["provider_specific_fields"] = ( - provider_specific_fields - ) + chat_completion_message[ + "provider_specific_fields" + ] = provider_specific_fields if reasoningContentBlocks is not None: - chat_completion_message["reasoning_content"] = ( - self._transform_reasoning_content(reasoningContentBlocks) - ) - chat_completion_message["thinking_blocks"] = ( - self._transform_thinking_blocks(reasoningContentBlocks) - ) + chat_completion_message[ + "reasoning_content" + ] = self._transform_reasoning_content(reasoningContentBlocks) + chat_completion_message[ + "thinking_blocks" + ] = self._transform_thinking_blocks(reasoningContentBlocks) chat_completion_message["content"] = content_str filtered_tools = self._filter_json_mode_tools( json_mode=json_mode, @@ -1983,7 +2022,10 @@ class AmazonConverseConfig(BaseConfig): chat_completion_message["tool_calls"] = filtered_tools ## CALCULATING USAGE - bedrock returns usage in the headers - usage = self._transform_usage(completion_response["usage"]) + usage = self._transform_usage( + completion_response["usage"], + reasoning_content=chat_completion_message.get("reasoning_content"), + ) ## HANDLE TOOL CALLS _message = Message(**chat_completion_message) diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 9b06e19820..1077731779 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -407,9 +407,9 @@ class BedrockLLM(BaseAWSLLM): # Claude 3+ indicators (all use Messages API) messages_api_indicators = [ - "claude-3", # Claude 3.x models - "claude-opus-4", # Claude Opus 4 - "claude-sonnet-4", # Claude Sonnet 4 + "claude-3", # Claude 3.x models + "claude-opus-4", # Claude Opus 4 + "claude-sonnet-4", # Claude Sonnet 4 "claude-haiku-4", # Claude Haiku 4 ] diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_mistral_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_mistral_transformation.py index 58dfa17a72..3992de4d4f 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_mistral_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_mistral_transformation.py @@ -87,7 +87,9 @@ class AmazonMistralConfig(AmazonInvokeConfig, BaseConfig): return optional_params @staticmethod - def get_outputText(completion_response: dict, model_response: "ModelResponse") -> str: + def get_outputText( + completion_response: dict, model_response: "ModelResponse" + ) -> str: """This function extracts the output text from a bedrock mistral completion. As a side effect, it updates the finish reason for a model response. @@ -101,11 +103,17 @@ class AmazonMistralConfig(AmazonInvokeConfig, BaseConfig): """ if "choices" in completion_response: outputText = completion_response["choices"][0]["message"]["content"] - model_response.choices[0].finish_reason = completion_response["choices"][0]["finish_reason"] + model_response.choices[0].finish_reason = completion_response["choices"][0][ + "finish_reason" + ] elif "outputs" in completion_response: outputText = completion_response["outputs"][0]["text"] - model_response.choices[0].finish_reason = completion_response["outputs"][0]["stop_reason"] + model_response.choices[0].finish_reason = completion_response["outputs"][0][ + "stop_reason" + ] else: - raise BedrockError(message="Unexpected mistral completion response", status_code=400) + raise BedrockError( + message="Unexpected mistral completion response", status_code=400 + ) return outputText diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py index e53410760d..3aeb65b58c 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py @@ -32,11 +32,11 @@ else: class AmazonMoonshotConfig(AmazonInvokeConfig, MoonshotChatConfig): """ Configuration for Bedrock Moonshot AI (Kimi K2) models. - + Reference: https://aws.amazon.com/about-aws/whats-new/2025/12/amazon-bedrock-fully-managed-open-weight-models/ https://platform.moonshot.ai/docs/api/chat - + Supported Params for the Amazon / Moonshot models: - `max_tokens` (integer) max tokens - `temperature` (float) temperature for model (0-1 for Moonshot) @@ -44,10 +44,10 @@ class AmazonMoonshotConfig(AmazonInvokeConfig, MoonshotChatConfig): - `stream` (bool) whether to stream responses - `tools` (list) tool definitions (supported on kimi-k2-thinking) - `tool_choice` (str|dict) tool choice specification (supported on kimi-k2-thinking) - + NOT Supported on Bedrock: - `stop` sequences (Bedrock doesn't support stopSequences field for this model) - + Note: The kimi-k2-thinking model DOES support tool calls, unlike kimi-thinking-preview. """ @@ -62,7 +62,7 @@ class AmazonMoonshotConfig(AmazonInvokeConfig, MoonshotChatConfig): def _get_model_id(self, model: str) -> str: """ Extract the actual model ID from the LiteLLM model name. - + Removes routing prefixes like: - bedrock/invoke/moonshot.kimi-k2-thinking -> moonshot.kimi-k2-thinking - invoke/moonshot.kimi-k2-thinking -> moonshot.kimi-k2-thinking @@ -71,39 +71,44 @@ class AmazonMoonshotConfig(AmazonInvokeConfig, MoonshotChatConfig): # Remove bedrock/ prefix if present if model.startswith("bedrock/"): model = model[8:] - + # Remove invoke/ prefix if present if model.startswith("invoke/"): model = model[7:] - + # Remove any provider prefix (e.g., moonshot/) if "/" in model and not model.startswith("arn:"): parts = model.split("/", 1) if len(parts) == 2: model = parts[1] - + return model def get_supported_openai_params(self, model: str) -> List[str]: """ Get the supported OpenAI params for Moonshot AI models on Bedrock. - + Bedrock-specific limitations: - stopSequences field is not supported on Bedrock (unlike native Moonshot API) - functions parameter is not supported (use tools instead) - tool_choice doesn't support "required" value - + Note: kimi-k2-thinking DOES support tool calls (unlike kimi-thinking-preview) The parent MoonshotChatConfig class handles the kimi-thinking-preview exclusion. """ - excluded_params: List[str] = ["functions", "stop"] # Bedrock doesn't support stopSequences - - base_openai_params = super(MoonshotChatConfig, self).get_supported_openai_params(model=model) + excluded_params: List[str] = [ + "functions", + "stop", + ] # Bedrock doesn't support stopSequences + + base_openai_params = super( + MoonshotChatConfig, self + ).get_supported_openai_params(model=model) final_params: List[str] = [] for param in base_openai_params: if param not in excluded_params: final_params.append(param) - + return final_params def map_openai_params( @@ -115,7 +120,7 @@ class AmazonMoonshotConfig(AmazonInvokeConfig, MoonshotChatConfig): ) -> dict: """ Map OpenAI parameters to Moonshot AI parameters for Bedrock. - + Handles Moonshot AI specific limitations: - tool_choice doesn't support "required" value - Temperature <0.3 limitation for n>1 @@ -139,7 +144,7 @@ class AmazonMoonshotConfig(AmazonInvokeConfig, MoonshotChatConfig): ) -> dict: """ Transform the request for Bedrock Moonshot AI models. - + Uses the Moonshot transformation logic which handles: - Converting content lists to strings (Moonshot doesn't support list format) - Adding tool_choice="required" message if needed @@ -148,10 +153,10 @@ class AmazonMoonshotConfig(AmazonInvokeConfig, MoonshotChatConfig): """ # Filter out AWS credentials using the existing method from BaseAWSLLM self._get_boto_credentials_from_optional_params(optional_params, model) - + # Strip routing prefixes to get the actual model ID clean_model_id = self._get_model_id(model) - + # Use Moonshot's transform_request which handles message transformation # and tool_choice="required" workaround return MoonshotChatConfig.transform_request( @@ -163,34 +168,34 @@ class AmazonMoonshotConfig(AmazonInvokeConfig, MoonshotChatConfig): headers=headers, ) - def _extract_reasoning_from_content(self, content: str) -> tuple[Optional[str], str]: + def _extract_reasoning_from_content( + self, content: str + ) -> tuple[Optional[str], str]: """ Extract reasoning content from tags in the response. - + Moonshot AI's Kimi K2 Thinking model returns reasoning in tags. This method extracts that content and returns it separately. - + Args: content: The full content string from the API response - + Returns: tuple: (reasoning_content, main_content) """ if not content: return None, content - + # Match ... tags reasoning_match = re.match( - r"(.*?)\s*(.*)", - content, - re.DOTALL + r"(.*?)\s*(.*)", content, re.DOTALL ) - + if reasoning_match: reasoning_content = reasoning_match.group(1).strip() main_content = reasoning_match.group(2).strip() return reasoning_content, main_content - + return None, content def transform_response( @@ -209,7 +214,7 @@ class AmazonMoonshotConfig(AmazonInvokeConfig, MoonshotChatConfig): ) -> "ModelResponse": """ Transform the response from Bedrock Moonshot AI models. - + Moonshot AI uses OpenAI-compatible response format, but returns reasoning content in tags. This method: 1. Calls parent class transformation @@ -231,22 +236,27 @@ class AmazonMoonshotConfig(AmazonInvokeConfig, MoonshotChatConfig): api_key=api_key, json_mode=json_mode, ) - + # Extract reasoning content from tags if model_response.choices and len(model_response.choices) > 0: for choice in model_response.choices: # Only process Choices (not StreamingChoices) which have message attribute - if isinstance(choice, Choices) and choice.message and choice.message.content: - reasoning_content, main_content = self._extract_reasoning_from_content( - choice.message.content - ) - + if ( + isinstance(choice, Choices) + and choice.message + and choice.message.content + ): + ( + reasoning_content, + main_content, + ) = self._extract_reasoning_from_content(choice.message.content) + if reasoning_content: # Set the reasoning_content field choice.message.reasoning_content = reasoning_content # Update the main content without reasoning tags choice.message.content = main_content - + return model_response def get_error_class( diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py index a438be1745..7b64c6066d 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py @@ -28,14 +28,14 @@ else: class AmazonBedrockOpenAIConfig(OpenAIGPTConfig, BaseAWSLLM): """ Configuration for Bedrock imported models that use OpenAI Chat Completions format. - + This class handles the transformation of requests and responses for Bedrock imported models that accept the OpenAI API format directly. - + Inherits from OpenAIGPTConfig to leverage standard OpenAI parameter handling and response transformation, while adding Bedrock-specific URL generation and AWS request signing. - + Usage: model = "bedrock/openai/arn:aws:bedrock:us-east-1:123456789012:imported-model/abc123" """ @@ -51,18 +51,18 @@ class AmazonBedrockOpenAIConfig(OpenAIGPTConfig, BaseAWSLLM): def _get_openai_model_id(self, model: str) -> str: """ Extract the actual model ID from the LiteLLM model name. - + Input format: bedrock/openai/ Returns: """ # Remove bedrock/ prefix if present if model.startswith("bedrock/"): model = model[8:] - + # Remove openai/ prefix if model.startswith("openai/"): model = model[7:] - + return model def get_complete_url( @@ -76,16 +76,16 @@ class AmazonBedrockOpenAIConfig(OpenAIGPTConfig, BaseAWSLLM): ) -> str: """ Get the complete URL for the Bedrock invoke endpoint. - + Uses the standard Bedrock invoke endpoint format. """ model_id = self._get_openai_model_id(model) - + # Get AWS region aws_region_name = self._get_aws_region_name( optional_params=optional_params, model=model ) - + # Get runtime endpoint aws_bedrock_runtime_endpoint = optional_params.get( "aws_bedrock_runtime_endpoint", None @@ -98,13 +98,15 @@ class AmazonBedrockOpenAIConfig(OpenAIGPTConfig, BaseAWSLLM): # Encode model ID for ARNs (e.g., :imported-model/ -> :imported-model%2F) model_id = CommonUtils.encode_bedrock_runtime_modelid_arn(model_id) - + # Build the invoke URL if stream: - endpoint_url = f"{endpoint_url}/model/{model_id}/invoke-with-response-stream" + endpoint_url = ( + f"{endpoint_url}/model/{model_id}/invoke-with-response-stream" + ) else: endpoint_url = f"{endpoint_url}/model/{model_id}/invoke" - + return endpoint_url def sign_request( @@ -143,20 +145,20 @@ class AmazonBedrockOpenAIConfig(OpenAIGPTConfig, BaseAWSLLM): ) -> dict: """ Transform the request to OpenAI Chat Completions format for Bedrock imported models. - + Removes AWS-specific params and stream param (handled separately in URL), then delegates to parent class for standard OpenAI request transformation. """ # Remove stream from optional_params as it's handled separately in URL optional_params.pop("stream", None) - + # Remove AWS-specific params that shouldn't be in the request body inference_params = { k: v for k, v in optional_params.items() if k not in self.aws_authentication_params } - + # Use parent class transform_request for OpenAI format return super().transform_request( model=self._get_openai_model_id(model), @@ -178,7 +180,7 @@ class AmazonBedrockOpenAIConfig(OpenAIGPTConfig, BaseAWSLLM): ) -> dict: """ Validate the environment and return headers. - + For Bedrock, we don't need Bearer token auth since we use AWS SigV4. """ return headers diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py index fe0fd40b55..c65e9e0b08 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py @@ -24,10 +24,10 @@ from litellm.types.utils import ModelResponse, Usage class AmazonQwen2Config(AmazonQwen3Config): """ Config for sending `qwen2` requests to `/bedrock/invoke/` - + Inherits from AmazonQwen3Config since Qwen2 and Qwen3 architectures are mostly similar. The main difference is in the response format: Qwen2 uses "text" field while Qwen3 uses "generation" field. - + Reference: https://docs.aws.amazon.com/bedrock/latest/userguide/invoke-imported-model.html """ @@ -47,30 +47,32 @@ class AmazonQwen2Config(AmazonQwen3Config): ) -> ModelResponse: """ Transform Qwen2 Bedrock response to OpenAI format - + Qwen2 uses "text" field, but we also support "generation" field for compatibility. """ try: - if hasattr(raw_response, 'json'): + if hasattr(raw_response, "json"): response_data = raw_response.json() else: response_data = raw_response - + # Extract the generated text - Qwen2 uses "text" field, but also support "generation" for compatibility - generated_text = response_data.get("generation", "") or response_data.get("text", "") - + generated_text = response_data.get("generation", "") or response_data.get( + "text", "" + ) + # Clean up the response (remove assistant start token if present) if generated_text.startswith("<|im_start|>assistant\n"): - generated_text = generated_text[len("<|im_start|>assistant\n"):] + generated_text = generated_text[len("<|im_start|>assistant\n") :] if generated_text.endswith("<|im_end|>"): - generated_text = generated_text[:-len("<|im_end|>")] - + generated_text = generated_text[: -len("<|im_end|>")] + # Set the content in the existing model_response structure - if hasattr(model_response, 'choices') and len(model_response.choices) > 0: + if hasattr(model_response, "choices") and len(model_response.choices) > 0: choice = model_response.choices[0] choice.message.content = generated_text choice.finish_reason = "stop" - + # Set usage information if available in response if "usage" in response_data: usage_data = response_data["usage"] @@ -83,9 +85,9 @@ class AmazonQwen2Config(AmazonQwen3Config): total_tokens=usage_data.get("total_tokens", 0), ), ) - + return model_response - + except Exception as e: if logging_obj: logging_obj.post_call( @@ -95,4 +97,3 @@ class AmazonQwen2Config(AmazonQwen3Config): additional_args={"error": str(e)}, ) raise e - diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py index 4be3e370fa..6325c38818 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py @@ -22,7 +22,7 @@ from litellm.types.utils import ModelResponse, Usage class AmazonQwen3Config(AmazonInvokeConfig, BaseConfig): """ Config for sending `qwen3` requests to `/bedrock/invoke/` - + Reference: https://docs.aws.amazon.com/bedrock/latest/userguide/invoke-imported-model.html """ @@ -91,12 +91,12 @@ class AmazonQwen3Config(AmazonInvokeConfig, BaseConfig): """ # Convert messages to prompt format prompt = self._convert_messages_to_prompt(messages) - + # Build the request body request_body = { "prompt": prompt, } - + # Add optional parameters if "max_tokens" in optional_params: request_body["max_gen_len"] = optional_params["max_tokens"] @@ -108,7 +108,7 @@ class AmazonQwen3Config(AmazonInvokeConfig, BaseConfig): request_body["top_k"] = optional_params["top_k"] if "stop" in optional_params: request_body["stop"] = optional_params["stop"] - + return request_body def _convert_messages_to_prompt(self, messages: List[AllMessageValues]) -> str: @@ -117,12 +117,12 @@ class AmazonQwen3Config(AmazonInvokeConfig, BaseConfig): Supports tool calls, multimodal content, and various message types """ prompt_parts = [] - + for message in messages: role = message.get("role", "") content = message.get("content", "") tool_calls = message.get("tool_calls", []) - + if role == "system": prompt_parts.append(f"<|im_start|>system\n{content}<|im_end|>") elif role == "user": @@ -134,7 +134,9 @@ class AmazonQwen3Config(AmazonInvokeConfig, BaseConfig): text_content.append(item.get("text", "")) elif item.get("type") == "image_url": # For Qwen3, we can include image placeholders - text_content.append("<|vision_start|><|image_pad|><|vision_end|>") + text_content.append( + "<|vision_start|><|image_pad|><|vision_end|>" + ) content = "".join(text_content) prompt_parts.append(f"<|im_start|>user\n{content}<|im_end|>") elif role == "assistant": @@ -142,17 +144,21 @@ class AmazonQwen3Config(AmazonInvokeConfig, BaseConfig): # Handle tool calls for tool_call in tool_calls: function_name = tool_call.get("function", {}).get("name", "") - function_args = tool_call.get("function", {}).get("arguments", "") - prompt_parts.append(f"<|im_start|>assistant\n\n{{\"name\": \"{function_name}\", \"arguments\": \"{function_args}\"}}\n<|im_end|>") + function_args = tool_call.get("function", {}).get( + "arguments", "" + ) + prompt_parts.append( + f'<|im_start|>assistant\n\n{{"name": "{function_name}", "arguments": "{function_args}"}}\n<|im_end|>' + ) else: prompt_parts.append(f"<|im_start|>assistant\n{content}<|im_end|>") elif role == "tool": # Handle tool responses prompt_parts.append(f"<|im_start|>tool\n{content}<|im_end|>") - + # Add assistant start token for response generation prompt_parts.append("<|im_start|>assistant\n") - + return "\n".join(prompt_parts) def transform_response( @@ -173,26 +179,26 @@ class AmazonQwen3Config(AmazonInvokeConfig, BaseConfig): Transform Qwen3 Bedrock response to OpenAI format """ try: - if hasattr(raw_response, 'json'): + if hasattr(raw_response, "json"): response_data = raw_response.json() else: response_data = raw_response - + # Extract the generated text - Qwen3 uses "generation" field generated_text = response_data.get("generation", "") - + # Clean up the response (remove assistant start token if present) if generated_text.startswith("<|im_start|>assistant\n"): - generated_text = generated_text[len("<|im_start|>assistant\n"):] + generated_text = generated_text[len("<|im_start|>assistant\n") :] if generated_text.endswith("<|im_end|>"): - generated_text = generated_text[:-len("<|im_end|>")] - + generated_text = generated_text[: -len("<|im_end|>")] + # Set the content in the existing model_response structure - if hasattr(model_response, 'choices') and len(model_response.choices) > 0: + if hasattr(model_response, "choices") and len(model_response.choices) > 0: choice = model_response.choices[0] choice.message.content = generated_text choice.finish_reason = "stop" - + # Set usage information if available in response if "usage" in response_data: usage_data = response_data["usage"] @@ -205,9 +211,9 @@ class AmazonQwen3Config(AmazonInvokeConfig, BaseConfig): total_tokens=usage_data.get("total_tokens", 0), ), ) - + return model_response - + except Exception as e: if logging_obj: logging_obj.post_call( diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py index 62e98f7472..889480d31a 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py @@ -70,12 +70,12 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): def _normalize_response_format(self, value: Any) -> Any: """Normalize response_format to TwelveLabs format. - + TwelveLabs expects: { "jsonSchema": {...} } - + But OpenAI format is: { "type": "json_schema", @@ -120,14 +120,14 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): for key in ("temperature", "maxOutputTokens"): if key in optional_params: request_data[key] = optional_params.get(key) - + # Handle responseFormat - transform to TwelveLabs format if "responseFormat" in optional_params: response_format = optional_params["responseFormat"] transformed_format = self._normalize_response_format(response_format) if transformed_format: request_data["responseFormat"] = transformed_format - + return request_data def _build_media_source(self, optional_params: dict) -> Optional[dict]: @@ -200,13 +200,13 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): ) -> ModelResponse: """ Transform TwelveLabs Pegasus response to LiteLLM format. - + TwelveLabs response format: { "message": "...", "finishReason": "stop" | "length" } - + LiteLLM format: ModelResponse with choices[0].message.content and finish_reason """ @@ -217,25 +217,26 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): message=f"Error parsing response: {raw_response.text}, error: {str(e)}", status_code=raw_response.status_code, ) - + verbose_logger.debug( "twelvelabs pegasus response: %s", json.dumps(completion_response, indent=4, default=str), ) - + # Extract message content message_content = completion_response.get("message", "") - + # Extract finish reason and map to LiteLLM format finish_reason_raw = completion_response.get("finishReason", "stop") finish_reason = map_finish_reason(finish_reason_raw) - + # Set the response content try: if ( message_content and hasattr(model_response.choices[0], "message") - and getattr(model_response.choices[0].message, "tool_calls", None) is None + and getattr(model_response.choices[0].message, "tool_calls", None) + is None ): model_response.choices[0].message.content = message_content # type: ignore model_response.choices[0].finish_reason = finish_reason @@ -246,7 +247,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): message=f"Error setting response content: {str(e)}. Response: {completion_response}", status_code=raw_response.status_code, ) - + # Calculate usage from headers bedrock_input_tokens = raw_response.headers.get( "x-amzn-bedrock-input-token-count", None @@ -254,11 +255,11 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): bedrock_output_tokens = raw_response.headers.get( "x-amzn-bedrock-output-token-count", None ) - + prompt_tokens = int( bedrock_input_tokens or litellm.token_counter(messages=messages) ) - + completion_tokens = int( bedrock_output_tokens or litellm.token_counter( @@ -266,7 +267,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): count_response_tokens=True, ) ) - + model_response.created = int(time.time()) model_response.model = model usage = Usage( @@ -275,6 +276,5 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): total_tokens=prompt_tokens + completion_tokens, ) setattr(model_response, "usage", usage) - - return model_response + return model_response diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 328c3a0b97..7936b6ea64 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -63,7 +63,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): if "response_format" in non_default_params: # Use a model name that forces tool-based approach model = "claude-3-sonnet-20240229" - + optional_params = AnthropicConfig.map_openai_params( self, non_default_params, @@ -71,12 +71,11 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): model, drop_params, ) - + # Restore original model name model = original_model - - return optional_params + return optional_params def transform_request( self, @@ -94,12 +93,12 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): if k not in self.aws_authentication_params } filtered_params = self._normalize_bedrock_tool_search_tools(filtered_params) - + _anthropic_request = AnthropicConfig.transform_request( self, model=model, messages=messages, - optional_params=filtered_params, + optional_params=filtered_params, litellm_params=litellm_params, headers=headers, ) @@ -130,15 +129,14 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): model=model, optional_params=optional_params, computer_tool_used=self.is_computer_tool_used(tools), - prompt_caching_set=False, + prompt_caching_set=False, file_id_used=self.is_file_id_used(messages), mcp_server_used=self.is_mcp_server_used(optional_params.get("mcp_servers")), ) beta_set.update(auto_betas) - if ( - tool_search_used - and not (programmatic_tool_calling_used or input_examples_used) + if tool_search_used and not ( + programmatic_tool_calling_used or input_examples_used ): beta_set.discard(ANTHROPIC_TOOL_SEARCH_BETA_HEADER) if "opus-4" in model.lower() or "opus_4" in model.lower(): diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 8e944988a9..9666aa68c9 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -455,7 +455,7 @@ def get_bedrock_base_model(model: str) -> str: stripped = model for rp in ["bedrock/converse/", "bedrock/", "converse/"]: if stripped.startswith(rp): - stripped = stripped[len(rp):] + stripped = stripped[len(rp) :] break if stripped.startswith("nova-2/"): return "amazon.nova-2-custom" @@ -638,7 +638,9 @@ class BedrockModelInfo(BaseLLMModelInfo): # Check for nova spec prefixes (nova/ and nova-2/) _model_after_bedrock = model.replace("bedrock/", "", 1) - if _model_after_bedrock.startswith("nova-2/") or _model_after_bedrock.startswith("nova/"): + if _model_after_bedrock.startswith( + "nova-2/" + ) or _model_after_bedrock.startswith("nova/"): return "converse" base_model = BedrockModelInfo.get_base_model(model) diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 772eb16968..eb7755574a 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -101,9 +101,7 @@ class BedrockTokenCounter(BaseTokenCounter): status_code=e.status_code, ) except Exception as e: - verbose_logger.warning( - f"Error calling Bedrock CountTokens API: {e}" - ) + verbose_logger.warning(f"Error calling Bedrock CountTokens API: {e}") return TokenCountResponse( total_tokens=0, request_model=request_model, diff --git a/litellm/llms/bedrock/count_tokens/handler.py b/litellm/llms/bedrock/count_tokens/handler.py index 9d2be6cca8..cfd32342d1 100644 --- a/litellm/llms/bedrock/count_tokens/handler.py +++ b/litellm/llms/bedrock/count_tokens/handler.py @@ -84,14 +84,16 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): api_key=api_key, ) - async_client = get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK) + async_client = get_async_httpx_client( + llm_provider=litellm.LlmProviders.BEDROCK + ) response = await async_client.post( - endpoint_url, - headers=signed_headers, - data=signed_body, - timeout=30.0, - ) + endpoint_url, + headers=signed_headers, + data=signed_body, + timeout=30.0, + ) verbose_logger.debug(f"Response status: {response.status_code}") diff --git a/litellm/llms/bedrock/count_tokens/transformation.py b/litellm/llms/bedrock/count_tokens/transformation.py index 64f1098e64..fe9ab80ced 100644 --- a/litellm/llms/bedrock/count_tokens/transformation.py +++ b/litellm/llms/bedrock/count_tokens/transformation.py @@ -91,7 +91,10 @@ class BedrockCountTokensConfig(BaseAWSLLM): # Transform messages user_messages = [] for message in messages: - transformed_message: Dict[str, Any] = {"role": message.get("role"), "content": []} + transformed_message: Dict[str, Any] = { + "role": message.get("role"), + "content": [], + } content = message.get("content", "") if isinstance(content, str): transformed_message["content"].append({"text": content}) @@ -121,10 +124,16 @@ class BedrockCountTokensConfig(BaseAWSLLM): return [{"text": system}] if isinstance(system, list): # Already in blocks format (e.g. [{"type": "text", "text": "..."}]) - return [{"text": block.get("text", "")} for block in system if isinstance(block, dict)] + return [ + {"text": block.get("text", "")} + for block in system + if isinstance(block, dict) + ] return [] - def _transform_tools(self, tools: Optional[List[Dict[str, Any]]]) -> Optional[Dict[str, Any]]: + def _transform_tools( + self, tools: Optional[List[Dict[str, Any]]] + ) -> Optional[Dict[str, Any]]: """Transform Anthropic tools to Bedrock toolConfig format.""" if not tools: return None @@ -139,15 +148,19 @@ class BedrockCountTokensConfig(BaseAWSLLM): name = name[:64] description = tool.get("description") or name - input_schema = tool.get("input_schema", {"type": "object", "properties": {}}) + input_schema = tool.get( + "input_schema", {"type": "object", "properties": {}} + ) - bedrock_tools.append({ - "toolSpec": { - "name": name, - "description": description, - "inputSchema": {"json": input_schema}, + bedrock_tools.append( + { + "toolSpec": { + "name": name, + "description": description, + "inputSchema": {"json": input_schema}, + } } - }) + ) return {"tools": bedrock_tools} diff --git a/litellm/llms/bedrock/embed/amazon_nova_transformation.py b/litellm/llms/bedrock/embed/amazon_nova_transformation.py index 40d2a21e1c..c20b52a6e0 100644 --- a/litellm/llms/bedrock/embed/amazon_nova_transformation.py +++ b/litellm/llms/bedrock/embed/amazon_nova_transformation.py @@ -14,13 +14,18 @@ Docs - https://docs.aws.amazon.com/bedrock/latest/userguide/nova-embed.html from typing import List, Optional -from litellm.types.utils import Embedding, EmbeddingResponse, PromptTokensDetailsWrapper, Usage +from litellm.types.utils import ( + Embedding, + EmbeddingResponse, + PromptTokensDetailsWrapper, + Usage, +) class AmazonNovaEmbeddingConfig: """ Reference: https://docs.aws.amazon.com/bedrock/latest/userguide/nova-embed.html - + Amazon Nova Multimodal Embeddings supports: - Text, image, video, and audio inputs - Synchronous (InvokeModel) and asynchronous (StartAsyncInvoke) APIs @@ -46,14 +51,14 @@ class AmazonNovaEmbeddingConfig: elif k in self.get_supported_openai_params(): optional_params[k] = v return optional_params - + def _parse_data_url(self, data_url: str) -> tuple: """ Parse a data URL to extract the media type and base64 data. - + Args: data_url: Data URL in format: data:image/jpeg;base64,/9j/4AAQ... - + Returns: tuple: (media_type, base64_data) media_type: e.g., "image/jpeg", "video/mp4", "audio/mpeg" @@ -61,23 +66,25 @@ class AmazonNovaEmbeddingConfig: """ if not data_url.startswith("data:"): raise ValueError(f"Invalid data URL format: {data_url[:50]}...") - + # Split by comma to separate metadata from data # Format: data:image/jpeg;base64, if "," not in data_url: - raise ValueError(f"Invalid data URL format (missing comma): {data_url[:50]}...") - + raise ValueError( + f"Invalid data URL format (missing comma): {data_url[:50]}..." + ) + metadata, base64_data = data_url.split(",", 1) - + # Extract media type from metadata # Remove 'data:' prefix and ';base64' suffix metadata = metadata[5:] # Remove 'data:' - + if ";" in metadata: media_type = metadata.split(";")[0] else: media_type = metadata - + return media_type, base64_data def _transform_request( @@ -90,111 +97,109 @@ class AmazonNovaEmbeddingConfig: ) -> dict: """ Transform OpenAI-style input to Nova format. - + Only handles OpenAI params (dimensions). All other Nova-specific params should be passed via inference_params and will be passed through as-is. - + Args: input: The input text or media reference inference_params: Additional parameters (will be passed through) async_invoke_route: Whether this is for async invoke model_id: Model ID (for async invoke) output_s3_uri: S3 URI for output (for async invoke) - + Returns: dict: Nova embedding request """ # Determine task type task_type = "SEGMENTED_EMBEDDING" if async_invoke_route else "SINGLE_EMBEDDING" - + # Build the base request structure request: dict = { "schemaVersion": "nova-multimodal-embed-v1", "taskType": task_type, } - + # Start with inference_params (user-provided params) embedding_params = inference_params.copy() - + embedding_params.pop("output_s3_uri", None) - + # Map OpenAI dimensions to embeddingDimension if provided if "dimensions" in embedding_params: embedding_params["embeddingDimension"] = embedding_params.pop("dimensions") elif "embedding_dimension" in embedding_params: - embedding_params["embeddingDimension"] = embedding_params.pop("embedding_dimension") - + embedding_params["embeddingDimension"] = embedding_params.pop( + "embedding_dimension" + ) + # Add required embeddingPurpose if not provided (required by Nova API) if "embeddingPurpose" not in embedding_params: embedding_params["embeddingPurpose"] = "GENERIC_INDEX" - + # Add required embeddingDimension if not provided (required by Nova API) if "embeddingDimension" not in embedding_params: embedding_params["embeddingDimension"] = 3072 - + # For text/media input, add basic structure if user hasn't provided text/image/video/audio - if "text" not in embedding_params and "image" not in embedding_params and "video" not in embedding_params and "audio" not in embedding_params: + if ( + "text" not in embedding_params + and "image" not in embedding_params + and "video" not in embedding_params + and "audio" not in embedding_params + ): # Check if input is a data URL (e.g., data:image/jpeg;base64,...) if input.startswith("data:"): # Parse the data URL to extract media type and base64 data media_type, base64_data = self._parse_data_url(input) - + if media_type.startswith("image/"): # Extract image format from MIME type (e.g., image/jpeg -> jpeg) image_format = media_type.split("/")[1].lower() # Nova API expects specific formats if image_format == "jpg": image_format = "jpeg" - + embedding_params["image"] = { "format": image_format, - "source": { - "bytes": base64_data - } + "source": {"bytes": base64_data}, } elif media_type.startswith("video/"): # Handle video data URLs video_format = media_type.split("/")[1].lower() embedding_params["video"] = { "format": video_format, - "source": { - "bytes": base64_data - } + "source": {"bytes": base64_data}, } elif media_type.startswith("audio/"): # Handle audio data URLs audio_format = media_type.split("/")[1].lower() embedding_params["audio"] = { "format": audio_format, - "source": { - "bytes": base64_data - } + "source": {"bytes": base64_data}, } else: # Fallback to text for unknown types - embedding_params["text"] = { - "value": input, - "truncationMode": "END" - } + embedding_params["text"] = {"value": input, "truncationMode": "END"} elif input.startswith("s3://"): # S3 URL - default to text for now, user should specify modality embedding_params["text"] = { "source": {"s3Location": {"uri": input}}, - "truncationMode": "END" # Required by Nova API + "truncationMode": "END", # Required by Nova API } else: # Plain text input embedding_params["text"] = { "value": input, - "truncationMode": "END" # Required by Nova API + "truncationMode": "END", # Required by Nova API } - + # Set the embedding params in the request if task_type == "SINGLE_EMBEDDING": request["singleEmbeddingParams"] = embedding_params else: request["segmentedEmbeddingParams"] = embedding_params - + # For async invoke, wrap in the async invoke format if async_invoke_route and model_id: return self._wrap_async_invoke_request( @@ -202,7 +207,7 @@ class AmazonNovaEmbeddingConfig: model_id=model_id, output_s3_uri=output_s3_uri, ) - + return request def _wrap_async_invoke_request( @@ -213,12 +218,12 @@ class AmazonNovaEmbeddingConfig: ) -> dict: """ Wrap the transformed request in the AWS Bedrock async invoke format. - + Args: model_input: The transformed Nova embedding request model_id: The model identifier (without async_invoke prefix) output_s3_uri: S3 URI for output data config - + Returns: dict: The wrapped async invoke request """ @@ -228,19 +233,15 @@ class AmazonNovaEmbeddingConfig: unquoted_model_id = urllib.parse.unquote(model_id) if unquoted_model_id.startswith("async_invoke/"): unquoted_model_id = unquoted_model_id.replace("async_invoke/", "") - + # Validate that the S3 URI is not empty if not output_s3_uri or output_s3_uri.strip() == "": raise ValueError("output_s3_uri is required for async invoke requests") - + return { "modelId": unquoted_model_id, "modelInput": model_input, - "outputDataConfig": { - "s3OutputDataConfig": { - "s3Uri": output_s3_uri - } - }, + "outputDataConfig": {"s3OutputDataConfig": {"s3Uri": output_s3_uri}}, } def _transform_response( @@ -326,36 +327,35 @@ class AmazonNovaEmbeddingConfig: ) -> EmbeddingResponse: """ Transform async invoke response (invocation ARN) to OpenAI format. - + AWS async invoke returns: { "invocationArn": "arn:aws:bedrock:us-east-1:123456789012:async-invoke/abc123" } - + We transform this to a job-like embedding response with the ARN in hidden params. """ invocation_arn = response.get("invocationArn", "") - + # Create a placeholder embedding object for the job embedding = Embedding( embedding=[], # Empty embedding for async jobs index=0, object="embedding", ) - + # Create usage object (empty for async jobs) usage = Usage(prompt_tokens=0, total_tokens=0) - + # Create hidden params with job ID from litellm.types.llms.base import HiddenParams - + hidden_params = HiddenParams() setattr(hidden_params, "_invocation_arn", invocation_arn) - + return EmbeddingResponse( data=[embedding], model=model, usage=usage, hidden_params=hidden_params, ) - diff --git a/litellm/llms/bedrock/embed/amazon_titan_multimodal_transformation.py b/litellm/llms/bedrock/embed/amazon_titan_multimodal_transformation.py index e59d3cbf77..07b04734c3 100644 --- a/litellm/llms/bedrock/embed/amazon_titan_multimodal_transformation.py +++ b/litellm/llms/bedrock/embed/amazon_titan_multimodal_transformation.py @@ -13,7 +13,12 @@ from litellm.types.llms.bedrock import ( AmazonTitanMultimodalEmbeddingRequest, AmazonTitanMultimodalEmbeddingResponse, ) -from litellm.types.utils import Embedding, EmbeddingResponse, PromptTokensDetailsWrapper, Usage +from litellm.types.utils import ( + Embedding, + EmbeddingResponse, + PromptTokensDetailsWrapper, + Usage, +) from litellm.utils import get_base64_str, is_base64_encoded diff --git a/litellm/llms/bedrock/embed/amazon_titan_v2_transformation.py b/litellm/llms/bedrock/embed/amazon_titan_v2_transformation.py index ff748b58e8..ca0b95cd64 100644 --- a/litellm/llms/bedrock/embed/amazon_titan_v2_transformation.py +++ b/litellm/llms/bedrock/embed/amazon_titan_v2_transformation.py @@ -30,7 +30,9 @@ class AmazonTitanV2Config: normalize: Optional[bool] = None dimensions: Optional[int] = None - def __init__(self, normalize: Optional[bool] = None, dimensions: Optional[int] = None) -> None: + def __init__( + self, normalize: Optional[bool] = None, dimensions: Optional[int] = None + ) -> None: locals_ = locals().copy() for key, value in locals_.items(): if key != "self" and value is not None: @@ -57,7 +59,9 @@ class AmazonTitanV2Config: def get_supported_openai_params(self) -> List[str]: return ["dimensions", "encoding_format"] - def map_openai_params(self, non_default_params: dict, optional_params: dict) -> dict: + def map_openai_params( + self, non_default_params: dict, optional_params: dict + ) -> dict: for k, v in non_default_params.items(): if k == "dimensions": optional_params["dimensions"] = v @@ -73,10 +77,14 @@ class AmazonTitanV2Config: optional_params["embeddingTypes"] = ["float"] return optional_params - def _transform_request(self, input: str, inference_params: dict) -> AmazonTitanV2EmbeddingRequest: + def _transform_request( + self, input: str, inference_params: dict + ) -> AmazonTitanV2EmbeddingRequest: return AmazonTitanV2EmbeddingRequest(inputText=input, **inference_params) # type: ignore - def _transform_response(self, response_list: List[dict], model: str) -> EmbeddingResponse: + def _transform_response( + self, response_list: List[dict], model: str + ) -> EmbeddingResponse: total_prompt_tokens = 0 transformed_responses: List[Embedding] = [] @@ -88,12 +96,16 @@ class AmazonTitanV2Config: # Otherwise, use float data from embeddingsByType or fallback to embedding field embedding_data: Union[List[float], List[int]] - if ("embeddingsByType" in _parsed_response and - "binary" in _parsed_response["embeddingsByType"]): + if ( + "embeddingsByType" in _parsed_response + and "binary" in _parsed_response["embeddingsByType"] + ): # Use binary data if available (for encoding_format="base64") embedding_data = _parsed_response["embeddingsByType"]["binary"] - elif ("embeddingsByType" in _parsed_response and - "float" in _parsed_response["embeddingsByType"]): + elif ( + "embeddingsByType" in _parsed_response + and "float" in _parsed_response["embeddingsByType"] + ): # Use float data from embeddingsByType embedding_data = _parsed_response["embeddingsByType"]["float"] elif "embedding" in _parsed_response: diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index 783345d78d..27dc785bf5 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -287,7 +287,9 @@ class BedrockEmbedding(BaseAWSLLM): "headers": prepped.headers, }, ) - headers_for_request = dict(prepped.headers) if hasattr(prepped, 'headers') else {} + headers_for_request = ( + dict(prepped.headers) if hasattr(prepped, "headers") else {} + ) response = self._make_sync_call( client=client, timeout=timeout, @@ -357,7 +359,9 @@ class BedrockEmbedding(BaseAWSLLM): ) # Convert CaseInsensitiveDict to regular dict for httpx compatibility # This ensures custom headers are properly forwarded, especially with IAM roles and custom api_base - headers_for_request = dict(prepped.headers) if hasattr(prepped, 'headers') else {} + headers_for_request = ( + dict(prepped.headers) if hasattr(prepped, "headers") else {} + ) response = await self._make_async_call( client=client, timeout=timeout, @@ -570,7 +574,9 @@ class BedrockEmbedding(BaseAWSLLM): ## ROUTING ## # Convert CaseInsensitiveDict to regular dict for httpx compatibility - headers_for_request = dict(prepped.headers) if hasattr(prepped, 'headers') else {} + headers_for_request = ( + dict(prepped.headers) if hasattr(prepped, "headers") else {} + ) return cohere_embedding( model=model, input=input, @@ -612,7 +618,6 @@ class BedrockEmbedding(BaseAWSLLM): aws_region_name=aws_region_name, ) - from urllib.parse import quote # Encode the ARN for use in URL path @@ -627,9 +632,7 @@ class BedrockEmbedding(BaseAWSLLM): from botocore.auth import SigV4Auth from botocore.awsrequest import AWSRequest except ImportError: - raise ImportError( - "Missing boto3 to call bedrock. Run 'pip install boto3'." - ) + raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") # Create AWSRequest with GET method and encoded URL request = AWSRequest( @@ -638,11 +641,11 @@ class BedrockEmbedding(BaseAWSLLM): data=None, # GET request, no body headers=headers, ) - + # Sign the request - SigV4Auth will create canonical string from request URL sigv4 = SigV4Auth(credentials, "bedrock", aws_region_name) sigv4.add_auth(request) - + # Prepare the request prepped = request.prepare() diff --git a/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py b/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py index c85c388eeb..56339ed223 100644 --- a/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py +++ b/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py @@ -93,7 +93,9 @@ class TwelveLabsMarengoEmbeddingConfig: # Get input_type or default to "text" input_type = cast( TWELVELABS_EMBEDDING_INPUT_TYPES, - inference_params.get("inputType") or inference_params.get("input_type") or "text" + inference_params.get("inputType") + or inference_params.get("input_type") + or "text", ) # Validate that async-invoke is used for video/audio @@ -130,6 +132,7 @@ class TwelveLabsMarengoEmbeddingConfig: else: # Direct base64 string from litellm.utils import get_base64_str + b64_str = get_base64_str(input) transformed_request["mediaSource"] = {"base64String": b64_str} diff --git a/litellm/llms/bedrock/files/handler.py b/litellm/llms/bedrock/files/handler.py index 0350271dc4..13bd87a1f0 100644 --- a/litellm/llms/bedrock/files/handler.py +++ b/litellm/llms/bedrock/files/handler.py @@ -18,7 +18,7 @@ from ..base_aws_llm import BaseAWSLLM class BedrockFilesHandler(BaseAWSLLM): """ Handles downloading files from S3 for Bedrock batch processing. - + This implementation downloads files from S3 buckets where Bedrock stores batch output files. """ @@ -32,14 +32,14 @@ class BedrockFilesHandler(BaseAWSLLM): def _extract_s3_uri_from_file_id(self, file_id: str) -> str: """ Extract S3 URI from encoded file ID. - + The file ID can be in two formats: 1. Base64-encoded unified file ID containing: llm_output_file_id,s3://bucket/path 2. Direct S3 URI: s3://bucket/path - + Args: file_id: Encoded file ID or direct S3 URI - + Returns: S3 URI (e.g., "s3://bucket-name/path/to/file") """ @@ -48,7 +48,7 @@ class BedrockFilesHandler(BaseAWSLLM): # Add padding if needed padded = file_id + "=" * (-len(file_id) % 4) decoded = base64.urlsafe_b64decode(padded).decode() - + # Check if it's a unified file ID format if decoded.startswith(SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value): # Extract llm_output_file_id from the decoded string @@ -57,36 +57,38 @@ class BedrockFilesHandler(BaseAWSLLM): return s3_uri except Exception: pass - + # If not base64 encoded or doesn't contain llm_output_file_id, assume it's already an S3 URI if file_id.startswith("s3://"): return file_id - + # If it doesn't start with s3://, assume it's a direct S3 URI and add the prefix return f"s3://{file_id}" def _parse_s3_uri(self, s3_uri: str) -> Tuple[str, str]: """ Parse S3 URI to extract bucket name and object key. - + Args: s3_uri: S3 URI (e.g., "s3://bucket-name/path/to/file") - + Returns: Tuple of (bucket_name, object_key) """ if not s3_uri.startswith("s3://"): - raise ValueError(f"Invalid S3 URI format: {s3_uri}. Expected format: s3://bucket-name/path/to/file") - + raise ValueError( + f"Invalid S3 URI format: {s3_uri}. Expected format: s3://bucket-name/path/to/file" + ) + # Remove 's3://' prefix path = s3_uri[5:] - + if "/" in path: bucket_name, object_key = path.split("/", 1) else: bucket_name = path object_key = "" - + return bucket_name, object_key async def afile_content( @@ -98,27 +100,27 @@ class BedrockFilesHandler(BaseAWSLLM): ) -> HttpxBinaryResponseContent: """ Download file content from S3 bucket for Bedrock files. - + Args: file_content_request: Contains file_id (encoded or S3 URI) optional_params: Optional parameters containing AWS credentials timeout: Request timeout max_retries: Max retry attempts - + Returns: HttpxBinaryResponseContent: Binary content wrapped in compatible response format """ import boto3 from botocore.credentials import Credentials - + file_id = file_content_request.get("file_id") if not file_id: raise ValueError("file_id is required in file_content_request") - + # Extract S3 URI from file ID s3_uri = self._extract_s3_uri_from_file_id(file_id) bucket_name, object_key = self._parse_s3_uri(s3_uri) - + # Get AWS credentials aws_region_name = self._get_aws_region_name( optional_params=optional_params, model="" @@ -134,7 +136,7 @@ class BedrockFilesHandler(BaseAWSLLM): aws_web_identity_token=optional_params.get("aws_web_identity_token"), aws_sts_endpoint=optional_params.get("aws_sts_endpoint"), ) - + # Create S3 client s3_client = boto3.client( "s3", @@ -144,14 +146,16 @@ class BedrockFilesHandler(BaseAWSLLM): region_name=aws_region_name, verify=self._get_ssl_verify(), ) - + # Download file from S3 try: response = s3_client.get_object(Bucket=bucket_name, Key=object_key) file_content = response["Body"].read() except Exception as e: - raise ValueError(f"Failed to download file from S3: {s3_uri}. Error: {str(e)}") - + raise ValueError( + f"Failed to download file from S3: {s3_uri}. Error: {str(e)}" + ) + # Create mock HTTP response mock_response = httpx.Response( status_code=200, @@ -159,7 +163,7 @@ class BedrockFilesHandler(BaseAWSLLM): headers={"content-type": "application/octet-stream"}, request=httpx.Request(method="GET", url=s3_uri), ) - + return HttpxBinaryResponseContent(response=mock_response) def file_content( @@ -176,7 +180,7 @@ class BedrockFilesHandler(BaseAWSLLM): """ Download file content from S3 bucket for Bedrock files. Supports both sync and async operations. - + Args: _is_async: Whether to run asynchronously file_content_request: Contains file_id (encoded or S3 URI) @@ -184,7 +188,7 @@ class BedrockFilesHandler(BaseAWSLLM): optional_params: Optional parameters containing AWS credentials timeout: Request timeout max_retries: Max retry attempts - + Returns: HttpxBinaryResponseContent or Coroutine: Binary content wrapped in compatible response format """ @@ -204,4 +208,3 @@ class BedrockFilesHandler(BaseAWSLLM): max_retries=max_retries, ) ) - diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index e29b07ca3a..096371749b 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -36,7 +36,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): """ Config for Bedrock Files - handles S3 uploads for Bedrock batch processing """ - + def __init__(self): self.jsonl_transformation = BedrockJsonlFilesTransformation() super().__init__() @@ -65,8 +65,6 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): # No additional headers needed for S3 uploads - AWS credentials handled by BaseAWSLLM return headers - - def _get_content_from_openai_file(self, openai_file_content: FileTypes) -> str: """ Helper to extract content from various OpenAI file types and return as string. @@ -117,10 +115,10 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): # Remove bedrock/ prefix if present if _model.startswith("bedrock/"): _model = _model[8:] - + # Replace colons with hyphens for Bedrock S3 URI compliance _model = _model.replace(":", "-") - + object_name = f"litellm-bedrock-files-{_model}-{uuid.uuid4()}.jsonl" return object_name @@ -167,12 +165,16 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): """ Get the complete S3 URL for the file upload request """ - bucket_name = litellm_params.get("s3_bucket_name") or os.getenv("AWS_S3_BUCKET_NAME") + bucket_name = litellm_params.get("s3_bucket_name") or os.getenv( + "AWS_S3_BUCKET_NAME" + ) if not bucket_name: - raise ValueError("S3 bucket_name is required. Set 's3_bucket_name' in litellm_params or AWS_S3_BUCKET_NAME env var") - + raise ValueError( + "S3 bucket_name is required. Set 's3_bucket_name' in litellm_params or AWS_S3_BUCKET_NAME env var" + ) + aws_region_name = self._get_aws_region_name(optional_params, model) - + file_data = data.get("file") purpose = data.get("purpose") if file_data is None: @@ -181,10 +183,13 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): raise ValueError("purpose is required") extracted_file_data = extract_file_data(file_data) object_name = self.get_object_name(extracted_file_data, purpose) - + # S3 endpoint URL format - s3_endpoint_url = optional_params.get("s3_endpoint_url") or f"https://s3.{aws_region_name}.amazonaws.com" - + s3_endpoint_url = ( + optional_params.get("s3_endpoint_url") + or f"https://s3.{aws_region_name}.amazonaws.com" + ) + return f"{s3_endpoint_url}/{bucket_name}/{object_name}" def get_supported_openai_params( @@ -201,7 +206,6 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): ) -> dict: return optional_params - # Providers whose InvokeModel body uses the Converse API format # (messages + inferenceConfig + image blocks). Nova is the primary # example; add others here as they adopt the same schema. @@ -286,24 +290,24 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): ) -> List[Dict[str, Any]]: """ Transforms OpenAI JSONL content to Bedrock batch format - + Bedrock batch format: { "recordId": "alphanumeric string", "modelInput": {JSON body} } Example: { - "recordId": "CALL0000001", + "recordId": "CALL0000001", "modelInput": { - "anthropic_version": "bedrock-2023-05-31", + "anthropic_version": "bedrock-2023-05-31", "max_tokens": 1024, - "messages": [ - { - "role": "user", + "messages": [ + { + "role": "user", "content": [{"type": "text", "text": "Hello"}] } ] } } """ - + bedrock_jsonl_content = [] for idx, _openai_jsonl_content in enumerate(openai_jsonl_content): # Extract the request body from OpenAI format @@ -312,28 +316,28 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): try: model, _, _, _ = get_llm_provider( - model=model, - custom_llm_provider=None, - ) + model=model, + custom_llm_provider=None, + ) except Exception as e: - verbose_logger.exception(f"litellm.llms.bedrock.files.transformation.py::_transform_openai_jsonl_content_to_bedrock_jsonl_content() - Error inferring custom_llm_provider - {str(e)}") - + verbose_logger.exception( + f"litellm.llms.bedrock.files.transformation.py::_transform_openai_jsonl_content_to_bedrock_jsonl_content() - Error inferring custom_llm_provider - {str(e)}" + ) + # Determine provider from model name provider = self.get_bedrock_invoke_provider(model) - + # Transform to Bedrock modelInput format model_input = self._map_openai_to_bedrock_params( - openai_request_body=openai_body, - provider=provider + openai_request_body=openai_body, provider=provider ) - + # Create Bedrock batch record - record_id = _openai_jsonl_content.get("custom_id", f"CALL{str(idx).zfill(7)}") - bedrock_record = { - "recordId": record_id, - "modelInput": model_input - } - + record_id = _openai_jsonl_content.get( + "custom_id", f"CALL{str(idx).zfill(7)}" + ) + bedrock_record = {"recordId": record_id, "modelInput": model_input} + bedrock_jsonl_content.append(bedrock_record) return bedrock_jsonl_content @@ -353,10 +357,10 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): raise ValueError("file is required") extracted_file_data = extract_file_data(file_data) extracted_file_data_content = extracted_file_data.get("content") - + if extracted_file_data_content is None: raise ValueError("file content is required") - + # Get and transform the file content if FilesAPIUtils.is_batch_jsonl_file( create_file_data=create_file_data, @@ -367,7 +371,9 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): extracted_file_data_content ) openai_jsonl_content = [ - json.loads(line) for line in original_file_content.splitlines() if line.strip() + json.loads(line) + for line in original_file_content.splitlines() + if line.strip() ] bedrock_jsonl_content = ( self._transform_openai_jsonl_content_to_bedrock_jsonl_content( @@ -376,12 +382,12 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): ) file_content = "\n".join(json.dumps(item) for item in bedrock_jsonl_content) elif isinstance(extracted_file_data_content, bytes): - file_content = extracted_file_data_content.decode('utf-8') + file_content = extracted_file_data_content.decode("utf-8") elif isinstance(extracted_file_data_content, str): file_content = extracted_file_data_content else: raise ValueError("Unsupported file content type") - + # Get the S3 URL for upload api_base = self.get_complete_file_url( api_base=None, @@ -391,7 +397,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): litellm_params=litellm_params, data=create_file_data, ) - + # Sign the request and return a pre-signed request object signed_headers, signed_body = self._sign_s3_request( content=file_content, @@ -400,7 +406,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): ) litellm_params["upload_url"] = api_base - + # Return a dict that tells the HTTP handler exactly what to do return { "method": "PUT", @@ -443,7 +449,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): aws_web_identity_token=optional_params.get("aws_web_identity_token"), aws_sts_endpoint=optional_params.get("aws_sts_endpoint"), ) - + # Calculate SHA256 hash of the content (REQUIRED for S3) content_hash = hashlib.sha256(content.encode("utf-8")).hexdigest() @@ -466,33 +472,33 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): data=prepped.body, headers=prepped.headers, ) - + # Get region name for non-LLM API calls (same as s3_v2.py) signing_region = self.get_aws_region_name_for_non_llm_api_calls( aws_region_name=aws_region_name ) - + SigV4Auth(credentials, "s3", signing_region).add_auth(aws_request) # Return signed headers and body signed_body = aws_request.body if isinstance(signed_body, bytes): - signed_body = signed_body.decode('utf-8') + signed_body = signed_body.decode("utf-8") elif signed_body is None: signed_body = content # Fallback to original content - + return dict(aws_request.headers), signed_body def _convert_https_url_to_s3_uri(self, https_url: str) -> tuple[str, str]: """ Convert HTTPS S3 URL to s3:// URI format. - + Args: https_url: HTTPS S3 URL (e.g., "https://s3.us-west-2.amazonaws.com/bucket/key") - + Returns: Tuple of (s3_uri, filename) - + Example: Input: "https://s3.us-west-2.amazonaws.com/litellm-proxy/file.jsonl" Output: ("s3://litellm-proxy/file.jsonl", "file.jsonl") @@ -502,13 +508,13 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): # Match HTTPS S3 URL patterns # Pattern 1: https://s3.region.amazonaws.com/bucket/key # Pattern 2: https://bucket.s3.region.amazonaws.com/key - + pattern1 = r"https://s3\.([^.]+)\.amazonaws\.com/([^/]+)/(.+)" pattern2 = r"https://([^.]+)\.s3\.([^.]+)\.amazonaws\.com/(.+)" - + match1 = re.match(pattern1, https_url) match2 = re.match(pattern2, https_url) - + if match1: # Pattern: https://s3.region.amazonaws.com/bucket/key region, bucket, key = match1.groups() @@ -520,17 +526,18 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): else: # Fallback: try to extract bucket and key from URL path from urllib.parse import urlparse + parsed = urlparse(https_url) - path_parts = parsed.path.lstrip('/').split('/', 1) + path_parts = parsed.path.lstrip("/").split("/", 1) if len(path_parts) >= 2: bucket, key = path_parts[0], path_parts[1] s3_uri = f"s3://{bucket}/{key}" else: raise ValueError(f"Unable to parse S3 URL: {https_url}") - + # Extract filename from key filename = key.split("/")[-1] if "/" in key else key - + return s3_uri, filename def transform_create_file_response( @@ -548,7 +555,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): # Extract S3 object information from the response # S3 PUT object returns ETag and other metadata in headers content_length = response_headers.get("Content-Length", "0") - + # Use the actual upload URL that was used for the S3 upload upload_url = litellm_params.get("upload_url") file_id: str = "" @@ -628,7 +635,9 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): optional_params: dict, litellm_params: dict, ) -> tuple[str, dict]: - raise NotImplementedError("BedrockFilesConfig does not support file content retrieval") + raise NotImplementedError( + "BedrockFilesConfig does not support file content retrieval" + ) def transform_file_content_response( self, @@ -636,7 +645,9 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): logging_obj: LiteLLMLoggingObj, litellm_params: dict, ) -> HttpxBinaryResponseContent: - raise NotImplementedError("BedrockFilesConfig does not support file content retrieval") + raise NotImplementedError( + "BedrockFilesConfig does not support file content retrieval" + ) class BedrockJsonlFilesTransformation: @@ -680,7 +691,9 @@ class BedrockJsonlFilesTransformation: Delegate to the main BedrockFilesConfig transformation method """ config = BedrockFilesConfig() - return config._transform_openai_jsonl_content_to_bedrock_jsonl_content(openai_jsonl_content) + return config._transform_openai_jsonl_content_to_bedrock_jsonl_content( + openai_jsonl_content + ) def _get_s3_object_name( self, @@ -698,8 +711,6 @@ class BedrockJsonlFilesTransformation: object_name = f"litellm-bedrock-files-{_model}-{uuid.uuid4()}.jsonl" return object_name - - def _get_content_from_openai_file(self, openai_file_content: FileTypes) -> str: """ Helper to extract content from various OpenAI file types and return as string. @@ -746,10 +757,10 @@ class BedrockJsonlFilesTransformation: # S3 response typically contains ETag, key, etc. object_key = s3_upload_response.get("Key", "") bucket_name = s3_upload_response.get("Bucket", "") - + # Extract filename from object key filename = object_key.split("/")[-1] if "/" in object_key else object_key - + return OpenAIFileObject( purpose=create_file_data.get("purpose", "batch"), id=f"s3://{bucket_name}/{object_key}", diff --git a/litellm/llms/bedrock/image_edit/__init__.py b/litellm/llms/bedrock/image_edit/__init__.py index f3a0e61067..ea6d13a676 100644 --- a/litellm/llms/bedrock/image_edit/__init__.py +++ b/litellm/llms/bedrock/image_edit/__init__.py @@ -7,4 +7,3 @@ Handles image edit operations for Bedrock stability models. from .handler import BedrockImageEdit __all__ = ["BedrockImageEdit"] - diff --git a/litellm/llms/bedrock/image_edit/handler.py b/litellm/llms/bedrock/image_edit/handler.py index ef441fa503..867944f879 100644 --- a/litellm/llms/bedrock/image_edit/handler.py +++ b/litellm/llms/bedrock/image_edit/handler.py @@ -307,4 +307,3 @@ class BedrockImageEdit(BaseAWSLLM): ) return model_response - diff --git a/litellm/llms/bedrock/image_edit/stability_transformation.py b/litellm/llms/bedrock/image_edit/stability_transformation.py index db4e3a0a7a..6a8b95e7e3 100644 --- a/litellm/llms/bedrock/image_edit/stability_transformation.py +++ b/litellm/llms/bedrock/image_edit/stability_transformation.py @@ -54,7 +54,7 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): def _is_stability_edit_model(cls, model: Optional[str] = None) -> bool: """ Returns True if the model is a Bedrock Stability edit model. - + Bedrock Stability edit models follow this pattern: stability.stable-conservative-upscale-v1:0 stability.stable-creative-upscale-v1:0 @@ -66,25 +66,25 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): """ if model: model_lower = model.lower() - if "stability." in model_lower and any([ - "upscale" in model_lower, - "outpaint" in model_lower, - "inpaint" in model_lower, - "erase" in model_lower, - "remove-background" in model_lower, - "search-recolor" in model_lower, - "search-replace" in model_lower, - "control-sketch" in model_lower, - "control-structure" in model_lower, - "style-guide" in model_lower, - "style-transfer" in model_lower, - ]): + if "stability." in model_lower and any( + [ + "upscale" in model_lower, + "outpaint" in model_lower, + "inpaint" in model_lower, + "erase" in model_lower, + "remove-background" in model_lower, + "search-recolor" in model_lower, + "search-replace" in model_lower, + "control-sketch" in model_lower, + "control-structure" in model_lower, + "style-guide" in model_lower, + "style-transfer" in model_lower, + ] + ): return True return False - def get_supported_openai_params( - self, model: str - ) -> list: + def get_supported_openai_params(self, model: str) -> list: """ Return list of OpenAI params supported by Bedrock Stability. """ @@ -149,7 +149,7 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): return mapped_params - def transform_image_edit_request( #noqa: PLR0915 + def transform_image_edit_request( # noqa: PLR0915 self, model: str, prompt: Optional[str], @@ -167,27 +167,27 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): data: Dict[str, Any] = { "output_format": "png", # Default to PNG } - + # Add prompt only if provided (some models don't require it) if prompt is not None and prompt != "": data["prompt"] = prompt - + # Convert image to base64 if provided if image is not None: image_b64: str - if hasattr(image, 'read') and callable(getattr(image, 'read', None)): + if hasattr(image, "read") and callable(getattr(image, "read", None)): # File-like object (e.g., BufferedReader from open()) image_bytes = image.read() # type: ignore - image_b64 = base64.b64encode(image_bytes).decode('utf-8') # type: ignore + image_b64 = base64.b64encode(image_bytes).decode("utf-8") # type: ignore elif isinstance(image, bytes): # Raw bytes - image_b64 = base64.b64encode(image).decode('utf-8') + image_b64 = base64.b64encode(image).decode("utf-8") elif isinstance(image, str): # Already a base64 string image_b64 = image else: # Try to handle as bytes - image_b64 = base64.b64encode(bytes(image)).decode('utf-8') # type: ignore + image_b64 = base64.b64encode(bytes(image)).decode("utf-8") # type: ignore # For style-transfer models, map image to init_image model_lower = model.lower() @@ -208,8 +208,10 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): file_value = value if isinstance(value, list) and len(value) > 0: file_value = value[0] - - if hasattr(file_value, 'read') and callable(getattr(file_value, 'read', None)): + + if hasattr(file_value, "read") and callable( + getattr(file_value, "read", None) + ): file_bytes = file_value.read() # type: ignore elif isinstance(file_value, bytes): file_bytes = file_value @@ -219,14 +221,14 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): continue else: file_bytes = file_value # type: ignore - + if isinstance(file_bytes, bytes): - file_b64 = base64.b64encode(file_bytes).decode('utf-8') + file_b64 = base64.b64encode(file_bytes).decode("utf-8") else: file_b64 = str(file_bytes) data[key] = file_b64 continue - + # Numeric fields that need to be converted to int/float numeric_int_fields = ["left", "right", "up", "down", "seed"] numeric_float_fields = [ @@ -239,7 +241,7 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): "style_strength", "change_strength", ] - + if key in numeric_int_fields: # Convert to int (these are pixel values for outpaint) try: @@ -329,13 +331,15 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): model_response._hidden_params = {} if "additional_headers" not in model_response._hidden_params: model_response._hidden_params["additional_headers"] = {} - + # Set cost based on model model_info = get_model_info(model, custom_llm_provider="bedrock") cost_per_image = model_info.get("output_cost_per_image", 0) if cost_per_image is not None: - model_response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = float(cost_per_image) - + model_response._hidden_params["additional_headers"][ + "llm_provider-x-litellm-response-cost" + ] = float(cost_per_image) + return model_response def use_multipart_form_data(self) -> bool: @@ -352,11 +356,11 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): ) -> str: """ Get the complete URL for the Bedrock Image Edit API. - + For Bedrock, this is handled by the handler which constructs the endpoint URL based on the model ID and AWS region. This method is required by the base class but the actual URL construction happens in BedrockImageEdit.image_edit(). - + Returns a placeholder - the real endpoint is constructed in the handler. """ # Bedrock URLs are constructed in the handler using boto3 @@ -371,25 +375,25 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): ) -> dict: """ Validate environment for Bedrock Stability image edit. - + For Bedrock, AWS credentials are managed by the BaseAWSLLM class. This method validates that headers are properly set up. - + Args: headers: The request headers to validate/update model: The model name being used api_key: Optional API key (not used for Bedrock, which uses AWS credentials) - + Returns: Updated headers dict """ if headers is None: headers = {} - + # Bedrock uses AWS credentials, not API keys # Headers are set up by the handler's get_request_headers() method # This just ensures basic headers are present if "Content-Type" not in headers: headers["Content-Type"] = "application/json" - + return headers diff --git a/litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py b/litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py index 1836699958..86c005bbfa 100644 --- a/litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py +++ b/litellm/llms/bedrock/image_generation/amazon_nova_canvas_transformation.py @@ -217,4 +217,4 @@ class AmazonNovaCanvasConfig: num_images: int = 0 if image_response.data: num_images = len(image_response.data) - return output_cost_per_image * num_images \ No newline at end of file + return output_cost_per_image * num_images diff --git a/litellm/llms/bedrock/image_generation/amazon_stability1_transformation.py b/litellm/llms/bedrock/image_generation/amazon_stability1_transformation.py index 07f82cec23..1d88aaf35f 100644 --- a/litellm/llms/bedrock/image_generation/amazon_stability1_transformation.py +++ b/litellm/llms/bedrock/image_generation/amazon_stability1_transformation.py @@ -115,7 +115,7 @@ class AmazonStabilityConfig: return { "text_prompts": [{"text": prompt, "weight": 1}], - **inference_params, + **inference_params, } @classmethod @@ -161,4 +161,4 @@ class AmazonStabilityConfig: num_images: int = 0 if image_response.data: num_images = len(image_response.data) - return output_cost_per_image * num_images \ No newline at end of file + return output_cost_per_image * num_images diff --git a/litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py b/litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py index 160d0af8e8..8aff24fe9a 100644 --- a/litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py +++ b/litellm/llms/bedrock/image_generation/amazon_stability3_transformation.py @@ -109,11 +109,11 @@ class AmazonStability3Config: @classmethod def cost_calculator( - cls, - model: str, - image_response: ImageResponse, - size: Optional[str] = None, - optional_params: Optional[dict] = None, + cls, + model: str, + image_response: ImageResponse, + size: Optional[str] = None, + optional_params: Optional[dict] = None, ) -> float: get_model_info = get_cached_model_info() model_info = get_model_info( diff --git a/litellm/llms/bedrock/image_generation/image_handler.py b/litellm/llms/bedrock/image_generation/image_handler.py index 7270b96ab8..d6053278cb 100644 --- a/litellm/llms/bedrock/image_generation/image_handler.py +++ b/litellm/llms/bedrock/image_generation/image_handler.py @@ -180,12 +180,12 @@ class BedrockImageGeneration(BaseAWSLLM): headers = {} guardrail_identifier = optional_params.pop("guardrailIdentifier", None) guardrail_version = optional_params.pop("guardrailVersion", None) - + if guardrail_identifier is not None: headers["x-amz-bedrock-guardrail-identifier"] = guardrail_identifier if guardrail_version is not None: headers["x-amz-bedrock-guardrail-version"] = guardrail_version - + return headers def _prepare_request( @@ -292,7 +292,9 @@ class BedrockImageGeneration(BaseAWSLLM): dict: The request body to use for the Bedrock Image Generation API """ config_class = self.get_config_class(model=model) - request_body = config_class.transform_request_body(text=prompt, optional_params=optional_params) + request_body = config_class.transform_request_body( + text=prompt, optional_params=optional_params + ) return dict(request_body) def _transform_response_dict_to_openai_response( diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index b11215e7f6..e31820d763 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -276,7 +276,7 @@ class AmazonAnthropicClaudeMessagesConfig( "opus_4.6", "opus-4-6", "opus_4_6", - #sonnet 4.6 + # sonnet 4.6 "sonnet-4.6", "sonnet_4.6", "sonnet-4-6", @@ -462,7 +462,7 @@ class AmazonAnthropicClaudeMessagesConfig( if "tool-search-tool-2025-10-19" in beta_set: beta_set.add("tool-examples-2025-10-29") - + if beta_set: anthropic_messages_request["anthropic_beta"] = list(beta_set) diff --git a/litellm/llms/bedrock/passthrough/transformation.py b/litellm/llms/bedrock/passthrough/transformation.py index 5efd3ba1d9..274b0282ac 100644 --- a/litellm/llms/bedrock/passthrough/transformation.py +++ b/litellm/llms/bedrock/passthrough/transformation.py @@ -27,28 +27,30 @@ class BedrockPassthroughConfig( def _encode_model_id_for_endpoint(self, model_id: str) -> str: """ Encode model_id (especially ARNs) for use in Bedrock endpoints. - + ARNs contain special characters like colons and slashes that need to be properly URL-encoded when used in HTTP request paths. For example: arn:aws:bedrock:us-east-1:123:application-inference-profile/abc123 becomes: arn:aws:bedrock:us-east-1:123:application-inference-profile%2Fabc123 - + Args: model_id: The model ID or ARN to encode - + Returns: The encoded model_id suitable for use in endpoint URLs """ from litellm.passthrough.utils import CommonUtils import re - + # Create a temporary endpoint with the model_id to check if encoding is needed temp_endpoint = f"/model/{model_id}/converse" - encoded_temp_endpoint = CommonUtils.encode_bedrock_runtime_modelid_arn(temp_endpoint) - + encoded_temp_endpoint = CommonUtils.encode_bedrock_runtime_modelid_arn( + temp_endpoint + ) + # Extract the encoded model_id from the temporary endpoint - encoded_model_id_match = re.search(r'/model/([^/]+)/', encoded_temp_endpoint) + encoded_model_id_match = re.search(r"/model/([^/]+)/", encoded_temp_endpoint) if encoded_model_id_match: return encoded_model_id_match.group(1) else: @@ -73,7 +75,9 @@ class BedrockPassthroughConfig( model_id=model_id, ) - aws_bedrock_runtime_endpoint = optional_params.get("aws_bedrock_runtime_endpoint") + aws_bedrock_runtime_endpoint = optional_params.get( + "aws_bedrock_runtime_endpoint" + ) endpoint_url, _ = self.get_runtime_endpoint( api_base=api_base, aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, @@ -85,13 +89,16 @@ class BedrockPassthroughConfig( # instead of the translated model name if model_id is not None: import re - + # Encode the model_id if it's an ARN to properly handle special characters encoded_model_id = self._encode_model_id_for_endpoint(model_id) - + # Replace the model name in the endpoint with the encoded model_id - endpoint = re.sub(r'model/[^/]+/', f'model/{encoded_model_id}/', endpoint) - return self.format_url(endpoint, endpoint_url, request_query_params or {}), endpoint_url + endpoint = re.sub(r"model/[^/]+/", f"model/{encoded_model_id}/", endpoint) + return ( + self.format_url(endpoint, endpoint_url, request_query_params or {}), + endpoint_url, + ) def sign_request( self, diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 9b6a80f4a2..cde9f3e6fc 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -97,8 +97,10 @@ class BedrockRealtime(BaseAWSLLM): try: # Initialize the bidirectional stream - bedrock_stream = await bedrock_client.invoke_model_with_bidirectional_stream( - InvokeModelWithBidirectionalStreamOperationInput(model_id=model) + bedrock_stream = ( + await bedrock_client.invoke_model_with_bidirectional_stream( + InvokeModelWithBidirectionalStreamOperationInput(model_id=model) + ) ) verbose_proxy_logger.debug( @@ -232,7 +234,7 @@ class BedrockRealtime(BaseAWSLLM): # Transform Bedrock format to OpenAI format from litellm.types.realtime import RealtimeResponseTransformInput - + realtime_response_transform_input: RealtimeResponseTransformInput = { "current_output_item_id": session_state.get( "current_output_item_id" @@ -251,13 +253,11 @@ class BedrockRealtime(BaseAWSLLM): ), } - transformed_response = ( - transformation_config.transform_realtime_response( - message=bedrock_response, - model=model, - logging_obj=logging_obj, - realtime_response_transform_input=realtime_response_transform_input, - ) + transformed_response = transformation_config.transform_realtime_response( + message=bedrock_response, + model=model, + logging_obj=logging_obj, + realtime_response_transform_input=realtime_response_transform_input, ) # Update session state diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index 1dde1b47fe..13d5bf3546 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -43,13 +43,13 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): self.prompt_name = str(uuid_lib.uuid4()) self.content_name = str(uuid_lib.uuid4()) self.audio_content_name = str(uuid_lib.uuid4()) - + # Default configuration values # Inference configuration self.max_tokens = 1024 self.top_p = 0.9 self.temperature = 0.7 - + # Audio output configuration self.output_sample_rate_hertz = 24000 self.output_sample_size_bits = 16 @@ -58,7 +58,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): self.output_encoding = "base64" self.output_audio_type = "SPEECH" self.output_media_type = "audio/lpcm" - + # Audio input configuration self.input_sample_rate_hertz = 16000 self.input_sample_size_bits = 16 @@ -66,7 +66,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): self.input_encoding = "base64" self.input_audio_type = "SPEECH" self.input_media_type = "audio/lpcm" - + # Text configuration self.text_media_type = "text/plain" @@ -86,7 +86,9 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): """Bedrock requires session configuration.""" return True - def session_configuration_request(self, model: str, tools: Optional[List[dict]] = None) -> str: + def session_configuration_request( + self, model: str, tools: Optional[List[dict]] = None + ) -> str: """ Create initial session configuration for Bedrock Nova Sonic. @@ -158,20 +160,22 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): "description": function.get("description", ""), "inputSchema": { "json": json.dumps(function.get("parameters", {})) - } + }, } } bedrock_tools.append(bedrock_tool) return bedrock_tools - def _map_audio_format_to_sample_rate(self, audio_format: str, is_output: bool = True) -> int: + def _map_audio_format_to_sample_rate( + self, audio_format: str, is_output: bool = True + ) -> int: """ Map OpenAI audio format to sample rate. - + Args: audio_format: OpenAI audio format (pcm16, g711_ulaw, g711_alaw) is_output: Whether this is for output (True) or input (False) - + Returns: Sample rate in Hz """ @@ -195,15 +199,15 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): """ verbose_logger.debug("Handling session.update") messages: List[str] = [] - + session_config = json_message.get("session", {}) - + # Update inference configuration from session if provided if "max_response_output_tokens" in session_config: self.max_tokens = session_config["max_response_output_tokens"] if "temperature" in session_config: self.temperature = session_config["temperature"] - + # Update audio output configuration from session if provided if "voice" in session_config: self.voice_id = session_config["voice"] @@ -212,14 +216,14 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): self.output_sample_rate_hertz = self._map_audio_format_to_sample_rate( output_format, is_output=True ) - + # Update audio input configuration from session if provided if "input_audio_format" in session_config: input_format = session_config["input_audio_format"] self.input_sample_rate_hertz = self._map_audio_format_to_sample_rate( input_format, is_output=False ) - + # Allow direct override of sample rates if provided (custom extension) if "output_sample_rate_hertz" in session_config: self.output_sample_rate_hertz = session_config["output_sample_rate_hertz"] @@ -313,7 +317,9 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): return messages - def transform_input_audio_buffer_append_event(self, json_message: dict) -> List[str]: + def transform_input_audio_buffer_append_event( + self, json_message: dict + ) -> List[str]: """ Transform input_audio_buffer.append event to Bedrock audio input. @@ -365,7 +371,9 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): return messages - def transform_input_audio_buffer_commit_event(self, json_message: dict) -> List[str]: + def transform_input_audio_buffer_commit_event( + self, json_message: dict + ) -> List[str]: """ Transform input_audio_buffer.commit event to Bedrock audio content end. @@ -410,7 +418,9 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): # Handle tool result if item_type == "function_call_output": - return self.transform_conversation_item_create_tool_result_event(json_message) + return self.transform_conversation_item_create_tool_result_event( + json_message + ) # Handle regular message if item_type == "message": @@ -549,14 +559,14 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): OpenAI session.created event """ verbose_logger.debug("Handling sessionStart") - + session = OpenAIRealtimeStreamSession( id=logging_obj.litellm_trace_id, modalities=["text", "audio"], ) if model is not None and isinstance(model, str): session["model"] = model - + return OpenAIRealtimeStreamSessionEvents( type="session.created", session=session, @@ -592,7 +602,13 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): role = content_start.get("role") if role != "ASSISTANT": - return [], current_response_id, current_output_item_id, current_conversation_id, None + return ( + [], + current_response_id, + current_output_item_id, + current_conversation_id, + None, + ) verbose_logger.debug("Handling ASSISTANT contentStart") @@ -606,7 +622,9 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): # Determine content type content_type = content_start.get("type", "TEXT") - current_delta_type: ALL_DELTA_TYPES = "text" if content_type == "TEXT" else "audio" + current_delta_type: ALL_DELTA_TYPES = ( + "text" if content_type == "TEXT" else "audio" + ) returned_messages: List[OpenAIRealtimeEvents] = [] @@ -850,7 +868,12 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): event: dict, current_response_id: Optional[str], current_conversation_id: Optional[str], - ) -> tuple[List[OpenAIRealtimeEvents], Optional[str], Optional[str], Optional[ALL_DELTA_TYPES]]: + ) -> tuple[ + List[OpenAIRealtimeEvents], + Optional[str], + Optional[str], + Optional[ALL_DELTA_TYPES], + ]: """ Transform Bedrock promptEnd event to OpenAI response.done. @@ -915,7 +938,11 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): tool_input = {} if "input" in tool_use: try: - tool_input = json.loads(tool_use["input"]) if isinstance(tool_use["input"], str) else tool_use["input"] + tool_input = ( + json.loads(tool_use["input"]) + if isinstance(tool_use["input"], str) + else tool_use["input"] + ) except json.JSONDecodeError: tool_input = {} @@ -925,6 +952,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): # Create a function call arguments done event # This is a custom event format that matches what clients expect from typing import cast + function_call_event: dict[str, Any] = { "type": "response.function_call_arguments.done", "event_id": f"event_{uuid.uuid4()}", @@ -936,9 +964,15 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): "arguments": json.dumps(tool_input), } - return [cast(OpenAIRealtimeEvents, function_call_event)], tool_call_id, tool_name + return ( + [cast(OpenAIRealtimeEvents, function_call_event)], + tool_call_id, + tool_name, + ) - def transform_conversation_item_create_tool_result_event(self, json_message: dict) -> List[str]: + def transform_conversation_item_create_tool_result_event( + self, json_message: dict + ) -> List[str]: """ Transform conversation.item.create with tool result to Bedrock format. @@ -969,10 +1003,8 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): "toolResultInputConfiguration": { "toolUseId": call_id, "type": "TEXT", - "textInputConfiguration": { - "mediaType": "text/plain" - } - } + "textInputConfiguration": {"mediaType": "text/plain"}, + }, } } } @@ -984,7 +1016,9 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): "toolResult": { "promptName": self.prompt_name, "contentName": tool_content_name, - "content": output if isinstance(output, str) else json.dumps(output) + "content": output + if isinstance(output, str) + else json.dumps(output), } } } @@ -1025,7 +1059,11 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): try: json_message = json.loads(message) except json.JSONDecodeError: - message_preview = message[:200].decode('utf-8', errors='replace') if isinstance(message, bytes) else message[:200] + message_preview = ( + message[:200].decode("utf-8", errors="replace") + if isinstance(message, bytes) + else message[:200] + ) verbose_logger.warning(f"Invalid JSON message: {message_preview}") return { "response": [], diff --git a/litellm/llms/bedrock/rerank/handler.py b/litellm/llms/bedrock/rerank/handler.py index 37167e7c33..812ca116c2 100644 --- a/litellm/llms/bedrock/rerank/handler.py +++ b/litellm/llms/bedrock/rerank/handler.py @@ -35,7 +35,12 @@ class BedrockRerankHandler(BaseAWSLLM): if client is None: client = get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK) try: - response = await client.post(url=prepared_request["endpoint_url"], headers=dict(prepared_request["prepped"].headers), data=prepared_request["body"], timeout=timeout) + response = await client.post( + url=prepared_request["endpoint_url"], + headers=dict(prepared_request["prepped"].headers), + data=prepared_request["body"], + timeout=timeout, + ) response.raise_for_status() except httpx.HTTPStatusError as err: error_code = err.response.status_code @@ -96,7 +101,12 @@ class BedrockRerankHandler(BaseAWSLLM): if client is None or not isinstance(client, HTTPHandler): client = _get_httpx_client() try: - response = client.post(url=prepared_request["endpoint_url"], headers=dict(prepared_request["prepped"].headers), data=prepared_request["body"], timeout=timeout) + response = client.post( + url=prepared_request["endpoint_url"], + headers=dict(prepared_request["prepped"].headers), + data=prepared_request["body"], + timeout=timeout, + ) response.raise_for_status() except httpx.HTTPStatusError as err: error_code = err.response.status_code diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index 72e1e1470d..4da0a7c779 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -152,7 +152,6 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): if param == "max_num_results": optional_params["numberOfResults"] = value elif param == "filters" and value is not None: - # map the openai filters to the aws kb filters format # openai filters = {"key": , "value": , "operator": } OR {"and" | "or": [{"key": , "value": , "operator": }]} # aws kb filters = {"operator": {"": }} OR {"andAll | orAll": [{"operator": {"": }}]} diff --git a/litellm/llms/black_forest_labs/image_edit/handler.py b/litellm/llms/black_forest_labs/image_edit/handler.py index 44a102ec48..dea2683a04 100644 --- a/litellm/llms/black_forest_labs/image_edit/handler.py +++ b/litellm/llms/black_forest_labs/image_edit/handler.py @@ -356,7 +356,12 @@ class BlackForestLabsImageEdit: if status == "Ready": return response - elif status in ["Error", "Failed", "Content Moderated", "Request Moderated"]: + elif status in [ + "Error", + "Failed", + "Content Moderated", + "Request Moderated", + ]: raise BlackForestLabsError( status_code=400, message=f"Image generation failed: {status}", @@ -436,7 +441,12 @@ class BlackForestLabsImageEdit: if status == "Ready": return response - elif status in ["Error", "Failed", "Content Moderated", "Request Moderated"]: + elif status in [ + "Error", + "Failed", + "Content Moderated", + "Request Moderated", + ]: raise BlackForestLabsError( status_code=400, message=f"Image generation failed: {status}", diff --git a/litellm/llms/black_forest_labs/image_edit/transformation.py b/litellm/llms/black_forest_labs/image_edit/transformation.py index 78898345bf..610fee1889 100644 --- a/litellm/llms/black_forest_labs/image_edit/transformation.py +++ b/litellm/llms/black_forest_labs/image_edit/transformation.py @@ -179,11 +179,7 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): """ Get the complete URL for the Black Forest Labs API request. """ - base_url: str = ( - api_base - or get_secret_str("BFL_API_BASE") - or DEFAULT_API_BASE - ) + base_url: str = api_base or get_secret_str("BFL_API_BASE") or DEFAULT_API_BASE base_url = base_url.rstrip("/") endpoint = self._get_model_endpoint(model) @@ -247,9 +243,18 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): # Add optional params (only BFL-recognized parameters) bfl_request_params = [ - "seed", "output_format", "safety_tolerance", "prompt_upsampling", - "aspect_ratio", "steps", "guidance", "grow_mask", - "top", "bottom", "left", "right", + "seed", + "output_format", + "safety_tolerance", + "prompt_upsampling", + "aspect_ratio", + "steps", + "guidance", + "grow_mask", + "top", + "bottom", + "left", + "right", ] for key, value in image_edit_optional_request_params.items(): if key in bfl_request_params and value is not None: diff --git a/litellm/llms/black_forest_labs/image_generation/handler.py b/litellm/llms/black_forest_labs/image_generation/handler.py index 99dc2feca3..5a1d885e52 100644 --- a/litellm/llms/black_forest_labs/image_generation/handler.py +++ b/litellm/llms/black_forest_labs/image_generation/handler.py @@ -342,7 +342,12 @@ class BlackForestLabsImageGeneration: if status == "Ready": return response - elif status in ["Error", "Failed", "Content Moderated", "Request Moderated"]: + elif status in [ + "Error", + "Failed", + "Content Moderated", + "Request Moderated", + ]: raise BlackForestLabsError( status_code=400, message=f"Image generation failed: {status}", @@ -422,7 +427,12 @@ class BlackForestLabsImageGeneration: if status == "Ready": return response - elif status in ["Error", "Failed", "Content Moderated", "Request Moderated"]: + elif status in [ + "Error", + "Failed", + "Content Moderated", + "Request Moderated", + ]: raise BlackForestLabsError( status_code=400, message=f"Image generation failed: {status}", diff --git a/litellm/llms/black_forest_labs/image_generation/transformation.py b/litellm/llms/black_forest_labs/image_generation/transformation.py index fd664b3ea7..a6ed77f535 100644 --- a/litellm/llms/black_forest_labs/image_generation/transformation.py +++ b/litellm/llms/black_forest_labs/image_generation/transformation.py @@ -203,9 +203,7 @@ class BlackForestLabsImageGenerationConfig(BaseImageGenerationConfig): """ Get the complete URL for the Black Forest Labs API request. """ - base_url: str = ( - api_base or get_secret_str("BFL_API_BASE") or DEFAULT_API_BASE - ) + base_url: str = api_base or get_secret_str("BFL_API_BASE") or DEFAULT_API_BASE base_url = base_url.rstrip("/") endpoint = self._get_model_endpoint(model) diff --git a/litellm/llms/bytez/chat/transformation.py b/litellm/llms/bytez/chat/transformation.py index ccd3c21645..a72f732a30 100644 --- a/litellm/llms/bytez/chat/transformation.py +++ b/litellm/llms/bytez/chat/transformation.py @@ -91,13 +91,11 @@ class BytezChatConfig(BaseConfig): model: str, drop_params: bool, ) -> dict: - adapted_params = {} all_params = {**non_default_params, **optional_params} for key, value in all_params.items(): - alias = self.openai_to_bytez_param_map.get(key) if alias is False: @@ -124,7 +122,6 @@ class BytezChatConfig(BaseConfig): api_key: Optional[str] = None, api_base: Optional[str] = None, ) -> dict: - headers.update( { "content-type": "application/json", @@ -141,7 +138,6 @@ class BytezChatConfig(BaseConfig): if not api_key: raise Exception("Missing api_key, make sure you pass in your api key") - return headers def get_complete_url( @@ -193,7 +189,6 @@ class BytezChatConfig(BaseConfig): api_key: Optional[str] = None, json_mode: Optional[bool] = None, ) -> ModelResponse: - json = raw_response.json() # noqa: F811 error = json.get("error") @@ -387,13 +382,11 @@ open_ai_to_bytez_content_item_map = { def adapt_messages_to_bytez_standard(messages: List[Dict]): - messages = _adapt_string_only_content_to_lists(messages) new_messages = [] for message in messages: - role = message["role"] content: list = message["content"] @@ -433,7 +426,6 @@ def _adapt_string_only_content_to_lists(messages: List[Dict]): new_messages = [] for message in messages: - role = message.get("role") content = message.get("content") @@ -446,7 +438,6 @@ def _adapt_string_only_content_to_lists(messages: List[Dict]): new_content.append(content) elif isinstance(content, list): - new_content_items = [] for content_item in content: if isinstance(content_item, str): diff --git a/litellm/llms/bytez/common_utils.py b/litellm/llms/bytez/common_utils.py index 2fedd2aad0..d6593a06b7 100644 --- a/litellm/llms/bytez/common_utils.py +++ b/litellm/llms/bytez/common_utils.py @@ -22,4 +22,4 @@ class BytezError(BaseLLMException): status_code=status_code, message=message, headers=headers, - ) \ No newline at end of file + ) diff --git a/litellm/llms/chatgpt/authenticator.py b/litellm/llms/chatgpt/authenticator.py index ff053730c3..e35b04a3fb 100644 --- a/litellm/llms/chatgpt/authenticator.py +++ b/litellm/llms/chatgpt/authenticator.py @@ -206,7 +206,9 @@ class Authenticator: "interval": str(interval or "5"), } - def _poll_for_authorization_code(self, device_code: Dict[str, str]) -> Dict[str, str]: + def _poll_for_authorization_code( + self, device_code: Dict[str, str] + ) -> Dict[str, str]: client = _get_httpx_client() interval = int(device_code.get("interval", "5")) start_time = time.time() @@ -284,7 +286,9 @@ class Authenticator: status_code=400, ) - if not all(key in data for key in ("access_token", "refresh_token", "id_token")): + if not all( + key in data for key in ("access_token", "refresh_token", "id_token") + ): raise GetAccessTokenError( message=f"Token exchange response missing fields: {data}", status_code=400, @@ -377,11 +381,11 @@ class Authenticator: auth_data = self._read_auth_file() if auth_data: access_token = auth_data.get("access_token") - if access_token and not self._is_token_expired( - auth_data, access_token - ): + if access_token and not self._is_token_expired(auth_data, access_token): return access_token - sleep_for = min(DEVICE_CODE_POLL_SLEEP_SECONDS, max(0.0, deadline - time.time())) + sleep_for = min( + DEVICE_CODE_POLL_SLEEP_SECONDS, max(0.0, deadline - time.time()) + ) if sleep_for <= 0: break time.sleep(sleep_for) diff --git a/litellm/llms/chatgpt/chat/streaming_utils.py b/litellm/llms/chatgpt/chat/streaming_utils.py index 3232b452a3..e9cf2d15c2 100644 --- a/litellm/llms/chatgpt/chat/streaming_utils.py +++ b/litellm/llms/chatgpt/chat/streaming_utils.py @@ -24,7 +24,9 @@ class ChatGPTToolCallNormalizer: self._stream = stream self._seen_ids: Dict[str, int] = {} # tool_call_id -> assigned_index self._next_index: int = 0 - self._last_id: Optional[str] = None # tracks which tool call the next delta belongs to + self._last_id: Optional[ + str + ] = None # tracks which tool call the next delta belongs to def __getattr__(self, name: str) -> Any: return getattr(self._stream, name) diff --git a/litellm/llms/chatgpt/common_utils.py b/litellm/llms/chatgpt/common_utils.py index d80487cde2..9cbcd6a4f4 100644 --- a/litellm/llms/chatgpt/common_utils.py +++ b/litellm/llms/chatgpt/common_utils.py @@ -145,9 +145,7 @@ def _safe_header_value(value: str) -> str: def _sanitize_user_agent_token(value: str) -> str: if not value: return "" - return "".join( - ch if (ch.isalnum() or ch in "-_./") else "_" for ch in value - ) + return "".join(ch if (ch.isalnum() or ch in "-_./") else "_" for ch in value) def _terminal_user_agent() -> str: @@ -159,9 +157,7 @@ def _terminal_user_agent() -> str: wezterm_version = os.getenv("WEZTERM_VERSION") if wezterm_version is not None: - token = ( - f"WezTerm/{wezterm_version}" if wezterm_version else "WezTerm" - ) + token = f"WezTerm/{wezterm_version}" if wezterm_version else "WezTerm" return _sanitize_user_agent_token(token) or "WezTerm" if ( @@ -182,9 +178,7 @@ def _terminal_user_agent() -> str: konsole_version = os.getenv("KONSOLE_VERSION") if konsole_version is not None: - token = ( - f"Konsole/{konsole_version}" if konsole_version else "Konsole" - ) + token = f"Konsole/{konsole_version}" if konsole_version else "Konsole" return _sanitize_user_agent_token(token) or "Konsole" if os.getenv("GNOME_TERMINAL_SCREEN"): diff --git a/litellm/llms/chatgpt/responses/transformation.py b/litellm/llms/chatgpt/responses/transformation.py index 66acd93341..3c59ca1658 100644 --- a/litellm/llms/chatgpt/responses/transformation.py +++ b/litellm/llms/chatgpt/responses/transformation.py @@ -77,9 +77,9 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): existing_instructions = request.get("instructions") if existing_instructions: if base_instructions not in existing_instructions: - request["instructions"] = ( - f"{base_instructions}\n\n{existing_instructions}" - ) + request[ + "instructions" + ] = f"{base_instructions}\n\n{existing_instructions}" else: request["instructions"] = base_instructions request["store"] = False diff --git a/litellm/llms/clarifai/chat/transformation.py b/litellm/llms/clarifai/chat/transformation.py index 48884ff013..d07f6eba05 100644 --- a/litellm/llms/clarifai/chat/transformation.py +++ b/litellm/llms/clarifai/chat/transformation.py @@ -25,6 +25,7 @@ class ClarifaiConfig(OpenAIGPTConfig): Configuration class for Clarifai chat completions. Since Clarifai is OpenAI-compatible, we extend OpenAIGPTConfig. """ + def get_supported_openai_params(self, model: str) -> list: """ Get the supported OpenAI params for the given model @@ -42,18 +43,15 @@ class ClarifaiConfig(OpenAIGPTConfig): "frequency_penalty", "stream_options", ] - + @staticmethod def get_api_key(api_key: Optional[str] = None) -> Optional[str]: - return ( - api_key - or get_secret_str("CLARIFAI_API_KEY") - ) - + return api_key or get_secret_str("CLARIFAI_API_KEY") + @staticmethod def get_api_base(api_base: Optional[str] = None) -> Optional[str]: return api_base or "https://api.clarifai.com/v2/ext/openai/v1" - + @staticmethod def get_base_model(model: Optional[str] = None) -> Optional[str]: if model: @@ -72,11 +70,15 @@ class ClarifaiConfig(OpenAIGPTConfig): api_base = api_base or "https://api.clarifai.com/v2/ext/openai/v1" dynamic_api_key = api_key or get_secret_str("CLARIFAI_API_KEY") or "" return api_base, dynamic_api_key - - def transform_request(self, model, messages, optional_params, litellm_params, headers): + + def transform_request( + self, model, messages, optional_params, litellm_params, headers + ): model = self.get_base_model(model) or model - return super().transform_request(model, messages, optional_params, litellm_params, headers) - + return super().transform_request( + model, messages, optional_params, litellm_params, headers + ) + def transform_response( self, model: str, @@ -95,7 +97,7 @@ class ClarifaiConfig(OpenAIGPTConfig): Transform the Clarifai response to a standard ModelResponse. Since Clarifai is OpenAI-compatible, we use OpenAI response transformation. """ - ## Logging + ## Logging logging_obj.post_call( input=messages, api_key=api_key, @@ -111,9 +113,9 @@ class ClarifaiConfig(OpenAIGPTConfig): message=f"Failed to parse Clarifai response: {str(e)}", headers=raw_response.headers, ) from e - + response = ModelResponse(**completion_response) - + if response.model is not None: response.model = "clarifai/" + model @@ -130,4 +132,4 @@ class ClarifaiConfig(OpenAIGPTConfig): status_code=status_code, message=error_message, headers=headers, - ) \ No newline at end of file + ) diff --git a/litellm/llms/cohere/chat/v2_transformation.py b/litellm/llms/cohere/chat/v2_transformation.py index 8f6dde1967..190491adfc 100644 --- a/litellm/llms/cohere/chat/v2_transformation.py +++ b/litellm/llms/cohere/chat/v2_transformation.py @@ -7,7 +7,7 @@ import litellm from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.cohere import CohereV2ChatResponse from litellm.types.llms.openai import ( - AllMessageValues, + AllMessageValues, ChatCompletionToolCallChunk, ChatCompletionAnnotation, ChatCompletionAnnotationURLCitation, @@ -172,8 +172,10 @@ class CohereV2ChatConfig(OpenAIGPTConfig): """ Cohere v2 chat api is in openai format, so we can use the openai transform request function to transform the request. """ - data = super().transform_request(model, messages, optional_params, litellm_params, headers) - + data = super().transform_request( + model, messages, optional_params, litellm_params, headers + ) + return data def transform_response( @@ -215,10 +217,13 @@ class CohereV2ChatConfig(OpenAIGPTConfig): ## ADD CITATIONS AS ANNOTATIONS annotations: Optional[List[ChatCompletionAnnotation]] = None citations = None - - if "message" in cohere_v2_chat_response and "citations" in cohere_v2_chat_response["message"]: + + if ( + "message" in cohere_v2_chat_response + and "citations" in cohere_v2_chat_response["message"] + ): citations = cohere_v2_chat_response["message"]["citations"] - + if citations: annotations = self._translate_citations_to_openai_annotations(citations) @@ -293,13 +298,15 @@ class CohereV2ChatConfig(OpenAIGPTConfig): ) -> BaseLLMException: return CohereError(status_code=status_code, message=error_message) - def _translate_citations_to_openai_annotations(self, citations: List[dict]) -> List[ChatCompletionAnnotation]: + def _translate_citations_to_openai_annotations( + self, citations: List[dict] + ) -> List[ChatCompletionAnnotation]: """ Transform Cohere citations to OpenAI annotations format. - + Creates separate annotations for each source in a citation, allowing multiple annotations with the same start/end index if they reference different sources. - + Args: citations: List of Cohere citation objects with format: { @@ -318,40 +325,40 @@ class CohereV2ChatConfig(OpenAIGPTConfig): } ] } - + Returns: List of OpenAI ChatCompletionAnnotation objects (one per source) """ annotations: List[ChatCompletionAnnotation] = [] - + for citation in citations: start_index = citation.get("start", 0) end_index = citation.get("end", 0) - + # Extract source information - loop through all sources sources = citation.get("sources", []) if not sources: continue - + # Create an annotation for each source for source in sources: if source.get("type") == "document" and "document" in source: document = source["document"] title = document.get("title", "") url = source.get("url") or f"source:{source.get('id', 'unknown')}" - + url_citation: ChatCompletionAnnotationURLCitation = { "start_index": start_index, "end_index": end_index, "title": title, "url": url, } - + annotation: ChatCompletionAnnotation = { "type": "url_citation", "url_citation": url_citation, } - + annotations.append(annotation) - - return annotations \ No newline at end of file + + return annotations diff --git a/litellm/llms/cohere/common_utils.py b/litellm/llms/cohere/common_utils.py index 333916fffa..05e3cec544 100644 --- a/litellm/llms/cohere/common_utils.py +++ b/litellm/llms/cohere/common_utils.py @@ -66,25 +66,26 @@ class CohereModelInfo(BaseLLMModelInfo): This function will return `anthropic.claude-3-opus-20240229-v1:0` """ pass - + @staticmethod def get_cohere_route(model: str) -> Literal["v1", "v2"]: """ Get the Cohere route for the given model. - + Args: model: The model name (e.g., "cohere_chat/v2/command-r-plus", "command-r-plus") - + Returns: "v2" for standard Cohere v2 API (default), "v1" for Cohere v1 API """ # Check for explicit v1 route if "v1/" in model: return "v1" - + # Default to v2 for all other cases return "v2" + def validate_environment( headers: dict, model: str, @@ -216,9 +217,10 @@ class ModelResponseIterator: except ValueError as e: raise RuntimeError(f"Error parsing chunk: {e},\nReceived chunk: {chunk}") + class CohereV2ModelResponseIterator: """V2-specific response iterator for Cohere streaming""" - + def __init__( self, streaming_response, sync_stream: bool, json_mode: Optional[bool] = False ): @@ -239,7 +241,9 @@ class CohereV2ModelResponseIterator: return content return "" - def _parse_tool_call_delta(self, chunk: dict) -> Optional[ChatCompletionToolCallChunk]: + def _parse_tool_call_delta( + self, chunk: dict + ) -> Optional[ChatCompletionToolCallChunk]: """Parse tool-call-delta chunks to extract tool calls.""" delta = chunk.get("delta", {}) tool_calls = delta.get("tool_calls", []) @@ -249,8 +253,8 @@ class CohereV2ModelResponseIterator: "type": "function", "function": { "name": tool_calls[0].get("name", ""), - "arguments": tool_calls[0].get("arguments", "") - } + "arguments": tool_calls[0].get("arguments", ""), + }, } # type: ignore return None @@ -276,18 +280,20 @@ class CohereV2ModelResponseIterator: "end": citations.get("end", 0), "text": citations.get("text", ""), "sources": citations.get("sources", []), - "type": citations.get("type", "TEXT_CONTENT") + "type": citations.get("type", "TEXT_CONTENT"), } return {"citations": [citation_data]} return None - def _parse_message_end(self, chunk: dict) -> Tuple[bool, str, Optional[ChatCompletionUsageBlock]]: + def _parse_message_end( + self, chunk: dict + ) -> Tuple[bool, str, Optional[ChatCompletionUsageBlock]]: """Parse message-end events to extract finish info and usage.""" data = chunk.get("data", {}) delta = data.get("delta", {}) is_finished = True finish_reason = delta.get("finish_reason", "stop") - + usage = None usage_data = delta.get("usage", {}) if usage_data: @@ -295,15 +301,16 @@ class CohereV2ModelResponseIterator: usage = ChatCompletionUsageBlock( prompt_tokens=tokens_data.get("input_tokens", 0), completion_tokens=tokens_data.get("output_tokens", 0), - total_tokens=tokens_data.get("input_tokens", 0) + tokens_data.get("output_tokens", 0) + total_tokens=tokens_data.get("input_tokens", 0) + + tokens_data.get("output_tokens", 0), ) - + return is_finished, finish_reason, usage def chunk_parser(self, chunk: dict) -> GenericStreamingChunk: """ Parse Cohere v2 streaming chunks. - + v2 format: - Content: chunk.type == "content-delta" -> chunk.delta.message.content.text - Tool calls: chunk.type == "tool-call-delta" -> chunk.delta.tool_calls @@ -408,4 +415,3 @@ class CohereV2ModelResponseIterator: raise StopAsyncIteration except ValueError as e: raise RuntimeError(f"Error parsing chunk: {e},\nReceived chunk: {chunk}") - diff --git a/litellm/llms/cohere/rerank/transformation.py b/litellm/llms/cohere/rerank/transformation.py index d085cb13c4..531b94d180 100644 --- a/litellm/llms/cohere/rerank/transformation.py +++ b/litellm/llms/cohere/rerank/transformation.py @@ -21,8 +21,8 @@ class CohereRerankConfig(BaseRerankConfig): pass def get_complete_url( - self, - api_base: Optional[str], + self, + api_base: Optional[str], model: str, optional_params: Optional[dict] = None, ) -> str: @@ -63,14 +63,16 @@ class CohereRerankConfig(BaseRerankConfig): No mapping required - returns all supported params """ - return dict(OptionalRerankParams( - query=query, - documents=documents, - top_n=top_n, - rank_fields=rank_fields, - return_documents=return_documents, - max_chunks_per_doc=max_chunks_per_doc, - )) + return dict( + OptionalRerankParams( + query=query, + documents=documents, + top_n=top_n, + rank_fields=rank_fields, + return_documents=return_documents, + max_chunks_per_doc=max_chunks_per_doc, + ) + ) def validate_environment( self, diff --git a/litellm/llms/cohere/rerank_v2/transformation.py b/litellm/llms/cohere/rerank_v2/transformation.py index 01309d937f..60d22ff4be 100644 --- a/litellm/llms/cohere/rerank_v2/transformation.py +++ b/litellm/llms/cohere/rerank_v2/transformation.py @@ -13,8 +13,8 @@ class CohereRerankV2Config(CohereRerankConfig): pass def get_complete_url( - self, - api_base: Optional[str], + self, + api_base: Optional[str], model: str, optional_params: Optional[dict] = None, ) -> str: @@ -55,14 +55,16 @@ class CohereRerankV2Config(CohereRerankConfig): No mapping required - returns all supported params """ - return dict(OptionalRerankParams( - query=query, - documents=documents, - top_n=top_n, - rank_fields=rank_fields, - return_documents=return_documents, - max_tokens_per_doc=max_tokens_per_doc, - )) + return dict( + OptionalRerankParams( + query=query, + documents=documents, + top_n=top_n, + rank_fields=rank_fields, + return_documents=return_documents, + max_tokens_per_doc=max_tokens_per_doc, + ) + ) def transform_rerank_request( self, diff --git a/litellm/llms/cometapi/chat/transformation.py b/litellm/llms/cometapi/chat/transformation.py index fedb8f61e5..1e15ee188c 100644 --- a/litellm/llms/cometapi/chat/transformation.py +++ b/litellm/llms/cometapi/chat/transformation.py @@ -21,11 +21,11 @@ from ..common_utils import CometAPIException class CometAPIConfig(OpenAIGPTConfig): """ CometAPI configuration class, inherits from OpenAIGPTConfig - + Since CometAPI is OpenAI-compatible API, we inherit from OpenAIGPTConfig and only need to override necessary methods to handle CometAPI-specific features """ - + def map_openai_params( self, non_default_params: dict, @@ -47,10 +47,10 @@ class CometAPIConfig(OpenAIGPTConfig): # custom_param = non_default_params.pop("custom_param", None) # if custom_param is not None: # extra_body["custom_param"] = custom_param - + if extra_body: mapped_openai_params["extra_body"] = extra_body - + return mapped_openai_params def remove_cache_control_flag_from_messages_and_tools( @@ -129,10 +129,7 @@ class CometAPIConfig(OpenAIGPTConfig): return f"{api_base}/{endpoint}" def get_error_class( - self, - error_message: str, - status_code: int, - headers: Union[dict, httpx.Headers] + self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] ) -> BaseLLMException: """ Return CometAPI-specific error class @@ -163,7 +160,7 @@ class CometAPIChatCompletionStreamingHandler(BaseModelResponseIterator): """ Handler for CometAPI streaming chat completion responses """ - + def chunk_parser(self, chunk: dict) -> ModelResponseStream: """ Parse individual chunks from streaming response @@ -186,9 +183,11 @@ class CometAPIChatCompletionStreamingHandler(BaseModelResponseIterator): for choice in chunk["choices"]: # Handle reasoning content if present if "delta" in choice and "reasoning" in choice["delta"]: - choice["delta"]["reasoning_content"] = choice["delta"].get("reasoning") + choice["delta"]["reasoning_content"] = choice["delta"].get( + "reasoning" + ) new_choices.append(choice) - + return ModelResponseStream( id=chunk["id"], object="chat.completion.chunk", diff --git a/litellm/llms/cometapi/common_utils.py b/litellm/llms/cometapi/common_utils.py index 2e5e3e5fab..8cb0a30402 100644 --- a/litellm/llms/cometapi/common_utils.py +++ b/litellm/llms/cometapi/common_utils.py @@ -3,4 +3,5 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException class CometAPIException(BaseLLMException): """CometAPI exception handling class""" + pass diff --git a/litellm/llms/cometapi/embed/transformation.py b/litellm/llms/cometapi/embed/transformation.py index 5cfd125314..d1972def8b 100644 --- a/litellm/llms/cometapi/embed/transformation.py +++ b/litellm/llms/cometapi/embed/transformation.py @@ -19,7 +19,7 @@ from ..common_utils import CometAPIException class CometAPIEmbeddingConfig(BaseEmbeddingConfig): """ Configuration class for CometAPI Embedding API. - + Since CometAPI is OpenAI-compatible, this class provides OpenAI-standard embedding functionality with CometAPI-specific authentication and endpoints. """ diff --git a/litellm/llms/cometapi/image_generation/cost_calculator.py b/litellm/llms/cometapi/image_generation/cost_calculator.py index b10c9d0908..987e79e18d 100644 --- a/litellm/llms/cometapi/image_generation/cost_calculator.py +++ b/litellm/llms/cometapi/image_generation/cost_calculator.py @@ -22,4 +22,6 @@ def cost_calculator( num_images = len(image_response.data) return output_cost_per_image * num_images else: - raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}") + raise ValueError( + f"image_response must be of type ImageResponse got type={type(image_response)}" + ) diff --git a/litellm/llms/cometapi/image_generation/transformation.py b/litellm/llms/cometapi/image_generation/transformation.py index bf1ca9ddde..bc6bd3f3ec 100644 --- a/litellm/llms/cometapi/image_generation/transformation.py +++ b/litellm/llms/cometapi/image_generation/transformation.py @@ -23,7 +23,7 @@ else: class CometAPIImageGenerationConfig(BaseImageGenerationConfig): DEFAULT_BASE_URL: str = "https://api.cometapi.com" IMAGE_GENERATION_ENDPOINT: str = "v1/images/generations" - + def get_supported_openai_params( self, model: str ) -> List[OpenAIImageGenerationOptionalParams]: @@ -37,7 +37,7 @@ class CometAPIImageGenerationConfig(BaseImageGenerationConfig): "size", "style", ] - + def map_openai_params( self, non_default_params: dict, @@ -46,7 +46,7 @@ class CometAPIImageGenerationConfig(BaseImageGenerationConfig): drop_params: bool, ) -> dict: supported_params = self.get_supported_openai_params(model) - + for k in non_default_params.keys(): if k not in optional_params.keys(): if k in supported_params: @@ -74,7 +74,7 @@ class CometAPIImageGenerationConfig(BaseImageGenerationConfig): Get the complete url for the request """ complete_url: str = ( - api_base + api_base or get_secret_str("COMETAPI_BASE_URL") or get_secret_str("COMETAPI_API_BASE") or self.DEFAULT_BASE_URL @@ -95,15 +95,15 @@ class CometAPIImageGenerationConfig(BaseImageGenerationConfig): api_base: Optional[str] = None, ) -> dict: final_api_key: Optional[str] = ( - api_key or - get_secret_str("COMETAPI_KEY") or - get_secret_str("COMETAPI_API_KEY") + api_key + or get_secret_str("COMETAPI_KEY") + or get_secret_str("COMETAPI_API_KEY") ) if not final_api_key: raise ValueError("COMETAPI_KEY or COMETAPI_API_KEY is not set") - + headers["Authorization"] = f"Bearer {final_api_key}" - headers["Content-Type"] = "application/json" + headers["Content-Type"] = "application/json" return headers def transform_image_generation_request( @@ -153,10 +153,10 @@ class CometAPIImageGenerationConfig(BaseImageGenerationConfig): status_code=raw_response.status_code, headers=raw_response.headers, ) - + if not model_response.data: model_response.data = [] - + # CometAPI returns OpenAI-compatible format # Expected format: {"created": timestamp, "data": [{"url": "...", "b64_json": "..."}]} if "data" in response_data: @@ -166,5 +166,5 @@ class CometAPIImageGenerationConfig(BaseImageGenerationConfig): url=image_data.get("url"), ) model_response.data.append(image_obj) - + return model_response diff --git a/litellm/llms/compactifai/__init__.py b/litellm/llms/compactifai/__init__.py index 16b0c04cda..d081dd7cf6 100644 --- a/litellm/llms/compactifai/__init__.py +++ b/litellm/llms/compactifai/__init__.py @@ -1 +1 @@ -# CompactifAI provider for LiteLLM \ No newline at end of file +# CompactifAI provider for LiteLLM diff --git a/litellm/llms/compactifai/chat/__init__.py b/litellm/llms/compactifai/chat/__init__.py index d1a4463166..221b0e0219 100644 --- a/litellm/llms/compactifai/chat/__init__.py +++ b/litellm/llms/compactifai/chat/__init__.py @@ -1 +1 @@ -# CompactifAI chat completions \ No newline at end of file +# CompactifAI chat completions diff --git a/litellm/llms/compactifai/chat/transformation.py b/litellm/llms/compactifai/chat/transformation.py index 5cb8cd9a4a..d4b9c5a83a 100644 --- a/litellm/llms/compactifai/chat/transformation.py +++ b/litellm/llms/compactifai/chat/transformation.py @@ -76,7 +76,9 @@ class CompactifAIChatConfig(OpenAIGPTConfig): # Convert tool calls to content for JSON mode tool_calls = message.get("tool_calls", []) if len(tool_calls) == 1: - message["content"] = tool_calls[0]["function"].get("arguments", "") + message["content"] = tool_calls[0]["function"].get( + "arguments", "" + ) message["tool_calls"] = None returned_response = ModelResponse(**response_json) @@ -97,4 +99,4 @@ class CompactifAIChatConfig(OpenAIGPTConfig): status_code=status_code, message=error_message, headers=headers, - ) \ No newline at end of file + ) diff --git a/litellm/llms/custom_httpx/aiohttp_transport.py b/litellm/llms/custom_httpx/aiohttp_transport.py index 60f34a2a82..132191c946 100644 --- a/litellm/llms/custom_httpx/aiohttp_transport.py +++ b/litellm/llms/custom_httpx/aiohttp_transport.py @@ -83,7 +83,9 @@ class AiohttpResponseStream(httpx.AsyncByteStream): async def __aiter__(self) -> typing.AsyncIterator[bytes]: try: - async for chunk in self._aiohttp_response.content.iter_chunked(self.CHUNK_SIZE): + async for chunk in self._aiohttp_response.content.iter_chunked( + self.CHUNK_SIZE + ): yield chunk except ( aiohttp.ClientPayloadError, @@ -101,7 +103,9 @@ class AiohttpResponseStream(httpx.AsyncByteStream): # with message "Connection closed.". Treat this as a graceful # end-of-stream so downstream consumers don't error. if "Connection closed" in str(e): - verbose_logger.debug("Upstream closed streaming connection; ending iterator gracefully") + verbose_logger.debug( + "Upstream closed streaming connection; ending iterator gracefully" + ) return raise except aiohttp.http_exceptions.TransferEncodingError as e: @@ -191,7 +195,11 @@ class LiteLLMAiohttpTransport(AiohttpTransport): current_loop = asyncio.get_running_loop() # If session is from a different or closed loop, recreate it - if session_loop is None or session_loop != current_loop or session_loop.is_closed(): + if ( + session_loop is None + or session_loop != current_loop + or session_loop.is_closed() + ): # Close old session to prevent leaks old_session = self.client try: @@ -200,7 +208,9 @@ class LiteLLMAiohttpTransport(AiohttpTransport): asyncio.create_task(old_session.close()) except RuntimeError: # Different event loop - can't schedule task, rely on GC - verbose_logger.debug("Old session from different loop, relying on GC") + verbose_logger.debug( + "Old session from different loop, relying on GC" + ) except Exception as e: verbose_logger.debug(f"Error closing old session: {e}") @@ -305,7 +315,9 @@ class LiteLLMAiohttpTransport(AiohttpTransport): except RuntimeError as e: # Handle the case where session was closed between our check and actual use if "Session is closed" in str(e): - verbose_logger.debug(f"Session closed during request, retrying with new session: {e}") + verbose_logger.debug( + f"Session closed during request, retrying with new session: {e}" + ) # Force creation of a new session if hasattr(self, "_client_factory") and callable(self._client_factory): self.client = self._client_factory() @@ -336,7 +348,10 @@ class LiteLLMAiohttpTransport(AiohttpTransport): async def _get_proxy_settings(self, request: httpx.Request): proxy = None - if not (litellm.disable_aiohttp_trust_env or str_to_bool(os.getenv("DISABLE_AIOHTTP_TRUST_ENV", "False"))): + if not ( + litellm.disable_aiohttp_trust_env + or str_to_bool(os.getenv("DISABLE_AIOHTTP_TRUST_ENV", "False")) + ): try: proxy = self._proxy_from_env(request.url) except Exception as e: # pragma: no cover - best effort diff --git a/litellm/llms/custom_httpx/async_client_cleanup.py b/litellm/llms/custom_httpx/async_client_cleanup.py index abbc61dc96..22629383ac 100644 --- a/litellm/llms/custom_httpx/async_client_cleanup.py +++ b/litellm/llms/custom_httpx/async_client_cleanup.py @@ -28,17 +28,17 @@ async def close_litellm_async_clients(): pass # Handle AsyncHTTPHandler instances (used by Gemini and other providers) - elif hasattr(handler, 'client'): + elif hasattr(handler, "client"): client = handler.client # Check if the httpx client has an aiohttp transport - if hasattr(client, '_transport') and hasattr(client._transport, 'aclose'): + if hasattr(client, "_transport") and hasattr(client._transport, "aclose"): try: await client._transport.aclose() except Exception: # Silently ignore errors during cleanup pass # Also close the httpx client itself - if hasattr(client, 'aclose') and not client.is_closed: + if hasattr(client, "aclose") and not client.is_closed: try: await client.aclose() except Exception: @@ -46,7 +46,7 @@ async def close_litellm_async_clients(): pass # Handle any other objects with aclose method - elif hasattr(handler, 'aclose'): + elif hasattr(handler, "aclose"): try: await handler.aclose() except Exception: @@ -55,9 +55,11 @@ async def close_litellm_async_clients(): # Close the global base_llm_aiohttp_handler instance (issue #12443) # This is used by Gemini and other providers that use aiohttp - if hasattr(litellm, 'base_llm_aiohttp_handler'): - base_handler = getattr(litellm, 'base_llm_aiohttp_handler', None) - if isinstance(base_handler, BaseLLMAIOHTTPHandler) and hasattr(base_handler, 'close'): + if hasattr(litellm, "base_llm_aiohttp_handler"): + base_handler = getattr(litellm, "base_llm_aiohttp_handler", None) + if isinstance(base_handler, BaseLLMAIOHTTPHandler) and hasattr( + base_handler, "close" + ): try: await base_handler.close() except Exception: diff --git a/litellm/llms/custom_httpx/container_handler.py b/litellm/llms/custom_httpx/container_handler.py index 73017eaaf3..3767949375 100644 --- a/litellm/llms/custom_httpx/container_handler.py +++ b/litellm/llms/custom_httpx/container_handler.py @@ -60,15 +60,15 @@ def _build_url( path_params: Dict[str, str], ) -> str: """Build the full URL by substituting path parameters. - + The api_base from get_complete_url already includes /containers, so we need to strip that prefix from the path_template. """ # api_base ends with /containers, path_template starts with /containers # So we need to strip /containers from the path if path_template.startswith("/containers"): - path_template = path_template[len("/containers"):] - + path_template = path_template[len("/containers") :] + url = f"{api_base.rstrip('/')}{path_template}" for param, value in path_params.items(): url = url.replace(f"{{{param}}}", value) @@ -94,36 +94,36 @@ def _prepare_multipart_file_upload( ) -> tuple: """ Prepare file and headers for multipart upload. - + Returns: Tuple of (files_dict, headers_without_content_type) """ from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_file_data, ) - + extracted = extract_file_data(file) filename = extracted.get("filename") or "file" content = extracted.get("content") or b"" content_type = extracted.get("content_type") or "application/octet-stream" files = {"file": (filename, content, content_type)} - + # Remove content-type header - httpx will set it automatically for multipart headers_copy = headers.copy() headers_copy.pop("content-type", None) headers_copy.pop("Content-Type", None) - + return files, headers_copy class GenericContainerHandler: """ Generic handler for container file API endpoints. - + This single handler can process any endpoint defined in endpoints.json, eliminating the need for individual handler methods per endpoint. """ - + def handle( self, endpoint_name: str, @@ -139,7 +139,7 @@ class GenericContainerHandler: ) -> Union[Any, Coroutine[Any, Any, Any]]: """ Generic handler for any container file endpoint. - + Args: endpoint_name: Name of the endpoint (e.g., "list_container_files") container_provider_config: Provider-specific configuration @@ -164,7 +164,7 @@ class GenericContainerHandler: client=client, **kwargs, ) - + return self._sync_handle( endpoint_name=endpoint_name, container_provider_config=container_provider_config, @@ -176,7 +176,7 @@ class GenericContainerHandler: client=client, **kwargs, ) - + def _sync_handle( self, endpoint_name: str, @@ -193,7 +193,7 @@ class GenericContainerHandler: endpoint_config = _get_endpoint_config(endpoint_name) if not endpoint_config: raise ValueError(f"Unknown endpoint: {endpoint_name}") - + # Get HTTP client if client is None or not isinstance(client, HTTPHandler): http_client = _get_httpx_client( @@ -201,7 +201,7 @@ class GenericContainerHandler: ) else: http_client = client - + # Build request headers = container_provider_config.validate_environment( headers=extra_headers or {}, @@ -209,21 +209,25 @@ class GenericContainerHandler: ) if extra_headers: headers.update(extra_headers) - + api_base = container_provider_config.get_complete_url( api_base=litellm_params.get("api_base", None), litellm_params=dict(litellm_params), ) - + # Build URL with path params - path_params = {p: kwargs.get(p, "") for p in endpoint_config.get("path_params", [])} + path_params = { + p: kwargs.get(p, "") for p in endpoint_config.get("path_params", []) + } url = _build_url(api_base, endpoint_config["path"], path_params) - + # Build query params - query_params = _build_query_params(endpoint_config.get("query_params", []), kwargs) + query_params = _build_query_params( + endpoint_config.get("query_params", []), kwargs + ) if extra_query: query_params.update(extra_query) - + # Log request logging_obj.pre_call( input="", @@ -234,50 +238,63 @@ class GenericContainerHandler: "params": query_params, }, ) - + # Make request method = endpoint_config["method"].upper() returns_binary = endpoint_config.get("returns_binary", False) is_multipart = endpoint_config.get("is_multipart", False) - + try: if method == "GET": - response = http_client.get(url=url, headers=headers, params=query_params) + response = http_client.get( + url=url, headers=headers, params=query_params + ) elif method == "DELETE": - response = http_client.delete(url=url, headers=headers, params=query_params) + response = http_client.delete( + url=url, headers=headers, params=query_params + ) elif method == "POST": if is_multipart and "file" in kwargs: - files, headers = _prepare_multipart_file_upload(kwargs["file"], headers) - response = http_client.post(url=url, headers=headers, params=query_params, files=files) + files, headers = _prepare_multipart_file_upload( + kwargs["file"], headers + ) + response = http_client.post( + url=url, headers=headers, params=query_params, files=files + ) else: - response = http_client.post(url=url, headers=headers, params=query_params) + response = http_client.post( + url=url, headers=headers, params=query_params + ) else: raise ValueError(f"Unsupported HTTP method: {method}") - + # For binary responses, return raw content if returns_binary: return response.content - + # Check for error response response_json = response.json() if "error" in response_json: from litellm.llms.base_llm.chat.transformation import BaseLLMException - error_msg = response_json.get("error", {}).get("message", str(response_json)) + + error_msg = response_json.get("error", {}).get( + "message", str(response_json) + ) raise BaseLLMException( status_code=response.status_code, message=error_msg, headers=dict(response.headers), ) - + # Parse response response_type = RESPONSE_TYPES.get(endpoint_config["response_type"]) if response_type: return response_type(**response_json) return response_json - + except Exception as e: raise e - + async def _async_handle( self, endpoint_name: str, @@ -294,7 +311,7 @@ class GenericContainerHandler: endpoint_config = _get_endpoint_config(endpoint_name) if not endpoint_config: raise ValueError(f"Unknown endpoint: {endpoint_name}") - + # Get HTTP client if client is None or not isinstance(client, AsyncHTTPHandler): http_client = get_async_httpx_client( @@ -303,7 +320,7 @@ class GenericContainerHandler: ) else: http_client = client - + # Build request headers = container_provider_config.validate_environment( headers=extra_headers or {}, @@ -311,21 +328,25 @@ class GenericContainerHandler: ) if extra_headers: headers.update(extra_headers) - + api_base = container_provider_config.get_complete_url( api_base=litellm_params.get("api_base", None), litellm_params=dict(litellm_params), ) - + # Build URL with path params - path_params = {p: kwargs.get(p, "") for p in endpoint_config.get("path_params", [])} + path_params = { + p: kwargs.get(p, "") for p in endpoint_config.get("path_params", []) + } url = _build_url(api_base, endpoint_config["path"], path_params) - + # Build query params - query_params = _build_query_params(endpoint_config.get("query_params", []), kwargs) + query_params = _build_query_params( + endpoint_config.get("query_params", []), kwargs + ) if extra_query: query_params.update(extra_query) - + # Log request logging_obj.pre_call( input="", @@ -336,51 +357,63 @@ class GenericContainerHandler: "params": query_params, }, ) - + # Make request method = endpoint_config["method"].upper() returns_binary = endpoint_config.get("returns_binary", False) is_multipart = endpoint_config.get("is_multipart", False) - + try: if method == "GET": - response = await http_client.get(url=url, headers=headers, params=query_params) + response = await http_client.get( + url=url, headers=headers, params=query_params + ) elif method == "DELETE": - response = await http_client.delete(url=url, headers=headers, params=query_params) + response = await http_client.delete( + url=url, headers=headers, params=query_params + ) elif method == "POST": if is_multipart and "file" in kwargs: - files, headers = _prepare_multipart_file_upload(kwargs["file"], headers) - response = await http_client.post(url=url, headers=headers, params=query_params, files=files) + files, headers = _prepare_multipart_file_upload( + kwargs["file"], headers + ) + response = await http_client.post( + url=url, headers=headers, params=query_params, files=files + ) else: - response = await http_client.post(url=url, headers=headers, params=query_params) + response = await http_client.post( + url=url, headers=headers, params=query_params + ) else: raise ValueError(f"Unsupported HTTP method: {method}") - + # For binary responses, return raw content if returns_binary: return response.content - + # Check for error response response_json = response.json() if "error" in response_json: from litellm.llms.base_llm.chat.transformation import BaseLLMException - error_msg = response_json.get("error", {}).get("message", str(response_json)) + + error_msg = response_json.get("error", {}).get( + "message", str(response_json) + ) raise BaseLLMException( status_code=response.status_code, message=error_msg, headers=dict(response.headers), ) - + # Parse response response_type = RESPONSE_TYPES.get(endpoint_config["response_type"]) if response_type: return response_type(**response_json) return response_json - + except Exception as e: raise e # Singleton instance generic_container_handler = GenericContainerHandler() - diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 3dfef07d42..001547557d 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -51,6 +51,7 @@ try: except Exception: version = "0.0.0" + def get_default_headers() -> dict: """ Get default headers for HTTP requests. @@ -64,6 +65,7 @@ def get_default_headers() -> dict: return {"User-Agent": f"litellm/{version}"} + # Initialize headers (User-Agent) headers = get_default_headers() @@ -1235,7 +1237,9 @@ def get_async_httpx_client( if params is not None: # Filter out params that are only used for cache key, not for AsyncHTTPHandler.__init__ - handler_params = {k: v for k, v in params.items() if k != "disable_aiohttp_transport"} + handler_params = { + k: v for k, v in params.items() if k != "disable_aiohttp_transport" + } handler_params["shared_session"] = shared_session _new_client = AsyncHTTPHandler(**handler_params) else: @@ -1284,7 +1288,9 @@ def _get_httpx_client(params: Optional[dict] = None) -> HTTPHandler: if params is not None: # Filter out params that are only used for cache key, not for HTTPHandler.__init__ - handler_params = {k: v for k, v in params.items() if k != "disable_aiohttp_transport"} + handler_params = { + k: v for k, v in params.items() if k != "disable_aiohttp_transport" + } _new_client = HTTPHandler(**handler_params) else: _new_client = HTTPHandler(timeout=httpx.Timeout(timeout=600.0, connect=5.0)) diff --git a/litellm/llms/custom_httpx/httpx_handler.py b/litellm/llms/custom_httpx/httpx_handler.py index 491cd97f7d..ce58794671 100644 --- a/litellm/llms/custom_httpx/httpx_handler.py +++ b/litellm/llms/custom_httpx/httpx_handler.py @@ -8,6 +8,7 @@ try: except Exception: version = "0.0.0" + def get_default_headers() -> dict: """ Get default headers for HTTP requests. @@ -21,6 +22,7 @@ def get_default_headers() -> dict: return {"User-Agent": f"litellm/{version}"} + class HTTPHandler: def __init__(self, concurrent_limit=1000): headers = get_default_headers() diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index d39e0fa886..4e6c3cba68 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -445,7 +445,9 @@ class BaseLLMHTTPHandler: # Check if stream was converted for WebSearch interception # This is set by the async_pre_request_hook in WebSearchInterceptionLogger if litellm_params.get("_websearch_interception_converted_stream", False): - logging_obj.model_call_details["websearch_interception_converted_stream"] = True + logging_obj.model_call_details[ + "websearch_interception_converted_stream" + ] = True if acompletion is True: if stream is True: @@ -1355,6 +1357,7 @@ class BaseLLMHTTPHandler: Returns: (headers, complete_url, data, files) """ from litellm.llms.base_llm.ocr.transformation import OCRRequestData + headers = provider_config.validate_environment( api_key=api_key, api_base=api_base, @@ -1847,9 +1850,11 @@ class BaseLLMHTTPHandler: Optional[litellm.types.utils.ProviderSpecificHeader], kwargs.get("provider_specific_header", None), ) - provider_specific_headers = ProviderSpecificHeaderUtils.get_provider_specific_headers( - provider_specific_header=provider_specific_header, - custom_llm_provider=custom_llm_provider, + provider_specific_headers = ( + ProviderSpecificHeaderUtils.get_provider_specific_headers( + provider_specific_header=provider_specific_header, + custom_llm_provider=custom_llm_provider, + ) ) forwarded_headers = kwargs.get("headers", None) # Also check for extra_headers in kwargs (from config or direct calls) @@ -1874,7 +1879,7 @@ class BaseLLMHTTPHandler: api_key=api_key, api_base=api_base, ) - + headers = update_headers_with_filtered_beta( headers=headers, provider=custom_llm_provider ) @@ -2848,12 +2853,12 @@ class BaseLLMHTTPHandler: ) -> tuple[Optional[str], Optional[dict]]: """ Extract upload URL from initial file creation response. - + Args: response: HTTP response from initial file creation request upload_url_location: Where to find URL ('headers' or 'body') upload_url_key: Key name for URL in response body (default: 'upload_url') - + Returns: Tuple of (upload_url, response_data) - upload_url: The extracted upload URL, or None if not found @@ -2934,7 +2939,10 @@ class BaseLLMHTTPHandler: else: sync_httpx_client = client - if isinstance(transformed_request, dict) and "initial_request" in transformed_request: + if ( + isinstance(transformed_request, dict) + and "initial_request" in transformed_request + ): # Handle two-step uploads (TwoStepFileUploadConfig) # Used by providers like Manus, Google Cloud Storage try: @@ -2950,24 +2958,33 @@ class BaseLLMHTTPHandler: ) # Extract upload URL from response - upload_url, initial_response_data = self._extract_upload_url_from_response( + ( + upload_url, + initial_response_data, + ) = self._extract_upload_url_from_response( response=initial_response, - upload_url_location=transformed_request.get("upload_url_location", "headers"), - upload_url_key=transformed_request.get("upload_url_key", "upload_url"), + upload_url_location=transformed_request.get( + "upload_url_location", "headers" + ), + upload_url_key=transformed_request.get( + "upload_url_key", "upload_url" + ), ) if not upload_url: raise ValueError("Failed to get upload URL from initial request") # Step 2: Upload the actual file - upload_method = transformed_request["upload_request"].get("method", "POST").lower() + upload_method = ( + transformed_request["upload_request"].get("method", "POST").lower() + ) upload_response = getattr(sync_httpx_client, upload_method)( url=upload_url, headers=transformed_request["upload_request"]["headers"], data=transformed_request["upload_request"]["data"], timeout=timeout, ) - + # Store initial response for transformation if initial_response_data: litellm_params["initial_file_response"] = initial_response_data @@ -2976,7 +2993,11 @@ class BaseLLMHTTPHandler: e=e, provider_config=provider_config, ) - elif isinstance(transformed_request, dict) and "method" in transformed_request and "initial_request" not in transformed_request: + elif ( + isinstance(transformed_request, dict) + and "method" in transformed_request + and "initial_request" not in transformed_request + ): # Handle pre-signed requests (e.g., from Bedrock S3 uploads) # Type narrowing: this is a plain dict, not TwoStepFileUploadConfig presigned_request = cast(Dict[str, Any], transformed_request) @@ -3012,8 +3033,19 @@ class BaseLLMHTTPHandler: data=transformed_request, timeout=timeout, ) + elif isinstance(transformed_request, dict) and "file" in transformed_request: + # Handle multipart form-data uploads (e.g., Anthropic Files API) + # The dict contains tuples suitable for httpx's `files` parameter + upload_response = sync_httpx_client.post( + url=api_base, + headers=headers, + files=transformed_request, + timeout=timeout, + ) else: - raise ValueError(f"Unsupported transformed_request type: {type(transformed_request)}") + raise ValueError( + f"Unsupported transformed_request type: {type(transformed_request)}" + ) # Store the upload URL in litellm_params for the transformation method # Honour the URL already set by transform_create_file_request (e.g. Bedrock pre-signed S3 uploads), @@ -3063,7 +3095,10 @@ class BaseLLMHTTPHandler: }, ) - if isinstance(transformed_request, dict) and "initial_request" in transformed_request: + if ( + isinstance(transformed_request, dict) + and "initial_request" in transformed_request + ): # Handle two-step uploads (TwoStepFileUploadConfig) # Used by providers like Manus, Google Cloud Storage try: @@ -3079,24 +3114,33 @@ class BaseLLMHTTPHandler: ) # Extract upload URL from response - upload_url, initial_response_data = self._extract_upload_url_from_response( + ( + upload_url, + initial_response_data, + ) = self._extract_upload_url_from_response( response=initial_response, - upload_url_location=transformed_request.get("upload_url_location", "headers"), - upload_url_key=transformed_request.get("upload_url_key", "upload_url"), + upload_url_location=transformed_request.get( + "upload_url_location", "headers" + ), + upload_url_key=transformed_request.get( + "upload_url_key", "upload_url" + ), ) if not upload_url: raise ValueError("Failed to get upload URL from initial request") # Step 2: Upload the actual file - upload_method = transformed_request["upload_request"].get("method", "POST").lower() + upload_method = ( + transformed_request["upload_request"].get("method", "POST").lower() + ) upload_response = await getattr(async_httpx_client, upload_method)( url=upload_url, headers=transformed_request["upload_request"]["headers"], data=transformed_request["upload_request"]["data"], timeout=timeout, ) - + # Store initial response for transformation if initial_response_data: litellm_params["initial_file_response"] = initial_response_data @@ -3106,7 +3150,11 @@ class BaseLLMHTTPHandler: e=e, provider_config=provider_config, ) - elif isinstance(transformed_request, dict) and "method" in transformed_request and "initial_request" not in transformed_request: + elif ( + isinstance(transformed_request, dict) + and "method" in transformed_request + and "initial_request" not in transformed_request + ): # Handle pre-signed requests (e.g., from Bedrock S3 uploads) # Type narrowing: this is a plain dict, not TwoStepFileUploadConfig presigned_request = cast(Dict[str, Any], transformed_request) @@ -3140,8 +3188,19 @@ class BaseLLMHTTPHandler: data=transformed_request, timeout=timeout, ) + elif isinstance(transformed_request, dict) and "file" in transformed_request: + # Handle multipart form-data uploads (e.g., Anthropic Files API) + # The dict contains tuples suitable for httpx's `files` parameter + upload_response = await async_httpx_client.post( + url=api_base, + headers=headers, + files=transformed_request, + timeout=timeout, + ) else: - raise ValueError(f"Unsupported transformed_request type: {type(transformed_request)}") + raise ValueError( + f"Unsupported transformed_request type: {type(transformed_request)}" + ) return provider_config.transform_create_file_response( model=None, @@ -3740,7 +3799,10 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - url, data = responses_api_provider_config.transform_compact_response_api_request( + ( + url, + data, + ) = responses_api_provider_config.transform_compact_response_api_request( model=model, input=input, response_api_optional_request_params=response_api_optional_request_params, @@ -3819,7 +3881,10 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - url, data = responses_api_provider_config.transform_compact_response_api_request( + ( + url, + data, + ) = responses_api_provider_config.transform_compact_response_api_request( model=model, input=input, response_api_optional_request_params=response_api_optional_request_params, @@ -3913,9 +3978,7 @@ class BaseLLMHTTPHandler: ) try: - response = sync_httpx_client.get( - url=url, headers=headers, params=params - ) + response = sync_httpx_client.get(url=url, headers=headers, params=params) except Exception as e: raise self._handle_error(e=e, provider_config=provider_config) @@ -4043,9 +4106,7 @@ class BaseLLMHTTPHandler: ) try: - response = sync_httpx_client.delete( - url=url, headers=headers, params=params - ) + response = sync_httpx_client.delete(url=url, headers=headers, params=params) except Exception as e: raise self._handle_error(e=e, provider_config=provider_config) @@ -4173,9 +4234,7 @@ class BaseLLMHTTPHandler: ) try: - response = sync_httpx_client.get( - url=url, headers=headers, params=params - ) + response = sync_httpx_client.get(url=url, headers=headers, params=params) except Exception as e: raise self._handle_error(e=e, provider_config=provider_config) @@ -4255,7 +4314,9 @@ class BaseLLMHTTPHandler: _is_async: bool = False, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, - ) -> Union["HttpxBinaryResponseContent", Coroutine[Any, Any, "HttpxBinaryResponseContent"]]: + ) -> Union[ + "HttpxBinaryResponseContent", Coroutine[Any, Any, "HttpxBinaryResponseContent"] + ]: """ Retrieve file content by ID """ @@ -4303,9 +4364,7 @@ class BaseLLMHTTPHandler: ) try: - response = sync_httpx_client.get( - url=url, headers=headers, params=params - ) + response = sync_httpx_client.get(url=url, headers=headers, params=params) except Exception as e: raise self._handle_error(e=e, provider_config=provider_config) @@ -4413,32 +4472,33 @@ class BaseLLMHTTPHandler: from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger - callbacks = litellm.callbacks + ( - logging_obj.dynamic_success_callbacks or [] - ) + callbacks = litellm.callbacks + (logging_obj.dynamic_success_callbacks or []) tools = anthropic_messages_optional_request_params.get("tools", []) for callback in callbacks: try: if isinstance(callback, CustomLogger): # First: Check if agentic loop should run - should_run, tool_calls = ( - await callback.async_should_run_agentic_loop( - response=response, - model=model, - messages=messages, - tools=tools, - stream=stream, - custom_llm_provider=custom_llm_provider, - kwargs=kwargs, - ) + ( + should_run, + tool_calls, + ) = await callback.async_should_run_agentic_loop( + response=response, + model=model, + messages=messages, + tools=tools, + stream=stream, + custom_llm_provider=custom_llm_provider, + kwargs=kwargs, ) if should_run: # Second: Execute agentic loop # Add custom_llm_provider to kwargs so the agentic loop can reconstruct the full model name kwargs_with_provider = kwargs.copy() if kwargs else {} - kwargs_with_provider["custom_llm_provider"] = custom_llm_provider + kwargs_with_provider[ + "custom_llm_provider" + ] = custom_llm_provider agentic_response = await callback.async_run_agentic_loop( tools=tool_calls, model=model, @@ -4458,7 +4518,9 @@ class BaseLLMHTTPHandler: verbose_logger.exception( "LiteLLM.AgenticHookError: Exception in agentic completion hooks " "[call_id=%s model=%s]: %s", - _call_id, model, str(e), + _call_id, + model, + str(e), ) # Check if we need to convert response to fake stream @@ -4467,11 +4529,13 @@ class BaseLLMHTTPHandler: # 2. No agentic loop ran (LLM didn't use the tool) # 3. We have a non-streaming response that needs to be converted to streaming websearch_converted_stream = ( - logging_obj.model_call_details.get("websearch_interception_converted_stream", False) + logging_obj.model_call_details.get( + "websearch_interception_converted_stream", False + ) if logging_obj is not None else False ) - + if websearch_converted_stream: from typing import cast @@ -4482,11 +4546,11 @@ class BaseLLMHTTPHandler: from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) - + verbose_logger.debug( "WebSearchInterception: No tool call made, converting non-streaming response to fake stream" ) - + # Convert the non-streaming response to a fake stream # The response should be an AnthropicMessagesResponse (dict) if isinstance(response, dict): @@ -4495,7 +4559,7 @@ class BaseLLMHTTPHandler: response=cast(AnthropicMessagesResponse, response) ) return fake_stream - + return None async def _call_agentic_chat_completion_hooks( @@ -4520,45 +4584,50 @@ class BaseLLMHTTPHandler: from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger - callbacks = litellm.callbacks + ( - logging_obj.dynamic_success_callbacks or [] - ) + callbacks = litellm.callbacks + (logging_obj.dynamic_success_callbacks or []) tools = optional_params.get("tools", []) for callback in callbacks: try: if isinstance(callback, CustomLogger): # Check if callback has the chat completion agentic loop method - if not hasattr(callback, "async_should_run_chat_completion_agentic_loop"): + if not hasattr( + callback, "async_should_run_chat_completion_agentic_loop" + ): continue # First: Check if agentic loop should run - should_run, tool_calls = ( - await callback.async_should_run_chat_completion_agentic_loop( - response=response, - model=model, - messages=messages, - tools=tools, - stream=stream, - custom_llm_provider=custom_llm_provider, - kwargs=kwargs, - ) + ( + should_run, + tool_calls, + ) = await callback.async_should_run_chat_completion_agentic_loop( + response=response, + model=model, + messages=messages, + tools=tools, + stream=stream, + custom_llm_provider=custom_llm_provider, + kwargs=kwargs, ) if should_run: # Second: Execute agentic loop # Add custom_llm_provider to kwargs so the agentic loop can reconstruct the full model name kwargs_with_provider = kwargs.copy() if kwargs else {} - kwargs_with_provider["custom_llm_provider"] = custom_llm_provider - agentic_response = await callback.async_run_chat_completion_agentic_loop( - tools=tool_calls, - model=model, - messages=messages, - response=response, - optional_params=optional_params, - logging_obj=logging_obj, - stream=stream, - kwargs=kwargs_with_provider, + kwargs_with_provider[ + "custom_llm_provider" + ] = custom_llm_provider + agentic_response = ( + await callback.async_run_chat_completion_agentic_loop( + tools=tool_calls, + model=model, + messages=messages, + response=response, + optional_params=optional_params, + logging_obj=logging_obj, + stream=stream, + kwargs=kwargs_with_provider, + ) ) # First hook that runs agentic loop wins return agentic_response @@ -4574,27 +4643,29 @@ class BaseLLMHTTPHandler: # 2. No agentic loop ran (LLM didn't use the tool) # 3. We have a non-streaming response that needs to be converted to streaming websearch_converted_stream = ( - logging_obj.model_call_details.get("websearch_interception_converted_stream", False) + logging_obj.model_call_details.get( + "websearch_interception_converted_stream", False + ) if logging_obj is not None else False ) - + if websearch_converted_stream: from litellm._logging import verbose_logger from litellm.llms.base_llm.base_model_iterator import ( convert_model_response_to_streaming, ) - + verbose_logger.debug( "WebSearchInterception: No tool call made, converting non-streaming chat completion to fake stream" ) - + # Convert the non-streaming ModelResponse to a fake stream if hasattr(response, "choices"): # Use the existing converter for ModelResponse fake_stream = convert_model_response_to_streaming(response) return fake_stream - + return None def _handle_error( @@ -4694,7 +4765,9 @@ class BaseLLMHTTPHandler: # (e.g. Gemini/Vertex AI Live needs a `setup` message before any realtime_input) _session_config: Optional[str] = None if provider_config.requires_session_configuration(): - _session_config = provider_config.session_configuration_request(model) + _session_config = provider_config.session_configuration_request( + model + ) if _session_config: await backend_ws.send(_session_config) @@ -4762,7 +4835,9 @@ class BaseLLMHTTPHandler: async_httpx_client = client if provider_config is not None: - url = provider_config.get_complete_url(api_base=api_base, model=model or "", api_version=api_version) + url = provider_config.get_complete_url( + api_base=api_base, model=model or "", api_version=api_version + ) headers: Dict[str, Any] = provider_config.validate_environment( headers={}, model=model or "", api_key=api_key ) @@ -4832,7 +4907,9 @@ class BaseLLMHTTPHandler: async_httpx_client = client if provider_config is not None: - url = provider_config.get_realtime_calls_url(api_base=api_base, model=model or "", api_version=api_version) + url = provider_config.get_realtime_calls_url( + api_base=api_base, model=model or "", api_version=api_version + ) headers: Dict[str, Any] = provider_config.get_realtime_calls_headers( ephemeral_key=openai_ephemeral_key ) @@ -4907,7 +4984,10 @@ class BaseLLMHTTPHandler: - Uses ManagedResponsesWebSocketHandler which makes HTTP streaming calls - Forwards events over the websocket connection """ - if responses_api_provider_config is None or not responses_api_provider_config.supports_native_websocket(): + if ( + responses_api_provider_config is None + or not responses_api_provider_config.supports_native_websocket() + ): from litellm.responses.streaming_iterator import ( ManagedResponsesWebSocketHandler, ) @@ -5016,10 +5096,7 @@ class BaseLLMHTTPHandler: _is_async: bool = False, fake_stream: bool = False, litellm_metadata: Optional[Dict[str, Any]] = None, - ) -> Union[ - ImageResponse, - Coroutine[Any, Any, ImageResponse], - ]: + ) -> Union[ImageResponse, Coroutine[Any, Any, ImageResponse],]: """ Handles image edit requests. @@ -5231,10 +5308,7 @@ class BaseLLMHTTPHandler: fake_stream: bool = False, litellm_metadata: Optional[Dict[str, Any]] = None, api_key: Optional[str] = None, - ) -> Union[ - ImageResponse, - Coroutine[Any, Any, ImageResponse], - ]: + ) -> Union[ImageResponse, Coroutine[Any, Any, ImageResponse],]: """ Handles image generation requests. When _is_async=True, returns a coroutine instead of making the call directly. @@ -5474,10 +5548,7 @@ class BaseLLMHTTPHandler: fake_stream: bool = False, litellm_metadata: Optional[Dict[str, Any]] = None, api_key: Optional[str] = None, - ) -> Union[ - VideoObject, - Coroutine[Any, Any, VideoObject], - ]: + ) -> Union[VideoObject, Coroutine[Any, Any, VideoObject],]: """ Handles video generation requests. When _is_async=True, returns a coroutine instead of making the call directly. @@ -5515,7 +5586,7 @@ class BaseLLMHTTPHandler: model=model, litellm_params=litellm_params, ) - + if extra_headers: headers.update(extra_headers) @@ -5525,7 +5596,11 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - data, files, api_base = video_generation_provider_config.transform_video_create_request( + ( + data, + files, + api_base, + ) = video_generation_provider_config.transform_video_create_request( model=model, prompt=prompt, video_create_optional_request_params=video_generation_optional_request_params, @@ -5626,7 +5701,11 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - data, files, api_base = video_generation_provider_config.transform_video_create_request( + ( + data, + files, + api_base, + ) = video_generation_provider_config.transform_video_create_request( model=model, prompt=prompt, api_base=api_base, @@ -5647,7 +5726,7 @@ class BaseLLMHTTPHandler: ) try: - #Use JSON when no files, otherwise use form data with files + # Use JSON when no files, otherwise use form data with files if files is None or len(files) == 0: response = await async_httpx_client.post( url=api_base, @@ -6301,7 +6380,10 @@ class BaseLLMHTTPHandler: ) # Transform the request using the provider config - url, data = video_status_provider_config.transform_video_status_retrieve_request( + ( + url, + data, + ) = video_status_provider_config.transform_video_status_retrieve_request( video_id=video_id, api_base=api_base, litellm_params=litellm_params, @@ -6335,10 +6417,12 @@ class BaseLLMHTTPHandler: headers=headers, ) - return video_status_provider_config.transform_video_status_retrieve_response( - raw_response=response, - logging_obj=logging_obj, - custom_llm_provider=custom_llm_provider, + return ( + video_status_provider_config.transform_video_status_retrieve_response( + raw_response=response, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + ) ) except Exception as e: @@ -6388,7 +6472,10 @@ class BaseLLMHTTPHandler: ) # Transform the request using the provider config - url, data = video_status_provider_config.transform_video_status_retrieve_request( + ( + url, + data, + ) = video_status_provider_config.transform_video_status_retrieve_request( video_id=video_id, api_base=api_base, litellm_params=litellm_params, @@ -6421,10 +6508,12 @@ class BaseLLMHTTPHandler: url=url, headers=headers, ) - return video_status_provider_config.transform_video_status_retrieve_response( - raw_response=response, - logging_obj=logging_obj, - custom_llm_provider=custom_llm_provider, + return ( + video_status_provider_config.transform_video_status_retrieve_response( + raw_response=response, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + ) ) except Exception as e: @@ -6432,7 +6521,7 @@ class BaseLLMHTTPHandler: e=e, provider_config=video_status_provider_config, ) - + ###### CONTAINER HANDLER ###### def container_create_handler( self, @@ -6472,7 +6561,7 @@ class BaseLLMHTTPHandler: headers=extra_headers or {}, api_key=litellm_params.get("api_key", None), ) - + # Add Content-Type header for JSON requests headers["Content-Type"] = "application/json" @@ -6522,7 +6611,7 @@ class BaseLLMHTTPHandler: e=e, provider_config=container_provider_config, ) - + async def async_container_create_handler( self, name: str, @@ -6548,7 +6637,7 @@ class BaseLLMHTTPHandler: headers=extra_headers or {}, api_key=litellm_params.get("api_key", None), ) - + # Add Content-Type header for JSON requests headers["Content-Type"] = "application/json" @@ -6598,7 +6687,7 @@ class BaseLLMHTTPHandler: e=e, provider_config=container_provider_config, ) - + def container_list_handler( self, container_provider_config: "BaseContainerConfig", @@ -6690,7 +6779,7 @@ class BaseLLMHTTPHandler: e=e, provider_config=container_provider_config, ) - + async def async_container_list_handler( self, container_provider_config: "BaseContainerConfig", @@ -6767,7 +6856,7 @@ class BaseLLMHTTPHandler: e=e, provider_config=container_provider_config, ) - + def container_retrieve_handler( self, container_id: str, @@ -6823,7 +6912,7 @@ class BaseLLMHTTPHandler: litellm_params=litellm_params, headers=headers, ) - + # Add any extra query parameters if extra_query: params.update(extra_query) @@ -6857,7 +6946,7 @@ class BaseLLMHTTPHandler: e=e, provider_config=container_provider_config, ) - + async def async_container_retrieve_handler( self, container_id: str, @@ -6900,7 +6989,7 @@ class BaseLLMHTTPHandler: litellm_params=litellm_params, headers=headers, ) - + # Add any extra query parameters if extra_query: params.update(extra_query) @@ -6934,7 +7023,7 @@ class BaseLLMHTTPHandler: e=e, provider_config=container_provider_config, ) - + def container_delete_handler( self, container_id: str, @@ -6990,7 +7079,7 @@ class BaseLLMHTTPHandler: litellm_params=litellm_params, headers=headers, ) - + # Add any extra query parameters if extra_query: params.update(extra_query) @@ -7024,7 +7113,7 @@ class BaseLLMHTTPHandler: e=e, provider_config=container_provider_config, ) - + async def async_container_delete_handler( self, container_id: str, @@ -7067,7 +7156,7 @@ class BaseLLMHTTPHandler: litellm_params=litellm_params, headers=headers, ) - + # Add any extra query parameters if extra_query: params.update(extra_query) @@ -7116,7 +7205,9 @@ class BaseLLMHTTPHandler: timeout: Union[float, httpx.Timeout] = 600, _is_async: bool = False, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, - ) -> Union["ContainerFileListResponse", Coroutine[Any, Any, "ContainerFileListResponse"]]: + ) -> Union[ + "ContainerFileListResponse", Coroutine[Any, Any, "ContainerFileListResponse"] + ]: if _is_async: return self.async_container_file_list_handler( container_id=container_id, @@ -7323,7 +7414,10 @@ class BaseLLMHTTPHandler: ) # Transform the request using the provider config - url, params = container_provider_config.transform_container_file_content_request( + ( + url, + params, + ) = container_provider_config.transform_container_file_content_request( container_id=container_id, file_id=file_id, api_base=api_base, @@ -7396,7 +7490,10 @@ class BaseLLMHTTPHandler: ) # Transform the request using the provider config - url, params = container_provider_config.transform_container_file_content_request( + ( + url, + params, + ) = container_provider_config.transform_container_file_content_request( container_id=container_id, file_id=file_id, api_base=api_base, @@ -7470,7 +7567,9 @@ class BaseLLMHTTPHandler: ) # Check if provider has async transform method - if hasattr(vector_store_provider_config, "atransform_search_vector_store_request"): + if hasattr( + vector_store_provider_config, "atransform_search_vector_store_request" + ): ( url, request_body, @@ -7518,7 +7617,6 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.post( url=url, headers=headers, @@ -7816,9 +7914,7 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.get( - url=url, headers=headers - ) + response = await async_httpx_client.get(url=url, headers=headers) except Exception as e: raise self._handle_error(e=e, provider_config=vector_store_provider_config) @@ -7929,7 +8025,7 @@ class BaseLLMHTTPHandler: ) url = api_base - + params = {} if after is not None: params["after"] = after @@ -8011,7 +8107,7 @@ class BaseLLMHTTPHandler: ) url = api_base - + params = {} if after is not None: params["after"] = after @@ -8073,14 +8169,15 @@ class BaseLLMHTTPHandler: ) url = f"{api_base}/{vector_store_id}" - + request_body = dict(vector_store_update_optional_params) - + # Clean metadata to only include string values (OpenAI requirement) if "metadata" in request_body and request_body["metadata"] is not None: from litellm.utils import add_openai_metadata + request_body["metadata"] = add_openai_metadata(request_body["metadata"]) - + if extra_body: request_body.update(extra_body) @@ -8155,14 +8252,15 @@ class BaseLLMHTTPHandler: ) url = f"{api_base}/{vector_store_id}" - + request_body = dict(vector_store_update_optional_params) - + # Clean metadata to only include string values (OpenAI requirement) if "metadata" in request_body and request_body["metadata"] is not None: from litellm.utils import add_openai_metadata + request_body["metadata"] = add_openai_metadata(request_body["metadata"]) - + if extra_body: request_body.update(extra_body) @@ -8650,12 +8748,13 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - url, request_params = ( - vector_store_files_provider_config.transform_retrieve_vector_store_file_request( - vector_store_id=vector_store_id, - file_id=file_id, - api_base=api_base, - ) + ( + url, + request_params, + ) = vector_store_files_provider_config.transform_retrieve_vector_store_file_request( + vector_store_id=vector_store_id, + file_id=file_id, + api_base=api_base, ) logging_obj.pre_call( @@ -8727,12 +8826,13 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - url, request_params = ( - vector_store_files_provider_config.transform_retrieve_vector_store_file_request( - vector_store_id=vector_store_id, - file_id=file_id, - api_base=api_base, - ) + ( + url, + request_params, + ) = vector_store_files_provider_config.transform_retrieve_vector_store_file_request( + vector_store_id=vector_store_id, + file_id=file_id, + api_base=api_base, ) logging_obj.pre_call( @@ -8791,12 +8891,13 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - url, request_params = ( - vector_store_files_provider_config.transform_retrieve_vector_store_file_content_request( - vector_store_id=vector_store_id, - file_id=file_id, - api_base=api_base, - ) + ( + url, + request_params, + ) = vector_store_files_provider_config.transform_retrieve_vector_store_file_content_request( + vector_store_id=vector_store_id, + file_id=file_id, + api_base=api_base, ) logging_obj.pre_call( @@ -8871,12 +8972,13 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - url, request_params = ( - vector_store_files_provider_config.transform_retrieve_vector_store_file_content_request( - vector_store_id=vector_store_id, - file_id=file_id, - api_base=api_base, - ) + ( + url, + request_params, + ) = vector_store_files_provider_config.transform_retrieve_vector_store_file_content_request( + vector_store_id=vector_store_id, + file_id=file_id, + api_base=api_base, ) logging_obj.pre_call( @@ -9094,12 +9196,13 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - url, request_params = ( - vector_store_files_provider_config.transform_delete_vector_store_file_request( - vector_store_id=vector_store_id, - file_id=file_id, - api_base=api_base, - ) + ( + url, + request_params, + ) = vector_store_files_provider_config.transform_delete_vector_store_file_request( + vector_store_id=vector_store_id, + file_id=file_id, + api_base=api_base, ) logging_obj.pre_call( @@ -9174,12 +9277,13 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - url, request_params = ( - vector_store_files_provider_config.transform_delete_vector_store_file_request( - vector_store_id=vector_store_id, - file_id=file_id, - api_base=api_base, - ) + ( + url, + request_params, + ) = vector_store_files_provider_config.transform_delete_vector_store_file_request( + vector_store_id=vector_store_id, + file_id=file_id, + api_base=api_base, ) logging_obj.pre_call( @@ -9675,29 +9779,29 @@ class BaseLLMHTTPHandler: ) -> tuple[Optional[Dict], Optional[list]]: """ Helper to prepare multipart/form-data request for skills API. - + Args: request_body: Request body containing files and other fields headers: Request headers - + Returns: Tuple of (data_dict, files_list) for multipart request, or (None, None) if no files """ if "files" not in request_body or not request_body["files"]: return None, None - + # Remove content-type header if present - httpx will set it automatically for multipart if "content-type" in headers: del headers["content-type"] - + # Prepare files for multipart upload files = [] for file_obj in request_body["files"]: files.append(("files[]", file_obj)) - + # Prepare data (non-file fields) data = {k: v for k, v in request_body.items() if k != "files"} - + return data, files def create_skill_handler( @@ -9753,7 +9857,7 @@ class BaseLLMHTTPHandler: data, files = self._prepare_skill_multipart_request( request_body=request_body, headers=headers ) - + if files is not None: response = sync_httpx_client.post( url=url, headers=headers, data=data, files=files, timeout=timeout @@ -9813,7 +9917,7 @@ class BaseLLMHTTPHandler: data, files = self._prepare_skill_multipart_request( request_body=request_body, headers=headers ) - + if files is not None: response = await async_httpx_client.post( url=url, headers=headers, data=data, files=files, timeout=timeout @@ -10037,9 +10141,7 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.get( - url=url, headers=headers - ) + response = await async_httpx_client.get(url=url, headers=headers) except Exception as e: raise self._handle_error( e=e, @@ -10477,9 +10579,7 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.get( - url=url, headers=headers - ) + response = await async_httpx_client.get(url=url, headers=headers) except Exception as e: raise self._handle_error( e=e, @@ -11136,9 +11236,7 @@ class BaseLLMHTTPHandler: ) try: - response = await async_httpx_client.get( - url=url, headers=headers - ) + response = await async_httpx_client.get(url=url, headers=headers) except Exception as e: raise self._handle_error( e=e, diff --git a/litellm/llms/custom_httpx/mock_transport.py b/litellm/llms/custom_httpx/mock_transport.py index 262d0dff12..c9844753e0 100644 --- a/litellm/llms/custom_httpx/mock_transport.py +++ b/litellm/llms/custom_httpx/mock_transport.py @@ -18,6 +18,7 @@ import httpx # Pre-built response templates # --------------------------------------------------------------------------- + def _mock_id() -> str: return f"chatcmpl-mock-{uuid.uuid4().hex[:8]}" diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index 7c2a9569c5..8ae02bd65e 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -60,6 +60,7 @@ from ...anthropic.chat.transformation import AnthropicConfig from ...openai_like.chat.transformation import OpenAILikeChatConfig from ..common_utils import DatabricksBase, DatabricksException + def _sanitize_empty_content(message_dict: dict[str, Any]) -> None: """ Remove or filter content so empty text blocks are not sent. @@ -330,8 +331,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): if "reasoning_effort" in non_default_params and "claude" in model: optional_params["thinking"] = AnthropicConfig._map_reasoning_effort( - reasoning_effort=non_default_params.get("reasoning_effort"), - model=model + reasoning_effort=non_default_params.get("reasoning_effort"), model=model ) optional_params.pop("reasoning_effort", None) ## handle thinking tokens diff --git a/litellm/llms/dataforseo/search/__init__.py b/litellm/llms/dataforseo/search/__init__.py index 28990c1af3..66f16d1e03 100644 --- a/litellm/llms/dataforseo/search/__init__.py +++ b/litellm/llms/dataforseo/search/__init__.py @@ -8,4 +8,3 @@ DataForSEO offers comprehensive search engine data with high accuracy. from .transformation import DataForSEOSearchConfig __all__ = ["DataForSEOSearchConfig"] - diff --git a/litellm/llms/dataforseo/search/transformation.py b/litellm/llms/dataforseo/search/transformation.py index 86b472f61b..940f1ca600 100644 --- a/litellm/llms/dataforseo/search/transformation.py +++ b/litellm/llms/dataforseo/search/transformation.py @@ -20,23 +20,25 @@ from litellm.secret_managers.main import get_secret_str class DataForSEOSearchConfig(BaseSearchConfig): """ Configuration for DataForSEO SERP API search. - + DataForSEO uses HTTP Basic Auth with login:password credentials. API endpoint: https://api.dataforseo.com/v3/serp/google/organic/live/advanced """ - - DATAFORSEO_API_BASE = "https://api.dataforseo.com/v3/serp/google/organic/live/advanced" - + + DATAFORSEO_API_BASE = ( + "https://api.dataforseo.com/v3/serp/google/organic/live/advanced" + ) + @staticmethod def ui_friendly_name() -> str: return "DataForSEO" - + def get_http_method(self) -> Literal["GET", "POST"]: """ DataForSEO uses POST requests with JSON body. """ return "POST" - + def validate_environment( self, headers: Dict, @@ -46,7 +48,7 @@ class DataForSEOSearchConfig(BaseSearchConfig): ) -> Dict: """ Validate DataForSEO environment and set up authentication. - + DataForSEO uses HTTP Basic Auth with login:password format. The credentials should be in DATAFORSEO_LOGIN and DATAFORSEO_PASSWORD env vars, or passed as api_key in "login:password" format. @@ -56,23 +58,27 @@ class DataForSEOSearchConfig(BaseSearchConfig): # Get login and password login = get_secret_str("DATAFORSEO_LOGIN") password = get_secret_str("DATAFORSEO_PASSWORD") - + # If api_key is provided in "login:password" format, use it if api_key and ":" in api_key: login, password = api_key.split(":", 1) - + if not login: - raise ValueError("DATAFORSEO_LOGIN is not set. Set `DATAFORSEO_LOGIN` environment variable or pass credentials in api_key parameter.") - + raise ValueError( + "DATAFORSEO_LOGIN is not set. Set `DATAFORSEO_LOGIN` environment variable or pass credentials in api_key parameter." + ) + if not password: - raise ValueError("DATAFORSEO_PASSWORD is not set. Set `DATAFORSEO_PASSWORD` environment variable or pass credentials in api_key parameter.") - + raise ValueError( + "DATAFORSEO_PASSWORD is not set. Set `DATAFORSEO_PASSWORD` environment variable or pass credentials in api_key parameter." + ) + # Create Basic Auth header credentials = f"{login}:{password}" encoded_credentials = base64.b64encode(credentials.encode()).decode() headers["Authorization"] = f"Basic {encoded_credentials}" headers["Content-Type"] = "application/json" - + return headers def get_complete_url( @@ -84,10 +90,14 @@ class DataForSEOSearchConfig(BaseSearchConfig): ) -> str: """ Get complete URL for DataForSEO SERP API endpoint. - + DataForSEO uses POST requests, so no query parameters in URL. """ - return api_base or get_secret_str("DATAFORSEO_API_BASE") or self.DATAFORSEO_API_BASE + return ( + api_base + or get_secret_str("DATAFORSEO_API_BASE") + or self.DATAFORSEO_API_BASE + ) def transform_search_request( self, @@ -98,7 +108,7 @@ class DataForSEOSearchConfig(BaseSearchConfig): ) -> Union[Dict, List[Dict]]: """ Transform Search request to DataForSEO SERP API format. - + Args: query: Search query (string or list of strings). DataForSEO supports single string queries. optional_params: Optional parameters for the request @@ -107,48 +117,54 @@ class DataForSEOSearchConfig(BaseSearchConfig): - search_domain_filter: Domain to filter results → maps to `domain` - Plus any DataForSEO-specific parameters (location_code, language_code, device, os, etc.) api_key: DataForSEO credentials (login:password format) - + Returns: List[Dict]: Request body for DataForSEO API (array of task objects as required by API) """ # DataForSEO expects an array of task objects task: Dict[str, Any] = {} - + # Convert query to string if it's a list if isinstance(query, list): query = query[0] if query else "" - + # Required field: keyword task["keyword"] = query - + # Map unified parameters to DataForSEO parameters if "max_results" in optional_params and optional_params["max_results"]: # DataForSEO uses 'depth' for number of results (max 700) depth = min(int(optional_params["max_results"]), 700) task["depth"] = depth - + if "country" in optional_params and optional_params["country"]: # DataForSEO uses location_code (e.g., 2840 for USA) # For simplicity, we'll use location_name which accepts country names task["location_name"] = optional_params["country"] - - if "search_domain_filter" in optional_params and optional_params["search_domain_filter"]: + + if ( + "search_domain_filter" in optional_params + and optional_params["search_domain_filter"] + ): # DataForSEO uses 'domain' parameter to filter by domain task["domain"] = optional_params["search_domain_filter"] - + # Add defaults if not specified if "language_code" not in task and "language_name" not in task: task["language_code"] = "en" - + # DataForSEO requires a location - use default from constants if not specified if "location_code" not in task and "location_name" not in task: task["location_code"] = DEFAULT_DATAFORSEO_LOCATION_CODE - + # Pass through all other parameters as-is for param, value in optional_params.items(): - if param not in self.get_supported_perplexity_optional_params() and param not in task: + if ( + param not in self.get_supported_perplexity_optional_params() + and param not in task + ): task[param] = value - + # DataForSEO API expects an array of tasks return [task] @@ -160,35 +176,35 @@ class DataForSEOSearchConfig(BaseSearchConfig): ) -> SearchResponse: """ Transform DataForSEO SERP API response to LiteLLM unified SearchResponse format. - + DataForSEO → LiteLLM mappings: - tasks[0].result[*].items[*].title → SearchResult.title - tasks[0].result[*].items[*].url → SearchResult.url - tasks[0].result[*].items[*].description → SearchResult.snippet - No date/last_updated fields in standard response (set to None) - + Args: raw_response: Raw httpx response from DataForSEO API logging_obj: Logging object for tracking - + Returns: SearchResponse with standardized format """ response_json = raw_response.json() - + # Transform results to SearchResult objects results = [] - + # DataForSEO wraps results in tasks array if "tasks" in response_json and len(response_json["tasks"]) > 0: task = response_json["tasks"][0] - + # Check if task was successful if task.get("status_code") == 20000 and "result" in task: # Result is an array, take first element if len(task["result"]) > 0: result = task["result"][0] - + # Items contain the actual search results for item in result.get("items", []): # Only process organic search results @@ -201,9 +217,8 @@ class DataForSEOSearchConfig(BaseSearchConfig): last_updated=None, ) results.append(search_result) - + return SearchResponse( results=results, object="search", ) - diff --git a/litellm/llms/deepinfra/chat/transformation.py b/litellm/llms/deepinfra/chat/transformation.py index 5198260a24..c36b490abc 100644 --- a/litellm/llms/deepinfra/chat/transformation.py +++ b/litellm/llms/deepinfra/chat/transformation.py @@ -14,6 +14,7 @@ class DeepInfraConfig(OpenAIGPTConfig): The class `DeepInfra` provides configuration for the DeepInfra's Chat Completions API interface. Below are the parameters: """ + @property def custom_llm_provider(self) -> Optional[str]: return "deepinfra" @@ -73,7 +74,7 @@ class DeepInfraConfig(OpenAIGPTConfig): "top_p", "response_format", "tools", - "tool_choice" + "tool_choice", ] if litellm.supports_reasoning( @@ -119,17 +120,19 @@ class DeepInfraConfig(OpenAIGPTConfig): optional_params[param] = value return optional_params - def _transform_tool_message_content(self, messages: List[AllMessageValues]) -> List[AllMessageValues]: + def _transform_tool_message_content( + self, messages: List[AllMessageValues] + ) -> List[AllMessageValues]: """ Transform tool message content from array to string format for DeepInfra compatibility. - + DeepInfra requires tool message content to be a string, not an array. This method converts tool message content from array format to string format. - + Example transformation: - Input: {"role": "tool", "content": [{"type": "text", "text": "20"}]} - Output: {"role": "tool", "content": "20"} - + Or if content is complex: - Input: {"role": "tool", "content": [{"type": "text", "text": "result"}]} - Output: {"role": "tool", "content": "[{\"type\": \"text\", \"text\": \"result\"}]"} @@ -137,13 +140,13 @@ class DeepInfraConfig(OpenAIGPTConfig): for message in messages: if message.get("role") == "tool": content = message.get("content") - + # If content is a list/array, convert it to string if isinstance(content, list): # Check if it's a simple single text item if ( - len(content) == 1 - and isinstance(content[0], dict) + len(content) == 1 + and isinstance(content[0], dict) and content[0].get("type") == "text" and "text" in content[0] ): @@ -152,7 +155,7 @@ class DeepInfraConfig(OpenAIGPTConfig): else: # For complex content, serialize the entire array as JSON string message["content"] = json.dumps(content) - + return messages @overload @@ -163,7 +166,10 @@ class DeepInfraConfig(OpenAIGPTConfig): @overload def _transform_messages( - self, messages: List[AllMessageValues], model: str, is_async: Literal[False] = False + self, + messages: List[AllMessageValues], + model: str, + is_async: Literal[False] = False, ) -> List[AllMessageValues]: ... @@ -183,6 +189,7 @@ class DeepInfraConfig(OpenAIGPTConfig): ) transformed_messages = await parent_result return self._transform_tool_message_content(transformed_messages) + return _async_transform() else: # Call parent with is_async=False (literal) for sync case diff --git a/litellm/llms/deepinfra/rerank/transformation.py b/litellm/llms/deepinfra/rerank/transformation.py index 47f47418cb..71e300d258 100644 --- a/litellm/llms/deepinfra/rerank/transformation.py +++ b/litellm/llms/deepinfra/rerank/transformation.py @@ -29,8 +29,8 @@ class DeepinfraRerankConfig(BaseRerankConfig): """ def get_complete_url( - self, - api_base: Optional[str], + self, + api_base: Optional[str], model: str, optional_params: Optional[dict] = None, ) -> str: diff --git a/litellm/llms/docker_model_runner/chat/transformation.py b/litellm/llms/docker_model_runner/chat/transformation.py index 3d84b24a01..4b81502bf8 100644 --- a/litellm/llms/docker_model_runner/chat/transformation.py +++ b/litellm/llms/docker_model_runner/chat/transformation.py @@ -18,7 +18,7 @@ from ...openai.chat.gpt_transformation import OpenAIGPTConfig class DockerModelRunnerChatConfig(OpenAIGPTConfig): """ Configuration for Docker Model Runner API. - + Docker Model Runner uses URLs in the format: /engines/{engine}/v1/chat/completions The engine name (e.g., "llama.cpp") is part of the API endpoint path. """ @@ -59,7 +59,7 @@ class DockerModelRunnerChatConfig(OpenAIGPTConfig): ) -> Tuple[Optional[str], Optional[str]]: """ Get API base and key for Docker Model Runner. - + Default API base: http://localhost:22088/engines/llama.cpp The engine path should be included in the api_base. """ @@ -69,7 +69,9 @@ class DockerModelRunnerChatConfig(OpenAIGPTConfig): or "http://localhost:22088/engines/llama.cpp" ) # type: ignore # Docker Model Runner may not require authentication for local instances - dynamic_api_key = api_key or get_secret_str("DOCKER_MODEL_RUNNER_API_KEY") or "dummy-key" + dynamic_api_key = ( + api_key or get_secret_str("DOCKER_MODEL_RUNNER_API_KEY") or "dummy-key" + ) return api_base, dynamic_api_key def get_complete_url( @@ -83,13 +85,13 @@ class DockerModelRunnerChatConfig(OpenAIGPTConfig): ) -> str: """ Build the complete URL for Docker Model Runner API. - + Docker Model Runner uses URLs in the format: /engines/{engine}/v1/chat/completions - + The engine name should be specified in the api_base: - api_base="http://model-runner.docker.internal/engines/llama.cpp" - Default: "http://localhost:22088/engines/llama.cpp" - + Args: api_base: Base URL for the Docker Model Runner instance including engine path api_key: API key (may not be required for local instances) @@ -97,26 +99,26 @@ class DockerModelRunnerChatConfig(OpenAIGPTConfig): optional_params: Optional parameters litellm_params: LiteLLM parameters stream: Whether streaming is enabled - + Returns: Complete URL for the API call """ if not api_base: api_base = "http://localhost:22088/engines/llama.cpp" - + # Remove trailing slashes from api_base api_base = api_base.rstrip("/") - + # Build the URL: {api_base}/v1/chat/completions # api_base is expected to already contain the engine path complete_url = f"{api_base}/v1/chat/completions" - + return complete_url def get_supported_openai_params(self, model: str) -> list: """ Get the supported OpenAI params for Docker Model Runner. - + Docker Model Runner is OpenAI-compatible and supports standard parameters. """ return super().get_supported_openai_params(model=model) @@ -130,7 +132,7 @@ class DockerModelRunnerChatConfig(OpenAIGPTConfig): ) -> dict: """ Map OpenAI parameters to Docker Model Runner parameters. - + Docker Model Runner is OpenAI-compatible, so most parameters map directly. """ supported_openai_params = self.get_supported_openai_params(model) @@ -141,4 +143,3 @@ class DockerModelRunnerChatConfig(OpenAIGPTConfig): optional_params[param] = value return optional_params - diff --git a/litellm/llms/duckduckgo/search/transformation.py b/litellm/llms/duckduckgo/search/transformation.py index 509d69041f..c754338153 100644 --- a/litellm/llms/duckduckgo/search/transformation.py +++ b/litellm/llms/duckduckgo/search/transformation.py @@ -19,6 +19,7 @@ from litellm.secret_managers.main import get_secret_str class _DuckDuckGoSearchRequestRequired(TypedDict): """Required fields for DuckDuckGo Search API request.""" + q: str # Required - search query @@ -27,6 +28,7 @@ class DuckDuckGoSearchRequest(_DuckDuckGoSearchRequestRequired, total=False): DuckDuckGo Instant Answer API request format. Based on: https://duckduckgo.com/api """ + format: str # Optional - output format ('json', 'xml'), default 'json' pretty: int # Optional - pretty print (0 or 1), default 1 no_redirect: int # Optional - skip HTTP redirects (0 or 1), default 0 @@ -36,21 +38,21 @@ class DuckDuckGoSearchRequest(_DuckDuckGoSearchRequestRequired, total=False): class DuckDuckGoSearchConfig(BaseSearchConfig): DUCKDUCKGO_API_BASE = "https://api.duckduckgo.com" - + @staticmethod def ui_friendly_name() -> str: return "DuckDuckGo" - + def get_http_method(self) -> Literal["GET", "POST"]: """ Get HTTP method for search requests. DuckDuckGo Instant Answer API uses GET requests. - + Returns: HTTP method 'GET' """ return "GET" - + def validate_environment( self, headers: Dict, @@ -77,16 +79,19 @@ class DuckDuckGoSearchConfig(BaseSearchConfig): Get complete URL for Search endpoint. DuckDuckGo uses query parameters, so we construct the URL with the query. """ - api_base = api_base or get_secret_str("DUCKDUCKGO_API_BASE") or self.DUCKDUCKGO_API_BASE - + api_base = ( + api_base + or get_secret_str("DUCKDUCKGO_API_BASE") + or self.DUCKDUCKGO_API_BASE + ) + # Build query parameters from the transformed request body if data and isinstance(data, dict) and "_duckduckgo_params" in data: params = data["_duckduckgo_params"] query_string = urlencode(params, doseq=True) return f"{api_base}/?{query_string}" - + return api_base - def transform_search_request( self, @@ -96,7 +101,7 @@ class DuckDuckGoSearchConfig(BaseSearchConfig): ) -> Dict: """ Transform Search request to DuckDuckGo API format. - + Args: query: Search query (string or list of strings). DuckDuckGo only supports single string queries. optional_params: Optional parameters for the request @@ -106,7 +111,7 @@ class DuckDuckGoSearchConfig(BaseSearchConfig): - no_redirect: Skip HTTP redirects (0 or 1) - no_html: Remove HTML from text (0 or 1) - skip_disambig: Skip disambiguation results (0 or 1) - + Returns: Dict with typed request data following DuckDuckGoSearchRequest spec """ @@ -118,19 +123,19 @@ class DuckDuckGoSearchConfig(BaseSearchConfig): "q": query, "format": "json", # Always use JSON format } - + # Convert to dict before dynamic key assignments result_data = dict(request_data) - + if "max_results" in optional_params: result_data["_max_results"] = optional_params["max_results"] - + # Pass through DuckDuckGo-specific parameters ddg_params = ["pretty", "no_redirect", "no_html", "skip_disambig"] for param in ddg_params: if param in optional_params: result_data[param] = optional_params[param] - + return { "_duckduckgo_params": result_data, } @@ -143,22 +148,22 @@ class DuckDuckGoSearchConfig(BaseSearchConfig): ) -> SearchResponse: """ Transform DuckDuckGo API response to LiteLLM unified SearchResponse format. - + DuckDuckGo → LiteLLM mappings: - RelatedTopics[].Text → SearchResult.title + snippet - RelatedTopics[].FirstURL → SearchResult.url - RelatedTopics[].Text → SearchResult.snippet - No date/last_updated fields in DuckDuckGo response (set to None) - + Args: raw_response: Raw httpx response from DuckDuckGo API logging_obj: Logging object for tracking - + Returns: SearchResponse with standardized format """ response_json = raw_response.json() - + # Extract max_results from the request URL params query_params = raw_response.request.url.params if raw_response.request else {} max_results = None @@ -167,13 +172,13 @@ class DuckDuckGoSearchConfig(BaseSearchConfig): max_results = int(query_params["_max_results"]) except (ValueError, TypeError): pass - + # Transform results to SearchResult objects results = [] - + # DuckDuckGo can return results in different fields # Priority: Abstract > Answer > RelatedTopics - + # Check if there's an Abstract with URL if response_json.get("AbstractURL") and response_json.get("AbstractText"): abstract_result = SearchResult( @@ -184,20 +189,20 @@ class DuckDuckGoSearchConfig(BaseSearchConfig): last_updated=None, ) results.append(abstract_result) - + # Process RelatedTopics related_topics = response_json.get("RelatedTopics", []) for topic in related_topics: # Stop if we've reached max_results if max_results is not None and len(results) >= max_results: break - + if isinstance(topic, dict): # Check if it's a direct result if "FirstURL" in topic and "Text" in topic: text = topic.get("Text", "") url = topic.get("FirstURL", "") - + # Try to split title and snippet if " - " in text: parts = text.split(" - ", 1) @@ -206,7 +211,7 @@ class DuckDuckGoSearchConfig(BaseSearchConfig): else: title = text[:50] + "..." if len(text) > 50 else text snippet = text - + search_result = SearchResult( title=title, url=url, @@ -215,7 +220,7 @@ class DuckDuckGoSearchConfig(BaseSearchConfig): last_updated=None, ) results.append(search_result) - + # Check if it contains nested topics elif "Topics" in topic: nested_topics = topic.get("Topics", []) @@ -223,11 +228,11 @@ class DuckDuckGoSearchConfig(BaseSearchConfig): # Stop if we've reached max_results if max_results is not None and len(results) >= max_results: break - + if "FirstURL" in nested_topic and "Text" in nested_topic: text = nested_topic.get("Text", "") url = nested_topic.get("FirstURL", "") - + # Try to split title and snippet if " - " in text: parts = text.split(" - ", 1) @@ -236,7 +241,7 @@ class DuckDuckGoSearchConfig(BaseSearchConfig): else: title = text[:50] + "..." if len(text) > 50 else text snippet = text - + search_result = SearchResult( title=title, url=url, @@ -245,7 +250,7 @@ class DuckDuckGoSearchConfig(BaseSearchConfig): last_updated=None, ) results.append(search_result) - + return SearchResponse( results=results, object="search", diff --git a/litellm/llms/elevenlabs/audio_transcription/transformation.py b/litellm/llms/elevenlabs/audio_transcription/transformation.py index e56e83b4de..8746e92d9f 100644 --- a/litellm/llms/elevenlabs/audio_transcription/transformation.py +++ b/litellm/llms/elevenlabs/audio_transcription/transformation.py @@ -66,20 +66,19 @@ class ElevenLabsAudioTranscriptionConfig(BaseAudioTranscriptionConfig): ) -> AudioTranscriptionRequestData: """ Transforms the audio transcription request for ElevenLabs API. - + Returns AudioTranscriptionRequestData with both form data and files. - + Returns: AudioTranscriptionRequestData: Structured data with form data and files """ - + # Use common utility to process the audio file processed_audio = process_audio_file(audio_file) - + # Prepare form data form_data = {"model_id": model} - ######################################################### # Add OpenAI Compatible Parameters ######################################################### @@ -87,29 +86,31 @@ class ElevenLabsAudioTranscriptionConfig(BaseAudioTranscriptionConfig): if key in self.get_supported_openai_params(model) and value is not None: # Convert values to strings for form data, but skip None values form_data[key] = str(value) - + ######################################################### # Add Provider Specific Parameters ######################################################### provider_specific_params = self.get_provider_specific_params( model=model, optional_params=optional_params, - openai_params=self.get_supported_openai_params(model) + openai_params=self.get_supported_openai_params(model), ) for key, value in provider_specific_params.items(): form_data[key] = str(value) ######################################################### ######################################################### - - # Prepare files - files = {"file": (processed_audio.filename, processed_audio.file_content, processed_audio.content_type)} - - return AudioTranscriptionRequestData( - data=form_data, - files=files - ) + # Prepare files + files = { + "file": ( + processed_audio.filename, + processed_audio.file_content, + processed_audio.content_type, + ) + } + + return AudioTranscriptionRequestData(data=form_data, files=files) def transform_audio_transcription_response( self, @@ -130,18 +131,20 @@ class ElevenLabsAudioTranscriptionConfig(BaseAudioTranscriptionConfig): # Add additional metadata matching OpenAI format response["task"] = "transcribe" response["language"] = response_json.get("language_code", "unknown") - + # Map ElevenLabs words to OpenAI format if "words" in response_json: response["words"] = [] for word_data in response_json["words"]: # Only include actual words, skip spacing and audio events if word_data.get("type") == "word": - response["words"].append({ - "word": word_data.get("text", ""), - "start": word_data.get("start", 0), - "end": word_data.get("end", 0) - }) + response["words"].append( + { + "word": word_data.get("text", ""), + "start": word_data.get("start", 0), + "end": word_data.get("end", 0), + } + ) # Store full response in hidden params response._hidden_params = response_json @@ -194,4 +197,4 @@ class ElevenLabsAudioTranscriptionConfig(BaseAudioTranscriptionConfig): } headers.update(auth_header) - return headers \ No newline at end of file + return headers diff --git a/litellm/llms/elevenlabs/common_utils.py b/litellm/llms/elevenlabs/common_utils.py index c1421b619f..d3221933eb 100644 --- a/litellm/llms/elevenlabs/common_utils.py +++ b/litellm/llms/elevenlabs/common_utils.py @@ -2,4 +2,4 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException class ElevenLabsException(BaseLLMException): - pass \ No newline at end of file + pass diff --git a/litellm/llms/elevenlabs/text_to_speech/transformation.py b/litellm/llms/elevenlabs/text_to_speech/transformation.py index b78d0bafc5..4dac2b8ba9 100644 --- a/litellm/llms/elevenlabs/text_to_speech/transformation.py +++ b/litellm/llms/elevenlabs/text_to_speech/transformation.py @@ -192,17 +192,17 @@ class ElevenLabsTextToSpeechConfig(BaseTextToSpeechConfig): "xi-api-key": api_key, "Content-Type": "application/json", } - ) - + ) + return headers - + def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, Headers] ) -> BaseLLMException: return ElevenLabsException( message=error_message, status_code=status_code, headers=headers ) - + def transform_text_to_speech_request( self, model: str, @@ -311,9 +311,7 @@ class ElevenLabsTextToSpeechConfig(BaseTextToSpeechConfig): Construct the ElevenLabs endpoint URL, including path voice_id and query params. """ base_url = ( - api_base - or get_secret_str("ELEVENLABS_API_BASE") - or self.TTS_BASE_URL + api_base or get_secret_str("ELEVENLABS_API_BASE") or self.TTS_BASE_URL ) base_url = base_url.rstrip("/") @@ -329,4 +327,4 @@ class ElevenLabsTextToSpeechConfig(BaseTextToSpeechConfig): if query_params: url = f"{url}?{urlencode(query_params)}" - return url \ No newline at end of file + return url diff --git a/litellm/llms/exa_ai/search/__init__.py b/litellm/llms/exa_ai/search/__init__.py index b647d2cd80..db1f080464 100644 --- a/litellm/llms/exa_ai/search/__init__.py +++ b/litellm/llms/exa_ai/search/__init__.py @@ -4,4 +4,3 @@ Exa AI Search API module. from litellm.llms.exa_ai.search.transformation import ExaAISearchConfig __all__ = ["ExaAISearchConfig"] - diff --git a/litellm/llms/exa_ai/search/transformation.py b/litellm/llms/exa_ai/search/transformation.py index 6b51c6cf25..fb352f3f93 100644 --- a/litellm/llms/exa_ai/search/transformation.py +++ b/litellm/llms/exa_ai/search/transformation.py @@ -18,6 +18,7 @@ from litellm.secret_managers.main import get_secret_str class _ExaAISearchRequestRequired(TypedDict): """Required fields for Exa AI Search API request.""" + query: str # Required - search query @@ -26,6 +27,7 @@ class ExaAISearchRequest(_ExaAISearchRequestRequired, total=False): Exa AI Search API request format. Based on: https://docs.exa.ai/reference/search """ + type: str # Optional - search type ('keyword', 'neural', 'fast', 'auto'), default 'auto' category: str # Optional - data category ('company', 'research paper', 'news', 'pdf', 'github', 'tweet', 'personal site', 'linkedin profile', 'financial report') userLocation: str # Optional - two-letter ISO country code @@ -37,7 +39,9 @@ class ExaAISearchRequest(_ExaAISearchRequestRequired, total=False): startPublishedDate: str # Optional - published date filter (ISO 8601 format) endPublishedDate: str # Optional - published date filter (ISO 8601 format) includeText: List[str] # Optional - strings that must be present in webpage text - excludeText: List[str] # Optional - strings that must not be present in webpage text + excludeText: List[ + str + ] # Optional - strings that must not be present in webpage text context: Union[bool, dict] # Optional - format results for LLMs moderation: bool # Optional - enable content moderation, default false contents: dict # Optional - content retrieval options @@ -45,11 +49,11 @@ class ExaAISearchRequest(_ExaAISearchRequestRequired, total=False): class ExaAISearchConfig(BaseSearchConfig): EXA_AI_API_BASE = "https://api.exa.ai" - + @staticmethod def ui_friendly_name() -> str: return "Exa AI" - + def validate_environment( self, headers: Dict, @@ -62,7 +66,9 @@ class ExaAISearchConfig(BaseSearchConfig): """ api_key = api_key or get_secret_str("EXA_API_KEY") if not api_key: - raise ValueError("EXA_API_KEY is not set. Set `EXA_API_KEY` environment variable.") + raise ValueError( + "EXA_API_KEY is not set. Set `EXA_API_KEY` environment variable." + ) headers["x-api-key"] = api_key headers["Content-Type"] = "application/json" return headers @@ -78,13 +84,12 @@ class ExaAISearchConfig(BaseSearchConfig): Get complete URL for Search endpoint. """ api_base = api_base or get_secret_str("EXA_API_BASE") or self.EXA_AI_API_BASE - + # Append "/search" to the api base if it's not already there if not api_base.endswith("/search"): api_base = f"{api_base}/search" return api_base - def transform_search_request( self, @@ -94,20 +99,20 @@ class ExaAISearchConfig(BaseSearchConfig): ) -> Dict: """ Transform Search request to Exa AI API format. - + Transforms Perplexity unified spec parameters: - query → query (same) - max_results → numResults - search_domain_filter → includeDomains - country → userLocation - max_tokens_per_page → (not applicable, ignored) - + All other Exa-specific parameters are passed through as-is. - + Args: query: Search query (string or list of strings). Exa AI only supports single string queries. optional_params: Optional parameters for the request - + Returns: Dict with typed request data following ExaAISearchRequest spec """ @@ -118,30 +123,33 @@ class ExaAISearchConfig(BaseSearchConfig): request_data: ExaAISearchRequest = { "query": query, } - + # Transform Perplexity unified spec parameters to Exa format if "max_results" in optional_params: request_data["numResults"] = optional_params["max_results"] - + if "search_domain_filter" in optional_params: request_data["includeDomains"] = optional_params["search_domain_filter"] - + if "country" in optional_params: request_data["userLocation"] = optional_params["country"] - + # Convert to dict before dynamic key assignments result_data = dict(request_data) - + # pass through all other parameters as-is for param, value in optional_params.items(): - if param not in self.get_supported_perplexity_optional_params() and param not in result_data: + if ( + param not in self.get_supported_perplexity_optional_params() + and param not in result_data + ): result_data[param] = value - + # By default, request text content if not explicitly specified # Exa AI doesn't return content/text unless explicitly requested if "contents" not in result_data: result_data["contents"] = {"text": True} - + return result_data def transform_search_response( @@ -152,23 +160,23 @@ class ExaAISearchConfig(BaseSearchConfig): ) -> SearchResponse: """ Transform Exa AI API response to LiteLLM unified SearchResponse format. - + Exa AI → LiteLLM mappings: - results[].title → SearchResult.title - results[].url → SearchResult.url - results[].text → SearchResult.snippet - results[].publishedDate → SearchResult.date - No last_updated field in Exa AI response (set to None) - + Args: raw_response: Raw httpx response from Exa AI API logging_obj: Logging object for tracking - + Returns: SearchResponse with standardized format """ response_json = raw_response.json() - + # Transform results to SearchResult objects results = [] for result in response_json.get("results", []): @@ -180,9 +188,8 @@ class ExaAISearchConfig(BaseSearchConfig): last_updated=None, # Exa AI doesn't provide last_updated in response ) results.append(search_result) - + return SearchResponse( results=results, object="search", ) - diff --git a/litellm/llms/fal_ai/__init__.py b/litellm/llms/fal_ai/__init__.py index 34cac014ce..0de526a8eb 100644 --- a/litellm/llms/fal_ai/__init__.py +++ b/litellm/llms/fal_ai/__init__.py @@ -25,4 +25,3 @@ __all__ = [ "FalAIStableDiffusionConfig", "get_fal_ai_image_generation_config", ] - diff --git a/litellm/llms/fal_ai/cost_calculator.py b/litellm/llms/fal_ai/cost_calculator.py index b7caae3834..9cdd0cd485 100644 --- a/litellm/llms/fal_ai/cost_calculator.py +++ b/litellm/llms/fal_ai/cost_calculator.py @@ -22,5 +22,6 @@ def cost_calculator( num_images = len(image_response.data) return output_cost_per_image * num_images else: - raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}") - + raise ValueError( + f"image_response must be of type ImageResponse got type={type(image_response)}" + ) diff --git a/litellm/llms/fal_ai/image_generation/__init__.py b/litellm/llms/fal_ai/image_generation/__init__.py index 27817ae5a5..9deeb403c4 100644 --- a/litellm/llms/fal_ai/image_generation/__init__.py +++ b/litellm/llms/fal_ai/image_generation/__init__.py @@ -35,15 +35,15 @@ __all__ = [ def get_fal_ai_image_generation_config(model: str) -> BaseImageGenerationConfig: """ Get the appropriate Fal AI image generation configuration based on the model. - + Args: model: The Fal AI model name (e.g., "fal-ai/imagen4/preview", "fal-ai/recraft/v3/text-to-image") - + Returns: The appropriate configuration class for the specified model """ model_lower = model.lower() - + # Map model names to their corresponding configuration classes if "imagen4" in model_lower or "imagen-4" in model_lower: return FalAIImagen4Config() @@ -55,7 +55,11 @@ def get_fal_ai_image_generation_config(model: str) -> BaseImageGenerationConfig: if "ultra" in model_lower: return FalAIFluxProV11UltraConfig() return FalAIFluxProV11Config() - elif "flux/schnell" in model_lower or "flux-schnell" in model_lower or "schnell" in model_lower: + elif ( + "flux/schnell" in model_lower + or "flux-schnell" in model_lower + or "schnell" in model_lower + ): return FalAIFluxSchnellConfig() elif "bytedance/seedream" in model_lower: return FalAIBytedanceSeedreamV3Config() @@ -65,7 +69,6 @@ def get_fal_ai_image_generation_config(model: str) -> BaseImageGenerationConfig: return FalAIIdeogramV3Config() elif "stable-diffusion" in model_lower: return FalAIStableDiffusionConfig() - + # Default to generic Fal AI configuration return FalAIImageGenerationConfig() - diff --git a/litellm/llms/fal_ai/image_generation/bria_transformation.py b/litellm/llms/fal_ai/image_generation/bria_transformation.py index cb5aa6b761..dd6e737324 100644 --- a/litellm/llms/fal_ai/image_generation/bria_transformation.py +++ b/litellm/llms/fal_ai/image_generation/bria_transformation.py @@ -18,15 +18,16 @@ else: class FalAIBriaConfig(FalAIBaseConfig): """ Configuration for Bria Text-to-Image 3.2 model. - + Bria 3.2 is a commercial-grade text-to-image model with prompt enhancement and multiple aspect ratio options. - + Model endpoint: bria/text-to-image/3.2 Documentation: https://fal.ai/models/bria/text-to-image/3.2 """ + IMAGE_GENERATION_ENDPOINT: str = "bria/text-to-image/3.2" - + def get_supported_openai_params( self, model: str ) -> List[OpenAIImageGenerationOptionalParams]: @@ -38,7 +39,7 @@ class FalAIBriaConfig(FalAIBaseConfig): "response_format", "size", ] - + def map_openai_params( self, non_default_params: dict, @@ -48,26 +49,26 @@ class FalAIBriaConfig(FalAIBaseConfig): ) -> dict: """ Map OpenAI parameters to Bria 3.2 parameters. - + Mappings: - size -> aspect_ratio (1:1, 2:3, 3:2, 3:4, 4:3, 4:5, 5:4, 9:16, 16:9) - response_format -> ignored (Bria returns URLs) - n -> ignored (Bria doesn't support multiple images in one call) """ supported_params = self.get_supported_openai_params(model) - + # Map OpenAI params to Bria params param_mapping = { "size": "aspect_ratio", } - + for k in non_default_params.keys(): if k not in optional_params.keys(): if k in supported_params: # Use mapped parameter name if exists mapped_key = param_mapping.get(k, k) mapped_value = non_default_params[k] - + # Transform specific parameters if k == "response_format": # Bria always returns URLs, so we can ignore this @@ -78,7 +79,7 @@ class FalAIBriaConfig(FalAIBaseConfig): elif k == "size": # Map OpenAI size format to Bria aspect ratio mapped_value = self._map_aspect_ratio(mapped_value) - + optional_params[mapped_key] = mapped_value elif drop_params: pass @@ -92,7 +93,7 @@ class FalAIBriaConfig(FalAIBaseConfig): def _map_aspect_ratio(self, size: str) -> str: """ Map OpenAI size format to Bria aspect ratio format. - + OpenAI format: "1024x1024", "1792x1024", etc. Bria format: "1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9" """ @@ -107,20 +108,20 @@ class FalAIBriaConfig(FalAIBaseConfig): "1280x960": "4:3", "960x1280": "3:4", } - + if size in size_to_aspect_ratio: return size_to_aspect_ratio[size] - + # Parse custom size format "WIDTHxHEIGHT" and calculate aspect ratio if "x" in size: try: width_str, height_str = size.split("x") width = int(width_str) height = int(height_str) - + # Calculate aspect ratio and find closest match ratio = width / height - + # Map to closest supported aspect ratio if 0.95 <= ratio <= 1.05: # Close to 1:1 return "1:1" @@ -142,7 +143,7 @@ class FalAIBriaConfig(FalAIBaseConfig): return "4:5" except (ValueError, AttributeError, ZeroDivisionError): pass - + # Default to 1:1 return "1:1" @@ -156,10 +157,10 @@ class FalAIBriaConfig(FalAIBaseConfig): ) -> dict: """ Transform the image generation request to Bria 3.2 request body. - + Required parameters: - prompt: Prompt for image generation - + Optional parameters: - aspect_ratio: "1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9" (default: "1:1") - prompt_enhancer: Improve the prompt (default: true) @@ -174,7 +175,7 @@ class FalAIBriaConfig(FalAIBaseConfig): "prompt": prompt, **optional_params, } - + return bria_request_body def transform_image_generation_response( @@ -192,7 +193,7 @@ class FalAIBriaConfig(FalAIBaseConfig): ) -> ImageResponse: """ Transform the Bria 3.2 response to litellm ImageResponse format. - + Expected response format: { "image": { @@ -213,10 +214,10 @@ class FalAIBriaConfig(FalAIBaseConfig): status_code=raw_response.status_code, headers=raw_response.headers, ) - + if not model_response.data: model_response.data = [] - + # Handle Bria response format - uses "image" (singular) not "images" image_data = response_data.get("image") if image_data and isinstance(image_data, dict): @@ -226,6 +227,5 @@ class FalAIBriaConfig(FalAIBaseConfig): b64_json=None, # Bria returns URLs only ) ) - - return model_response + return model_response diff --git a/litellm/llms/fal_ai/image_generation/bytedance_transformation.py b/litellm/llms/fal_ai/image_generation/bytedance_transformation.py index d6aa242edc..b52d08dd9e 100644 --- a/litellm/llms/fal_ai/image_generation/bytedance_transformation.py +++ b/litellm/llms/fal_ai/image_generation/bytedance_transformation.py @@ -102,5 +102,3 @@ class FalAIBytedanceDreaminaV31Config(FalAIBytedanceBaseConfig): """ IMAGE_GENERATION_ENDPOINT: str = "fal-ai/bytedance/dreamina/v3.1/text-to-image" - - diff --git a/litellm/llms/fal_ai/image_generation/flux_pro_v11_transformation.py b/litellm/llms/fal_ai/image_generation/flux_pro_v11_transformation.py index 682ee0c267..5226419a29 100644 --- a/litellm/llms/fal_ai/image_generation/flux_pro_v11_transformation.py +++ b/litellm/llms/fal_ai/image_generation/flux_pro_v11_transformation.py @@ -87,5 +87,3 @@ class FalAIFluxProV11Config(FalAIFluxProV11UltraConfig): pass return "landscape_4_3" - - diff --git a/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py b/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py index 664f11d40d..fef292d331 100644 --- a/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py +++ b/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py @@ -18,15 +18,16 @@ else: class FalAIFluxProV11UltraConfig(FalAIBaseConfig): """ Configuration for Fal AI Flux Pro v1.1-ultra model. - + FLUX Pro v1.1-ultra is a high-quality text-to-image model with enhanced detail and support for image prompts. - + Model endpoint: fal-ai/flux-pro/v1.1-ultra Documentation: https://fal.ai/models/fal-ai/flux-pro/v1.1-ultra """ + IMAGE_GENERATION_ENDPOINT: str = "fal-ai/flux-pro/v1.1-ultra" - + def get_supported_openai_params( self, model: str ) -> List[OpenAIImageGenerationOptionalParams]: @@ -38,7 +39,7 @@ class FalAIFluxProV11UltraConfig(FalAIBaseConfig): "response_format", "size", ] - + def map_openai_params( self, non_default_params: dict, @@ -48,28 +49,28 @@ class FalAIFluxProV11UltraConfig(FalAIBaseConfig): ) -> dict: """ Map OpenAI parameters to Flux Pro v1.1-ultra parameters. - + Mappings: - n -> num_images (1-4, default 1) - response_format -> output_format (jpeg or png) - size -> aspect_ratio (21:9, 16:9, 4:3, 3:2, 1:1, 2:3, 3:4, 9:16, 9:21) """ supported_params = self.get_supported_openai_params(model) - + # Map OpenAI params to Flux Pro v1.1-ultra params param_mapping = { "n": "num_images", "response_format": "output_format", "size": "aspect_ratio", } - + for k in non_default_params.keys(): if k not in optional_params.keys(): if k in supported_params: # Use mapped parameter name if exists mapped_key = param_mapping.get(k, k) mapped_value = non_default_params[k] - + # Transform specific parameters if k == "response_format": # Map OpenAI response formats to image formats @@ -78,7 +79,7 @@ class FalAIFluxProV11UltraConfig(FalAIBaseConfig): elif k == "size": # Map OpenAI size format to Flux aspect ratio mapped_value = self._map_aspect_ratio(mapped_value) - + optional_params[mapped_key] = mapped_value elif drop_params: pass @@ -92,10 +93,10 @@ class FalAIFluxProV11UltraConfig(FalAIBaseConfig): def _map_aspect_ratio(self, size: str) -> str: """ Map OpenAI size format to Flux Pro aspect ratio format. - + OpenAI format: "1024x1024", "1792x1024", etc. Flux format: "21:9", "16:9", "4:3", "3:2", "1:1", "2:3", "3:4", "9:16", "9:21" - + Default: "16:9" """ # Map common OpenAI sizes to Flux aspect ratios @@ -111,20 +112,20 @@ class FalAIFluxProV11UltraConfig(FalAIBaseConfig): "2048x876": "21:9", "876x2048": "9:21", } - + if size in size_to_aspect_ratio: return size_to_aspect_ratio[size] - + # Parse custom size format "WIDTHxHEIGHT" and calculate aspect ratio if "x" in size: try: width_str, height_str = size.split("x") width = int(width_str) height = int(height_str) - + # Calculate aspect ratio and find closest match ratio = width / height - + # Map to closest supported aspect ratio if 0.95 <= ratio <= 1.05: # Close to 1:1 return "1:1" @@ -146,7 +147,7 @@ class FalAIFluxProV11UltraConfig(FalAIBaseConfig): return "9:21" except (ValueError, AttributeError, ZeroDivisionError): pass - + # Default to 16:9 return "16:9" @@ -160,10 +161,10 @@ class FalAIFluxProV11UltraConfig(FalAIBaseConfig): ) -> dict: """ Transform the image generation request to Flux Pro v1.1-ultra request body. - + Required parameters: - prompt: The prompt to generate an image from - + Optional parameters: - num_images: Number of images (1-4, default: 1) - aspect_ratio: Aspect ratio (default: "16:9") @@ -181,7 +182,7 @@ class FalAIFluxProV11UltraConfig(FalAIBaseConfig): "prompt": prompt, **optional_params, } - + return flux_pro_request_body def transform_image_generation_response( @@ -199,7 +200,7 @@ class FalAIFluxProV11UltraConfig(FalAIBaseConfig): ) -> ImageResponse: """ Transform the Flux Pro v1.1-ultra response to litellm ImageResponse format. - + Expected response format: { "images": [ @@ -224,10 +225,10 @@ class FalAIFluxProV11UltraConfig(FalAIBaseConfig): status_code=raw_response.status_code, headers=raw_response.headers, ) - + if not model_response.data: model_response.data = [] - + # Handle Flux Pro v1.1-ultra response format images = response_data.get("images", []) if isinstance(images, list): @@ -247,7 +248,7 @@ class FalAIFluxProV11UltraConfig(FalAIBaseConfig): b64_json=None, ) ) - + # Add additional metadata from Flux Pro response if hasattr(model_response, "_hidden_params"): if "seed" in response_data: @@ -258,6 +259,5 @@ class FalAIFluxProV11UltraConfig(FalAIBaseConfig): model_response._hidden_params["has_nsfw_concepts"] = response_data[ "has_nsfw_concepts" ] - - return model_response + return model_response diff --git a/litellm/llms/fal_ai/image_generation/flux_schnell_transformation.py b/litellm/llms/fal_ai/image_generation/flux_schnell_transformation.py index ed6ed37fb4..7a59fae6c1 100644 --- a/litellm/llms/fal_ai/image_generation/flux_schnell_transformation.py +++ b/litellm/llms/fal_ai/image_generation/flux_schnell_transformation.py @@ -85,4 +85,3 @@ class FalAIFluxSchnellConfig(FalAIFluxProV11UltraConfig): pass return "landscape_4_3" - diff --git a/litellm/llms/fal_ai/image_generation/ideogram_v3_transformation.py b/litellm/llms/fal_ai/image_generation/ideogram_v3_transformation.py index f05ffa888e..14e136d5d6 100644 --- a/litellm/llms/fal_ai/image_generation/ideogram_v3_transformation.py +++ b/litellm/llms/fal_ai/image_generation/ideogram_v3_transformation.py @@ -189,5 +189,3 @@ class FalAIIdeogramV3Config(FalAIBaseConfig): model_response._hidden_params["seed"] = response_data["seed"] return model_response - - diff --git a/litellm/llms/fal_ai/image_generation/imagen4_transformation.py b/litellm/llms/fal_ai/image_generation/imagen4_transformation.py index 4e7708c9f4..ea6e7c1f3c 100644 --- a/litellm/llms/fal_ai/image_generation/imagen4_transformation.py +++ b/litellm/llms/fal_ai/image_generation/imagen4_transformation.py @@ -18,18 +18,19 @@ else: class FalAIImagen4Config(FalAIBaseConfig): """ Configuration for Fal AI Imagen4 model. - + Google's highest quality image generation model available through Fal AI. - + Model variants: - fal-ai/imagen4/preview (Standard): $0.05 per image - fal-ai/imagen4/preview/fast (Fast): $0.02 per image - fal-ai/imagen4/preview/ultra (Ultra): $0.06 per image - + Documentation: https://fal.ai/models/fal-ai/imagen4/preview """ + IMAGE_GENERATION_ENDPOINT: str = "fal-ai/imagen4/preview" - + def get_supported_openai_params( self, model: str ) -> List[OpenAIImageGenerationOptionalParams]: @@ -41,7 +42,7 @@ class FalAIImagen4Config(FalAIBaseConfig): "response_format", "size", ] - + def map_openai_params( self, non_default_params: dict, @@ -51,27 +52,27 @@ class FalAIImagen4Config(FalAIBaseConfig): ) -> dict: """ Map OpenAI parameters to Imagen4 parameters. - + Mappings: - n -> num_images (1-4, default 1) - size -> aspect_ratio (1:1, 16:9, 9:16, 3:4, 4:3) - response_format -> ignored (Imagen4 returns URLs) """ supported_params = self.get_supported_openai_params(model) - + # Map OpenAI params to Imagen4 params param_mapping = { "n": "num_images", "size": "aspect_ratio", } - + for k in non_default_params.keys(): if k not in optional_params.keys(): if k in supported_params: # Use mapped parameter name if exists mapped_key = param_mapping.get(k, k) mapped_value = non_default_params[k] - + # Transform specific parameters if k == "response_format": # Imagen4 always returns URLs, so we can ignore this @@ -79,7 +80,7 @@ class FalAIImagen4Config(FalAIBaseConfig): elif k == "size": # Map OpenAI size format to Imagen4 aspect ratio mapped_value = self._map_aspect_ratio(mapped_value) - + optional_params[mapped_key] = mapped_value elif drop_params: pass @@ -93,10 +94,10 @@ class FalAIImagen4Config(FalAIBaseConfig): def _map_aspect_ratio(self, size: str) -> str: """ Map OpenAI size format to Imagen4 aspect ratio format. - + OpenAI format: "1024x1024", "1792x1024", etc. Imagen4 format: "1:1", "16:9", "9:16", "3:4", "4:3" - + Available aspect ratios: - 1:1 (default) - 16:9 @@ -113,20 +114,20 @@ class FalAIImagen4Config(FalAIBaseConfig): "1024x768": "4:3", "768x1024": "3:4", } - + if size in size_to_aspect_ratio: return size_to_aspect_ratio[size] - + # Parse custom size format "WIDTHxHEIGHT" and calculate aspect ratio if "x" in size: try: width_str, height_str = size.split("x") width = int(width_str) height = int(height_str) - + # Calculate aspect ratio and find closest match ratio = width / height - + # Map to closest supported aspect ratio if 0.95 <= ratio <= 1.05: # Close to 1:1 return "1:1" @@ -140,7 +141,7 @@ class FalAIImagen4Config(FalAIBaseConfig): return "3:4" except (ValueError, AttributeError, ZeroDivisionError): pass - + # Default to 1:1 return "1:1" @@ -154,10 +155,10 @@ class FalAIImagen4Config(FalAIBaseConfig): ) -> dict: """ Transform the image generation request to Imagen4 request body. - + Required parameters: - prompt: The text prompt describing what you want to see - + Optional parameters: - aspect_ratio: "1:1", "16:9", "9:16", "3:4", "4:3" (default: "1:1") - num_images: Number of images (1-4, default: 1) @@ -169,7 +170,7 @@ class FalAIImagen4Config(FalAIBaseConfig): "prompt": prompt, **optional_params, } - + return imagen4_request_body def transform_image_generation_response( @@ -187,7 +188,7 @@ class FalAIImagen4Config(FalAIBaseConfig): ) -> ImageResponse: """ Transform the Imagen4 response to litellm ImageResponse format. - + Expected response format: { "images": [ @@ -209,10 +210,10 @@ class FalAIImagen4Config(FalAIBaseConfig): status_code=raw_response.status_code, headers=raw_response.headers, ) - + if not model_response.data: model_response.data = [] - + # Handle Imagen4 response format images = response_data.get("images", []) if isinstance(images, list): @@ -232,11 +233,10 @@ class FalAIImagen4Config(FalAIBaseConfig): b64_json=None, ) ) - + # Add seed metadata from Imagen4 response if hasattr(model_response, "_hidden_params"): if "seed" in response_data: model_response._hidden_params["seed"] = response_data["seed"] - - return model_response + return model_response diff --git a/litellm/llms/fal_ai/image_generation/recraft_v3_transformation.py b/litellm/llms/fal_ai/image_generation/recraft_v3_transformation.py index 572a8a0f1c..72ee165b51 100644 --- a/litellm/llms/fal_ai/image_generation/recraft_v3_transformation.py +++ b/litellm/llms/fal_ai/image_generation/recraft_v3_transformation.py @@ -18,15 +18,16 @@ else: class FalAIRecraftV3Config(FalAIBaseConfig): """ Configuration for Fal AI Recraft v3 Text-to-Image model. - + Recraft v3 is a text-to-image model with multiple style options including realistic images, digital illustrations, and vector illustrations. - + Model endpoint: fal-ai/recraft/v3/text-to-image Documentation: https://fal.ai/models/fal-ai/recraft/v3/text-to-image """ + IMAGE_GENERATION_ENDPOINT: str = "fal-ai/recraft/v3/text-to-image" - + def get_supported_openai_params( self, model: str ) -> List[OpenAIImageGenerationOptionalParams]: @@ -38,7 +39,7 @@ class FalAIRecraftV3Config(FalAIBaseConfig): "response_format", "size", ] - + def map_openai_params( self, non_default_params: dict, @@ -48,26 +49,26 @@ class FalAIRecraftV3Config(FalAIBaseConfig): ) -> dict: """ Map OpenAI parameters to Recraft v3 parameters. - + Mappings: - size -> image_size (can be preset or custom width/height) - response_format -> ignored (Recraft returns URLs) - n -> ignored (Recraft doesn't support multiple images) """ supported_params = self.get_supported_openai_params(model) - + # Map OpenAI params to Recraft v3 params param_mapping = { "size": "image_size", } - + for k in non_default_params.keys(): if k not in optional_params.keys(): if k in supported_params: # Use mapped parameter name if exists mapped_key = param_mapping.get(k, k) mapped_value = non_default_params[k] - + # Transform specific parameters if k == "response_format": # Recraft always returns URLs, so we can ignore this @@ -78,7 +79,7 @@ class FalAIRecraftV3Config(FalAIBaseConfig): elif k == "size": # Map OpenAI size format to Recraft image_size mapped_value = self._map_image_size(mapped_value) - + optional_params[mapped_key] = mapped_value elif drop_params: pass @@ -92,10 +93,10 @@ class FalAIRecraftV3Config(FalAIBaseConfig): def _map_image_size(self, size: str) -> Any: """ Map OpenAI size format to Recraft v3 image_size format. - + OpenAI format: "1024x1024", "1792x1024", etc. Recraft format: Can be preset strings or {"width": int, "height": int} - + Available presets: - square_hd (default) - square @@ -113,10 +114,10 @@ class FalAIRecraftV3Config(FalAIBaseConfig): "1024x768": "landscape_4_3", "1024x576": "landscape_16_9", } - + if size in size_mapping: return size_mapping[size] - + # Parse custom size format "WIDTHxHEIGHT" if "x" in size: try: @@ -127,7 +128,7 @@ class FalAIRecraftV3Config(FalAIBaseConfig): } except (ValueError, AttributeError): pass - + # Default to square_hd return "square_hd" @@ -141,10 +142,10 @@ class FalAIRecraftV3Config(FalAIBaseConfig): ) -> dict: """ Transform the image generation request to Recraft v3 request body. - + Required parameters: - prompt: Text prompt (max 1000 characters) - + Optional parameters: - image_size: Preset or {"width": int, "height": int} (default: "square_hd") - style: Style preset (default: "realistic_image") @@ -152,14 +153,14 @@ class FalAIRecraftV3Config(FalAIBaseConfig): - colors: Array of RGB color objects [{"r": 0-255, "g": 0-255, "b": 0-255}] - enable_safety_checker: Enable safety checker (default: false) - style_id: UUID for custom style reference - + Note: Vector illustrations cost 2X as much. """ recraft_request_body = { "prompt": prompt, **optional_params, } - + return recraft_request_body def transform_image_generation_response( @@ -177,7 +178,7 @@ class FalAIRecraftV3Config(FalAIBaseConfig): ) -> ImageResponse: """ Transform the Recraft v3 response to litellm ImageResponse format. - + Expected response format: { "images": [ @@ -198,10 +199,10 @@ class FalAIRecraftV3Config(FalAIBaseConfig): status_code=raw_response.status_code, headers=raw_response.headers, ) - + if not model_response.data: model_response.data = [] - + # Handle Recraft v3 response format images = response_data.get("images", []) if isinstance(images, list): @@ -221,6 +222,5 @@ class FalAIRecraftV3Config(FalAIBaseConfig): b64_json=None, ) ) - - return model_response + return model_response diff --git a/litellm/llms/fal_ai/image_generation/stable_diffusion_transformation.py b/litellm/llms/fal_ai/image_generation/stable_diffusion_transformation.py index 10e2c6b416..f0077c6a67 100644 --- a/litellm/llms/fal_ai/image_generation/stable_diffusion_transformation.py +++ b/litellm/llms/fal_ai/image_generation/stable_diffusion_transformation.py @@ -18,17 +18,18 @@ else: class FalAIStableDiffusionConfig(FalAIBaseConfig): """ Configuration for Fal AI Stable Diffusion models. - + Supports Stable Diffusion v3.5 variants and other Stable Diffusion models on Fal AI. - + Example models: - fal-ai/stable-diffusion-v35-medium - fal-ai/stable-diffusion-v35-large - + Documentation: https://fal.ai/models/fal-ai/stable-diffusion-v35-medium """ + IMAGE_GENERATION_ENDPOINT: str = "" # Will be set from model name - + def get_complete_url( self, api_base: Optional[str], @@ -40,19 +41,17 @@ class FalAIStableDiffusionConfig(FalAIBaseConfig): ) -> str: """ Get the complete url for the request. - + For Stable Diffusion models, extract the endpoint from the model name. """ from litellm.secret_managers.main import get_secret_str - + complete_url: str = ( - api_base - or get_secret_str("FAL_AI_API_BASE") - or self.DEFAULT_BASE_URL + api_base or get_secret_str("FAL_AI_API_BASE") or self.DEFAULT_BASE_URL ) - + complete_url = complete_url.rstrip("/") - + # Extract endpoint from model name # e.g., "fal-ai/stable-diffusion-v35-medium" or "stable-diffusion-v35-medium" endpoint = model @@ -62,10 +61,10 @@ class FalAIStableDiffusionConfig(FalAIBaseConfig): elif not model.startswith("fal-ai/"): # If model is just "stable-diffusion-v35-medium", prepend fal-ai endpoint = f"fal-ai/{model}" - + complete_url = f"{complete_url}/{endpoint}" return complete_url - + def get_supported_openai_params( self, model: str ) -> List[OpenAIImageGenerationOptionalParams]: @@ -77,7 +76,7 @@ class FalAIStableDiffusionConfig(FalAIBaseConfig): "response_format", "size", ] - + def map_openai_params( self, non_default_params: dict, @@ -87,28 +86,28 @@ class FalAIStableDiffusionConfig(FalAIBaseConfig): ) -> dict: """ Map OpenAI parameters to Stable Diffusion parameters. - + Mappings: - n -> num_images (1-4, default 1) - response_format -> output_format (jpeg or png) - size -> image_size (can be preset or custom width/height) """ supported_params = self.get_supported_openai_params(model) - + # Map OpenAI params to Stable Diffusion params param_mapping = { "n": "num_images", "response_format": "output_format", "size": "image_size", } - + for k in non_default_params.keys(): if k not in optional_params.keys(): if k in supported_params: # Use mapped parameter name if exists mapped_key = param_mapping.get(k, k) mapped_value = non_default_params[k] - + # Transform specific parameters if k == "response_format": # Map OpenAI response formats to image formats @@ -117,7 +116,7 @@ class FalAIStableDiffusionConfig(FalAIBaseConfig): elif k == "size": # Map OpenAI size format to Stable Diffusion image_size mapped_value = self._map_image_size(mapped_value) - + optional_params[mapped_key] = mapped_value elif drop_params: pass @@ -131,10 +130,10 @@ class FalAIStableDiffusionConfig(FalAIBaseConfig): def _map_image_size(self, size: str) -> Any: """ Map OpenAI size format to Stable Diffusion image_size format. - + OpenAI format: "1024x1024", "1792x1024", etc. Stable Diffusion format: Can be preset strings or {"width": int, "height": int} - + Available presets: - square_hd - square @@ -152,10 +151,10 @@ class FalAIStableDiffusionConfig(FalAIBaseConfig): "1024x768": "landscape_4_3", "1024x576": "landscape_16_9", } - + if size in size_mapping: return size_mapping[size] - + # Parse custom size format "WIDTHxHEIGHT" if "x" in size: try: @@ -166,7 +165,7 @@ class FalAIStableDiffusionConfig(FalAIBaseConfig): } except (ValueError, AttributeError): pass - + # Default to landscape_4_3 return "landscape_4_3" @@ -180,10 +179,10 @@ class FalAIStableDiffusionConfig(FalAIBaseConfig): ) -> dict: """ Transform the image generation request to Stable Diffusion request body. - + Required parameters: - prompt: The prompt to generate an image from - + Optional parameters: - num_images: Number of images (1-4, default: 1) - image_size: Size preset or {"width": int, "height": int} (default: landscape_4_3) @@ -199,7 +198,7 @@ class FalAIStableDiffusionConfig(FalAIBaseConfig): "prompt": prompt, **optional_params, } - + return stable_diffusion_request_body def transform_image_generation_response( @@ -217,7 +216,7 @@ class FalAIStableDiffusionConfig(FalAIBaseConfig): ) -> ImageResponse: """ Transform the Stable Diffusion response to litellm ImageResponse format. - + Expected response format: { "images": [ @@ -242,10 +241,10 @@ class FalAIStableDiffusionConfig(FalAIBaseConfig): status_code=raw_response.status_code, headers=raw_response.headers, ) - + if not model_response.data: model_response.data = [] - + # Handle Stable Diffusion response format images = response_data.get("images", []) if isinstance(images, list): @@ -265,7 +264,7 @@ class FalAIStableDiffusionConfig(FalAIBaseConfig): b64_json=None, ) ) - + # Add additional metadata from Stable Diffusion response if hasattr(model_response, "_hidden_params"): if "seed" in response_data: @@ -276,6 +275,5 @@ class FalAIStableDiffusionConfig(FalAIBaseConfig): model_response._hidden_params["has_nsfw_concepts"] = response_data[ "has_nsfw_concepts" ] - - return model_response + return model_response diff --git a/litellm/llms/fal_ai/image_generation/transformation.py b/litellm/llms/fal_ai/image_generation/transformation.py index 04b7b16752..4a0dea48a1 100644 --- a/litellm/llms/fal_ai/image_generation/transformation.py +++ b/litellm/llms/fal_ai/image_generation/transformation.py @@ -25,6 +25,7 @@ class FalAIBaseConfig(BaseImageGenerationConfig): Base configuration for Fal AI image generation models. Handles common functionality like URL construction and authentication. """ + DEFAULT_BASE_URL: str = "https://fal.run" IMAGE_GENERATION_ENDPOINT: str = "" @@ -43,9 +44,7 @@ class FalAIBaseConfig(BaseImageGenerationConfig): Some providers need `model` in `api_base` """ complete_url: str = ( - api_base - or get_secret_str("FAL_AI_API_BASE") - or self.DEFAULT_BASE_URL + api_base or get_secret_str("FAL_AI_API_BASE") or self.DEFAULT_BASE_URL ) complete_url = complete_url.rstrip("/") @@ -63,14 +62,11 @@ class FalAIBaseConfig(BaseImageGenerationConfig): api_key: Optional[str] = None, api_base: Optional[str] = None, ) -> dict: - final_api_key: Optional[str] = ( - api_key or - get_secret_str("FAL_AI_API_KEY") - ) + final_api_key: Optional[str] = api_key or get_secret_str("FAL_AI_API_KEY") if not final_api_key: raise ValueError("FAL_AI_API_KEY is not set") - - headers["Authorization"] = f"Key {final_api_key}" + + headers["Authorization"] = f"Key {final_api_key}" return headers def transform_image_generation_response( @@ -99,23 +95,27 @@ class FalAIBaseConfig(BaseImageGenerationConfig): ) if not model_response.data: model_response.data = [] - + # Handle fal.ai response format images = response_data.get("images", []) if isinstance(images, list): for image_data in images: if isinstance(image_data, dict): - model_response.data.append(ImageObject( - url=image_data.get("url", None), - b64_json=image_data.get("b64_json", None), - )) + model_response.data.append( + ImageObject( + url=image_data.get("url", None), + b64_json=image_data.get("b64_json", None), + ) + ) elif isinstance(image_data, str): # If images is just a list of URLs - model_response.data.append(ImageObject( - url=image_data, - b64_json=None, - )) - + model_response.data.append( + ImageObject( + url=image_data, + b64_json=None, + ) + ) + return model_response @@ -123,7 +123,7 @@ class FalAIImageGenerationConfig(FalAIBaseConfig): """ Default Fal AI image generation configuration for generic models. """ - + def get_supported_openai_params( self, model: str ) -> List[OpenAIImageGenerationOptionalParams]: @@ -135,7 +135,7 @@ class FalAIImageGenerationConfig(FalAIBaseConfig): "response_format", "size", ] - + def map_openai_params( self, non_default_params: dict, @@ -173,4 +173,3 @@ class FalAIImageGenerationConfig(FalAIBaseConfig): **optional_params, } return fal_ai_image_generation_request_body - diff --git a/litellm/llms/firecrawl/__init__.py b/litellm/llms/firecrawl/__init__.py index bacf1eac07..b43d2da321 100644 --- a/litellm/llms/firecrawl/__init__.py +++ b/litellm/llms/firecrawl/__init__.py @@ -4,4 +4,3 @@ Firecrawl API integration module. from litellm.llms.firecrawl.search.transformation import FirecrawlSearchConfig __all__ = ["FirecrawlSearchConfig"] - diff --git a/litellm/llms/firecrawl/search/__init__.py b/litellm/llms/firecrawl/search/__init__.py index 999dce655d..46619d05b6 100644 --- a/litellm/llms/firecrawl/search/__init__.py +++ b/litellm/llms/firecrawl/search/__init__.py @@ -4,4 +4,3 @@ Firecrawl Search API module. from litellm.llms.firecrawl.search.transformation import FirecrawlSearchConfig __all__ = ["FirecrawlSearchConfig"] - diff --git a/litellm/llms/firecrawl/search/transformation.py b/litellm/llms/firecrawl/search/transformation.py index af501a8eac..61b589218c 100644 --- a/litellm/llms/firecrawl/search/transformation.py +++ b/litellm/llms/firecrawl/search/transformation.py @@ -18,6 +18,7 @@ from litellm.secret_managers.main import get_secret_str class _FirecrawlSearchRequestRequired(TypedDict): """Required fields for Firecrawl Search API request.""" + query: str # Required - search query @@ -26,9 +27,14 @@ class FirecrawlSearchRequest(_FirecrawlSearchRequestRequired, total=False): Firecrawl Search API request format. Based on: https://docs.firecrawl.dev/api-reference/endpoint/search """ + limit: int # Optional - maximum number of results to return (default 5, max 100) - sources: List[str] # Optional - sources to search ('web', 'images', 'news'), default ['web'] - categories: List[Dict[str, str]] # Optional - categories to filter by (github, research, pdf) + sources: List[ + str + ] # Optional - sources to search ('web', 'images', 'news'), default ['web'] + categories: List[ + Dict[str, str] + ] # Optional - categories to filter by (github, research, pdf) tbs: str # Optional - time-based search parameter location: str # Optional - location parameter for geo-targeting country: str # Optional - ISO country code (default 'US') @@ -39,11 +45,11 @@ class FirecrawlSearchRequest(_FirecrawlSearchRequestRequired, total=False): class FirecrawlSearchConfig(BaseSearchConfig): FIRECRAWL_API_BASE = "https://api.firecrawl.dev/v2" - + @staticmethod def ui_friendly_name() -> str: return "Firecrawl" - + def validate_environment( self, headers: Dict, @@ -56,7 +62,9 @@ class FirecrawlSearchConfig(BaseSearchConfig): """ api_key = api_key or get_secret_str("FIRECRAWL_API_KEY") if not api_key: - raise ValueError("FIRECRAWL_API_KEY is not set. Set `FIRECRAWL_API_KEY` environment variable.") + raise ValueError( + "FIRECRAWL_API_KEY is not set. Set `FIRECRAWL_API_KEY` environment variable." + ) headers["Authorization"] = f"Bearer {api_key}" headers["Content-Type"] = "application/json" return headers @@ -71,14 +79,15 @@ class FirecrawlSearchConfig(BaseSearchConfig): """ Get complete URL for Search endpoint. """ - api_base = api_base or get_secret_str("FIRECRAWL_API_BASE") or self.FIRECRAWL_API_BASE - + api_base = ( + api_base or get_secret_str("FIRECRAWL_API_BASE") or self.FIRECRAWL_API_BASE + ) + # Append "/search" to the api base if it's not already there if not api_base.endswith("/search"): api_base = f"{api_base}/search" return api_base - def transform_search_request( self, @@ -88,20 +97,20 @@ class FirecrawlSearchConfig(BaseSearchConfig): ) -> Dict: """ Transform Search request to Firecrawl API format. - + Transforms Perplexity unified spec parameters: - query → query (same) - max_results → limit - search_domain_filter → (not directly supported, can use scrapeOptions) - country → country - max_tokens_per_page → (not applicable, ignored) - + All other Firecrawl-specific parameters are passed through as-is. - + Args: query: Search query (string or list of strings). Firecrawl only supports single string queries. optional_params: Optional parameters for the request - + Returns: Dict with typed request data following FirecrawlSearchRequest spec """ @@ -112,30 +121,33 @@ class FirecrawlSearchConfig(BaseSearchConfig): request_data: FirecrawlSearchRequest = { "query": query, } - + # Transform Perplexity unified spec parameters to Firecrawl format if "max_results" in optional_params: request_data["limit"] = optional_params["max_results"] - + if "country" in optional_params: request_data["country"] = optional_params["country"] - + # Convert to dict before dynamic key assignments result_data = dict(request_data) - + # pass through all other parameters as-is for param, value in optional_params.items(): - if param not in self.get_supported_perplexity_optional_params() and param not in result_data: + if ( + param not in self.get_supported_perplexity_optional_params() + and param not in result_data + ): result_data[param] = value - + # By default, request markdown content if not explicitly specified # Firecrawl doesn't return content unless explicitly requested via scrapeOptions if "scrapeOptions" not in result_data: result_data["scrapeOptions"] = { "formats": ["markdown"], - "onlyMainContent": True + "onlyMainContent": True, } - + return result_data def transform_search_response( @@ -146,37 +158,37 @@ class FirecrawlSearchConfig(BaseSearchConfig): ) -> SearchResponse: """ Transform Firecrawl API response to LiteLLM unified SearchResponse format. - + Firecrawl → LiteLLM mappings: - data.web[].title → SearchResult.title - data.web[].url → SearchResult.url - data.web[].description OR data.web[].markdown → SearchResult.snippet - No date field in web results (set to None) - No last_updated field in Firecrawl response (set to None) - + Note: Firecrawl v2 returns results organized by source type (web, images, news). We primarily use web results for the unified format. - + Args: raw_response: Raw httpx response from Firecrawl API logging_obj: Logging object for tracking - + Returns: SearchResponse with standardized format """ response_json = raw_response.json() - + # Transform results to SearchResult objects results = [] - + # Process web results (primary source) data = response_json.get("data", {}) web_results = data.get("web", []) - + for result in web_results: # Use markdown if available, otherwise fall back to description snippet = result.get("markdown") or result.get("description", "") - + search_result = SearchResult( title=result.get("title", ""), url=result.get("url", ""), @@ -185,12 +197,12 @@ class FirecrawlSearchConfig(BaseSearchConfig): last_updated=None, # Firecrawl doesn't provide last_updated in response ) results.append(search_result) - + # Process news results if available (they have date field) news_results = data.get("news", []) for result in news_results: snippet = result.get("markdown") or result.get("snippet", "") - + search_result = SearchResult( title=result.get("title", ""), url=result.get("url", ""), @@ -199,9 +211,8 @@ class FirecrawlSearchConfig(BaseSearchConfig): last_updated=None, ) results.append(search_result) - + return SearchResponse( results=results, object="search", ) - diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 7ec32fecc4..8407e8ab69 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -257,30 +257,33 @@ class FireworksAIConfig(OpenAIGPTConfig): "gpt-oss-120b", "gpt-oss-20b", ] - + # Normalize model name - remove prefix if present normalized_model = model if model.startswith("fireworks_ai/"): normalized_model = model.replace("fireworks_ai/", "") if normalized_model.startswith("accounts/fireworks/models/"): - normalized_model = normalized_model.replace("accounts/fireworks/models/", "") - + normalized_model = normalized_model.replace( + "accounts/fireworks/models/", "" + ) + # Check if model supports reasoning supports_reasoning_value = any( - reasoning_model in normalized_model for reasoning_model in reasoning_supported_models + reasoning_model in normalized_model + for reasoning_model in reasoning_supported_models ) - + provider_specific_model_info: ProviderSpecificModelInfo = { "supports_function_calling": True, "supports_prompt_caching": True, # https://docs.fireworks.ai/guides/prompt-caching "supports_pdf_input": True, # via document inlining "supports_vision": True, # via document inlining } - + # Only include supports_reasoning if True if supports_reasoning_value: provider_specific_model_info["supports_reasoning"] = True - + return provider_specific_model_info def transform_request( @@ -426,8 +429,11 @@ class FireworksAIConfig(OpenAIGPTConfig): "FIREWORKS_ACCOUNT_ID is not set. Please set the environment variable, to query Fireworks AI's `/models` endpoint." ) + base = api_base.rstrip("/") + if base.endswith("/v1"): + base = base[: -len("/v1")] response = litellm.module_level_client.get( - url=f"{api_base}/v1/accounts/{account_id}/models", + url=f"{base}/v1/accounts/{account_id}/models", headers={"Authorization": f"Bearer {api_key}"}, ) diff --git a/litellm/llms/fireworks_ai/rerank/__init__.py b/litellm/llms/fireworks_ai/rerank/__init__.py index b8e99317a2..2312d016ab 100644 --- a/litellm/llms/fireworks_ai/rerank/__init__.py +++ b/litellm/llms/fireworks_ai/rerank/__init__.py @@ -1,2 +1 @@ # Fireworks AI Rerank - diff --git a/litellm/llms/fireworks_ai/rerank/transformation.py b/litellm/llms/fireworks_ai/rerank/transformation.py index e2893464bd..eb92399a05 100644 --- a/litellm/llms/fireworks_ai/rerank/transformation.py +++ b/litellm/llms/fireworks_ai/rerank/transformation.py @@ -75,26 +75,26 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig): "query": query, "documents": documents, } - + if top_n is not None: params["top_n"] = top_n - + if return_documents is not None: params["return_documents"] = return_documents - + # Fireworks AI doesn't support these params if rank_fields is not None: # Silently ignore rank_fields as Fireworks AI doesn't support it pass - + if max_chunks_per_doc is not None: # Silently ignore max_chunks_per_doc as Fireworks AI doesn't support it pass - + if max_tokens_per_doc is not None: # Silently ignore max_tokens_per_doc as Fireworks AI doesn't support it pass - + return params def validate_environment( # type: ignore[override] @@ -140,7 +140,7 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig): # Remove fireworks_ai/ prefix if present if model.startswith("fireworks_ai/"): model = model.replace("fireworks_ai/", "") - + # If model doesn't start with "fireworks/", add it # But don't add if it already has the prefix if not model.startswith("fireworks/"): @@ -152,11 +152,19 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig): "documents": optional_rerank_params["documents"], } - if "top_n" in optional_rerank_params and optional_rerank_params["top_n"] is not None: + if ( + "top_n" in optional_rerank_params + and optional_rerank_params["top_n"] is not None + ): request_data["top_n"] = optional_rerank_params["top_n"] - if "return_documents" in optional_rerank_params and optional_rerank_params["return_documents"] is not None: - request_data["return_documents"] = optional_rerank_params["return_documents"] + if ( + "return_documents" in optional_rerank_params + and optional_rerank_params["return_documents"] is not None + ): + request_data["return_documents"] = optional_rerank_params[ + "return_documents" + ] return request_data @@ -191,7 +199,7 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig): # { # "index": 0, # "relevance_score": 0.95, - # "document": "..." + # "document": "..." # } # ], # "usage": { @@ -203,9 +211,7 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig): # Extract usage information usage = raw_response_json.get("usage", {}) - _billed_units = RerankBilledUnits( - search_units=usage.get("total_tokens", 0) - ) + _billed_units = RerankBilledUnits(search_units=usage.get("total_tokens", 0)) _tokens = RerankTokens( input_tokens=usage.get("prompt_tokens", 0), output_tokens=usage.get("completion_tokens", 0), @@ -213,7 +219,9 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig): rerank_meta = RerankResponseMeta(billed_units=_billed_units, tokens=_tokens) # Extract results - Fireworks AI uses "data" instead of "results" - _results: Optional[List[dict]] = raw_response_json.get("data") or raw_response_json.get("results") + _results: Optional[List[dict]] = raw_response_json.get( + "data" + ) or raw_response_json.get("results") if _results is None: raise ValueError(f"No results found in the response={raw_response_json}") @@ -251,11 +259,14 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig): rerank_results.append(rerank_result) # Use model name as id if no id is provided - response_id = raw_response_json.get("id") or raw_response_json.get("model") or str(uuid.uuid4()) + response_id = ( + raw_response_json.get("id") + or raw_response_json.get("model") + or str(uuid.uuid4()) + ) return RerankResponse( id=response_id, results=rerank_results, meta=rerank_meta, ) - diff --git a/litellm/llms/gemini/chat/transformation.py b/litellm/llms/gemini/chat/transformation.py index d5a5ab667a..5f8dead204 100644 --- a/litellm/llms/gemini/chat/transformation.py +++ b/litellm/llms/gemini/chat/transformation.py @@ -126,13 +126,15 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig): image_obj = convert_to_anthropic_image_obj( _image_url, format=format ) - converted_image_url = convert_generic_image_chunk_to_openai_image_obj( - image_obj + converted_image_url = ( + convert_generic_image_chunk_to_openai_image_obj( + image_obj + ) ) if detail is not None: img_element["image_url"] = { # type: ignore "url": converted_image_url, - "detail": detail + "detail": detail, } else: img_element["image_url"] = converted_image_url # type: ignore diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 17b9c78123..87c107fab3 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -45,7 +45,11 @@ class GeminiModelInfo(BaseLLMModelInfo): @staticmethod def get_api_key(api_key: Optional[str] = None) -> Optional[str]: - return api_key or (get_secret_str("GOOGLE_API_KEY")) or (get_secret_str("GEMINI_API_KEY")) + return ( + api_key + or (get_secret_str("GOOGLE_API_KEY")) + or (get_secret_str("GEMINI_API_KEY")) + ) @staticmethod def get_base_model(model: str) -> Optional[str]: @@ -90,11 +94,11 @@ class GeminiModelInfo(BaseLLMModelInfo): return GeminiError( status_code=status_code, message=error_message, headers=headers ) - + def get_token_counter(self) -> Optional[BaseTokenCounter]: """ Factory method to create a token counter for this provider. - + Returns: Optional TokenCounterInterface implementation for this provider, or None if token counting is not supported. @@ -152,13 +156,15 @@ def get_api_key_from_env() -> Optional[str]: class GoogleAIStudioTokenCounter(BaseTokenCounter): """Token counter implementation for Google AI Studio provider.""" + def should_use_token_counting_api( - self, + self, custom_llm_provider: Optional[str] = None, ) -> bool: from litellm.types.utils import LlmProviders + return custom_llm_provider == LlmProviders.GEMINI.value - + async def count_tokens( self, model_to_use: str, @@ -172,8 +178,11 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): import copy from litellm.llms.gemini.count_tokens.handler import GoogleAIStudioTokenCounter + deployment = deployment or {} - count_tokens_params_request = copy.deepcopy(deployment.get("litellm_params", {})) + count_tokens_params_request = copy.deepcopy( + deployment.get("litellm_params", {}) + ) count_tokens_params = { "model": model_to_use, "contents": contents, @@ -182,7 +191,7 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): result = await GoogleAIStudioTokenCounter().acount_tokens( **count_tokens_params_request, ) - + if result is not None: return TokenCountResponse( total_tokens=result.get("totalTokens", 0), @@ -191,5 +200,5 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): tokenizer_type=result.get("tokenizer_used", ""), original_response=result, ) - - return None \ No newline at end of file + + return None diff --git a/litellm/llms/gemini/cost_calculator.py b/litellm/llms/gemini/cost_calculator.py index 79242fe01d..45850e0d66 100644 --- a/litellm/llms/gemini/cost_calculator.py +++ b/litellm/llms/gemini/cost_calculator.py @@ -21,7 +21,10 @@ def cost_per_token( from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token return generic_cost_per_token( - model=model, usage=usage, custom_llm_provider="gemini", service_tier=service_tier + model=model, + usage=usage, + custom_llm_provider="gemini", + service_tier=service_tier, ) diff --git a/litellm/llms/gemini/files/transformation.py b/litellm/llms/gemini/files/transformation.py index cc799cfd6a..bdfb0ee1e5 100644 --- a/litellm/llms/gemini/files/transformation.py +++ b/litellm/llms/gemini/files/transformation.py @@ -52,8 +52,10 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): """ resolved_api_key = self.get_api_key(api_key) if not resolved_api_key: - raise ValueError("GEMINI_API_KEY is required for Google AI Studio file operations") - + raise ValueError( + "GEMINI_API_KEY is required for Google AI Studio file operations" + ) + headers["x-goog-api-key"] = resolved_api_key return headers @@ -206,7 +208,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): ) -> tuple[str, dict]: """ Get the URL to retrieve a file from Google AI Studio. - + We expect file_id to be the URI (e.g. https://generativelanguage.googleapis.com/v1beta/files/...) as returned by the upload response. """ @@ -218,7 +220,10 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): url = "{}?key={}".format(file_id, api_key) else: # Fallback for just file name (files/...) - api_base = self.get_api_base(litellm_params.get("api_base")) or "https://generativelanguage.googleapis.com" + api_base = ( + self.get_api_base(litellm_params.get("api_base")) + or "https://generativelanguage.googleapis.com" + ) api_base = api_base.rstrip("/") url = "{}/v1beta/{}?key={}".format(api_base, file_id, api_key) @@ -236,7 +241,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): """ try: response_json = raw_response.json() - + # Map Gemini state to OpenAI status gemini_state = response_json.get("state", "STATE_UNSPECIFIED") # Explicitly type status as the Literal union @@ -246,7 +251,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): status = "error" else: status = "uploaded" - + return OpenAIFileObject( id=response_json.get("uri", ""), bytes=int(response_json.get("sizeBytes", 0)), @@ -262,7 +267,9 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): object="file", purpose="user_data", status=status, - status_details=str(response_json.get("error", "")) if gemini_state == "FAILED" else None, + status_details=str(response_json.get("error", "")) + if gemini_state == "FAILED" + else None, ) except Exception as e: verbose_logger.exception(f"Error parsing file retrieve response: {str(e)}") @@ -276,24 +283,24 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): ) -> tuple[str, dict]: """ Transform delete file request for Google AI Studio. - + Args: file_id: The file URI (e.g., "files/abc123" or full URI) optional_params: Optional parameters litellm_params: LiteLLM parameters containing api_key - + Returns: tuple[str, dict]: (url, params) for the DELETE request """ api_base = self.get_api_base(litellm_params.get("api_base")) if not api_base: raise ValueError("api_base is required") - + # Get API key from multiple sources (same pattern as get_complete_url) api_key = litellm_params.get("api_key") or self.get_api_key() if not api_key: raise ValueError("api_key is required") - + # Extract file name from URI if full URI is provided # file_id could be "files/abc123" or "https://generativelanguage.googleapis.com/v1beta/files/abc123" if file_id.startswith("http"): @@ -301,13 +308,13 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): file_name = file_id.split("/v1beta/")[-1] else: file_name = file_id if file_id.startswith("files/") else f"files/{file_id}" - + # Construct the delete URL url = f"{api_base}/v1beta/{file_name}" - + # Add API key as header (Google AI Studio uses x-goog-api-key header) params: dict = {} - + return url, params def transform_delete_file_response( @@ -318,7 +325,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): ) -> FileDeleted: """ Transform Gemini's file delete response into OpenAI-style FileDeleted. - + Google AI Studio returns an empty JSON object {} on successful deletion. """ try: @@ -333,12 +340,8 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): # Add the files/ prefix if not present if not file_id.startswith("files/"): file_id = f"files/{file_id}" - - return FileDeleted( - id=file_id, - deleted=True, - object="file" - ) + + return FileDeleted(id=file_id, deleted=True, object="file") else: raise ValueError(f"Failed to delete file: {raw_response.text}") except Exception as e: @@ -351,7 +354,9 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): optional_params: dict, litellm_params: dict, ) -> tuple[str, dict]: - raise NotImplementedError("GoogleAIStudioFilesHandler does not support file listing") + raise NotImplementedError( + "GoogleAIStudioFilesHandler does not support file listing" + ) def transform_list_files_response( self, @@ -359,7 +364,9 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): logging_obj: LiteLLMLoggingObj, litellm_params: dict, ) -> List[OpenAIFileObject]: - raise NotImplementedError("GoogleAIStudioFilesHandler does not support file listing") + raise NotImplementedError( + "GoogleAIStudioFilesHandler does not support file listing" + ) def transform_file_content_request( self, @@ -367,7 +374,9 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): optional_params: dict, litellm_params: dict, ) -> tuple[str, dict]: - raise NotImplementedError("GoogleAIStudioFilesHandler does not support file content retrieval") + raise NotImplementedError( + "GoogleAIStudioFilesHandler does not support file content retrieval" + ) def transform_file_content_response( self, @@ -375,4 +384,6 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): logging_obj: LiteLLMLoggingObj, litellm_params: dict, ) -> HttpxBinaryResponseContent: - raise NotImplementedError("GoogleAIStudioFilesHandler does not support file content retrieval") + raise NotImplementedError( + "GoogleAIStudioFilesHandler does not support file content retrieval" + ) diff --git a/litellm/llms/gemini/google_genai/transformation.py b/litellm/llms/gemini/google_genai/transformation.py index 48046dd9df..7c4c7dba62 100644 --- a/litellm/llms/gemini/google_genai/transformation.py +++ b/litellm/llms/gemini/google_genai/transformation.py @@ -75,7 +75,7 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): "seed", "response_mime_type", "response_schema", - "response_json_schema", + "response_json_schema", "routing_config", "model_selection_config", "safety_settings", @@ -111,29 +111,33 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): _camel_to_snake, _snake_to_camel, ) - + _generate_content_config_dict: Dict[str, Any] = {} supported_google_genai_params = ( self.get_supported_generate_content_optional_params(model) ) # Create a set with both camelCase and snake_case versions for faster lookup supported_params_set = set(supported_google_genai_params) - supported_params_set.update(_snake_to_camel(p) for p in supported_google_genai_params) - supported_params_set.update(_camel_to_snake(p) for p in supported_google_genai_params if "_" not in p) - + supported_params_set.update( + _snake_to_camel(p) for p in supported_google_genai_params + ) + supported_params_set.update( + _camel_to_snake(p) for p in supported_google_genai_params if "_" not in p + ) + for param, value in generate_content_config_dict.items(): # Google GenAI API expects camelCase, so we'll always output in camelCase # Check if param (or its variants) is supported param_snake = _camel_to_snake(param) param_camel = _snake_to_camel(param) - + # Check if param is supported in any format is_supported = ( - param in supported_google_genai_params or - param_snake in supported_google_genai_params or - param_camel in supported_google_genai_params + param in supported_google_genai_params + or param_snake in supported_google_genai_params + or param_camel in supported_google_genai_params ) - + if is_supported: # Always output in camelCase for Google GenAI API output_key = param_camel if param != param_camel else param @@ -234,9 +238,11 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): """ Sync version of get_auth_token_and_url. """ - vertex_credentials, vertex_project, vertex_location = ( - self._get_common_auth_components(litellm_params) - ) + ( + vertex_credentials, + vertex_project, + vertex_location, + ) = self._get_common_auth_components(litellm_params) _auth_header, vertex_project = self._ensure_access_token( credentials=vertex_credentials, @@ -273,9 +279,11 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): Returns: Tuple of headers and API base """ - vertex_credentials, vertex_project, vertex_location = ( - self._get_common_auth_components(litellm_params) - ) + ( + vertex_credentials, + vertex_project, + vertex_location, + ) = self._get_common_auth_components(litellm_params) _auth_header, vertex_project = await self._ensure_access_token_async( credentials=vertex_credentials, @@ -315,7 +323,7 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): ) request_dict = cast(dict, typed_generate_content_request) - + if system_instruction is not None: request_dict["systemInstruction"] = system_instruction return request_dict @@ -359,9 +367,13 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): """ if "candidates" in response: for candidate in response["candidates"]: - if "citationMetadata" in candidate and isinstance(candidate["citationMetadata"], dict): + if "citationMetadata" in candidate and isinstance( + candidate["citationMetadata"], dict + ): citation_metadata = candidate["citationMetadata"] # Transform citationSources to citations to match expected schema if "citationSources" in citation_metadata: - citation_metadata["citations"] = citation_metadata.pop("citationSources") - return response \ No newline at end of file + citation_metadata["citations"] = citation_metadata.pop( + "citationSources" + ) + return response diff --git a/litellm/llms/gemini/image_edit/__init__.py b/litellm/llms/gemini/image_edit/__init__.py index 6181015b81..cb097d3eee 100644 --- a/litellm/llms/gemini/image_edit/__init__.py +++ b/litellm/llms/gemini/image_edit/__init__.py @@ -8,4 +8,3 @@ __all__ = ["GeminiImageEditConfig", "get_gemini_image_edit_config", "cost_calcul def get_gemini_image_edit_config(model: str) -> BaseImageEditConfig: return GeminiImageEditConfig() - diff --git a/litellm/llms/gemini/image_edit/cost_calculator.py b/litellm/llms/gemini/image_edit/cost_calculator.py index 31f35345d8..2e332a7fc0 100644 --- a/litellm/llms/gemini/image_edit/cost_calculator.py +++ b/litellm/llms/gemini/image_edit/cost_calculator.py @@ -32,4 +32,3 @@ def cost_calculator( num_images = len(image_response.data or []) return output_cost_per_image * num_images - diff --git a/litellm/llms/gemini/image_edit/transformation.py b/litellm/llms/gemini/image_edit/transformation.py index c3ea63ad43..5d9b1255d0 100644 --- a/litellm/llms/gemini/image_edit/transformation.py +++ b/litellm/llms/gemini/image_edit/transformation.py @@ -73,7 +73,9 @@ class GeminiImageEditConfig(BaseImageEditConfig): api_base: Optional[str], litellm_params: dict, ) -> str: - base_url = api_base or get_secret_str("GEMINI_API_BASE") or self.DEFAULT_BASE_URL + base_url = ( + api_base or get_secret_str("GEMINI_API_BASE") or self.DEFAULT_BASE_URL + ) base_url = base_url.rstrip("/") return f"{base_url}/models/{model}:generateContent" @@ -109,9 +111,9 @@ class GeminiImageEditConfig(BaseImageEditConfig): # Move aspectRatio into imageConfig inside generationConfig if "imageConfig" not in generation_config: generation_config["imageConfig"] = {} - generation_config["imageConfig"]["aspectRatio"] = image_edit_optional_request_params[ + generation_config["imageConfig"][ "aspectRatio" - ] + ] = image_edit_optional_request_params["aspectRatio"] if generation_config: request_body["generationConfig"] = generation_config @@ -206,4 +208,4 @@ class GeminiImageEditConfig(BaseImageEditConfig): data = image.read() image.seek(current_pos) return data - raise ValueError("Unsupported image type for Gemini image edit.") \ No newline at end of file + raise ValueError("Unsupported image type for Gemini image edit.") diff --git a/litellm/llms/gemini/image_generation/cost_calculator.py b/litellm/llms/gemini/image_generation/cost_calculator.py index 941ab0d50f..3c8e69374a 100644 --- a/litellm/llms/gemini/image_generation/cost_calculator.py +++ b/litellm/llms/gemini/image_generation/cost_calculator.py @@ -39,4 +39,6 @@ def cost_calculator( num_images = len(image_response.data) return output_cost_per_image * num_images else: - raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}") + raise ValueError( + f"image_response must be of type ImageResponse got type={type(image_response)}" + ) diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index 73aef15e4c..3e3f6162fc 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/litellm/llms/gemini/image_generation/transformation.py @@ -28,7 +28,7 @@ else: class GoogleImageGenConfig(BaseImageGenerationConfig): DEFAULT_BASE_URL: str = "https://generativelanguage.googleapis.com/v1beta" - + def get_supported_openai_params( self, model: str ) -> List[OpenAIImageGenerationOptionalParams]: @@ -36,11 +36,8 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): Google AI Imagen API supported parameters https://ai.google.dev/gemini-api/docs/imagen """ - return [ - "n", - "size" - ] - + return ["n", "size"] + def map_openai_params( self, non_default_params: dict, @@ -50,7 +47,7 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): ) -> dict: supported_params = self.get_supported_openai_params(model) mapped_params = {} - + for k, v in non_default_params.items(): if k not in optional_params.keys(): if k in supported_params: @@ -61,9 +58,8 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): # Map OpenAI size format to Google aspectRatio mapped_params["aspectRatio"] = self._map_size_to_aspect_ratio(v) else: - mapped_params[k] = v + mapped_params[k] = v return mapped_params - def _map_size_to_aspect_ratio(self, size: str) -> str: """ @@ -72,13 +68,13 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): """ aspect_ratio_map = { "1024x1024": "1:1", - "1792x1024": "16:9", + "1792x1024": "16:9", "1024x1792": "9:16", "1280x896": "4:3", "896x1280": "3:4", } return aspect_ratio_map.get(size, "1:1") - + def _transform_image_usage(self, usage_metadata: dict) -> ImageUsage: """ Transform Gemini usageMetadata to ImageUsage format @@ -87,7 +83,7 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): image_tokens=0, text_tokens=0, ) - + # Extract detailed token counts from promptTokensDetails tokens_details = usage_metadata.get("promptTokensDetails", []) for details in tokens_details: @@ -98,7 +94,7 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): input_tokens_details.text_tokens = token_count elif modality == "IMAGE": input_tokens_details.image_tokens = token_count - + return ImageUsage( input_tokens=usage_metadata.get("promptTokenCount", 0), input_tokens_details=input_tokens_details, @@ -122,9 +118,7 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): Other Imagen models: :predict """ complete_url: str = ( - api_base - or get_secret_str("GEMINI_API_BASE") - or self.DEFAULT_BASE_URL + api_base or get_secret_str("GEMINI_API_BASE") or self.DEFAULT_BASE_URL ) complete_url = complete_url.rstrip("/") @@ -148,13 +142,10 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): api_key: Optional[str] = None, api_base: Optional[str] = None, ) -> dict: - final_api_key: Optional[str] = ( - api_key or - get_secret_str("GEMINI_API_KEY") - ) + final_api_key: Optional[str] = api_key or get_secret_str("GEMINI_API_KEY") if not final_api_key: raise ValueError("GEMINI_API_KEY is not set") - + headers["x-goog-api-key"] = final_api_key headers["Content-Type"] = "application/json" return headers @@ -187,16 +178,8 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): # For Gemini Flash Image Preview models, use standard Gemini format if "gemini" in model: request_body: dict = { - "contents": [ - { - "parts": [ - {"text": prompt} - ] - } - ], - "generationConfig": { - "response_modalities": ["IMAGE", "TEXT"] - } + "contents": [{"parts": [{"text": prompt}]}], + "generationConfig": {"response_modalities": ["IMAGE", "TEXT"]}, } return request_body else: @@ -205,13 +188,12 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): GeminiImageGenerationInstance, GeminiImageGenerationParameters, ) - request_body_obj: GeminiImageGenerationRequest = GeminiImageGenerationRequest( - instances=[ - GeminiImageGenerationInstance( - prompt=prompt - ) - ], - parameters=GeminiImageGenerationParameters(**optional_params) + + request_body_obj: GeminiImageGenerationRequest = ( + GeminiImageGenerationRequest( + instances=[GeminiImageGenerationInstance(prompt=prompt)], + parameters=GeminiImageGenerationParameters(**optional_params), + ) ) return request_body_obj.model_dump(exclude_none=True) @@ -239,7 +221,7 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): status_code=raw_response.status_code, headers=raw_response.headers, ) - + if not model_response.data: model_response.data = [] @@ -256,22 +238,32 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): inline_data = part["inlineData"] if "data" in inline_data: thought_sig = part.get("thoughtSignature") - model_response.data.append(ImageObject( - b64_json=inline_data["data"], - url=None, - provider_specific_fields={"thought_signature": thought_sig} if thought_sig else None, - )) - + model_response.data.append( + ImageObject( + b64_json=inline_data["data"], + url=None, + provider_specific_fields={ + "thought_signature": thought_sig + } + if thought_sig + else None, + ) + ) + # Extract usage metadata for Gemini models if "usageMetadata" in response_data: - model_response.usage = self._transform_image_usage(response_data["usageMetadata"]) + model_response.usage = self._transform_image_usage( + response_data["usageMetadata"] + ) else: # Original Imagen format - predictions with generated images predictions = response_data.get("predictions", []) for prediction in predictions: # Google AI returns base64 encoded images in the prediction - model_response.data.append(ImageObject( - b64_json=prediction.get("bytesBase64Encoded", None), - url=None, # Google AI returns base64, not URLs - )) - return model_response \ No newline at end of file + model_response.data.append( + ImageObject( + b64_json=prediction.get("bytesBase64Encoded", None), + url=None, # Google AI returns base64, not URLs + ) + ) + return model_response diff --git a/litellm/llms/gemini/interactions/transformation.py b/litellm/llms/gemini/interactions/transformation.py index d21775eb23..772530342e 100644 --- a/litellm/llms/gemini/interactions/transformation.py +++ b/litellm/llms/gemini/interactions/transformation.py @@ -39,7 +39,7 @@ else: class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): """ Configuration for Google AI Studio Interactions API. - + Minimal config - we follow the OpenAPI spec directly with no transformation. """ @@ -54,9 +54,18 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): def get_supported_params(self, model: str) -> List[str]: """Per OpenAPI spec CreateModelInteractionParams.""" return [ - "model", "agent", "input", "tools", "system_instruction", - "generation_config", "stream", "store", "background", - "response_modalities", "response_format", "response_mime_type", + "model", + "agent", + "input", + "tools", + "system_instruction", + "generation_config", + "stream", + "store", + "background", + "response_modalities", + "response_format", + "response_mime_type", "previous_interaction_id", ] @@ -83,16 +92,16 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): litellm_params = litellm_params or {} api_base = GeminiModelInfo.get_api_base(api_base) api_key = GeminiModelInfo.get_api_key(litellm_params.get("api_key")) - + if not api_key: raise ValueError( "Google API key is required. Set GOOGLE_API_KEY or GEMINI_API_KEY environment variable." ) - + query_params = f"key={api_key}" if stream: query_params += "&alt=sse" - + return f"{api_base}/{self.api_version}/interactions?{query_params}" def transform_request( @@ -108,7 +117,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): Build request body per OpenAPI spec - minimal transformation. """ request_body: Dict[str, Any] = {} - + # Model or Agent (one required) if model: request_body["model"] = GeminiModelInfo.get_base_model(model) or model @@ -116,21 +125,28 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): request_body["agent"] = agent else: raise ValueError("Either 'model' or 'agent' must be provided") - + # Input if input is not None: request_body["input"] = input - + # Pass through optional params directly (they match the spec) optional_keys = [ - "tools", "system_instruction", "generation_config", "stream", "store", - "background", "response_modalities", "response_format", - "response_mime_type", "previous_interaction_id", + "tools", + "system_instruction", + "generation_config", + "stream", + "store", + "background", + "response_modalities", + "response_format", + "response_mime_type", + "previous_interaction_id", ] for key in optional_keys: if optional_params.get(key) is not None: request_body[key] = optional_params[key] - + return request_body def transform_response( @@ -152,13 +168,15 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): status_code=raw_response.status_code, headers=dict(raw_response.headers), ) - + verbose_logger.debug("Google AI Interactions response: %s", raw_json) - + response = InteractionsAPIResponse(**raw_json) response._hidden_params["headers"] = dict(raw_response.headers) - response._hidden_params["additional_headers"] = process_response_headers(dict(raw_response.headers)) - + response._hidden_params["additional_headers"] = process_response_headers( + dict(raw_response.headers) + ) + return response def transform_streaming_response( @@ -172,7 +190,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): return InteractionsAPIStreamingResponse(**parsed_chunk) # GET / DELETE / CANCEL - just build URLs, responses match spec directly - + def transform_get_interaction_request( self, interaction_id: str, @@ -185,7 +203,10 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): api_key = GeminiModelInfo.get_api_key(litellm_params.api_key) if not api_key: raise ValueError("Google API key is required") - return f"{resolved_api_base}/{self.api_version}/interactions/{interaction_id}?key={api_key}", {} + return ( + f"{resolved_api_base}/{self.api_version}/interactions/{interaction_id}?key={api_key}", + {}, + ) def transform_get_interaction_response( self, @@ -216,7 +237,10 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): api_key = GeminiModelInfo.get_api_key(litellm_params.api_key) if not api_key: raise ValueError("Google API key is required") - return f"{resolved_api_base}/{self.api_version}/interactions/{interaction_id}?key={api_key}", {} + return ( + f"{resolved_api_base}/{self.api_version}/interactions/{interaction_id}?key={api_key}", + {}, + ) def transform_delete_interaction_response( self, @@ -244,7 +268,10 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): api_key = GeminiModelInfo.get_api_key(litellm_params.api_key) if not api_key: raise ValueError("Google API key is required") - return f"{resolved_api_base}/{self.api_version}/interactions/{interaction_id}:cancel?key={api_key}", {} + return ( + f"{resolved_api_base}/{self.api_version}/interactions/{interaction_id}:cancel?key={api_key}", + {}, + ) def transform_cancel_interaction_response( self, diff --git a/litellm/llms/gemini/realtime/transformation.py b/litellm/llms/gemini/realtime/transformation.py index a3eedd36a6..2bb7bcd8b4 100644 --- a/litellm/llms/gemini/realtime/transformation.py +++ b/litellm/llms/gemini/realtime/transformation.py @@ -186,10 +186,10 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): ) vertex_gemini_config = VertexGeminiConfig() - optional_params["generationConfig"]["tools"] = ( - vertex_gemini_config._map_function( - value=value, optional_params=optional_params - ) + optional_params["generationConfig"][ + "tools" + ] = vertex_gemini_config._map_function( + value=value, optional_params=optional_params ) elif key == "input_audio_transcription" and value is not None: optional_params["inputAudioTranscription"] = {} @@ -201,10 +201,10 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): if ( len(transformed_audio_activity_config) > 0 ): # if the config is not empty, add it to the optional params - optional_params["realtimeInputConfig"] = ( - BidiGenerateContentRealtimeInputConfig( - automaticActivityDetection=transformed_audio_activity_config - ) + optional_params[ + "realtimeInputConfig" + ] = BidiGenerateContentRealtimeInputConfig( + automaticActivityDetection=transformed_audio_activity_config ) if len(optional_params["generationConfig"]) == 0: optional_params.pop("generationConfig") @@ -235,9 +235,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): optional_params={}, non_default_params=json_message["session"] ) client_session_configuration_request["model"] = f"models/{model}" - messages.append( - json.dumps({"setup": client_session_configuration_request}) - ) + messages.append(json.dumps({"setup": client_session_configuration_request})) return messages ## HANDLE response.create — Gemini responds automatically; nothing to forward ## @@ -320,7 +318,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): if "/models/" in _model: session["model"] = _model.split("/models/")[-1] elif _model.startswith("models/"): - session["model"] = _model[len("models/"):] + session["model"] = _model[len("models/") :] else: session["model"] = _model @@ -779,8 +777,12 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): # Use IDs from the done event — transform_content_done_event may have # generated UUID fallbacks when the originals were None. - resolved_item_id = transformed_content_done_event.get("item_id") or current_output_item_id - resolved_response_id = transformed_content_done_event.get("response_id") or current_response_id + resolved_item_id = ( + transformed_content_done_event.get("item_id") or current_output_item_id + ) + resolved_response_id = ( + transformed_content_done_event.get("response_id") or current_response_id + ) additional_items = self.return_additional_content_done_events( current_output_item_id=resolved_item_id, @@ -862,9 +864,9 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): "session_configuration_request" ] current_item_chunks = realtime_response_transform_input["current_item_chunks"] - current_delta_type: Optional[ALL_DELTA_TYPES] = ( - realtime_response_transform_input["current_delta_type"] - ) + current_delta_type: Optional[ + ALL_DELTA_TYPES + ] = realtime_response_transform_input["current_delta_type"] returned_message: List[OpenAIRealtimeEvents] = [] # Handle transcription events that arrive independently from model @@ -875,32 +877,45 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): input_tx = server_content.get("inputTranscription") if isinstance(input_tx, dict) and input_tx.get("text"): returned_message.append( - cast(OpenAIRealtimeEvents, { - "type": "conversation.item.input_audio_transcription.completed", - "event_id": "event_{}".format(uuid.uuid4()), - "transcript": input_tx["text"], - "item_id": "item_{}".format(uuid.uuid4()), - "content_index": 0, - }) + cast( + OpenAIRealtimeEvents, + { + "type": "conversation.item.input_audio_transcription.completed", + "event_id": "event_{}".format(uuid.uuid4()), + "transcript": input_tx["text"], + "item_id": "item_{}".format(uuid.uuid4()), + "content_index": 0, + }, + ) ) output_tx = server_content.get("outputTranscription") if isinstance(output_tx, dict) and output_tx.get("text"): returned_message.append( - cast(OpenAIRealtimeEvents, { - "type": "response.audio_transcript.delta", - "event_id": "event_{}".format(uuid.uuid4()), - "delta": output_tx["text"], - "item_id": current_output_item_id or "item_{}".format(uuid.uuid4()), - "response_id": current_response_id or "resp_{}".format(uuid.uuid4()), - "output_index": 0, - "content_index": 0, - }) + cast( + OpenAIRealtimeEvents, + { + "type": "response.audio_transcript.delta", + "event_id": "event_{}".format(uuid.uuid4()), + "delta": output_tx["text"], + "item_id": current_output_item_id + or "item_{}".format(uuid.uuid4()), + "response_id": current_response_id + or "resp_{}".format(uuid.uuid4()), + "output_index": 0, + "content_index": 0, + }, + ) ) # If serverContent only contained transcription(s) and no model # content, return early — the main loop would fail on unknown keys. - _model_content_keys = {"modelTurn", "turnComplete", "interrupted", "generationComplete"} + _model_content_keys = { + "modelTurn", + "turnComplete", + "interrupted", + "generationComplete", + } if not any(k in server_content for k in _model_content_keys): return { "response": returned_message, diff --git a/litellm/llms/gemini/vector_stores/__init__.py b/litellm/llms/gemini/vector_stores/__init__.py index 613b5775b6..b2d276ac21 100644 --- a/litellm/llms/gemini/vector_stores/__init__.py +++ b/litellm/llms/gemini/vector_stores/__init__.py @@ -3,4 +3,3 @@ from .transformation import GeminiVectorStoreConfig __all__ = ["GeminiVectorStoreConfig"] - diff --git a/litellm/llms/gemini/vector_stores/transformation.py b/litellm/llms/gemini/vector_stores/transformation.py index 4d76f691e5..11fd77aeca 100644 --- a/litellm/llms/gemini/vector_stores/transformation.py +++ b/litellm/llms/gemini/vector_stores/transformation.py @@ -54,7 +54,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): def get_vector_store_endpoints_by_type(self) -> VectorStoreIndexEndpoints: """ Gemini File Search endpoints. - + Note: Search is done via generateContent with file_search tool, not a dedicated search endpoint. """ @@ -79,22 +79,22 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): api_key = litellm_params.get("api_key") or get_api_key_from_env() if api_key: self._cached_api_key = api_key - + return headers def get_complete_url(self, api_base: Optional[str], litellm_params: dict) -> str: """ Get the complete base URL for Gemini API. - + Note: This returns the base URL WITHOUT the API key. The API key will be appended to specific endpoint URLs in the transform methods. """ if api_base is None: api_base = GeminiModelInfo.get_api_base() - + if api_base is None: raise ValueError("GEMINI_API_BASE is not set") - + # Ensure we're using the v1beta version for File Search api_version = "v1beta" return f"{api_base}/{api_version}" @@ -120,7 +120,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): ) -> Tuple[str, Dict]: """ Transform search request to Gemini's generateContent format. - + Gemini File Search works by calling generateContent with a file_search tool. """ # Convert query list to single string if needed @@ -157,23 +157,15 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): if isinstance(value, str): filter_parts.append(f'{key} = "{value}"') else: - filter_parts.append(f'{key} = {value}') + filter_parts.append(f"{key} = {value}") file_search_config["metadata_filter"] = " AND ".join(filter_parts) else: file_search_config["metadata_filter"] = metadata_filter # Build request body request_body: Dict[str, Any] = { - "contents": [ - { - "parts": [{"text": query}] - } - ], - "tools": [ - { - "file_search": file_search_config - } - ], + "contents": [{"parts": [{"text": query}]}], + "tools": [{"file_search": file_search_config}], } # Add max_num_results if specified @@ -193,7 +185,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): ) -> VectorStoreSearchResponse: """ Transform Gemini's generateContent response to standard format. - + Extracts grounding metadata and citations from the response. """ try: @@ -202,28 +194,30 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): # Extract candidates and grounding metadata candidates = response_data.get("candidates", []) - + for candidate in candidates: grounding_metadata = candidate.get("groundingMetadata", {}) grounding_chunks = grounding_metadata.get("groundingChunks", []) - + # Process each grounding chunk for chunk in grounding_chunks: retrieved_context = chunk.get("retrievedContext") - + if retrieved_context: # This is from file search text = retrieved_context.get("text", "") uri = retrieved_context.get("uri", "") title = retrieved_context.get("title", "") - + # Extract file_id from URI if available file_id = uri if uri else None - + results.append( VectorStoreSearchResult( score=None, # Gemini doesn't provide explicit scores - content=[VectorStoreResultContent(text=text, type="text")], + content=[ + VectorStoreResultContent(text=text, type="text") + ], file_id=file_id, filename=title if title else None, attributes={ @@ -238,13 +232,13 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): for support in grounding_supports: segment = support.get("segment", {}) text = segment.get("text", "") - + grounding_chunk_indices = support.get("groundingChunkIndices", []) confidence_scores = support.get("confidenceScores", []) - + # Use first confidence score as relevance score score = confidence_scores[0] if confidence_scores else None - + # Only add if we have meaningful text and it's not a duplicate if text: already_exists = False @@ -258,7 +252,9 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): results.append( VectorStoreSearchResult( score=score, - content=[VectorStoreResultContent(text=text, type="text")], + content=[ + VectorStoreResultContent(text=text, type="text") + ], attributes={ "grounding_chunk_indices": grounding_chunk_indices, }, @@ -266,7 +262,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): ) query = litellm_logging_obj.model_call_details.get("query", "") - + return VectorStoreSearchResponse( object="vector_store.search_results.page", search_query=query, @@ -289,7 +285,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): Transform create request to Gemini's fileSearchStores format. """ url = f"{api_base}/fileSearchStores" - + # Append API key as query parameter (required by Gemini) api_key = self._cached_api_key or get_api_key_from_env() if api_key: @@ -312,7 +308,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): """ try: response_data = response.json() - + # Extract store name (format: fileSearchStores/xxxxxxx) store_name = response_data.get("name", "") display_name = response_data.get("displayName", "") @@ -320,10 +316,13 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): # Convert ISO timestamp to Unix timestamp import datetime + created_at = None if create_time: try: - dt = datetime.datetime.fromisoformat(create_time.replace("Z", "+00:00")) + dt = datetime.datetime.fromisoformat( + create_time.replace("Z", "+00:00") + ) created_at = int(dt.timestamp()) except Exception: created_at = None @@ -354,4 +353,3 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): status_code=response.status_code, headers=response.headers, ) - diff --git a/litellm/llms/gemini/videos/__init__.py b/litellm/llms/gemini/videos/__init__.py index c5aed2db2d..b8e0452cb0 100644 --- a/litellm/llms/gemini/videos/__init__.py +++ b/litellm/llms/gemini/videos/__init__.py @@ -2,4 +2,3 @@ from .transformation import GeminiVideoConfig __all__ = ["GeminiVideoConfig"] - diff --git a/litellm/llms/gemini/videos/transformation.py b/litellm/llms/gemini/videos/transformation.py index 7daeb75b65..c16b20fe57 100644 --- a/litellm/llms/gemini/videos/transformation.py +++ b/litellm/llms/gemini/videos/transformation.py @@ -13,7 +13,12 @@ from litellm.types.videos.utils import ( ) from litellm.images.utils import ImageEditRequestUtils import litellm -from litellm.types.llms.gemini import GeminiLongRunningOperationResponse, GeminiVideoGenerationInstance, GeminiVideoGenerationParameters, GeminiVideoGenerationRequest +from litellm.types.llms.gemini import ( + GeminiLongRunningOperationResponse, + GeminiVideoGenerationInstance, + GeminiVideoGenerationParameters, + GeminiVideoGenerationRequest, +) from litellm.constants import DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS from litellm.llms.base_llm.videos.transformation import BaseVideoConfig @@ -31,30 +36,27 @@ else: def _convert_image_to_gemini_format(image_file) -> Dict[str, str]: """ Convert image file to Gemini format with base64 encoding and MIME type. - + Args: image_file: File-like object opened in binary mode (e.g., open("path", "rb")) - + Returns: Dict with bytesBase64Encoded and mimeType """ mime_type = ImageEditRequestUtils.get_image_content_type(image_file) - - if hasattr(image_file, 'seek'): + + if hasattr(image_file, "seek"): image_file.seek(0) image_bytes = image_file.read() base64_encoded = base64.b64encode(image_bytes).decode("utf-8") - - return { - "bytesBase64Encoded": base64_encoded, - "mimeType": mime_type - } + + return {"bytesBase64Encoded": base64_encoded, "mimeType": mime_type} class GeminiVideoConfig(BaseVideoConfig): """ Configuration class for Gemini (Veo) video generation. - + Veo uses a long-running operation model: 1. POST to :predictLongRunning returns operation name 2. Poll operation until done=true @@ -70,13 +72,7 @@ class GeminiVideoConfig(BaseVideoConfig): Get the list of supported OpenAI parameters for Veo video generation. Veo supports minimal parameters compared to OpenAI. """ - return [ - "model", - "prompt", - "input_reference", - "seconds", - "size" - ] + return ["model", "prompt", "input_reference", "seconds", "size"] def map_openai_params( self, @@ -86,28 +82,29 @@ class GeminiVideoConfig(BaseVideoConfig): ) -> Dict[str, Any]: """ Map OpenAI-style parameters to Veo format. - + Mappings: - prompt → prompt - input_reference → image - size → aspectRatio (e.g., "1280x720" → "16:9") - seconds → durationSeconds (defaults to 4 seconds if not provided) - + All other params are passed through as-is to support Gemini-specific parameters. """ mapped_params: Dict[str, Any] = {} - + # Get supported OpenAI params (exclude "model" and "prompt" which are handled separately) supported_openai_params = self.get_supported_openai_params(model) openai_params_to_map = { - param for param in supported_openai_params + param + for param in supported_openai_params if param not in {"model", "prompt"} } - + # Map input_reference to image if "input_reference" in video_create_optional_params: mapped_params["image"] = video_create_optional_params["input_reference"] - + # Map size to aspectRatio if "size" in video_create_optional_params: size = video_create_optional_params["size"] @@ -115,7 +112,7 @@ class GeminiVideoConfig(BaseVideoConfig): aspect_ratio = self._convert_size_to_aspect_ratio(size) if aspect_ratio: mapped_params["aspectRatio"] = aspect_ratio - + # Map seconds to durationSeconds, default to 4 seconds (matching OpenAI) if "seconds" in video_create_optional_params: seconds = video_create_optional_params["seconds"] @@ -126,34 +123,33 @@ class GeminiVideoConfig(BaseVideoConfig): except (ValueError, TypeError): # If conversion fails, use default pass - + # Pass through any other params that weren't mapped (Gemini-specific params) for key, value in video_create_optional_params.items(): if key not in openai_params_to_map and key not in mapped_params: mapped_params[key] = value - + return mapped_params - + def _convert_size_to_aspect_ratio(self, size: str) -> Optional[str]: """ Convert OpenAI size format to Veo aspectRatio format. - + https://cloud.google.com/vertex-ai/generative-ai/docs/image/generate-videos - + Supported aspect ratios: 9:16 (portrait), 16:9 (landscape) """ if not size: return None - + aspect_ratio_map = { "1280x720": "16:9", "1920x1080": "16:9", "720x1280": "9:16", "1080x1920": "9:16", } - - return aspect_ratio_map.get(size, "16:9") + return aspect_ratio_map.get(size, "16:9") def validate_environment( self, @@ -169,24 +165,26 @@ class GeminiVideoConfig(BaseVideoConfig): # Use api_key from litellm_params if available, otherwise fall back to other sources if litellm_params and litellm_params.api_key: api_key = api_key or litellm_params.api_key - + api_key = ( api_key or litellm.api_key or get_secret_str("GOOGLE_API_KEY") or get_secret_str("GEMINI_API_KEY") ) - + if not api_key: raise ValueError( "GEMINI_API_KEY or GOOGLE_API_KEY is required for Veo video generation. " "Set it via environment variable or pass it as api_key parameter." ) - - headers.update({ - "x-goog-api-key": api_key, - "Content-Type": "application/json", - }) + + headers.update( + { + "x-goog-api-key": api_key, + "Content-Type": "application/json", + } + ) return headers def get_complete_url( @@ -201,14 +199,17 @@ class GeminiVideoConfig(BaseVideoConfig): For status/delete: returns base URL only """ if api_base is None: - api_base = get_secret_str("GEMINI_API_BASE") or "https://generativelanguage.googleapis.com" - + api_base = ( + get_secret_str("GEMINI_API_BASE") + or "https://generativelanguage.googleapis.com" + ) + if not model or model == "": - return api_base.rstrip('/') - + return api_base.rstrip("/") + model_name = model.replace("gemini/", "") url = f"{api_base.rstrip('/')}/v1beta/models/{model_name}:predictLongRunning" - + return url def transform_video_create_request( @@ -222,7 +223,7 @@ class GeminiVideoConfig(BaseVideoConfig): ) -> Tuple[Dict, RequestFiles, str]: """ Transform the video creation request for Veo API. - + Veo expects: { "instances": [ @@ -238,22 +239,21 @@ class GeminiVideoConfig(BaseVideoConfig): } """ instance = GeminiVideoGenerationInstance(prompt=prompt) - + params_copy = video_create_optional_request_params.copy() - + if "image" in params_copy and params_copy["image"] is not None: image_data = _convert_image_to_gemini_format(params_copy["image"]) params_copy["image"] = image_data - + parameters = GeminiVideoGenerationParameters(**params_copy) - + request_body_obj = GeminiVideoGenerationRequest( - instances=[instance], - parameters=parameters + instances=[instance], parameters=parameters ) - + request_data = request_body_obj.model_dump(exclude_none=True) - + return request_data, [], api_base def transform_video_create_response( @@ -266,7 +266,7 @@ class GeminiVideoConfig(BaseVideoConfig): ) -> VideoObject: """ Transform the Veo video creation response. - + Veo returns: { "name": "operations/generate_1234567890", @@ -274,46 +274,51 @@ class GeminiVideoConfig(BaseVideoConfig): "done": false, "error": {...} } - + We return this as a VideoObject with: - id: operation name (used for polling) - status: "processing" - usage: includes duration_seconds for cost calculation - """ + """ response_data = raw_response.json() - + # Parse response using Pydantic model for type safety try: operation_response = GeminiLongRunningOperationResponse(**response_data) except Exception as e: raise ValueError(f"Failed to parse operation response: {e}") - + operation_name = operation_response.name if not operation_name: raise ValueError(f"No operation name in Veo response: {response_data}") - + if custom_llm_provider: - video_id = encode_video_id_with_provider(operation_name, custom_llm_provider, model) + video_id = encode_video_id_with_provider( + operation_name, custom_llm_provider, model + ) else: video_id = operation_name - + video_obj = VideoObject( id=video_id, object="video", status="processing", model=model, ) - + usage_data = {} if request_data: parameters = request_data.get("parameters", {}) - duration = parameters.get("durationSeconds") or DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS + duration = ( + parameters.get("durationSeconds") + or DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS + ) if duration is not None: try: usage_data["duration_seconds"] = float(duration) except (ValueError, TypeError): pass - + video_obj.usage = usage_data return video_obj @@ -326,14 +331,14 @@ class GeminiVideoConfig(BaseVideoConfig): ) -> Tuple[str, Dict]: """ Transform the video status retrieve request for Veo API. - + Veo polls operations at: GET https://generativelanguage.googleapis.com/v1beta/{operation_name} """ operation_name = extract_original_video_id(video_id) url = f"{api_base.rstrip('/')}/v1beta/{operation_name}" params: Dict[str, Any] = {} - + return url, params def transform_video_status_retrieve_response( @@ -344,13 +349,13 @@ class GeminiVideoConfig(BaseVideoConfig): ) -> VideoObject: """ Transform the Veo operation status response. - + Veo returns: { "name": "operations/generate_1234567890", "done": false # or true when complete } - + When done=true: { "name": "operations/generate_1234567890", @@ -367,23 +372,25 @@ class GeminiVideoConfig(BaseVideoConfig): } } } - """ + """ response_data = raw_response.json() # Parse response using Pydantic model for type safety operation_response = GeminiLongRunningOperationResponse(**response_data) - + operation_name = operation_response.name is_done = operation_response.done - + if custom_llm_provider: - video_id = encode_video_id_with_provider(operation_name, custom_llm_provider, None) + video_id = encode_video_id_with_provider( + operation_name, custom_llm_provider, None + ) else: video_id = operation_name - + video_obj = VideoObject( id=video_id, object="video", - status="processing" if not is_done else "completed" + status="processing" if not is_done else "completed", ) return video_obj @@ -401,15 +408,15 @@ class GeminiVideoConfig(BaseVideoConfig): For Veo, we need to: 1. Get operation status to extract video URI 2. Return download URL for the video - """ + """ operation_name = extract_original_video_id(video_id) - + status_url = f"{api_base.rstrip('/')}/v1beta/{operation_name}" client = litellm.module_level_client status_response = client.get(url=status_url, headers=headers) status_response.raise_for_status() response_data = status_response.json() - + operation_response = GeminiLongRunningOperationResponse(**response_data) if not operation_response.done: @@ -417,15 +424,17 @@ class GeminiVideoConfig(BaseVideoConfig): "Video generation is not complete yet. " "Please check status with video_status() before downloading." ) - + if not operation_response.response: raise ValueError("No response data in completed operation") - - generated_samples = operation_response.response.generateVideoResponse.generatedSamples + + generated_samples = ( + operation_response.response.generateVideoResponse.generatedSamples + ) download_url = generated_samples[0].video.uri - + params: Dict[str, Any] = {} - + return download_url, params def transform_video_content_response( @@ -525,4 +534,3 @@ class GeminiVideoConfig(BaseVideoConfig): message=error_message, headers=headers, ) - diff --git a/litellm/llms/gigachat/authenticator.py b/litellm/llms/gigachat/authenticator.py index e61015a4a2..59942a9c03 100644 --- a/litellm/llms/gigachat/authenticator.py +++ b/litellm/llms/gigachat/authenticator.py @@ -104,7 +104,9 @@ def get_access_token( token, expires_at = _request_token_sync(credentials, scope, auth_url) # Cache token - ttl_seconds = max(0, (expires_at - TOKEN_EXPIRY_BUFFER_MS - time.time() * 1000) / 1000) + ttl_seconds = max( + 0, (expires_at - TOKEN_EXPIRY_BUFFER_MS - time.time() * 1000) / 1000 + ) if ttl_seconds > 0: _token_cache.set_cache(cache_key, (token, expires_at), ttl=ttl_seconds) @@ -140,7 +142,9 @@ async def get_access_token_async( token, expires_at = await _request_token_async(credentials, scope, auth_url) # Cache token - ttl_seconds = max(0, (expires_at - TOKEN_EXPIRY_BUFFER_MS - time.time() * 1000) / 1000) + ttl_seconds = max( + 0, (expires_at - TOKEN_EXPIRY_BUFFER_MS - time.time() * 1000) / 1000 + ) if ttl_seconds > 0: _token_cache.set_cache(cache_key, (token, expires_at), ttl=ttl_seconds) diff --git a/litellm/llms/gigachat/chat/streaming.py b/litellm/llms/gigachat/chat/streaming.py index 3565559e43..4f10f8bb65 100644 --- a/litellm/llms/gigachat/chat/streaming.py +++ b/litellm/llms/gigachat/chat/streaming.py @@ -6,7 +6,10 @@ import json import uuid from typing import Any, Optional -from litellm.types.llms.openai import ChatCompletionToolCallChunk, ChatCompletionToolCallFunctionChunk +from litellm.types.llms.openai import ( + ChatCompletionToolCallChunk, + ChatCompletionToolCallFunctionChunk, +) from litellm.types.utils import GenericStreamingChunk diff --git a/litellm/llms/gigachat/chat/transformation.py b/litellm/llms/gigachat/chat/transformation.py index f546f356e1..cef8076876 100644 --- a/litellm/llms/gigachat/chat/transformation.py +++ b/litellm/llms/gigachat/chat/transformation.py @@ -254,7 +254,7 @@ class GigaChatConfig(BaseConfig): func_name = tool_choice.get("function", {}).get("name") if func_name: return {"name": func_name} - + # Default to None (don't set function_call) return None diff --git a/litellm/llms/github_copilot/authenticator.py b/litellm/llms/github_copilot/authenticator.py index 7d7ef522a4..85c22516f9 100644 --- a/litellm/llms/github_copilot/authenticator.py +++ b/litellm/llms/github_copilot/authenticator.py @@ -357,7 +357,6 @@ class Authenticator: print( # noqa: T201 f"Please visit {verification_uri} and enter code {user_code} to authenticate.", - # When this is running in docker, it may not be flushed immediately # so we force flush to ensure the user sees the message flush=True, diff --git a/litellm/llms/github_copilot/common_utils.py b/litellm/llms/github_copilot/common_utils.py index 7870f56b84..d3169e3ca9 100644 --- a/litellm/llms/github_copilot/common_utils.py +++ b/litellm/llms/github_copilot/common_utils.py @@ -15,6 +15,7 @@ USER_AGENT = f"GitHubCopilotChat/{COPILOT_VERSION}" API_VERSION = "2025-04-01" GITHUB_COPILOT_API_BASE = "https://api.githubcopilot.com" + class GithubCopilotError(BaseLLMException): def __init__( self, diff --git a/litellm/llms/github_copilot/embedding/transformation.py b/litellm/llms/github_copilot/embedding/transformation.py index 0146601027..fa7bd4e322 100644 --- a/litellm/llms/github_copilot/embedding/transformation.py +++ b/litellm/llms/github_copilot/embedding/transformation.py @@ -100,9 +100,7 @@ class GithubCopilotEmbeddingConfig(BaseEmbeddingConfig): """ # Use provided api_base or fall back to authenticator's base or default api_base = ( - self.authenticator.get_api_base() - or api_base - or GITHUB_COPILOT_API_BASE + self.authenticator.get_api_base() or api_base or GITHUB_COPILOT_API_BASE ) # Remove trailing slashes @@ -121,7 +119,7 @@ class GithubCopilotEmbeddingConfig(BaseEmbeddingConfig): """ Transform embedding request to GitHub Copilot format. """ - + # Ensure input is a list if isinstance(input, str): input = [input] @@ -151,10 +149,10 @@ class GithubCopilotEmbeddingConfig(BaseEmbeddingConfig): Transform embedding response from GitHub Copilot format. """ logging_obj.post_call(original_response=raw_response.text) - + # GitHub Copilot returns standard OpenAI-compatible embedding response response_json = raw_response.json() - + return convert_to_model_response_object( response_object=response_json, model_response_object=model_response, @@ -189,4 +187,3 @@ class GithubCopilotEmbeddingConfig(BaseEmbeddingConfig): return OpenAIConfig().get_error_class( error_message=error_message, status_code=status_code, headers=headers ) - diff --git a/litellm/llms/github_copilot/responses/transformation.py b/litellm/llms/github_copilot/responses/transformation.py index 73240d4651..46efc124b1 100644 --- a/litellm/llms/github_copilot/responses/transformation.py +++ b/litellm/llms/github_copilot/responses/transformation.py @@ -166,9 +166,7 @@ class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig): """ # Use provided api_base or fall back to authenticator's base or default api_base = ( - api_base - or self.authenticator.get_api_base() - or GITHUB_COPILOT_API_BASE + api_base or self.authenticator.get_api_base() or GITHUB_COPILOT_API_BASE ) # Remove trailing slashes @@ -308,7 +306,9 @@ class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig): # Check arrays if isinstance(value, list): return any( - self._contains_vision_content(item, depth=depth + 1, max_depth=max_depth) + self._contains_vision_content( + item, depth=depth + 1, max_depth=max_depth + ) for item in value ) @@ -324,7 +324,9 @@ class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig): # Check content field recursively if "content" in value and isinstance(value["content"], list): return any( - self._contains_vision_content(item, depth=depth + 1, max_depth=max_depth) + self._contains_vision_content( + item, depth=depth + 1, max_depth=max_depth + ) for item in value["content"] ) diff --git a/litellm/llms/google_pse/search/__init__.py b/litellm/llms/google_pse/search/__init__.py index cda3f360f9..0fcfff82c3 100644 --- a/litellm/llms/google_pse/search/__init__.py +++ b/litellm/llms/google_pse/search/__init__.py @@ -4,5 +4,3 @@ Google Programmable Search Engine (PSE) API module. from litellm.llms.google_pse.search.transformation import GooglePSESearchConfig __all__ = ["GooglePSESearchConfig"] - - diff --git a/litellm/llms/google_pse/search/transformation.py b/litellm/llms/google_pse/search/transformation.py index c1ba9cfe62..2fabbc5d16 100644 --- a/litellm/llms/google_pse/search/transformation.py +++ b/litellm/llms/google_pse/search/transformation.py @@ -18,6 +18,7 @@ from litellm.secret_managers.main import get_secret_str class _GooglePSESearchRequestRequired(TypedDict): """Required fields for Google PSE Search API request.""" + q: str # Required - search query cx: str # Required - Programmable Search Engine ID key: str # Required - API key @@ -28,6 +29,7 @@ class GooglePSESearchRequest(_GooglePSESearchRequestRequired, total=False): Google Programmable Search Engine API request format. Based on: https://developers.google.com/custom-search/v1/reference/rest/v1/cse/list """ + num: int # Optional - number of results (1-10), default 10 start: int # Optional - index of first result (default 1) cr: str # Optional - country restrict (e.g., 'countryUS', 'countryGB') @@ -54,17 +56,17 @@ class GooglePSESearchRequest(_GooglePSESearchRequestRequired, total=False): class GooglePSESearchConfig(BaseSearchConfig): GOOGLE_PSE_API_BASE = "https://www.googleapis.com/customsearch/v1" - + @staticmethod def ui_friendly_name() -> str: return "Google PSE" - + def get_http_method(self) -> Literal["GET", "POST"]: """ Google PSE uses GET requests with query parameters. """ return "GET" - + def validate_environment( self, headers: Dict, @@ -74,19 +76,25 @@ class GooglePSESearchConfig(BaseSearchConfig): ) -> Dict: """ Validate environment and return headers. - + Google PSE uses API key as a query parameter, not in headers. This method is called but headers are not used for authentication. """ api_key = api_key or get_secret_str("GOOGLE_PSE_API_KEY") if not api_key: - raise ValueError("GOOGLE_PSE_API_KEY is not set. Set `GOOGLE_PSE_API_KEY` environment variable.") - + raise ValueError( + "GOOGLE_PSE_API_KEY is not set. Set `GOOGLE_PSE_API_KEY` environment variable." + ) + # Also check for search engine ID - search_engine_id = kwargs.get("search_engine_id") or get_secret_str("GOOGLE_PSE_ENGINE_ID") + search_engine_id = kwargs.get("search_engine_id") or get_secret_str( + "GOOGLE_PSE_ENGINE_ID" + ) if not search_engine_id: - raise ValueError("GOOGLE_PSE_ENGINE_ID is not set. Set `GOOGLE_PSE_ENGINE_ID` environment variable or pass `search_engine_id` parameter.") - + raise ValueError( + "GOOGLE_PSE_ENGINE_ID is not set. Set `GOOGLE_PSE_ENGINE_ID` environment variable or pass `search_engine_id` parameter." + ) + headers["Content-Type"] = "application/json" return headers @@ -99,22 +107,25 @@ class GooglePSESearchConfig(BaseSearchConfig): ) -> str: """ Get complete URL for Search endpoint with query parameters. - + Google PSE uses GET requests, so we build the full URL with query params here. The transformed request body (data) contains the parameters needed for the URL. """ from urllib.parse import urlencode - - api_base = api_base or get_secret_str("GOOGLE_PSE_API_BASE") or self.GOOGLE_PSE_API_BASE - + + api_base = ( + api_base + or get_secret_str("GOOGLE_PSE_API_BASE") + or self.GOOGLE_PSE_API_BASE + ) + # Build query parameters from the transformed request body if data and isinstance(data, dict) and "_google_pse_params" in data: params = data["_google_pse_params"] query_string = urlencode(params) return f"{api_base}?{query_string}" - + return api_base - def transform_search_request( self, @@ -126,22 +137,22 @@ class GooglePSESearchConfig(BaseSearchConfig): ) -> Dict: """ Transform Search request to Google PSE API format. - + Transforms Perplexity unified spec parameters: - query → q (same) - max_results → num - search_domain_filter → siteSearch - country → gl - max_tokens_per_page → (not applicable, ignored) - + All other Google PSE-specific parameters are passed through as-is. - + Args: query: Search query (string or list of strings). Google PSE supports single string queries. optional_params: Optional parameters for the request api_key: Google API key search_engine_id: Google Programmable Search Engine ID (cx parameter) - + Returns: Dict with typed request data following GooglePSESearchRequest spec """ @@ -152,7 +163,7 @@ class GooglePSESearchConfig(BaseSearchConfig): # Get API credentials api_key = api_key or get_secret_str("GOOGLE_PSE_API_KEY") search_engine_id = search_engine_id or get_secret_str("GOOGLE_PSE_ENGINE_ID") - + if not api_key: raise ValueError("GOOGLE_PSE_API_KEY is required") if not search_engine_id: @@ -163,13 +174,13 @@ class GooglePSESearchConfig(BaseSearchConfig): "cx": search_engine_id, "key": api_key, } - + # Transform unified spec parameters to Google PSE format if "max_results" in optional_params: # Google PSE supports 1-10 results per request num_results = min(optional_params["max_results"], 10) request_data["num"] = num_results - + if "search_domain_filter" in optional_params: # Convert list to single domain (take first if multiple) domains = optional_params["search_domain_filter"] @@ -179,19 +190,22 @@ class GooglePSESearchConfig(BaseSearchConfig): elif isinstance(domains, str): request_data["siteSearch"] = domains request_data["siteSearchFilter"] = "i" # include - + if "country" in optional_params: # Google PSE uses 2-letter country codes for gl parameter request_data["gl"] = optional_params["country"].upper() - + # Convert to dict before dynamic key assignments result_data = dict(request_data) - + # Pass through all other parameters as-is for param, value in optional_params.items(): - if param not in self.get_supported_perplexity_optional_params() and param not in result_data: + if ( + param not in self.get_supported_perplexity_optional_params() + and param not in result_data + ): result_data[param] = value - + # Store params in special key for URL building (Google PSE uses GET not POST) # Return a wrapper dict that stores params for get_complete_url to use return { @@ -206,22 +220,22 @@ class GooglePSESearchConfig(BaseSearchConfig): ) -> SearchResponse: """ Transform Google PSE API response to LiteLLM unified SearchResponse format. - + Google PSE → LiteLLM mappings: - items[].title → SearchResult.title - items[].link → SearchResult.url - items[].snippet → SearchResult.snippet - No date/last_updated fields in Google PSE response (set to None) - + Args: raw_response: Raw httpx response from Google PSE API logging_obj: Logging object for tracking - + Returns: SearchResponse with standardized format """ response_json = raw_response.json() - + # Transform results to SearchResult objects results = [] for item in response_json.get("items", []): @@ -233,10 +247,8 @@ class GooglePSESearchConfig(BaseSearchConfig): last_updated=None, # Google PSE doesn't provide last_updated in response ) results.append(search_result) - + return SearchResponse( results=results, object="search", ) - - diff --git a/litellm/llms/gradient_ai/chat/transformation.py b/litellm/llms/gradient_ai/chat/transformation.py index d631affdef..1bc5e8896b 100644 --- a/litellm/llms/gradient_ai/chat/transformation.py +++ b/litellm/llms/gradient_ai/chat/transformation.py @@ -12,7 +12,6 @@ GRADIENT_AI_SERVERLESS_ENDPOINT = "https://inference.do-ai.run" class GradientAIConfig(OpenAILikeChatConfig): - k: Optional[int] = None kb_filters: Optional[List[Dict]] = None filter_kb_content_by_query_metadata: Optional[bool] = None @@ -21,7 +20,9 @@ class GradientAIConfig(OpenAILikeChatConfig): include_retrieval_info: Optional[bool] = None include_guardrails_info: Optional[bool] = None provide_citations: Optional[bool] = None - retrieval_method: Optional[Literal["rewrite", "step_back", "sub_queries", "none"]] = None + retrieval_method: Optional[ + Literal["rewrite", "step_back", "sub_queries", "none"] + ] = None def __init__( self, @@ -76,14 +77,16 @@ class GradientAIConfig(OpenAILikeChatConfig): ] return supported_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): + 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, + ): api_key = api_key or get_secret_str("GRADIENT_AI_API_KEY") if api_key is None: raise ValueError("GradientAI API key not found") @@ -107,7 +110,10 @@ class GradientAIConfig(OpenAILikeChatConfig): if api_base and api_base != GRADIENT_AI_SERVERLESS_ENDPOINT: complete_url = f"{api_base}/api/v1/chat/completions" - elif gradient_ai_endpoint and gradient_ai_endpoint != GRADIENT_AI_SERVERLESS_ENDPOINT: + elif ( + gradient_ai_endpoint + and gradient_ai_endpoint != GRADIENT_AI_SERVERLESS_ENDPOINT + ): complete_url = f"{gradient_ai_endpoint}/api/v1/chat/completions" return complete_url @@ -139,9 +145,10 @@ class GradientAIConfig(OpenAILikeChatConfig): optional_params[param] = value elif not drop_params: from litellm.utils import UnsupportedParamsError + raise UnsupportedParamsError( status_code=400, - message=f"GradientAI does not support parameter '{param}'. To drop unsupported params, set `drop_params=True`." + message=f"GradientAI does not support parameter '{param}'. To drop unsupported params, set `drop_params=True`.", ) return optional_params diff --git a/litellm/llms/heroku/chat/transformation.py b/litellm/llms/heroku/chat/transformation.py index a64d8afe63..d95e953636 100644 --- a/litellm/llms/heroku/chat/transformation.py +++ b/litellm/llms/heroku/chat/transformation.py @@ -12,10 +12,12 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( from litellm.types.llms.openai import AllMessageValues from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig + # Base error class for Heroku class HerokuError(Exception): pass + class HerokuChatConfig(OpenAIGPTConfig): @overload def _transform_messages( @@ -49,19 +51,31 @@ class HerokuChatConfig(OpenAIGPTConfig): messages=messages, model=model, is_async=False ) - def _get_openai_compatible_provider_info(self, api_base: Optional[str], api_key: Optional[str]) -> Tuple[Optional[str], Optional[str]]: + def _get_openai_compatible_provider_info( + self, api_base: Optional[str], api_key: Optional[str] + ) -> Tuple[Optional[str], Optional[str]]: api_base = api_base or os.getenv("HEROKU_API_BASE") api_key = api_key or os.getenv("HEROKU_API_KEY") - + return api_base, api_key - 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: + 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_openai_compatible_provider_info(api_base, api_key) if not api_base: - raise HerokuError("No api base was set. Please provide an api_base, or set the HEROKU_API_BASE environment variable.") - - if not api_base.endswith("/v1/chat/completions"): - api_base = f"{api_base}/v1/chat/completions" + raise HerokuError( + "No api base was set. Please provide an api_base, or set the HEROKU_API_BASE environment variable." + ) - return api_base \ No newline at end of file + if not api_base.endswith("/v1/chat/completions"): + api_base = f"{api_base}/v1/chat/completions" + + return api_base diff --git a/litellm/llms/hosted_vllm/chat/transformation.py b/litellm/llms/hosted_vllm/chat/transformation.py index 35dfa8a385..05db1544a2 100644 --- a/litellm/llms/hosted_vllm/chat/transformation.py +++ b/litellm/llms/hosted_vllm/chat/transformation.py @@ -153,9 +153,7 @@ class HostedVLLMChatConfig(OpenAIGPTConfig): ] existing_content = message.get("content") if isinstance(existing_content, str): - new_content.append( - {"type": "text", "text": existing_content} - ) + new_content.append({"type": "text", "text": existing_content}) elif isinstance(existing_content, list): new_content.extend(existing_content) message["content"] = new_content # type: ignore diff --git a/litellm/llms/hosted_vllm/rerank/transformation.py b/litellm/llms/hosted_vllm/rerank/transformation.py index 8316e923df..8066e53afc 100644 --- a/litellm/llms/hosted_vllm/rerank/transformation.py +++ b/litellm/llms/hosted_vllm/rerank/transformation.py @@ -38,8 +38,8 @@ class HostedVLLMRerankConfig(BaseRerankConfig): pass def get_complete_url( - self, - api_base: Optional[str], + self, + api_base: Optional[str], model: str, optional_params: Optional[dict] = None, ) -> str: @@ -82,14 +82,16 @@ class HostedVLLMRerankConfig(BaseRerankConfig): """ if max_chunks_per_doc is not None: raise ValueError("Hosted VLLM does not support max_chunks_per_doc") - - return dict(OptionalRerankParams( - query=query, - documents=documents, - top_n=top_n, - rank_fields=rank_fields, - return_documents=return_documents, - )) + + return dict( + OptionalRerankParams( + query=query, + documents=documents, + top_n=top_n, + rank_fields=rank_fields, + return_documents=return_documents, + ) + ) def validate_environment( self, @@ -124,7 +126,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig): raise ValueError("query is required for Hosted VLLM rerank") if "documents" not in optional_rerank_params: raise ValueError("documents is required for Hosted VLLM rerank") - + rerank_request = RerankRequest( model=model, query=optional_rerank_params["query"], @@ -161,12 +163,16 @@ class HostedVLLMRerankConfig(BaseRerankConfig): def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] ) -> BaseLLMException: - return HostedVLLMRerankError(message=error_message, status_code=status_code, headers=headers) + return HostedVLLMRerankError( + message=error_message, status_code=status_code, headers=headers + ) def _transform_response(self, response: dict) -> RerankResponse: # Extract usage information usage_data = response.get("usage", {}) - _billed_units = RerankBilledUnits(total_tokens=usage_data.get("total_tokens", 0)) + _billed_units = RerankBilledUnits( + total_tokens=usage_data.get("total_tokens", 0) + ) _tokens = RerankTokens(input_tokens=usage_data.get("total_tokens", 0)) rerank_meta = RerankResponseMeta(billed_units=_billed_units, tokens=_tokens) @@ -207,4 +213,4 @@ class HostedVLLMRerankConfig(BaseRerankConfig): id=response.get("id") or str(uuid.uuid4()), results=rerank_results, meta=rerank_meta, - ) \ No newline at end of file + ) diff --git a/litellm/llms/huggingface/embedding/transformation.py b/litellm/llms/huggingface/embedding/transformation.py index 88d42cfcdc..03088d6e15 100644 --- a/litellm/llms/huggingface/embedding/transformation.py +++ b/litellm/llms/huggingface/embedding/transformation.py @@ -40,17 +40,17 @@ class HuggingFaceEmbeddingConfig(BaseConfig): Reference: https://huggingface.github.io/text-generation-inference/#/Text%20Generation%20Inference/compat_generate """ - hf_task: Optional[hf_tasks] = ( - None # litellm-specific param, used to know the api spec to use when calling huggingface api - ) + hf_task: Optional[ + hf_tasks + ] = None # litellm-specific param, used to know the api spec to use when calling huggingface api best_of: Optional[int] = None decoder_input_details: Optional[bool] = None details: Optional[bool] = True # enables returning logprobs + best of max_new_tokens: Optional[int] = None repetition_penalty: Optional[float] = None - return_full_text: Optional[bool] = ( - False # by default don't return the input as part of the output - ) + return_full_text: Optional[ + bool + ] = False # by default don't return the input as part of the output seed: Optional[int] = None temperature: Optional[float] = None top_k: Optional[int] = None @@ -120,9 +120,9 @@ class HuggingFaceEmbeddingConfig(BaseConfig): optional_params["top_p"] = value if param == "n": optional_params["best_of"] = value - optional_params["do_sample"] = ( - True # Need to sample if you want best of for hf inference endpoints - ) + optional_params[ + "do_sample" + ] = True # Need to sample if you want best of for hf inference endpoints if param == "stream": optional_params["stream"] = value if param == "stop": @@ -363,9 +363,9 @@ class HuggingFaceEmbeddingConfig(BaseConfig): "content-type": "application/json", } if api_key is not None: - default_headers["Authorization"] = ( - f"Bearer {api_key}" # Huggingface Inference Endpoint default is to accept bearer tokens - ) + default_headers[ + "Authorization" + ] = f"Bearer {api_key}" # Huggingface Inference Endpoint default is to accept bearer tokens headers = {**headers, **default_headers} return headers diff --git a/litellm/llms/huggingface/rerank/transformation.py b/litellm/llms/huggingface/rerank/transformation.py index b386daf1c8..3f83b8e422 100644 --- a/litellm/llms/huggingface/rerank/transformation.py +++ b/litellm/llms/huggingface/rerank/transformation.py @@ -61,8 +61,8 @@ class HuggingFaceRerankConfig(BaseRerankConfig): return "https://api-inference.huggingface.co" def get_complete_url( - self, - api_base: Optional[str], + self, + api_base: Optional[str], model: str, optional_params: Optional[dict] = None, ) -> str: diff --git a/litellm/llms/infinity/common_utils.py b/litellm/llms/infinity/common_utils.py index 089818c829..67c54caff9 100644 --- a/litellm/llms/infinity/common_utils.py +++ b/litellm/llms/infinity/common_utils.py @@ -6,11 +6,8 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException class InfinityError(BaseLLMException): def __init__( - self, - status_code: int, - message: str, - headers: Union[dict, httpx.Headers] = {} - ): + self, status_code: int, message: str, headers: Union[dict, httpx.Headers] = {} + ): self.status_code = status_code self.message = message self.request = httpx.Request( diff --git a/litellm/llms/infinity/rerank/transformation.py b/litellm/llms/infinity/rerank/transformation.py index 1c15de714b..314bf2f8a3 100644 --- a/litellm/llms/infinity/rerank/transformation.py +++ b/litellm/llms/infinity/rerank/transformation.py @@ -27,8 +27,8 @@ from ..common_utils import InfinityError class InfinityRerankConfig(CohereRerankConfig): def get_complete_url( - self, - api_base: Optional[str], + self, + api_base: Optional[str], model: str, optional_params: Optional[dict] = None, ) -> str: diff --git a/litellm/llms/jina_ai/rerank/transformation.py b/litellm/llms/jina_ai/rerank/transformation.py index 0fddd754a9..48d876f8ea 100644 --- a/litellm/llms/jina_ai/rerank/transformation.py +++ b/litellm/llms/jina_ai/rerank/transformation.py @@ -51,13 +51,15 @@ class JinaAIRerankConfig(BaseRerankConfig): for k, v in non_default_params.items(): if k in supported_params: optional_params[k] = v - return dict(OptionalRerankParams( - **optional_params, - )) + return dict( + OptionalRerankParams( + **optional_params, + ) + ) def get_complete_url( - self, - api_base: Optional[str], + self, + api_base: Optional[str], model: str, optional_params: Optional[dict] = None, ) -> str: @@ -127,9 +129,9 @@ class JinaAIRerankConfig(BaseRerankConfig): ) # Return response def validate_environment( - self, - headers: Dict, - model: str, + self, + headers: Dict, + model: str, api_key: Optional[str] = None, optional_params: Optional[dict] = None, ) -> Dict: diff --git a/litellm/llms/lambda_ai/chat/transformation.py b/litellm/llms/lambda_ai/chat/transformation.py index 2d481d6682..262a189428 100644 --- a/litellm/llms/lambda_ai/chat/transformation.py +++ b/litellm/llms/lambda_ai/chat/transformation.py @@ -13,7 +13,7 @@ class LambdaAIChatConfig(OpenAILikeChatConfig): """ Lambda AI is OpenAI-compatible with standard endpoints """ - + @property def custom_llm_provider(self) -> Optional[str]: return "lambda_ai" @@ -28,4 +28,4 @@ class LambdaAIChatConfig(OpenAILikeChatConfig): or "https://api.lambda.ai/v1" # Default Lambda API base URL ) # type: ignore dynamic_api_key = api_key or get_secret_str("LAMBDA_API_KEY") - return api_base, dynamic_api_key \ No newline at end of file + return api_base, dynamic_api_key diff --git a/litellm/llms/langgraph/__init__.py b/litellm/llms/langgraph/__init__.py index aa075dc96c..6d7b490ed6 100644 --- a/litellm/llms/langgraph/__init__.py +++ b/litellm/llms/langgraph/__init__.py @@ -1,4 +1,3 @@ from litellm.llms.langgraph.chat.transformation import LangGraphConfig __all__ = ["LangGraphConfig"] - diff --git a/litellm/llms/langgraph/chat/__init__.py b/litellm/llms/langgraph/chat/__init__.py index aa075dc96c..6d7b490ed6 100644 --- a/litellm/llms/langgraph/chat/__init__.py +++ b/litellm/llms/langgraph/chat/__init__.py @@ -1,4 +1,3 @@ from litellm.llms.langgraph.chat.transformation import LangGraphConfig __all__ = ["LangGraphConfig"] - diff --git a/litellm/llms/langgraph/chat/sse_iterator.py b/litellm/llms/langgraph/chat/sse_iterator.py index cf81998055..2eb17b4d4b 100644 --- a/litellm/llms/langgraph/chat/sse_iterator.py +++ b/litellm/llms/langgraph/chat/sse_iterator.py @@ -232,4 +232,3 @@ class LangGraphSSEStreamIterator: except Exception as e: verbose_logger.error(f"Error in LangGraph SSE stream: {str(e)}") raise StopAsyncIteration - diff --git a/litellm/llms/langgraph/chat/transformation.py b/litellm/llms/langgraph/chat/transformation.py index b6afa5ab1a..00cc3a8f51 100644 --- a/litellm/llms/langgraph/chat/transformation.py +++ b/litellm/llms/langgraph/chat/transformation.py @@ -66,9 +66,7 @@ class LangGraphConfig(BaseConfig): from litellm.secret_managers.main import get_secret_str api_base = ( - api_base - or get_secret_str("LANGGRAPH_API_BASE") - or "http://localhost:2024" + api_base or get_secret_str("LANGGRAPH_API_BASE") or "http://localhost:2024" ) api_key = api_key or get_secret_str("LANGGRAPH_API_KEY") @@ -166,7 +164,7 @@ class LangGraphConfig(BaseConfig): # Handle content that might be a list if isinstance(content, list): content = convert_content_list_to_str(msg) - + # Ensure content is a string if not isinstance(content, str): content = str(content) @@ -510,4 +508,3 @@ class LangGraphConfig(BaseConfig): LangGraph has native streaming support, so we don't need to fake stream. """ return False - diff --git a/litellm/llms/lemonade/chat/transformation.py b/litellm/llms/lemonade/chat/transformation.py index 8cba844435..a9039388a4 100644 --- a/litellm/llms/lemonade/chat/transformation.py +++ b/litellm/llms/lemonade/chat/transformation.py @@ -63,20 +63,20 @@ class LemonadeChatConfig(OpenAILikeChatConfig): def get_models(self, api_key: Optional[str] = None, api_base: Optional[str] = None): """ Get available models from Lemonade API. - + This method queries the Lemonade /models endpoint to retrieve the list of available models. - + Args: api_key: Optional API key (Lemonade doesn't require authentication) api_base: Optional API base URL (defaults to LEMONADE_API_BASE env var or http://localhost:8000) - + Returns: List of model names prefixed with "lemonade/" """ api_base, api_key = self._get_openai_compatible_provider_info( api_base=api_base, api_key=api_key ) - + if api_base is None: raise ValueError( "LEMONADE_API_BASE is not set. Please set the environment variable to query Lemonade's /models endpoint." @@ -113,7 +113,6 @@ class LemonadeChatConfig(OpenAILikeChatConfig): key = "lemonade" return api_base, key - def transform_response( self, model: str, @@ -146,4 +145,3 @@ class LemonadeChatConfig(OpenAILikeChatConfig): setattr(model_response, "model", "lemonade/" + model) return model_response - \ No newline at end of file diff --git a/litellm/llms/lemonade/cost_calculator.py b/litellm/llms/lemonade/cost_calculator.py index 27e1ca275f..2042f6d0d4 100644 --- a/litellm/llms/lemonade/cost_calculator.py +++ b/litellm/llms/lemonade/cost_calculator.py @@ -15,21 +15,21 @@ def cost_per_token( ) -> Tuple[float, float]: """ Calculate cost per token for Lemonade models. - + Since Lemonade is a local/self-hosted deployment, there are no per-token costs. This function returns (0.0, 0.0) for all models to allow cost tracking to work without errors for any Lemonade model, regardless of whether it's in the model_prices_and_context_window.json file. - + Args: model: The model name (with or without "lemonade/" prefix) usage: Usage object containing token counts - + Returns: Tuple of (prompt_cost, completion_cost) - always (0.0, 0.0) for Lemonade """ # Lemonade is self-hosted/local, so cost is always 0 prompt_cost = 0.0 completion_cost = 0.0 - + return prompt_cost, completion_cost diff --git a/litellm/llms/linkup/__init__.py b/litellm/llms/linkup/__init__.py index b1553a1737..c761584b07 100644 --- a/litellm/llms/linkup/__init__.py +++ b/litellm/llms/linkup/__init__.py @@ -4,4 +4,3 @@ Linkup API integration module. from litellm.llms.linkup.search.transformation import LinkupSearchConfig __all__ = ["LinkupSearchConfig"] - diff --git a/litellm/llms/linkup/search/__init__.py b/litellm/llms/linkup/search/__init__.py index b47af3f305..667c463023 100644 --- a/litellm/llms/linkup/search/__init__.py +++ b/litellm/llms/linkup/search/__init__.py @@ -4,4 +4,3 @@ Linkup Search API module. from litellm.llms.linkup.search.transformation import LinkupSearchConfig __all__ = ["LinkupSearchConfig"] - diff --git a/litellm/llms/linkup/search/transformation.py b/litellm/llms/linkup/search/transformation.py index bbe76664b4..0554b8ab34 100644 --- a/litellm/llms/linkup/search/transformation.py +++ b/litellm/llms/linkup/search/transformation.py @@ -79,9 +79,7 @@ class LinkupSearchConfig(BaseSearchConfig): """ Get complete URL for Search endpoint. """ - api_base = ( - api_base or get_secret_str("LINKUP_API_BASE") or self.LINKUP_API_BASE - ) + api_base = api_base or get_secret_str("LINKUP_API_BASE") or self.LINKUP_API_BASE # Append "/search" to the api base if it's not already there if not api_base.endswith("/search"): @@ -203,4 +201,3 @@ class LinkupSearchConfig(BaseSearchConfig): results=results, object="search", ) - diff --git a/litellm/llms/litellm_proxy/image_generation/transformation.py b/litellm/llms/litellm_proxy/image_generation/transformation.py index 6174424154..3932070e96 100644 --- a/litellm/llms/litellm_proxy/image_generation/transformation.py +++ b/litellm/llms/litellm_proxy/image_generation/transformation.py @@ -8,6 +8,7 @@ from litellm.secret_managers.main import get_secret_str class LiteLLMProxyImageGenerationConfig(GPTImageGenerationConfig): """Configuration for image generation requests routed through LiteLLM Proxy.""" + def validate_environment( self, headers: dict, diff --git a/litellm/llms/litellm_proxy/responses/transformation.py b/litellm/llms/litellm_proxy/responses/transformation.py index a122b76875..e5bbaa78d1 100644 --- a/litellm/llms/litellm_proxy/responses/transformation.py +++ b/litellm/llms/litellm_proxy/responses/transformation.py @@ -15,7 +15,7 @@ from litellm.types.utils import LlmProviders class LiteLLMProxyResponsesAPIConfig(OpenAIResponsesAPIConfig): """ Configuration for LiteLLM Proxy Responses API support. - + Extends OpenAI's config since the proxy follows OpenAI's API spec, but uses LITELLM_PROXY_API_BASE for the base URL. """ @@ -31,11 +31,11 @@ class LiteLLMProxyResponsesAPIConfig(OpenAIResponsesAPIConfig): ) -> str: """ Get the endpoint for LiteLLM Proxy responses API. - + Uses LITELLM_PROXY_API_BASE environment variable if api_base is not provided. """ api_base = api_base or get_secret_str("LITELLM_PROXY_API_BASE") - + if api_base is None: raise ValueError( "api_base not set for LiteLLM Proxy responses API. " diff --git a/litellm/llms/litellm_proxy/skills/code_execution.py b/litellm/llms/litellm_proxy/skills/code_execution.py index d307b8b36d..2b567f0376 100644 --- a/litellm/llms/litellm_proxy/skills/code_execution.py +++ b/litellm/llms/litellm_proxy/skills/code_execution.py @@ -22,17 +22,18 @@ from litellm._logging import verbose_logger class LiteLLMInternalTools(str, Enum): """ Enum for internal LiteLLM tools that are injected into requests. - + These tools are handled automatically by LiteLLM hooks and are not passed to the underlying LLM provider directly. """ + CODE_EXECUTION = "litellm_code_execution" def get_litellm_code_execution_tool() -> Dict[str, Any]: """ Returns the litellm_code_execution tool definition in OpenAI format. - + This tool enables automatic code execution in a sandboxed environment when skills include executable Python code. """ @@ -44,21 +45,18 @@ def get_litellm_code_execution_tool() -> Dict[str, Any]: "parameters": { "type": "object", "properties": { - "code": { - "type": "string", - "description": "Python code to execute" - } + "code": {"type": "string", "description": "Python code to execute"} }, - "required": ["code"] - } - } + "required": ["code"], + }, + }, } def get_litellm_code_execution_tool_anthropic() -> Dict[str, Any]: """ Returns the litellm_code_execution tool definition in Anthropic/messages API format. - + This tool enables automatic code execution in a sandboxed environment when skills include executable Python code. """ @@ -68,13 +66,10 @@ def get_litellm_code_execution_tool_anthropic() -> Dict[str, Any]: "input_schema": { "type": "object", "properties": { - "code": { - "type": "string", - "description": "Python code to execute" - } + "code": {"type": "string", "description": "Python code to execute"} }, - "required": ["code"] - } + "required": ["code"], + }, } @@ -85,12 +80,12 @@ LITELLM_CODE_EXECUTION_TOOL = get_litellm_code_execution_tool() class CodeExecutionHandler: """ Handles automatic code execution for LiteLLM skills. - + When enabled, this handler intercepts LLM responses with code execution tool calls, executes them in a sandbox, and continues the conversation automatically until completion. """ - + def __init__( self, max_iterations: Optional[int] = None, @@ -100,10 +95,10 @@ class CodeExecutionHandler: DEFAULT_MAX_ITERATIONS, DEFAULT_SANDBOX_TIMEOUT, ) - + self.max_iterations = max_iterations or DEFAULT_MAX_ITERATIONS self.sandbox_timeout = sandbox_timeout or DEFAULT_SANDBOX_TIMEOUT - + async def execute_with_code_execution( self, model: str, @@ -115,14 +110,14 @@ class CodeExecutionHandler: ) -> Dict[str, Any]: """ Execute an LLM call with automatic code execution handling. - + This method: 1. Makes the initial LLM call 2. If model calls litellm_code_execution, executes the code 3. Continues conversation with results 4. Repeats until model stops calling tools 5. Returns final response with generated files inline - + Args: model: Model to use messages: Initial messages @@ -130,7 +125,7 @@ class CodeExecutionHandler: skill_files: Dict of skill files for execution skill_id: Optional skill ID for tracking **kwargs: Additional args for litellm.acompletion - + Returns: Dict with: - response: Final LLM response @@ -141,19 +136,19 @@ class CodeExecutionHandler: from litellm.llms.litellm_proxy.skills.sandbox_executor import ( SkillsSandboxExecutor, ) - + current_messages = list(messages) generated_files: List[Dict[str, Any]] = [] # Files returned directly execution_results: List[Dict] = [] - + executor = SkillsSandboxExecutor(timeout=self.sandbox_timeout) response: Any = None # Initialize to avoid possibly unbound error - + for iteration in range(self.max_iterations): verbose_logger.debug( f"CodeExecutionHandler: Iteration {iteration + 1}/{self.max_iterations}" ) - + # Make LLM call response = await litellm.acompletion( model=model, @@ -161,10 +156,10 @@ class CodeExecutionHandler: tools=tools, **kwargs, ) - + assistant_message = response.choices[0].message # type: ignore stop_reason = response.choices[0].finish_reason # type: ignore - + # Build assistant message for conversation history assistant_msg_dict: Dict[str, Any] = { "role": "assistant", @@ -177,13 +172,13 @@ class CodeExecutionHandler: "type": "function", "function": { "name": tc.function.name, - "arguments": tc.function.arguments - } + "arguments": tc.function.arguments, + }, } for tc in assistant_message.tool_calls ] current_messages.append(assistant_msg_dict) - + # Check if we're done (no tool calls or not tool_calls finish reason) if stop_reason != "tool_calls" or not assistant_message.tool_calls: verbose_logger.debug( @@ -195,21 +190,21 @@ class CodeExecutionHandler: "execution_results": execution_results, "messages": current_messages, } - + # Handle tool calls for tool_call in assistant_message.tool_calls: tool_name = tool_call.function.name - + if tool_name == LiteLLMInternalTools.CODE_EXECUTION.value: # Execute code in sandbox try: args = json.loads(tool_call.function.arguments) code = args.get("code", "") - + verbose_logger.debug( f"CodeExecutionHandler: Executing code ({len(code)} chars)" ) - + exec_result = executor.execute( code=code, skill_files=skill_files, @@ -218,62 +213,74 @@ class CodeExecutionHandler: verbose_logger.debug( f"CodeExecutionHandler: Execution result: {exec_result}" ) - - execution_results.append({ - "iteration": iteration, - "success": exec_result["success"], - "output": exec_result["output"], - "error": exec_result["error"], - "files": [f["name"] for f in exec_result["files"]], - }) - + + execution_results.append( + { + "iteration": iteration, + "success": exec_result["success"], + "output": exec_result["output"], + "error": exec_result["error"], + "files": [f["name"] for f in exec_result["files"]], + } + ) + # Build tool result content tool_result = exec_result["output"] or "" - + # Collect generated files (returned directly, no storage) if exec_result["files"]: tool_result += "\n\nGenerated files:" for f in exec_result["files"]: file_content = base64.b64decode(f["content_base64"]) # Add to generated files list (returned in response) - generated_files.append({ - "name": f["name"], - "mime_type": f["mime_type"], - "content_base64": f["content_base64"], - "size": len(file_content), - }) - tool_result += f"\n- {f['name']} ({len(file_content)} bytes)" - + generated_files.append( + { + "name": f["name"], + "mime_type": f["mime_type"], + "content_base64": f["content_base64"], + "size": len(file_content), + } + ) + tool_result += ( + f"\n- {f['name']} ({len(file_content)} bytes)" + ) + verbose_logger.debug( f"CodeExecutionHandler: Generated file {f['name']} ({len(file_content)} bytes)" ) - + if exec_result["error"]: tool_result += f"\n\nError:\n{exec_result['error']}" - + except Exception as e: tool_result = f"Code execution failed: {str(e)}" - execution_results.append({ - "iteration": iteration, - "success": False, - "error": str(e), - }) - + execution_results.append( + { + "iteration": iteration, + "success": False, + "error": str(e), + } + ) + # Add tool result to messages - current_messages.append({ - "role": "tool", - "tool_call_id": tool_call.id, - "content": tool_result, - }) + current_messages.append( + { + "role": "tool", + "tool_call_id": tool_call.id, + "content": tool_result, + } + ) else: # Non-code-execution tool - pass through # In a full implementation, this would call other tool handlers - current_messages.append({ - "role": "tool", - "tool_call_id": tool_call.id, - "content": f"Tool '{tool_name}' not handled by code execution handler", - }) - + current_messages.append( + { + "role": "tool", + "tool_call_id": tool_call.id, + "content": f"Tool '{tool_name}' not handled by code execution handler", + } + ) + # Max iterations reached verbose_logger.warning( f"CodeExecutionHandler: Max iterations ({self.max_iterations}) reached" @@ -308,4 +315,3 @@ def add_code_execution_tool(tools: Optional[List[Dict]]) -> List[Dict]: # Global handler instance code_execution_handler = CodeExecutionHandler() - diff --git a/litellm/llms/litellm_proxy/skills/constants.py b/litellm/llms/litellm_proxy/skills/constants.py index a2be6961db..a8c2697fce 100644 --- a/litellm/llms/litellm_proxy/skills/constants.py +++ b/litellm/llms/litellm_proxy/skills/constants.py @@ -10,4 +10,3 @@ DEFAULT_MAX_ITERATIONS: int = 10 DEFAULT_SANDBOX_TIMEOUT: int = 120 """Default timeout in seconds for sandbox code execution.""" - diff --git a/litellm/llms/litellm_proxy/skills/handler.py b/litellm/llms/litellm_proxy/skills/handler.py index f44ac4cda9..8e5070c272 100644 --- a/litellm/llms/litellm_proxy/skills/handler.py +++ b/litellm/llms/litellm_proxy/skills/handler.py @@ -15,13 +15,13 @@ from litellm.proxy._types import LiteLLM_SkillsTable, NewSkillRequest def _prisma_skill_to_litellm(prisma_skill) -> LiteLLM_SkillsTable: """ Convert a Prisma skill record to LiteLLM_SkillsTable. - + Handles Base64 decoding of file_content field. """ import base64 data = prisma_skill.model_dump() - + # Decode Base64 file_content back to bytes # model_dump() converts Base64 field to base64-encoded string if data.get("file_content") is not None: @@ -30,7 +30,7 @@ def _prisma_skill_to_litellm(prisma_skill) -> LiteLLM_SkillsTable: elif isinstance(data["file_content"], bytes): # Already bytes, no conversion needed pass - + return LiteLLM_SkillsTable(**data) diff --git a/litellm/llms/litellm_proxy/skills/prompt_injection.py b/litellm/llms/litellm_proxy/skills/prompt_injection.py index 17469274c1..2b86f74122 100644 --- a/litellm/llms/litellm_proxy/skills/prompt_injection.py +++ b/litellm/llms/litellm_proxy/skills/prompt_injection.py @@ -16,7 +16,7 @@ from litellm.proxy._types import LiteLLM_SkillsTable class SkillPromptInjectionHandler: """ Handles skill content extraction and system prompt injection. - + Responsibilities: - Extract SKILL.md content from skill ZIP files - Extract ALL files from ZIP for code execution @@ -27,19 +27,19 @@ class SkillPromptInjectionHandler: def extract_skill_content(self, skill: LiteLLM_SkillsTable) -> Optional[str]: """ Extract skill content from the stored zip file. - + Looks for SKILL.md or README.md in the zip and returns its content. This content describes the skill's capabilities and instructions. - + Args: skill: The skill from LiteLLM database - + Returns: The skill content as a string, or None if not available """ if not skill.file_content: return skill.instructions - + try: zip_buffer = BytesIO(skill.file_content) with zipfile.ZipFile(zip_buffer, "r") as zf: @@ -49,14 +49,14 @@ class SkillPromptInjectionHandler: content = zf.read(name).decode("utf-8") if content: return f"## Skill: {skill.display_title or skill.skill_id}\n\n{content}" - + # Fall back to README.md for name in zf.namelist(): if name.endswith("README.md"): content = zf.read(name).decode("utf-8") if content: return f"## Skill: {skill.display_title or skill.skill_id}\n\n{content}" - + # Fall back to any .md file for name in zf.namelist(): if name.endswith(".md"): @@ -67,27 +67,27 @@ class SkillPromptInjectionHandler: verbose_logger.warning( f"SkillPromptInjectionHandler: Error extracting content from skill {skill.skill_id}: {e}" ) - + return skill.instructions def extract_all_files(self, skill: LiteLLM_SkillsTable) -> Dict[str, bytes]: """ Extract ALL files from skill ZIP for code execution. - + Returns a dict mapping file paths to their binary content. The paths have the skill folder prefix removed (e.g., "slack-gif-creator/core/..." -> "core/..."). - + Args: skill: The skill from LiteLLM database - + Returns: Dict mapping file paths to binary content """ files: Dict[str, bytes] = {} - + if not skill.file_content: return files - + try: zip_buffer = BytesIO(skill.file_content) with zipfile.ZipFile(zip_buffer, "r") as zf: @@ -95,21 +95,21 @@ class SkillPromptInjectionHandler: # Skip directories if name.endswith("/"): continue - + # Remove skill folder prefix (first path component) parts = name.split("/") if len(parts) > 1: clean_path = "/".join(parts[1:]) else: clean_path = name - + if clean_path: files[clean_path] = zf.read(name) except Exception as e: verbose_logger.warning( f"SkillPromptInjectionHandler: Error extracting files from skill {skill.skill_id}: {e}" ) - + return files def inject_skill_content_to_messages( @@ -117,27 +117,29 @@ class SkillPromptInjectionHandler: ) -> dict: """ Inject skill content into the system prompt. - + For Anthropic messages API (use_anthropic_format=True): - Injects into top-level 'system' parameter (not in messages array) - + For OpenAI-style APIs (use_anthropic_format=False): - Injects into messages array with role="system" - + Args: data: The request data dict skill_contents: List of skill content strings to inject use_anthropic_format: If True, use top-level 'system' param for Anthropic - + Returns: Modified data dict with skill content in system prompt """ if not skill_contents: return data - + # Build the skill injection text - skill_section = "\n\n---\n\n# Available Skills\n\n" + "\n\n---\n\n".join(skill_contents) - + skill_section = "\n\n---\n\n# Available Skills\n\n" + "\n\n---\n\n".join( + skill_contents + ) + if use_anthropic_format: # Anthropic messages API: use top-level 'system' parameter current_system = data.get("system", "") @@ -146,19 +148,19 @@ class SkillPromptInjectionHandler: else: data["system"] = skill_section.strip() return data - + # OpenAI-style: inject into messages array messages = data.get("messages", []) if not messages: return data - + # Find or create system message system_msg_idx = None for i, msg in enumerate(messages): if isinstance(msg, dict) and msg.get("role") == "system": system_msg_idx = i break - + if system_msg_idx is not None: # Append to existing system message current_content = messages[system_msg_idx].get("content", "") @@ -166,20 +168,20 @@ class SkillPromptInjectionHandler: else: # Create new system message at the beginning messages.insert(0, {"role": "system", "content": skill_section.strip()}) - + data["messages"] = messages return data def create_execute_code_tool(self, skill_modules: List[str]) -> Dict[str, Any]: """ Create the execute_code tool definition. - + This tool allows the model to execute Python code with access to the skill's modules (e.g., 'from core.gif_builder import GIFBuilder'). - + Args: skill_modules: List of available module paths (e.g., ["core/gif_builder.py"]) - + Returns: OpenAI-style tool definition """ @@ -190,11 +192,11 @@ class SkillPromptInjectionHandler: # Convert path to import: "core/gif_builder.py" -> "from core.gif_builder import ..." import_path = mod.replace("/", ".").replace(".py", "") module_examples.append(f"from {import_path} import ...") - + module_hint = "" if module_examples: module_hint = f" Available modules: {', '.join(module_examples)}" - + return { "type": "function", "function": { @@ -205,12 +207,12 @@ class SkillPromptInjectionHandler: "properties": { "code": { "type": "string", - "description": "Python code to execute. You can import skill modules and use standard libraries." + "description": "Python code to execute. You can import skill modules and use standard libraries.", } }, - "required": ["code"] - } - } + "required": ["code"], + }, + }, } def convert_skill_to_tool(self, skill: LiteLLM_SkillsTable) -> Dict[str, Any]: @@ -263,7 +265,9 @@ class SkillPromptInjectionHandler: return tool - def convert_skill_to_anthropic_tool(self, skill: LiteLLM_SkillsTable) -> Dict[str, Any]: + def convert_skill_to_anthropic_tool( + self, skill: LiteLLM_SkillsTable + ) -> Dict[str, Any]: """ Convert a LiteLLM skill to an Anthropic-style tool (messages API format). @@ -302,4 +306,3 @@ class SkillPromptInjectionHandler: "description": description, "input_schema": input_schema, } - diff --git a/litellm/llms/litellm_proxy/skills/sandbox_executor.py b/litellm/llms/litellm_proxy/skills/sandbox_executor.py index 7676ade5cd..a5c0a539c9 100644 --- a/litellm/llms/litellm_proxy/skills/sandbox_executor.py +++ b/litellm/llms/litellm_proxy/skills/sandbox_executor.py @@ -15,7 +15,7 @@ from litellm._logging import verbose_logger class SkillsSandboxExecutor: """ Executes skill code in llm-sandbox Docker container. - + Responsibilities: - Create sandbox session with skill files - Install requirements @@ -31,7 +31,7 @@ class SkillsSandboxExecutor: ): """ Initialize the sandbox executor. - + Args: timeout: Maximum execution time in seconds backend: Sandbox backend ("docker", "podman", "kubernetes") @@ -50,12 +50,12 @@ class SkillsSandboxExecutor: ) -> Dict[str, Any]: """ Execute code with skill files in sandbox. - + Args: code: Python code to execute skill_files: Dict mapping file paths to binary content requirements: Optional requirements.txt content - + Returns: { "success": bool, @@ -84,10 +84,10 @@ class SkillsSandboxExecutor: "lang": "python", "verbose": False, } - + if self.image: session_kwargs["image"] = self.image - + with SandboxSession(**session_kwargs) as session: # 1. Copy skill files into sandbox using copy_to_runtime import tempfile @@ -100,15 +100,15 @@ class SkillsSandboxExecutor: os.makedirs(os.path.dirname(local_path), exist_ok=True) with open(local_path, "wb") as f: f.write(content) - + # Copy to sandbox sandbox_path = f"/sandbox/{path}" session.copy_to_runtime(local_path, sandbox_path) - + verbose_logger.debug( f"SkillsSandboxExecutor: Copied {len(skill_files)} files to sandbox" ) - + # 2. Install requirements if present req_packages = None if requirements: @@ -116,7 +116,7 @@ class SkillsSandboxExecutor: elif "requirements.txt" in skill_files: req_content = skill_files["requirements.txt"].decode("utf-8") req_packages = req_content.strip().replace("\n", " ") - + if req_packages: # Run pip install as code pip_code = f""" @@ -127,7 +127,7 @@ subprocess.run(['pip', 'install'] + '{req_packages}'.split(), check=True) verbose_logger.debug( "SkillsSandboxExecutor: Installed requirements" ) - + # 3. Execute the code # Wrap code to run from /sandbox directory wrapped_code = f""" @@ -139,11 +139,11 @@ sys.path.insert(0, '/sandbox') {code} """ result = session.run(wrapped_code) - + success = result.exit_code == 0 output = result.stdout or "" error = result.stderr or "" - + if success: verbose_logger.debug( "SkillsSandboxExecutor: Code execution succeeded" @@ -158,21 +158,19 @@ sys.path.insert(0, '/sandbox') verbose_logger.debug( f"SkillsSandboxExecutor: stdout: {output[:500] if output else 'No stdout'}" ) - + # 4. Collect generated files generated_files = self._collect_generated_files(session, skill_files) - + return { "success": success, "output": output, "error": error, "files": generated_files, } - + except Exception as e: - verbose_logger.error( - f"SkillsSandboxExecutor: Execution failed: {e}" - ) + verbose_logger.error(f"SkillsSandboxExecutor: Execution failed: {e}") return { "success": False, "output": "", @@ -187,19 +185,19 @@ sys.path.insert(0, '/sandbox') ) -> List[Dict[str, Any]]: """ Collect files generated during execution. - + Looks for new files in /sandbox that weren't in the original skill files. Focuses on common output types: GIF, PNG, JPG, PDF, CSV, etc. - + Args: session: The sandbox session original_files: Original skill files (to exclude) - + Returns: List of generated files with base64 content """ generated_files: List[Dict[str, Any]] = [] - + try: import tempfile @@ -215,43 +213,46 @@ for root, dirs, filenames in os.walk('/sandbox'): print(json.dumps(files)) """ result = session.run(list_code) - + if result.exit_code == 0 and result.stdout: import json + try: filepaths = json.loads(result.stdout.strip()) except json.JSONDecodeError: filepaths = [] - + for filepath in filepaths: if not filepath: continue - + # Get relative path rel_path = filepath.replace("/sandbox/", "") - + # Skip if it was an original file if rel_path in original_files: continue - + # Copy file from sandbox using copy_from_runtime with tempfile.NamedTemporaryFile(delete=False) as tmp: tmp_path = tmp.name - + try: session.copy_from_runtime(filepath, tmp_path) - + with open(tmp_path, "rb") as f: content = f.read() - + content_b64 = base64.b64encode(content).decode("utf-8") - generated_files.append({ - "name": os.path.basename(filepath), - "path": rel_path, - "content_base64": content_b64, - "mime_type": self._get_mime_type(filepath), - }) - + generated_files.append( + { + "name": os.path.basename(filepath), + "path": rel_path, + "content_base64": content_b64, + "mime_type": self._get_mime_type(filepath), + } + ) + verbose_logger.debug( f"SkillsSandboxExecutor: Collected generated file: {rel_path}" ) @@ -262,12 +263,12 @@ print(json.dumps(files)) finally: if os.path.exists(tmp_path): os.unlink(tmp_path) - + except Exception as e: verbose_logger.warning( f"SkillsSandboxExecutor: Error collecting generated files: {e}" ) - + return generated_files def _get_mime_type(self, filename: str) -> str: @@ -283,4 +284,3 @@ print(json.dumps(files)) "json": "application/json", "txt": "text/plain", }.get(ext, "application/octet-stream") - diff --git a/litellm/llms/litellm_proxy/skills/transformation.py b/litellm/llms/litellm_proxy/skills/transformation.py index e7c999eace..cd000829ca 100644 --- a/litellm/llms/litellm_proxy/skills/transformation.py +++ b/litellm/llms/litellm_proxy/skills/transformation.py @@ -23,7 +23,7 @@ if TYPE_CHECKING: class LiteLLMSkillsTransformationHandler: """ Transformation handler for skills API requests to LiteLLM database operations. - + This is used when custom_llm_provider="litellm_proxy" to store/retrieve skills from the LiteLLM proxy database instead of calling an external API. """ @@ -51,7 +51,7 @@ class LiteLLMSkillsTransformationHandler: ) -> Union[Skill, Coroutine[Any, Any, Skill]]: """ Create a skill in LiteLLM database. - + Args: display_title: Display title for the skill description: Description of the skill @@ -63,7 +63,7 @@ class LiteLLMSkillsTransformationHandler: metadata: Additional metadata user_id: User ID for tracking _is_async: Whether to return a coroutine - + Returns: Skill object or coroutine that returns Skill """ @@ -84,7 +84,9 @@ class LiteLLMSkillsTransformationHandler: if isinstance(first_file, tuple) and len(first_file) >= 2: file_name = first_file[0] file_content = first_file[1] - file_type = first_file[2] if len(first_file) > 2 else "application/zip" + file_type = ( + first_file[2] if len(first_file) > 2 else "application/zip" + ) if _is_async: return self._async_create_skill( @@ -97,8 +99,9 @@ class LiteLLMSkillsTransformationHandler: metadata=metadata, user_id=user_id, ) - + import asyncio + return asyncio.get_event_loop().run_until_complete( self._async_create_skill( display_title=display_title, @@ -156,14 +159,14 @@ class LiteLLMSkillsTransformationHandler: ) -> Union[ListSkillsResponse, Coroutine[Any, Any, ListSkillsResponse]]: """ List skills from LiteLLM database. - + Args: limit: Maximum number of skills to return offset: Number of skills to skip _is_async: Whether to return a coroutine logging_obj: LiteLLM logging object litellm_call_id: Call ID for logging - + Returns: ListSkillsResponse or coroutine that returns ListSkillsResponse """ @@ -178,8 +181,9 @@ class LiteLLMSkillsTransformationHandler: if _is_async: return self._async_list_skills(limit=limit, offset=offset) - + import asyncio + return asyncio.get_event_loop().run_until_complete( self._async_list_skills(limit=limit, offset=offset) ) @@ -215,13 +219,13 @@ class LiteLLMSkillsTransformationHandler: ) -> Union[Skill, Coroutine[Any, Any, Skill]]: """ Get a skill from LiteLLM database. - + Args: skill_id: The skill ID to retrieve _is_async: Whether to return a coroutine logging_obj: LiteLLM logging object litellm_call_id: Call ID for logging - + Returns: Skill or coroutine that returns Skill """ @@ -236,8 +240,9 @@ class LiteLLMSkillsTransformationHandler: if _is_async: return self._async_get_skill(skill_id=skill_id) - + import asyncio + return asyncio.get_event_loop().run_until_complete( self._async_get_skill(skill_id=skill_id) ) @@ -260,13 +265,13 @@ class LiteLLMSkillsTransformationHandler: ) -> Union[DeleteSkillResponse, Coroutine[Any, Any, DeleteSkillResponse]]: """ Delete a skill from LiteLLM database. - + Args: skill_id: The skill ID to delete _is_async: Whether to return a coroutine logging_obj: LiteLLM logging object litellm_call_id: Call ID for logging - + Returns: DeleteSkillResponse or coroutine that returns DeleteSkillResponse """ @@ -281,8 +286,9 @@ class LiteLLMSkillsTransformationHandler: if _is_async: return self._async_delete_skill(skill_id=skill_id) - + import asyncio + return asyncio.get_event_loop().run_until_complete( self._async_delete_skill(skill_id=skill_id) ) @@ -301,16 +307,16 @@ class LiteLLMSkillsTransformationHandler: def _db_skill_to_response(self, db_skill: Any) -> Skill: """ Convert a database skill record to Anthropic-compatible Skill response. - + Args: db_skill: LiteLLM_SkillsTable record - + Returns: Skill object """ created_at = "" updated_at = "" - + if hasattr(db_skill, "created_at") and db_skill.created_at: created_at = ( db_skill.created_at.isoformat() @@ -333,4 +339,3 @@ class LiteLLMSkillsTransformationHandler: source=db_skill.source or "custom", type="skill", ) - diff --git a/litellm/llms/llamafile/chat/transformation.py b/litellm/llms/llamafile/chat/transformation.py index b0f8cd3fc3..3387a0eb6a 100644 --- a/litellm/llms/llamafile/chat/transformation.py +++ b/litellm/llms/llamafile/chat/transformation.py @@ -15,7 +15,9 @@ class LlamafileChatConfig(OpenAIGPTConfig): If both are None, a fake API key is returned. """ - return api_key or get_secret_str("LLAMAFILE_API_KEY") or "fake-api-key" # llamafile does not require an API key + return ( + api_key or get_secret_str("LLAMAFILE_API_KEY") or "fake-api-key" + ) # llamafile does not require an API key @staticmethod def _resolve_api_base(api_base: Optional[str] = None) -> Optional[str]: @@ -25,13 +27,10 @@ class LlamafileChatConfig(OpenAIGPTConfig): If both are None, a default Llamafile server URL is returned. See: https://github.com/Mozilla-Ocho/llamafile/blob/bd1bbe9aabb1ee12dbdcafa8936db443c571eb9d/README.md#L61 """ - return api_base or get_secret_str("LLAMAFILE_API_BASE") or "http://127.0.0.1:8080/v1" # type: ignore - + return api_base or get_secret_str("LLAMAFILE_API_BASE") or "http://127.0.0.1:8080/v1" # type: ignore def _get_openai_compatible_provider_info( - self, - api_base: Optional[str], - api_key: Optional[str] + self, api_base: Optional[str], api_key: Optional[str] ) -> Tuple[Optional[str], Optional[str]]: """Attempts to ensure that the API base and key are set, preferring user-provided values, before falling back to secret manager values (``LLAMAFILE_API_BASE`` and ``LLAMAFILE_API_KEY`` diff --git a/litellm/llms/lm_studio/chat/transformation.py b/litellm/llms/lm_studio/chat/transformation.py index 7b188ff33f..64ed38467d 100644 --- a/litellm/llms/lm_studio/chat/transformation.py +++ b/litellm/llms/lm_studio/chat/transformation.py @@ -18,7 +18,7 @@ class LMStudioChatConfig(OpenAIGPTConfig): api_key or get_secret_str("LM_STUDIO_API_KEY") or "fake-api-key" ) # LM Studio does not require an api key, but OpenAI client requires non-None value return api_base, dynamic_api_key - + def map_openai_params( self, non_default_params: dict, @@ -46,4 +46,4 @@ class LMStudioChatConfig(OpenAIGPTConfig): optional_params=optional_params, model=model, drop_params=drop_params, - ) \ No newline at end of file + ) diff --git a/litellm/llms/manus/__init__.py b/litellm/llms/manus/__init__.py index 81eef02546..03f1707d44 100644 --- a/litellm/llms/manus/__init__.py +++ b/litellm/llms/manus/__init__.py @@ -1,2 +1 @@ # Manus provider implementation - diff --git a/litellm/llms/manus/files/__init__.py b/litellm/llms/manus/files/__init__.py index 66d23ca034..3659eef17c 100644 --- a/litellm/llms/manus/files/__init__.py +++ b/litellm/llms/manus/files/__init__.py @@ -1,2 +1 @@ # Manus Files API implementation - diff --git a/litellm/llms/manus/files/transformation.py b/litellm/llms/manus/files/transformation.py index a796501196..3381a5327e 100644 --- a/litellm/llms/manus/files/transformation.py +++ b/litellm/llms/manus/files/transformation.py @@ -74,11 +74,7 @@ class ManusFilesConfig(BaseFilesConfig): Manus uses API_KEY header instead of Authorization: Bearer. For file uploads, don't set Content-Type - httpx will set it for multipart. """ - api_key = ( - api_key - or litellm.api_key - or get_secret_str("MANUS_API_KEY") - ) + api_key = api_key or litellm.api_key or get_secret_str("MANUS_API_KEY") if not api_key: raise ValueError( @@ -194,14 +190,14 @@ class ManusFilesConfig(BaseFilesConfig): optional_params=optional_params, litellm_params=litellm_params, ) - + # Get API key api_key = ( litellm_params.get("api_key") or litellm.api_key or get_secret_str("MANUS_API_KEY") ) - + if not api_key: raise ValueError( "Manus API key is required. Set MANUS_API_KEY environment variable or pass api_key parameter." @@ -436,4 +432,3 @@ class ManusFilesConfig(BaseFilesConfig): ) -> HttpxBinaryResponseContent: """Transform file content response.""" return HttpxBinaryResponseContent(response=raw_response) - diff --git a/litellm/llms/manus/responses/__init__.py b/litellm/llms/manus/responses/__init__.py index e8cabc5426..7df60c923b 100644 --- a/litellm/llms/manus/responses/__init__.py +++ b/litellm/llms/manus/responses/__init__.py @@ -1,2 +1 @@ # Manus Responses API implementation - diff --git a/litellm/llms/manus/responses/transformation.py b/litellm/llms/manus/responses/transformation.py index bf1a6fab50..510c41304a 100644 --- a/litellm/llms/manus/responses/transformation.py +++ b/litellm/llms/manus/responses/transformation.py @@ -33,12 +33,12 @@ MANUS_API_BASE = "https://api.manus.im" class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): """ Configuration for Manus API's Responses API. - + Manus API is OpenAI-compatible but has some differences: - API key passed via `API_KEY` header (not `Authorization: Bearer`) - Model format: `manus/{agent_profile}` (e.g., `manus/manus-1.6`) - Requires `extra_body` with `task_mode: "agent"` and `agent_profile` - + Reference: https://open.manus.im/docs/openai-compatibility """ @@ -62,10 +62,10 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): def _extract_agent_profile(self, model: str) -> str: """ Extract agent profile from model name. - + Model format: `manus/{agent_profile}` Examples: `manus/manus-1.6`, `manus/manus-1.6-lite`, `manus/manus-1.6-max` - + Returns: str: The agent profile (e.g., "manus-1.6") """ @@ -79,21 +79,19 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): ) -> dict: """ Validate environment and set up headers for Manus API. - + Manus uses `API_KEY` header instead of `Authorization: Bearer`. """ litellm_params = litellm_params or GenericLiteLLMParams() api_key = ( - litellm_params.api_key - or litellm.api_key - or get_secret_str("MANUS_API_KEY") + litellm_params.api_key or litellm.api_key or get_secret_str("MANUS_API_KEY") ) - + if not api_key: raise ValueError( "Manus API key is required. Set MANUS_API_KEY environment variable or pass api_key parameter." ) - + # Manus uses API_KEY header, not Authorization: Bearer # Content-Type is required for all requests (including GET) headers.update( @@ -111,7 +109,7 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): ) -> str: """ Get the complete URL for Manus Responses API endpoint. - + Returns: str: The full URL for the Manus /v1/responses endpoint """ @@ -121,10 +119,10 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): or get_secret_str("MANUS_API_BASE") or MANUS_API_BASE ) - + # Remove trailing slashes api_base = api_base.rstrip("/") - + # Manus API uses /v1/responses endpoint (OpenAI-compatible) if api_base.endswith("/v1"): return f"{api_base}/responses" @@ -140,7 +138,7 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): ) -> Dict: """ Transform the request for Manus API. - + Manus requires: - `task_mode: "agent"` in the request body - `agent_profile` extracted from model name in the request body @@ -153,24 +151,24 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): litellm_params=litellm_params, headers=headers, ) - + # Extract agent profile from model name agent_profile = self._extract_agent_profile(model=model) - + # Add Manus-specific parameters directly to the request body # These will be sent as part of the request base_request["task_mode"] = "agent" base_request["agent_profile"] = agent_profile - + # Merge any existing extra_body into the request extra_body = response_api_optional_request_params.get("extra_body", {}) or {} if extra_body: base_request.update(extra_body) - + verbose_logger.debug( f"Manus: Using agent_profile={agent_profile}, task_mode=agent" ) - + return base_request def transform_response_api_response( @@ -181,7 +179,7 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): ) -> ResponsesAPIResponse: """ Transform Manus API response to OpenAI-compatible format. - + Manus uses camelCase (createdAt) instead of snake_case (created_at). """ try: @@ -190,13 +188,16 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): additional_args={"complete_input_dict": {}}, ) raw_response_json = raw_response.json() - + # Manus uses camelCase "createdAt" instead of snake_case "created_at" - if "createdAt" in raw_response_json and "created_at" not in raw_response_json: + if ( + "createdAt" in raw_response_json + and "created_at" not in raw_response_json + ): raw_response_json["created_at"] = _safe_convert_created_field( raw_response_json["createdAt"] ) - + # Ensure created_at is set if "created_at" in raw_response_json: raw_response_json["created_at"] = _safe_convert_created_field( @@ -206,20 +207,23 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): raise OpenAIError( message=raw_response.text, status_code=raw_response.status_code ) - + raw_response_headers = dict(raw_response.headers) processed_headers = process_response_headers(raw_response_headers) - + # Ensure reasoning is an empty dict if not present, OpenAI SDK does not allow None - if "reasoning" not in raw_response_json or raw_response_json.get("reasoning") is None: + if ( + "reasoning" not in raw_response_json + or raw_response_json.get("reasoning") is None + ): raw_response_json["reasoning"] = {} - + if "text" not in raw_response_json or raw_response_json.get("text") is None: raw_response_json["text"] = {} - + if "output" not in raw_response_json or raw_response_json.get("output") is None: raw_response_json["output"] = [] - + # Ensure usage is present with default values if not provided if "usage" not in raw_response_json or raw_response_json.get("usage") is None: raw_response_json["usage"] = ResponseAPIUsage( @@ -227,13 +231,13 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): output_tokens=0, total_tokens=0, ) - + # Ensure id is present - failed responses may not include it if "id" not in raw_response_json or raw_response_json.get("id") is None: # Generate a placeholder id for failed responses # This allows the response object to be created even when the API doesn't return an id raw_response_json["id"] = f"unknown-{uuid.uuid4().hex[:8]}" - + try: response = ResponsesAPIResponse(**raw_response_json) except Exception: @@ -241,7 +245,7 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): f"Error constructing ResponsesAPIResponse: {raw_response_json}, using model_construct" ) response = ResponsesAPIResponse.model_construct(**raw_response_json) - + # Store processed headers in additional_headers so they get returned to the client response._hidden_params["additional_headers"] = processed_headers response._hidden_params["headers"] = raw_response_headers @@ -260,10 +264,10 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): ) -> Tuple[str, Dict]: """ Transform the get response API request into a URL and data. - + Manus API follows OpenAI-compatible format: - GET /v1/responses/{response_id} - + Reference: https://open.manus.im/docs/openai-compatibility """ url = f"{api_base}/{response_id}" @@ -277,7 +281,7 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): ) -> ResponsesAPIResponse: """ Transform Manus API GET response to OpenAI-compatible format. - + Manus uses camelCase (createdAt) instead of snake_case (created_at). Same transformation as transform_response_api_response. """ @@ -287,13 +291,16 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): additional_args={"complete_input_dict": {}}, ) raw_response_json = raw_response.json() - + # Manus uses camelCase "createdAt" instead of snake_case "created_at" - if "createdAt" in raw_response_json and "created_at" not in raw_response_json: + if ( + "createdAt" in raw_response_json + and "created_at" not in raw_response_json + ): raw_response_json["created_at"] = _safe_convert_created_field( raw_response_json["createdAt"] ) - + # Ensure created_at is set if "created_at" in raw_response_json: raw_response_json["created_at"] = _safe_convert_created_field( @@ -303,32 +310,35 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): raise OpenAIError( message=raw_response.text, status_code=raw_response.status_code ) - + raw_response_headers = dict(raw_response.headers) processed_headers = process_response_headers(raw_response_headers) - + # Ensure reasoning, text, output, and usage are present with defaults - if "reasoning" not in raw_response_json or raw_response_json.get("reasoning") is None: + if ( + "reasoning" not in raw_response_json + or raw_response_json.get("reasoning") is None + ): raw_response_json["reasoning"] = {} - + if "text" not in raw_response_json or raw_response_json.get("text") is None: raw_response_json["text"] = {} - + if "output" not in raw_response_json or raw_response_json.get("output") is None: raw_response_json["output"] = [] - + if "usage" not in raw_response_json or raw_response_json.get("usage") is None: raw_response_json["usage"] = ResponseAPIUsage( input_tokens=0, output_tokens=0, total_tokens=0, ) - + # Ensure id is present - failed responses may not include it if "id" not in raw_response_json or raw_response_json.get("id") is None: # Generate a placeholder id for failed responses raw_response_json["id"] = f"unknown-{uuid.uuid4().hex[:8]}" - + try: response = ResponsesAPIResponse(**raw_response_json) except Exception: @@ -336,9 +346,8 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): f"Error constructing ResponsesAPIResponse: {raw_response_json}, using model_construct" ) response = ResponsesAPIResponse.model_construct(**raw_response_json) - + # Store processed headers in additional_headers so they get returned to the client response._hidden_params["additional_headers"] = processed_headers response._hidden_params["headers"] = raw_response_headers return response - diff --git a/litellm/llms/minimax/__init__.py b/litellm/llms/minimax/__init__.py index 19093c2dad..e1b0e602e9 100644 --- a/litellm/llms/minimax/__init__.py +++ b/litellm/llms/minimax/__init__.py @@ -11,4 +11,3 @@ __all__ = [ "MinimaxTextToSpeechConfig", "MinimaxException", ] - diff --git a/litellm/llms/minimax/chat/__init__.py b/litellm/llms/minimax/chat/__init__.py index 45bcfd03b4..eeeba74326 100644 --- a/litellm/llms/minimax/chat/__init__.py +++ b/litellm/llms/minimax/chat/__init__.py @@ -1,4 +1,3 @@ """ MiniMax OpenAI-compatible chat API """ - diff --git a/litellm/llms/minimax/chat/transformation.py b/litellm/llms/minimax/chat/transformation.py index 3e9dc0209f..4095e57a8a 100644 --- a/litellm/llms/minimax/chat/transformation.py +++ b/litellm/llms/minimax/chat/transformation.py @@ -15,7 +15,7 @@ class MinimaxChatConfig(OpenAIGPTConfig): MiniMax provides an OpenAI-compatible API at: - International: https://api.minimax.io/v1 - China: https://api.minimaxi.com/v1 - + Supported models: - MiniMax-M2.1 - MiniMax-M2.1-lightning @@ -27,11 +27,7 @@ class MinimaxChatConfig(OpenAIGPTConfig): """ Get MiniMax API key from environment or parameters. """ - return ( - api_key - or get_secret_str("MINIMAX_API_KEY") - or litellm.api_key - ) + return api_key or get_secret_str("MINIMAX_API_KEY") or litellm.api_key @staticmethod def get_api_base( @@ -63,7 +59,7 @@ class MinimaxChatConfig(OpenAIGPTConfig): """ # Get the base URL (either provided or default MiniMax endpoint) base_url = self.get_api_base(api_base=api_base) - + # Ensure it ends with /chat/completions if base_url.endswith("/chat/completions"): return base_url @@ -94,13 +90,12 @@ class MinimaxChatConfig(OpenAIGPTConfig): """ base_params = super().get_supported_openai_params(model=model) additional_params = ["reasoning_split"] - + # Add thinking parameter if model supports reasoning try: if litellm.supports_reasoning(model=model, custom_llm_provider="minimax"): additional_params.append("thinking") except Exception: pass - - return base_params + additional_params + return base_params + additional_params diff --git a/litellm/llms/minimax/messages/transformation.py b/litellm/llms/minimax/messages/transformation.py index 27d28f02d8..13ed6ad386 100644 --- a/litellm/llms/minimax/messages/transformation.py +++ b/litellm/llms/minimax/messages/transformation.py @@ -16,7 +16,7 @@ class MinimaxMessagesConfig(AnthropicMessagesConfig): MiniMax provides an Anthropic-compatible API at: - International: https://api.minimax.io/anthropic - China: https://api.minimaxi.com/anthropic - + Supported models: - MiniMax-M2.1 - MiniMax-M2.1-lightning @@ -32,11 +32,7 @@ class MinimaxMessagesConfig(AnthropicMessagesConfig): """ Get MiniMax API key from environment or parameters. """ - return ( - api_key - or get_secret_str("MINIMAX_API_KEY") - or litellm.api_key - ) + return api_key or get_secret_str("MINIMAX_API_KEY") or litellm.api_key @staticmethod def get_api_base( @@ -68,14 +64,13 @@ class MinimaxMessagesConfig(AnthropicMessagesConfig): """ # Get the base URL (either provided or default MiniMax endpoint) base_url = self.get_api_base(api_base=api_base) - + # If the base URL already includes the full path, return it if base_url.endswith("/v1/messages"): return base_url - + # Otherwise append the messages endpoint if base_url.endswith("/"): return f"{base_url}v1/messages" else: return f"{base_url}/v1/messages" - diff --git a/litellm/llms/minimax/text_to_speech/__init__.py b/litellm/llms/minimax/text_to_speech/__init__.py index e3fcddeb05..bf4ac9010a 100644 --- a/litellm/llms/minimax/text_to_speech/__init__.py +++ b/litellm/llms/minimax/text_to_speech/__init__.py @@ -5,4 +5,3 @@ MiniMax Text-to-Speech module from .transformation import MinimaxException, MinimaxTextToSpeechConfig __all__ = ["MinimaxTextToSpeechConfig", "MinimaxException"] - diff --git a/litellm/llms/minimax/text_to_speech/transformation.py b/litellm/llms/minimax/text_to_speech/transformation.py index a3a75d220f..2a7d6897ed 100644 --- a/litellm/llms/minimax/text_to_speech/transformation.py +++ b/litellm/llms/minimax/text_to_speech/transformation.py @@ -43,7 +43,7 @@ class MinimaxTextToSpeechConfig(BaseTextToSpeechConfig): Configuration for MiniMax Text-to-Speech Reference: https://platform.minimax.io/docs - + MiniMax TTS API supports both WebSocket and HTTP endpoints. This implementation uses the HTTP endpoint for simplicity. """ @@ -186,11 +186,7 @@ class MinimaxTextToSpeechConfig(BaseTextToSpeechConfig): """ Validate MiniMax environment and set up authentication headers """ - api_key = ( - api_key - or litellm.api_key - or get_secret_str("MINIMAX_API_KEY") - ) + api_key = api_key or litellm.api_key or get_secret_str("MINIMAX_API_KEY") if api_key is None: raise ValueError( @@ -224,7 +220,7 @@ class MinimaxTextToSpeechConfig(BaseTextToSpeechConfig): ) -> TextToSpeechRequestData: """ Build the MiniMax TTS request payload. - + MiniMax uses a different structure than OpenAI: - model: The TTS model to use - text: The input text @@ -237,16 +233,18 @@ class MinimaxTextToSpeechConfig(BaseTextToSpeechConfig): voice_id = params.pop("voice_id", voice or "male-qn-qingse") speed = params.pop("speed", 1.0) audio_format = params.pop("format", "mp3") - + # Extract additional voice settings vol = params.pop("vol", 1.0) # Volume (0.1 to 10) pitch = params.pop("pitch", 0) # Pitch adjustment (-12 to 12) - + # Extract audio settings sample_rate = params.pop("sample_rate", 32000) # 16000, 24000, 32000 - bitrate = params.pop("bitrate", 128000) # For MP3: 64000, 128000, 192000, 256000 + bitrate = params.pop( + "bitrate", 128000 + ) # For MP3: 64000, 128000, 192000, 256000 channel = params.pop("channel", 1) # 1 for mono, 2 for stereo - + # Output format: 'url' or 'hex' (default is 'hex') output_format = params.pop("output_format", "hex") @@ -289,14 +287,14 @@ class MinimaxTextToSpeechConfig(BaseTextToSpeechConfig): ) -> "HttpxBinaryResponseContent": """ Transform MiniMax response to standard format. - + MiniMax returns JSON with base64-encoded audio data: { "base_resp": {"status_code": 0, "status_msg": "success"}, "audio_file": "", "extra_info": {...} } - + We need to decode the base64 audio and return it as binary content. """ import base64 @@ -307,12 +305,12 @@ class MinimaxTextToSpeechConfig(BaseTextToSpeechConfig): try: # Parse JSON response response_json = raw_response.json() - + # MiniMax API response format check # The API can return different structures: # 1. {"data": {"audio": "..."}, "status": 0, ...} for HTTP endpoint # 2. {"base_resp": {"status_code": 0, ...}, "audio_file": "..."} for older versions - + # Check for errors - MiniMax uses "status" field in HTTP endpoint response # status: 0 = success, 2 = invalid api key, etc. status = response_json.get("status") @@ -324,11 +322,11 @@ class MinimaxTextToSpeechConfig(BaseTextToSpeechConfig): message=f"MiniMax TTS error: {error_detail}", headers=dict(raw_response.headers), ) - + # Extract audio data # MiniMax returns audio in "data" field data = response_json.get("data", {}) - + # Check if response contains a URL (output_format='url') audio_url = data.get("audio_url", None) if audio_url: @@ -339,17 +337,17 @@ class MinimaxTextToSpeechConfig(BaseTextToSpeechConfig): message=f"URL output format is not yet supported. Use 'hex' format or fetch from URL: {audio_url}", headers=dict(raw_response.headers), ) - + # Get hex-encoded audio data audio_hex = data.get("audio", "") or response_json.get("audio_file", "") - + if not audio_hex: raise MinimaxException( status_code=500, message=f"No audio data in MiniMax response. Response keys: {list(response_json.keys())}", headers=dict(raw_response.headers), ) - + # MiniMax returns hex-encoded audio by default # Try hex decoding first, fall back to base64 if that fails try: @@ -364,15 +362,15 @@ class MinimaxTextToSpeechConfig(BaseTextToSpeechConfig): message=f"Failed to decode audio data: {str(e)}", headers=dict(raw_response.headers), ) - + # Create a new response with binary audio content # We need to create a response that contains the decoded audio bytes # Remove gzip encoding headers to avoid decompression issues clean_headers = dict(raw_response.headers) - clean_headers.pop('content-encoding', None) - clean_headers.pop('transfer-encoding', None) - clean_headers['content-length'] = str(len(audio_bytes)) - + clean_headers.pop("content-encoding", None) + clean_headers.pop("transfer-encoding", None) + clean_headers["content-length"] = str(len(audio_bytes)) + # Create a new response object with the binary content binary_response = httpx.Response( status_code=200, @@ -380,9 +378,9 @@ class MinimaxTextToSpeechConfig(BaseTextToSpeechConfig): content=audio_bytes, request=raw_response.request, ) - + return HttpxBinaryResponseContent(binary_response) - + except json.JSONDecodeError as e: raise MinimaxException( status_code=500, @@ -407,15 +405,10 @@ class MinimaxTextToSpeechConfig(BaseTextToSpeechConfig): """ Construct the MiniMax endpoint URL. """ - base_url = ( - api_base - or get_secret_str("MINIMAX_API_BASE") - or self.TTS_BASE_URL - ) + base_url = api_base or get_secret_str("MINIMAX_API_BASE") or self.TTS_BASE_URL base_url = base_url.rstrip("/") # MiniMax uses a simple endpoint path url = f"{base_url}{self.TTS_ENDPOINT_PATH}" return url - diff --git a/litellm/llms/mistral/audio_transcription/transformation.py b/litellm/llms/mistral/audio_transcription/transformation.py index fd84d63c4f..4d29406349 100644 --- a/litellm/llms/mistral/audio_transcription/transformation.py +++ b/litellm/llms/mistral/audio_transcription/transformation.py @@ -60,9 +60,7 @@ class MistralAudioTranscriptionConfig(BaseAudioTranscriptionConfig): stream: Optional[bool] = None, ) -> str: api_base = ( - "https://api.mistral.ai/v1" - if api_base is None - else api_base.rstrip("/") + "https://api.mistral.ai/v1" if api_base is None else api_base.rstrip("/") ) return f"{api_base}/audio/transcriptions" @@ -121,7 +119,9 @@ class MistralAudioTranscriptionConfig(BaseAudioTranscriptionConfig): openai_params=self.get_supported_openai_params(model), ) for key, value in provider_specific_params.items(): - form_fields[key] = str(value).lower() if isinstance(value, bool) else str(value) + form_fields[key] = ( + str(value).lower() if isinstance(value, bool) else str(value) + ) files = { "file": ( diff --git a/litellm/llms/mistral/chat/transformation.py b/litellm/llms/mistral/chat/transformation.py index 2673862337..23fbe467fc 100644 --- a/litellm/llms/mistral/chat/transformation.py +++ b/litellm/llms/mistral/chat/transformation.py @@ -244,7 +244,7 @@ class MistralConfig(OpenAIGPTConfig): - if `name` is passed, then drop it for mistral API: https://github.com/BerriAI/litellm/issues/6696 Motivation: mistral api doesn't support content as a list. - The above statement is not valid now. Need to plan to remove all the #1,2,3 + The above statement is not valid now. Need to plan to remove all the #1,2,3 Mistral API supports content as a list. """ ## 1. If 'image_url' or 'file' in content, then transform with base class and mistral-specific handling @@ -276,8 +276,8 @@ class MistralConfig(OpenAIGPTConfig): else: return super()._transform_messages(new_messages, model, False) - async def _transform_messages_async(self, - messages: List[AllMessageValues], model: str + async def _transform_messages_async( + self, messages: List[AllMessageValues], model: str ) -> List[AllMessageValues]: """ Handle modification of messages for Mistral API in an async context. @@ -288,11 +288,10 @@ class MistralConfig(OpenAIGPTConfig): messages = self._handle_message_with_file(messages) return messages - def _transform_messages_sync(self, - messages: List[AllMessageValues], model: str + def _transform_messages_sync( + self, messages: List[AllMessageValues], model: str ) -> List[AllMessageValues]: - """ Handle modification of messages for Mistral API in a sync context. - """ + """Handle modification of messages for Mistral API in a sync context.""" # Call parent sync method to handle basic transformations # and then apply Mistral-specific handling for files # This is the sync version of the async method above @@ -301,23 +300,25 @@ class MistralConfig(OpenAIGPTConfig): return messages def _handle_message_with_file( - self, - messages: List[AllMessageValues]) -> List[AllMessageValues]: + self, messages: List[AllMessageValues] + ) -> List[AllMessageValues]: """ Mistral API supports only 'file_id' in message content with type 'file'. """ for m in messages: _content_block = m.get("content") - if _content_block and isinstance(_content_block, list): + if _content_block and isinstance(_content_block, list): if any(c.get("type") == "file" for c in _content_block): # If file content is present, we get file_id from 'file' attribute of content block # then replace 'file' with 'file_id' and assign the value of 'file_id' attribute to it. - file_contents = [c for c in _content_block if c.get("type") == "file"] + file_contents = [ + c for c in _content_block if c.get("type") == "file" + ] for file_content in file_contents: file_id = file_content.get("file", {}).get("file_id") if file_id: # Replace 'file' with 'file_id' - file_content["file_id"] = file_id # type: ignore + file_content["file_id"] = file_id # type: ignore file_content.pop("file", None) return messages @@ -343,9 +344,9 @@ class MistralConfig(OpenAIGPTConfig): # Handle both string and list content, preserving original format if isinstance(existing_content, str): # String content - prepend reasoning prompt - new_content: Union[str, list] = ( - f"{reasoning_prompt}\n\n{existing_content}" - ) + new_content: Union[ + str, list + ] = f"{reasoning_prompt}\n\n{existing_content}" elif isinstance(existing_content, list): # List content - prepend reasoning prompt as text block new_content = [ @@ -679,5 +680,7 @@ class MistralChatResponseIterator(OpenAIChatCompletionStreamingHandler): text_segments.append(block.get("text", "")) normalized_text = "".join(text_segments) if text_segments else None - reasoning_content = "\n".join(reasoning_segments) if reasoning_segments else None + reasoning_content = ( + "\n".join(reasoning_segments) if reasoning_segments else None + ) return normalized_text, thinking_blocks, reasoning_content diff --git a/litellm/llms/mistral/embedding.py b/litellm/llms/mistral/embedding.py index 0aae35ad7f..4861674a19 100644 --- a/litellm/llms/mistral/embedding.py +++ b/litellm/llms/mistral/embedding.py @@ -1,4 +1,4 @@ """ Calls handled in openai/ as mistral is an openai-compatible endpoint. -""" \ No newline at end of file +""" diff --git a/litellm/llms/mistral/ocr/__init__.py b/litellm/llms/mistral/ocr/__init__.py index 40cc62696b..54eed416c1 100644 --- a/litellm/llms/mistral/ocr/__init__.py +++ b/litellm/llms/mistral/ocr/__init__.py @@ -1,2 +1 @@ """Mistral OCR transformation module.""" - diff --git a/litellm/llms/mistral/ocr/guardrail_translation/handler.py b/litellm/llms/mistral/ocr/guardrail_translation/handler.py index 87d79a3ce6..697bd2daa3 100644 --- a/litellm/llms/mistral/ocr/guardrail_translation/handler.py +++ b/litellm/llms/mistral/ocr/guardrail_translation/handler.py @@ -108,9 +108,7 @@ class OCRHandler(BaseTranslation): Modified OCRResponse with guardrailed page text """ if not hasattr(response, "pages") or not response.pages: - verbose_proxy_logger.debug( - "OCR guardrail: No pages found in OCR response" - ) + verbose_proxy_logger.debug("OCR guardrail: No pages found in OCR response") return response # Extract markdown text from all pages diff --git a/litellm/llms/mistral/ocr/transformation.py b/litellm/llms/mistral/ocr/transformation.py index ed5e235939..11848f8acf 100644 --- a/litellm/llms/mistral/ocr/transformation.py +++ b/litellm/llms/mistral/ocr/transformation.py @@ -18,7 +18,7 @@ from litellm.secret_managers.main import get_secret_str class MistralOCRConfig(BaseOCRConfig): """ Mistral OCR transformation configuration. - + Reference: https://docs.mistral.ai/api/#tag/ocr """ @@ -28,7 +28,7 @@ class MistralOCRConfig(BaseOCRConfig): def get_supported_ocr_params(self, model: str) -> list: """ Get supported OCR parameters for Mistral OCR. - + Mistral OCR supports: - pages: List of page numbers to process - include_image_base64: Whether to include base64 encoded images @@ -45,7 +45,7 @@ class MistralOCRConfig(BaseOCRConfig): "bbox_annotation_format", "document_annotation_format", ] - + def map_ocr_params( self, non_default_params: dict, @@ -54,18 +54,18 @@ class MistralOCRConfig(BaseOCRConfig): ) -> dict: """ Map OCR parameters to Mistral-specific format. - + Mistral accepts these parameters directly, so no transformation needed. Just filter out unsupported params. """ supported_params = self.get_supported_ocr_params(model=model) - + # Only include params that are in the supported list mapped_params = {} for param, value in non_default_params.items(): if param in supported_params: mapped_params[param] = value - + return mapped_params def validate_environment( @@ -82,9 +82,7 @@ class MistralOCRConfig(BaseOCRConfig): """ # Get API key from environment if not provided if api_key is None: - api_key = ( - get_secret_str("MISTRAL_API_KEY") - ) + api_key = get_secret_str("MISTRAL_API_KEY") if api_key is None: raise ValueError( @@ -95,7 +93,7 @@ class MistralOCRConfig(BaseOCRConfig): "Authorization": f"Bearer {api_key}", **headers, } - + # Don't set Content-Type for multipart/form-data - httpx will handle it return headers @@ -110,7 +108,7 @@ class MistralOCRConfig(BaseOCRConfig): ) -> str: """ Get complete URL for Mistral OCR endpoint. - + Returns: https://api.mistral.ai/v1/ocr """ if api_base is None: @@ -118,14 +116,13 @@ class MistralOCRConfig(BaseOCRConfig): # Ensure no trailing slash api_base = api_base.rstrip("/") - + # Remove /v1 if it's already in the base to avoid duplication if api_base.endswith("/v1"): return f"{api_base}/ocr" return f"{api_base}/v1/ocr" - def transform_ocr_request( self, model: str, @@ -136,7 +133,7 @@ class MistralOCRConfig(BaseOCRConfig): ) -> OCRRequestData: """ Transform OCR request to Mistral-specific format. - + Mistral OCR API accepts: { "model": "mistral-ocr-latest", @@ -148,32 +145,32 @@ class MistralOCRConfig(BaseOCRConfig): "include_image_base64": false, # optional ... } - + Args: model: Model name (e.g., "mistral-ocr-latest") document: Document dict from user (Mistral format) - already validated in main.py optional_params: Already mapped optional parameters headers: Request headers - + Returns: OCRRequestData with JSON data """ verbose_logger.debug(f"Mistral OCR transform_ocr_request - model: {model}") - + # Document parameter is the Mistral-format dict from the user # Just pass it through as-is to the Mistral API if not isinstance(document, dict): raise ValueError(f"Expected document dict, got {type(document)}") - + # Build request data - use document dict directly data = { "model": model, "document": document, # Pass through the Mistral-format document dict } - + # Add all optional parameters from the already-mapped optional_params data.update(optional_params) - + # No multipart files - using JSON return OCRRequestData(data=data, files=None) @@ -186,10 +183,10 @@ class MistralOCRConfig(BaseOCRConfig): ) -> OCRResponse: """ Return Mistral OCR response in native format. - + Mistral OCR is the standard format for LiteLLM OCR responses. No transformation needed - return native response. - + Mistral OCR returns: { "pages": [ @@ -208,9 +205,9 @@ class MistralOCRConfig(BaseOCRConfig): """ try: response_json = raw_response.json() - + verbose_logger.debug(f"Mistral OCR response keys: {response_json.keys()}") - + # Return native Mistral format - no transformation return OCRResponse( pages=response_json.get("pages", []), @@ -222,4 +219,3 @@ class MistralOCRConfig(BaseOCRConfig): except Exception as e: verbose_logger.error(f"Error parsing Mistral OCR response: {e}") raise e - diff --git a/litellm/llms/moonshot/chat/transformation.py b/litellm/llms/moonshot/chat/transformation.py index 72c51bf74f..3ed08f51c8 100644 --- a/litellm/llms/moonshot/chat/transformation.py +++ b/litellm/llms/moonshot/chat/transformation.py @@ -95,24 +95,24 @@ class MoonshotChatConfig(OpenAIGPTConfig): def get_supported_openai_params(self, model: str) -> list: """ Get the supported OpenAI params for Moonshot AI models - + Moonshot AI limitations: - functions parameter is not supported (use tools instead) - tool_choice doesn't support "required" value - kimi-thinking-preview doesn't support tool calls at all """ excluded_params: List[str] = ["functions"] - + # kimi-thinking-preview has additional limitations if "kimi-thinking-preview" in model: excluded_params.extend(["tools", "tool_choice"]) - + base_openai_params = super().get_supported_openai_params(model=model) final_params: List[str] = [] for param in base_openai_params: if param not in excluded_params: final_params.append(param) - + return final_params def map_openai_params( @@ -124,7 +124,7 @@ class MoonshotChatConfig(OpenAIGPTConfig): ) -> dict: """ Map OpenAI parameters to Moonshot AI parameters - + Handles Moonshot AI specific limitations: - tool_choice doesn't support "required" value - Temperature <0.3 limitation for n>1 @@ -139,7 +139,7 @@ class MoonshotChatConfig(OpenAIGPTConfig): ########################################## # temperature limitations # 1. `temperature` on KIMI API is [0, 1] but OpenAI is [0, 2] - # 2. If temperature < 0.3 and n > 1, KIMI will raise an exception. + # 2. If temperature < 0.3 and n > 1, KIMI will raise an exception. # If we enter this condition, we set the temperature to 0.3 as suggested by Moonshot AI ########################################## if "temperature" in optional_params: @@ -148,7 +148,6 @@ class MoonshotChatConfig(OpenAIGPTConfig): if optional_params["temperature"] < 0.3 and optional_params.get("n", 1) > 1: optional_params["temperature"] = 0.3 return optional_params - def transform_request( self, @@ -178,17 +177,20 @@ class MoonshotChatConfig(OpenAIGPTConfig): litellm_params=litellm_params, headers=headers, ) - - def _add_tool_choice_required_message(self, messages: List[AllMessageValues], optional_params: dict) -> List[AllMessageValues]: + def _add_tool_choice_required_message( + self, messages: List[AllMessageValues], optional_params: dict + ) -> List[AllMessageValues]: """ Add a message to the messages list to indicate that the tool choice is required. https://platform.moonshot.ai/docs/guide/migrating-from-openai-to-kimi#about-tool_choice """ - messages.append({ - "role": "user", - "content": "Please select a tool to handle the current issue.", # Usually, the Kimi large language model understands the intention to invoke a tool and selects one for invocation - }) + messages.append( + { + "role": "user", + "content": "Please select a tool to handle the current issue.", # Usually, the Kimi large language model understands the intention to invoke a tool and selects one for invocation + } + ) optional_params.pop("tool_choice") return messages diff --git a/litellm/llms/nvidia_nim/rerank/common_utils.py b/litellm/llms/nvidia_nim/rerank/common_utils.py index 2bd8c123c9..738fb09364 100644 --- a/litellm/llms/nvidia_nim/rerank/common_utils.py +++ b/litellm/llms/nvidia_nim/rerank/common_utils.py @@ -6,13 +6,13 @@ Common utilities for NVIDIA NIM rerank provider. def get_nvidia_nim_rerank_config(model: str): """ Get the appropriate NVIDIA NIM rerank config based on the model. - + Args: model: The model string (e.g., "nvidia/llama-3.2-nv-rerankqa-1b-v2" or "ranking/nvidia/llama-3.2-nv-rerankqa-1b-v2") - + Returns: NvidiaNimRankingConfig if model starts with "ranking/", else NvidiaNimRerankConfig - + Example: - "ranking/nvidia/llama-3.2-nv-rerankqa-1b-v2" -> NvidiaNimRankingConfig - "nvidia/llama-3.2-nv-rerankqa-1b-v2" -> NvidiaNimRerankConfig @@ -25,4 +25,3 @@ def get_nvidia_nim_rerank_config(model: str): if model.startswith("ranking/"): return NvidiaNimRankingConfig() return NvidiaNimRerankConfig() - diff --git a/litellm/llms/nvidia_nim/rerank/ranking_transformation.py b/litellm/llms/nvidia_nim/rerank/ranking_transformation.py index d97c47bcb2..757d874bf3 100644 --- a/litellm/llms/nvidia_nim/rerank/ranking_transformation.py +++ b/litellm/llms/nvidia_nim/rerank/ranking_transformation.py @@ -31,10 +31,10 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig): """Strip 'nvidia_nim/' and 'ranking/' prefixes from model name.""" # First strip nvidia_nim/ prefix if present if model.startswith("nvidia_nim/"): - model = model[len("nvidia_nim/"):] + model = model[len("nvidia_nim/") :] # Then strip ranking/ prefix if present if model.startswith("ranking/"): - model = model[len("ranking/"):] + model = model[len("ranking/") :] return model def get_complete_url( @@ -45,7 +45,7 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig): ) -> str: """ Construct the Nvidia NIM ranking URL. - + Format: {api_base}/v1/ranking """ if not api_base: @@ -76,4 +76,3 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig): optional_rerank_params=optional_rerank_params, headers=headers, ) - diff --git a/litellm/llms/nvidia_nim/rerank/transformation.py b/litellm/llms/nvidia_nim/rerank/transformation.py index c7b1b249da..bd5abac60c 100644 --- a/litellm/llms/nvidia_nim/rerank/transformation.py +++ b/litellm/llms/nvidia_nim/rerank/transformation.py @@ -45,11 +45,12 @@ class NvidiaNimRerankResponse(TypedDict): class NvidiaNimRerankConfig(BaseRerankConfig): """ Reference: https://docs.api.nvidia.com/nim/reference/nvidia-llama-3_2-nv-rerankqa-1b-v2-infer - + Nvidia NIM rerank API uses a different format: - query is an object with 'text' field - documents are called 'passages' and have 'text' field """ + DEFAULT_NIM_RERANK_API_BASE = "https://ai.api.nvidia.com" def __init__(self) -> None: @@ -58,39 +59,39 @@ class NvidiaNimRerankConfig(BaseRerankConfig): def _get_clean_model_name(self, model: str) -> str: """Strip 'nvidia_nim/' prefix from model name if present.""" if model.startswith("nvidia_nim/"): - return model[len("nvidia_nim/"):] + return model[len("nvidia_nim/") :] return model def get_complete_url( - self, - api_base: Optional[str], + self, + api_base: Optional[str], model: str, optional_params: Optional[dict] = None, ) -> str: """ Construct the Nvidia NIM rerank URL. - + Format: {api_base}/v1/retrieval/{model}/reranking - + If the user provides a full URL (e.g., {api_base}/v1/retrieval/{model}/reranking), it will be used as-is. """ if not api_base: api_base = self.DEFAULT_NIM_RERANK_API_BASE - + api_base = api_base.rstrip("/") - + # Check if user already provided the full URL with /retrieval/ path if "/retrieval/" in api_base: return api_base - + # Ensure we don't have duplicate /v1 if api_base.endswith("/v1"): api_base = api_base[:-3] - + # Strip nvidia_nim/ prefix from model name if present clean_model = self._get_clean_model_name(model) - + return f"{api_base}/v1/retrieval/{clean_model}/reranking" def get_supported_cohere_rerank_params(self, model: str) -> list: @@ -119,10 +120,10 @@ class NvidiaNimRerankConfig(BaseRerankConfig): ) -> Dict: """ Map Cohere/OpenAI rerank params to Nvidia NIM format. - + Parameter mapping: - top_n (Cohere) -> top_k (Nvidia) - + Nvidia NIM specific params (passed through as-is from non_default_params): - truncate: How to truncate input if too long (NONE, END) """ @@ -130,11 +131,11 @@ class NvidiaNimRerankConfig(BaseRerankConfig): "query": query, "documents": documents, } - + # Map Cohere's top_n to Nvidia's top_k if top_n is not None: optional_nvidia_nim_rerank_params["top_k"] = top_n - + # Pass through Nvidia-specific params from non_default_params if non_default_params: optional_nvidia_nim_rerank_params.update(non_default_params) @@ -151,10 +152,7 @@ class NvidiaNimRerankConfig(BaseRerankConfig): Validate that the Nvidia NIM API key is present. """ if api_key is None: - api_key = ( - get_secret_str("NVIDIA_NIM_API_KEY") - or litellm.api_key - ) + api_key = get_secret_str("NVIDIA_NIM_API_KEY") or litellm.api_key if api_key is None: raise ValueError( @@ -182,12 +180,12 @@ class NvidiaNimRerankConfig(BaseRerankConfig): ) -> dict: """ Transform request to Nvidia NIM format. - + Nvidia NIM expects: - query as {text: "..."} - documents as passages: [{text: "..."}, ...] - Optional: truncate (NONE or END), top_k - + Note: optional_rerank_params may contain provider-specific params like 'top_k' and 'truncate' that aren't in the OptionalRerankParams TypedDict but are passed through at runtime. The mapping from Cohere's 'top_n' to Nvidia's 'top_k' already happened in map_cohere_rerank_params. @@ -199,10 +197,10 @@ class NvidiaNimRerankConfig(BaseRerankConfig): query = optional_rerank_params["query"] documents = optional_rerank_params["documents"] - + # Transform query to object format query_obj: NvidiaNimQueryObject = {"text": query} - + # Transform documents to passages format passages: List[NvidiaNimPassageObject] = [] for doc in documents: @@ -215,35 +213,36 @@ class NvidiaNimRerankConfig(BaseRerankConfig): else: # Otherwise, stringify the dict import json + passages.append({"text": json.dumps(doc)}) else: passages.append({"text": str(doc)}) - + # Strip nvidia_nim/ prefix from model name if present clean_model = self._get_clean_model_name(model) - + # Note: URL path uses underscores (llama-3_2) but JSON body uses periods (llama-3.2) # Convert underscores back to periods for the model field in request body model_for_body = clean_model.replace("_", ".") - + # Build request using TypedDict request_data: NvidiaNimRerankRequest = { "model": model_for_body, "query": query_obj, "passages": passages, } - + # Add optional top_k parameter if provided (already mapped from top_n in map_cohere_rerank_params) if "top_k" in optional_rerank_params and optional_rerank_params.get("top_k") is not None: # type: ignore request_data["top_k"] = optional_rerank_params.get("top_k") # type: ignore - + # Add Nvidia-specific truncate parameter if provided # This is passed through from non_default_params, not in base OptionalRerankParams if "truncate" in optional_rerank_params and optional_rerank_params.get("truncate") is not None: # type: ignore truncate_value = optional_rerank_params.get("truncate") # type: ignore if truncate_value in ["NONE", "END"]: request_data["truncate"] = truncate_value # type: ignore - + return dict(request_data) def transform_rerank_response( @@ -259,7 +258,7 @@ class NvidiaNimRerankConfig(BaseRerankConfig): ) -> RerankResponse: """ Transform Nvidia NIM rerank response to LiteLLM format. - + Nvidia NIM returns (NvidiaNimRerankResponse): { "rankings": [ @@ -269,7 +268,7 @@ class NvidiaNimRerankConfig(BaseRerankConfig): } ] } - + LiteLLM expects (RerankResponse): { "results": [ @@ -292,40 +291,40 @@ class NvidiaNimRerankConfig(BaseRerankConfig): # Parse as NvidiaNimRerankResponse nvidia_response: NvidiaNimRerankResponse = raw_response_json - + # Transform Nvidia NIM response to LiteLLM format results: List[RerankResponseResult] = [] rankings = nvidia_response.get("rankings", []) - + # Get original documents from request if we need to include them - original_passages: List[NvidiaNimPassageObject] = request_data.get("passages", []) - + original_passages: List[NvidiaNimPassageObject] = request_data.get( + "passages", [] + ) + for ranking in rankings: result_item: RerankResponseResult = { "index": ranking["index"], "relevance_score": ranking["logit"], } - + # Include document if it was in the original request index: int = ranking["index"] if index < len(original_passages): result_item["document"] = {"text": original_passages[index]["text"]} # type: ignore - + results.append(result_item) - + # Construct metadata with billed_units # Nvidia NIM uses "usage" field with "total_tokens" usage = raw_response_json.get("usage", {}) total_tokens = usage.get("total_tokens", 0) - + billed_units: RerankBilledUnits = { "total_tokens": total_tokens if total_tokens > 0 else len(results) } - - meta: RerankResponseMeta = { - "billed_units": billed_units - } - + + meta: RerankResponseMeta = {"billed_units": billed_units} + return RerankResponse( id=raw_response_json.get("id") or str(uuid.uuid4()), results=results, @@ -340,4 +339,3 @@ class NvidiaNimRerankConfig(BaseRerankConfig): message=error_message, headers=headers, ) - diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index 1c22602b48..b1af7ed2ec 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -3,7 +3,17 @@ import datetime import hashlib import json from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, List, Optional, Protocol, Tuple, Union +from typing import ( + TYPE_CHECKING, + Any, + AsyncIterator, + Dict, + List, + Optional, + Protocol, + Tuple, + Union, +) from urllib.parse import urlparse import httpx @@ -74,7 +84,9 @@ class OCISignerProtocol(Protocol): See: https://docs.oracle.com/en-us/iaas/tools/python/latest/api/signing.html """ - def do_request_sign(self, request: Any, *, enforce_content_headers: bool = False) -> None: + def do_request_sign( + self, request: Any, *, enforce_content_headers: bool = False + ) -> None: """ Sign an HTTP request by adding authentication headers. @@ -93,6 +105,7 @@ class OCIRequestWrapper: This class wraps request data in a format compatible with OCI SDK signers, which expect objects with method, url, headers, body, and path_url attributes. """ + method: str url: str headers: dict @@ -222,7 +235,9 @@ class OCIChatConfig(BaseConfig): } # Cohere and Gemini use the same parameter mapping as GENERIC - self.openai_to_oci_cohere_param_map = self.openai_to_oci_generic_param_map.copy() + self.openai_to_oci_cohere_param_map = ( + self.openai_to_oci_generic_param_map.copy() + ) def get_supported_openai_params(self, model: str) -> List[str]: supported_params = [] @@ -310,14 +325,13 @@ class OCIChatConfig(BaseConfig): prepared_headers.setdefault("content-length", str(len(body))) request_wrapper = OCIRequestWrapper( - method=method, - url=api_base, - headers=prepared_headers, - body=body + method=method, url=api_base, headers=prepared_headers, body=body ) if oci_signer is None: - raise ValueError("oci_signer cannot be None when calling _sign_with_oci_signer") + raise ValueError( + "oci_signer cannot be None when calling _sign_with_oci_signer" + ) try: oci_signer.do_request_sign(request_wrapper, enforce_content_headers=True) @@ -329,7 +343,7 @@ class OCIChatConfig(BaseConfig): "The signer must implement the OCI SDK Signer interface with a " "do_request_sign(request, enforce_content_headers=True) method. " "See: https://docs.oracle.com/en-us/iaas/tools/python/latest/api/signing.html" - ) + ), ) from e headers.update(request_wrapper.headers) @@ -442,7 +456,9 @@ class OCIChatConfig(BaseConfig): private_key = ( load_private_key_from_str(oci_key_content) if oci_key_content - else load_private_key_from_file(oci_key_file) if oci_key_file else None + else load_private_key_from_file(oci_key_file) + if oci_key_file + else None ) if private_key is None: @@ -539,10 +555,14 @@ class OCIChatConfig(BaseConfig): # If a signer is provided, use it for request signing if oci_signer is not None: - return self._sign_with_oci_signer(headers, optional_params, request_data, api_base) + return self._sign_with_oci_signer( + headers, optional_params, request_data, api_base + ) # Standard manual credential signing - return self._sign_with_manual_credentials(headers, optional_params, request_data, api_base) + return self._sign_with_manual_credentials( + headers, optional_params, request_data, api_base + ) def validate_environment( self, @@ -653,7 +673,7 @@ class OCIChatConfig(BaseConfig): "temperature": 1, "topK": 0, "topP": 0.75, - "frequencyPenalty": 0 + "frequencyPenalty": 0, } else: open_ai_to_oci_param_map = self.openai_to_oci_generic_param_map @@ -665,7 +685,11 @@ class OCIChatConfig(BaseConfig): # Also check for already-mapped OCI params (for backward compatibility) for oci_value in open_ai_to_oci_param_map.values(): - if oci_value and oci_value in optional_params and oci_value not in selected_params: + if ( + oci_value + and oci_value in optional_params + and oci_value not in selected_params + ): selected_params[oci_value] = optional_params[oci_value] # type: ignore[index] if "tools" in selected_params: @@ -709,7 +733,9 @@ class OCIChatConfig(BaseConfig): return selected_params - def adapt_messages_to_cohere_standard(self, messages: List[AllMessageValues]) -> List[CohereMessage]: + def adapt_messages_to_cohere_standard( + self, messages: List[AllMessageValues] + ) -> List[CohereMessage]: """Build chat history for Cohere models.""" chat_history = [] for msg in messages[:-1]: # All messages except the last one @@ -720,7 +746,10 @@ class OCIChatConfig(BaseConfig): # Extract text from content array text_content = "" for content_item in content: - if isinstance(content_item, dict) and content_item.get("type") == "text": + if ( + isinstance(content_item, dict) + and content_item.get("type") == "text" + ): text_content += content_item.get("text", "") content = text_content @@ -734,7 +763,9 @@ class OCIChatConfig(BaseConfig): tool_calls = [] for tool_call in msg["tool_calls"]: # type: ignore[union-attr,typeddict-item] # Parse arguments if they're a JSON string - raw_arguments: Any = tool_call.get("function", {}).get("arguments", {}) + raw_arguments: Any = tool_call.get("function", {}).get( + "arguments", {} + ) if isinstance(raw_arguments, str): try: arguments: Dict[str, Any] = json.loads(raw_arguments) @@ -743,26 +774,34 @@ class OCIChatConfig(BaseConfig): else: arguments = raw_arguments - tool_calls.append(CohereToolCall( - name=str(tool_call.get("function", {}).get("name", "")), - parameters=arguments - )) + tool_calls.append( + CohereToolCall( + name=str(tool_call.get("function", {}).get("name", "")), + parameters=arguments, + ) + ) if role == "user": chat_history.append(CohereMessage(role="USER", message=content)) elif role == "assistant": - chat_history.append(CohereMessage(role="CHATBOT", message=content, toolCalls=tool_calls)) + chat_history.append( + CohereMessage(role="CHATBOT", message=content, toolCalls=tool_calls) + ) elif role == "tool": # Tool messages need special handling - chat_history.append(CohereMessage( - role="TOOL", - message=content, - toolCalls=None # Tool messages don't have tool calls - )) + chat_history.append( + CohereMessage( + role="TOOL", + message=content, + toolCalls=None, # Tool messages don't have tool calls + ) + ) return chat_history - def adapt_tool_definitions_to_cohere_standard(self, tools: List[Dict[str, Any]]) -> List[CohereTool]: + def adapt_tool_definitions_to_cohere_standard( + self, tools: List[Dict[str, Any]] + ) -> List[CohereTool]: """Adapt tool definitions to Cohere format.""" cohere_tools = [] for tool in tools: @@ -775,14 +814,16 @@ class OCIChatConfig(BaseConfig): parameter_definitions[param_name] = CohereParameterDefinition( description=param_schema.get("description", ""), type=param_schema.get("type", "string"), - isRequired=param_name in required + isRequired=param_name in required, ) - cohere_tools.append(CohereTool( - name=function_def.get("name", ""), - description=function_def.get("description", ""), - parameterDefinitions=parameter_definitions - )) + cohere_tools.append( + CohereTool( + name=function_def.get("name", ""), + description=function_def.get("description", ""), + parameterDefinitions=parameter_definitions, + ) + ) return cohere_tools @@ -793,7 +834,10 @@ class OCIChatConfig(BaseConfig): elif isinstance(content, list): text_content = "" for content_item in content: - if isinstance(content_item, dict) and content_item.get("type") == "text": + if ( + isinstance(content_item, dict) + and content_item.get("type") == "text" + ): text_content += content_item.get("text", "") return text_content return str(content) @@ -843,25 +887,28 @@ class OCIChatConfig(BaseConfig): preamble_override = None if system_messages: preamble = "\n".join( - self._extract_text_content(msg["content"]) for msg in system_messages + self._extract_text_content(msg["content"]) + for msg in system_messages ) if preamble: preamble_override = preamble # Create Cohere-specific chat request - optional_cohere_params = self._get_optional_params(OCIVendors.COHERE, optional_params) + optional_cohere_params = self._get_optional_params( + OCIVendors.COHERE, optional_params + ) chat_request = CohereChatRequest( apiFormat="COHERE", message=self._extract_text_content(user_messages[-1]["content"]), chatHistory=self.adapt_messages_to_cohere_standard(messages), preambleOverride=preamble_override, - **optional_cohere_params + **optional_cohere_params, ) data = OCICompletionPayload( compartmentId=oci_compartment_id, servingMode=servingMode, - chatRequest=chat_request + chatRequest=chat_request, ) else: # Use generic format for other vendors @@ -878,10 +925,7 @@ class OCIChatConfig(BaseConfig): return data.model_dump(exclude_none=True) def _handle_cohere_response( - self, - json_response: dict, - model: str, - model_response: ModelResponse + self, json_response: dict, model: str, model_response: ModelResponse ) -> ModelResponse: """Handle Cohere-specific response format.""" cohere_response = CohereChatResult(**json_response) @@ -909,35 +953,39 @@ class OCIChatConfig(BaseConfig): if cohere_response.chatResponse.toolCalls: tool_calls = [] for tool_call in cohere_response.chatResponse.toolCalls: - tool_calls.append({ - "id": f"call_{len(tool_calls)}", # Generate a simple ID - "type": "function", - "function": { - "name": tool_call.name, - "arguments": json.dumps(tool_call.parameters) + tool_calls.append( + { + "id": f"call_{len(tool_calls)}", # Generate a simple ID + "type": "function", + "function": { + "name": tool_call.name, + "arguments": json.dumps(tool_call.parameters), + }, } - }) + ) # Create choice from litellm.types.utils import Choices + choice = Choices( index=0, message={ "role": "assistant", "content": response_text, - "tool_calls": tool_calls + "tool_calls": tool_calls, }, - finish_reason=finish_reason + finish_reason=finish_reason, ) model_response.choices = [choice] # Extract usage info usage_info = cohere_response.chatResponse.usage from litellm.types.utils import Usage + model_response.usage = Usage( # type: ignore[attr-defined] prompt_tokens=usage_info.promptTokens, # type: ignore[union-attr] completion_tokens=usage_info.completionTokens, # type: ignore[union-attr] - total_tokens=usage_info.totalTokens # type: ignore[union-attr] + total_tokens=usage_info.totalTokens, # type: ignore[union-attr] ) return model_response @@ -947,7 +995,7 @@ class OCIChatConfig(BaseConfig): json: dict, model: str, model_response: ModelResponse, - raw_response: httpx.Response + raw_response: httpx.Response, ) -> ModelResponse: """Handle generic OCI response format.""" try: @@ -1018,7 +1066,9 @@ class OCIChatConfig(BaseConfig): if vendor == OCIVendors.COHERE: model_response = self._handle_cohere_response(json, model, model_response) else: - model_response = self._handle_generic_response(json, model, model_response, raw_response) + model_response = self._handle_generic_response( + json, model, model_response, raw_response + ) model_response._hidden_params["additional_headers"] = raw_response.headers @@ -1174,7 +1224,9 @@ def adapt_messages_to_generic_oci_standard_content_message( if isinstance(image_url, dict): image_url = image_url.get("url") if not isinstance(image_url, str): - raise Exception("Prop `image_url` must be a string or an object with a `url` property") + raise Exception( + "Prop `image_url` must be a string or an object with a `url` property" + ) new_content.append(OCIImageContentPart(imageUrl=OCIImageUrl(url=image_url))) return OCIMessage( diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index bc5aa654aa..3d9618dfed 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -396,7 +396,6 @@ class OllamaChatConfig(BaseConfig): model_response.choices[0].message = message # type: ignore model_response.choices[0].finish_reason = "tool_calls" else: - _message = litellm.Message(**response_json_message) model_response.choices[0].message = _message # type: ignore # Set finish_reason to "tool_calls" when tool_calls are present @@ -505,7 +504,10 @@ class OllamaChatCompletionResponseIterator(BaseModelResponseIterator): reasoning_content = chunk["message"].get("thinking") self.started_reasoning_content = True elif chunk["message"].get("content") is not None: - if self.started_reasoning_content and not self.finished_reasoning_content: + if ( + self.started_reasoning_content + and not self.finished_reasoning_content + ): self.finished_reasoning_content = True message_content = chunk["message"].get("content") diff --git a/litellm/llms/ollama/common_utils.py b/litellm/llms/ollama/common_utils.py index 166ceee27f..8aedd9b350 100644 --- a/litellm/llms/ollama/common_utils.py +++ b/litellm/llms/ollama/common_utils.py @@ -71,7 +71,6 @@ class OllamaModelInfo(BaseLLMModelInfo): or get_secret_str("OLLAMA_API_KEY") ) - @staticmethod def get_api_base(api_base: Optional[str] = None) -> str: from litellm.secret_managers.main import get_secret_str @@ -86,7 +85,7 @@ class OllamaModelInfo(BaseLLMModelInfo): base = self.get_api_base(api_base) api_key = self.get_api_key() - headers = { "Authorization": f"Bearer {api_key}" } if api_key else {} + headers = {"Authorization": f"Bearer {api_key}"} if api_key else {} names: set[str] = set() try: diff --git a/litellm/llms/ollama/completion/handler.py b/litellm/llms/ollama/completion/handler.py index 71956158f5..97e4f13b56 100644 --- a/litellm/llms/ollama/completion/handler.py +++ b/litellm/llms/ollama/completion/handler.py @@ -13,9 +13,8 @@ from litellm.types.utils import EmbeddingResponse def _prepare_ollama_embedding_payload( model: str, prompts: List[str], optional_params: Dict[str, Any] ) -> Dict[str, Any]: - data: Dict[str, Any] = {"model": model, "input": prompts} - special_optional_params = ["truncate", "options", "keep_alive","dimensions"] + special_optional_params = ["truncate", "options", "keep_alive", "dimensions"] for k, v in optional_params.items(): if k in special_optional_params: diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index ed14b6a331..6a03325e6c 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -93,9 +93,9 @@ class OllamaConfig(BaseConfig): repeat_penalty: Optional[float] = None temperature: Optional[float] = None seed: Optional[int] = None - stop: Optional[list] = ( - None # stop is a list based on this - https://github.com/ollama/ollama/pull/442 - ) + stop: Optional[ + list + ] = None # stop is a list based on this - https://github.com/ollama/ollama/pull/442 tfs_z: Optional[float] = None num_predict: Optional[int] = None top_k: Optional[int] = None @@ -234,9 +234,7 @@ class OllamaConfig(BaseConfig): if model.startswith("ollama/") or model.startswith("ollama_chat/"): model = model.split("/", 1)[1] api_base = ( - api_base - or get_secret_str("OLLAMA_API_BASE") - or "http://localhost:11434" + api_base or get_secret_str("OLLAMA_API_BASE") or "http://localhost:11434" ) api_key = self.get_api_key() headers = {"Authorization": f"Bearer {api_key}"} if api_key else {} @@ -598,7 +596,7 @@ class OllamaTextCompletionResponseIterator(BaseModelResponseIterator): ) else: # In this case, 'thinking' is not present in the chunk, chunk["done"] is false, - # and chunk["response"] is falsy (None or empty string), + # and chunk["response"] is falsy (None or empty string), # but Ollama is just starting to stream, so it should be processed as a normal dict return ModelResponseStream( choices=[ diff --git a/litellm/llms/openai/chat/gpt_5_transformation.py b/litellm/llms/openai/chat/gpt_5_transformation.py index f7d7c437cb..675242ae68 100644 --- a/litellm/llms/openai/chat/gpt_5_transformation.py +++ b/litellm/llms/openai/chat/gpt_5_transformation.py @@ -25,6 +25,22 @@ def _normalize_reasoning_effort_for_chat_completion( return None +def _get_effort_level(value: Union[str, dict, None]) -> Optional[str]: + """Extract the effective effort level from reasoning_effort (string or dict). + + Use this for guards that compare effort level (e.g. xhigh validation, "none" checks). + Ensures dict inputs like {"effort": "none", "summary": "detailed"} are correctly + treated as effort="none" for validation purposes. + """ + if value is None: + return None + if isinstance(value, str): + return value + if isinstance(value, dict) and "effort" in value: + return value["effort"] + return None + + class OpenAIGPT5Config(OpenAIGPTConfig): """Configuration for gpt-5 models including GPT-5-Codex variants. @@ -70,6 +86,19 @@ class OpenAIGPT5Config(OpenAIGPTConfig): model_name = model.split("/")[-1] return model_name.startswith("gpt-5.4") + @classmethod + def is_model_gpt_5_4_plus_model(cls, model: str) -> bool: + """Check if the model is gpt-5.4 or newer (5.4, 5.5, 5.6, etc., including pro).""" + model_name = model.split("/")[-1] + if not model_name.startswith("gpt-5."): + return False + try: + version_str = model_name.replace("gpt-5.", "").split("-")[0] + major = version_str.split(".")[0] + return int(major) >= 4 + except (ValueError, IndexError): + return False + @classmethod def _supports_reasoning_effort_level(cls, model: str, level: str) -> bool: """Check if the model supports a specific reasoning_effort level. @@ -150,21 +179,35 @@ class OpenAIGPT5Config(OpenAIGPTConfig): drop_params=drop_params, ) - # Normalize reasoning_effort: chat completion API expects a string, not a dict - # (e.g. {'effort': 'high', 'summary': 'detailed'} -> 'high') - raw_reasoning_effort = ( + # Get raw reasoning_effort and effective effort level for all guards. + # Use effective_effort (extracted string) for xhigh validation, "none" checks, and + # tool/sampling guards — dict inputs like {"effort": "none", "summary": "detailed"} + # must be treated as effort="none" to avoid incorrect tool-drop or sampling errors. + raw_reasoning_effort = non_default_params.get( + "reasoning_effort" + ) or optional_params.get("reasoning_effort") + effective_effort = _get_effort_level(raw_reasoning_effort) + + # Normalize to string for Chat Completions API when dict has only "effort". + # Preserve full dict (e.g. {"effort": "high", "summary": "detailed"}) for Responses API. + if isinstance(raw_reasoning_effort, dict) and set( + raw_reasoning_effort.keys() + ) <= {"effort"}: + normalized = _normalize_reasoning_effort_for_chat_completion( + raw_reasoning_effort + ) + if normalized is not None: + if "reasoning_effort" in non_default_params: + non_default_params["reasoning_effort"] = normalized + if "reasoning_effort" in optional_params: + optional_params["reasoning_effort"] = normalized + + reasoning_effort = ( non_default_params.get("reasoning_effort") or optional_params.get("reasoning_effort") + or raw_reasoning_effort ) - normalized = _normalize_reasoning_effort_for_chat_completion(raw_reasoning_effort) - if raw_reasoning_effort is not None and normalized is not None: - if "reasoning_effort" in non_default_params: - non_default_params["reasoning_effort"] = normalized - if "reasoning_effort" in optional_params: - optional_params["reasoning_effort"] = normalized - - reasoning_effort = normalized or raw_reasoning_effort - if reasoning_effort is not None and reasoning_effort == "xhigh": + if effective_effort is not None and effective_effort == "xhigh": if not self._supports_reasoning_effort_level(model, "xhigh"): if litellm.drop_params or drop_params: non_default_params.pop("reasoning_effort", None) @@ -185,23 +228,22 @@ class OpenAIGPT5Config(OpenAIGPTConfig): "max_tokens" ) - # gpt-5.4: function calls not supported when reasoning_effort != "none" - # Drop reasoning_effort when tools are present (small minority of volume) + # gpt-5.4: reasoning_effort + tools is only supported in the Responses API + # Drop reasoning_effort when tools are present in chat completions if self.is_model_gpt_5_4_model(model): has_tools = bool( non_default_params.get("tools") or optional_params.get("tools") ) - if has_tools and reasoning_effort not in (None, "none"): + if has_tools and effective_effort is not None: non_default_params.pop("reasoning_effort", None) optional_params.pop("reasoning_effort", None) - reasoning_effort = None # noqa: F841 # gpt-5.1/5.2 support logprobs, top_p, top_logprobs only when reasoning_effort="none" supports_none = self._supports_reasoning_effort_level(model, "none") if supports_none: sampling_params = ["logprobs", "top_logprobs", "top_p"] has_sampling = any(p in non_default_params for p in sampling_params) - if has_sampling and reasoning_effort not in (None, "none"): + if has_sampling and effective_effort not in (None, "none"): if litellm.drop_params or drop_params: for p in sampling_params: non_default_params.pop(p, None) @@ -211,7 +253,7 @@ class OpenAIGPT5Config(OpenAIGPTConfig): "gpt-5.1/5.2/5.4 only support logprobs, top_p, top_logprobs when " "reasoning_effort='none'. Current reasoning_effort='{}'. " "To drop unsupported params set `litellm.drop_params = True`" - ).format(reasoning_effort), + ).format(effective_effort), status_code=400, ) @@ -219,7 +261,9 @@ class OpenAIGPT5Config(OpenAIGPTConfig): temperature_value: Optional[float] = non_default_params.pop("temperature") if temperature_value is not None: # models supporting reasoning_effort="none" also support flexible temperature - if supports_none and (reasoning_effort == "none" or reasoning_effort is None): + if supports_none and ( + effective_effort == "none" or effective_effort is None + ): optional_params["temperature"] = temperature_value elif temperature_value == 1: optional_params["temperature"] = temperature_value diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index d19210d31a..63beb82ded 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -174,7 +174,9 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): model_specific_params.append("response_format") # Normalize model name for responses API (e.g., "responses/gpt-4.1" -> "gpt-4.1") - model_for_check = model.split("responses/", 1)[1] if "responses/" in model else model + model_for_check = ( + model.split("responses/", 1)[1] if "responses/" in model else model + ) if ( model_for_check in litellm.open_ai_chat_completion_models ) or model_for_check in litellm.open_ai_text_completion_models: @@ -367,10 +369,10 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): List[OpenAIMessageContentListBlock], message_content ) for i, content_item in enumerate(message_content_types): - message_content_types[i] = ( - await self._async_transform_content_item( - cast(OpenAIMessageContentListBlock, content_item), - ) + message_content_types[ + i + ] = await self._async_transform_content_item( + cast(OpenAIMessageContentListBlock, content_item), ) return messages @@ -457,12 +459,13 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): transformed_messages = await self._transform_messages( messages=messages, model=model, is_async=True ) - transformed_messages, tools = ( - self.remove_cache_control_flag_from_messages_and_tools( - model=model, - messages=transformed_messages, - tools=optional_params.get("tools", []), - ) + ( + transformed_messages, + tools, + ) = self.remove_cache_control_flag_from_messages_and_tools( + model=model, + messages=transformed_messages, + tools=optional_params.get("tools", []), ) if tools is not None and len(tools) > 0: optional_params["tools"] = tools @@ -592,9 +595,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): ) translated_choice.finish_reason = map_finish_reason( - self._get_finish_reason( - translated_message, choice["finish_reason"] - ) + self._get_finish_reason(translated_message, choice["finish_reason"]) ) transformed_choices.append(translated_choice) @@ -783,13 +784,13 @@ class OpenAIChatCompletionStreamingHandler(BaseModelResponseIterator): def _map_reasoning_to_reasoning_content(self, choices: list) -> list: """ Map 'reasoning' field to 'reasoning_content' field in delta. - - Some OpenAI-compatible providers (e.g., GLM-5, hosted_vllm) return + + Some OpenAI-compatible providers (e.g., GLM-5, hosted_vllm) return delta.reasoning, but LiteLLM expects delta.reasoning_content. - + Args: choices: List of choice objects from the streaming chunk - + Returns: List of choices with reasoning field mapped to reasoning_content """ @@ -798,12 +799,12 @@ class OpenAIChatCompletionStreamingHandler(BaseModelResponseIterator): if "reasoning" in delta: delta["reasoning_content"] = delta.pop("reasoning") return choices - + def chunk_parser(self, chunk: dict) -> ModelResponseStream: try: choices = chunk.get("choices", []) choices = self._map_reasoning_to_reasoning_content(choices) - + kwargs = { "id": chunk["id"], "object": "chat.completion.chunk", diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 10b0b58b6a..bab4c3b5eb 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -558,7 +558,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation): for streaming_choice in response.choices: if isinstance(streaming_choice, litellm.StreamingChoices): # Check for text content - if streaming_choice.delta.content and isinstance(streaming_choice.delta.content, str): + if streaming_choice.delta.content and isinstance( + streaming_choice.delta.content, str + ): return True # Check for tool calls if streaming_choice.delta.tool_calls and isinstance( diff --git a/litellm/llms/openai/chat/o_series_transformation.py b/litellm/llms/openai/chat/o_series_transformation.py index 0c5ee90b33..fe8aec9bc2 100644 --- a/litellm/llms/openai/chat/o_series_transformation.py +++ b/litellm/llms/openai/chat/o_series_transformation.py @@ -132,7 +132,9 @@ class OpenAIOSeriesConfig(OpenAIGPTConfig): def is_model_o_series_model(self, model: str) -> bool: model = model.split("/")[-1] # could be "openai/o3" or "o3" return ( - len(model) > 1 and model[0] == "o" and model[1].isdigit() + len(model) > 1 + and model[0] == "o" + and model[1].isdigit() and model in litellm.open_ai_chat_completion_models ) @@ -174,4 +176,4 @@ class OpenAIOSeriesConfig(OpenAIGPTConfig): else: return super()._transform_messages( messages, model, is_async=cast(Literal[False], False) - ) \ No newline at end of file + ) diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index b6b302782e..35723ccd63 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -7,7 +7,17 @@ import inspect import json import os import ssl -from typing import TYPE_CHECKING, Any, Dict, List, Literal, NamedTuple, Optional, Tuple, Union +from typing import ( + TYPE_CHECKING, + Any, + Dict, + List, + Literal, + NamedTuple, + Optional, + Tuple, + Union, +) import httpx import openai @@ -271,10 +281,7 @@ def get_openai_credentials( or None ) resolved_api_key = ( - api_key - or litellm.api_key - or litellm.openai_key - or os.getenv("OPENAI_API_KEY") + api_key or litellm.api_key or litellm.openai_key or os.getenv("OPENAI_API_KEY") ) return OpenAICredentials( api_base=resolved_api_base, diff --git a/litellm/llms/openai/completion/transformation.py b/litellm/llms/openai/completion/transformation.py index 77dc0b54fe..44a4949d45 100644 --- a/litellm/llms/openai/completion/transformation.py +++ b/litellm/llms/openai/completion/transformation.py @@ -111,9 +111,9 @@ class OpenAITextCompletionConfig(BaseTextCompletionConfig, OpenAIGPTConfig): if "model" in response_object: model_response_object.model = response_object["model"] - model_response_object._hidden_params["original_response"] = ( - response_object # track original response, if users make a litellm.text_completion() request, we can return the original response - ) + model_response_object._hidden_params[ + "original_response" + ] = response_object # track original response, if users make a litellm.text_completion() request, we can return the original response return model_response_object except Exception as e: raise e diff --git a/litellm/llms/openai/containers/transformation.py b/litellm/llms/openai/containers/transformation.py index b89204230a..645538fdd9 100644 --- a/litellm/llms/openai/containers/transformation.py +++ b/litellm/llms/openai/containers/transformation.py @@ -31,15 +31,13 @@ else: class OpenAIContainerConfig(BaseContainerConfig): - """Configuration class for OpenAI container API. - """ + """Configuration class for OpenAI container API.""" def __init__(self): super().__init__() def get_supported_openai_params(self) -> list: - """Get the list of supported OpenAI parameters for container API. - """ + """Get the list of supported OpenAI parameters for container API.""" return [ "name", "expires_after", @@ -78,8 +76,7 @@ class OpenAIContainerConfig(BaseContainerConfig): api_base: Optional[str], litellm_params: dict, ) -> str: - """Get the complete URL for OpenAI container API. - """ + """Get the complete URL for OpenAI container API.""" api_base = ( api_base or litellm.api_base @@ -97,11 +94,11 @@ class OpenAIContainerConfig(BaseContainerConfig): litellm_params: GenericLiteLLMParams, headers: dict, ) -> Dict: - """Transform the container creation request for OpenAI API. - """ + """Transform the container creation request for OpenAI API.""" # Remove extra_headers from optional params as they're handled separately container_create_optional_request_params = { - k: v for k, v in container_create_optional_request_params.items() + k: v + for k, v in container_create_optional_request_params.items() if k not in ["extra_headers"] } @@ -118,8 +115,7 @@ class OpenAIContainerConfig(BaseContainerConfig): raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, ) -> ContainerObject: - """Transform the OpenAI container creation response. - """ + """Transform the OpenAI container creation response.""" response_data = raw_response.json() # Transform the response data @@ -132,12 +128,17 @@ class OpenAIContainerConfig(BaseContainerConfig): sessions=1, provider="openai", ) - - if not hasattr(container_obj, "_hidden_params") or container_obj._hidden_params is None: + + if ( + not hasattr(container_obj, "_hidden_params") + or container_obj._hidden_params is None + ): container_obj._hidden_params = {} if "additional_headers" not in container_obj._hidden_params: container_obj._hidden_params["additional_headers"] = {} - container_obj._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = container_cost + container_obj._hidden_params["additional_headers"][ + "llm_provider-x-litellm-response-cost" + ] = container_cost return container_obj @@ -152,7 +153,7 @@ class OpenAIContainerConfig(BaseContainerConfig): extra_query: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: """Transform the container list request for OpenAI API. - + OpenAI API expects the following request: - GET /v1/containers """ @@ -179,8 +180,7 @@ class OpenAIContainerConfig(BaseContainerConfig): raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, ) -> ContainerListResponse: - """Transform the OpenAI container list response. - """ + """Transform the OpenAI container list response.""" response_data = raw_response.json() # Transform the response data @@ -195,8 +195,7 @@ class OpenAIContainerConfig(BaseContainerConfig): litellm_params: GenericLiteLLMParams, headers: dict, ) -> Tuple[str, Dict]: - """Transform the OpenAI container retrieve request. - """ + """Transform the OpenAI container retrieve request.""" # For container retrieve, we just need to construct the URL url = f"{api_base.rstrip('/')}/{container_id}" @@ -210,8 +209,7 @@ class OpenAIContainerConfig(BaseContainerConfig): raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, ) -> ContainerObject: - """Transform the OpenAI container retrieve response. - """ + """Transform the OpenAI container retrieve response.""" response_data = raw_response.json() # Transform the response data container_obj = ContainerObject(**response_data) # type: ignore[arg-type] @@ -226,7 +224,7 @@ class OpenAIContainerConfig(BaseContainerConfig): headers: dict, ) -> Tuple[str, Dict]: """Transform the container delete request for OpenAI API. - + OpenAI API expects the following request: - DELETE /v1/containers/{container_id} """ @@ -243,8 +241,7 @@ class OpenAIContainerConfig(BaseContainerConfig): raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, ) -> DeleteContainerResult: - """Transform the OpenAI container delete response. - """ + """Transform the OpenAI container delete response.""" response_data = raw_response.json() # Transform the response data @@ -264,7 +261,7 @@ class OpenAIContainerConfig(BaseContainerConfig): extra_query: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: """Transform the container file list request for OpenAI API. - + OpenAI API expects the following request: - GET /v1/containers/{container_id}/files """ @@ -291,8 +288,7 @@ class OpenAIContainerConfig(BaseContainerConfig): raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, ) -> ContainerFileListResponse: - """Transform the OpenAI container file list response. - """ + """Transform the OpenAI container file list response.""" response_data = raw_response.json() # Transform the response data @@ -309,7 +305,7 @@ class OpenAIContainerConfig(BaseContainerConfig): headers: dict, ) -> Tuple[str, Dict]: """Transform the container file content request for OpenAI API. - + OpenAI API expects the following request: - GET /v1/containers/{container_id}/files/{file_id}/content """ @@ -327,13 +323,16 @@ class OpenAIContainerConfig(BaseContainerConfig): logging_obj: LiteLLMLoggingObj, ) -> bytes: """Transform the OpenAI container file content response. - + Returns the raw binary content of the file. """ return raw_response.content def get_error_class( - self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers], + self, + error_message: str, + status_code: int, + headers: Union[dict, httpx.Headers], ) -> BaseLLMException: from ...base_llm.chat.transformation import BaseLLMException @@ -342,4 +341,3 @@ class OpenAIContainerConfig(BaseContainerConfig): message=error_message, headers=headers, ) - diff --git a/litellm/llms/openai/image_edit/__init__.py b/litellm/llms/openai/image_edit/__init__.py index c1898326b7..5d933b8186 100644 --- a/litellm/llms/openai/image_edit/__init__.py +++ b/litellm/llms/openai/image_edit/__init__.py @@ -3,24 +3,27 @@ from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from .dalle2_transformation import DallE2ImageEditConfig from .transformation import OpenAIImageEditConfig -__all__ = ["OpenAIImageEditConfig", "DallE2ImageEditConfig", "get_openai_image_edit_config"] +__all__ = [ + "OpenAIImageEditConfig", + "DallE2ImageEditConfig", + "get_openai_image_edit_config", +] def get_openai_image_edit_config(model: str) -> BaseImageEditConfig: """ Get the appropriate OpenAI image edit config based on the model. - + Args: model: The model name (e.g., "dall-e-2", "gpt-image-1") - + Returns: The appropriate config instance for the model """ model_normalized = model.lower().replace("-", "").replace("_", "") - + if model_normalized == "dalle2": return DallE2ImageEditConfig() else: # Default to standard OpenAI config for gpt-image-1 and other models return OpenAIImageEditConfig() - diff --git a/litellm/llms/openai/image_edit/dalle2_transformation.py b/litellm/llms/openai/image_edit/dalle2_transformation.py index fd697b210e..04995ce951 100644 --- a/litellm/llms/openai/image_edit/dalle2_transformation.py +++ b/litellm/llms/openai/image_edit/dalle2_transformation.py @@ -22,7 +22,7 @@ else: class DallE2ImageEditConfig(OpenAIImageEditConfig): """ DALL-E-2 specific configuration for image edit API. - + DALL-E-2 only supports editing a single image (not an array). Uses "image" field name instead of "image[]". """ @@ -40,7 +40,7 @@ class DallE2ImageEditConfig(OpenAIImageEditConfig): Transform image edit request for DALL-E-2. DALL-E-2 only accepts a single image with field name "image" (not "image[]"). - """ + """ request_params = { "model": model, **image_edit_optional_request_params, @@ -49,11 +49,10 @@ class DallE2ImageEditConfig(OpenAIImageEditConfig): request_params["image"] = image if prompt is not None: request_params["prompt"] = prompt - + request = ImageEditRequestParams(**request_params) request_dict = cast(Dict, request) - ######################################################### # Separate images and masks as `files` and send other parameters as `data` ######################################################### @@ -103,4 +102,3 @@ class DallE2ImageEditConfig(OpenAIImageEditConfig): files_list.append(("mask", ("mask.png", _mask, mask_content_type))) return data_without_files, files_list - diff --git a/litellm/llms/openai/image_edit/transformation.py b/litellm/llms/openai/image_edit/transformation.py index a92a89eac6..6917e8d799 100644 --- a/litellm/llms/openai/image_edit/transformation.py +++ b/litellm/llms/openai/image_edit/transformation.py @@ -101,7 +101,7 @@ class OpenAIImageEditConfig(BaseImageEditConfig): request_params["image"] = image if prompt is not None: request_params["prompt"] = prompt - + request = ImageEditRequestParams(**request_params) request_dict = cast(Dict, request) diff --git a/litellm/llms/openai/image_generation/cost_calculator.py b/litellm/llms/openai/image_generation/cost_calculator.py index 988d562613..8bca75172f 100644 --- a/litellm/llms/openai/image_generation/cost_calculator.py +++ b/litellm/llms/openai/image_generation/cost_calculator.py @@ -47,8 +47,8 @@ def cost_calculator( # ImageUsage has the same format as ResponseAPIUsage from litellm.responses.utils import ResponseAPILoggingUtils - chat_usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - usage + chat_usage = ( + ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) ) # Use generic_cost_per_token for cost calculation diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 5a8b4aafe0..be54267748 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -522,17 +522,14 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger - callbacks = litellm.callbacks + ( - logging_obj.dynamic_success_callbacks or [] - ) + callbacks = litellm.callbacks + (logging_obj.dynamic_success_callbacks or []) # Avoid logging full callback objects to prevent leaking sensitive data - verbose_logger.debug( - "LiteLLM.AgenticHooks: callbacks_count=%s", len(callbacks) - ) + verbose_logger.debug("LiteLLM.AgenticHooks: callbacks_count=%s", len(callbacks)) tools = optional_params.get("tools", []) # Avoid logging full tools payloads; they may contain sensitive parameters verbose_logger.debug( - "LiteLLM.AgenticHooks: tools_count=%s", len(tools) if isinstance(tools, list) else 1 if tools else 0 + "LiteLLM.AgenticHooks: tools_count=%s", + len(tools) if isinstance(tools, list) else 1 if tools else 0, ) # Get custom_llm_provider from litellm_params custom_llm_provider = litellm_params.get("custom_llm_provider", "openai") @@ -541,37 +538,46 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): try: if isinstance(callback, CustomLogger): # Check if the callback has the chat completion agentic loop methods - if not hasattr(callback, 'async_should_run_chat_completion_agentic_loop'): + if not hasattr( + callback, "async_should_run_chat_completion_agentic_loop" + ): continue - + # First: Check if agentic loop should run (using chat completion method) - should_run, tool_calls = ( - await callback.async_should_run_chat_completion_agentic_loop( - response=response, - model=model, - messages=messages, - tools=tools, - stream=stream, - custom_llm_provider=custom_llm_provider, - kwargs=litellm_params, - ) + ( + should_run, + tool_calls, + ) = await callback.async_should_run_chat_completion_agentic_loop( + response=response, + model=model, + messages=messages, + tools=tools, + stream=stream, + custom_llm_provider=custom_llm_provider, + kwargs=litellm_params, ) if should_run: # Second: Execute agentic loop - kwargs_with_provider = litellm_params.copy() if litellm_params else {} - kwargs_with_provider["custom_llm_provider"] = custom_llm_provider - + kwargs_with_provider = ( + litellm_params.copy() if litellm_params else {} + ) + kwargs_with_provider[ + "custom_llm_provider" + ] = custom_llm_provider + # For OpenAI Chat Completions, use the chat completion agentic loop method - agentic_response = await callback.async_run_chat_completion_agentic_loop( - tools=tool_calls, - model=model, - messages=messages, - response=response, - optional_params=optional_params, - logging_obj=logging_obj, - stream=stream, - kwargs=kwargs_with_provider, + agentic_response = ( + await callback.async_run_chat_completion_agentic_loop( + tools=tool_calls, + model=model, + messages=messages, + response=response, + optional_params=optional_params, + logging_obj=logging_obj, + stream=stream, + kwargs=kwargs_with_provider, + ) ) # First hook that runs agentic loop wins return agentic_response @@ -951,7 +957,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): stream=False, litellm_params=litellm_params, ) - + if agentic_response is not None: final_response_obj = agentic_response diff --git a/litellm/llms/openai/realtime/handler.py b/litellm/llms/openai/realtime/handler.py index 05915e36a6..c04857fc25 100644 --- a/litellm/llms/openai/realtime/handler.py +++ b/litellm/llms/openai/realtime/handler.py @@ -18,28 +18,28 @@ from ..openai import OpenAIChatCompletion class OpenAIRealtime(OpenAIChatCompletion): """ Base handler for OpenAI-compatible realtime WebSocket connections. - + Subclasses can override template methods to customize: - _get_default_api_base(): Default API base URL - _get_additional_headers(): Extra headers beyond Authorization - _get_ssl_config(): SSL configuration for WebSocket connection """ - + def _get_default_api_base(self) -> str: """ Get the default API base URL for this provider. Override this in subclasses to set provider-specific defaults. """ return "https://api.openai.com/" - + def _get_additional_headers(self, api_key: str) -> dict: """ Get additional headers beyond Authorization. Override this in subclasses to customize headers (e.g., remove OpenAI-Beta). - + Args: api_key: API key for authentication - + Returns: Dictionary of additional headers """ @@ -47,31 +47,31 @@ class OpenAIRealtime(OpenAIChatCompletion): "Authorization": f"Bearer {api_key}", "OpenAI-Beta": "realtime=v1", } - + def _get_ssl_config(self, url: str) -> Any: """ Get SSL configuration for WebSocket connection. Override this in subclasses to customize SSL behavior. - + Args: url: WebSocket URL (ws:// or wss://) - + Returns: SSL configuration (None, True, or SSLContext) """ if url.startswith("ws://"): return None - + # Use the shared SSL context which respects custom CA certs and SSL settings ssl_config = get_shared_realtime_ssl_context() - + # If ssl_config is False (ssl_verify=False), websockets library needs True instead # to establish connection without verification (False would fail) if ssl_config is False: return True - + return ssl_config - + def _construct_url(self, api_base: str, query_params: RealtimeQueryParams) -> str: """ Construct the backend websocket URL with all query parameters (including 'model'). @@ -104,7 +104,7 @@ class OpenAIRealtime(OpenAIChatCompletion): ): import websockets from websockets.asyncio.client import ClientConnection - + if api_base is None: api_base = self._get_default_api_base() if api_key is None: @@ -118,10 +118,10 @@ class OpenAIRealtime(OpenAIChatCompletion): try: # Get provider-specific SSL configuration ssl_config = self._get_ssl_config(url) - + # Get provider-specific headers headers = self._get_additional_headers(api_key) - + # Log a masked request preview consistent with other endpoints. logging_obj.pre_call( input=None, diff --git a/litellm/llms/openai/realtime/http_transformation.py b/litellm/llms/openai/realtime/http_transformation.py index ff69ef987d..1663fcd1fc 100644 --- a/litellm/llms/openai/realtime/http_transformation.py +++ b/litellm/llms/openai/realtime/http_transformation.py @@ -25,13 +25,17 @@ class OpenAIRealtimeHTTPConfig(BaseRealtimeHTTPConfig): or "" ) - def get_complete_url(self, api_base: Optional[str], model: str, api_version: Optional[str] = None) -> str: + def get_complete_url( + self, api_base: Optional[str], model: str, api_version: Optional[str] = None + ) -> str: base = self.get_api_base(api_base).rstrip("/") if base.endswith("/v1"): base = base[:-3] return f"{base}/v1/realtime/client_secrets" - def get_realtime_calls_url(self, api_base: Optional[str], model: str, api_version: Optional[str] = None) -> str: + def get_realtime_calls_url( + self, api_base: Optional[str], model: str, api_version: Optional[str] = None + ) -> str: base = self.get_api_base(api_base).rstrip("/") if base.endswith("/v1"): base = base[:-3] diff --git a/litellm/llms/openai/responses/count_tokens/handler.py b/litellm/llms/openai/responses/count_tokens/handler.py index 721d07796e..7fb5f6dad7 100644 --- a/litellm/llms/openai/responses/count_tokens/handler.py +++ b/litellm/llms/openai/responses/count_tokens/handler.py @@ -66,7 +66,9 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig): llm_provider=litellm.LlmProviders.OPENAI ) - request_timeout = timeout if timeout is not None else litellm.request_timeout + request_timeout = ( + timeout if timeout is not None else litellm.request_timeout + ) response = await async_client.post( endpoint_url, diff --git a/litellm/llms/openai/responses/count_tokens/transformation.py b/litellm/llms/openai/responses/count_tokens/transformation.py index 3893775fc0..41d1a01ec6 100644 --- a/litellm/llms/openai/responses/count_tokens/transformation.py +++ b/litellm/llms/openai/responses/count_tokens/transformation.py @@ -52,9 +52,7 @@ class OpenAICountTokensConfig: "Authorization": f"Bearer {api_key}", } - def validate_request( - self, model: str, input: Union[str, List[Any]] - ) -> None: + def validate_request(self, model: str, input: Union[str, List[Any]]) -> None: if not model: raise ValueError("model parameter is required") @@ -139,20 +137,24 @@ class OpenAICountTokensConfig: if tool_calls: for tc in tool_calls: func = tc.get("function", {}) - input_items.append({ - "type": "function_call", - "call_id": tc.get("id", ""), - "name": func.get("name", ""), - "arguments": func.get("arguments", ""), - }) + input_items.append( + { + "type": "function_call", + "call_id": tc.get("id", ""), + "name": func.get("name", ""), + "arguments": func.get("arguments", ""), + } + ) elif not content: input_items.append({"role": "assistant", "content": content}) elif role == "tool": - input_items.append({ - "type": "function_call_output", - "call_id": msg.get("tool_call_id", ""), - "output": content if isinstance(content, str) else str(content), - }) + input_items.append( + { + "type": "function_call_output", + "call_id": msg.get("tool_call_id", ""), + "output": content if isinstance(content, str) else str(content), + } + ) instructions = "\n".join(instructions_parts) if instructions_parts else None return input_items, instructions diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 7c3354cf88..466e2e76f1 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -30,22 +30,27 @@ Output: response.output is List[GenericResponseOutputItem] where each has: from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast -from openai.types.responses.response_function_tool_call import \ - ResponseFunctionToolCall +from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall from pydantic import BaseModel from litellm._logging import verbose_proxy_logger from litellm.completion_extras.litellm_responses_transformation.transformation import ( LiteLLMResponsesTransformationHandler, - OpenAiResponsesToChatCompletionStreamIterator) -from litellm.llms.base_llm.guardrail_translation.base_translation import \ - BaseTranslation -from litellm.responses.litellm_completion_transformation.transformation import \ - LiteLLMCompletionResponsesConfig -from litellm.types.llms.openai import (ChatCompletionToolCallChunk, - ChatCompletionToolParam) -from litellm.types.responses.main import (GenericResponseOutputItem, - OutputFunctionToolCall, OutputText) + OpenAiResponsesToChatCompletionStreamIterator, +) +from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, +) +from litellm.types.llms.openai import ( + ChatCompletionToolCallChunk, + ChatCompletionToolParam, +) +from litellm.types.responses.main import ( + GenericResponseOutputItem, + OutputFunctionToolCall, + OutputText, +) from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 2808010366..9d909fd401 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -181,7 +181,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): f"Error constructing ResponsesAPIResponse: {raw_response_json}, using model_construct" ) response = ResponsesAPIResponse.model_construct(**raw_response_json) - + # Store processed headers in additional_headers so they get returned to the client response._hidden_params["additional_headers"] = processed_headers response._hidden_params["headers"] = raw_response_headers @@ -411,7 +411,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): ) -> ResponsesAPIResponse: """ Transform the get response API response into a ResponsesAPIResponse - """ + """ try: raw_response_json = raw_response.json() except Exception: @@ -423,7 +423,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): response = ResponsesAPIResponse(**raw_response_json) response._hidden_params["additional_headers"] = processed_headers response._hidden_params["headers"] = raw_response_headers - + return response ######################################################### @@ -503,11 +503,11 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): ) raw_response_headers = dict(raw_response.headers) processed_headers = process_response_headers(raw_response_headers) - + response = ResponsesAPIResponse(**raw_response_json) response._hidden_params["additional_headers"] = processed_headers response._hidden_params["headers"] = raw_response_headers - + return response ######################################################### @@ -532,14 +532,14 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): parsed_url = httpx.URL(api_base) compact_path = parsed_url.path.rstrip("/") + "/compact" url = str(parsed_url.copy_with(path=compact_path)) - + input = self._validate_input_param(input) data = dict( ResponsesAPIRequestParams( model=model, input=input, **response_api_optional_request_params ) ) - + return url, data def transform_compact_response_api_response( @@ -565,7 +565,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): ) raw_response_headers = dict(raw_response.headers) processed_headers = process_response_headers(raw_response_headers) - + try: response = ResponsesAPIResponse(**raw_response_json) except Exception: @@ -573,8 +573,8 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): f"Error constructing ResponsesAPIResponse: {raw_response_json}, using model_construct" ) response = ResponsesAPIResponse.model_construct(**raw_response_json) - + response._hidden_params["additional_headers"] = processed_headers response._hidden_params["headers"] = raw_response_headers - + return response diff --git a/litellm/llms/openai/transcriptions/handler.py b/litellm/llms/openai/transcriptions/handler.py index 397b4c9956..e079a17087 100644 --- a/litellm/llms/openai/transcriptions/handler.py +++ b/litellm/llms/openai/transcriptions/handler.py @@ -37,7 +37,6 @@ class OpenAIAudioTranscription(OpenAIChatCompletion): - call openai_aclient.audio.transcriptions.create by default """ try: - raw_response = ( await openai_aclient.audio.transcriptions.with_raw_response.create( **data, timeout=timeout diff --git a/litellm/llms/openai/transcriptions/whisper_transformation.py b/litellm/llms/openai/transcriptions/whisper_transformation.py index fa507e1bc2..1a7f47ae56 100644 --- a/litellm/llms/openai/transcriptions/whisper_transformation.py +++ b/litellm/llms/openai/transcriptions/whisper_transformation.py @@ -110,9 +110,9 @@ class OpenAIWhisperAudioTranscriptionConfig(BaseAudioTranscriptionConfig): if "response_format" not in data or ( data["response_format"] == "text" or data["response_format"] == "json" ): - data["response_format"] = ( - "verbose_json" # ensures 'duration' is received - used for cost calculation - ) + data[ + "response_format" + ] = "verbose_json" # ensures 'duration' is received - used for cost calculation return AudioTranscriptionRequestData( data=data, diff --git a/litellm/llms/openai/vector_store_files/transformation.py b/litellm/llms/openai/vector_store_files/transformation.py index 8953e404f3..cd5f10251b 100644 --- a/litellm/llms/openai/vector_store_files/transformation.py +++ b/litellm/llms/openai/vector_store_files/transformation.py @@ -41,9 +41,9 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): } } - def get_vector_store_file_endpoints_by_type(self) -> Dict[ - str, Tuple[Tuple[str, str], ...] - ]: + def get_vector_store_file_endpoints_by_type( + self, + ) -> Dict[str, Tuple[Tuple[str, str], ...]]: return { "read": ( ("GET", "/vector_stores/{vector_store_id}/files"), diff --git a/litellm/llms/openai/videos/transformation.py b/litellm/llms/openai/videos/transformation.py index 5c880ab665..e224097fb0 100644 --- a/litellm/llms/openai/videos/transformation.py +++ b/litellm/llms/openai/videos/transformation.py @@ -69,7 +69,7 @@ class OpenAIVideoConfig(BaseVideoConfig): # Use api_key from litellm_params if available, otherwise fall back to other sources if litellm_params and litellm_params.api_key: api_key = api_key or litellm_params.api_key - + api_key = ( api_key or litellm.api_key @@ -94,7 +94,7 @@ class OpenAIVideoConfig(BaseVideoConfig): """ if api_base is None: api_base = "https://api.openai.com/v1" - + return f"{api_base.rstrip('/')}/videos" def transform_video_create_request( @@ -111,15 +111,14 @@ class OpenAIVideoConfig(BaseVideoConfig): """ # Remove model and extra_headers from optional params as they're handled separately video_create_optional_request_params = { - k: v for k, v in video_create_optional_request_params.items() + k: v + for k, v in video_create_optional_request_params.items() if k not in ["model", "extra_headers", "prompt"] } - + # Create the request data video_create_request = CreateVideoRequest( - model=model, - prompt=prompt, - **video_create_optional_request_params + model=model, prompt=prompt, **video_create_optional_request_params ) request_dict = cast(Dict, video_create_request) @@ -149,21 +148,23 @@ class OpenAIVideoConfig(BaseVideoConfig): ) -> VideoObject: """Transform the OpenAI video creation response.""" response_data = raw_response.json() - + video_obj = VideoObject(**response_data) # type: ignore[arg-type] - + if custom_llm_provider and video_obj.id: - video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, model) - + video_obj.id = encode_video_id_with_provider( + video_obj.id, custom_llm_provider, model + ) + usage_data = {} if video_obj: - if hasattr(video_obj, 'seconds') and video_obj.seconds: + if hasattr(video_obj, "seconds") and video_obj.seconds: try: usage_data["duration_seconds"] = float(video_obj.seconds) except (ValueError, TypeError): pass video_obj.usage = usage_data - + return video_obj def transform_video_content_request( @@ -204,24 +205,24 @@ class OpenAIVideoConfig(BaseVideoConfig): ) -> Tuple[str, Dict]: """ Transform the video remix request for OpenAI API. - + OpenAI API expects the following request: - POST /v1/videos/{video_id}/remix """ original_video_id = extract_original_video_id(video_id) - + # Construct the URL for video remix url = f"{api_base.rstrip('/')}/{original_video_id}/remix" - + # Prepare the request data data = {"prompt": prompt} - + # Add any extra body parameters if extra_body: data.update(extra_body) - + return url, data - + def transform_video_content_response( self, raw_response: httpx.Response, @@ -240,18 +241,20 @@ class OpenAIVideoConfig(BaseVideoConfig): Transform the OpenAI video remix response. """ response_data = raw_response.json() - + # Transform the response data video_obj = VideoObject(**response_data) # type: ignore[arg-type] - + if custom_llm_provider and video_obj.id: - video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, None) - + video_obj.id = encode_video_id_with_provider( + video_obj.id, custom_llm_provider, None + ) + # Create usage object with duration information for cost calculation # Video remix API doesn't provide usage, so we create one with duration usage_data = {} if video_obj: - if hasattr(video_obj, 'seconds') and video_obj.seconds: + if hasattr(video_obj, "seconds") and video_obj.seconds: try: usage_data["duration_seconds"] = float(video_obj.seconds) except (ValueError, TypeError): @@ -346,18 +349,18 @@ class OpenAIVideoConfig(BaseVideoConfig): ) -> Tuple[str, Dict]: """ Transform the video delete request for OpenAI API. - + OpenAI API expects the following request: - DELETE /v1/videos/{video_id} """ original_video_id = extract_original_video_id(video_id) - + # Construct the URL for video delete url = f"{api_base.rstrip('/')}/{original_video_id}" - + # No data needed for DELETE request data: Dict[str, Any] = {} - + return url, data def transform_video_delete_response( @@ -369,7 +372,7 @@ class OpenAIVideoConfig(BaseVideoConfig): Transform the OpenAI video delete response. """ response_data = raw_response.json() - + # Transform the response data video_obj = VideoObject(**response_data) # type: ignore[arg-type] # type: ignore[arg-type] @@ -387,13 +390,13 @@ class OpenAIVideoConfig(BaseVideoConfig): """ # Extract the original video_id (remove provider encoding if present) original_video_id = extract_original_video_id(video_id) - + # For video retrieve, we just need to construct the URL url = f"{api_base.rstrip('/')}/{original_video_id}" - + # No additional data needed for GET request data: Dict[str, Any] = {} - + return url, data def transform_video_status_retrieve_response( @@ -408,9 +411,11 @@ class OpenAIVideoConfig(BaseVideoConfig): response_data = raw_response.json() # Transform the response data video_obj = VideoObject(**response_data) # type: ignore[arg-type] - + if custom_llm_provider and video_obj.id: - video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, None) + video_obj.id = encode_video_id_with_provider( + video_obj.id, custom_llm_provider, None + ) return video_obj @@ -437,4 +442,6 @@ class OpenAIVideoConfig(BaseVideoConfig): if isinstance(image, BufferedReader): files_list.append((field_name, (image.name, image, image_content_type))) else: - files_list.append((field_name, ("input_reference.png", image, image_content_type))) + files_list.append( + (field_name, ("input_reference.png", image, image_content_type)) + ) diff --git a/litellm/llms/openai_like/dynamic_config.py b/litellm/llms/openai_like/dynamic_config.py index 8be749f34a..3d66556e52 100644 --- a/litellm/llms/openai_like/dynamic_config.py +++ b/litellm/llms/openai_like/dynamic_config.py @@ -44,11 +44,11 @@ def create_config_class(provider: SimpleProviderConfig): self, messages: List[AllMessageValues], model: str, is_async: bool = False ) -> Union[List[AllMessageValues], Coroutine[Any, Any, List[AllMessageValues]]]: """Transform messages based on special_handling config""" - + # Handle content list to string conversion if configured if provider.special_handling.get("convert_content_list_to_string"): messages = handle_messages_with_content_list_to_str_conversion(messages) - + if is_async: return super()._transform_messages( messages=messages, model=model, is_async=True @@ -108,7 +108,13 @@ def create_config_class(provider: SimpleProviderConfig): ) if not _supports_fc: - tool_params = ["tools", "tool_choice", "function_call", "functions", "parallel_tool_calls"] + tool_params = [ + "tools", + "tool_choice", + "function_call", + "functions", + "parallel_tool_calls", + ] for param in tool_params: if param in supported_params: supported_params.remove(param) @@ -129,7 +135,7 @@ def create_config_class(provider: SimpleProviderConfig): """Apply parameter mappings and constraints""" supported_params = self.get_supported_openai_params(model) - + # Apply supported params for param, value in non_default_params.items(): # Check parameter mappings first @@ -197,10 +203,7 @@ def create_responses_config_class(provider: SimpleProviderConfig): litellm_params: Optional[GenericLiteLLMParams], ) -> dict: litellm_params = litellm_params or GenericLiteLLMParams() - api_key = ( - litellm_params.api_key - or get_secret_str(provider.api_key_env) - ) + api_key = litellm_params.api_key or get_secret_str(provider.api_key_env) if api_key: headers["Authorization"] = f"Bearer {api_key}" return headers @@ -217,9 +220,7 @@ def create_responses_config_class(provider: SimpleProviderConfig): api_base = provider.base_url if api_base is None: - raise ValueError( - f"api_base is required for provider {provider.slug}" - ) + raise ValueError(f"api_base is required for provider {provider.slug}") api_base = api_base.rstrip("/") return f"{api_base}/responses" diff --git a/litellm/llms/openai_like/embedding/handler.py b/litellm/llms/openai_like/embedding/handler.py index d0d26d5959..e3884fa56d 100644 --- a/litellm/llms/openai_like/embedding/handler.py +++ b/litellm/llms/openai_like/embedding/handler.py @@ -105,7 +105,9 @@ class OpenAILikeEmbeddingHandler(OpenAILikeBase): custom_endpoint=custom_endpoint, ) model = model - filtered_optional_params = {k: v for k, v in optional_params.items() if v not in (None, '')} + filtered_optional_params = { + k: v for k, v in optional_params.items() if v not in (None, "") + } data = {"model": model, "input": input, **filtered_optional_params} ## LOGGING diff --git a/litellm/llms/openai_like/json_loader.py b/litellm/llms/openai_like/json_loader.py index 8b55fe4b61..c6ff0f7a39 100644 --- a/litellm/llms/openai_like/json_loader.py +++ b/litellm/llms/openai_like/json_loader.py @@ -37,7 +37,7 @@ class JSONProviderRegistry: return json_path = Path(__file__).parent / "providers.json" - + if not json_path.exists(): # No JSON file yet, that's okay cls._loaded = True @@ -52,7 +52,9 @@ class JSONProviderRegistry: cls._loaded = True except Exception as e: - verbose_logger.warning(f"Warning: Failed to load JSON provider configs: {e}") + verbose_logger.warning( + f"Warning: Failed to load JSON provider configs: {e}" + ) cls._loaded = True @classmethod diff --git a/litellm/llms/openrouter/chat/transformation.py b/litellm/llms/openrouter/chat/transformation.py index e3770dbbf4..86e63fd0c4 100644 --- a/litellm/llms/openrouter/chat/transformation.py +++ b/litellm/llms/openrouter/chat/transformation.py @@ -24,6 +24,7 @@ from ..common_utils import OpenRouterException class CacheControlSupportedModels(str, Enum): """Models that support cache_control in content blocks.""" + CLAUDE = "claude" GEMINI = "gemini" MINIMAX = "minimax" @@ -69,15 +70,15 @@ class OpenrouterConfig(OpenAIGPTConfig): extra_body["models"] = models if route is not None: extra_body["route"] = route - mapped_openai_params["extra_body"] = ( - extra_body # openai client supports `extra_body` param - ) + mapped_openai_params[ + "extra_body" + ] = extra_body # openai client supports `extra_body` param return mapped_openai_params def _supports_cache_control_in_content(self, model: str) -> bool: """ Check if the model supports cache_control in content blocks. - + Returns: bool: True if model supports cache_control (Claude or Gemini models) """ @@ -106,7 +107,7 @@ class OpenrouterConfig(OpenAIGPTConfig): """ Move cache_control from message level to content blocks. OpenRouter requires cache_control to be inside content blocks, not at message level. - + To avoid exceeding Anthropic's limit of 4 cache breakpoints, cache_control is only added to the LAST content block in each message. """ @@ -114,10 +115,10 @@ class OpenrouterConfig(OpenAIGPTConfig): for message in messages: message_dict = dict(message) cache_control = message_dict.pop("cache_control", None) - + if cache_control is not None: content = message_dict.get("content") - + if isinstance(content, list): # Content is already a list, add cache_control only to the last block if len(content) > 0: @@ -138,10 +139,10 @@ class OpenrouterConfig(OpenAIGPTConfig): "cache_control": cache_control, } ] - + # Cast back to AllMessageValues after modification transformed_messages.append(cast(AllMessageValues, message_dict)) - + return transformed_messages def transform_request( @@ -160,7 +161,7 @@ class OpenrouterConfig(OpenAIGPTConfig): """ if self._supports_cache_control_in_content(model): messages = self._move_cache_control_to_content(messages) - + extra_body = optional_params.pop("extra_body", {}) response = super().transform_request( model, messages, optional_params, litellm_params, headers @@ -223,7 +224,9 @@ class OpenrouterConfig(OpenAIGPTConfig): model_response._hidden_params = {} if "additional_headers" not in model_response._hidden_params: model_response._hidden_params["additional_headers"] = {} - model_response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = float(response_cost) + model_response._hidden_params["additional_headers"][ + "llm_provider-x-litellm-response-cost" + ] = float(response_cost) except Exception: # If we can't extract cost, continue without it - don't fail the response pass diff --git a/litellm/llms/openrouter/image_edit/transformation.py b/litellm/llms/openrouter/image_edit/transformation.py index 7a4cef1798..9e5e313aad 100644 --- a/litellm/llms/openrouter/image_edit/transformation.py +++ b/litellm/llms/openrouter/image_edit/transformation.py @@ -55,7 +55,13 @@ from litellm.llms.openrouter.common_utils import OpenRouterException from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams from litellm.types.router import GenericLiteLLMParams -from litellm.types.utils import FileTypes, ImageObject, ImageResponse, ImageUsage, ImageUsageInputTokensDetails +from litellm.types.utils import ( + FileTypes, + ImageObject, + ImageResponse, + ImageUsage, + ImageUsageInputTokensDetails, +) if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -91,7 +97,9 @@ class OpenRouterImageEditConfig(BaseImageEditConfig): if key == "size": if "image_config" not in mapped_params: mapped_params["image_config"] = {} - mapped_params["image_config"]["aspect_ratio"] = self._map_size_to_aspect_ratio(cast(str, value)) + mapped_params["image_config"][ + "aspect_ratio" + ] = self._map_size_to_aspect_ratio(cast(str, value)) elif key == "quality": image_size = self._map_quality_to_image_size(cast(str, value)) if image_size: @@ -109,11 +117,7 @@ class OpenRouterImageEditConfig(BaseImageEditConfig): model: str, api_key: Optional[str] = None, ) -> dict: - api_key = ( - api_key - or litellm.api_key - or get_secret_str("OPENROUTER_API_KEY") - ) + api_key = api_key or litellm.api_key or get_secret_str("OPENROUTER_API_KEY") if not api_key: raise ValueError("OPENROUTER_API_KEY is not set") headers.update( @@ -133,7 +137,11 @@ class OpenRouterImageEditConfig(BaseImageEditConfig): api_base: Optional[str], litellm_params: dict, ) -> str: - base_url = api_base or get_secret_str("OPENROUTER_API_BASE") or "https://openrouter.ai/api/v1" + base_url = ( + api_base + or get_secret_str("OPENROUTER_API_BASE") + or "https://openrouter.ai/api/v1" + ) base_url = base_url.rstrip("/") if not base_url.endswith("/chat/completions"): return f"{base_url}/chat/completions" @@ -162,9 +170,7 @@ class OpenRouterImageEditConfig(BaseImageEditConfig): content_parts.append( { "type": "image_url", - "image_url": { - "url": f"data:{mime_type};base64,{b64_data}" - }, + "image_url": {"url": f"data:{mime_type};base64,{b64_data}"}, } ) @@ -344,7 +350,9 @@ class OpenRouterImageEditConfig(BaseImageEditConfig): if cost_details: if "response_cost_details" not in model_response._hidden_params: model_response._hidden_params["response_cost_details"] = {} - model_response._hidden_params["response_cost_details"].update(cost_details) + model_response._hidden_params["response_cost_details"].update( + cost_details + ) model_response._hidden_params["model"] = response_json.get("model", model) diff --git a/litellm/llms/openrouter/image_generation/__init__.py b/litellm/llms/openrouter/image_generation/__init__.py index f2d06439d4..af5dc036e4 100644 --- a/litellm/llms/openrouter/image_generation/__init__.py +++ b/litellm/llms/openrouter/image_generation/__init__.py @@ -10,4 +10,4 @@ __all__ = [ def get_openrouter_image_generation_config(model: str) -> BaseImageGenerationConfig: - return OpenRouterImageGenerationConfig() \ No newline at end of file + return OpenRouterImageGenerationConfig() diff --git a/litellm/llms/openrouter/image_generation/transformation.py b/litellm/llms/openrouter/image_generation/transformation.py index 92084b533a..a55716a5e5 100644 --- a/litellm/llms/openrouter/image_generation/transformation.py +++ b/litellm/llms/openrouter/image_generation/transformation.py @@ -37,8 +37,16 @@ from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) from litellm.secret_managers.main import get_secret_str -from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams, AllMessageValues -from litellm.types.utils import ImageObject, ImageResponse, ImageUsage, ImageUsageInputTokensDetails +from litellm.types.llms.openai import ( + OpenAIImageGenerationOptionalParams, + AllMessageValues, +) +from litellm.types.utils import ( + ImageObject, + ImageResponse, + ImageUsage, + ImageUsageInputTokensDetails, +) from litellm.llms.openrouter.common_utils import OpenRouterException @@ -51,7 +59,7 @@ else: class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): """ Configuration for OpenRouter image generation via chat completions. - + OpenRouter uses chat completion endpoints for image generation, so we need to transform image generation requests to chat format and extract images from chat responses. @@ -62,7 +70,7 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): ) -> List[OpenAIImageGenerationOptionalParams]: """ Get supported OpenAI parameters for OpenRouter image generation. - + Since OpenRouter uses chat completions for image generation, we support standard image generation params. """ @@ -81,13 +89,13 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): ) -> dict: """ Map image generation params to OpenRouter chat completion format. - + Maps OpenAI parameters to OpenRouter's image_config format: - size -> image_config.aspect_ratio - quality -> image_config.image_size """ supported_params = self.get_supported_openai_params(model) - + for key, value in non_default_params.items(): if key in supported_params: if key == "size": @@ -109,13 +117,13 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): elif not drop_params: # If not supported and drop_params is False, pass through optional_params[key] = value - + return optional_params - + def _map_size_to_aspect_ratio(self, size: str) -> str: """ Map OpenAI size format to OpenRouter aspect_ratio format. - + OpenAI sizes: - 1024x1024 (square) - 1536x1024 (landscape) @@ -124,7 +132,7 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): - 1024x1792 (tall portrait, dall-e-3) - 256x256, 512x512 (dall-e-2) - auto (default) - + OpenRouter aspect_ratios: - 1:1 → 1024×1024 (default) - 2:3 → 832×1248 @@ -152,16 +160,16 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): "auto": "1:1", } return size_to_aspect_ratio.get(size, "1:1") - + def _map_quality_to_image_size(self, quality: str) -> Optional[str]: """ Map OpenAI quality to OpenRouter image_size format. - + OpenAI quality values: - auto (default) - automatically select best quality - high, medium, low - for GPT image models - hd, standard - for dall-e-3 - + OpenRouter image_size values (Gemini only): - 1K → Standard resolution (default) - 2K → Higher resolution @@ -178,7 +186,7 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): "auto": "1K", } return quality_to_image_size.get(quality) - + def _set_usage_and_cost( self, model_response: ImageResponse, @@ -187,7 +195,7 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): ) -> None: """ Extract and set usage and cost information from OpenRouter response. - + Args: model_response: ImageResponse object to populate response_json: Parsed JSON response from OpenRouter @@ -197,10 +205,10 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): if usage_data: prompt_tokens = usage_data.get("prompt_tokens", 0) total_tokens = usage_data.get("total_tokens", 0) - + completion_tokens_details = usage_data.get("completion_tokens_details", {}) image_tokens = completion_tokens_details.get("image_tokens", 0) - + model_response.usage = ImageUsage( input_tokens=prompt_tokens, input_tokens_details=ImageUsageInputTokensDetails( @@ -210,7 +218,7 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): output_tokens=image_tokens, total_tokens=total_tokens, ) - + cost = usage_data.get("cost") if cost is not None: if not hasattr(model_response, "_hidden_params"): @@ -220,13 +228,15 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): model_response._hidden_params["additional_headers"][ "llm_provider-x-litellm-response-cost" ] = float(cost) - + cost_details = usage_data.get("cost_details", {}) if cost_details: if "response_cost_details" not in model_response._hidden_params: model_response._hidden_params["response_cost_details"] = {} - model_response._hidden_params["response_cost_details"].update(cost_details) - + model_response._hidden_params["response_cost_details"].update( + cost_details + ) + model_response._hidden_params["model"] = response_json.get("model", model) def get_complete_url( @@ -240,7 +250,7 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): ) -> str: """ Get the complete URL for OpenRouter image generation. - + OpenRouter uses chat completions endpoint for image generation. Default: https://openrouter.ai/api/v1/chat/completions """ @@ -249,7 +259,7 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): api_base = api_base.rstrip("/") return f"{api_base}/chat/completions" return api_base - + return "https://openrouter.ai/api/v1/chat/completions" def validate_environment( @@ -262,11 +272,7 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): api_key: Optional[str] = None, api_base: Optional[str] = None, ) -> dict: - api_key = ( - api_key - or litellm.api_key - or get_secret_str("OPENROUTER_API_KEY") - ) + api_key = api_key or litellm.api_key or get_secret_str("OPENROUTER_API_KEY") headers.update( { "Authorization": f"Bearer {api_key}", @@ -284,32 +290,27 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): ) -> dict: """ Transform image generation request to OpenRouter chat completion format. - + Args: model: The model name prompt: The image generation prompt optional_params: Optional parameters (including image_config) litellm_params: LiteLLM parameters headers: Request headers - + Returns: dict: Request body in chat completion format with image_config """ request_body = { "model": model, - "messages": [ - { - "role": "user", - "content": prompt - } - ] + "messages": [{"role": "user", "content": prompt}], } - + # These will be passed through to OpenRouter for key, value in optional_params.items(): if key not in ["model", "messages", "modalities"]: request_body[key] = value - + return request_body def transform_image_generation_response( @@ -327,9 +328,9 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): ) -> ImageResponse: """ Transform OpenRouter chat completion response to ImageResponse format. - + Extracts images from the message content and maps usage/cost information. - + Args: model: The model name raw_response: Raw HTTP response from OpenRouter @@ -341,7 +342,7 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): encoding: Encoding api_key: API key json_mode: JSON mode flag - + Returns: ImageResponse: Populated image response """ @@ -353,28 +354,28 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): status_code=raw_response.status_code, headers=raw_response.headers, ) - + if not model_response.data: model_response.data = [] - + try: choices = response_json.get("choices", []) - + for choice in choices: message = choice.get("message", {}) images = message.get("images", []) - + for image_data in images: image_url_obj = image_data.get("image_url", {}) image_url = image_url_obj.get("url") - + if image_url: if image_url.startswith("data:"): # Extract base64 data # Format: data:image/png;base64, parts = image_url.split(",", 1) b64_data = parts[1] if len(parts) > 1 else None - + model_response.data.append( ImageObject( b64_json=b64_data, @@ -390,12 +391,12 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): revised_prompt=None, ) ) - + # Extract and set usage and cost information self._set_usage_and_cost(model_response, response_json, model) - + return model_response - + except Exception as e: raise OpenRouterException( message=f"Error transforming OpenRouter image generation response: {str(e)}", diff --git a/litellm/llms/ovhcloud/audio_transcription/transformation.py b/litellm/llms/ovhcloud/audio_transcription/transformation.py index 7233d911b0..7ff6dc986b 100644 --- a/litellm/llms/ovhcloud/audio_transcription/transformation.py +++ b/litellm/llms/ovhcloud/audio_transcription/transformation.py @@ -31,7 +31,13 @@ class OVHCloudAudioTranscriptionConfig(BaseAudioTranscriptionConfig): ) -> List[OpenAIAudioTranscriptionOptionalParams]: # OVHCloud implements the OpenAI-compatible Whisper interface. # We pass through the same optional params as the OpenAI Whisper API. - return ["language", "prompt", "response_format", "timestamp_granularities", "temperature"] + return [ + "language", + "prompt", + "response_format", + "timestamp_granularities", + "temperature", + ] def map_openai_params( self, @@ -152,5 +158,3 @@ class OVHCloudAudioTranscriptionConfig(BaseAudioTranscriptionConfig): response._hidden_params = response_json return response - - diff --git a/litellm/llms/ovhcloud/chat/transformation.py b/litellm/llms/ovhcloud/chat/transformation.py index e9dc5be3ee..e2a9fea789 100644 --- a/litellm/llms/ovhcloud/chat/transformation.py +++ b/litellm/llms/ovhcloud/chat/transformation.py @@ -15,6 +15,7 @@ from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.openai import AllMessageValues + class OVHCloudChatConfig(OpenAIGPTConfig): @property def custom_llm_provider(self) -> Optional[str]: @@ -45,7 +46,7 @@ class OVHCloudChatConfig(OpenAIGPTConfig): optional_params.remove("function_call") optional_params.remove("response_format") return optional_params - + def get_complete_url( self, api_base: Optional[str], @@ -55,15 +56,16 @@ class OVHCloudChatConfig(OpenAIGPTConfig): litellm_params: dict, stream: Optional[bool] = None, ) -> str: - api_base = "https://oai.endpoints.kepler.ai.cloud.ovh.net/v1" if api_base is None else api_base.rstrip("/") + api_base = ( + "https://oai.endpoints.kepler.ai.cloud.ovh.net/v1" + if api_base is None + else api_base.rstrip("/") + ) complete_url = f"{api_base}/chat/completions" return complete_url - + def get_error_class( - self, - error_message: str, - status_code: int, - headers: Union[dict, httpx.Headers] + self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] ) -> BaseLLMException: return OVHCloudException( message=error_message, @@ -82,7 +84,7 @@ class OVHCloudChatConfig(OpenAIGPTConfig): non_default_params, optional_params, model, drop_params ) return mapped_openai_params - + def transform_request( self, model: str, @@ -98,6 +100,7 @@ class OVHCloudChatConfig(OpenAIGPTConfig): response.update(extra_body) return response + class OVHCloudChatCompletionStreamingHandler(BaseModelResponseIterator): """ Handler for OVHCloud AI Endpoints streaming chat completion responses @@ -122,7 +125,9 @@ class OVHCloudChatCompletionStreamingHandler(BaseModelResponseIterator): new_choices = [] for choice in chunk["choices"]: if "delta" in choice and "reasoning" in choice["delta"]: - choice["delta"]["reasoning_content"] = choice["delta"].get("reasoning") + choice["delta"]["reasoning_content"] = choice["delta"].get( + "reasoning" + ) new_choices.append(choice) return ModelResponseStream( @@ -140,4 +145,4 @@ class OVHCloudChatCompletionStreamingHandler(BaseModelResponseIterator): headers={"Content-Type": "application/json"}, ) except Exception as e: - raise e \ No newline at end of file + raise e diff --git a/litellm/llms/ovhcloud/embedding/transformation.py b/litellm/llms/ovhcloud/embedding/transformation.py index 1266f74c0a..38e88da125 100644 --- a/litellm/llms/ovhcloud/embedding/transformation.py +++ b/litellm/llms/ovhcloud/embedding/transformation.py @@ -29,7 +29,11 @@ class OVHCloudEmbeddingConfig(BaseEmbeddingConfig): litellm_params: dict, stream: Optional[bool] = None, ) -> str: - api_base = "https://oai.endpoints.kepler.ai.cloud.ovh.net/v1" if api_base is None else api_base.rstrip("/") + api_base = ( + "https://oai.endpoints.kepler.ai.cloud.ovh.net/v1" + if api_base is None + else api_base.rstrip("/") + ) complete_url = f"{api_base}/embeddings" return complete_url diff --git a/litellm/llms/ovhcloud/utils.py b/litellm/llms/ovhcloud/utils.py index 9ae4dfb1ef..046df4bca1 100644 --- a/litellm/llms/ovhcloud/utils.py +++ b/litellm/llms/ovhcloud/utils.py @@ -3,4 +3,5 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException class OVHCloudException(BaseLLMException): """OVHCloud AI Endpoints exception handling class""" - pass \ No newline at end of file + + pass diff --git a/litellm/llms/parallel_ai/search/__init__.py b/litellm/llms/parallel_ai/search/__init__.py index cc2ff91ea3..b96914f13d 100644 --- a/litellm/llms/parallel_ai/search/__init__.py +++ b/litellm/llms/parallel_ai/search/__init__.py @@ -4,4 +4,3 @@ Parallel AI Search API module. from litellm.llms.parallel_ai.search.transformation import ParallelAISearchConfig __all__ = ["ParallelAISearchConfig"] - diff --git a/litellm/llms/parallel_ai/search/transformation.py b/litellm/llms/parallel_ai/search/transformation.py index 95919b85c2..e19bc5400d 100644 --- a/litellm/llms/parallel_ai/search/transformation.py +++ b/litellm/llms/parallel_ai/search/transformation.py @@ -18,12 +18,14 @@ from litellm.secret_managers.main import get_secret_str class _ParallelAISourcePolicy(TypedDict, total=False): """Source policy for Parallel AI search results.""" + allowed_domains: List[str] # Optional - list of allowed domains disallowed_domains: List[str] # Optional - list of disallowed domains class _ParallelAISearchRequestRequired(TypedDict): """Required fields for Parallel AI Search API request.""" + # Note: At least one of objective or search_queries must be provided pass @@ -33,6 +35,7 @@ class ParallelAISearchRequest(_ParallelAISearchRequestRequired, total=False): Parallel AI Search API request format. Based on: https://docs.parallel.ai/api-reference/search-and-extract-api-beta/search """ + objective: str # Optional - natural-language description of search goal search_queries: List[str] # Optional - list of keyword search queries processor: str # Optional - search processor ('base', 'pro'), default 'base' @@ -44,11 +47,11 @@ class ParallelAISearchRequest(_ParallelAISearchRequestRequired, total=False): class ParallelAISearchConfig(BaseSearchConfig): PARALLEL_AI_API_BASE = "https://api.parallel.ai" PARALLEL_HEADER_SEARCH_EXTRACT_VALUE = "search-extract-2025-10-10" - + @staticmethod def ui_friendly_name() -> str: return "Parallel AI" - + def validate_environment( self, headers: Dict, @@ -59,9 +62,15 @@ class ParallelAISearchConfig(BaseSearchConfig): """ Validate environment and return headers. """ - api_key = api_key or get_secret_str("PARALLEL_AI_API_KEY") or get_secret_str("PARALLEL_API_KEY") + api_key = ( + api_key + or get_secret_str("PARALLEL_AI_API_KEY") + or get_secret_str("PARALLEL_API_KEY") + ) if not api_key: - raise ValueError("PARALLEL_API_KEY is not set. Set `PARALLEL_API_KEY` environment variable.") + raise ValueError( + "PARALLEL_API_KEY is not set. Set `PARALLEL_API_KEY` environment variable." + ) headers["x-api-key"] = api_key headers["Content-Type"] = "application/json" headers["parallel-beta"] = self.PARALLEL_HEADER_SEARCH_EXTRACT_VALUE @@ -77,8 +86,12 @@ class ParallelAISearchConfig(BaseSearchConfig): """ Get complete URL for Search endpoint. """ - api_base = api_base or get_secret_str("PARALLEL_AI_API_BASE") or self.PARALLEL_AI_API_BASE - + api_base = ( + api_base + or get_secret_str("PARALLEL_AI_API_BASE") + or self.PARALLEL_AI_API_BASE + ) + # Parallel AI search endpoint is at /v1beta/search if not api_base.endswith("/v1beta/search"): if api_base.endswith("/"): @@ -87,7 +100,7 @@ class ParallelAISearchConfig(BaseSearchConfig): api_base = f"{api_base}/v1beta/search" return api_base - + def _transform_query_to_objective(self, query: Union[str, List[str]]) -> str: """ Transform query to objective. @@ -95,7 +108,6 @@ class ParallelAISearchConfig(BaseSearchConfig): if isinstance(query, list): return " ".join(query) return query - def transform_search_request( self, @@ -105,7 +117,7 @@ class ParallelAISearchConfig(BaseSearchConfig): ) -> Dict: """ Transform Search request to Parallel AI API format. - + Args: query: Search query (string or list of strings) - If string: maps to `objective` (natural language) @@ -116,42 +128,45 @@ class ParallelAISearchConfig(BaseSearchConfig): - exclude_domains: List of domains to exclude -> maps to `source_policy.disallowed_domains` - processor: Search processor ('base', 'pro') - max_chars_per_result: Max characters per result excerpt - + Returns: Dict with typed request data following ParallelAISearchRequest spec """ request_data: ParallelAISearchRequest = {} - + # Map query to objective (string or list both become objective) if isinstance(query, list): request_data["objective"] = self._transform_query_to_objective(query) else: request_data["objective"] = query - + # Transform Perplexity unified spec parameters to Parallel AI format if "max_results" in optional_params: request_data["max_results"] = optional_params["max_results"] - + # Map domain filters to source_policy source_policy: _ParallelAISourcePolicy = {} - + if "search_domain_filter" in optional_params: source_policy["allowed_domains"] = optional_params["search_domain_filter"] - + if "exclude_domains" in optional_params: source_policy["disallowed_domains"] = optional_params["exclude_domains"] - + if source_policy: request_data["source_policy"] = source_policy - + # Convert to dict before dynamic key assignments result_data = dict(request_data) - + # pass through all other parameters as-is for param, value in optional_params.items(): - if param not in self.get_supported_perplexity_optional_params() and param not in result_data: + if ( + param not in self.get_supported_perplexity_optional_params() + and param not in result_data + ): result_data[param] = value - + return result_data def transform_search_response( @@ -162,29 +177,29 @@ class ParallelAISearchConfig(BaseSearchConfig): ) -> SearchResponse: """ Transform Parallel AI API response to LiteLLM unified SearchResponse format. - + Parallel AI → LiteLLM mappings: - results[].title → SearchResult.title - results[].url → SearchResult.url - results[].excerpts (array) → SearchResult.snippet (joined string) - No date/last_updated fields in Parallel AI response (set to None) - + Args: raw_response: Raw httpx response from Parallel AI API logging_obj: Logging object for tracking - + Returns: SearchResponse with standardized format """ response_json = raw_response.json() - + # Transform results to SearchResult objects results = [] for result in response_json.get("results", []): # Join excerpts array into a single snippet string excerpts = result.get("excerpts", []) snippet = " ... ".join(excerpts) if excerpts else "" - + search_result = SearchResult( title=result.get("title", ""), url=result.get("url", ""), @@ -193,9 +208,8 @@ class ParallelAISearchConfig(BaseSearchConfig): last_updated=None, # Parallel AI doesn't provide last_updated in response ) results.append(search_result) - + return SearchResponse( results=results, object="search", ) - diff --git a/litellm/llms/perplexity/chat/transformation.py b/litellm/llms/perplexity/chat/transformation.py index 27e6415ff8..48299529ff 100644 --- a/litellm/llms/perplexity/chat/transformation.py +++ b/litellm/llms/perplexity/chat/transformation.py @@ -61,7 +61,7 @@ class PerplexityChatConfig(OpenAIGPTConfig): base_openai_params.append("reasoning_effort") except Exception as e: verbose_logger.debug(f"Error checking if model supports reasoning: {e}") - + try: if litellm.supports_web_search( model=model, custom_llm_provider=self.custom_llm_provider @@ -69,7 +69,7 @@ class PerplexityChatConfig(OpenAIGPTConfig): base_openai_params.append("web_search_options") except Exception as e: verbose_logger.debug(f"Error checking if model supports web search: {e}") - + return base_openai_params def transform_response( @@ -109,7 +109,9 @@ class PerplexityChatConfig(OpenAIGPTConfig): ) self._add_citations_as_annotations(model_response, raw_response_json) except Exception as e: - verbose_logger.debug(f"Error extracting Perplexity-specific usage fields: {e}") + verbose_logger.debug( + f"Error extracting Perplexity-specific usage fields: {e}" + ) return model_response @@ -123,9 +125,7 @@ class PerplexityChatConfig(OpenAIGPTConfig): if not hasattr(model_response, "usage") or model_response.usage is None: # Create a usage object if it doesn't exist (when usage was None) model_response.usage = Usage( # type: ignore[attr-defined] - prompt_tokens=0, - completion_tokens=0, - total_tokens=0 + prompt_tokens=0, completion_tokens=0, total_tokens=0 ) usage = model_response.usage # type: ignore[attr-defined] @@ -146,7 +146,7 @@ class PerplexityChatConfig(OpenAIGPTConfig): # Extract search queries count from usage or response metadata # Perplexity might include this in the usage object or as separate metadata perplexity_usage = raw_response_json.get("usage", {}) - + # Try to extract search queries from usage field first, then root level num_search_queries = perplexity_usage.get("num_search_queries") if num_search_queries is None: @@ -155,18 +155,18 @@ class PerplexityChatConfig(OpenAIGPTConfig): num_search_queries = perplexity_usage.get("search_queries") if num_search_queries is None: num_search_queries = raw_response_json.get("search_queries") - + # Create or update prompt_tokens_details to include web search requests and citation tokens if citation_tokens > 0 or ( num_search_queries is not None and num_search_queries > 0 ): if usage.prompt_tokens_details is None: usage.prompt_tokens_details = PromptTokensDetailsWrapper() - + # Store citation tokens count for cost calculation if citation_tokens > 0: setattr(usage, "citation_tokens", citation_tokens) - + # Store search queries count in the standard web_search_requests field if num_search_queries is not None and num_search_queries > 0: usage.prompt_tokens_details.web_search_requests = num_search_queries @@ -248,4 +248,4 @@ class PerplexityChatConfig(OpenAIGPTConfig): if citations: setattr(model_response, "citations", citations) if search_results: - setattr(model_response, "search_results", search_results) \ No newline at end of file + setattr(model_response, "search_results", search_results) diff --git a/litellm/llms/perplexity/cost_calculator.py b/litellm/llms/perplexity/cost_calculator.py index 463d897901..0f9c3cad84 100644 --- a/litellm/llms/perplexity/cost_calculator.py +++ b/litellm/llms/perplexity/cost_calculator.py @@ -34,7 +34,9 @@ def cost_per_token(model: str, usage: Usage) -> Tuple[float, float]: ## GET MODEL INFO model_info = get_model_info(model=model, custom_llm_provider="perplexity") - def _safe_float_cast(value: Union[str, int, float, None, object], default: float = 0.0) -> float: + def _safe_float_cast( + value: Union[str, int, float, None, object], default: float = 0.0 + ) -> float: """Safely cast a value to float with proper type handling for mypy.""" if value is None: return default @@ -61,9 +63,15 @@ def cost_per_token(model: str, usage: Usage) -> Tuple[float, float]: ## ADD REASONING TOKENS COST (if present) reasoning_tokens = getattr(usage, "reasoning_tokens", 0) or 0 # Also check completion_tokens_details if reasoning_tokens is not directly available - if reasoning_tokens == 0 and hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details: - reasoning_tokens = getattr(usage.completion_tokens_details, "reasoning_tokens", 0) or 0 - + if ( + reasoning_tokens == 0 + and hasattr(usage, "completion_tokens_details") + and usage.completion_tokens_details + ): + reasoning_tokens = ( + getattr(usage.completion_tokens_details, "reasoning_tokens", 0) or 0 + ) + reasoning_cost_value = model_info.get("output_cost_per_reasoning_token") if reasoning_tokens > 0 and reasoning_cost_value is not None: reasoning_cost_per_token = _safe_float_cast(reasoning_cost_value) @@ -72,19 +80,26 @@ def cost_per_token(model: str, usage: Usage) -> Tuple[float, float]: ## ADD SEARCH QUERIES COST (if present) num_search_queries = 0 if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details: - num_search_queries = getattr(usage.prompt_tokens_details, "web_search_requests", 0) or 0 - + num_search_queries = ( + getattr(usage.prompt_tokens_details, "web_search_requests", 0) or 0 + ) + # Check both possible keys for search cost (legacy and current) - search_cost_value = model_info.get("search_queries_cost_per_query") or model_info.get("search_context_cost_per_query") + search_cost_value = model_info.get( + "search_queries_cost_per_query" + ) or model_info.get("search_context_cost_per_query") if num_search_queries > 0 and search_cost_value is not None: # Handle both dict and float formats if isinstance(search_cost_value, dict): # Use the "low" size as default - tests expect 0.005 / 1000 - search_cost_per_query = _safe_float_cast(search_cost_value.get("search_context_size_low", 0)) / 1000 + search_cost_per_query = ( + _safe_float_cast(search_cost_value.get("search_context_size_low", 0)) + / 1000 + ) else: search_cost_per_query = _safe_float_cast(search_cost_value) search_cost = num_search_queries * search_cost_per_query # Add search cost to completion cost (similar to how other providers handle it) completion_cost += search_cost - return prompt_cost, completion_cost \ No newline at end of file + return prompt_cost, completion_cost diff --git a/litellm/llms/perplexity/responses/transformation.py b/litellm/llms/perplexity/responses/transformation.py index f365ef07a6..cacdcdb9d7 100644 --- a/litellm/llms/perplexity/responses/transformation.py +++ b/litellm/llms/perplexity/responses/transformation.py @@ -23,7 +23,6 @@ from litellm.types.utils import LlmProviders class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): - def get_supported_openai_params(self, model: str) -> list: """Ref: https://docs.perplexity.ai/api-reference/responses-post""" return [ @@ -55,7 +54,11 @@ class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): return headers def get_complete_url(self, api_base: Optional[str], litellm_params: dict) -> str: - api_base = api_base or get_secret_str("PERPLEXITY_API_BASE") or "https://api.perplexity.ai" + api_base = ( + api_base + or get_secret_str("PERPLEXITY_API_BASE") + or "https://api.perplexity.ai" + ) return f"{api_base.rstrip('/')}/v1/responses" def _ensure_message_type( @@ -86,7 +89,7 @@ class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): if model.startswith("preset/"): input = self._validate_input_param(input) data: Dict = { - "preset": model[len("preset/"):], + "preset": model[len("preset/") :], "input": input, } data.update(response_api_optional_request_params) diff --git a/litellm/llms/perplexity/search/transformation.py b/litellm/llms/perplexity/search/transformation.py index f1dc0909b4..f89d556549 100644 --- a/litellm/llms/perplexity/search/transformation.py +++ b/litellm/llms/perplexity/search/transformation.py @@ -16,6 +16,7 @@ from litellm.secret_managers.main import get_secret_str class _PerplexitySearchRequestRequired(TypedDict): """Required fields for Perplexity Search API request.""" + query: Union[str, List[str]] # Required - search query or queries @@ -24,6 +25,7 @@ class PerplexitySearchRequest(_PerplexitySearchRequestRequired, total=False): Perplexity Search API request format. Based on: https://docs.perplexity.ai/api-reference/search-post """ + max_results: int # Optional - maximum number of results (1-20), default 10 search_domain_filter: List[str] # Optional - list of domains to filter (max 20) max_tokens_per_page: int # Optional - max tokens per page, default 1024 @@ -32,11 +34,11 @@ class PerplexitySearchRequest(_PerplexitySearchRequestRequired, total=False): class PerplexitySearchConfig(BaseSearchConfig): PERPLEXITY_API_BASE = "https://api.perplexity.ai" - + @staticmethod def ui_friendly_name() -> str: return "Perplexity" - + def validate_environment( self, headers: Dict, @@ -49,7 +51,9 @@ class PerplexitySearchConfig(BaseSearchConfig): """ api_key = api_key or get_secret_str("PERPLEXITYAI_API_KEY") if not api_key: - raise ValueError("PERPLEXITYAI_API_KEY is not set. Set `PERPLEXITYAI_API_KEY` environment variable.") + raise ValueError( + "PERPLEXITYAI_API_KEY is not set. Set `PERPLEXITYAI_API_KEY` environment variable." + ) headers["Authorization"] = f"Bearer {api_key}" headers["Content-Type"] = "application/json" return headers @@ -64,14 +68,17 @@ class PerplexitySearchConfig(BaseSearchConfig): """ Get complete URL for Search endpoint. """ - api_base = api_base or get_secret_str("PERPLEXITY_API_BASE") or self.PERPLEXITY_API_BASE - + api_base = ( + api_base + or get_secret_str("PERPLEXITY_API_BASE") + or self.PERPLEXITY_API_BASE + ) + # append "/search" to the api base if it's not already there if not api_base.endswith("/search"): api_base = f"{api_base}/search" return api_base - def transform_search_request( self, @@ -85,9 +92,9 @@ class PerplexitySearchConfig(BaseSearchConfig): Note: LiteLLM's native spec is the perplexity search spec. There's no transformation needed for the request data. - + https://docs.perplexity.ai/api-reference/search-post - + Args: query: Search query (string or list of strings) optional_params: Optional parameters for the request @@ -95,31 +102,31 @@ class PerplexitySearchConfig(BaseSearchConfig): - search_domain_filter: List of domains to filter (max 20) - max_tokens_per_page: Max tokens per page (default 1024) - country: Country code filter (e.g., 'US', 'GB', 'DE') - + Returns: Dict with typed request data following PerplexitySearchRequest spec """ request_data: PerplexitySearchRequest = { "query": query, } - + # Add optional parameters following Perplexity API spec (only if not None) max_results = optional_params.get("max_results") if max_results is not None: request_data["max_results"] = max_results - + search_domain_filter = optional_params.get("search_domain_filter") if search_domain_filter is not None: request_data["search_domain_filter"] = search_domain_filter - + max_tokens_per_page = optional_params.get("max_tokens_per_page") if max_tokens_per_page is not None: request_data["max_tokens_per_page"] = max_tokens_per_page - + country = optional_params.get("country") if country is not None: request_data["country"] = country - + return dict(request_data) def transform_search_response( @@ -130,16 +137,16 @@ class PerplexitySearchConfig(BaseSearchConfig): ) -> SearchResponse: """ Transform Perplexity API response to standard SearchResponse format. - + Args: raw_response: Raw httpx response from Perplexity API logging_obj: Logging object for tracking - + Returns: SearchResponse with standardized format """ response_json = raw_response.json() - + # Transform results to SearchResult objects results = [] for result in response_json.get("results", []): @@ -151,9 +158,8 @@ class PerplexitySearchConfig(BaseSearchConfig): last_updated=result.get("last_updated"), ) results.append(search_result) - + return SearchResponse( results=results, object="search", ) - diff --git a/litellm/llms/pg_vector/vector_stores/transformation.py b/litellm/llms/pg_vector/vector_stores/transformation.py index 5d10faeba5..ba87a8f2b0 100644 --- a/litellm/llms/pg_vector/vector_stores/transformation.py +++ b/litellm/llms/pg_vector/vector_stores/transformation.py @@ -10,6 +10,7 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any + class PGVectorStoreConfig(OpenAIVectorStoreConfig): """ PG Vector Store configuration that inherits from OpenAI since it's OpenAI-compatible. @@ -19,7 +20,7 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig): https://github.com/BerriAI/litellm-pgvector You just need to connect litellm proxy to this deployed server. - + Requires: - api_base: The base URL for the PG vector service - api_key: API key for authentication with the PG vector service @@ -32,16 +33,15 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig): Validate environment and set headers for PG vector service authentication """ litellm_params = litellm_params or GenericLiteLLMParams() - + # Get API key from various sources - api_key = ( - litellm_params.api_key - or get_secret_str("PG_VECTOR_API_KEY") - ) - + api_key = litellm_params.api_key or get_secret_str("PG_VECTOR_API_KEY") + if not api_key: - raise ValueError("PG Vector API key is required. Set PG_VECTOR_API_KEY environment variable or pass api_key in litellm_params.") - + raise ValueError( + "PG Vector API key is required. Set PG_VECTOR_API_KEY environment variable or pass api_key in litellm_params." + ) + headers.update( { "Authorization": f"Bearer {api_key}", @@ -60,19 +60,17 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig): Get the complete URL for PG vector service endpoints """ # Get API base from various sources - api_base = ( - api_base - or get_secret_str("PG_VECTOR_API_BASE") - ) - + api_base = api_base or get_secret_str("PG_VECTOR_API_BASE") + if not api_base: - raise ValueError("PG Vector API base URL is required. Set PG_VECTOR_API_BASE environment variable or pass api_base in litellm_params.") + raise ValueError( + "PG Vector API base URL is required. Set PG_VECTOR_API_BASE environment variable or pass api_base in litellm_params." + ) # Remove trailing slashes api_base = api_base.rstrip("/") - return f"{api_base}/v1/vector_stores" - + return f"{api_base}/v1/vector_stores" def transform_search_vector_store_request( self, @@ -83,7 +81,7 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, ) -> Tuple[str, Dict]: - url = f"{api_base}/{vector_store_id}/search" + url = f"{api_base}/{vector_store_id}/search" _, request_body = super().transform_search_vector_store_request( vector_store_id=vector_store_id, query=query, @@ -92,4 +90,4 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig): litellm_logging_obj=litellm_logging_obj, litellm_params=litellm_params, ) - return url, request_body \ No newline at end of file + return url, request_body diff --git a/litellm/llms/ragflow/__init__.py b/litellm/llms/ragflow/__init__.py index 17d12bed31..3ca54e3855 100644 --- a/litellm/llms/ragflow/__init__.py +++ b/litellm/llms/ragflow/__init__.py @@ -5,4 +5,3 @@ RAGFlow provides OpenAI-compatible APIs with unique path structures: - Chat endpoint: /api/v1/chats_openai/{chat_id}/chat/completions - Agent endpoint: /api/v1/agents_openai/{agent_id}/chat/completions """ - diff --git a/litellm/llms/ragflow/chat/__init__.py b/litellm/llms/ragflow/chat/__init__.py index 0e0f47d07b..4f84cce42b 100644 --- a/litellm/llms/ragflow/chat/__init__.py +++ b/litellm/llms/ragflow/chat/__init__.py @@ -1,4 +1,3 @@ """ RAGFlow chat completion configuration. """ - diff --git a/litellm/llms/ragflow/chat/transformation.py b/litellm/llms/ragflow/chat/transformation.py index 58fbfa83c9..d49a5fd370 100644 --- a/litellm/llms/ragflow/chat/transformation.py +++ b/litellm/llms/ragflow/chat/transformation.py @@ -21,7 +21,7 @@ from litellm.types.llms.openai import AllMessageValues class RAGFlowConfig(OpenAIConfig): """ Configuration for RAGFlow OpenAI-compatible API. - + Handles both chat and agent endpoints by parsing the model name format: - ragflow/chat/{chat_id}/{model_name} for chat endpoints - ragflow/agent/{agent_id}/{model_name} for agent endpoints @@ -30,13 +30,13 @@ class RAGFlowConfig(OpenAIConfig): def _parse_ragflow_model(self, model: str) -> Tuple[str, str, str]: """ Parse RAGFlow model name format: ragflow/{endpoint_type}/{id}/{model_name} - + Args: model: Model name in format ragflow/chat/{chat_id}/{model} or ragflow/agent/{agent_id}/{model} - + Returns: Tuple of (endpoint_type, id, model_name) - + Raises: ValueError: If model format is invalid """ @@ -46,21 +46,23 @@ class RAGFlowConfig(OpenAIConfig): f"Invalid RAGFlow model format: {model}. " f"Expected format: ragflow/chat/{{chat_id}}/{{model}} or ragflow/agent/{{agent_id}}/{{model}}" ) - + if parts[0] != "ragflow": raise ValueError( f"Invalid RAGFlow model format: {model}. Must start with 'ragflow/'" ) - + endpoint_type = parts[1] if endpoint_type not in ["chat", "agent"]: raise ValueError( f"Invalid RAGFlow endpoint type: {endpoint_type}. Must be 'chat' or 'agent'" ) - + entity_id = parts[2] - model_name = "/".join(parts[3:]) # Handle model names that might contain slashes - + model_name = "/".join( + parts[3:] + ) # Handle model names that might contain slashes + return endpoint_type, entity_id, model_name def get_complete_url( @@ -74,11 +76,11 @@ class RAGFlowConfig(OpenAIConfig): ) -> str: """ Get the complete URL for the RAGFlow API call. - + Constructs URL based on endpoint type: - Chat: /api/v1/chats_openai/{chat_id}/chat/completions - Agent: /api/v1/agents_openai/{agent_id}/chat/completions - + Args: api_base: Base API URL (e.g., http://ragflow-server:port or http://ragflow-server:port/v1) api_key: API key (not used in URL construction) @@ -86,47 +88,53 @@ class RAGFlowConfig(OpenAIConfig): optional_params: Optional parameters litellm_params: LiteLLM parameters (may contain api_base) stream: Whether streaming is enabled - + Returns: Complete URL for the API call """ # Get api_base from multiple sources: input param, litellm_params, environment, or global litellm setting - if litellm_params and hasattr(litellm_params, 'api_base') and litellm_params.api_base: + if ( + litellm_params + and hasattr(litellm_params, "api_base") + and litellm_params.api_base + ): api_base = api_base or litellm_params.api_base - + api_base = ( api_base or litellm.api_base or get_secret("RAGFLOW_API_BASE") or get_secret_str("RAGFLOW_API_BASE") ) - + if api_base is None: - raise ValueError("api_base is required for RAGFlow provider. Set it via api_base parameter, RAGFLOW_API_BASE environment variable, or litellm.api_base") - + raise ValueError( + "api_base is required for RAGFlow provider. Set it via api_base parameter, RAGFLOW_API_BASE environment variable, or litellm.api_base" + ) + # Parse model name to extract endpoint type and ID endpoint_type, entity_id, _ = self._parse_ragflow_model(model) - + # Remove trailing slash from api_base if present api_base = api_base.rstrip("/") - + # Strip /v1 or /api/v1 from api_base if present, since we'll add the full path # Check /api/v1 first because /api/v1 ends with /v1 if api_base.endswith("/api/v1"): api_base = api_base[:-7] # Remove /api/v1 elif api_base.endswith("/v1"): api_base = api_base[:-3] # Remove /v1 - + # Construct the RAGFlow-specific path if endpoint_type == "chat": path = f"/api/v1/chats_openai/{entity_id}/chat/completions" else: # agent path = f"/api/v1/agents_openai/{entity_id}/chat/completions" - + # Ensure path starts with / if not path.startswith("/"): path = "/" + path - + return f"{api_base}{path}" def _get_openai_compatible_provider_info( @@ -138,20 +146,20 @@ class RAGFlowConfig(OpenAIConfig): ) -> Tuple[Optional[str], Optional[str], str]: """ Get OpenAI-compatible provider information for RAGFlow. - + Args: model: Model name (will be parsed to extract actual model name) api_base: Base API URL (from input params) api_key: API key (from input params) custom_llm_provider: Custom LLM provider name - + Returns: Tuple of (api_base, api_key, custom_llm_provider) """ # Parse model to extract the actual model name # The model name will be stored in litellm_params for use in requests _, _, actual_model = self._parse_ragflow_model(model) - + # Get api_base from multiple sources: input param, environment, or global litellm setting dynamic_api_base = ( api_base @@ -159,14 +167,12 @@ class RAGFlowConfig(OpenAIConfig): or get_secret("RAGFLOW_API_BASE") or get_secret_str("RAGFLOW_API_BASE") ) - + # Get api_key from multiple sources: input param, environment, or global litellm setting dynamic_api_key = ( - api_key - or litellm.api_key - or get_secret_str("RAGFLOW_API_KEY") + api_key or litellm.api_key or get_secret_str("RAGFLOW_API_KEY") ) - + return dynamic_api_base, dynamic_api_key, custom_llm_provider def validate_environment( @@ -181,7 +187,7 @@ class RAGFlowConfig(OpenAIConfig): ) -> dict: """ Validate environment and set up headers for RAGFlow API. - + Args: headers: Request headers model: Model name @@ -190,28 +196,28 @@ class RAGFlowConfig(OpenAIConfig): litellm_params: LiteLLM parameters (may contain api_key) api_key: API key (from input params) api_base: Base API URL - + Returns: Updated headers dictionary """ # Use api_key from litellm_params if available, otherwise fall back to other sources - if litellm_params and hasattr(litellm_params, 'api_key') and litellm_params.api_key: + if ( + litellm_params + and hasattr(litellm_params, "api_key") + and litellm_params.api_key + ): api_key = api_key or litellm_params.api_key - + # Get api_key from multiple sources: input param, litellm_params, environment, or global litellm setting - api_key = ( - api_key - or litellm.api_key - or get_secret_str("RAGFLOW_API_KEY") - ) - + api_key = api_key or litellm.api_key or get_secret_str("RAGFLOW_API_KEY") + if api_key is not None: headers["Authorization"] = f"Bearer {api_key}" - + # Ensure Content-Type is set to application/json if "content-type" not in headers and "Content-Type" not in headers: headers["Content-Type"] = "application/json" - + # Parse model to extract actual model name and store it # The actual model name should be used in the request body try: @@ -221,7 +227,7 @@ class RAGFlowConfig(OpenAIConfig): except ValueError: # If parsing fails, use the original model name pass - + return headers def transform_request( @@ -234,16 +240,16 @@ class RAGFlowConfig(OpenAIConfig): ) -> dict: """ Transform request for RAGFlow API. - + Uses the actual model name extracted from the RAGFlow model format. - + Args: model: Model name in RAGFlow format messages: Chat messages optional_params: Optional parameters litellm_params: LiteLLM parameters (may contain _ragflow_actual_model) headers: Request headers - + Returns: Transformed request dictionary """ @@ -256,9 +262,8 @@ class RAGFlowConfig(OpenAIConfig): except ValueError: # If parsing fails, use the original model name actual_model = model - + # Use parent's transform_request with the actual model name return super().transform_request( actual_model, messages, optional_params, litellm_params, headers ) - diff --git a/litellm/llms/ragflow/vector_stores/__init__.py b/litellm/llms/ragflow/vector_stores/__init__.py index 3be29310b3..f36e35f168 100644 --- a/litellm/llms/ragflow/vector_stores/__init__.py +++ b/litellm/llms/ragflow/vector_stores/__init__.py @@ -1,2 +1 @@ # RAGFlow vector stores module - diff --git a/litellm/llms/ragflow/vector_stores/transformation.py b/litellm/llms/ragflow/vector_stores/transformation.py index b6401a4b8d..ed5397eef0 100644 --- a/litellm/llms/ragflow/vector_stores/transformation.py +++ b/litellm/llms/ragflow/vector_stores/transformation.py @@ -32,7 +32,9 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): # Try to get from environment variable api_key = get_secret_str("RAGFLOW_API_KEY") if api_key is None: - raise ValueError("api_key is required (set RAGFLOW_API_KEY env var or pass in litellm_params)") + raise ValueError( + "api_key is required (set RAGFLOW_API_KEY env var or pass in litellm_params)" + ) return { "headers": { "Authorization": f"Bearer {api_key}", @@ -51,14 +53,13 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): ) -> dict: """Validate environment and set headers for RAGFlow API.""" litellm_params = litellm_params or GenericLiteLLMParams() - api_key = ( - litellm_params.api_key - or get_secret_str("RAGFLOW_API_KEY") - ) - + api_key = litellm_params.api_key or get_secret_str("RAGFLOW_API_KEY") + if api_key is None: - raise ValueError("RAGFLOW_API_KEY is required (set env var or pass in litellm_params)") - + raise ValueError( + "RAGFLOW_API_KEY is required (set env var or pass in litellm_params)" + ) + headers.update( { "Authorization": f"Bearer {api_key}", @@ -74,7 +75,7 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): ) -> str: """ Get the complete URL for RAGFlow datasets API. - + Supports: - RAGFLOW_API_BASE env var - api_base in litellm_params @@ -122,22 +123,22 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): ) -> Tuple[str, Dict]: """ Transform create request to RAGFlow POST /api/v1/datasets format. - + Maps LiteLLM params to RAGFlow dataset creation parameters. RAGFlow-specific fields can be passed via metadata. """ url = api_base # Already includes /api/v1/datasets from get_complete_url - + # Extract name (required by RAGFlow) name = vector_store_create_optional_params.get("name") if not name: raise ValueError("name is required for RAGFlow dataset creation") - + # Build request body request_body: Dict[str, Any] = { "name": name, } - + # Extract RAGFlow-specific fields from metadata metadata = vector_store_create_optional_params.get("metadata") if metadata: @@ -152,22 +153,22 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): "parse_type", "pipeline_id", ] - + for field in ragflow_fields: if field in metadata: request_body[field] = metadata[field] - + # Validate: chunk_method and pipeline_id are mutually exclusive if "chunk_method" in request_body and "pipeline_id" in request_body: raise ValueError( "chunk_method and pipeline_id are mutually exclusive. " "Specify either chunk_method or pipeline_id, not both." ) - + # If neither chunk_method nor pipeline_id is specified, default to naive if "chunk_method" not in request_body and "pipeline_id" not in request_body: request_body["chunk_method"] = "naive" - + return url, request_body def transform_create_vector_store_response( @@ -175,7 +176,7 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): ) -> VectorStoreCreateResponse: """ Transform RAGFlow response to VectorStoreCreateResponse format. - + RAGFlow response format: { "code": 0, @@ -189,7 +190,7 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): """ try: response_json = response.json() - + # Check for RAGFlow error response if response_json.get("code") != 0: error_message = response_json.get("message", "Unknown error") @@ -198,21 +199,21 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): status_code=response.status_code, headers=response.headers, ) - + data = response_json.get("data", {}) - + # Extract dataset ID dataset_id = data.get("id") if not dataset_id: raise ValueError("RAGFlow response missing dataset id") - + # Extract name name = data.get("name") - + # Convert create_time from milliseconds to seconds (Unix timestamp) create_time_ms = data.get("create_time", 0) created_at = int(create_time_ms / 1000) if create_time_ms else None - + # Build VectorStoreCreateResponse return VectorStoreCreateResponse( id=dataset_id, @@ -246,4 +247,3 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): status_code=response.status_code, headers=response.headers, ) - diff --git a/litellm/llms/recraft/cost_calculator.py b/litellm/llms/recraft/cost_calculator.py index 5ab47e9395..27b9108e5f 100644 --- a/litellm/llms/recraft/cost_calculator.py +++ b/litellm/llms/recraft/cost_calculator.py @@ -22,4 +22,6 @@ def cost_calculator( num_images = len(image_response.data) return output_cost_per_image * num_images else: - raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}") + raise ValueError( + f"image_response must be of type ImageResponse got type={type(image_response)}" + ) diff --git a/litellm/llms/recraft/image_edit/transformation.py b/litellm/llms/recraft/image_edit/transformation.py index d2a5623681..4c199bc78d 100644 --- a/litellm/llms/recraft/image_edit/transformation.py +++ b/litellm/llms/recraft/image_edit/transformation.py @@ -25,19 +25,17 @@ class RecraftImageEditConfig(BaseImageEditConfig): DEFAULT_BASE_URL: str = "https://external.api.recraft.ai" IMAGE_EDIT_ENDPOINT: str = "v1/images/imageToImage" DEFAULT_STRENGTH: float = 0.2 - - def get_supported_openai_params( - self, model: str - ) -> List: + + def get_supported_openai_params(self, model: str) -> List: """ Supported OpenAI parameters that can be mapped to Recraft image edit API. - + Based on Recraft API docs: https://www.recraft.ai/docs#image-to-image """ return [ - "n", # Maps to n (number of images) - "response_format", # Maps to response_format (url or b64_json) - "style" # Maps to style parameter + "n", # Maps to n (number of images) + "response_format", # Maps to response_format (url or b64_json) + "style", # Maps to style parameter ] def map_openai_params( @@ -52,14 +50,13 @@ class RecraftImageEditConfig(BaseImageEditConfig): """ # Start with all params like OpenAI does all_params = dict(image_edit_optional_params) - + # Filter to only supported Recraft parameters supported_params = self.get_supported_openai_params(model) filtered_params = {k: v for k, v in all_params.items() if k in supported_params} - + return filtered_params - def get_complete_url( self, model: str, @@ -72,9 +69,7 @@ class RecraftImageEditConfig(BaseImageEditConfig): Some providers need `model` in `api_base` """ complete_url: str = ( - api_base - or get_secret_str("RECRAFT_API_BASE") - or self.DEFAULT_BASE_URL + api_base or get_secret_str("RECRAFT_API_BASE") or self.DEFAULT_BASE_URL ) complete_url = complete_url.rstrip("/") @@ -87,16 +82,12 @@ class RecraftImageEditConfig(BaseImageEditConfig): model: str, api_key: Optional[str] = None, ) -> dict: - final_api_key: Optional[str] = ( - api_key or - get_secret_str("RECRAFT_API_KEY") - ) + final_api_key: Optional[str] = api_key or get_secret_str("RECRAFT_API_KEY") if not final_api_key: raise ValueError("RECRAFT_API_KEY is not set") - - headers["Authorization"] = f"Bearer {final_api_key}" - return headers + headers["Authorization"] = f"Bearer {final_api_key}" + return headers def transform_image_edit_request( self, @@ -113,32 +104,35 @@ class RecraftImageEditConfig(BaseImageEditConfig): https://www.recraft.ai/docs#image-to-image """ - + request_params = { "model": model, - "strength": image_edit_optional_request_params.pop("strength", self.DEFAULT_STRENGTH), + "strength": image_edit_optional_request_params.pop( + "strength", self.DEFAULT_STRENGTH + ), **image_edit_optional_request_params, } if prompt is not None: request_params["prompt"] = prompt - + request_body = RecraftImageEditRequestParams(**request_params) request_dict = cast(Dict, request_body) ######################################################### # Reuse OpenAI logic: Separate images as `files` and send other parameters as `data` ######################################################### - files_list = self._get_image_files_for_request(image=image) if image is not None else [] + files_list = ( + self._get_image_files_for_request(image=image) if image is not None else [] + ) data_without_images = {k: v for k, v in request_dict.items() if k != "image"} - + return data_without_images, files_list - def _get_image_files_for_request( self, image: Optional[FileTypes], ) -> List[Tuple[str, Any]]: files_list: List[Tuple[str, Any]] = [] - + # Handle single image (Recraft expects single image, not array) if image: # OpenAI wraps images in arrays, but for Recraft we need single image @@ -146,9 +140,11 @@ class RecraftImageEditConfig(BaseImageEditConfig): _image = image[0] if image else None # Take first image for Recraft else: _image = image - + if _image is not None: - image_content_type: str = ImageEditRequestUtils.get_image_content_type(_image) + image_content_type: str = ImageEditRequestUtils.get_image_content_type( + _image + ) if isinstance(_image, BufferedReader): files_list.append( ("image", (_image.name, _image, image_content_type)) @@ -159,7 +155,7 @@ class RecraftImageEditConfig(BaseImageEditConfig): ) return files_list - + def transform_image_edit_response( self, model: str, @@ -177,11 +173,13 @@ class RecraftImageEditConfig(BaseImageEditConfig): ) if not model_response.data: model_response.data = [] - + for image_data in response_data["data"]: - model_response.data.append(ImageObject( - url=image_data.get("url", None), - b64_json=image_data.get("b64_json", None), - )) - - return model_response \ No newline at end of file + model_response.data.append( + ImageObject( + url=image_data.get("url", None), + b64_json=image_data.get("b64_json", None), + ) + ) + + return model_response diff --git a/litellm/llms/recraft/image_generation/transformation.py b/litellm/llms/recraft/image_generation/transformation.py index f632b49f3a..4a00512dfb 100644 --- a/litellm/llms/recraft/image_generation/transformation.py +++ b/litellm/llms/recraft/image_generation/transformation.py @@ -24,20 +24,15 @@ else: class RecraftImageGenerationConfig(BaseImageGenerationConfig): DEFAULT_BASE_URL: str = "https://external.api.recraft.ai" IMAGE_GENERATION_ENDPOINT: str = "v1/images/generations" - + def get_supported_openai_params( self, model: str ) -> List[OpenAIImageGenerationOptionalParams]: """ https://www.recraft.ai/docs#generate-image """ - return [ - "n", - "response_format", - "size", - "style" - ] - + return ["n", "response_format", "size", "style"] + def map_openai_params( self, non_default_params: dict, @@ -74,9 +69,7 @@ class RecraftImageGenerationConfig(BaseImageGenerationConfig): Some providers need `model` in `api_base` """ complete_url: str = ( - api_base - or get_secret_str("RECRAFT_API_BASE") - or self.DEFAULT_BASE_URL + api_base or get_secret_str("RECRAFT_API_BASE") or self.DEFAULT_BASE_URL ) complete_url = complete_url.rstrip("/") @@ -93,18 +86,13 @@ class RecraftImageGenerationConfig(BaseImageGenerationConfig): api_key: Optional[str] = None, api_base: Optional[str] = None, ) -> dict: - final_api_key: Optional[str] = ( - api_key or - get_secret_str("RECRAFT_API_KEY") - ) + final_api_key: Optional[str] = api_key or get_secret_str("RECRAFT_API_KEY") if not final_api_key: raise ValueError("RECRAFT_API_KEY is not set") - - headers["Authorization"] = f"Bearer {final_api_key}" + + headers["Authorization"] = f"Bearer {final_api_key}" return headers - - def transform_image_generation_request( self, model: str, @@ -118,10 +106,12 @@ class RecraftImageGenerationConfig(BaseImageGenerationConfig): https://www.recraft.ai/docs#generate-image """ - recratft_image_generation_request_body: RecraftImageGenerationRequestParams = RecraftImageGenerationRequestParams( - prompt=prompt, - model=model, - **optional_params, + recratft_image_generation_request_body: RecraftImageGenerationRequestParams = ( + RecraftImageGenerationRequestParams( + prompt=prompt, + model=model, + **optional_params, + ) ) return dict(recratft_image_generation_request_body) @@ -153,11 +143,13 @@ class RecraftImageGenerationConfig(BaseImageGenerationConfig): ) if not model_response.data: model_response.data = [] - + for image_data in response_data["data"]: - model_response.data.append(ImageObject( - url=image_data.get("url", None), - b64_json=image_data.get("b64_json", None), - )) - - return model_response \ No newline at end of file + model_response.data.append( + ImageObject( + url=image_data.get("url", None), + b64_json=image_data.get("b64_json", None), + ) + ) + + return model_response diff --git a/litellm/llms/replicate/chat/handler.py b/litellm/llms/replicate/chat/handler.py index c37473b318..cc4c61e397 100644 --- a/litellm/llms/replicate/chat/handler.py +++ b/litellm/llms/replicate/chat/handler.py @@ -88,7 +88,9 @@ async def async_handle_prediction_response_streaming( response_data = response.json() status = response_data.get("status", "") # Check that "output" exists and is not None or empty - output_present = "output" in response_data and response_data["output"] is not None + output_present = ( + "output" in response_data and response_data["output"] is not None + ) if output_present: try: # If output is None or not a list, treat as empty string @@ -219,10 +221,10 @@ def completion( litellm.DEFAULT_REPLICATE_POLLING_DELAY_SECONDS + 2 * retry ) # wait to allow response to be generated by replicate - else partial output is generated with status=="processing" response = httpx_client.get(url=prediction_url, headers=headers) - if ( - response.status_code == 200 - and response.json().get("status") in ["processing", "starting"] - ): + if response.status_code == 200 and response.json().get("status") in [ + "processing", + "starting", + ]: continue return litellm.ReplicateConfig().transform_response( model=model, @@ -290,10 +292,10 @@ async def async_completion( litellm.DEFAULT_REPLICATE_POLLING_DELAY_SECONDS + 2 * retry ) # wait to allow response to be generated by replicate - else partial output is generated with status=="processing" response = await async_handler.get(url=prediction_url, headers=headers) - if ( - response.status_code == 200 - and response.json().get("status") in ["processing", "starting"] - ): + if response.status_code == 200 and response.json().get("status") in [ + "processing", + "starting", + ]: continue return litellm.ReplicateConfig().transform_response( model=model, diff --git a/litellm/llms/runwayml/cost_calculator.py b/litellm/llms/runwayml/cost_calculator.py index fa3cd26d08..35b6086f19 100644 --- a/litellm/llms/runwayml/cost_calculator.py +++ b/litellm/llms/runwayml/cost_calculator.py @@ -10,7 +10,7 @@ def cost_calculator( ) -> float: """ RunwayML image generation cost calculator. - + RunwayML charges per image generated, not per pixel. Pricing is stored in model_prices_and_context_window.json with output_cost_per_image. """ @@ -28,4 +28,3 @@ def cost_calculator( raise ValueError( f"image_response must be of type ImageResponse, got type={type(image_response)}" ) - diff --git a/litellm/llms/runwayml/image_generation/transformation.py b/litellm/llms/runwayml/image_generation/transformation.py index e92ffa8e9c..448dcd4a67 100644 --- a/litellm/llms/runwayml/image_generation/transformation.py +++ b/litellm/llms/runwayml/image_generation/transformation.py @@ -31,6 +31,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): """ Configuration for RunwayML image generation models. """ + DEFAULT_BASE_URL: str = "https://api.dev.runwayml.com" IMAGE_GENERATION_ENDPOINT: str = "v1/text_to_image" @@ -49,9 +50,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): Some providers need `model` in `api_base` """ complete_url: str = ( - api_base - or get_secret_str("RUNWAYML_API_BASE") - or self.DEFAULT_BASE_URL + api_base or get_secret_str("RUNWAYML_API_BASE") or self.DEFAULT_BASE_URL ) complete_url = complete_url.rstrip("/") @@ -70,14 +69,14 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): api_base: Optional[str] = None, ) -> dict: final_api_key: Optional[str] = ( - api_key or - get_secret_str("RUNWAYML_API_SECRET") or - get_secret_str("RUNWAYML_API_KEY") + api_key + or get_secret_str("RUNWAYML_API_SECRET") + or get_secret_str("RUNWAYML_API_KEY") ) if not final_api_key: raise ValueError("RUNWAYML_API_SECRET or RUNWAYML_API_KEY is not set") - - headers["Authorization"] = f"Bearer {final_api_key}" + + headers["Authorization"] = f"Bearer {final_api_key}" headers["X-Runway-Version"] = RUNWAYML_DEFAULT_API_VERSION return headers @@ -88,7 +87,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): ) -> ImageResponse: """ Transform RunwayML response format to OpenAI ImageResponse format. - + RunwayML response format (after polling): { "id": "task_123...", @@ -96,7 +95,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): "output": ["https://cloudfront.net/.../image.png"], "completedAt": "2025-11-13T..." } - + OpenAI ImageResponse format: { "data": [ @@ -106,47 +105,51 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): } ] } - + Args: response_data: JSON response from RunwayML (after polling completes) model_response: ImageResponse object to populate - + Returns: Populated ImageResponse in OpenAI format """ if not model_response.data: model_response.data = [] - + # Handle RunwayML response format # Response contains task.output with image URL(s) output = response_data.get("output", []) - + if isinstance(output, list): for image_item in output: if isinstance(image_item, str): # If output is a list of URL strings - model_response.data.append(ImageObject( - url=image_item, - b64_json=None, - )) + model_response.data.append( + ImageObject( + url=image_item, + b64_json=None, + ) + ) elif isinstance(image_item, dict): # If output contains dict with url/b64_json - model_response.data.append(ImageObject( - url=image_item.get("url", None), - b64_json=image_item.get("b64_json", None), - )) - + model_response.data.append( + ImageObject( + url=image_item.get("url", None), + b64_json=image_item.get("b64_json", None), + ) + ) + return model_response @staticmethod def _check_timeout(start_time: float, timeout_secs: float) -> None: """ Check if operation has timed out. - + Args: start_time: Start time of the operation timeout_secs: Timeout duration in seconds - + Raises: TimeoutError: If operation has exceeded timeout """ @@ -159,22 +162,22 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): def _check_task_status(response_data: Dict[str, Any]) -> str: """ Check RunwayML task status from response. - + RunwayML statuses: PENDING, RUNNING, SUCCEEDED, FAILED, CANCELLED, THROTTLED - + Args: response_data: JSON response from RunwayML task endpoint - + Returns: Normalized status string: "running", "succeeded", or raises on failure - + Raises: ValueError: If task failed or status is unknown """ status = response_data.get("status", "").upper() - + verbose_logger.debug(f"RunwayML task status: {status}") - + if status == "SUCCEEDED": return "succeeded" elif status == "FAILED": @@ -199,16 +202,16 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): ) -> httpx.Response: """ Poll RunwayML task until completion (sync). - + RunwayML POST returns immediately with a task that has status PENDING/RUNNING. We need to poll GET /v1/tasks/{task_id} until status is SUCCEEDED or FAILED. - + Args: task_id: The task ID to poll api_base: Base URL for RunwayML API headers: Request headers (including auth) timeout_secs: Total timeout in seconds (default: 600s = 10 minutes) - + Returns: Final response with completed task """ @@ -216,25 +219,25 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): client = _get_httpx_client() start_time = time.time() - + # Build task status URL api_base = api_base.rstrip("/") task_url = f"{api_base}/v1/tasks/{task_id}" - + verbose_logger.debug(f"Polling RunwayML task: {task_url}") - + while True: self._check_timeout(start_time=start_time, timeout_secs=timeout_secs) - + # Poll the task status response = client.get(url=task_url, headers=headers) response.raise_for_status() - + response_data = response.json() - + # Check task status status = self._check_task_status(response_data=response_data) - + if status == "succeeded": return response elif status == "running": @@ -250,13 +253,13 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): ) -> httpx.Response: """ Poll RunwayML task until completion (async). - + Args: task_id: The task ID to poll api_base: Base URL for RunwayML API headers: Request headers (including auth) timeout_secs: Total timeout in seconds (default: 600s = 10 minutes) - + Returns: Final response with completed task """ @@ -265,25 +268,25 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): client = get_async_httpx_client(llm_provider=litellm.LlmProviders.RUNWAYML) start_time = time.time() - + # Build task status URL api_base = api_base.rstrip("/") task_url = f"{api_base}/v1/tasks/{task_id}" - + verbose_logger.debug(f"Polling RunwayML task (async): {task_url}") - + while True: self._check_timeout(start_time=start_time, timeout_secs=timeout_secs) - + # Poll the task status response = await client.get(url=task_url, headers=headers) response.raise_for_status() - + response_data = response.json() - + # Check task status status = self._check_task_status(response_data=response_data) - + if status == "succeeded": return response elif status == "running": @@ -305,17 +308,17 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): ) -> ImageResponse: """ Transform the image generation response to the litellm image response. - + RunwayML returns a task immediately with status PENDING/RUNNING. We need to poll the task until it completes (status SUCCEEDED). - + Initial response: { "id": "task_123...", "status": "PENDING" | "RUNNING", "createdAt": "2025-11-13T..." } - + After polling: { "id": "task_123...", @@ -332,23 +335,22 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): status_code=raw_response.status_code, headers=raw_response.headers, ) - - verbose_logger.debug( - "RunwayML starting polling..." - ) - + verbose_logger.debug("RunwayML starting polling...") + # Get task ID task_id = response_data.get("id") if not task_id: raise ValueError("RunwayML response missing task ID") - + # Get headers for polling (need auth) poll_headers = { "Authorization": raw_response.request.headers.get("Authorization", ""), - "X-Runway-Version": raw_response.request.headers.get("X-Runway-Version", RUNWAYML_DEFAULT_API_VERSION), + "X-Runway-Version": raw_response.request.headers.get( + "X-Runway-Version", RUNWAYML_DEFAULT_API_VERSION + ), } - + # Poll until task completes raw_response = self._poll_task_sync( task_id=task_id, @@ -356,12 +358,12 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): headers=poll_headers, timeout_secs=RUNWAYML_POLLING_TIMEOUT, ) - + # Update response_data with polled result response_data = raw_response.json() - + verbose_logger.debug("RunwayML polling complete, transforming to OpenAI format") - + # Transform RunwayML response to OpenAI format return self._transform_runwayml_response_to_openai( response_data=response_data, @@ -383,7 +385,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): ) -> ImageResponse: """ Async transform the image generation response to the litellm image response. - + RunwayML returns a task immediately with status PENDING/RUNNING. We need to poll the task until it completes (status SUCCEEDED) using async polling. """ @@ -395,22 +397,22 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): status_code=raw_response.status_code, headers=raw_response.headers, ) - - verbose_logger.debug( - "RunwayML starting polling (async)..." - ) - + + verbose_logger.debug("RunwayML starting polling (async)...") + # Get task ID task_id = response_data.get("id") if not task_id: raise ValueError("RunwayML response missing task ID") - + # Get headers for polling (need auth) poll_headers = { "Authorization": raw_response.request.headers.get("Authorization", ""), - "X-Runway-Version": raw_response.request.headers.get("X-Runway-Version", RUNWAYML_DEFAULT_API_VERSION), + "X-Runway-Version": raw_response.request.headers.get( + "X-Runway-Version", RUNWAYML_DEFAULT_API_VERSION + ), } - + # Poll until task completes (async) raw_response = await self._poll_task_async( task_id=task_id, @@ -418,18 +420,20 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): headers=poll_headers, timeout_secs=RUNWAYML_POLLING_TIMEOUT, ) - + # Update response_data with polled result response_data = raw_response.json() - - verbose_logger.debug("RunwayML polling complete (async), transforming to OpenAI format") - + + verbose_logger.debug( + "RunwayML polling complete (async), transforming to OpenAI format" + ) + # Transform RunwayML response to OpenAI format return self._transform_runwayml_response_to_openai( response_data=response_data, model_response=model_response, ) - + def get_supported_openai_params( self, model: str ) -> List[OpenAIImageGenerationOptionalParams]: @@ -439,7 +443,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): return [ "size", ] - + def map_openai_params( self, non_default_params: dict, @@ -448,7 +452,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): drop_params: bool, ) -> dict: supported_params = self.get_supported_openai_params(model) - + # Map OpenAI 'size' parameter to RunwayML 'ratio' parameter if "size" in non_default_params: size = non_default_params["size"] @@ -461,7 +465,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): "1080x1920": "1080:1920", } optional_params["ratio"] = size_to_ratio_map.get(size, "1920:1080") - + for k in non_default_params.keys(): if k not in optional_params.keys(): if k in supported_params: @@ -485,7 +489,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): ) -> dict: """ Transform the image generation request to the RunwayML image generation request body - + RunwayML expects: - model: The model to use (e.g., 'gen4_image') - promptText: The text prompt @@ -495,7 +499,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): "model": model or "gen4_image", "promptText": prompt, } - + # Add any RunwayML-specific parameters if "ratio" in optional_params: runwayml_request_body["ratio"] = optional_params["ratio"] @@ -503,11 +507,9 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): # Set default ratio if not provided runwayml_request_body["ratio"] = "1920:1080" - # Add any other optional parameters for k, v in optional_params.items(): if k not in runwayml_request_body and k not in ["size"]: runwayml_request_body[k] = v - - return runwayml_request_body + return runwayml_request_body diff --git a/litellm/llms/runwayml/text_to_speech/__init__.py b/litellm/llms/runwayml/text_to_speech/__init__.py index 491e8449e0..98337a8321 100644 --- a/litellm/llms/runwayml/text_to_speech/__init__.py +++ b/litellm/llms/runwayml/text_to_speech/__init__.py @@ -2,4 +2,3 @@ from .transformation import RunwayMLTextToSpeechConfig __all__ = ["RunwayMLTextToSpeechConfig"] - diff --git a/litellm/llms/runwayml/text_to_speech/transformation.py b/litellm/llms/runwayml/text_to_speech/transformation.py index ac926beb22..dfcb92bc68 100644 --- a/litellm/llms/runwayml/text_to_speech/transformation.py +++ b/litellm/llms/runwayml/text_to_speech/transformation.py @@ -32,25 +32,25 @@ else: class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): """ Configuration for RunwayML Text-to-Speech - + Reference: https://api.dev.runwayml.com/v1/text_to_speech """ - + DEFAULT_BASE_URL: str = "https://api.dev.runwayml.com" TTS_ENDPOINT_PATH: str = "v1/text_to_speech" DEFAULT_MODEL: str = "eleven_multilingual_v2" DEFAULT_VOICE_TYPE: str = "runway-preset" DEFAULT_VOICE_PRESET_ID: str = "Bernard" - + # Voice mappings from OpenAI voices to RunwayML preset IDs # OpenAI voices mapped to similar-sounding RunwayML voices VOICE_MAPPINGS = { - "alloy": "Maya", # Neutral, balanced female voice - "echo": "James", # Male voice - "fable": "Bernard", # Warm, storytelling voice - "onyx": "Vincent", # Deep male voice - "nova": "Serene", # Warm, expressive female voice - "shimmer": "Ella", # Clear, friendly female voice + "alloy": "Maya", # Neutral, balanced female voice + "echo": "James", # Male voice + "fable": "Bernard", # Warm, storytelling voice + "onyx": "Vincent", # Deep male voice + "nova": "Serene", # Warm, expressive female voice + "shimmer": "Ella", # Clear, friendly female voice } def dispatch_text_to_speech( @@ -74,9 +74,9 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): ]: """ Dispatch method to handle RunwayML TTS requests - + This method encapsulates RunwayML-specific credential resolution and parameter handling - + Args: base_llm_http_handler: The BaseLLMHTTPHandler instance from main.py """ @@ -88,7 +88,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): or get_secret_str("RUNWAYML_API_BASE") or self.DEFAULT_BASE_URL ) - + # Resolve api_key from multiple sources api_key = ( api_key @@ -97,7 +97,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): or get_secret_str("RUNWAYML_API_SECRET") or get_secret_str("RUNWAYML_API_KEY") ) - + # Convert voice to appropriate format voice_param: Optional[Union[str, Dict]] = voice if isinstance(voice, str): @@ -106,12 +106,14 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): elif isinstance(voice, dict): # Already in dict format, pass through voice_param = voice - - litellm_params_dict.update({ - "api_key": api_key, - "api_base": api_base, - }) - + + litellm_params_dict.update( + { + "api_key": api_key, + "api_base": api_base, + } + ) + # Call the text_to_speech_handler response = base_llm_http_handler.text_to_speech_handler( model=model, @@ -127,7 +129,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): client=None, _is_async=aspeech, ) - + return response def get_supported_openai_params(self, model: str) -> list: @@ -146,15 +148,15 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): ) -> Tuple[Optional[str], Dict]: """ Map OpenAI parameters to RunwayML TTS parameters - + Returns: Tuple of (mapped_voice_string, mapped_params) - + Note: Since RunwayML requires voice as a dict, we store it in mapped_params["runwayml_voice"] and return None for the voice string. """ mapped_params = {} - + # Map voice parameter to RunwayML format dict voice_dict: Optional[Dict] = None if isinstance(voice, str): @@ -174,14 +176,14 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): elif isinstance(voice, dict): # Already in RunwayML format, use as-is voice_dict = voice - + # Store the voice dict in optional_params for later use if voice_dict is not None: mapped_params["runwayml_voice"] = voice_dict - + # No other OpenAI params are currently supported by RunwayML TTS # (response_format, speed, etc. are not supported) - + # Return None for voice string since RunwayML uses dict format return None, mapped_params @@ -196,20 +198,20 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): Validate RunwayML environment and set up authentication headers """ validated_headers = headers.copy() - + final_api_key = ( - api_key - or get_secret_str("RUNWAYML_API_SECRET") + api_key + or get_secret_str("RUNWAYML_API_SECRET") or get_secret_str("RUNWAYML_API_KEY") ) - + if not final_api_key: raise ValueError("RUNWAYML_API_SECRET or RUNWAYML_API_KEY is not set") - + validated_headers["Authorization"] = f"Bearer {final_api_key}" validated_headers["X-Runway-Version"] = RUNWAYML_DEFAULT_API_VERSION validated_headers["Content-Type"] = "application/json" - + return validated_headers def get_complete_url( @@ -222,11 +224,9 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): Get the complete URL for RunwayML TTS request """ complete_url = ( - api_base - or get_secret_str("RUNWAYML_API_BASE") - or self.DEFAULT_BASE_URL + api_base or get_secret_str("RUNWAYML_API_BASE") or self.DEFAULT_BASE_URL ) - + complete_url = complete_url.rstrip("/") return f"{complete_url}/{self.TTS_ENDPOINT_PATH}" @@ -234,11 +234,11 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): def _check_timeout(start_time: float, timeout_secs: float) -> None: """ Check if operation has timed out. - + Args: start_time: Start time of the operation timeout_secs: Timeout duration in seconds - + Raises: TimeoutError: If operation has exceeded timeout """ @@ -251,22 +251,22 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): def _check_task_status(response_data: Dict[str, Any]) -> str: """ Check RunwayML task status from response. - + RunwayML statuses: PENDING, RUNNING, SUCCEEDED, FAILED, CANCELLED, THROTTLED - + Args: response_data: JSON response from RunwayML task endpoint - + Returns: Normalized status string: "running", "succeeded", or raises on failure - + Raises: ValueError: If task failed or status is unknown """ status = response_data.get("status", "").upper() - + verbose_logger.debug(f"RunwayML TTS task status: {status}") - + if status == "SUCCEEDED": return "succeeded" elif status == "FAILED": @@ -291,16 +291,16 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): ) -> httpx.Response: """ Poll RunwayML task until completion (sync). - + RunwayML POST returns immediately with a task that has status PENDING/RUNNING. We need to poll GET /v1/tasks/{task_id} until status is SUCCEEDED or FAILED. - + Args: task_id: The task ID to poll api_base: Base URL for RunwayML API headers: Request headers (including auth) timeout_secs: Total timeout in seconds (default: 600s = 10 minutes) - + Returns: Final response with completed task """ @@ -308,25 +308,25 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): client = _get_httpx_client() start_time = time.time() - + # Build task status URL api_base = api_base.rstrip("/") task_url = f"{api_base}/v1/tasks/{task_id}" - + verbose_logger.debug(f"Polling RunwayML TTS task: {task_url}") - + while True: self._check_timeout(start_time=start_time, timeout_secs=timeout_secs) - + # Poll the task status response = client.get(url=task_url, headers=headers) response.raise_for_status() - + response_data = response.json() - + # Check task status status = self._check_task_status(response_data=response_data) - + if status == "succeeded": return response elif status == "running": @@ -342,13 +342,13 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): ) -> httpx.Response: """ Poll RunwayML task until completion (async). - + Args: task_id: The task ID to poll api_base: Base URL for RunwayML API headers: Request headers (including auth) timeout_secs: Total timeout in seconds (default: 600s = 10 minutes) - + Returns: Final response with completed task """ @@ -356,25 +356,25 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): client = get_async_httpx_client(llm_provider=litellm.LlmProviders.RUNWAYML) start_time = time.time() - + # Build task status URL api_base = api_base.rstrip("/") task_url = f"{api_base}/v1/tasks/{task_id}" - + verbose_logger.debug(f"Polling RunwayML TTS task (async): {task_url}") - + while True: self._check_timeout(start_time=start_time, timeout_secs=timeout_secs) - + # Poll the task status response = await client.get(url=task_url, headers=headers) response.raise_for_status() - + response_data = response.json() - + # Check task status status = self._check_task_status(response_data=response_data) - + if status == "succeeded": return response elif status == "running": @@ -392,7 +392,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): ) -> TextToSpeechRequestData: """ Transform OpenAI TTS request to RunwayML TTS format - + RunwayML expects: - model: The model to use (e.g., 'eleven_multilingual_v2') - promptText: The text to convert to speech @@ -401,7 +401,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): "type": "runway-preset", "presetId": "Bernard" } - + Returns: TextToSpeechRequestData: Contains JSON body and headers """ @@ -413,19 +413,19 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): "type": self.DEFAULT_VOICE_TYPE, "presetId": self.DEFAULT_VOICE_PRESET_ID, } - + # Build request body request_body = { "model": model or self.DEFAULT_MODEL, "promptText": input, "voice": runwayml_voice, } - + # Add any other optional parameters (except runwayml_voice which we already used) for k, v in optional_params.items(): if k not in request_body and k != "runwayml_voice": request_body[k] = v - + return { "dict_body": request_body, "headers": headers, @@ -439,17 +439,17 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): ) -> "HttpxBinaryResponseContent": """ Transform RunwayML TTS response to standard format - + RunwayML returns a task immediately with status PENDING/RUNNING. We need to poll the task until it completes, then download the audio. - + Initial response: { "id": "task_123...", "status": "PENDING" | "RUNNING", "createdAt": "2025-11-13T..." } - + After polling: { "id": "task_123...", @@ -468,14 +468,14 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): status_code=raw_response.status_code, headers=dict(raw_response.headers), ) - + verbose_logger.debug("RunwayML TTS starting polling...") - + # Get task ID task_id = response_data.get("id") if not task_id: raise ValueError("RunwayML TTS response missing task ID") - + # Get headers for polling (need auth) poll_headers = { "Authorization": raw_response.request.headers.get("Authorization", ""), @@ -483,7 +483,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): "X-Runway-Version", RUNWAYML_DEFAULT_API_VERSION ), } - + # Poll until task completes polled_response = self._poll_task_sync( task_id=task_id, @@ -491,30 +491,30 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): headers=poll_headers, timeout_secs=RUNWAYML_POLLING_TIMEOUT, ) - + # Get the completed task data task_data = polled_response.json() - + verbose_logger.debug("RunwayML TTS polling complete, downloading audio") - + # Get audio URL from output output = task_data.get("output", []) if not output or not isinstance(output, list) or len(output) == 0: raise ValueError("RunwayML TTS response missing audio URL in output") - + audio_url = output[0] if not isinstance(audio_url, str): raise ValueError(f"RunwayML TTS audio URL is not a string: {audio_url}") - + # Download the audio file from litellm.llms.custom_httpx.http_handler import _get_httpx_client client = _get_httpx_client() audio_response = client.get(url=audio_url) audio_response.raise_for_status() - + verbose_logger.debug("RunwayML TTS audio downloaded successfully") - + # Return the audio data wrapped in HttpxBinaryResponseContent return HttpxBinaryResponseContent(audio_response) @@ -526,7 +526,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): ) -> "HttpxBinaryResponseContent": """ Async transform RunwayML TTS response to standard format - + Same as sync version but uses async polling and download """ from litellm.types.llms.openai import HttpxBinaryResponseContent @@ -539,14 +539,14 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): status_code=raw_response.status_code, headers=dict(raw_response.headers), ) - + verbose_logger.debug("RunwayML TTS starting polling (async)...") - + # Get task ID task_id = response_data.get("id") if not task_id: raise ValueError("RunwayML TTS response missing task ID") - + # Get headers for polling (need auth) poll_headers = { "Authorization": raw_response.request.headers.get("Authorization", ""), @@ -554,7 +554,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): "X-Runway-Version", RUNWAYML_DEFAULT_API_VERSION ), } - + # Poll until task completes (async) polled_response = await self._poll_task_async( task_id=task_id, @@ -562,30 +562,29 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): headers=poll_headers, timeout_secs=RUNWAYML_POLLING_TIMEOUT, ) - + # Get the completed task data task_data = polled_response.json() - + verbose_logger.debug("RunwayML TTS polling complete (async), downloading audio") - + # Get audio URL from output output = task_data.get("output", []) if not output or not isinstance(output, list) or len(output) == 0: raise ValueError("RunwayML TTS response missing audio URL in output") - + audio_url = output[0] if not isinstance(audio_url, str): raise ValueError(f"RunwayML TTS audio URL is not a string: {audio_url}") - + # Download the audio file (async) from litellm.llms.custom_httpx.http_handler import get_async_httpx_client client = get_async_httpx_client(llm_provider=litellm.LlmProviders.RUNWAYML) audio_response = await client.get(url=audio_url) audio_response.raise_for_status() - + verbose_logger.debug("RunwayML TTS audio downloaded successfully (async)") - + # Return the audio data wrapped in HttpxBinaryResponseContent return HttpxBinaryResponseContent(audio_response) - diff --git a/litellm/llms/runwayml/videos/__init__.py b/litellm/llms/runwayml/videos/__init__.py index 9c72dec29a..6d6f2b65e9 100644 --- a/litellm/llms/runwayml/videos/__init__.py +++ b/litellm/llms/runwayml/videos/__init__.py @@ -1,2 +1 @@ # RunwayML video generation - diff --git a/litellm/llms/runwayml/videos/transformation.py b/litellm/llms/runwayml/videos/transformation.py index 318a732dc2..3fc656a92b 100644 --- a/litellm/llms/runwayml/videos/transformation.py +++ b/litellm/llms/runwayml/videos/transformation.py @@ -33,7 +33,7 @@ else: class RunwayMLVideoConfig(BaseVideoConfig): """ Configuration class for RunwayML video generation. - + RunwayML uses a task-based API where: 1. POST /v1/image_to_video creates a task 2. The task returns immediately with a task ID @@ -70,43 +70,47 @@ class RunwayMLVideoConfig(BaseVideoConfig): ) -> Dict: """ Map OpenAI parameters to RunwayML format. - + Mappings: - prompt -> promptText - - input_reference -> promptImage + - input_reference -> promptImage - size -> ratio (convert "WIDTHxHEIGHT" to "WIDTH:HEIGHT") - seconds -> duration (convert to integer) """ mapped_params: Dict[str, Any] = {} - + # Handle input_reference parameter - map to promptImage if "input_reference" in video_create_optional_params: input_reference = video_create_optional_params["input_reference"] # RunwayML supports URLs and data URIs directly mapped_params["promptImage"] = input_reference - + # Handle size parameter - convert "1280x720" to "1280:720" if "size" in video_create_optional_params: size = video_create_optional_params["size"] if isinstance(size, str) and "x" in size: mapped_params["ratio"] = size.replace("x", ":") - + # Handle seconds parameter - convert to integer if "seconds" in video_create_optional_params: seconds = video_create_optional_params["seconds"] if seconds is not None: try: - mapped_params["duration"] = int(float(seconds)) if isinstance(seconds, str) else int(seconds) + mapped_params["duration"] = ( + int(float(seconds)) + if isinstance(seconds, str) + else int(seconds) + ) except (ValueError, TypeError): # If conversion fails, use default duration pass - + # Pass through other parameters that aren't OpenAI-specific supported_openai_params = self.get_supported_openai_params(model) for key, value in video_create_optional_params.items(): if key not in supported_openai_params: mapped_params[key] = value - + return mapped_params def validate_environment( @@ -123,25 +127,27 @@ class RunwayMLVideoConfig(BaseVideoConfig): # Use api_key from litellm_params if available, otherwise fall back to other sources if litellm_params and litellm_params.api_key: api_key = api_key or litellm_params.api_key - + api_key = ( api_key or litellm.api_key or get_secret_str("RUNWAYML_API_SECRET") or get_secret_str("RUNWAYML_API_KEY") ) - + if api_key is None: raise ValueError( "RunwayML API key is required. Set RUNWAYML_API_SECRET environment variable " "or pass api_key parameter." ) - - headers.update({ - "Authorization": f"Bearer {api_key}", - "X-Runway-Version": RUNWAYML_DEFAULT_API_VERSION, - "Content-Type": "application/json", - }) + + headers.update( + { + "Authorization": f"Bearer {api_key}", + "X-Runway-Version": RUNWAYML_DEFAULT_API_VERSION, + "Content-Type": "application/json", + } + ) return headers def get_complete_url( @@ -156,8 +162,8 @@ class RunwayMLVideoConfig(BaseVideoConfig): """ if api_base is None: api_base = "https://api.dev.runwayml.com/v1" - - return api_base.rstrip('/') + + return api_base.rstrip("/") def transform_video_create_request( self, @@ -170,7 +176,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): ) -> Tuple[Dict, RequestFiles, str]: """ Transform the video creation request for RunwayML API. - + RunwayML expects: { "model": "gen4_turbo", @@ -179,22 +185,22 @@ class RunwayMLVideoConfig(BaseVideoConfig): "ratio": "1280:720", "duration": 5 } - """ + """ # Build the request data request_data: Dict[str, Any] = { "model": model, "promptText": prompt, } - + # Add mapped parameters request_data.update(video_create_optional_request_params) - + # RunwayML uses JSON body, no files multipart files_list: List[Tuple[str, Any]] = [] - + # Append the specific endpoint for video generation full_api_base = f"{api_base}/image_to_video" - + return request_data, files_list, full_api_base def transform_video_create_response( @@ -207,18 +213,18 @@ class RunwayMLVideoConfig(BaseVideoConfig): ) -> VideoObject: """ Transform the RunwayML video creation response. - + RunwayML returns a task object that looks like: { "id": "task_123...", "status": "PENDING" | "RUNNING" | "SUCCEEDED" | "FAILED", "output": ["https://...video.mp4"] (when succeeded) } - + We map this to OpenAI VideoObject format. """ response_data = raw_response.json() - + # Map RunwayML task response to VideoObject format video_data: Dict[str, Any] = { "id": response_data.get("id", ""), @@ -226,21 +232,27 @@ class RunwayMLVideoConfig(BaseVideoConfig): "status": self._map_runway_status(response_data.get("status", "pending")), "created_at": self._parse_runway_timestamp(response_data.get("createdAt")), } - + # Add optional fields if present if "output" in response_data and response_data["output"]: # RunwayML returns output as array of URLs when task succeeds - video_data["output_url"] = response_data["output"][0] if isinstance(response_data["output"], list) else response_data["output"] - + video_data["output_url"] = ( + response_data["output"][0] + if isinstance(response_data["output"], list) + else response_data["output"] + ) + if "completedAt" in response_data: - video_data["completed_at"] = self._parse_runway_timestamp(response_data.get("completedAt")) - + video_data["completed_at"] = self._parse_runway_timestamp( + response_data.get("completedAt") + ) + if "failureCode" in response_data or "failure" in response_data: video_data["error"] = { "code": response_data.get("failureCode", "unknown"), - "message": response_data.get("failure", "Video generation failed") + "message": response_data.get("failure", "Video generation failed"), } - + # Add model and size info if available from request if request_data: if "model" in request_data: @@ -252,27 +264,29 @@ class RunwayMLVideoConfig(BaseVideoConfig): video_data["size"] = ratio.replace(":", "x") if "duration" in request_data: video_data["seconds"] = str(request_data["duration"]) - + video_obj = VideoObject(**video_data) # type: ignore[arg-type] - + if custom_llm_provider and video_obj.id: - video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, model) - + video_obj.id = encode_video_id_with_provider( + video_obj.id, custom_llm_provider, model + ) + # Add usage data for cost tracking usage_data = {} - if video_obj and hasattr(video_obj, 'seconds') and video_obj.seconds: + if video_obj and hasattr(video_obj, "seconds") and video_obj.seconds: try: usage_data["duration_seconds"] = float(video_obj.seconds) except (ValueError, TypeError): pass video_obj.usage = usage_data - + return video_obj def _map_runway_status(self, runway_status: str) -> str: """ Map RunwayML status to OpenAI status format. - + RunwayML statuses: PENDING, RUNNING, SUCCEEDED, FAILED, CANCELLED OpenAI statuses: queued, in_progress, completed, failed """ @@ -285,20 +299,20 @@ class RunwayMLVideoConfig(BaseVideoConfig): "THROTTLED": "queued", } return status_map.get(runway_status.upper(), "queued") - + def _parse_runway_timestamp(self, timestamp_str: Optional[str]) -> int: """ Convert RunwayML ISO 8601 timestamp to Unix timestamp. - + RunwayML returns timestamps like: "2025-11-11T21:48:50.448Z" We need to convert to Unix timestamp (seconds since epoch). """ if not timestamp_str: return 0 - + try: # Parse ISO 8601 timestamp - dt = datetime.fromisoformat(timestamp_str.replace('Z', '+00:00')) + dt = datetime.fromisoformat(timestamp_str.replace("Z", "+00:00")) # Convert to Unix timestamp return int(dt.timestamp()) except (ValueError, AttributeError): @@ -320,12 +334,12 @@ class RunwayMLVideoConfig(BaseVideoConfig): We'll retrieve the task and extract the video URL. """ original_video_id = extract_original_video_id(video_id) - + # Get task status to retrieve video URL url = f"{api_base}/tasks/{original_video_id}" - + params: Dict[str, Any] = {} - + return url, params def _extract_video_url_from_response(self, response_data: Dict[str, Any]) -> str: @@ -338,18 +352,22 @@ class RunwayMLVideoConfig(BaseVideoConfig): if "output" in response_data and response_data["output"]: output = response_data["output"] video_url = output[0] if isinstance(output, list) else output - + if not video_url: # Check if the video generation failed or is still processing status = response_data.get("status", "UNKNOWN") if status in ["PENDING", "RUNNING", "THROTTLED"]: - raise ValueError(f"Video is still processing (status: {status}). Please wait and try again.") + raise ValueError( + f"Video is still processing (status: {status}). Please wait and try again." + ) elif status == "FAILED": failure_reason = response_data.get("failure", "Unknown error") raise ValueError(f"Video generation failed: {failure_reason}") else: - raise ValueError("Video URL not found in response. Video may not be ready yet.") - + raise ValueError( + "Video URL not found in response. Video may not be ready yet." + ) + return video_url def transform_video_content_response( @@ -359,10 +377,10 @@ class RunwayMLVideoConfig(BaseVideoConfig): ) -> bytes: """ Transform the RunwayML video content download response (synchronous). - + RunwayML's task endpoint returns JSON with a video URL in the output field. We need to extract the URL and download the video. - + Example response: { "id":"63fd0f13-f29d-4e58-99d3-1cb9efa14a5b", @@ -373,12 +391,12 @@ class RunwayMLVideoConfig(BaseVideoConfig): """ response_data = raw_response.json() video_url = self._extract_video_url_from_response(response_data) - + # Download the video from the CloudFront URL synchronously httpx_client: HTTPHandler = _get_httpx_client() video_response = httpx_client.get(video_url) video_response.raise_for_status() - + return video_response.content async def async_transform_video_content_response( @@ -388,10 +406,10 @@ class RunwayMLVideoConfig(BaseVideoConfig): ) -> bytes: """ Transform the RunwayML video content download response (asynchronous). - + RunwayML's task endpoint returns JSON with a video URL in the output field. We need to extract the URL and download the video asynchronously. - + Example response: { "id":"63fd0f13-f29d-4e58-99d3-1cb9efa14a5b", @@ -402,14 +420,14 @@ class RunwayMLVideoConfig(BaseVideoConfig): """ response_data = raw_response.json() video_url = self._extract_video_url_from_response(response_data) - + # Download the video from the CloudFront URL asynchronously async_httpx_client: AsyncHTTPHandler = get_async_httpx_client( llm_provider=litellm.LlmProviders.RUNWAYML, ) video_response = await async_httpx_client.get(video_url) video_response.raise_for_status() - + return video_response.content def transform_video_remix_request( @@ -423,7 +441,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): ) -> Tuple[str, Dict]: """ Transform the video remix request for RunwayML API. - + RunwayML doesn't have a direct remix endpoint in their current API. This would need to be implemented when/if they add this feature. """ @@ -450,7 +468,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): ) -> Tuple[str, Dict]: """ Transform the video list request for RunwayML API. - + RunwayML doesn't expose a list endpoint in their public API yet. """ raise NotImplementedError("Video listing is not yet supported by RunwayML API") @@ -473,16 +491,16 @@ class RunwayMLVideoConfig(BaseVideoConfig): ) -> Tuple[str, Dict]: """ Transform the video delete request for RunwayML API. - + RunwayML uses task cancellation. """ original_video_id = extract_original_video_id(video_id) - + # Construct the URL for task cancellation url = f"{api_base}/tasks/{original_video_id}/cancel" - + data: Dict[str, Any] = {} - + return url, data def transform_video_delete_response( @@ -492,7 +510,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): ) -> VideoObject: """Transform the RunwayML video delete/cancel response.""" response_data = raw_response.json() - + video_obj = VideoObject( id=response_data.get("id", ""), object="video", @@ -511,17 +529,17 @@ class RunwayMLVideoConfig(BaseVideoConfig): ) -> Tuple[str, Dict]: """ Transform the RunwayML video status retrieve request. - + RunwayML uses GET /v1/tasks/{task_id} to retrieve task status. """ original_video_id = extract_original_video_id(video_id) - + # Construct the full URL for task status retrieval url = f"{api_base}/tasks/{original_video_id}" - + # Empty dict for GET request (no body) data: Dict[str, Any] = {} - + return url, data def transform_video_status_retrieve_response( @@ -534,7 +552,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): Transform the RunwayML video status retrieve response. """ response_data = raw_response.json() - + # Map RunwayML task response to VideoObject format video_data: Dict[str, Any] = { "id": response_data.get("id", ""), @@ -542,27 +560,35 @@ class RunwayMLVideoConfig(BaseVideoConfig): "status": self._map_runway_status(response_data.get("status", "pending")), "created_at": self._parse_runway_timestamp(response_data.get("createdAt")), } - + # Add optional fields if present if "output" in response_data and response_data["output"]: - video_data["output_url"] = response_data["output"][0] if isinstance(response_data["output"], list) else response_data["output"] - + video_data["output_url"] = ( + response_data["output"][0] + if isinstance(response_data["output"], list) + else response_data["output"] + ) + if "completedAt" in response_data: - video_data["completed_at"] = self._parse_runway_timestamp(response_data.get("completedAt")) - + video_data["completed_at"] = self._parse_runway_timestamp( + response_data.get("completedAt") + ) + if "progress" in response_data: video_data["progress"] = response_data["progress"] - + if "failureCode" in response_data or "failure" in response_data: video_data["error"] = { "code": response_data.get("failureCode", "unknown"), - "message": response_data.get("failure", "Video generation failed") + "message": response_data.get("failure", "Video generation failed"), } - + video_obj = VideoObject(**video_data) # type: ignore[arg-type] - + if custom_llm_provider and video_obj.id: - video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, None) + video_obj.id = encode_video_id_with_provider( + video_obj.id, custom_llm_provider, None + ) return video_obj @@ -576,4 +602,3 @@ class RunwayMLVideoConfig(BaseVideoConfig): message=error_message, headers=headers, ) - diff --git a/litellm/llms/s3_vectors/vector_stores/transformation.py b/litellm/llms/s3_vectors/vector_stores/transformation.py index df81a78289..11836e361e 100644 --- a/litellm/llms/s3_vectors/vector_stores/transformation.py +++ b/litellm/llms/s3_vectors/vector_stores/transformation.py @@ -82,13 +82,15 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): # If not in that format, try to construct it from litellm_params bucket_name: str index_name: str - + if ":" in vector_store_id: bucket_name, index_name = vector_store_id.split(":", 1) else: # Try to get bucket_name from litellm_params bucket_name_from_params = litellm_params.get("vector_bucket_name") - if not bucket_name_from_params or not isinstance(bucket_name_from_params, str): + if not bucket_name_from_params or not isinstance( + bucket_name_from_params, str + ): raise ValueError( "vector_store_id must be in format 'bucket_name:index_name' for S3 Vectors, " "or vector_bucket_name must be provided in litellm_params" @@ -100,10 +102,15 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): query = " ".join(query) # Generate embedding for the query - embedding_model = litellm_params.get("embedding_model", "text-embedding-3-small") - + embedding_model = litellm_params.get( + "embedding_model", "text-embedding-3-small" + ) + import litellm as litellm_module - embedding_response = litellm_module.embedding(model=embedding_model, input=[query]) + + embedding_response = litellm_module.embedding( + model=embedding_model, input=[query] + ) query_embedding = embedding_response.data[0]["embedding"] url = f"{api_base}/QueryVectors" @@ -112,7 +119,9 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): "vectorBucketName": bucket_name, "indexName": index_name, "queryVector": {"float32": query_embedding}, - "topK": vector_store_search_optional_params.get("max_num_results", 5), # Default to 5 + "topK": vector_store_search_optional_params.get( + "max_num_results", 5 + ), # Default to 5 "returnDistance": True, "returnMetadata": True, } @@ -134,13 +143,15 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): # If not in that format, try to construct it from litellm_params bucket_name: str index_name: str - + if ":" in vector_store_id: bucket_name, index_name = vector_store_id.split(":", 1) else: # Try to get bucket_name from litellm_params bucket_name_from_params = litellm_params.get("vector_bucket_name") - if not bucket_name_from_params or not isinstance(bucket_name_from_params, str): + if not bucket_name_from_params or not isinstance( + bucket_name_from_params, str + ): raise ValueError( "vector_store_id must be in format 'bucket_name:index_name' for S3 Vectors, " "or vector_bucket_name must be provided in litellm_params" @@ -152,10 +163,15 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): query = " ".join(query) # Generate embedding for the query asynchronously - embedding_model = litellm_params.get("embedding_model", "text-embedding-3-small") - + embedding_model = litellm_params.get( + "embedding_model", "text-embedding-3-small" + ) + import litellm as litellm_module - embedding_response = await litellm_module.aembedding(model=embedding_model, input=[query]) + + embedding_response = await litellm_module.aembedding( + model=embedding_model, input=[query] + ) query_embedding = embedding_response.data[0]["embedding"] url = f"{api_base}/QueryVectors" @@ -164,7 +180,9 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): "vectorBucketName": bucket_name, "indexName": index_name, "queryVector": {"float32": query_embedding}, - "topK": vector_store_search_optional_params.get("max_num_results", 5), # Default to 5 + "topK": vector_store_search_optional_params.get( + "max_num_results", 5 + ), # Default to 5 "returnDistance": True, "returnMetadata": True, } @@ -223,7 +241,9 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): results.append( VectorStoreSearchResult( score=score, - content=[VectorStoreResultContent(text=source_text, type="text")], + content=[ + VectorStoreResultContent(text=source_text, type="text") + ], file_id=file_id, filename=filename, attributes=metadata, diff --git a/litellm/llms/sagemaker/completion/handler.py b/litellm/llms/sagemaker/completion/handler.py index 2a30dc5ef3..efbb218f57 100644 --- a/litellm/llms/sagemaker/completion/handler.py +++ b/litellm/llms/sagemaker/completion/handler.py @@ -583,35 +583,17 @@ class SagemakerLLM(BaseAWSLLM): ### BOTO3 INIT import boto3 - # pop aws_secret_access_key, aws_access_key_id, aws_region_name from kwargs, since completion calls fail with them - aws_secret_access_key = optional_params.pop("aws_secret_access_key", None) - aws_access_key_id = optional_params.pop("aws_access_key_id", None) - aws_region_name = optional_params.pop("aws_region_name", None) + # Use _load_credentials to support role assumption (aws_role_name, aws_session_name) + credentials, aws_region_name = self._load_credentials(optional_params) - if aws_access_key_id is not None: - # uses auth params passed to completion - # aws_access_key_id is not None, assume user is trying to auth using litellm.completion - client = boto3.client( - service_name="sagemaker-runtime", - aws_access_key_id=aws_access_key_id, - aws_secret_access_key=aws_secret_access_key, - region_name=aws_region_name, - ) - else: - # aws_access_key_id is None, assume user is trying to auth using env variables - # boto3 automaticaly reads env variables - - # we need to read region name from env - # I assume majority of users use .env for auth - region_name = ( - get_secret("AWS_REGION_NAME") - or aws_region_name # get region from config file if specified - or "us-west-2" # default to us-west-2 if region not specified - ) - client = boto3.client( - service_name="sagemaker-runtime", - region_name=region_name, - ) + # Create boto3 session with the loaded credentials + session = boto3.Session( + aws_access_key_id=credentials.access_key, + aws_secret_access_key=credentials.secret_key, + aws_session_token=credentials.token, + region_name=aws_region_name, + ) + client = session.client(service_name="sagemaker-runtime") # pop streaming if it's in the optional params as 'stream' raises an error with sagemaker inference_params = deepcopy(optional_params) @@ -628,7 +610,9 @@ class SagemakerLLM(BaseAWSLLM): #### EMBEDDING LOGIC # Transform request based on model type provider_config = SagemakerEmbeddingConfig.get_model_config(model) - request_data = provider_config.transform_embedding_request(model, input, optional_params, {}) + request_data = provider_config.transform_embedding_request( + model, input, optional_params, {} + ) data = json.dumps(request_data).encode("utf-8") ## LOGGING @@ -673,19 +657,19 @@ class SagemakerLLM(BaseAWSLLM): ) print_verbose(f"raw model_response: {response}") - + # Transform response based on model type from httpx import Response as HttpxResponse - + # Create a mock httpx Response object for the transformation mock_response = HttpxResponse( status_code=200, - content=json.dumps(response).encode('utf-8'), - headers={"content-type": "application/json"} + content=json.dumps(response).encode("utf-8"), + headers={"content-type": "application/json"}, ) - + model_response = EmbeddingResponse() - + # Use the request_data that was already transformed above return provider_config.transform_embedding_response( model=model, @@ -695,5 +679,5 @@ class SagemakerLLM(BaseAWSLLM): api_key=None, request_data=request_data, optional_params=optional_params, - litellm_params=litellm_params or {} + litellm_params=litellm_params or {}, ) diff --git a/litellm/llms/sagemaker/completion/transformation.py b/litellm/llms/sagemaker/completion/transformation.py index 42202bbf07..dd7cb60390 100644 --- a/litellm/llms/sagemaker/completion/transformation.py +++ b/litellm/llms/sagemaker/completion/transformation.py @@ -68,7 +68,15 @@ class SagemakerConfig(BaseConfig): ) def get_supported_openai_params(self, model: str) -> List: - return ["stream", "temperature", "max_tokens", "max_completion_tokens", "top_p", "stop", "n"] + return [ + "stream", + "temperature", + "max_tokens", + "max_completion_tokens", + "top_p", + "stop", + "n", + ] def map_openai_params( self, @@ -278,5 +286,3 @@ class SagemakerConfig(BaseConfig): headers = {"Content-Type": "application/json", **headers} return headers - - diff --git a/litellm/llms/sagemaker/embedding/transformation.py b/litellm/llms/sagemaker/embedding/transformation.py index 04b201380f..0443017118 100644 --- a/litellm/llms/sagemaker/embedding/transformation.py +++ b/litellm/llms/sagemaker/embedding/transformation.py @@ -23,7 +23,7 @@ class SagemakerEmbeddingConfig(BaseEmbeddingConfig): """ SageMaker embedding configuration factory for supporting embedding parameters """ - + def __init__(self) -> None: pass @@ -31,10 +31,10 @@ class SagemakerEmbeddingConfig(BaseEmbeddingConfig): def get_model_config(cls, model: str) -> "BaseEmbeddingConfig": """ Factory method to get the appropriate embedding config based on model type - + Args: model: The model name - + Returns: Appropriate embedding config instance """ @@ -57,7 +57,6 @@ class SagemakerEmbeddingConfig(BaseEmbeddingConfig): model: str, drop_params: bool, ) -> dict: - return optional_params def get_error_class( @@ -98,8 +97,8 @@ class SagemakerEmbeddingConfig(BaseEmbeddingConfig): response_data = raw_response.json() except Exception as e: raise SagemakerError( - message=f"Failed to parse response: {str(e)}", - status_code=raw_response.status_code + message=f"Failed to parse response: {str(e)}", + status_code=raw_response.status_code, ) # Handle both raw array format (TEI) and wrapped format (standard HF) diff --git a/litellm/llms/sambanova/chat.py b/litellm/llms/sambanova/chat.py index 2218c80872..3c4003f72e 100644 --- a/litellm/llms/sambanova/chat.py +++ b/litellm/llms/sambanova/chat.py @@ -117,10 +117,11 @@ class SambanovaConfig(OpenAIGPTConfig): ) -> Union[List[AllMessageValues], Coroutine[Any, Any, List[AllMessageValues]]]: """ Transform messages to handle content list conversion. - + SambaNova API doesn't support content as a list - only string content. This converts content lists like [{"type": "text", "text": "..."}] to strings. """ + async def _async_transform(): return handle_messages_with_content_list_to_str_conversion(messages) diff --git a/litellm/llms/sap/chat/handler.py b/litellm/llms/sap/chat/handler.py index 1390b2a478..713143d895 100755 --- a/litellm/llms/sap/chat/handler.py +++ b/litellm/llms/sap/chat/handler.py @@ -251,11 +251,8 @@ class AsyncSAPStreamIterator: # ------------------------------- class GenAIHubOrchestration(BaseLLMHTTPHandler): def _add_stream_param_to_request_body( - self, - data: dict, - provider_config: BaseConfig, - fake_stream: bool - ): + self, data: dict, provider_config: BaseConfig, fake_stream: bool + ): if data.get("config", {}).get("stream", None) is not None: data["config"]["stream"]["enabled"] = True else: diff --git a/litellm/llms/sap/chat/models.py b/litellm/llms/sap/chat/models.py index 1b09ce9a75..8ca2aa7a69 100644 --- a/litellm/llms/sap/chat/models.py +++ b/litellm/llms/sap/chat/models.py @@ -7,21 +7,22 @@ def validate_different_content(v: Union[str, dict, list]) -> str: if v in ((), {}, []): return "" elif isinstance(v, dict) and "text" in v: - return v['text'] + return v["text"] elif isinstance(v, list): new_v = [] for item in v: if isinstance(item, dict) and "text" in item: - if item['text']: - new_v.append(item['text']) + if item["text"]: + new_v.append(item["text"]) elif isinstance(item, str): new_v.append(item) - return '\n'.join(new_v) + return "\n".join(new_v) elif isinstance(v, str): return v raise ValueError("Content must be a string") return v + class TextContent(BaseModel): type_: Literal["text"] = Field(default="text", alias="type") text: str @@ -80,7 +81,9 @@ class SAPMessage(BaseModel): role: Literal["system", "developer"] = "system" content: str - _content_validator = field_validator("content", mode="before")(validate_different_content) + _content_validator = field_validator("content", mode="before")( + validate_different_content + ) class SAPUserMessage(BaseModel): @@ -96,8 +99,9 @@ class SAPAssistantMessage(BaseModel): refusal: str = "" tool_calls: list[MessageToolCall] = [] - _content_validator = field_validator("content", mode="before")(validate_different_content) - + _content_validator = field_validator("content", mode="before")( + validate_different_content + ) class SAPToolChatMessage(BaseModel): @@ -105,7 +109,9 @@ class SAPToolChatMessage(BaseModel): tool_call_id: str content: str - _content_validator = field_validator("content", mode="before")(validate_different_content) + _content_validator = field_validator("content", mode="before")( + validate_different_content + ) class ResponseFormat(BaseModel): diff --git a/litellm/llms/sap/chat/transformation.py b/litellm/llms/sap/chat/transformation.py index a019ba1767..7f6bab4a1d 100755 --- a/litellm/llms/sap/chat/transformation.py +++ b/litellm/llms/sap/chat/transformation.py @@ -1,7 +1,17 @@ """ Translate from OpenAI's `/v1/chat/completions` to SAP Generative AI Hub's Orchestration Service`v2/completion` """ -from typing import List, Optional, Union, Dict, Tuple, Any, TYPE_CHECKING, Iterator, AsyncIterator +from typing import ( + List, + Optional, + Union, + Dict, + Tuple, + Any, + TYPE_CHECKING, + Iterator, + AsyncIterator, +) from functools import cached_property import litellm import httpx @@ -29,7 +39,12 @@ from .models import ( ResponseFormat, SAPUserMessage, ) -from .handler import GenAIHubOrchestrationError, AsyncSAPStreamIterator, SAPStreamIterator +from .handler import ( + GenAIHubOrchestrationError, + AsyncSAPStreamIterator, + SAPStreamIterator, +) + def validate_dict(data: dict, model) -> dict: return model(**data).model_dump(by_alias=True) @@ -77,16 +92,15 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig): def run_env_setup(self, service_key: Optional[str] = None) -> None: try: - self.token_creator, self._base_url, self._resource_group = get_token_creator(service_key) # type: ignore + self.token_creator, self._base_url, self._resource_group = get_token_creator(service_key) # type: ignore except ValueError as err: raise GenAIHubOrchestrationError(status_code=400, message=err.args[0]) - @property def headers(self) -> Dict[str, str]: if self.token_creator is None: self.run_env_setup() - access_token = self.token_creator() # type: ignore + access_token = self.token_creator() # type: ignore return { "Authorization": access_token, "AI-Resource-Group": self.resource_group, @@ -98,14 +112,13 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig): def base_url(self) -> str: if self._base_url is None: self.run_env_setup() - return self._base_url # type: ignore - + return self._base_url # type: ignore @property def resource_group(self) -> str: if self._resource_group is None: self.run_env_setup() - return self._resource_group # type: ignore + return self._resource_group # type: ignore @cached_property def deployment_url(self) -> str: @@ -169,7 +182,6 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig): params.remove("tool_choice") return params - def validate_environment( self, headers: dict, @@ -185,13 +197,13 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig): return self.headers 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, + self, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: dict, + litellm_params: dict, + stream: Optional[bool] = None, ): api_base_ = f"{self.deployment_url}/v2/completion" return api_base_ @@ -199,7 +211,7 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig): def transform_request( self, model: str, - messages: List[Dict[str, str]], # type: ignore + messages: List[Dict[str, str]], # type: ignore optional_params: dict, litellm_params: dict, headers: dict, @@ -240,8 +252,10 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig): response_format = model_params.pop("response_format", {}) resp_type = response_format.get("type", None) if resp_type: - if resp_type== "json_schema": - response_format = validate_dict(response_format, ResponseFormatJSONSchema) + if resp_type == "json_schema": + response_format = validate_dict( + response_format, ResponseFormatJSONSchema + ) else: response_format = validate_dict(response_format, ResponseFormat) response_format = {"response_format": response_format} @@ -259,11 +273,7 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig): "config": { "modules": { "prompt_templating": { - "prompt": { - "template": template, - **tools, - **response_format - }, + "prompt": {"template": template, **tools, **response_format}, "model": { "name": model, "params": model_params, @@ -278,18 +288,18 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig): return config def transform_response( - 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, - encoding: Any, - api_key: Optional[str] = None, - json_mode: Optional[bool] = None, + 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, + encoding: Any, + api_key: Optional[str] = None, + json_mode: Optional[bool] = None, ) -> ModelResponse: logging_obj.post_call( input=messages, @@ -323,17 +333,17 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig): if choice.message and choice.message.content: content = choice.message.content.strip() # Match ```json ... ``` or ``` ... ``` - match = re.match(r'^```(?:json)?\s*\n?(.*?)\n?```$', content, re.DOTALL) + match = re.match(r"^```(?:json)?\s*\n?(.*?)\n?```$", content, re.DOTALL) if match: choice.message.content = match.group(1).strip() return response def get_model_response_iterator( - self, - streaming_response: Union[Iterator[str], AsyncIterator[str], "ModelResponse"], - sync_stream: bool, - json_mode: Optional[bool] = False, + self, + streaming_response: Union[Iterator[str], AsyncIterator[str], "ModelResponse"], + sync_stream: bool, + json_mode: Optional[bool] = False, ): if sync_stream: return SAPStreamIterator(response=streaming_response) # type: ignore diff --git a/litellm/llms/sap/credentials.py b/litellm/llms/sap/credentials.py index e10bcbf7ea..aeae51bf0b 100644 --- a/litellm/llms/sap/credentials.py +++ b/litellm/llms/sap/credentials.py @@ -180,7 +180,9 @@ def _resolve_value( return cred.default -def fetch_credentials(service_key: Optional[str] = None, profile: Optional[str] = None, **kwargs) -> Dict[str, str]: +def fetch_credentials( + service_key: Optional[str] = None, profile: Optional[str] = None, **kwargs +) -> Dict[str, str]: """ Resolution order per key: kwargs @@ -196,8 +198,11 @@ def fetch_credentials(service_key: Optional[str] = None, profile: Optional[str] if not config: # Prefer AICORE_SERVICE_KEY if present; otherwise fall back to the VCAP service. - service_like = service_key or sap_service_key or _load_json_env(SERVICE_KEY_ENV_VAR) or _get_vcap_service( - VCAP_AICORE_SERVICE_NAME + service_like = ( + service_key + or sap_service_key + or _load_json_env(SERVICE_KEY_ENV_VAR) + or _get_vcap_service(VCAP_AICORE_SERVICE_NAME) ) out: Dict[str, str] = {} @@ -241,7 +246,9 @@ def get_token_creator( """ # Resolve credentials using your helper - credentials: Dict[str, str] = fetch_credentials(service_key=service_key, profile=profile, **overrides) + credentials: Dict[str, str] = fetch_credentials( + service_key=service_key, profile=profile, **overrides + ) auth_url = credentials.get("auth_url") client_id = credentials.get("client_id") diff --git a/litellm/llms/searchapi/search/transformation.py b/litellm/llms/searchapi/search/transformation.py index f3333bb20c..92b2814018 100644 --- a/litellm/llms/searchapi/search/transformation.py +++ b/litellm/llms/searchapi/search/transformation.py @@ -19,6 +19,7 @@ from litellm.secret_managers.main import get_secret_str class _SearchAPIRequestRequired(TypedDict): """Required fields for SearchAPI.io request.""" + engine: str # Required - search engine (e.g., 'google') q: str # Required - search query @@ -28,6 +29,7 @@ class SearchAPIRequest(_SearchAPIRequestRequired, total=False): SearchAPI.io request format for Google Search. Based on: https://www.searchapi.io/docs/google """ + kgmid: str # Optional - Knowledge Graph identifier device: str # Optional - device type ('desktop', 'mobile', 'tablet') location: str # Optional - geographic location @@ -50,17 +52,17 @@ class SearchAPIRequest(_SearchAPIRequestRequired, total=False): class SearchAPIConfig(BaseSearchConfig): SEARCHAPI_API_BASE = "https://www.searchapi.io/api/v1/search" - + @staticmethod def ui_friendly_name() -> str: return "SearchAPI.io (Google Search)" - + def get_http_method(self) -> Literal["GET", "POST"]: """ SearchAPI.io uses GET requests for search. """ return "GET" - + def validate_environment( self, headers: Dict, @@ -72,14 +74,14 @@ class SearchAPIConfig(BaseSearchConfig): Validate environment and return headers. """ api_key = api_key or get_secret_str("SEARCHAPI_API_KEY") - + if not api_key: raise ValueError( "SEARCHAPI_API_KEY is not set. Set `SEARCHAPI_API_KEY` environment variable." ) - + headers["Content-Type"] = "application/json" - + return headers def get_complete_url( @@ -94,7 +96,9 @@ class SearchAPIConfig(BaseSearchConfig): SearchAPI.io uses GET requests and includes api_key in query params. """ - api_base = api_base or get_secret_str("SEARCHAPI_API_BASE") or self.SEARCHAPI_API_BASE + api_base = ( + api_base or get_secret_str("SEARCHAPI_API_BASE") or self.SEARCHAPI_API_BASE + ) # Build query parameters from the transformed request body if data and isinstance(data, dict) and "_searchapi_params" in data: @@ -197,7 +201,7 @@ class SearchAPIConfig(BaseSearchConfig): ) -> SearchResponse: """ Transform SearchAPI.io response to LiteLLM unified SearchResponse format. - + SearchAPI.io → LiteLLM mappings: - organic_results[].title → SearchResult.title - organic_results[].link → SearchResult.url @@ -215,7 +219,7 @@ class SearchAPIConfig(BaseSearchConfig): url = result.get("link", "") snippet = result.get("snippet", "") date = result.get("date") # SearchAPI.io provides date in some results - + search_result = SearchResult( title=title, url=url, diff --git a/litellm/llms/searxng/__init__.py b/litellm/llms/searxng/__init__.py index 91d237a8a0..f7ad1978c7 100644 --- a/litellm/llms/searxng/__init__.py +++ b/litellm/llms/searxng/__init__.py @@ -4,4 +4,3 @@ SearXNG API integration module. from litellm.llms.searxng.search.transformation import SearXNGSearchConfig __all__ = ["SearXNGSearchConfig"] - diff --git a/litellm/llms/searxng/search/__init__.py b/litellm/llms/searxng/search/__init__.py index cb6fccfa9d..88ac5dc629 100644 --- a/litellm/llms/searxng/search/__init__.py +++ b/litellm/llms/searxng/search/__init__.py @@ -4,4 +4,3 @@ SearXNG Search API module. from litellm.llms.searxng.search.transformation import SearXNGSearchConfig __all__ = ["SearXNGSearchConfig"] - diff --git a/litellm/llms/searxng/search/transformation.py b/litellm/llms/searxng/search/transformation.py index 00ad9d1948..bbd3b76501 100644 --- a/litellm/llms/searxng/search/transformation.py +++ b/litellm/llms/searxng/search/transformation.py @@ -18,6 +18,7 @@ from litellm.secret_managers.main import get_secret_str class _SearXNGSearchRequestRequired(TypedDict): """Required fields for SearXNG Search API request.""" + q: str # Required - search query @@ -26,6 +27,7 @@ class SearXNGSearchRequest(_SearXNGSearchRequestRequired, total=False): SearXNG Search API request format. Based on: https://docs.searxng.org/dev/search_api.html """ + categories: str # Optional - comma-separated list of categories engines: str # Optional - comma-separated list of engines language: str # Optional - language code @@ -35,17 +37,16 @@ class SearXNGSearchRequest(_SearXNGSearchRequestRequired, total=False): class SearXNGSearchConfig(BaseSearchConfig): - @staticmethod def ui_friendly_name() -> str: return "SearXNG" - + def get_http_method(self): """ SearXNG supports both GET and POST, but we'll use GET for simplicity. """ return "GET" - + def validate_environment( self, headers: Dict, @@ -74,27 +75,27 @@ class SearXNGSearchConfig(BaseSearchConfig): ) -> str: """ Get complete URL for Search endpoint with query parameters. - + SearXNG uses GET requests, so we build the full URL with query params here. The transformed request body (data) contains the parameters needed for the URL. """ from urllib.parse import urlencode - + api_base = api_base or get_secret_str("SEARXNG_API_BASE") - + if not api_base: raise ValueError( "SEARXNG_API_BASE is not set. Please set the `SEARXNG_API_BASE` environment variable " "or pass `api_base` parameter. Example: os.environ['SEARXNG_API_BASE'] = 'https://your-searxng-instance.com'" ) - + # Append "/search" to the api base if it's not already there if not api_base.endswith("/search"): if api_base.endswith("/"): api_base = f"{api_base}search" else: api_base = f"{api_base}/search" - + # Build query parameters from the transformed request body if data and isinstance(data, dict) and "_searxng_params" in data: params = data["_searxng_params"] @@ -102,7 +103,6 @@ class SearXNGSearchConfig(BaseSearchConfig): return f"{api_base}?{query_string}" return api_base - def transform_search_request( self, @@ -112,20 +112,20 @@ class SearXNGSearchConfig(BaseSearchConfig): ) -> Dict: """ Transform Search request to SearXNG API format. - + Transforms Perplexity unified spec parameters: - query → q - max_results → (handled via pageno, SearXNG returns ~20 results per page) - search_domain_filter → (not directly supported) - country → language (approximate mapping) - max_tokens_per_page → (not applicable, ignored) - + All other SearXNG-specific parameters are passed through as-is. - + Args: query: Search query (string or list of strings). SearXNG only supports single string queries. optional_params: Optional parameters for the request - + Returns: Dict with typed request data following SearXNGSearchRequest spec """ @@ -137,7 +137,7 @@ class SearXNGSearchConfig(BaseSearchConfig): "q": query, "format": "json", # Always request JSON format } - + # Transform Perplexity unified spec parameters to SearXNG format if "country" in optional_params: # Map country code to language (approximate) @@ -154,22 +154,25 @@ class SearXNGSearchConfig(BaseSearchConfig): request_data["language"] = "ja" else: request_data["language"] = country # Pass through as-is - + # Handle max_results via pagination (SearXNG returns ~20 results per page by default) # For simplicity, we'll just use page 1 and let SearXNG return its default number of results if "max_results" in optional_params: # Note: We could calculate pageno based on max_results, but for now we'll ignore this # and let SearXNG return its default results pass - + # Convert to dict before dynamic key assignments result_data = dict(request_data) - + # Pass through all other SearXNG-specific parameters as-is for param, value in optional_params.items(): - if param not in self.get_supported_perplexity_optional_params() and param not in result_data: + if ( + param not in self.get_supported_perplexity_optional_params() + and param not in result_data + ): result_data[param] = value - + # Store params in special key for GET request URL building # This will be used by get_complete_url to build the query string return {"_searxng_params": result_data} @@ -182,23 +185,23 @@ class SearXNGSearchConfig(BaseSearchConfig): ) -> SearchResponse: """ Transform SearXNG API response to LiteLLM unified SearchResponse format. - + SearXNG → LiteLLM mappings: - results[].title → SearchResult.title - results[].url → SearchResult.url - results[].content → SearchResult.snippet - results[].publishedDate OR results[].pubdate → SearchResult.date - No last_updated field in SearXNG response (set to None) - + Args: raw_response: Raw httpx response from SearXNG API logging_obj: Logging object for tracking - + Returns: SearchResponse with standardized format """ response_json = raw_response.json() - + # Transform results to SearchResult objects # Note: SearXNG doesn't natively support limiting results via API params # It returns ~20 results per page by default @@ -206,7 +209,7 @@ class SearXNGSearchConfig(BaseSearchConfig): for result in response_json.get("results", []): # Get date from either publishedDate or pubdate field date = result.get("publishedDate") or result.get("pubdate") - + search_result = SearchResult( title=result.get("title", ""), url=result.get("url", ""), @@ -215,9 +218,8 @@ class SearXNGSearchConfig(BaseSearchConfig): last_updated=None, # SearXNG doesn't provide last_updated in response ) results.append(search_result) - + return SearchResponse( results=results, object="search", ) - diff --git a/litellm/llms/serper/search/transformation.py b/litellm/llms/serper/search/transformation.py index 63526ea8ab..34e726dc77 100644 --- a/litellm/llms/serper/search/transformation.py +++ b/litellm/llms/serper/search/transformation.py @@ -18,6 +18,7 @@ from litellm.secret_managers.main import get_secret_str class _SerperSearchRequestRequired(TypedDict): """Required fields for Serper Search API request.""" + q: str # Required - search query @@ -26,6 +27,7 @@ class SerperSearchRequest(_SerperSearchRequestRequired, total=False): Serper Search API request format. Based on: https://serper.dev """ + num: int # Optional - number of results to return, default 10 page: int # Optional - page number (default 1) gl: str # Optional - country/geolocation code (e.g., "us", "gb") @@ -37,11 +39,11 @@ class SerperSearchRequest(_SerperSearchRequestRequired, total=False): class SerperSearchConfig(BaseSearchConfig): SERPER_API_BASE = "https://google.serper.dev" - + @staticmethod def ui_friendly_name() -> str: return "Serper" - + def validate_environment( self, headers: Dict, @@ -54,7 +56,9 @@ class SerperSearchConfig(BaseSearchConfig): """ api_key = api_key or get_secret_str("SERPER_API_KEY") if not api_key: - raise ValueError("SERPER_API_KEY is not set. Set `SERPER_API_KEY` environment variable.") + raise ValueError( + "SERPER_API_KEY is not set. Set `SERPER_API_KEY` environment variable." + ) headers["X-API-KEY"] = api_key headers["Content-Type"] = "application/json" return headers @@ -71,7 +75,7 @@ class SerperSearchConfig(BaseSearchConfig): """ api_base = api_base or get_secret_str("SERPER_API_BASE") or self.SERPER_API_BASE api_base = api_base.rstrip("/") - + if not api_base.endswith("/search"): api_base = f"{api_base}/search" @@ -85,14 +89,14 @@ class SerperSearchConfig(BaseSearchConfig): ) -> Dict: """ Transform Search request to Serper API format. - + Args: query: Search query (string or list of strings). Serper only supports single string queries. optional_params: Optional parameters for the request - max_results: Maximum number of search results -> maps to `num` - search_domain_filter: List of domains -> appended as site: clauses to `q` - country: Country code filter (e.g., 'US', 'GB') -> maps to `gl` (lowercased) - + Returns: Dict with typed request data following SerperSearchRequest spec """ @@ -102,27 +106,30 @@ class SerperSearchConfig(BaseSearchConfig): request_data: SerperSearchRequest = { "q": query, } - + if "max_results" in optional_params: request_data["num"] = optional_params["max_results"] - + if "country" in optional_params: request_data["gl"] = optional_params["country"].lower() - + if "search_domain_filter" in optional_params: domains = optional_params["search_domain_filter"] if isinstance(domains, list) and len(domains) > 0: domain_clauses = " OR ".join(f"site:{d}" for d in domains) request_data["q"] = f"({request_data['q']}) ({domain_clauses})" - + # Convert to dict before dynamic key assignments result_data = dict(request_data) - + # pass through all other parameters as-is for param, value in optional_params.items(): - if param not in self.get_supported_perplexity_optional_params() and param not in result_data: + if ( + param not in self.get_supported_perplexity_optional_params() + and param not in result_data + ): result_data[param] = value - + return result_data def transform_search_response( @@ -133,22 +140,22 @@ class SerperSearchConfig(BaseSearchConfig): ) -> SearchResponse: """ Transform Serper API response to LiteLLM unified SearchResponse format. - + Serper -> LiteLLM mappings: - organic[].title -> SearchResult.title - organic[].link -> SearchResult.url - organic[].snippet -> SearchResult.snippet - organic[].date -> SearchResult.date (optional, not always present) - + Args: raw_response: Raw httpx response from Serper API logging_obj: Logging object for tracking - + Returns: SearchResponse with standardized format """ response_json = raw_response.json() - + results = [] for result in response_json.get("organic", []): search_result = SearchResult( @@ -159,9 +166,8 @@ class SerperSearchConfig(BaseSearchConfig): last_updated=None, ) results.append(search_result) - + return SearchResponse( results=results, object="search", ) - diff --git a/litellm/llms/stability/image_edit/transformations.py b/litellm/llms/stability/image_edit/transformations.py index 53bdc825dd..eb400a2526 100644 --- a/litellm/llms/stability/image_edit/transformations.py +++ b/litellm/llms/stability/image_edit/transformations.py @@ -40,9 +40,7 @@ class StabilityImageEditConfig(BaseImageEditConfig): DEFAULT_BASE_URL: str = "https://api.stability.ai" - def get_supported_openai_params( - self, model: str - ) -> List[str]: + def get_supported_openai_params(self, model: str) -> List[str]: """ Return list of OpenAI params supported by Stability AI. @@ -52,7 +50,7 @@ class StabilityImageEditConfig(BaseImageEditConfig): "n", # Number of images (Stability always returns 1, we can loop) "size", # Maps to aspect_ratio "response_format", # b64_json or url (Stability only returns b64) - "mask" + "mask", ] def map_openai_params( @@ -188,7 +186,7 @@ class StabilityImageEditConfig(BaseImageEditConfig): data: Dict[str, Any] = { "output_format": "png", # Default to PNG } - + # Add prompt only if provided (some Stability endpoints don't require it) if prompt is not None and prompt != "": data["prompt"] = prompt @@ -241,7 +239,7 @@ class StabilityImageEditConfig(BaseImageEditConfig): "select_prompt", "control_strength", "composition_fidelity", - "change_strength" + "change_strength", ]: data[key] = value # type: ignore @@ -310,7 +308,9 @@ class StabilityImageEditConfig(BaseImageEditConfig): model_info = get_model_info(model, custom_llm_provider="stability") cost_per_image = model_info.get("output_cost_per_image", 0) if cost_per_image is not None: - model_response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = float(cost_per_image) + model_response._hidden_params["additional_headers"][ + "llm_provider-x-litellm-response-cost" + ] = float(cost_per_image) return model_response def use_multipart_form_data(self) -> bool: diff --git a/litellm/llms/stability/image_generation/transformation.py b/litellm/llms/stability/image_generation/transformation.py index d69dd399b2..ac63548bf5 100644 --- a/litellm/llms/stability/image_generation/transformation.py +++ b/litellm/llms/stability/image_generation/transformation.py @@ -80,9 +80,9 @@ class StabilityImageGenerationConfig(BaseImageGenerationConfig): if k in supported_params: # Map size to aspect_ratio if k == "size" and v in OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO: - optional_params["aspect_ratio"] = ( - OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO[v] - ) + optional_params[ + "aspect_ratio" + ] = OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO[v] elif k == "n": # Store n for later, but don't pass to Stability optional_params["_n"] = v @@ -132,9 +132,7 @@ class StabilityImageGenerationConfig(BaseImageGenerationConfig): Get the complete URL for the Stability AI API request. """ base_url: str = ( - api_base - or get_secret_str("STABILITY_API_BASE") - or self.DEFAULT_BASE_URL + api_base or get_secret_str("STABILITY_API_BASE") or self.DEFAULT_BASE_URL ) base_url = base_url.rstrip("/") diff --git a/litellm/llms/tavily/search/__init__.py b/litellm/llms/tavily/search/__init__.py index 4753928806..6e3fe1163c 100644 --- a/litellm/llms/tavily/search/__init__.py +++ b/litellm/llms/tavily/search/__init__.py @@ -4,4 +4,3 @@ Tavily Search API module. from litellm.llms.tavily.search.transformation import TavilySearchConfig __all__ = ["TavilySearchConfig"] - diff --git a/litellm/llms/tavily/search/transformation.py b/litellm/llms/tavily/search/transformation.py index 7fc33416a0..1228433b53 100644 --- a/litellm/llms/tavily/search/transformation.py +++ b/litellm/llms/tavily/search/transformation.py @@ -18,6 +18,7 @@ from litellm.secret_managers.main import get_secret_str class _TavilySearchRequestRequired(TypedDict): """Required fields for Tavily Search API request.""" + query: str # Required - search query @@ -26,6 +27,7 @@ class TavilySearchRequest(_TavilySearchRequestRequired, total=False): Tavily Search API request format. Based on: https://docs.tavily.com/documentation/api-reference/endpoint/search """ + max_results: int # Optional - maximum number of results (0-20), default 5 include_domains: List[str] # Optional - list of domains to include (max 300) exclude_domains: List[str] # Optional - list of domains to exclude (max 150) @@ -44,11 +46,11 @@ class TavilySearchRequest(_TavilySearchRequestRequired, total=False): class TavilySearchConfig(BaseSearchConfig): TAVILY_API_BASE = "https://api.tavily.com" - + @staticmethod def ui_friendly_name() -> str: return "Tavily" - + def validate_environment( self, headers: Dict, @@ -61,7 +63,9 @@ class TavilySearchConfig(BaseSearchConfig): """ api_key = api_key or get_secret_str("TAVILY_API_KEY") if not api_key: - raise ValueError("TAVILY_API_KEY is not set. Set `TAVILY_API_KEY` environment variable.") + raise ValueError( + "TAVILY_API_KEY is not set. Set `TAVILY_API_KEY` environment variable." + ) headers["Authorization"] = f"Bearer {api_key}" headers["Content-Type"] = "application/json" return headers @@ -77,13 +81,12 @@ class TavilySearchConfig(BaseSearchConfig): Get complete URL for Search endpoint. """ api_base = api_base or get_secret_str("TAVILY_API_BASE") or self.TAVILY_API_BASE - + # Append "/search" to the api base if it's not already there if not api_base.endswith("/search"): api_base = f"{api_base}/search" return api_base - def transform_search_request( self, @@ -93,7 +96,7 @@ class TavilySearchConfig(BaseSearchConfig): ) -> Dict: """ Transform Search request to Tavily API format. - + Args: query: Search query (string or list of strings). Tavily only supports single string queries. optional_params: Optional parameters for the request @@ -111,7 +114,7 @@ class TavilySearchConfig(BaseSearchConfig): - start_date: Start date filter (YYYY-MM-DD) - end_date: End date filter (YYYY-MM-DD) - country: Country code filter (e.g., 'US', 'GB', 'DE') - + Returns: Dict with typed request data following TavilySearchRequest spec """ @@ -122,26 +125,29 @@ class TavilySearchConfig(BaseSearchConfig): request_data: TavilySearchRequest = { "query": query, } - + # Transform Perplexity unified spec parameters to Tavily format if "max_results" in optional_params: request_data["max_results"] = optional_params["max_results"] - + if "search_domain_filter" in optional_params: request_data["include_domains"] = optional_params["search_domain_filter"] - + if "country" in optional_params: # Tavily expects lowercase country names request_data["country"] = optional_params["country"].lower() - + # Convert to dict before dynamic key assignments result_data = dict(request_data) - + # pass through all other parameters as-is for param, value in optional_params.items(): - if param not in self.get_supported_perplexity_optional_params() and param not in result_data: + if ( + param not in self.get_supported_perplexity_optional_params() + and param not in result_data + ): result_data[param] = value - + return result_data def transform_search_response( @@ -152,36 +158,37 @@ class TavilySearchConfig(BaseSearchConfig): ) -> SearchResponse: """ Transform Tavily API response to LiteLLM unified SearchResponse format. - + Tavily → LiteLLM mappings: - results[].title → SearchResult.title - results[].url → SearchResult.url - results[].content → SearchResult.snippet - No date/last_updated fields in Tavily response (set to None) - + Args: raw_response: Raw httpx response from Tavily API logging_obj: Logging object for tracking - + Returns: SearchResponse with standardized format """ response_json = raw_response.json() - + # Transform results to SearchResult objects results = [] for result in response_json.get("results", []): search_result = SearchResult( title=result.get("title", ""), url=result.get("url", ""), - snippet=result.get("content", ""), # Tavily uses "content" instead of "snippet" + snippet=result.get( + "content", "" + ), # Tavily uses "content" instead of "snippet" date=None, # Tavily doesn't provide date in response last_updated=None, # Tavily doesn't provide last_updated in response ) results.append(search_result) - + return SearchResponse( results=results, object="search", ) - diff --git a/litellm/llms/v0/chat/transformation.py b/litellm/llms/v0/chat/transformation.py index 1417e5f5ae..7b65cec9d3 100644 --- a/litellm/llms/v0/chat/transformation.py +++ b/litellm/llms/v0/chat/transformation.py @@ -13,7 +13,7 @@ class V0ChatConfig(OpenAILikeChatConfig): """ v0 is OpenAI-compatible with standard endpoints """ - + @property def custom_llm_provider(self) -> Optional[str]: return "v0" @@ -36,9 +36,9 @@ class V0ChatConfig(OpenAILikeChatConfig): Reference: https://v0.dev/docs/v0-model-api#request-body """ return [ - "messages", # Required - "model", # Required - "stream", # Optional - "tools", # Optional + "messages", # Required + "model", # Required + "stream", # Optional + "tools", # Optional "tool_choice", # Optional - ] \ No newline at end of file + ] diff --git a/litellm/llms/vercel_ai_gateway/chat/transformation.py b/litellm/llms/vercel_ai_gateway/chat/transformation.py index 13a8837748..81a1688b90 100644 --- a/litellm/llms/vercel_ai_gateway/chat/transformation.py +++ b/litellm/llms/vercel_ai_gateway/chat/transformation.py @@ -33,14 +33,13 @@ class VercelAIGatewayConfig(OpenAIGPTConfig): def _get_openai_compatible_provider_info( self, api_base: Optional[str], api_key: Optional[str] ) -> Tuple[Optional[str], Optional[str]]: - api_base = ( api_base or get_secret_str("VERCEL_AI_GATEWAY_API_BASE") or "https://ai-gateway.vercel.sh/v1" ) user_api_key = ( - api_key + api_key or get_secret_str("VERCEL_AI_GATEWAY_API_KEY") or get_secret_str("VERCEL_OIDC_TOKEN") ) @@ -60,11 +59,13 @@ class VercelAIGatewayConfig(OpenAIGPTConfig): # Vercel AI Gateway-only parameters extra_body = {} provider_options = non_default_params.pop("providerOptions", None) - + if provider_options is not None: extra_body["providerOptions"] = provider_options - - mapped_openai_params["extra_body"] = extra_body # openai client supports `extra_body` param + + mapped_openai_params[ + "extra_body" + ] = extra_body # openai client supports `extra_body` param return mapped_openai_params def transform_request( @@ -98,10 +99,10 @@ class VercelAIGatewayConfig(OpenAIGPTConfig): self, api_key: Optional[str] = None, api_base: Optional[str] = None ) -> List[str]: api_base, _ = self._get_openai_compatible_provider_info(api_base, api_key) - + if api_base is None: api_base = "https://ai-gateway.vercel.sh/v1" - + models_url = f"{api_base}/models" response = litellm.module_level_client.get(url=models_url) diff --git a/litellm/llms/vertex_ai/agent_engine/__init__.py b/litellm/llms/vertex_ai/agent_engine/__init__.py index de891f8560..a790d8b6bd 100644 --- a/litellm/llms/vertex_ai/agent_engine/__init__.py +++ b/litellm/llms/vertex_ai/agent_engine/__init__.py @@ -10,4 +10,3 @@ from litellm.llms.vertex_ai.agent_engine.transformation import ( ) __all__ = ["VertexAgentEngineConfig", "VertexAgentEngineError"] - diff --git a/litellm/llms/vertex_ai/agent_engine/transformation.py b/litellm/llms/vertex_ai/agent_engine/transformation.py index 42032079f9..0707a7b4c2 100644 --- a/litellm/llms/vertex_ai/agent_engine/transformation.py +++ b/litellm/llms/vertex_ai/agent_engine/transformation.py @@ -120,7 +120,9 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase): # Get project and location from litellm_params or environment vertex_project = self.safe_get_vertex_ai_project(litellm_params) - vertex_location = self.safe_get_vertex_ai_location(litellm_params) or "us-central1" + vertex_location = ( + self.safe_get_vertex_ai_location(litellm_params) or "us-central1" + ) # Build the full resource path if only engine_id was provided if not resource_path.startswith("projects/"): @@ -156,7 +158,9 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase): project_id=vertex_project, ) - verbose_logger.debug(f"Vertex Agent Engine: Authenticated for project {project_id}") + verbose_logger.debug( + f"Vertex Agent Engine: Authenticated for project {project_id}" + ) return { "Authorization": f"Bearer {access_token}", @@ -300,7 +304,9 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase): """ try: content_type = raw_response.headers.get("content-type", "").lower() - verbose_logger.debug(f"Vertex Agent Engine response Content-Type: {content_type}") + verbose_logger.debug( + f"Vertex Agent Engine response Content-Type: {content_type}" + ) # Parse the SSE response response_text = raw_response.text @@ -340,7 +346,9 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase): return model_response except Exception as e: - verbose_logger.error(f"Error processing Vertex Agent Engine response: {str(e)}") + verbose_logger.error( + f"Error processing Vertex Agent Engine response: {str(e)}" + ) raise VertexAgentEngineError( message=f"Error processing response: {str(e)}", status_code=raw_response.status_code, @@ -398,7 +406,9 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase): ) # Create iterator for SSE stream - completion_stream = self.get_streaming_response(model=model, raw_response=response) + completion_stream = self.get_streaming_response( + model=model, raw_response=response + ) streaming_response = CustomStreamWrapper( completion_stream=completion_stream, @@ -505,4 +515,3 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase): ) -> bool: """Agent Engine always returns SSE streams, so we use real streaming.""" return False - diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 5f1fefca96..f0b181c9a6 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -118,7 +118,8 @@ class VertexAIBatchPrediction(VertexLLM): error_body = e.response.text litellm.verbose_logger.error( "Vertex AI batch create failed: status=%s, body=%s", - e.response.status_code, error_body[:1000], + e.response.status_code, + error_body[:1000], ) raise if response.status_code != 200: @@ -202,6 +203,7 @@ class VertexAIBatchPrediction(VertexLLM): # Log the request using logging_obj if available if logging_obj is not None: from litellm.litellm_core_utils.litellm_logging import Logging + if isinstance(logging_obj, Logging): logging_obj.pre_call( input="", @@ -243,10 +245,11 @@ class VertexAIBatchPrediction(VertexLLM): client = get_async_httpx_client( llm_provider=litellm.LlmProviders.VERTEX_AI, ) - + # Log the request using logging_obj if available if logging_obj is not None: from litellm.litellm_core_utils.litellm_logging import Logging + if isinstance(logging_obj, Logging): logging_obj.pre_call( input="", @@ -264,7 +267,7 @@ class VertexAIBatchPrediction(VertexLLM): ), }, ) - + response = await client.get( url=api_base, headers=headers, diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index c02d63414c..5895a91f3a 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -39,8 +39,10 @@ class VertexAIModelRoute(str, Enum): OPENAI_COMPATIBLE = "openai" AGENT_ENGINE = "agent_engine" + VERTEX_AI_MODEL_ROUTES = [f"{route.value}/" for route in VertexAIModelRoute] + def get_vertex_ai_model_route( model: str, litellm_params: Optional[dict] = None ) -> VertexAIModelRoute: @@ -66,7 +68,7 @@ def get_vertex_ai_model_route( >>> get_vertex_ai_model_route("openai/gpt-oss-120b") VertexAIModelRoute.MODEL_GARDEN - + >>> get_vertex_ai_model_route("1234567890", {"api_base": "http://10.96.32.8"}) VertexAIModelRoute.GEMINI # Numeric endpoints with api_base use HTTP path """ @@ -82,20 +84,20 @@ def get_vertex_ai_model_route( # Check for agent_engine models (Reasoning Engines) if "agent_engine/" in model: return VertexAIModelRoute.AGENT_ENGINE - + # Check if numeric endpoint ID with custom api_base (PSC endpoint) # Route to GEMINI (HTTP path) to support PSC endpoints properly if model.isdigit() and litellm_params and litellm_params.get("api_base"): return VertexAIModelRoute.GEMINI - + # Check for partner models (llama, mistral, claude, etc.) if VertexAIPartnerModels.is_vertex_partner_model(model=model): return VertexAIModelRoute.PARTNER_MODELS - + # Check for BGE models if "bge/" in model or "bge" in model.lower(): return VertexAIModelRoute.BGE - + # Check for gemma models if "gemma/" in model: return VertexAIModelRoute.GEMMA @@ -189,27 +191,27 @@ all_gemini_url_modes = Literal[ def get_vertex_base_model_name(model: str) -> str: """ Strip routing prefixes from model name for PSC/endpoint URL construction. - - Patterns like "bge/", "gemma/", "openai/" are used for internal routing but + + Patterns like "bge/", "gemma/", "openai/" are used for internal routing but should not appear in the actual endpoint URL. Routing prefixes are derived from VertexAIModelRoute enum values. - + Args: model: The model name with potential prefix (e.g., "bge/123456", "gemma/gemma-3-12b-it") - + Returns: str: The model name without routing prefix (e.g., "123456", "gemma-3-12b-it") - + Examples: >>> get_vertex_base_model_name("bge/378943383978115072") "378943383978115072" - + >>> get_vertex_base_model_name("gemma/gemma-3-12b-it") "gemma-3-12b-it" - + >>> get_vertex_base_model_name("openai/gpt-oss-120b") "gpt-oss-120b" - + >>> get_vertex_base_model_name("1234567890") "1234567890" """ @@ -218,7 +220,7 @@ def get_vertex_base_model_name(model: str) -> str: for route in VERTEX_AI_MODEL_ROUTES: if model.startswith(route): return model.replace(route, "", 1) - + return model @@ -242,16 +244,16 @@ def _get_embedding_url( ) -> Tuple[str, str]: """ Get URL for embedding models. - + Handles special patterns: - bge/endpoint_id -> strips to endpoint_id for endpoints/ routing - numeric model -> routes to endpoints/ - regular model -> routes to publishers/google/models/ - models with uses_embed_content flag -> use embedContent endpoint instead of predict - """ + """ original_model = model model = get_vertex_base_model_name(model=model) - + try: model_info = litellm.get_model_info( model=original_model, @@ -260,16 +262,16 @@ def _get_embedding_url( uses_embed_content = model_info.get("uses_embed_content", False) except Exception: uses_embed_content = False - + endpoint = "embedContent" if uses_embed_content else "predict" - + base_url = get_vertex_base_url(vertex_location) - + if model.isdigit(): url = f"{base_url}/{vertex_api_version}/projects/{vertex_project}/locations/{vertex_location}/endpoints/{model}:{endpoint}" else: url = f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:{endpoint}" - + return url, endpoint @@ -285,15 +287,15 @@ def _get_vertex_url( endpoint: Optional[str] = None model = litellm.VertexGeminiConfig.get_model_for_vertex_ai_url(model=model) - + if mode == "chat": ### SET RUNTIME ENDPOINT ### endpoint = "generateContent" base_url = get_vertex_base_url(vertex_location) - + if stream is True: endpoint = "streamGenerateContent" - + # if model is only numeric chars then it's a fine tuned gemini model # model = 4965075652664360960 # send to this url: url = f"{base_url}/{version}/projects/{vertex_project}/locations/{vertex_location}/endpoints/{model}:{endpoint}" @@ -303,7 +305,7 @@ def _get_vertex_url( else: # Regular model - use publishers/google/models/ path url = f"{base_url}/{vertex_api_version}/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:{endpoint}" - + if stream is True: url += "?alt=sse" elif mode == "embedding": @@ -340,10 +342,12 @@ def _get_gemini_url( from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) - + _gemini_model_name = "models/{}".format(model) - api_version = "v1alpha" if VertexGeminiConfig._is_gemini_3_or_newer(model) else "v1beta" - + api_version = ( + "v1alpha" if VertexGeminiConfig._is_gemini_3_or_newer(model) else "v1beta" + ) + if mode == "chat": endpoint = "generateContent" if stream is True: @@ -352,10 +356,8 @@ def _get_gemini_url( api_version, _gemini_model_name, endpoint, gemini_api_key ) else: - url = ( - "https://generativelanguage.googleapis.com/{}/{}:{}?key={}".format( - api_version, _gemini_model_name, endpoint, gemini_api_key - ) + url = "https://generativelanguage.googleapis.com/{}/{}:{}?key={}".format( + api_version, _gemini_model_name, endpoint, gemini_api_key ) elif mode == "embedding": endpoint = "embedContent" @@ -717,7 +719,12 @@ def convert_anyof_null_to_nullable(schema, depth=0): def add_object_type(schema): # Gemini requires all function parameters to be type OBJECT # Handle case where schema has no properties and no type (e.g. tools with no arguments) - if "type" not in schema and "anyOf" not in schema and "oneOf" not in schema and "allOf" not in schema: + if ( + "type" not in schema + and "anyOf" not in schema + and "oneOf" not in schema + and "allOf" not in schema + ): schema["type"] = "object" properties = schema.get("properties", None) @@ -794,10 +801,19 @@ def _convert_schema_types(schema, depth=0): if "type" in schema: type_val = schema["type"] if isinstance(type_val, list) and len(type_val) > 1: - # Convert type arrays to anyOf format + # Convert type arrays to anyOf format # Fields that are specific to object/array types and should move into anyOf - type_specific_fields = {"properties", "required", "additionalProperties", "items", "minItems", "maxItems", "minProperties", "maxProperties"} - + type_specific_fields = { + "properties", + "required", + "additionalProperties", + "items", + "minItems", + "maxItems", + "minProperties", + "maxProperties", + } + any_of: List[Dict[str, Any]] = [] for t in type_val: if not isinstance(t, str): @@ -806,7 +822,7 @@ def _convert_schema_types(schema, depth=0): # Keep null entry minimal so we can strip it later. any_of.append({"type": "null"}) continue - + # For object/array types, include type-specific fields if t in ("object", "array"): item_schema = {"type": t} @@ -818,13 +834,15 @@ def _convert_schema_types(schema, depth=0): else: # For primitive types, only include the type any_of.append({"type": t}) - + # Remove type-specific fields from parent if we moved them into anyOf - has_object_or_array = any(t in ("object", "array") for t in type_val if isinstance(t, str)) + has_object_or_array = any( + t in ("object", "array") for t in type_val if isinstance(t, str) + ) if has_object_or_array: for field in type_specific_fields: schema.pop(field, None) - + schema["anyOf"] = any_of schema.pop("type") elif isinstance(type_val, list) and len(type_val) == 1: @@ -944,26 +962,6 @@ def construct_target_url( return updated_url -def is_global_only_vertex_model(model: str) -> bool: - """ - Check if a model is only available in the global region. - - Args: - model: The model name to check - - Returns: - True if the model is only available in global region, False otherwise - """ - from litellm.utils import get_supported_regions - - supported_regions = get_supported_regions( - model=model, custom_llm_provider="vertex_ai" - ) - if supported_regions is None: - return False - return "global" in supported_regions - - class VertexAIModelInfo(BaseLLMModelInfo): def get_token_counter(self) -> Optional[BaseTokenCounter]: """ @@ -1057,15 +1055,16 @@ class VertexAITokenCounter(BaseTokenCounter): vertex_project = count_tokens_params_request.get( "vertex_project" ) or count_tokens_params_request.get("vertex_ai_project") - + vertex_location = count_tokens_params_request.get( "vertex_location" ) or count_tokens_params_request.get("vertex_ai_location") # Count tokens not available on global location: https://docs.cloud.google.com/vertex-ai/generative-ai/docs/partner-models/claude/count-tokens - vertex_location = count_tokens_params_request.get( - "vertex_count_tokens_location" - ) or vertex_location + vertex_location = ( + count_tokens_params_request.get("vertex_count_tokens_location") + or vertex_location + ) vertex_credentials = count_tokens_params_request.get( "vertex_credentials" @@ -1110,4 +1109,4 @@ class VertexAITokenCounter(BaseTokenCounter): original_response=result, ) - return None \ No newline at end of file + return None diff --git a/litellm/llms/vertex_ai/context_caching/transformation.py b/litellm/llms/vertex_ai/context_caching/transformation.py index bc5c1b451f..950edbeb47 100644 --- a/litellm/llms/vertex_ai/context_caching/transformation.py +++ b/litellm/llms/vertex_ai/context_caching/transformation.py @@ -51,37 +51,37 @@ def get_first_continuous_block_idx( def extract_ttl_from_cached_messages(messages: List[AllMessageValues]) -> Optional[str]: """ Extract TTL from cached messages. Returns the first valid TTL found. - + Args: messages: List of messages to extract TTL from - + Returns: Optional[str]: TTL string in format "3600s" or None if not found/invalid """ for message in messages: if not is_cached_message(message): continue - + content = message.get("content") if not content or isinstance(content, str): continue - + for content_item in content: # Type check to ensure content_item is a dictionary before calling .get() if not isinstance(content_item, dict): continue - + cache_control = content_item.get("cache_control") if not cache_control or not isinstance(cache_control, dict): continue - + if cache_control.get("type") != "ephemeral": continue - + ttl = cache_control.get("ttl") if ttl and _is_valid_ttl_format(ttl): return str(ttl) - + return None @@ -89,23 +89,23 @@ def _is_valid_ttl_format(ttl: str) -> bool: """ Validate TTL format. Should be a string ending with 's' for seconds. Examples: "3600s", "7200s", "1.5s" - + Args: ttl: TTL string to validate - + Returns: bool: True if valid format, False otherwise """ if not isinstance(ttl, str): return False - + # TTL should end with 's' and contain a valid number before it - pattern = r'^([0-9]*\.?[0-9]+)s$' + pattern = r"^([0-9]*\.?[0-9]+)s$" match = re.match(pattern, ttl) - + if not match: return False - + try: # Ensure the numeric part is valid and positive numeric_part = float(match.group(1)) @@ -164,7 +164,7 @@ def transform_openai_messages_to_gemini_context_caching( ) -> CachedContentRequestBody: # Extract TTL from cached messages BEFORE system message transformation ttl = extract_ttl_from_cached_messages(messages) - + supports_system_message = get_supports_system_message( model=model, custom_llm_provider=custom_llm_provider ) @@ -173,8 +173,10 @@ def transform_openai_messages_to_gemini_context_caching( supports_system_message=supports_system_message, messages=messages ) - transformed_messages = _gemini_convert_messages_with_history(messages=new_messages, model=model) - + transformed_messages = _gemini_convert_messages_with_history( + messages=new_messages, model=model + ) + model_name = "models/{}".format(model) if custom_llm_provider == "vertex_ai" or custom_llm_provider == "vertex_ai_beta": @@ -185,11 +187,11 @@ def transform_openai_messages_to_gemini_context_caching( model=model_name, displayName=cache_key, ) - + # Add TTL if present and valid if ttl: data["ttl"] = ttl - + if transformed_system_messages is not None: data["system_instruction"] = transformed_system_messages diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index 4450ae5834..db6be9499a 100644 --- a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py +++ b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py @@ -81,7 +81,6 @@ class ContextCachingEndpoints(VertexBase): else: url = f"https://{vertex_location}-aiplatform.googleapis.com/v1beta1/projects/{vertex_project}/locations/{vertex_location}/{endpoint}" - return self._check_custom_proxy( api_base=api_base, custom_llm_provider=custom_llm_provider, @@ -93,7 +92,9 @@ class ContextCachingEndpoints(VertexBase): model=None, vertex_project=vertex_project, vertex_location=vertex_location, - vertex_api_version="v1beta1" if custom_llm_provider == "vertex_ai_beta" else "v1", + vertex_api_version="v1beta1" + if custom_llm_provider == "vertex_ai_beta" + else "v1", ) def check_cache( @@ -126,7 +127,7 @@ class ContextCachingEndpoints(VertexBase): api_base=api_base, vertex_project=vertex_project, vertex_location=vertex_location, - vertex_auth_header=vertex_auth_header + vertex_auth_header=vertex_auth_header, ) page_token: Optional[str] = None @@ -199,7 +200,7 @@ class ContextCachingEndpoints(VertexBase): custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"], vertex_project: Optional[str], vertex_location: Optional[str], - vertex_auth_header: Optional[str] + vertex_auth_header: Optional[str], ) -> Optional[str]: """ Checks if content already cached. @@ -218,7 +219,7 @@ class ContextCachingEndpoints(VertexBase): api_base=api_base, vertex_project=vertex_project, vertex_location=vertex_location, - vertex_auth_header=vertex_auth_header + vertex_auth_header=vertex_auth_header, ) page_token: Optional[str] = None @@ -340,7 +341,7 @@ class ContextCachingEndpoints(VertexBase): api_base=api_base, vertex_project=vertex_project, vertex_location=vertex_location, - vertex_auth_header=vertex_auth_header + vertex_auth_header=vertex_auth_header, ) headers = { @@ -375,7 +376,7 @@ class ContextCachingEndpoints(VertexBase): custom_llm_provider=custom_llm_provider, vertex_project=vertex_project, vertex_location=vertex_location, - vertex_auth_header=vertex_auth_header + vertex_auth_header=vertex_auth_header, ) if google_cache_name: return non_cached_messages, optional_params, google_cache_name @@ -486,7 +487,7 @@ class ContextCachingEndpoints(VertexBase): api_base=api_base, vertex_project=vertex_project, vertex_location=vertex_location, - vertex_auth_header=vertex_auth_header + vertex_auth_header=vertex_auth_header, ) headers = { @@ -518,7 +519,7 @@ class ContextCachingEndpoints(VertexBase): custom_llm_provider=custom_llm_provider, vertex_project=vertex_project, vertex_location=vertex_location, - vertex_auth_header=vertex_auth_header + vertex_auth_header=vertex_auth_header, ) if google_cache_name: @@ -574,4 +575,4 @@ class ContextCachingEndpoints(VertexBase): pass async def async_get_cache(self): - pass \ No newline at end of file + pass diff --git a/litellm/llms/vertex_ai/count_tokens/handler.py b/litellm/llms/vertex_ai/count_tokens/handler.py index d95c6801e5..9a175371a2 100644 --- a/litellm/llms/vertex_ai/count_tokens/handler.py +++ b/litellm/llms/vertex_ai/count_tokens/handler.py @@ -17,7 +17,9 @@ class VertexAITokenCounter(GoogleAIStudioTokenCounter, VertexBase): Returns a Tuple of headers and url for the Vertex AI countTokens endpoint. """ litellm_params = litellm_params or {} - vertex_credentials = self.get_vertex_ai_credentials(litellm_params=litellm_params) + vertex_credentials = self.get_vertex_ai_credentials( + litellm_params=litellm_params + ) vertex_project = self.get_vertex_ai_project(litellm_params=litellm_params) vertex_location = self.get_vertex_ai_location(litellm_params=litellm_params) should_use_v1beta1_features = self.is_using_v1beta1_features(litellm_params) @@ -43,4 +45,4 @@ class VertexAITokenCounter(GoogleAIStudioTokenCounter, VertexBase): headers = { "Authorization": f"Bearer {auth_header}", } - return headers, api_base \ No newline at end of file + return headers, api_base diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index bf3ed5e6ac..070ec50828 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -165,7 +165,11 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): """ Get the complete url for the request """ - bucket_name = litellm_params.get("bucket_name") or litellm_params.get("litellm_metadata", {}).pop("gcs_bucket_name", None) or os.getenv("GCS_BUCKET_NAME") + bucket_name = ( + litellm_params.get("bucket_name") + or litellm_params.get("litellm_metadata", {}).pop("gcs_bucket_name", None) + or os.getenv("GCS_BUCKET_NAME") + ) if not bucket_name: raise ValueError("GCS bucket_name is required") file_data = data.get("file") @@ -410,6 +414,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): url = str(raw_response.request.url) if "/b/" in url and "/o/" in url: import urllib.parse + bucket_part = url.split("/b/")[-1].split("/o/")[0] encoded_name = url.split("/o/")[-1].split("?")[0] file_id = f"gs://{bucket_part}/{urllib.parse.unquote(encoded_name)}" diff --git a/litellm/llms/vertex_ai/fine_tuning/handler.py b/litellm/llms/vertex_ai/fine_tuning/handler.py index e2cd052fff..77891e245c 100644 --- a/litellm/llms/vertex_ai/fine_tuning/handler.py +++ b/litellm/llms/vertex_ai/fine_tuning/handler.py @@ -65,9 +65,9 @@ class VertexFineTuningAPI(VertexLLM): ) if create_fine_tuning_job_data.validation_file: - supervised_tuning_spec["validation_dataset"] = ( - create_fine_tuning_job_data.validation_file - ) + supervised_tuning_spec[ + "validation_dataset" + ] = create_fine_tuning_job_data.validation_file _vertex_hyperparameters = ( self._transform_openai_hyperparameters_to_vertex_hyperparameters( @@ -332,7 +332,7 @@ class VertexFineTuningAPI(VertexLLM): } base_url = get_vertex_base_url(vertex_location) - + url = None if request_route == "/tuningJobs": url = f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}/tuningJobs" @@ -349,9 +349,9 @@ class VertexFineTuningAPI(VertexLLM): elif "cachedContents" in request_route: _model = request_data.get("model") if _model is not None and "/publishers/google/models/" not in _model: - request_data["model"] = ( - f"projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{_model}" - ) + request_data[ + "model" + ] = f"projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{_model}" url = f"{base_url}/v1beta1/projects/{vertex_project}/locations/{vertex_location}{request_route}" else: diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 48477f2f3a..d7b96b4db7 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -335,7 +335,9 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 image_url = img_element["image_url"]["url"] format = img_element["image_url"].get("format") detail = img_element["image_url"].get("detail") - media_resolution_enum = _convert_detail_to_media_resolution_enum(detail) + media_resolution_enum = ( + _convert_detail_to_media_resolution_enum(detail) + ) else: image_url = img_element["image_url"] _part = _process_gemini_media( @@ -384,7 +386,9 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 ) # Convert detail to media_resolution_enum - media_resolution_enum = _convert_detail_to_media_resolution_enum(detail) + media_resolution_enum = ( + _convert_detail_to_media_resolution_enum(detail) + ) try: _part = _process_gemini_media( @@ -402,10 +406,7 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 ) ) user_content.extend(_parts) - elif ( - _message_content is not None - and isinstance(_message_content, str) - ): + elif _message_content is not None and isinstance(_message_content, str): _part = PartType(text=_message_content) user_content.append(_part) @@ -473,19 +474,26 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 _parts.append(_part) assistant_content.extend(_parts) - elif ( - _message_content is not None - and isinstance(_message_content, str) - ): + elif _message_content is not None and isinstance(_message_content, str): assistant_text = _message_content # Check if message has thought_signatures in provider_specific_fields - provider_specific_fields = assistant_msg.get("provider_specific_fields") + provider_specific_fields = assistant_msg.get( + "provider_specific_fields" + ) thought_signatures = None - if provider_specific_fields and isinstance(provider_specific_fields, dict): - thought_signatures = provider_specific_fields.get("thought_signatures") - + if provider_specific_fields and isinstance( + provider_specific_fields, dict + ): + thought_signatures = provider_specific_fields.get( + "thought_signatures" + ) + # If we have thought signatures, add them to the part - if thought_signatures and isinstance(thought_signatures, list) and len(thought_signatures) > 0: + if ( + thought_signatures + and isinstance(thought_signatures, list) + and len(thought_signatures) > 0 + ): # Use the first signature for the text part (Gemini expects one signature per part) assistant_content.append(PartType(text=assistant_text, thoughtSignature=thought_signatures[0])) # type: ignore else: @@ -502,7 +510,9 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 assistant_image_url = image_url_obj.get("url") format = image_url_obj.get("format") detail = image_url_obj.get("detail") - media_resolution_enum = _convert_detail_to_media_resolution_enum(detail) + media_resolution_enum = ( + _convert_detail_to_media_resolution_enum(detail) + ) if assistant_image_url: _part = _process_gemini_media( image_url=assistant_image_url, @@ -583,13 +593,23 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 raise e +# Keys that LiteLLM consumes internally and must never be forwarded to the +_LITELLM_INTERNAL_EXTRA_BODY_KEYS: frozenset = frozenset({"cache", "tags"}) + + def _pop_and_merge_extra_body(data: RequestBody, optional_params: dict) -> None: """Pop extra_body from optional_params and shallow-merge into data, deep-merging dict values.""" extra_body: Optional[dict] = optional_params.pop("extra_body", None) if extra_body is not None: data_dict: dict = data # type: ignore[assignment] for k, v in extra_body.items(): - if k in data_dict and isinstance(data_dict[k], dict) and isinstance(v, dict): + if k in _LITELLM_INTERNAL_EXTRA_BODY_KEYS: + continue + if ( + k in data_dict + and isinstance(data_dict[k], dict) + and isinstance(v, dict) + ): data_dict[k].update(v) else: data_dict[k] = v @@ -665,7 +685,9 @@ def _transform_request_body( # noqa: PLR0915 labels = {k: v for k, v in rm.items() if isinstance(v, str)} filtered_params = { - k: v for k, v in optional_params.items() if _get_equivalent_key(k, set(config_fields)) + k: v + for k, v in optional_params.items() + if _get_equivalent_key(k, set(config_fields)) } generation_config: Optional[GenerationConfig] = GenerationConfig( @@ -682,7 +704,9 @@ def _transform_request_body( # noqa: PLR0915 max_media_resolution ) if media_resolution_value and generation_config is not None: - generation_config["mediaResolution"] = media_resolution_value["level"] + generation_config["mediaResolution"] = media_resolution_value[ + "level" + ] data = RequestBody(contents=content) if system_instructions is not None: @@ -728,9 +752,9 @@ def sync_transform_request_body( context_caching_endpoints = ContextCachingEndpoints() ( - messages, - optional_params, - cached_content, + messages, + optional_params, + cached_content, ) = context_caching_endpoints.check_and_create_cache( messages=messages, optional_params=optional_params, @@ -748,7 +772,6 @@ def sync_transform_request_body( vertex_auth_header=vertex_auth_header, ) - return _transform_request_body( messages=messages, model=model, @@ -780,9 +803,9 @@ async def async_transform_request_body( context_caching_endpoints = ContextCachingEndpoints() ( - messages, - optional_params, - cached_content, + messages, + optional_params, + cached_content, ) = await context_caching_endpoints.async_check_and_create_cache( messages=messages, optional_params=optional_params, diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 6cd430d6cb..3f1bccaccf 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -498,9 +498,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): value = _remove_strict_from_schema(value) for tool in value: - openai_function_object: Optional[ChatCompletionToolParamFunctionChunk] = ( - None - ) + openai_function_object: Optional[ + ChatCompletionToolParamFunctionChunk + ] = None if "function" in tool: # tools list _openai_function_object = ChatCompletionToolParamFunctionChunk( # type: ignore **tool["function"] @@ -632,15 +632,15 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): _tools_list.append(search_tool) if googleSearchRetrieval is not None: retrieval_tool = Tools() - retrieval_tool[VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value] = ( - googleSearchRetrieval - ) + retrieval_tool[ + VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value + ] = googleSearchRetrieval _tools_list.append(retrieval_tool) if enterpriseWebSearch is not None: enterprise_tool = Tools() - enterprise_tool[VertexToolName.ENTERPRISE_WEB_SEARCH.value] = ( - enterpriseWebSearch - ) + enterprise_tool[ + VertexToolName.ENTERPRISE_WEB_SEARCH.value + ] = enterpriseWebSearch _tools_list.append(enterprise_tool) if code_execution is not None: code_tool = Tools() @@ -802,12 +802,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): # Check if this is gemini-3-flash which supports MINIMAL thinking level # Covers gemini-3-flash, gemini-3-flash-preview, gemini-3.1-flash, gemini-3.1-flash-lite-preview, etc. is_gemini3flash = model and ( - "gemini-3-flash" in model.lower() - or "gemini-3.1-flash" in model.lower() - ) - is_gemini31pro = model and ( - "gemini-3.1-pro-preview" in model.lower() + "gemini-3-flash" in model.lower() or "gemini-3.1-flash" in model.lower() ) + is_gemini31pro = model and ("gemini-3.1-pro-preview" in model.lower()) if reasoning_effort == "minimal": if is_gemini3flash: return {"thinkingLevel": "minimal", "includeThoughts": True} @@ -1090,16 +1087,16 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): param_description="thinking_budget", ) if VertexGeminiConfig._is_gemini_3_or_newer(model): - optional_params["thinkingConfig"] = ( - VertexGeminiConfig._map_reasoning_effort_to_thinking_level( - effort_value, model - ) + optional_params[ + "thinkingConfig" + ] = VertexGeminiConfig._map_reasoning_effort_to_thinking_level( + effort_value, model ) else: - optional_params["thinkingConfig"] = ( - VertexGeminiConfig._map_reasoning_effort_to_thinking_budget( - effort_value, model - ) + optional_params[ + "thinkingConfig" + ] = VertexGeminiConfig._map_reasoning_effort_to_thinking_budget( + effort_value, model ) elif param == "thinking": # Validate no conflict with thinking_level @@ -1108,11 +1105,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): param_name="thinking", param_description="thinking_budget", ) - optional_params["thinkingConfig"] = ( - VertexGeminiConfig._map_thinking_param( - cast(AnthropicThinkingParam, value), - model=model, - ) + optional_params[ + "thinkingConfig" + ] = VertexGeminiConfig._map_thinking_param( + cast(AnthropicThinkingParam, value), + model=model, ) elif param == "modalities" and isinstance(value, list): response_modalities = self.map_response_modalities(value) @@ -1230,12 +1227,25 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "IMAGE_PROHIBITED_CONTENT": "The token generation was stopped as the response was flagged for prohibited image content.", } - _GEMINI_FINISH_REASON_KEYS = frozenset({ - "STOP", "MAX_TOKENS", "SAFETY", "RECITATION", "FINISH_REASON_UNSPECIFIED", - "MALFORMED_FUNCTION_CALL", "LANGUAGE", "OTHER", "BLOCKLIST", - "PROHIBITED_CONTENT", "SPII", "IMAGE_SAFETY", "IMAGE_PROHIBITED_CONTENT", - "TOO_MANY_TOOL_CALLS", "MALFORMED_RESPONSE", - }) + _GEMINI_FINISH_REASON_KEYS = frozenset( + { + "STOP", + "MAX_TOKENS", + "SAFETY", + "RECITATION", + "FINISH_REASON_UNSPECIFIED", + "MALFORMED_FUNCTION_CALL", + "LANGUAGE", + "OTHER", + "BLOCKLIST", + "PROHIBITED_CONTENT", + "SPII", + "IMAGE_SAFETY", + "IMAGE_PROHIBITED_CONTENT", + "TOO_MANY_TOOL_CALLS", + "MALFORMED_RESPONSE", + } + ) @staticmethod def get_finish_reason_mapping() -> Dict[str, OpenAIChatCompletionFinishReason]: @@ -1458,10 +1468,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): _tool_response_chunk["provider_specific_fields"] = { # type: ignore "thought_signature": thought_signature } - _tool_response_chunk["id"] = ( - _encode_tool_call_id_with_signature( - _tool_response_chunk["id"] or "", thought_signature - ) + _tool_response_chunk[ + "id" + ] = _encode_tool_call_id_with_signature( + _tool_response_chunk["id"] or "", thought_signature ) _tools.append(_tool_response_chunk) cumulative_tool_call_idx += 1 @@ -2271,35 +2281,37 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ## ADD METADATA TO RESPONSE ## setattr(model_response, "vertex_ai_grounding_metadata", grounding_metadata) - model_response._hidden_params["vertex_ai_grounding_metadata"] = ( - grounding_metadata - ) + model_response._hidden_params[ + "vertex_ai_grounding_metadata" + ] = grounding_metadata setattr( model_response, "vertex_ai_url_context_metadata", url_context_metadata ) - model_response._hidden_params["vertex_ai_url_context_metadata"] = ( - url_context_metadata - ) + model_response._hidden_params[ + "vertex_ai_url_context_metadata" + ] = url_context_metadata setattr(model_response, "vertex_ai_safety_results", safety_ratings) - model_response._hidden_params["vertex_ai_safety_results"] = ( - safety_ratings # older approach - maintaining to prevent regressions - ) + model_response._hidden_params[ + "vertex_ai_safety_results" + ] = safety_ratings # older approach - maintaining to prevent regressions ## ADD CITATION METADATA ## setattr(model_response, "vertex_ai_citation_metadata", citation_metadata) - model_response._hidden_params["vertex_ai_citation_metadata"] = ( - citation_metadata # older approach - maintaining to prevent regressions - ) + model_response._hidden_params[ + "vertex_ai_citation_metadata" + ] = citation_metadata # older approach - maintaining to prevent regressions ## ADD TRAFFIC TYPE ## traffic_type = completion_response.get("usageMetadata", {}).get( "trafficType" ) if traffic_type: - model_response._hidden_params.setdefault("provider_specific_fields", {})["traffic_type"] = traffic_type + model_response._hidden_params.setdefault( + "provider_specific_fields", {} + )["traffic_type"] = traffic_type except Exception as e: raise VertexAIError( @@ -2955,7 +2967,11 @@ class ModelResponseIterator: # to correctly set finish_reason="tool_calls" per the OpenAI spec. if not self.has_seen_tool_calls: for choice in model_response.choices: - if hasattr(choice, "delta") and choice.delta and choice.delta.tool_calls: + if ( + hasattr(choice, "delta") + and choice.delta + and choice.delta.tool_calls + ): self.has_seen_tool_calls = True break @@ -2971,8 +2987,10 @@ class ModelResponseIterator: if self.has_seen_tool_calls: mapped_finish_reason = "tool_calls" else: - mapped_finish_reason = VertexGeminiConfig._check_finish_reason( - None, finish_reason_str + mapped_finish_reason = ( + VertexGeminiConfig._check_finish_reason( + None, finish_reason_str + ) ) choice = StreamingChoices( finish_reason=mapped_finish_reason, @@ -3005,7 +3023,9 @@ class ModelResponseIterator: "trafficType" ) if traffic_type: - model_response._hidden_params.setdefault("provider_specific_fields", {})["traffic_type"] = traffic_type + model_response._hidden_params.setdefault( + "provider_specific_fields", {} + )["traffic_type"] = traffic_type setattr(model_response, "usage", usage) # type: ignore diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py index 68901340c7..3ec0bdf22a 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py @@ -40,37 +40,37 @@ class GoogleBatchEmbeddings(VertexLLM): ) -> Dict[str, Dict[str, str]]: """ Resolve Gemini file references (files/...) to get mime_type and uri. - + Args: input: EmbeddingInput that may contain file references api_key: Gemini API key sync_handler: HTTP client - + Returns: Dict mapping file name to {mime_type, uri} """ input_list = [input] if isinstance(input, str) else input resolved_files: Dict[str, Dict[str, str]] = {} - + for element in input_list: if isinstance(element, str) and _is_file_reference(element): url = f"https://generativelanguage.googleapis.com/v1beta/{element}" headers = {"x-goog-api-key": api_key} response = sync_handler.get(url=url, headers=headers) - + if response.status_code != 200: raise Exception( f"Error fetching file {element}: {response.status_code} {response.text}" ) - + file_data = response.json() resolved_files[element] = { "mime_type": file_data.get("mimeType", ""), "uri": file_data.get("uri", element), } - + return resolved_files - + async def _async_resolve_file_references( self, input: EmbeddingInput, @@ -79,37 +79,37 @@ class GoogleBatchEmbeddings(VertexLLM): ) -> Dict[str, Dict[str, str]]: """ Async version of _resolve_file_references. - + Args: input: EmbeddingInput that may contain file references api_key: Gemini API key async_handler: Async HTTP client - + Returns: Dict mapping file name to {mime_type, uri} """ input_list = [input] if isinstance(input, str) else input resolved_files: Dict[str, Dict[str, str]] = {} - + for element in input_list: if isinstance(element, str) and _is_file_reference(element): url = f"https://generativelanguage.googleapis.com/v1beta/{element}" headers = {"x-goog-api-key": api_key} response = await async_handler.get(url=url, headers=headers) - + if response.status_code != 200: raise Exception( f"Error fetching file {element}: {response.status_code} {response.text}" ) - + file_data = response.json() resolved_files[element] = { "mime_type": file_data.get("mimeType", ""), "uri": file_data.get("uri", element), } - + return resolved_files - + def batch_embeddings( self, model: str, @@ -238,7 +238,7 @@ class GoogleBatchEmbeddings(VertexLLM): raise Exception(f"Error: {response.status_code} {response.text}") _json_response = response.json() - + if use_embed_content: return process_embed_content_response( input=input, @@ -327,7 +327,7 @@ class GoogleBatchEmbeddings(VertexLLM): raise Exception(f"Error: {response.status_code} {response.text}") _json_response = response.json() - + if use_embed_content: return process_embed_content_response( input=input, diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py index 41f477d9db..0f6d85525d 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py @@ -43,13 +43,13 @@ def _is_gcs_url(s: str) -> bool: def _infer_mime_type_from_gcs_url(gcs_url: str) -> str: """ Infer MIME type from GCS URL file extension. - + Args: gcs_url: GCS URL like gs://bucket/path/to/file.png - + Returns: str: Inferred MIME type - + Raises: ValueError: If file extension is not supported """ @@ -63,12 +63,12 @@ def _infer_mime_type_from_gcs_url(gcs_url: str) -> str: ".mov": "video/quicktime", ".pdf": "application/pdf", } - + gcs_url_lower = gcs_url.lower() for ext, mime_type in extension_to_mime.items(): if gcs_url_lower.endswith(ext): return mime_type - + raise ValueError( f"Unable to infer MIME type from GCS URL: {gcs_url}. " f"Supported extensions: {', '.join(extension_to_mime.keys())}" @@ -78,49 +78,49 @@ def _infer_mime_type_from_gcs_url(gcs_url: str) -> str: def _parse_data_url(data_url: str) -> Tuple[str, str]: """ Parse a data URL to extract the media type and base64 data. - + Args: data_url: Data URL in format: data:image/jpeg;base64,/9j/4AAQ... - + Returns: tuple: (media_type, base64_data) media_type: e.g., "image/jpeg", "video/mp4", "audio/mpeg" base64_data: The base64-encoded data without the prefix - + Raises: ValueError: If data URL format is invalid or MIME type is unsupported """ if not data_url.startswith("data:"): raise ValueError(f"Invalid data URL format: {data_url[:50]}...") - + if "," not in data_url: raise ValueError(f"Invalid data URL format (missing comma): {data_url[:50]}...") - + metadata, base64_data = data_url.split(",", 1) - + metadata = metadata[5:] - + if ";" in metadata: media_type = metadata.split(";")[0] else: media_type = metadata - + if media_type not in SUPPORTED_EMBEDDING_MIME_TYPES: raise ValueError( f"Unsupported MIME type for embedding: {media_type}. " f"Supported types: {', '.join(sorted(SUPPORTED_EMBEDDING_MIME_TYPES))}" ) - + return media_type, base64_data def _is_multimodal_input(input: EmbeddingInput) -> bool: """ Check if the input contains multimodal data (data URIs, file references, or GCS URLs). - + Args: input: EmbeddingInput (str or List[str]) - + Returns: bool: True if any element is a data URI, file reference, or GCS URL """ @@ -128,7 +128,7 @@ def _is_multimodal_input(input: EmbeddingInput) -> bool: input_list = [input] else: input_list = input - + for element in input_list: if isinstance(element, str): if element.startswith("data:") and ";base64," in element: @@ -137,7 +137,7 @@ def _is_multimodal_input(input: EmbeddingInput) -> bool: return True if _is_gcs_url(element): return True - + return False @@ -148,17 +148,17 @@ def transform_openai_input_gemini_content( The content to embed. Only the parts.text fields will be counted. """ gemini_model_name = "models/{}".format(model) - + gemini_params = optional_params.copy() if "dimensions" in gemini_params: gemini_params["outputDimensionality"] = gemini_params.pop("dimensions") - + requests: List[EmbedContentRequest] = [] if isinstance(input, str): request = EmbedContentRequest( model=gemini_model_name, content=ContentType(parts=[PartType(text=input)]), - **gemini_params + **gemini_params, ) requests.append(request) else: @@ -166,7 +166,7 @@ def transform_openai_input_gemini_content( request = EmbedContentRequest( model=gemini_model_name, content=ContentType(parts=[PartType(text=i)]), - **gemini_params + **gemini_params, ) requests.append(request) @@ -181,29 +181,29 @@ def transform_openai_input_gemini_embed_content( ) -> dict: """ Transform OpenAI embedding input to Gemini embedContent format (multimodal). - + Args: input: EmbeddingInput (str or List[str]) with text, data URIs, or file references model: Model name optional_params: Additional parameters (taskType, outputDimensionality, etc.) resolved_files: Dict mapping file names (files/abc) to {mime_type, uri} - + Returns: dict: Gemini embedContent request body with content.parts """ resolved_files = resolved_files or {} - + gemini_params = optional_params.copy() if "dimensions" in gemini_params: gemini_params["outputDimensionality"] = gemini_params.pop("dimensions") - + input_list = [input] if isinstance(input, str) else input parts: List[PartType] = [] - + for element in input_list: if not isinstance(element, str): raise ValueError(f"Unsupported input type: {type(element)}") - + if element.startswith("data:") and ";base64," in element: mime_type, base64_data = _parse_data_url(element) blob: BlobType = {"mime_type": mime_type, "data": base64_data} @@ -226,12 +226,12 @@ def transform_openai_input_gemini_embed_content( parts.append(PartType(file_data=file_data_ref)) else: parts.append(PartType(text=element)) - + request_body: dict = { "content": ContentType(parts=parts), **gemini_params, } - + return request_body @@ -243,30 +243,32 @@ def process_embed_content_response( ) -> EmbeddingResponse: """ Process Gemini embedContent response (single embedding for multimodal input). - + Args: input: Original input model_response: EmbeddingResponse to populate model: Model name response_json: Raw JSON response from embedContent endpoint - + Returns: EmbeddingResponse with single embedding """ if "embedding" not in response_json: - raise ValueError(f"embedContent response missing 'embedding' field: {response_json}") - + raise ValueError( + f"embedContent response missing 'embedding' field: {response_json}" + ) + embedding_data = response_json["embedding"] - + openai_embedding = Embedding( embedding=embedding_data["values"], index=0, object="embedding", ) - + model_response.data = [openai_embedding] model_response.model = model - + if _is_multimodal_input(input): prompt_tokens = 0 else: @@ -275,7 +277,7 @@ def process_embed_content_response( model_response.usage = Usage( prompt_tokens=prompt_tokens, total_tokens=prompt_tokens ) - + return model_response diff --git a/litellm/llms/vertex_ai/image_edit/__init__.py b/litellm/llms/vertex_ai/image_edit/__init__.py index 44914e861a..51bb151165 100644 --- a/litellm/llms/vertex_ai/image_edit/__init__.py +++ b/litellm/llms/vertex_ai/image_edit/__init__.py @@ -1,35 +1,38 @@ from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig -from litellm.llms.vertex_ai.common_utils import VertexAIModelRoute, get_vertex_ai_model_route +from litellm.llms.vertex_ai.common_utils import ( + VertexAIModelRoute, + get_vertex_ai_model_route, +) from .cost_calculator import cost_calculator from .vertex_gemini_transformation import VertexAIGeminiImageEditConfig from .vertex_imagen_transformation import VertexAIImagenImageEditConfig __all__ = [ - "VertexAIGeminiImageEditConfig", + "VertexAIGeminiImageEditConfig", "VertexAIImagenImageEditConfig", - "get_vertex_ai_image_edit_config", - "cost_calculator" + "get_vertex_ai_image_edit_config", + "cost_calculator", ] def get_vertex_ai_image_edit_config(model: str) -> BaseImageEditConfig: """ Get the appropriate image edit config for a Vertex AI model. - + Routes to the correct transformation class based on the model type: - Gemini models use generateContent API (VertexAIGeminiImageEditConfig) - Imagen models use predict API (VertexAIImagenImageEditConfig) - + Args: model: The model name (e.g., "gemini-2.5-flash", "imagegeneration@006") - + Returns: BaseImageEditConfig: The appropriate configuration class """ # Determine the model route model_route = get_vertex_ai_model_route(model) - + if model_route == VertexAIModelRoute.GEMINI: # Gemini models use generateContent API return VertexAIGeminiImageEditConfig() diff --git a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py index 8fcd285824..de7f234a86 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py @@ -28,9 +28,10 @@ else: class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): """ Vertex AI Gemini Image Edit Configuration - + Uses generateContent API for Gemini models on Vertex AI """ + SUPPORTED_PARAMS: List[str] = ["size"] def __init__(self) -> None: @@ -99,17 +100,23 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): ) -> dict: headers = headers or {} litellm_params = litellm_params or {} - + # If a custom api_base is provided, skip credential validation # This allows users to use proxies or mock endpoints without needing Vertex AI credentials _api_base = litellm_params.get("api_base") or api_base if _api_base is not None: return headers - + # First check litellm_params (where vertex_ai_project/vertex_ai_credentials are passed) # then fall back to environment variables and other sources - vertex_project = self.safe_get_vertex_ai_project(litellm_params) or self._resolve_vertex_project() - vertex_credentials = self.safe_get_vertex_ai_credentials(litellm_params) or self._resolve_vertex_credentials() + vertex_project = ( + self.safe_get_vertex_ai_project(litellm_params) + or self._resolve_vertex_project() + ) + vertex_credentials = ( + self.safe_get_vertex_ai_credentials(litellm_params) + or self._resolve_vertex_credentials() + ) access_token, _ = self._ensure_access_token( credentials=vertex_credentials, project_id=vertex_project, @@ -138,11 +145,19 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): # First check litellm_params (where vertex_ai_project/vertex_ai_location are passed) # then fall back to environment variables and other sources - vertex_project = self.safe_get_vertex_ai_project(litellm_params) or self._resolve_vertex_project() - vertex_location = self.safe_get_vertex_ai_location(litellm_params) or self._resolve_vertex_location() + vertex_project = ( + self.safe_get_vertex_ai_project(litellm_params) + or self._resolve_vertex_project() + ) + vertex_location = ( + self.safe_get_vertex_ai_location(litellm_params) + or self._resolve_vertex_location() + ) if not vertex_project or not vertex_location: - raise ValueError("vertex_project and vertex_location are required for Vertex AI") + raise ValueError( + "vertex_project and vertex_location are required for Vertex AI" + ) base_url = get_vertex_base_url(vertex_location) @@ -167,23 +182,20 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): parts.append({"text": prompt}) # Correct format for Vertex AI Gemini image editing - contents = { - "role": "USER", - "parts": parts - } + contents = {"role": "USER", "parts": parts} request_body: Dict[str, Any] = {"contents": contents} # Generation config with proper structure for image editing - generation_config: Dict[str, Any] = { - "response_modalities": ["IMAGE"] - } + generation_config: Dict[str, Any] = {"response_modalities": ["IMAGE"]} # Add image-specific configuration image_config: Dict[str, Any] = {} if "aspectRatio" in image_edit_optional_request_params: - image_config["aspect_ratio"] = image_edit_optional_request_params["aspectRatio"] - + image_config["aspect_ratio"] = image_edit_optional_request_params[ + "aspectRatio" + ] + if image_config: generation_config["image_config"] = image_config @@ -191,7 +203,9 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): payload: Any = json.dumps(request_body) empty_files = cast(RequestFiles, []) - return cast(Tuple[Dict[str, Any], Optional[RequestFiles]], (payload, empty_files)) + return cast( + Tuple[Dict[str, Any], Optional[RequestFiles]], (payload, empty_files) + ) def transform_image_edit_response( self, diff --git a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py index b58825e1fa..7979e0e790 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py @@ -29,9 +29,10 @@ else: class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): """ Vertex AI Imagen Image Edit Configuration - + Uses predict API for Imagen models on Vertex AI """ + SUPPORTED_PARAMS: List[str] = ["n", "size", "mask"] def __init__(self) -> None: @@ -59,12 +60,12 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): # Map OpenAI parameters to Imagen format if "n" in filtered_params: mapped_params["sampleCount"] = filtered_params["n"] - + if "size" in filtered_params: mapped_params["aspectRatio"] = self._map_size_to_aspect_ratio( filtered_params["size"] # type: ignore[arg-type] ) - + if "mask" in filtered_params: mapped_params["mask"] = filtered_params["mask"] @@ -126,7 +127,9 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): vertex_location = self._resolve_vertex_location() if not vertex_project or not vertex_location: - raise ValueError("vertex_project and vertex_location are required for Vertex AI") + raise ValueError( + "vertex_project and vertex_location are required for Vertex AI" + ) # Use the model name as provided, handling vertex_ai prefix model_name = model @@ -151,35 +154,34 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): ) -> Tuple[Dict[str, Any], Optional[RequestFiles]]: # Prepare reference images in the correct Imagen format if image is None: - raise ValueError("Vertex AI Imagen image edit requires at least one reference image.") - reference_images = self._prepare_reference_images(image, image_edit_optional_request_params) + raise ValueError( + "Vertex AI Imagen image edit requires at least one reference image." + ) + reference_images = self._prepare_reference_images( + image, image_edit_optional_request_params + ) if not reference_images: - raise ValueError("Vertex AI Imagen image edit requires at least one reference image.") + raise ValueError( + "Vertex AI Imagen image edit requires at least one reference image." + ) if prompt is None: raise ValueError("Vertex AI Imagen image edit requires a prompt.") # Correct Imagen instances format - instances = [ - { - "prompt": prompt, - "referenceImages": reference_images - } - ] + instances = [{"prompt": prompt, "referenceImages": reference_images}] # Extract OpenAI parameters and set sensible defaults for Vertex AI-specific parameters sample_count = image_edit_optional_request_params.get("sampleCount", 1) # Use sensible defaults for Vertex AI-specific parameters (not exposed to users) edit_mode = "EDIT_MODE_INPAINT_INSERTION" # Default edit mode base_steps = 50 # Default number of steps - + # Imagen parameters with correct structure parameters = { "sampleCount": sample_count, "editMode": edit_mode, - "editConfig": { - "baseSteps": base_steps - } + "editConfig": {"baseSteps": base_steps}, } # Set default values for Vertex AI-specific parameters (not configurable by users via OpenAI API) @@ -188,12 +190,14 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): request_body: Dict[str, Any] = { "instances": instances, - "parameters": parameters + "parameters": parameters, } payload: Any = json.dumps(request_body) empty_files = cast(RequestFiles, []) - return cast(Tuple[Dict[str, Any], Optional[RequestFiles]], (payload, empty_files)) + return cast( + Tuple[Dict[str, Any], Optional[RequestFiles]], (payload, empty_files) + ) def transform_image_edit_response( self, @@ -231,7 +235,7 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): """Map OpenAI size format to Imagen aspect ratio format""" aspect_ratio_map = { "1024x1024": "1:1", - "1792x1024": "16:9", + "1792x1024": "16:9", "1024x1792": "9:16", "1280x896": "4:3", "896x1280": "3:4", @@ -239,8 +243,9 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): return aspect_ratio_map.get(size, "1:1") def _prepare_reference_images( - self, image: Union[FileTypes, List[FileTypes]], - image_edit_optional_request_params: Dict[str, Any] + self, + image: Union[FileTypes, List[FileTypes]], + image_edit_optional_request_params: Dict[str, Any], ) -> List[Dict[str, Any]]: """ Prepare reference images in the correct Imagen API format @@ -252,41 +257,37 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): images = [image] reference_images: List[Dict[str, Any]] = [] - + for idx, img in enumerate(images): if img is None: continue image_bytes = self._read_all_bytes(img) base64_data = base64.b64encode(image_bytes).decode("utf-8") - + # Create reference image structure reference_image = { "referenceType": "REFERENCE_TYPE_RAW", "referenceId": idx + 1, - "referenceImage": { - "bytesBase64Encoded": base64_data - } + "referenceImage": {"bytesBase64Encoded": base64_data}, } - + reference_images.append(reference_image) - + # Handle mask image if provided (for inpainting) mask_image = image_edit_optional_request_params.get("mask") if mask_image is not None: mask_bytes = self._read_all_bytes(mask_image) mask_base64 = base64.b64encode(mask_bytes).decode("utf-8") - + mask_reference = { "referenceType": "REFERENCE_TYPE_MASK", "referenceId": len(reference_images) + 1, - "referenceImage": { - "bytesBase64Encoded": mask_base64 - }, + "referenceImage": {"bytesBase64Encoded": mask_base64}, "maskImageConfig": { "maskMode": "MASK_MODE_USER_PROVIDED", - "dilation": 0.03 # Default dilation value (not configurable via OpenAI API) - } + "dilation": 0.03, # Default dilation value (not configurable via OpenAI API) + }, } reference_images.append(mask_reference) @@ -303,7 +304,9 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): if isinstance(image, (list, tuple)): for item in image: if item is not None: - return self._read_all_bytes(item, depth=depth + 1, max_depth=max_depth) + return self._read_all_bytes( + item, depth=depth + 1, max_depth=max_depth + ) raise ValueError("Unsupported image type for Vertex AI Imagen image edit.") if isinstance(image, dict): @@ -315,9 +318,13 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): return base64.b64decode(value) except Exception: continue - return self._read_all_bytes(value, depth=depth + 1, max_depth=max_depth) + return self._read_all_bytes( + value, depth=depth + 1, max_depth=max_depth + ) if "path" in image: - return self._read_all_bytes(image["path"], depth=depth + 1, max_depth=max_depth) + return self._read_all_bytes( + image["path"], depth=depth + 1, max_depth=max_depth + ) if isinstance(image, bytes): return image diff --git a/litellm/llms/vertex_ai/image_generation/__init__.py b/litellm/llms/vertex_ai/image_generation/__init__.py index a6f6156167..9445660dba 100644 --- a/litellm/llms/vertex_ai/image_generation/__init__.py +++ b/litellm/llms/vertex_ai/image_generation/__init__.py @@ -10,29 +10,29 @@ from .vertex_gemini_transformation import VertexAIGeminiImageGenerationConfig from .vertex_imagen_transformation import VertexAIImagenImageGenerationConfig __all__ = [ - "VertexAIGeminiImageGenerationConfig", + "VertexAIGeminiImageGenerationConfig", "VertexAIImagenImageGenerationConfig", - "get_vertex_ai_image_generation_config", + "get_vertex_ai_image_generation_config", ] def get_vertex_ai_image_generation_config(model: str) -> BaseImageGenerationConfig: """ Get the appropriate image generation config for a Vertex AI model. - + Routes to the correct transformation class based on the model type: - Gemini image generation models use generateContent API (VertexAIGeminiImageGenerationConfig) - Imagen models use predict API (VertexAIImagenImageGenerationConfig) - + Args: model: The model name (e.g., "gemini-2.5-flash-image", "imagegeneration@006") - + Returns: BaseImageGenerationConfig: The appropriate configuration class """ # Determine the model route model_route = get_vertex_ai_model_route(model) - + if model_route == VertexAIModelRoute.GEMINI: # Gemini models use generateContent API return VertexAIGeminiImageGenerationConfig() @@ -40,4 +40,3 @@ def get_vertex_ai_image_generation_config(model: str) -> BaseImageGenerationConf # Default to Imagen for other models (imagegeneration, etc.) # This includes NON_GEMINI models like imagegeneration@006 return VertexAIImagenImageGenerationConfig() - diff --git a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py index 447612877f..98e02743bd 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -29,18 +29,16 @@ else: class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): """ Vertex AI Gemini Image Generation Configuration - + Uses generateContent API for Gemini image generation models on Vertex AI Supports models like gemini-2.5-flash-image, gemini-3-pro-image-preview, etc. """ - + def __init__(self) -> None: BaseImageGenerationConfig.__init__(self) VertexLLM.__init__(self) - - def get_supported_openai_params( - self, model: str - ) -> list: + + def get_supported_openai_params(self, model: str) -> list: """ Gemini image generation supported parameters @@ -55,7 +53,7 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): "imageSize", "image_size", ] - + def map_openai_params( self, non_default_params: dict, @@ -65,7 +63,7 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): ) -> dict: supported_params = self.get_supported_openai_params(model) mapped_params = {} - + for k, v in non_default_params.items(): if k not in optional_params.keys(): if k in supported_params: @@ -81,22 +79,22 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): mapped_params["imageSize"] = v else: mapped_params[k] = v - + return mapped_params - + def _map_size_to_aspect_ratio(self, size: str) -> str: """ Map OpenAI size format to Gemini aspect ratio format """ aspect_ratio_map = { "1024x1024": "1:1", - "1792x1024": "16:9", + "1792x1024": "16:9", "1024x1792": "9:16", "1280x896": "4:3", - "896x1280": "3:4" + "896x1280": "3:4", } return aspect_ratio_map.get(size, "1:1") - + def _resolve_vertex_project(self) -> Optional[str]: return ( getattr(self, "_vertex_project", None) @@ -148,11 +146,19 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): # First check litellm_params (where vertex_ai_project/vertex_ai_location are passed) # then fall back to environment variables and other sources - vertex_project = self.safe_get_vertex_ai_project(litellm_params) or self._resolve_vertex_project() - vertex_location = self.safe_get_vertex_ai_location(litellm_params) or self._resolve_vertex_location() + vertex_project = ( + self.safe_get_vertex_ai_project(litellm_params) + or self._resolve_vertex_project() + ) + vertex_location = ( + self.safe_get_vertex_ai_location(litellm_params) + or self._resolve_vertex_location() + ) if not vertex_project or not vertex_location: - raise ValueError("vertex_project and vertex_location are required for Vertex AI") + raise ValueError( + "vertex_project and vertex_location are required for Vertex AI" + ) base_url = get_vertex_base_url(vertex_location) @@ -169,17 +175,23 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): api_base: Optional[str] = None, ) -> dict: headers = headers or {} - + # If a custom api_base is provided, skip credential validation # This allows users to use proxies or mock endpoints without needing Vertex AI credentials _api_base = litellm_params.get("api_base") or api_base if _api_base is not None: return headers - + # First check litellm_params (where vertex_ai_project/vertex_ai_credentials are passed) # then fall back to environment variables and other sources - vertex_project = self.safe_get_vertex_ai_project(litellm_params) or self._resolve_vertex_project() - vertex_credentials = self.safe_get_vertex_ai_credentials(litellm_params) or self._resolve_vertex_credentials() + vertex_project = ( + self.safe_get_vertex_ai_project(litellm_params) + or self._resolve_vertex_project() + ) + vertex_credentials = ( + self.safe_get_vertex_ai_credentials(litellm_params) + or self._resolve_vertex_credentials() + ) access_token, _ = self._ensure_access_token( credentials=vertex_credentials, project_id=vertex_project, @@ -197,51 +209,44 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): ) -> dict: """ Transform the image generation request to Gemini format - + Uses generateContent API with responseModalities: ["IMAGE"] """ # Prepare messages with the prompt - contents = [ - { - "role": "user", - "parts": [{"text": prompt}] - } - ] - + contents = [{"role": "user", "parts": [{"text": prompt}]}] + # Prepare generation config - generation_config: Dict[str, Any] = { - "responseModalities": ["IMAGE"] - } - + generation_config: Dict[str, Any] = {"responseModalities": ["IMAGE"]} + # Handle image-specific config parameters image_config: Dict[str, Any] = {} - + # Map aspectRatio if "aspectRatio" in optional_params: image_config["aspectRatio"] = optional_params["aspectRatio"] elif "aspect_ratio" in optional_params: image_config["aspectRatio"] = optional_params["aspect_ratio"] - + # Map imageSize (for Gemini 3 Pro) if "imageSize" in optional_params: image_config["imageSize"] = optional_params["imageSize"] elif "image_size" in optional_params: image_config["imageSize"] = optional_params["image_size"] - + if image_config: generation_config["imageConfig"] = image_config - + # Handle candidate_count (n parameter) if "candidate_count" in optional_params: generation_config["candidateCount"] = optional_params["candidate_count"] elif "n" in optional_params: generation_config["candidateCount"] = optional_params["n"] - + request_body: Dict[str, Any] = { "contents": contents, - "generationConfig": generation_config + "generationConfig": generation_config, } - + return request_body def _transform_image_usage(self, usage: dict) -> ImageUsage: @@ -289,7 +294,7 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): status_code=raw_response.status_code, headers=raw_response.headers, ) - + if not model_response.data: model_response.data = [] @@ -304,14 +309,19 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): inline_data = part["inlineData"] if "data" in inline_data: thought_sig = part.get("thoughtSignature") - model_response.data.append(ImageObject( - b64_json=inline_data["data"], - url=None, - provider_specific_fields={"thought_signature": thought_sig} if thought_sig else None, - )) + model_response.data.append( + ImageObject( + b64_json=inline_data["data"], + url=None, + provider_specific_fields={ + "thought_signature": thought_sig + } + if thought_sig + else None, + ) + ) if usage_metadata := response_data.get("usageMetadata", None): model_response.usage = self._transform_image_usage(usage_metadata) - - return model_response + return model_response diff --git a/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py index 6f9e387417..1c7696d55a 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py @@ -27,26 +27,23 @@ else: class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): """ Vertex AI Imagen Image Generation Configuration - + Uses predict API for Imagen models on Vertex AI Supports models like imagegeneration@006 """ - + def __init__(self) -> None: BaseImageGenerationConfig.__init__(self) VertexLLM.__init__(self) - + def get_supported_openai_params( self, model: str ) -> List[OpenAIImageGenerationOptionalParams]: """ Imagen API supported parameters """ - return [ - "n", - "size" - ] - + return ["n", "size"] + def map_openai_params( self, non_default_params: dict, @@ -56,7 +53,7 @@ class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): ) -> dict: supported_params = self.get_supported_openai_params(model) mapped_params = {} - + for k, v in non_default_params.items(): if k not in optional_params.keys(): if k in supported_params: @@ -68,22 +65,22 @@ class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): mapped_params["aspectRatio"] = self._map_size_to_aspect_ratio(v) else: mapped_params[k] = v - + return mapped_params - + def _map_size_to_aspect_ratio(self, size: str) -> str: """ Map OpenAI size format to Imagen aspect ratio format """ aspect_ratio_map = { "1024x1024": "1:1", - "1792x1024": "16:9", + "1792x1024": "16:9", "1024x1792": "9:16", "1280x896": "4:3", - "896x1280": "3:4" + "896x1280": "3:4", } return aspect_ratio_map.get(size, "1:1") - + def _resolve_vertex_project(self) -> Optional[str]: return ( getattr(self, "_vertex_project", None) @@ -135,11 +132,19 @@ class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): # First check litellm_params (where vertex_ai_project/vertex_ai_location are passed) # then fall back to environment variables and other sources - vertex_project = self.safe_get_vertex_ai_project(litellm_params) or self._resolve_vertex_project() - vertex_location = self.safe_get_vertex_ai_location(litellm_params) or self._resolve_vertex_location() + vertex_project = ( + self.safe_get_vertex_ai_project(litellm_params) + or self._resolve_vertex_project() + ) + vertex_location = ( + self.safe_get_vertex_ai_location(litellm_params) + or self._resolve_vertex_location() + ) if not vertex_project or not vertex_location: - raise ValueError("vertex_project and vertex_location are required for Vertex AI") + raise ValueError( + "vertex_project and vertex_location are required for Vertex AI" + ) base_url = get_vertex_base_url(vertex_location) @@ -156,17 +161,23 @@ class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): api_base: Optional[str] = None, ) -> dict: headers = headers or {} - + # If a custom api_base is provided, skip credential validation # This allows users to use proxies or mock endpoints without needing Vertex AI credentials _api_base = litellm_params.get("api_base") or api_base if _api_base is not None: return headers - + # First check litellm_params (where vertex_ai_project/vertex_ai_credentials are passed) # then fall back to environment variables and other sources - vertex_project = self.safe_get_vertex_ai_project(litellm_params) or self._resolve_vertex_project() - vertex_credentials = self.safe_get_vertex_ai_credentials(litellm_params) or self._resolve_vertex_credentials() + vertex_project = ( + self.safe_get_vertex_ai_project(litellm_params) + or self._resolve_vertex_project() + ) + vertex_credentials = ( + self.safe_get_vertex_ai_credentials(litellm_params) + or self._resolve_vertex_credentials() + ) access_token, _ = self._ensure_access_token( credentials=vertex_credentials, project_id=vertex_project, @@ -184,22 +195,22 @@ class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): ) -> dict: """ Transform the image generation request to Imagen format - + Uses predict API with instances and parameters """ # Default parameters default_params = { "sampleCount": 1, } - + # Merge with optional params parameters = {**default_params, **optional_params} - + request_body = { "instances": [{"prompt": prompt}], "parameters": parameters, } - + return request_body def transform_image_generation_response( @@ -226,7 +237,7 @@ class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): status_code=raw_response.status_code, headers=raw_response.headers, ) - + if not model_response.data: model_response.data = [] @@ -235,10 +246,11 @@ class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): for prediction in predictions: # Imagen returns images as bytesBase64Encoded if "bytesBase64Encoded" in prediction: - model_response.data.append(ImageObject( - b64_json=prediction["bytesBase64Encoded"], - url=None, - )) - - return model_response + model_response.data.append( + ImageObject( + b64_json=prediction["bytesBase64Encoded"], + url=None, + ) + ) + return model_response diff --git a/litellm/llms/vertex_ai/ocr/__init__.py b/litellm/llms/vertex_ai/ocr/__init__.py index fa8c85da9c..15da24f308 100644 --- a/litellm/llms/vertex_ai/ocr/__init__.py +++ b/litellm/llms/vertex_ai/ocr/__init__.py @@ -2,4 +2,3 @@ from .transformation import VertexAIOCRConfig __all__ = ["VertexAIOCRConfig"] - diff --git a/litellm/llms/vertex_ai/ocr/common_utils.py b/litellm/llms/vertex_ai/ocr/common_utils.py index dc2c07420b..3e5fbe2344 100644 --- a/litellm/llms/vertex_ai/ocr/common_utils.py +++ b/litellm/llms/vertex_ai/ocr/common_utils.py @@ -14,20 +14,20 @@ if TYPE_CHECKING: def get_vertex_ai_ocr_config(model: str) -> Optional["BaseOCRConfig"]: """ Determine which Vertex AI OCR configuration to use based on the model name. - + Vertex AI supports multiple OCR services: - Vertex AI OCR: vertex_ai/ - + Args: model: The model name (e.g., "vertex_ai/ocr/") - + Returns: OCR configuration instance for the specified model - + Examples: >>> get_vertex_ai_ocr_config("vertex_ai/deepseek-ai/deepseek-ocr-maas") - + >>> get_vertex_ai_ocr_config("vertex_ai/ocr/mistral-ocr-maas") """ @@ -35,7 +35,7 @@ def get_vertex_ai_ocr_config(model: str) -> Optional["BaseOCRConfig"]: VertexAIDeepSeekOCRConfig, ) from litellm.llms.vertex_ai.ocr.transformation import VertexAIOCRConfig + if "deepseek" in model: return VertexAIDeepSeekOCRConfig() return VertexAIOCRConfig() - diff --git a/litellm/llms/vertex_ai/ocr/deepseek_transformation.py b/litellm/llms/vertex_ai/ocr/deepseek_transformation.py index b16f73af3f..953bb51fd1 100644 --- a/litellm/llms/vertex_ai/ocr/deepseek_transformation.py +++ b/litellm/llms/vertex_ai/ocr/deepseek_transformation.py @@ -26,7 +26,7 @@ else: class VertexAIDeepSeekOCRConfig(BaseOCRConfig): """ Vertex AI DeepSeek OCR transformation configuration. - + Vertex AI DeepSeek OCR uses the chat completion API format through the openapi endpoint. This transformation converts OCR requests to chat completion format and vice versa. """ @@ -46,16 +46,20 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): ) -> Dict: """ Validate environment and return headers for Vertex AI OCR. - + Vertex AI uses Bearer token authentication with access token from credentials. """ # Extract Vertex AI parameters using safe helpers from VertexBase # Use safe_get_* methods that don't mutate litellm_params dict litellm_params = litellm_params or {} - - vertex_project = VertexBase.safe_get_vertex_ai_project(litellm_params=litellm_params) - vertex_credentials = VertexBase.safe_get_vertex_ai_credentials(litellm_params=litellm_params) - + + vertex_project = VertexBase.safe_get_vertex_ai_project( + litellm_params=litellm_params + ) + vertex_credentials = VertexBase.safe_get_vertex_ai_credentials( + litellm_params=litellm_params + ) + # Get access token from Vertex credentials access_token, project_id = self.vertex_base.get_access_token( credentials=vertex_credentials, @@ -80,25 +84,29 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): ) -> str: """ Get complete URL for Vertex AI DeepSeek OCR endpoint. - - Vertex AI endpoint format: + + Vertex AI endpoint format: https://{location}-aiplatform.googleapis.com/v1/projects/{project}/locations/{location}/endpoints/openapi/chat/completions - + Args: api_base: Vertex AI API base URL (optional) model: Model name (e.g., "deepseek-ai/deepseek-ocr-maas") optional_params: Optional parameters litellm_params: LiteLLM parameters containing vertex_project, vertex_location - + Returns: Complete URL for Vertex AI OCR endpoint """ # Extract Vertex AI parameters using safe helpers from VertexBase # Use safe_get_* methods that don't mutate litellm_params dict litellm_params = litellm_params or {} - - vertex_project = VertexBase.safe_get_vertex_ai_project(litellm_params=litellm_params) - vertex_location = VertexBase.safe_get_vertex_ai_location(litellm_params=litellm_params) - + + vertex_project = VertexBase.safe_get_vertex_ai_project( + litellm_params=litellm_params + ) + vertex_location = VertexBase.safe_get_vertex_ai_location( + litellm_params=litellm_params + ) + if vertex_project is None: raise ValueError( "Missing vertex_project - Set VERTEXAI_PROJECT environment variable or pass vertex_project parameter" @@ -113,7 +121,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): # Ensure no trailing slash api_base = api_base.rstrip("/") - + # Vertex AI DeepSeek OCR endpoint format # Format: https://{region}-aiplatform.googleapis.com/v1/projects/{project}/locations/{region}/endpoints/openapi/chat/completions return f"{api_base}/v1/projects/{vertex_project}/locations/{vertex_location}/endpoints/openapi/chat/completions" @@ -128,63 +136,56 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): ) -> OCRRequestData: """ Transform OCR request to chat completion format for Vertex AI DeepSeek OCR. - + Converts OCR document format to chat completion messages format: - Input: {"type": "image_url", "image_url": "gs://..."} - Output: {"model": "deepseek-ai/deepseek-ocr-maas", "messages": [{"role": "user", "content": [{"type": "image_url", "image_url": "gs://..."}]}]} - + Args: model: Model name (e.g., "deepseek-ai/deepseek-ocr-maas") document: Document dict from user (Mistral OCR format) optional_params: Already mapped optional parameters headers: Request headers **kwargs: Additional arguments - + Returns: OCRRequestData with JSON data in chat completion format """ - verbose_logger.debug("Vertex AI DeepSeek OCR transform_ocr_request (sync) called") - + verbose_logger.debug( + "Vertex AI DeepSeek OCR transform_ocr_request (sync) called" + ) + if not isinstance(document, dict): raise ValueError(f"Expected document dict, got {type(document)}") - + # Extract document type and URL doc_type = document.get("type") image_url = None document_url = None - + if doc_type == "image_url": image_url = document.get("image_url", "") elif doc_type == "document_url": document_url = document.get("document_url", "") else: - raise ValueError(f"Unsupported document type: {doc_type}. Expected 'image_url' or 'document_url'") - + raise ValueError( + f"Unsupported document type: {doc_type}. Expected 'image_url' or 'document_url'" + ) + # Build chat completion message content content_item = {} if image_url: - content_item = { - "type": "image_url", - "image_url": image_url - } + content_item = {"type": "image_url", "image_url": image_url} elif document_url: # For document URLs, we use image_url type as well (Vertex AI supports both) - content_item = { - "type": "image_url", - "image_url": document_url - } - + content_item = {"type": "image_url", "image_url": document_url} + # Build chat completion request data = { "model": "deepseek-ai/" + model, - "messages": [ - { - "role": "user", - "content": [content_item] - } - ] + "messages": [{"role": "user", "content": [content_item]}], } - + # Add optional parameters (stream, temperature, etc.) # Filter out OCR-specific params that don't apply to chat completion chat_completion_params = {} @@ -192,11 +193,13 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): # Include common chat completion params if key in ["stream", "temperature", "max_tokens", "top_p", "n", "stop"]: chat_completion_params[key] = value - + data.update(chat_completion_params) - - verbose_logger.debug("Vertex AI DeepSeek OCR: Transformed request to chat completion format") - + + verbose_logger.debug( + "Vertex AI DeepSeek OCR: Transformed request to chat completion format" + ) + return OCRRequestData(data=data, files=None) async def async_transform_ocr_request( @@ -209,16 +212,16 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): ) -> OCRRequestData: """ Transform OCR request to chat completion format for Vertex AI DeepSeek OCR (async). - + Same as sync version - no async-specific logic needed. - + Args: model: Model name document: Document dict from user optional_params: Already mapped optional parameters headers: Request headers **kwargs: Additional arguments - + Returns: OCRRequestData with JSON data in chat completion format """ @@ -239,7 +242,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): ) -> OCRResponse: """ Transform chat completion response to OCR format. - + Vertex AI DeepSeek OCR returns chat completion format: { "id": "...", @@ -252,35 +255,35 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): }], "usage": {...} } - + We need to extract the content and convert it to OCRResponse format. - + Args: model: Model name raw_response: Raw HTTP response from Vertex AI logging_obj: Logging object **kwargs: Additional arguments - + Returns: OCRResponse in standard format """ verbose_logger.debug("Vertex AI DeepSeek OCR transform_ocr_response called") verbose_logger.debug(f"Raw response: {raw_response.text}") - + try: response_json = raw_response.json() - + # Extract content from chat completion response choices = response_json.get("choices", []) if not choices: raise ValueError("No choices in chat completion response") - + message = choices[0].get("message", {}) content = message.get("content", "") - + if not content: raise ValueError("No content in chat completion response") - + # Try to parse content as JSON (OCR result might be JSON string) ocr_data = None try: @@ -292,28 +295,18 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): else: # If content is markdown text, create a single page with the markdown ocr_data = { - "pages": [ - { - "index": 0, - "markdown": content - } - ], + "pages": [{"index": 0, "markdown": content}], "model": model, - "usage_info": response_json.get("usage", {}) + "usage_info": response_json.get("usage", {}), } except json.JSONDecodeError: # If JSON parsing fails, treat content as markdown ocr_data = { - "pages": [ - { - "index": 0, - "markdown": content - } - ], + "pages": [{"index": 0, "markdown": content}], "model": model, - "usage_info": response_json.get("usage", {}) + "usage_info": response_json.get("usage", {}), } - + # Ensure we have the expected structure if "pages" not in ocr_data: # If OCR data doesn't have pages, wrap the content in a page @@ -321,20 +314,24 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): "pages": [ { "index": 0, - "markdown": content if isinstance(content, str) else json.dumps(content) + "markdown": content + if isinstance(content, str) + else json.dumps(content), } ], "model": ocr_data.get("model", model), - "usage_info": ocr_data.get("usage_info", response_json.get("usage", {})) + "usage_info": ocr_data.get( + "usage_info", response_json.get("usage", {}) + ), } - + # Convert usage info if present usage_info = None if "usage_info" in ocr_data: usage_dict = ocr_data["usage_info"] if isinstance(usage_dict, dict): usage_info = OCRUsageInfo(**usage_dict) - + # Build OCRResponse pages = [] for page_data in ocr_data.get("pages", []): @@ -344,14 +341,18 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): index=page_data.get("index", 0), markdown=page_data.get("markdown", ""), images=page_data.get("images"), - dimensions=page_data.get("dimensions") + dimensions=page_data.get("dimensions"), ) pages.append(page) - + if not pages: # Create a default page if none exist - pages = [OCRPage(index=0, markdown=content if isinstance(content, str) else "")] - + pages = [ + OCRPage( + index=0, markdown=content if isinstance(content, str) else "" + ) + ] + return OCRResponse( pages=pages, model=ocr_data.get("model", model), @@ -359,7 +360,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): usage_info=usage_info, object="ocr", ) - + except Exception as e: verbose_logger.error(f"Error parsing Vertex AI DeepSeek OCR response: {e}") raise e @@ -373,15 +374,15 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): ) -> OCRResponse: """ Async transform chat completion response to OCR format. - + Same as sync version - no async-specific logic needed. - + Args: model: Model name raw_response: Raw HTTP response logging_obj: Logging object **kwargs: Additional arguments - + Returns: OCRResponse in standard format """ @@ -391,4 +392,3 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): logging_obj=logging_obj, **kwargs, ) - diff --git a/litellm/llms/vertex_ai/ocr/transformation.py b/litellm/llms/vertex_ai/ocr/transformation.py index 849e332dae..6fe88459ea 100644 --- a/litellm/llms/vertex_ai/ocr/transformation.py +++ b/litellm/llms/vertex_ai/ocr/transformation.py @@ -17,12 +17,12 @@ from litellm.llms.vertex_ai.vertex_llm_base import VertexBase class VertexAIOCRConfig(MistralOCRConfig): """ Vertex AI Mistral OCR transformation configuration. - + Vertex AI uses Mistral's OCR API format through the Mistral publisher endpoint. Inherits transformation logic from MistralOCRConfig since they use the same format. - + Reference: Vertex AI Mistral OCR documentation - + Important: Vertex AI OCR only supports base64 data URIs (data:image/..., data:application/pdf;base64,...). Regular URLs are not supported. """ @@ -42,16 +42,20 @@ class VertexAIOCRConfig(MistralOCRConfig): ) -> Dict: """ Validate environment and return headers for Vertex AI OCR. - + Vertex AI uses Bearer token authentication with access token from credentials. """ # Extract Vertex AI parameters using safe helpers from VertexBase # Use safe_get_* methods that don't mutate litellm_params dict litellm_params = litellm_params or {} - - vertex_project = VertexBase.safe_get_vertex_ai_project(litellm_params=litellm_params) - vertex_credentials = VertexBase.safe_get_vertex_ai_credentials(litellm_params=litellm_params) - + + vertex_project = VertexBase.safe_get_vertex_ai_project( + litellm_params=litellm_params + ) + vertex_credentials = VertexBase.safe_get_vertex_ai_credentials( + litellm_params=litellm_params + ) + # Get access token from Vertex credentials access_token, project_id = self.vertex_base.get_access_token( credentials=vertex_credentials, @@ -76,25 +80,29 @@ class VertexAIOCRConfig(MistralOCRConfig): ) -> str: """ Get complete URL for Vertex AI OCR endpoint. - - Vertex AI endpoint format: + + Vertex AI endpoint format: https://{location}-aiplatform.googleapis.com/v1/projects/{project}/locations/{location}/publishers/mistralai/ocr - + Args: api_base: Vertex AI API base URL (optional) model: Model name (not used in URL construction) optional_params: Optional parameters litellm_params: LiteLLM parameters containing vertex_project, vertex_location - + Returns: Complete URL for Vertex AI OCR endpoint """ # Extract Vertex AI parameters using safe helpers from VertexBase # Use safe_get_* methods that don't mutate litellm_params dict litellm_params = litellm_params or {} - - vertex_project = VertexBase.safe_get_vertex_ai_project(litellm_params=litellm_params) - vertex_location = VertexBase.safe_get_vertex_ai_location(litellm_params=litellm_params) - + + vertex_project = VertexBase.safe_get_vertex_ai_project( + litellm_params=litellm_params + ) + vertex_location = VertexBase.safe_get_vertex_ai_location( + litellm_params=litellm_params + ) + if vertex_project is None: raise ValueError( "Missing vertex_project - Set VERTEXAI_PROJECT environment variable or pass vertex_project parameter" @@ -109,7 +117,7 @@ class VertexAIOCRConfig(MistralOCRConfig): # Ensure no trailing slash api_base = api_base.rstrip("/") - + # Vertex AI OCR endpoint format for Mistral publisher # Format: https://{region}-aiplatform.googleapis.com/v1/projects/{project}/locations/{region}/publishers/mistralai/models/{model}:rawPredict return f"{api_base}/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/mistralai/models/{model}:rawPredict" @@ -117,47 +125,55 @@ class VertexAIOCRConfig(MistralOCRConfig): def _convert_url_to_data_uri_sync(self, url: str) -> str: """ Synchronously convert a URL to a base64 data URI. - + Vertex AI OCR doesn't have internet access, so we need to fetch URLs and convert them to base64 data URIs. - + Args: url: The URL to convert - + Returns: Base64 data URI string """ - verbose_logger.debug(f"Vertex AI OCR: Converting URL to base64 data URI (sync): {url}") - + verbose_logger.debug( + f"Vertex AI OCR: Converting URL to base64 data URI (sync): {url}" + ) + # Fetch and convert to base64 data URI # convert_url_to_base64 already returns a full data URI like "data:image/jpeg;base64,..." data_uri = convert_url_to_base64(url=url) - - verbose_logger.debug(f"Vertex AI OCR: Converted URL to data URI (length: {len(data_uri)})") - + + verbose_logger.debug( + f"Vertex AI OCR: Converted URL to data URI (length: {len(data_uri)})" + ) + return data_uri async def _convert_url_to_data_uri_async(self, url: str) -> str: """ Asynchronously convert a URL to a base64 data URI. - + Vertex AI OCR doesn't have internet access, so we need to fetch URLs and convert them to base64 data URIs. - + Args: url: The URL to convert - + Returns: Base64 data URI string """ - verbose_logger.debug(f"Vertex AI OCR: Converting URL to base64 data URI (async): {url}") - + verbose_logger.debug( + f"Vertex AI OCR: Converting URL to base64 data URI (async): {url}" + ) + # Fetch and convert to base64 data URI asynchronously # async_convert_url_to_base64 already returns a full data URI like "data:image/jpeg;base64,..." data_uri = await async_convert_url_to_base64(url=url) - - verbose_logger.debug(f"Vertex AI OCR: Converted URL to data URI (length: {len(data_uri)})") - + + verbose_logger.debug( + f"Vertex AI OCR: Converted URL to data URI (length: {len(data_uri)})" + ) + return data_uri def transform_ocr_request( @@ -170,29 +186,29 @@ class VertexAIOCRConfig(MistralOCRConfig): ) -> OCRRequestData: """ Transform OCR request for Vertex AI, converting URLs to base64 data URIs (sync). - + Vertex AI OCR doesn't have internet access, so we automatically fetch any URLs and convert them to base64 data URIs synchronously. - + Args: model: Model name document: Document dict from user optional_params: Already mapped optional parameters headers: Request headers **kwargs: Additional arguments - + Returns: OCRRequestData with JSON data """ verbose_logger.debug("Vertex AI OCR transform_ocr_request (sync) called") - + if not isinstance(document, dict): raise ValueError(f"Expected document dict, got {type(document)}") - + # Check if we need to convert URL to base64 doc_type = document.get("type") transformed_document = document.copy() - + if doc_type == "document_url": document_url = document.get("document_url", "") # If it's not already a data URI, convert it @@ -211,7 +227,7 @@ class VertexAIOCRConfig(MistralOCRConfig): ) data_uri = self._convert_url_to_data_uri_sync(url=image_url) transformed_document["image_url"] = data_uri - + # Call parent's transform to build the request return super().transform_ocr_request( model=model, @@ -231,29 +247,31 @@ class VertexAIOCRConfig(MistralOCRConfig): ) -> OCRRequestData: """ Transform OCR request for Vertex AI, converting URLs to base64 data URIs (async). - + Vertex AI OCR doesn't have internet access, so we automatically fetch any URLs and convert them to base64 data URIs asynchronously. - + Args: model: Model name document: Document dict from user optional_params: Already mapped optional parameters headers: Request headers **kwargs: Additional arguments - + Returns: OCRRequestData with JSON data """ - verbose_logger.debug(f"Vertex AI OCR async_transform_ocr_request - model: {model}") - + verbose_logger.debug( + f"Vertex AI OCR async_transform_ocr_request - model: {model}" + ) + if not isinstance(document, dict): raise ValueError(f"Expected document dict, got {type(document)}") - + # Check if we need to convert URL to base64 doc_type = document.get("type") transformed_document = document.copy() - + if doc_type == "document_url": document_url = document.get("document_url", "") # If it's not already a data URI, convert it @@ -272,7 +290,7 @@ class VertexAIOCRConfig(MistralOCRConfig): ) data_uri = await self._convert_url_to_data_uri_async(url=image_url) transformed_document["image_url"] = data_uri - + # Call parent's transform to build the request return super().transform_ocr_request( model=model, @@ -281,4 +299,3 @@ class VertexAIOCRConfig(MistralOCRConfig): headers=headers, **kwargs, ) - diff --git a/litellm/llms/vertex_ai/rag_engine/__init__.py b/litellm/llms/vertex_ai/rag_engine/__init__.py index 2a88b43f5a..79b9e2c132 100644 --- a/litellm/llms/vertex_ai/rag_engine/__init__.py +++ b/litellm/llms/vertex_ai/rag_engine/__init__.py @@ -11,4 +11,3 @@ __all__ = [ "VertexAIRAGIngestion", "VertexAIRAGTransformation", ] - diff --git a/litellm/llms/vertex_ai/rag_engine/ingestion.py b/litellm/llms/vertex_ai/rag_engine/ingestion.py index 6b435a46bc..2ec6166779 100644 --- a/litellm/llms/vertex_ai/rag_engine/ingestion.py +++ b/litellm/llms/vertex_ai/rag_engine/ingestion.py @@ -79,10 +79,9 @@ class VertexAIRAGIngestion(BaseRAGIngestion): ) # GCP config - self.vertex_project = ( - self.vector_store_config.get("vertex_project") - or get_secret_str("VERTEXAI_PROJECT") - ) + self.vertex_project = self.vector_store_config.get( + "vertex_project" + ) or get_secret_str("VERTEXAI_PROJECT") self.vertex_location = ( self.vector_store_config.get("vertex_location") or get_secret_str("VERTEXAI_LOCATION") @@ -91,9 +90,8 @@ class VertexAIRAGIngestion(BaseRAGIngestion): self.vertex_credentials = self.vector_store_config.get("vertex_credentials") # GCS bucket for file uploads - self.gcs_bucket = ( - self.vector_store_config.get("gcs_bucket") - or os.environ.get("GCS_BUCKET_NAME") + self.gcs_bucket = self.vector_store_config.get("gcs_bucket") or os.environ.get( + "GCS_BUCKET_NAME" ) if not self.gcs_bucket: raise ValueError( @@ -312,4 +310,3 @@ class VertexAIRAGIngestion(BaseRAGIngestion): raise RuntimeError(f"Failed to import file into RAG corpus: {e}") from e return str(self.corpus_id), gcs_uri - diff --git a/litellm/llms/vertex_ai/rag_engine/transformation.py b/litellm/llms/vertex_ai/rag_engine/transformation.py index 7e70202fb7..ed5154bbdf 100644 --- a/litellm/llms/vertex_ai/rag_engine/transformation.py +++ b/litellm/llms/vertex_ai/rag_engine/transformation.py @@ -121,9 +121,7 @@ class VertexAIRAGTransformation(VertexBase): return { "import_rag_files_config": { - "gcs_source": { - "uris": [gcs_uri] - }, + "gcs_source": {"uris": [gcs_uri]}, "rag_file_transformation_config": transformation_config, } } @@ -153,4 +151,3 @@ class VertexAIRAGTransformation(VertexBase): "Authorization": f"Bearer {access_token}", "Content-Type": "application/json", } - diff --git a/litellm/llms/vertex_ai/realtime/transformation.py b/litellm/llms/vertex_ai/realtime/transformation.py index 5eae143175..2b4746b174 100644 --- a/litellm/llms/vertex_ai/realtime/transformation.py +++ b/litellm/llms/vertex_ai/realtime/transformation.py @@ -35,7 +35,10 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig): # ------------------------------------------------------------------ def get_complete_url( - self, api_base: Optional[str], model: str, api_key: Optional[str] = None # noqa: ARG002 + self, + api_base: Optional[str], + model: str, + api_key: Optional[str] = None, # noqa: ARG002 ) -> str: """ Build the Vertex AI Live WSS endpoint URL. diff --git a/litellm/llms/vertex_ai/rerank/transformation.py b/litellm/llms/vertex_ai/rerank/transformation.py index 953c6c84ea..5365183967 100644 --- a/litellm/llms/vertex_ai/rerank/transformation.py +++ b/litellm/llms/vertex_ai/rerank/transformation.py @@ -13,14 +13,18 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.secret_managers.main import get_secret_str -from litellm.types.rerank import RerankResponse, RerankResponseMeta, RerankBilledUnits, RerankResponseResult - +from litellm.types.rerank import ( + RerankResponse, + RerankResponseMeta, + RerankBilledUnits, + RerankResponseResult, +) class VertexAIRerankConfig(BaseRerankConfig, VertexBase): """ Configuration for Vertex AI Discovery Engine Rerank API - + Reference: https://cloud.google.com/generative-ai-app-builder/docs/ranking#rank_or_rerank_a_set_of_records_according_to_a_query """ @@ -28,8 +32,8 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase): super().__init__() def get_complete_url( - self, - api_base: Optional[str], + self, + api_base: Optional[str], model: str, optional_params: Optional[Dict] = None, ) -> str: @@ -38,11 +42,11 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase): """ # Try to get project ID from optional_params first (e.g., vertex_project parameter) params = optional_params or {} - + # Get credentials to extract project ID if needed vertex_credentials = self.safe_get_vertex_ai_credentials(params.copy()) vertex_project = self.safe_get_vertex_ai_project(params.copy()) - + # Use _ensure_access_token to extract project_id from credentials # This is the same method used in vertex embeddings _, vertex_project = self._ensure_access_token( @@ -50,19 +54,19 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase): project_id=vertex_project, custom_llm_provider="vertex_ai", ) - + # Fallback to environment or litellm config project_id = ( vertex_project - or get_secret_str("VERTEXAI_PROJECT") + or get_secret_str("VERTEXAI_PROJECT") or litellm.vertex_project ) - + if not project_id: raise ValueError( "Vertex AI project ID is required. Please set 'VERTEXAI_PROJECT', 'litellm.vertex_project', or pass 'vertex_project' parameter" ) - + return f"https://discoveryengine.googleapis.com/v1/projects/{project_id}/locations/global/rankingConfigs/default_ranking_config:rank" def validate_environment( @@ -79,14 +83,14 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase): litellm_params = optional_params.copy() if optional_params else {} vertex_credentials = self.safe_get_vertex_ai_credentials(litellm_params) vertex_project = self.safe_get_vertex_ai_project(litellm_params) - + # Get access token using the base class method access_token, project_id = self._ensure_access_token( credentials=vertex_credentials, project_id=vertex_project, custom_llm_provider="vertex_ai", ) - + default_headers = { "Authorization": f"Bearer {access_token}", "Content-Type": "application/json", @@ -113,12 +117,12 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase): raise ValueError("query is required for Vertex AI rerank") if "documents" not in optional_rerank_params: raise ValueError("documents is required for Vertex AI rerank") - + query = optional_rerank_params["query"] documents = optional_rerank_params["documents"] top_n = optional_rerank_params.get("top_n", None) return_documents = optional_rerank_params.get("return_documents", True) - + # Convert documents to records format records = [] for idx, document in enumerate(documents): @@ -129,26 +133,18 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase): # Handle dict format content = document.get("text", str(document)) title = document.get("title", " ".join(content.split()[:3])) - - records.append({ - "id": str(idx), - "title": title, - "content": content - }) - - request_data = { - "model": model, - "query": query, - "records": records - } - + + records.append({"id": str(idx), "title": title, "content": content}) + + request_data = {"model": model, "query": query, "records": records} + if top_n is not None: request_data["topN"] = top_n - + # Map return_documents to ignoreRecordDetailsInResponse # When return_documents is False, we want to ignore record details (return only IDs) request_data["ignoreRecordDetailsInResponse"] = not return_documents - + return request_data def transform_rerank_response( @@ -172,54 +168,55 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase): # Extract records from response records = raw_response_json.get("records", []) - + # Convert to Cohere format results = [] for record in records: # Handle both cases: with full details and with only IDs if "score" in record: # Full response with score and details - results.append({ - "index": int(record["id"]), - "relevance_score": record.get("score", 0.0) - }) + results.append( + { + "index": int(record["id"]), + "relevance_score": record.get("score", 0.0), + } + ) else: # Response with only IDs (when ignoreRecordDetailsInResponse=true) # We can't provide a relevance score, so we'll use a default - results.append({ - "index": int(record["id"]), - "relevance_score": 1.0 # Default score when details are ignored - }) - + results.append( + { + "index": int(record["id"]), + "relevance_score": 1.0, # Default score when details are ignored + } + ) + # Sort by relevance score (descending) results.sort(key=lambda x: x["relevance_score"], reverse=True) - - # Create response in Cohere format + + # Create response in Cohere format # Convert results to proper RerankResponseResult objects rerank_results = [] for result in results: - rerank_results.append(RerankResponseResult( - index=result["index"], - relevance_score=result["relevance_score"] - )) - + rerank_results.append( + RerankResponseResult( + index=result["index"], relevance_score=result["relevance_score"] + ) + ) + # Create meta object meta = RerankResponseMeta( - billed_units=RerankBilledUnits( - search_units=len(records) - ) + billed_units=RerankBilledUnits(search_units=len(records)) ) - + return RerankResponse( - id=f"vertex_ai_rerank_{model}", - results=rerank_results, - meta=meta + id=f"vertex_ai_rerank_{model}", results=rerank_results, meta=meta ) def get_supported_cohere_rerank_params(self, model: str) -> list: return [ "query", - "documents", + "documents", "top_n", "return_documents", ] @@ -249,4 +246,3 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase): } result.update(non_default_params) return result - diff --git a/litellm/llms/vertex_ai/text_to_speech/transformation.py b/litellm/llms/vertex_ai/text_to_speech/transformation.py index 18ca077c4d..be7bcfcadd 100644 --- a/litellm/llms/vertex_ai/text_to_speech/transformation.py +++ b/litellm/llms/vertex_ai/text_to_speech/transformation.py @@ -164,12 +164,14 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): voice_str = voice.get("name") if voice else None # Store credentials in litellm_params for use in transform methods - litellm_params_dict.update({ - "vertex_credentials": vertex_credentials, - "vertex_project": vertex_project, - "vertex_location": vertex_location, - "api_base": api_base, - }) + litellm_params_dict.update( + { + "vertex_credentials": vertex_credentials, + "vertex_project": vertex_project, + "vertex_location": vertex_location, + "api_base": api_base, + } + ) # Call the text_to_speech_handler response = base_llm_http_handler.text_to_speech_handler( @@ -328,7 +330,9 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): if not input_data: raise ValueError("Either 'text' or 'ssml' must be provided.") if "text" in input_data and "ssml" in input_data: - raise ValueError("Only one of 'text' or 'ssml' should be provided, not both.") + raise ValueError( + "Only one of 'text' or 'ssml' should be provided, not both." + ) return input_data @@ -389,9 +393,8 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): # Check for voice dict stored in: # 1. litellm_params by dispatch method # 2. optional_params by map_openai_params - voice_dict = ( - litellm_params.get("vertex_voice_dict") - or optional_params.get("vertex_voice_dict") + voice_dict = litellm_params.get("vertex_voice_dict") or optional_params.get( + "vertex_voice_dict" ) if voice_dict is not None and isinstance(voice_dict, dict): vertex_voice = VertexTextToSpeechVoice(**voice_dict) @@ -414,12 +417,16 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): ) # Build audio configuration - audio_encoding = optional_params.get("audioEncoding", self.DEFAULT_AUDIO_ENCODING) + audio_encoding = optional_params.get( + "audioEncoding", self.DEFAULT_AUDIO_ENCODING + ) speaking_rate = optional_params.get("speakingRate", self.DEFAULT_SPEAKING_RATE) # Check for full audioConfig in optional_params if "audioConfig" in optional_params: - vertex_audio_config = VertexTextToSpeechAudioConfig(**optional_params["audioConfig"]) + vertex_audio_config = VertexTextToSpeechAudioConfig( + **optional_params["audioConfig"] + ) else: vertex_audio_config = VertexTextToSpeechAudioConfig( audioEncoding=audio_encoding, diff --git a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py index 1be9cd820a..4baa5774c4 100644 --- a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py @@ -162,7 +162,6 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): Transform Vertex AI RAG API response to standard vector store search response """ try: - response_json = response.json() # Extract contexts from Vertex AI response - handle nested structure contexts = response_json.get("contexts", {}).get("contexts", []) diff --git a/litellm/llms/vertex_ai/vertex_ai_aws_wif.py b/litellm/llms/vertex_ai/vertex_ai_aws_wif.py index 44a0016e4e..a03a4e37a2 100644 --- a/litellm/llms/vertex_ai/vertex_ai_aws_wif.py +++ b/litellm/llms/vertex_ai/vertex_ai_aws_wif.py @@ -18,18 +18,20 @@ GOOGLE_IMPORT_ERROR_MESSAGE = ( # AWS params recognized in WIF credential JSON for explicit auth. # These match the kwargs accepted by BaseAWSLLM.get_credentials(). -_AWS_CREDENTIAL_KEYS = frozenset({ - "aws_access_key_id", - "aws_secret_access_key", - "aws_session_token", - "aws_region_name", - "aws_session_name", - "aws_profile_name", - "aws_role_name", - "aws_web_identity_token", - "aws_sts_endpoint", - "aws_external_id", -}) +_AWS_CREDENTIAL_KEYS = frozenset( + { + "aws_access_key_id", + "aws_secret_access_key", + "aws_session_token", + "aws_region_name", + "aws_session_name", + "aws_profile_name", + "aws_role_name", + "aws_web_identity_token", + "aws_sts_endpoint", + "aws_external_id", + } +) class VertexAIAwsWifAuth: @@ -46,11 +48,7 @@ class VertexAIAwsWifAuth: Returns a dict of {param_name: value} for any recognized aws_* keys found in the JSON. Returns empty dict if none are present. """ - return { - key: json_obj[key] - for key in _AWS_CREDENTIAL_KEYS - if key in json_obj - } + return {key: json_obj[key] for key in _AWS_CREDENTIAL_KEYS if key in json_obj} @staticmethod def credentials_from_explicit_aws(json_obj, aws_params, scopes): diff --git a/litellm/llms/vertex_ai/vertex_ai_non_gemini.py b/litellm/llms/vertex_ai/vertex_ai_non_gemini.py index 54cb83bb0b..cfbab584f6 100644 --- a/litellm/llms/vertex_ai/vertex_ai_non_gemini.py +++ b/litellm/llms/vertex_ai/vertex_ai_non_gemini.py @@ -145,11 +145,9 @@ def completion( # noqa: PLR0915 json_obj = json.loads(vertex_credentials) - creds = ( - google.oauth2.service_account.Credentials.from_service_account_info( - json_obj, - scopes=["https://www.googleapis.com/auth/cloud-platform"], - ) + creds = google.oauth2.service_account.Credentials.from_service_account_info( + json_obj, + scopes=["https://www.googleapis.com/auth/cloud-platform"], ) else: creds, _ = google.auth.default(quota_project_id=vertex_project) diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py index 6bede1a235..d3b0217d04 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py @@ -33,10 +33,12 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert """ vertex_ai_project = VertexBase.safe_get_vertex_ai_project(litellm_params) vertex_ai_location = VertexBase.safe_get_vertex_ai_location(litellm_params) - + project_id: Optional[str] = None if "Authorization" not in headers: - vertex_credentials = VertexBase.safe_get_vertex_ai_credentials(litellm_params) + vertex_credentials = VertexBase.safe_get_vertex_ai_credentials( + litellm_params + ) access_token, project_id = self._ensure_access_token( credentials=vertex_credentials, @@ -62,11 +64,11 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert ) headers["content-type"] = "application/json" - + # Add beta headers for Vertex AI tools = optional_params.get("tools", []) beta_values: set[str] = set() - + # Get existing beta headers if any existing_beta = headers.get("anthropic-beta") if existing_beta: @@ -79,36 +81,42 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert edits = context_management_param.get("edits", []) has_compact = False has_other = False - + for edit in edits: edit_type = edit.get("type", "") if edit_type == "compact_20260112": has_compact = True else: has_other = True - + # Add compact header if any compact edits exist if has_compact: beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value) - + # Add context management header if any other edits exist if has_other: - beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value) + beta_values.add( + ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value + ) # Check for web search tool for tool in tools: - if isinstance(tool, dict) and tool.get("type", "").startswith(ANTHROPIC_HOSTED_TOOLS.WEB_SEARCH.value): - beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.WEB_SEARCH_2025_03_05.value) + if isinstance(tool, dict) and tool.get("type", "").startswith( + ANTHROPIC_HOSTED_TOOLS.WEB_SEARCH.value + ): + beta_values.add( + ANTHROPIC_BETA_HEADER_VALUES.WEB_SEARCH_2025_03_05.value + ) break - + # Check for tool search tools - Vertex AI uses different beta header anthropic_model_info = AnthropicModelInfo() if anthropic_model_info.is_tool_search_used(tools): beta_values.add(get_tool_search_beta_header("vertex_ai")) - + if beta_values: headers["anthropic-beta"] = ",".join(beta_values) - + return headers, api_base def get_complete_url( diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py index 4e2c2895f9..504914c479 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py @@ -107,7 +107,7 @@ class VertexAIAnthropicConfig(AnthropicConfig): # VertexAI doesn't support output_format parameter, remove it if present data.pop("output_format", None) - + # VertexAI doesn't support output_config parameter, remove it if present data.pop("output_config", None) diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/transformation.py index 86e36e802e..47c388f0a5 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/transformation.py @@ -8,9 +8,10 @@ class VertexAIGPTOSSTransformation(OpenAIGPTConfig): https://console.cloud.google.com/vertex-ai/publishers/openai/model-garden/gpt-oss-120b-maas?hl=id """ + def __init__(self): super().__init__() - + def get_supported_openai_params(self, model: str) -> list: base_gpt_series_params = super().get_supported_openai_params(model=model) gpt_oss_only_params = ["reasoning_effort"] @@ -20,8 +21,16 @@ class VertexAIGPTOSSTransformation(OpenAIGPTConfig): # VertexAI - GPT-OSS does not support tool calls ######################################################### if litellm.supports_function_calling(model=model) is False: - TOOL_CALLING_PARAMS_TO_REMOVE = ["tool", "tool_choice", "function_call", "functions"] - base_gpt_series_params = [param for param in base_gpt_series_params if param not in TOOL_CALLING_PARAMS_TO_REMOVE] + TOOL_CALLING_PARAMS_TO_REMOVE = [ + "tool", + "tool_choice", + "function_call", + "functions", + ] + base_gpt_series_params = [ + param + for param in base_gpt_series_params + if param not in TOOL_CALLING_PARAMS_TO_REMOVE + ] return base_gpt_series_params - diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py index 51310e4fa8..b38b453506 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py @@ -151,12 +151,12 @@ class VertexAILlama3StreamingHandler(OpenAIChatCompletionStreamingHandler): """ Vertex AI Llama models may not include role in streaming chunk deltas. This handler ensures the first chunk always has role="assistant". - + When Vertex AI returns a single chunk with both role and finish_reason (empty response), this handler splits it into two chunks: 1. First chunk: role="assistant", content="", finish_reason=None 2. Second chunk: role=None, content=None, finish_reason="stop" - + This matches OpenAI's streaming format where the first chunk has role and the final chunk has finish_reason but no role. """ @@ -171,7 +171,7 @@ class VertexAILlama3StreamingHandler(OpenAIChatCompletionStreamingHandler): if not self.sent_role and result.choices: delta = result.choices[0].delta finish_reason = result.choices[0].finish_reason - + # If this is both the first chunk AND the final chunk (has finish_reason), # we need to split it into two chunks to match OpenAI format if finish_reason is not None: @@ -202,7 +202,9 @@ class VertexAILlama3StreamingHandler(OpenAIChatCompletionStreamingHandler): elif delta.role is None: delta.role = "assistant" # If the first chunk has empty content, ensure it's still emitted - if (delta.content == "" or delta.content is None) and delta.provider_specific_fields is None: + if ( + delta.content == "" or delta.content is None + ) and delta.provider_specific_fields is None: delta.provider_specific_fields = {} self.sent_role = True return result diff --git a/litellm/llms/vertex_ai/vertex_embeddings/bge.py b/litellm/llms/vertex_ai/vertex_embeddings/bge.py index 2eff0ba96d..e3f25b425f 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/bge.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/bge.py @@ -25,14 +25,14 @@ from .types import ( class VertexBGEConfig: """ Configuration and transformation logic for BGE models on Vertex AI. - + BGE (BAAI General Embedding) models use a different request format where the input field is named "prompt" instead of "content". - + Supported model patterns (after provider split in main.py): - "bge-small-en-v1.5" (model name) - "bge/204379420394258432" (endpoint ID pattern) - + Note: Model name transformation (bge/ -> numeric ID) is handled automatically in common_utils._get_vertex_url(). This class focuses on request/response format only. """ @@ -41,14 +41,14 @@ class VertexBGEConfig: def is_bge_model(model: str) -> bool: """ Check if the model is a BGE (BAAI General Embedding) model. - + After provider split in main.py, supports: - "bge-small-en-v1.5" (model name) - "bge/204379420394258432" (endpoint ID pattern) - + Args: model: The model name after provider split - + Returns: bool: True if the model is a BGE model """ @@ -62,14 +62,14 @@ class VertexBGEConfig: ) -> VertexEmbeddingRequest: """ Transforms an OpenAI request to a Vertex BGE embedding request. - + BGE models use "prompt" instead of "content" as the input field. - + Args: input: The input text(s) to embed optional_params: Optional parameters for the request model: The model name - + Returns: VertexEmbeddingRequest: The transformed request """ @@ -124,7 +124,7 @@ class VertexBGEConfig: ) -> EmbeddingResponse: """ Transforms a Vertex BGE embedding response to OpenAI format. - + BGE models return embeddings directly as arrays in predictions: { "predictions": [ @@ -132,26 +132,28 @@ class VertexBGEConfig: [0.003, 0.022, ...] ] } - + Args: response: The raw response from Vertex AI model: The model name model_response: The EmbeddingResponse object to populate - + Returns: EmbeddingResponse: The transformed response in OpenAI format - + Raises: KeyError: If response doesn't contain 'predictions' ValueError: If predictions is not a list or contains invalid data """ if "predictions" not in response: raise KeyError("Response missing 'predictions' field") - + _predictions = response["predictions"] - + if not isinstance(_predictions, list): - raise ValueError(f"Expected 'predictions' to be a list, got {type(_predictions)}") + raise ValueError( + f"Expected 'predictions' to be a list, got {type(_predictions)}" + ) embedding_response = [] # BGE models don't return token counts, so we estimate or set to 0 @@ -162,7 +164,7 @@ class VertexBGEConfig: raise ValueError( f"Expected embedding at index {idx} to be a list, got {type(embedding_values)}" ) - + embedding_response.append( { "object": "embedding", @@ -179,4 +181,3 @@ class VertexBGEConfig: ) setattr(model_response, "usage", usage) return model_response - diff --git a/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py b/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py index 8a03738ad7..5fffd983c2 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py @@ -74,7 +74,7 @@ class VertexEmbedding(VertexBase): ) # Extract use_psc_endpoint_format from optional_params use_psc_endpoint_format = optional_params.get("use_psc_endpoint_format", False) - + auth_header, api_base = self._get_token_and_url( model=model, gemini_api_key=gemini_api_key, @@ -90,10 +90,8 @@ class VertexEmbedding(VertexBase): use_psc_endpoint_format=use_psc_endpoint_format, ) headers = self.set_headers(auth_header=auth_header, extra_headers=extra_headers) - vertex_request: VertexEmbeddingRequest = ( - litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request( - input=input, optional_params=optional_params, model=model - ) + vertex_request: VertexEmbeddingRequest = litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request( + input=input, optional_params=optional_params, model=model ) _client_params = {} @@ -170,7 +168,7 @@ class VertexEmbedding(VertexBase): ) # Extract use_psc_endpoint_format from optional_params use_psc_endpoint_format = optional_params.get("use_psc_endpoint_format", False) - + auth_header, api_base = self._get_token_and_url( model=model, gemini_api_key=gemini_api_key, @@ -186,10 +184,8 @@ class VertexEmbedding(VertexBase): use_psc_endpoint_format=use_psc_endpoint_format, ) headers = self.set_headers(auth_header=auth_header, extra_headers=extra_headers) - vertex_request: VertexEmbeddingRequest = ( - litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request( - input=input, optional_params=optional_params, model=model - ) + vertex_request: VertexEmbeddingRequest = litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request( + input=input, optional_params=optional_params, model=model ) _async_client_params = {} diff --git a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py index 5a3a4a7188..132f29987a 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py @@ -107,6 +107,7 @@ class VertexAITextEmbeddingConfig(BaseModel): """ # Import here to avoid circular import issues with litellm.__init__ from litellm.llms.vertex_ai.vertex_embeddings.bge import VertexBGEConfig + if model.isdigit(): return self._transform_openai_request_to_fine_tuned_embedding_request( input, optional_params, model @@ -174,7 +175,10 @@ class VertexAITextEmbeddingConfig(BaseModel): **optional_params ) # Remove 'shared_session' from parameters if present - if vertex_request["parameters"] is not None and "shared_session" in vertex_request["parameters"]: + if ( + vertex_request["parameters"] is not None + and "shared_session" in vertex_request["parameters"] + ): del vertex_request["parameters"]["shared_session"] # type: ignore[typeddict-item] return vertex_request @@ -215,10 +219,10 @@ class VertexAITextEmbeddingConfig(BaseModel): return self._transform_vertex_response_to_openai_for_fine_tuned_models( response, model, model_response ) - + # Import here to avoid circular import issues with litellm.__init__ from litellm.llms.vertex_ai.vertex_embeddings.bge import VertexBGEConfig - + if VertexBGEConfig.is_bge_model(model): return VertexBGEConfig.transform_response( response=response, model=model, model_response=model_response diff --git a/litellm/llms/vertex_ai/vertex_embeddings/types.py b/litellm/llms/vertex_ai/vertex_embeddings/types.py index fa9794d79a..317b9c4fb8 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/types.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/types.py @@ -50,7 +50,11 @@ class EmbeddingParameters(TypedDict, total=False): class VertexEmbeddingRequest(TypedDict, total=False): - instances: Union[List[TextEmbeddingInput], List[TextEmbeddingBGEInput], List[TextEmbeddingFineTunedInput]] + instances: Union[ + List[TextEmbeddingInput], + List[TextEmbeddingBGEInput], + List[TextEmbeddingFineTunedInput], + ] parameters: Optional[Union[EmbeddingParameters, TextEmbeddingFineTunedParameters]] diff --git a/litellm/llms/vertex_ai/vertex_gemma_models/__init__.py b/litellm/llms/vertex_ai/vertex_gemma_models/__init__.py index d06c7a5cd7..92106ab7c2 100644 --- a/litellm/llms/vertex_ai/vertex_gemma_models/__init__.py +++ b/litellm/llms/vertex_ai/vertex_gemma_models/__init__.py @@ -1,2 +1 @@ """Vertex AI Gemma-AI Models Handler""" - diff --git a/litellm/llms/vertex_ai/vertex_gemma_models/main.py b/litellm/llms/vertex_ai/vertex_gemma_models/main.py index 41bd6b5431..82cfe6de98 100644 --- a/litellm/llms/vertex_ai/vertex_gemma_models/main.py +++ b/litellm/llms/vertex_ai/vertex_gemma_models/main.py @@ -82,7 +82,6 @@ class VertexAIGemmaModels(VertexBase): message="""Upgrade vertex ai. Run `pip install "google-cloud-aiplatform>=1.38"`""", ) try: - model = get_vertex_base_model_name(model=model) vertex_httpx_logic = VertexLLM() @@ -143,4 +142,3 @@ class VertexAIGemmaModels(VertexBase): if hasattr(e, "status_code"): raise e raise VertexAIError(status_code=500, message=str(e)) - diff --git a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py index 24b53f0ba4..6c6446958b 100644 --- a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py +++ b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py @@ -21,7 +21,7 @@ from litellm.types.utils import ModelResponse class VertexGemmaConfig(OpenAIGPTConfig): """ Configuration and transformation class for Vertex AI Gemma models - + Extends OpenAIGPTConfig to wrap/unwrap the instances/predictions format used by Vertex AI's Gemma deployment endpoint. """ @@ -48,16 +48,17 @@ class VertexGemmaConfig(OpenAIGPTConfig): ) -> Union[ModelResponse, Any]: """ Helper method to return fake stream iterator if streaming is requested. - + Args: model_response: The completed model response stream: Whether streaming was requested - + Returns: MockResponseIterator if stream=True, otherwise the model_response """ if stream: from litellm.llms.base_llm.base_model_iterator import MockResponseIterator + return MockResponseIterator(model_response=model_response) return model_response @@ -71,7 +72,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): ) -> dict: """ Transform request to Vertex Gemma format. - + Uses parent class to create OpenAI-compatible request, then wraps it in the Vertex Gemma instances format. """ @@ -83,12 +84,14 @@ class VertexGemmaConfig(OpenAIGPTConfig): litellm_params=litellm_params, headers=headers, ) - + # Remove params not needed/supported by Vertex Gemma openai_request.pop("model", None) - openai_request.pop("stream", None) # Streaming not supported, will be faked client-side + openai_request.pop( + "stream", None + ) # Streaming not supported, will be faked client-side openai_request.pop("stream_options", None) # Stream options not supported - + # Wrap in Vertex Gemma format return { "instances": [ @@ -105,7 +108,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): ) -> Dict[str, Any]: """ Unwrap the Vertex Gemma predictions format to OpenAI format. - + Vertex Gemma wraps the OpenAI-compatible response in a 'predictions' field. This method extracts it so the parent class can process it normally. """ @@ -114,7 +117,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): status_code=422, message="Invalid response format: missing 'predictions' field", ) - + return response_json["predictions"] def completion( @@ -189,7 +192,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): # Check if streaming is requested (will be faked) stream = optional_params.get("stream", False) - + # Transform the request using parent class methods request_data = self.transform_request( model=model, @@ -198,7 +201,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): litellm_params=litellm_params, headers={}, ) - + # Set up headers headers = { "Authorization": f"Bearer {api_key}", @@ -231,10 +234,10 @@ class VertexGemmaConfig(OpenAIGPTConfig): ) response_json = response.json() - + # Unwrap predictions to get OpenAI-compatible response openai_response = self._unwrap_predictions_response(response_json) - + # Use litellm's standard response converter model_response = cast( ModelResponse, @@ -244,10 +247,10 @@ class VertexGemmaConfig(OpenAIGPTConfig): _response_headers={}, ), ) - + # Ensure model is set correctly model_response.model = model - + # Log the response logging_obj.post_call( input=messages, @@ -255,9 +258,11 @@ class VertexGemmaConfig(OpenAIGPTConfig): original_response=response_json, additional_args={"complete_input_dict": request_data}, ) - + # Return fake stream iterator if streaming was requested - return self._handle_fake_stream_response(model_response=model_response, stream=stream) + return self._handle_fake_stream_response( + model_response=model_response, stream=stream + ) async def _async_completion( self, @@ -280,7 +285,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): # Check if streaming is requested (will be faked) stream = optional_params.get("stream", False) - + # Transform the request using parent class async methods request_data = await self.async_transform_request( model=model, @@ -289,7 +294,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): litellm_params=litellm_params, headers={}, ) - + # Set up headers headers = { "Authorization": f"Bearer {api_key}", @@ -324,10 +329,10 @@ class VertexGemmaConfig(OpenAIGPTConfig): ) response_json = response.json() - + # Unwrap predictions to get OpenAI-compatible response openai_response = self._unwrap_predictions_response(response_json) - + # Use litellm's standard response converter model_response = cast( ModelResponse, @@ -337,10 +342,10 @@ class VertexGemmaConfig(OpenAIGPTConfig): _response_headers={}, ), ) - + # Ensure model is set correctly model_response.model = model - + # Log the response logging_obj.post_call( input=messages, @@ -348,7 +353,8 @@ class VertexGemmaConfig(OpenAIGPTConfig): original_response=response_json, additional_args={"complete_input_dict": request_data}, ) - - # Return fake stream iterator if streaming was requested - return self._handle_fake_stream_response(model_response=model_response, stream=stream) + # Return fake stream iterator if streaming was requested + return self._handle_fake_stream_response( + model_response=model_response, stream=stream + ) diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 86e14a30df..1a29ba82ea 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -21,7 +21,6 @@ from .common_utils import ( all_gemini_url_modes, get_vertex_base_model_name, get_vertex_base_url, - is_global_only_vertex_model, ) GOOGLE_IMPORT_ERROR_MESSAGE = ( @@ -49,8 +48,32 @@ class VertexBase: self.async_handler: Optional[AsyncHTTPHandler] = None def get_vertex_region(self, vertex_region: Optional[str], model: str) -> str: - if is_global_only_vertex_model(model): - return "global" + import litellm + + # Try to get supported_regions directly from model_cost + # Check both with and without vertex_ai/ prefix + model_key = ( + f"vertex_ai/{model}" if not model.startswith("vertex_ai/") else model + ) + model_info = litellm.model_cost.get(model_key, {}) + supported_regions = model_info.get("supported_regions") + + if supported_regions and len(supported_regions) > 0: + # If user didn't specify region, use the first supported region + if vertex_region is None: + return supported_regions[0] + # If user specified a region not supported by this model, override it + if vertex_region not in supported_regions: + verbose_logger.warning( + "Vertex AI model '%s' does not support region '%s' " + "(supported: %s). Routing to '%s'.", + model, + vertex_region, + supported_regions, + supported_regions[0], + ) + return supported_regions[0] + return vertex_region return vertex_region or "us-central1" def load_auth( @@ -214,7 +237,9 @@ class VertexBase: ) -> str: if api_base: return api_base - return get_vertex_base_url(vertex_location or self.get_default_vertex_location()) + return get_vertex_base_url( + vertex_location or self.get_default_vertex_location() + ) @staticmethod def create_vertex_url( diff --git a/litellm/llms/vertex_ai/videos/__init__.py b/litellm/llms/vertex_ai/videos/__init__.py index 1dcdbdf4de..7e00770787 100644 --- a/litellm/llms/vertex_ai/videos/__init__.py +++ b/litellm/llms/vertex_ai/videos/__init__.py @@ -7,4 +7,3 @@ This module provides support for Vertex AI's Veo video generation API. from .transformation import VertexAIVideoConfig __all__ = ["VertexAIVideoConfig"] - diff --git a/litellm/llms/vertex_ai/videos/transformation.py b/litellm/llms/vertex_ai/videos/transformation.py index 60852c1bf0..e61f2f46ec 100644 --- a/litellm/llms/vertex_ai/videos/transformation.py +++ b/litellm/llms/vertex_ai/videos/transformation.py @@ -78,11 +78,11 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): def extract_model_from_operation_name(operation_name: str) -> Optional[str]: """ Extract the model name from a Vertex AI operation name. - + Args: operation_name: Operation name in format: projects/PROJECT/locations/LOCATION/publishers/google/models/MODEL/operations/OPERATION_ID - + Returns: Model name (e.g., "veo-2.0-generate-001") or None if extraction fails """ @@ -174,17 +174,23 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): ) -> dict: """ Validate environment and return headers for Vertex AI OCR. - + Vertex AI uses Bearer token authentication with access token from credentials. """ # Extract Vertex AI parameters using safe helpers from VertexBase # Use safe_get_* methods that don't mutate litellm_params dict # Ensure litellm_params is a dict for type checking - params_dict: Dict[str, Any] = cast(Dict[str, Any], litellm_params) if litellm_params is not None else {} - - vertex_project = VertexBase.safe_get_vertex_ai_project(litellm_params=params_dict) - vertex_credentials = VertexBase.safe_get_vertex_ai_credentials(litellm_params=params_dict) - + params_dict: Dict[str, Any] = ( + cast(Dict[str, Any], litellm_params) if litellm_params is not None else {} + ) + + vertex_project = VertexBase.safe_get_vertex_ai_project( + litellm_params=params_dict + ) + vertex_credentials = VertexBase.safe_get_vertex_ai_credentials( + litellm_params=params_dict + ) + # Get access token from Vertex credentials access_token, project_id = self.get_access_token( credentials=vertex_credentials, @@ -353,24 +359,23 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): else: video_id = operation_name - video_obj = VideoObject( - id=video_id, - object="video", - status="processing", - model=model + id=video_id, object="video", status="processing", model=model ) usage_data = {} if request_data: parameters = request_data.get("parameters", {}) - duration = parameters.get("durationSeconds") or DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS + duration = ( + parameters.get("durationSeconds") + or DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS + ) if duration is not None: try: usage_data["duration_seconds"] = float(duration) except (ValueError, TypeError): pass - + video_obj.usage = usage_data return video_obj @@ -388,7 +393,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): """ operation_name = extract_original_video_id(video_id) model = self.extract_model_from_operation_name(operation_name) - + if not model: raise ValueError( f"Invalid operation name format: {operation_name}. " @@ -500,7 +505,9 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): Since we need to make an HTTP call here, we'll use the same fetchPredictOperation approach as status retrieval. """ - return self.transform_video_status_retrieve_request(video_id, api_base, litellm_params, headers) + return self.transform_video_status_retrieve_request( + video_id, api_base, litellm_params, headers + ) def transform_video_content_response( self, @@ -627,4 +634,3 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): message=error_message, headers=headers, ) - diff --git a/litellm/llms/volcengine/chat/transformation.py b/litellm/llms/volcengine/chat/transformation.py index 6df1cd3826..7395f9ce75 100644 --- a/litellm/llms/volcengine/chat/transformation.py +++ b/litellm/llms/volcengine/chat/transformation.py @@ -7,6 +7,7 @@ class VolcEngineChatConfig(OpenAILikeChatConfig): """ Reference: https://www.volcengine.com/docs/82379/1494384 """ + frequency_penalty: Optional[int] = None function_call: Optional[Union[str, dict]] = None functions: Optional[list] = None @@ -95,10 +96,13 @@ class VolcEngineChatConfig(OpenAILikeChatConfig): if ( thinking_value is not None and isinstance(thinking_value, dict) - and thinking_value.get("type", None) in ["enabled", "disabled", "auto"] # legal values, see docs + and thinking_value.get("type", None) + in ["enabled", "disabled", "auto"] # legal values, see docs ): # Add thinking parameter to extra_body for all legal cases - optional_params.setdefault("extra_body", {})["thinking"] = thinking_value + optional_params.setdefault("extra_body", {})[ + "thinking" + ] = thinking_value else: # Skip adding thinking parameter when it's not set or has invalid value pass diff --git a/litellm/llms/volcengine/embedding/transformation.py b/litellm/llms/volcengine/embedding/transformation.py index 20747b7672..cb497c9f15 100644 --- a/litellm/llms/volcengine/embedding/transformation.py +++ b/litellm/llms/volcengine/embedding/transformation.py @@ -59,7 +59,7 @@ class VolcEngineEmbeddingConfig(BaseEmbeddingConfig): ) -> str: """ Get the complete URL for volcengine embedding API calls. - + Args: api_base: Optional custom API base URL api_key: API key (not used for URL construction) @@ -67,7 +67,7 @@ class VolcEngineEmbeddingConfig(BaseEmbeddingConfig): optional_params: Optional parameters (not used for URL construction) litellm_params: LiteLLM parameters (not used for URL construction) stream: Stream parameter (not used for URL construction) - + Returns: Complete URL for the embedding API endpoint """ @@ -117,8 +117,6 @@ class VolcEngineEmbeddingConfig(BaseEmbeddingConfig): return optional_params - - def transform_embedding_request( self, model: str, @@ -175,7 +173,7 @@ class VolcEngineEmbeddingConfig(BaseEmbeddingConfig): # Add id if present if "id" in response_json: transformed_response["id"] = response_json["id"] - + # Create EmbeddingResponse from transformed data return EmbeddingResponse(**transformed_response) @@ -201,6 +199,7 @@ class VolcEngineEmbeddingConfig(BaseEmbeddingConfig): ) -> BaseLLMException: """Get error class for Volcengine errors""" from ..common_utils import VolcEngineError + # Convert dict to httpx.Headers if needed if isinstance(headers, dict): headers = httpx.Headers(headers) diff --git a/litellm/llms/volcengine/responses/transformation.py b/litellm/llms/volcengine/responses/transformation.py index f9ed93f680..f6dda4dd25 100644 --- a/litellm/llms/volcengine/responses/transformation.py +++ b/litellm/llms/volcengine/responses/transformation.py @@ -92,7 +92,9 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] ) -> VolcEngineError: typed_headers: httpx.Headers = ( - headers if isinstance(headers, httpx.Headers) else httpx.Headers(headers or {}) + headers + if isinstance(headers, httpx.Headers) + else httpx.Headers(headers or {}) ) return VolcEngineError( status_code=status_code, @@ -193,7 +195,9 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): allowed = set(self._SUPPORTED_OPTIONAL_PARAMS) sanitized_optional = { - k: v for k, v in response_api_optional_request_params.items() if k in allowed + k: v + for k, v in response_api_optional_request_params.items() + if k in allowed } # Ensure metadata never reaches provider sanitized_optional.pop("metadata", None) @@ -203,7 +207,9 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): # leaking unsupported params to the provider. if isinstance(sanitized_optional.get("extra_body"), dict): filtered_body = { - k: v for k, v in sanitized_optional["extra_body"].items() if k in allowed + k: v + for k, v in sanitized_optional["extra_body"].items() + if k in allowed } if filtered_body: sanitized_optional["extra_body"] = filtered_body @@ -438,9 +444,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): return False @staticmethod - def _fill_missing_fields( - chunk: Any, event_model: Any - ) -> Dict[str, Any]: + def _fill_missing_fields(chunk: Any, event_model: Any) -> Dict[str, Any]: """ Heuristically fill missing required fields with safe defaults based on the event model's field annotations. This keeps parsing tolerant of providers that @@ -460,7 +464,10 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): continue # Explicit default or factory - if field.default is not pyd_fields.PydanticUndefined and field.default is not None: + if ( + field.default is not pyd_fields.PydanticUndefined + and field.default is not None + ): patched[name] = field.default continue if ( diff --git a/litellm/llms/voyage/rerank/transformation.py b/litellm/llms/voyage/rerank/transformation.py index a6fe38c0cd..521dae980d 100644 --- a/litellm/llms/voyage/rerank/transformation.py +++ b/litellm/llms/voyage/rerank/transformation.py @@ -23,7 +23,6 @@ from ..embedding.transformation import VoyageError class VoyageRerankConfig(BaseRerankConfig): - def get_supported_cohere_rerank_params(self, model: str) -> list: return ["query", "documents", "top_n", "return_documents"] @@ -137,12 +136,17 @@ class VoyageRerankConfig(BaseRerankConfig): optional_params: Optional[dict] = None, ) -> Dict: if api_key is None: - api_key = get_secret_str("VOYAGE_API_KEY") or get_secret_str("VOYAGE_AI_API_KEY") + api_key = get_secret_str("VOYAGE_API_KEY") or get_secret_str( + "VOYAGE_AI_API_KEY" + ) if api_key is None: raise ValueError( "Voyage AI API key is required. Set via `api_key` parameter or `VOYAGE_API_KEY` env var." ) - return {"Authorization": f"Bearer {api_key}", "content-type": "application/json"} + return { + "Authorization": f"Bearer {api_key}", + "content-type": "application/json", + } def calculate_rerank_cost( self, @@ -166,4 +170,6 @@ class VoyageRerankConfig(BaseRerankConfig): def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] ): - return VoyageError(message=error_message, status_code=status_code, headers=headers) + return VoyageError( + message=error_message, status_code=status_code, headers=headers + ) diff --git a/litellm/llms/watsonx/rerank/transformation.py b/litellm/llms/watsonx/rerank/transformation.py index 7b4c2a07c3..4f8e196f25 100644 --- a/litellm/llms/watsonx/rerank/transformation.py +++ b/litellm/llms/watsonx/rerank/transformation.py @@ -42,7 +42,9 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig): params = optional_params or {} - complete_url = self._add_api_version_to_url(url=url, api_version=(params.get("api_version", None))) + complete_url = self._add_api_version_to_url( + url=url, api_version=(params.get("api_version", None)) + ) return complete_url def get_supported_cohere_rerank_params(self, model: str) -> list: @@ -76,7 +78,8 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig): ) zen_api_key = cast( Optional[str], - optional_params.pop("zen_api_key", None) or get_secret_str("WATSONX_ZENAPIKEY"), + optional_params.pop("zen_api_key", None) + or get_secret_str("WATSONX_ZENAPIKEY"), ) if token: headers["Authorization"] = f"Bearer {token}" @@ -115,11 +118,17 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig): {"text": el} if isinstance(el, str) else el for el in v ] elif k == "top_n" and v is not None: - optional_rerank_params.setdefault("parameters", {}).setdefault("return_options", {})["top_n"] = v + optional_rerank_params.setdefault("parameters", {}).setdefault( + "return_options", {} + )["top_n"] = v elif k == "return_documents" and v is not None and isinstance(v, bool): - optional_rerank_params.setdefault("parameters", {}).setdefault("return_options", {})["inputs"] = v + optional_rerank_params.setdefault("parameters", {}).setdefault( + "return_options", {} + )["inputs"] = v elif k == "max_tokens_per_doc" and v is not None: - optional_rerank_params.setdefault("parameters", {})["truncate_input_tokens"] = v + optional_rerank_params.setdefault("parameters", {})[ + "truncate_input_tokens" + ] = v # IBM watsonx.ai require one of below parameters elif k == "project_id" and v is not None: @@ -189,7 +198,11 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig): transformed_results.append(transformed_result) - response_id = raw_response_json.get("id") or raw_response_json.get("model_id") or str(uuid.uuid4()) + response_id = ( + raw_response_json.get("id") + or raw_response_json.get("model_id") + or str(uuid.uuid4()) + ) # Extract usage information _tokens = RerankTokens( diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index aa2dee354c..bfa55105a6 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -62,14 +62,13 @@ class XAIChatConfig(OpenAIGPTConfig): ######################################################### if self._supports_stop_reason(model): base_openai_params.append("stop") - ######################################################### # frequency penalty check ######################################################### if self._supports_frequency_penalty(model): base_openai_params.append("frequency_penalty") - + ######################################################### # reasoning check ######################################################### @@ -82,7 +81,7 @@ class XAIChatConfig(OpenAIGPTConfig): verbose_logger.debug(f"Error checking if model supports reasoning: {e}") return base_openai_params - + def _supports_stop_reason(self, model: str) -> bool: if "grok-3-mini" in model: return False @@ -91,7 +90,7 @@ class XAIChatConfig(OpenAIGPTConfig): elif "grok-code-fast" in model: return False return True - + def _supports_frequency_penalty(self, model: str) -> bool: """ From manual testing grok-4 does not support `frequency_penalty` @@ -162,13 +161,15 @@ class XAIChatConfig(OpenAIGPTConfig): def _fix_choice_finish_reason_for_tool_calls(choice: Choices) -> None: """ Helper to fix finish_reason for tool calls when XAI API returns empty string. - + XAI API returns empty string for finish_reason when using tools, so we need to set it to "tool_calls" when tool_calls are present. """ - if (choice.finish_reason == "" and - choice.message.tool_calls and - len(choice.message.tool_calls) > 0): + if ( + choice.finish_reason == "" + and choice.message.tool_calls + and len(choice.message.tool_calls) > 0 + ): choice.finish_reason = "tool_calls" def transform_response( @@ -187,13 +188,13 @@ class XAIChatConfig(OpenAIGPTConfig): ) -> ModelResponse: """ Transform the response from the XAI API. - + XAI API returns empty string for finish_reason when using tools, so we need to fix this after the standard OpenAI transformation. - + Also handles X.AI web search usage tracking by extracting num_sources_used. """ - + # First, let the parent class handle the standard transformation response = super().transform_response( model=model, @@ -237,12 +238,12 @@ class XAIChatConfig(OpenAIGPTConfig): response_usage = raw_response_json.get("usage", {}) if isinstance(response_usage, dict) and "num_sources_used" in response_usage: num_sources_used = response_usage.get("num_sources_used") - + # Map num_sources_used to web_search_requests for cost detection if num_sources_used is not None and num_sources_used > 0: if usage.prompt_tokens_details is None: usage.prompt_tokens_details = PromptTokensDetailsWrapper() - + usage.prompt_tokens_details.web_search_requests = int(num_sources_used) setattr(usage, "num_sources_used", int(num_sources_used)) verbose_logger.debug(f"X.AI web search sources used: {num_sources_used}") @@ -252,10 +253,10 @@ class XAIChatCompletionStreamingHandler(OpenAIChatCompletionStreamingHandler): def chunk_parser(self, chunk: dict) -> ModelResponseStream: """ Handle xAI-specific streaming behavior. - + xAI Grok sends a final chunk with empty choices array but with usage data when stream_options={"include_usage": True} is set. - + Example from xAI API: {"id":"...","object":"chat.completion.chunk","created":...,"model":"grok-4-1-fast-non-reasoning", "choices":[],"usage":{"prompt_tokens":171,"completion_tokens":2,"total_tokens":173,...}} @@ -266,5 +267,5 @@ class XAIChatCompletionStreamingHandler(OpenAIChatCompletionStreamingHandler): # xAI sends usage in a chunk with empty choices array # Add a dummy choice with empty delta to ensure proper processing chunk["choices"] = [{"index": 0, "delta": {}, "finish_reason": None}] - + return super().chunk_parser(chunk) diff --git a/litellm/llms/xai/cost_calculator.py b/litellm/llms/xai/cost_calculator.py index 91ad87e0b8..0cfcfe9841 100644 --- a/litellm/llms/xai/cost_calculator.py +++ b/litellm/llms/xai/cost_calculator.py @@ -30,22 +30,22 @@ def cost_per_token(model: str, usage: Usage) -> Tuple[float, float]: completion_tokens = int(getattr(usage, "completion_tokens", 0) or 0) reasoning_tokens = 0 if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details: - reasoning_tokens = int(getattr(usage.completion_tokens_details, "reasoning_tokens", 0) or 0) + reasoning_tokens = int( + getattr(usage.completion_tokens_details, "reasoning_tokens", 0) or 0 + ) total_completion_tokens = completion_tokens + reasoning_tokens - + modified_usage = Usage( prompt_tokens=usage.prompt_tokens, completion_tokens=total_completion_tokens, total_tokens=usage.total_tokens, prompt_tokens_details=usage.prompt_tokens_details, - completion_tokens_details=None + completion_tokens_details=None, ) - + prompt_cost, completion_cost = generic_cost_per_token( - model=model, - usage=modified_usage, - custom_llm_provider="xai" + model=model, usage=modified_usage, custom_llm_provider="xai" ) return prompt_cost, completion_cost @@ -54,30 +54,30 @@ def cost_per_token(model: str, usage: Usage) -> Tuple[float, float]: def cost_per_web_search_request(usage: "Usage", model_info: "ModelInfo") -> float: """ Calculate the cost of web search requests for X.AI models. - + X.AI Live Search costs $25 per 1,000 sources used. Each source costs $0.025. - + The number of sources is stored in prompt_tokens_details.web_search_requests by the transformation layer to be compatible with the existing detection system. """ # Cost per source used: $25 per 1,000 sources = $0.025 per source cost_per_source = 25.0 / 1000.0 # $0.025 - + num_sources_used = 0 - + if ( - hasattr(usage, "prompt_tokens_details") + hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details is not None and hasattr(usage.prompt_tokens_details, "web_search_requests") and usage.prompt_tokens_details.web_search_requests is not None ): num_sources_used = int(usage.prompt_tokens_details.web_search_requests) - + # Fallback: try to get from num_sources_used if set directly elif hasattr(usage, "num_sources_used") and usage.num_sources_used is not None: num_sources_used = int(usage.num_sources_used) total_cost = cost_per_source * num_sources_used - + return total_cost diff --git a/litellm/llms/xai/realtime/handler.py b/litellm/llms/xai/realtime/handler.py index c79477ba1d..805cce5a26 100644 --- a/litellm/llms/xai/realtime/handler.py +++ b/litellm/llms/xai/realtime/handler.py @@ -15,19 +15,19 @@ from ...openai.realtime.handler import OpenAIRealtime class XAIRealtime(OpenAIRealtime): """ Handler for xAI Grok Voice Agent API. - + xAI's Realtime API uses the same WebSocket protocol as OpenAI but with: - Different endpoint: wss://api.x.ai/v1/realtime (via _get_default_api_base) - No OpenAI-Beta header required (via _get_additional_headers) - Model: grok-4-1-fast-non-reasoning - + All WebSocket logic is inherited from OpenAIRealtime. """ - + def _get_default_api_base(self) -> str: """xAI uses a different API base URL.""" return XAI_API_BASE - + def _get_additional_headers(self, api_key: str) -> dict: """ xAI does NOT require the OpenAI-Beta header. diff --git a/litellm/llms/xai/responses/transformation.py b/litellm/llms/xai/responses/transformation.py index 3c69b7d08b..23aee3a120 100644 --- a/litellm/llms/xai/responses/transformation.py +++ b/litellm/llms/xai/responses/transformation.py @@ -21,13 +21,13 @@ else: class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): """ Configuration for XAI's Responses API. - + Inherits from OpenAIResponsesAPIConfig since XAI's Responses API is largely compatible with OpenAI's, with a few differences: - Does not support the 'instructions' parameter - Requires code_interpreter tools to have 'container' field removed - Recommends store=false when sending images - + Reference: https://docs.x.ai/docs/api-reference#create-new-response """ @@ -38,60 +38,64 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): def get_supported_openai_params(self, model: str) -> list: """ Get supported parameters for XAI Responses API. - + XAI supports most OpenAI Responses API params except 'instructions'. """ supported_params = super().get_supported_openai_params(model) - + # Remove 'instructions' as it's not supported by XAI if "instructions" in supported_params: supported_params.remove("instructions") - + return supported_params - def _transform_web_search_tool(self, tool: Dict[str, Any]) -> Union[XAIWebSearchTool, Dict[str, Any]]: + def _transform_web_search_tool( + self, tool: Dict[str, Any] + ) -> Union[XAIWebSearchTool, Dict[str, Any]]: """ Transform web_search tool to XAI format. - + XAI supports web_search with specific filters: - allowed_domains (max 5) - excluded_domains (max 5) - enable_image_understanding - + XAI does NOT support search_context_size (OpenAI-specific). """ xai_tool: Dict[str, Any] = {"type": "web_search"} - + # Remove search_context_size if present (not supported by XAI) if "search_context_size" in tool: verbose_logger.info( "XAI does not support 'search_context_size' parameter. Removing it from web_search tool." ) - + # Handle filters (XAI-specific structure) filters = {} if "allowed_domains" in tool: allowed_domains = tool["allowed_domains"] filters["allowed_domains"] = allowed_domains - + if "excluded_domains" in tool: excluded_domains = tool["excluded_domains"] filters["excluded_domains"] = excluded_domains - + # Add filters if any were specified if filters: xai_tool["filters"] = filters - + # Handle enable_image_understanding (top-level in XAI format) if "enable_image_understanding" in tool: xai_tool["enable_image_understanding"] = tool["enable_image_understanding"] - + return xai_tool - - def _transform_x_search_tool(self, tool: Dict[str, Any]) -> Union[XAIXSearchTool, Dict[str, Any]]: + + def _transform_x_search_tool( + self, tool: Dict[str, Any] + ) -> Union[XAIXSearchTool, Dict[str, Any]]: """ Transform x_search tool to XAI format. - + XAI supports x_search with specific parameters: - allowed_x_handles (max 10) - excluded_x_handles (max 10) @@ -101,31 +105,31 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): - enable_video_understanding """ xai_tool: Dict[str, Any] = {"type": "x_search"} - + # Handle allowed_x_handles if "allowed_x_handles" in tool: allowed_handles = tool["allowed_x_handles"] xai_tool["allowed_x_handles"] = allowed_handles - + # Handle excluded_x_handles if "excluded_x_handles" in tool: excluded_handles = tool["excluded_x_handles"] xai_tool["excluded_x_handles"] = excluded_handles - + # Handle date range if "from_date" in tool: xai_tool["from_date"] = tool["from_date"] - + if "to_date" in tool: xai_tool["to_date"] = tool["to_date"] - + # Handle media understanding flags if "enable_image_understanding" in tool: xai_tool["enable_image_understanding"] = tool["enable_image_understanding"] - + if "enable_video_understanding" in tool: xai_tool["enable_video_understanding"] = tool["enable_video_understanding"] - + return xai_tool def map_openai_params( @@ -136,7 +140,7 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): ) -> Dict: """ Map parameters for XAI Responses API. - + Handles XAI-specific transformations: 1. Drops 'instructions' parameter (not supported) 2. Transforms code_interpreter tools to remove 'container' field @@ -145,61 +149,61 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): 5. Sets store=false when images are detected (recommended by XAI) """ params = dict(response_api_optional_params) - + # Drop instructions parameter (not supported by XAI) if "instructions" in params: verbose_logger.debug( "XAI Responses API does not support 'instructions' parameter. Dropping it." ) params.pop("instructions") - + if "metadata" in params: verbose_logger.debug( "XAI Responses API does not support 'metadata' parameter. Dropping it." ) params.pop("metadata") - + # Transform tools if "tools" in params and params["tools"]: tools_list = params["tools"] # Ensure tools is a list for iteration if not isinstance(tools_list, list): tools_list = [tools_list] - + transformed_tools: List[Any] = [] for tool in tools_list: if isinstance(tool, dict): tool_type = tool.get("type") - + if tool_type == "code_interpreter": # XAI supports code_interpreter but doesn't use the container field verbose_logger.debug( "XAI: Transforming code_interpreter tool, removing container field" ) transformed_tools.append({"type": "code_interpreter"}) - + elif tool_type == "web_search": # Transform web_search to XAI format verbose_logger.debug( "XAI: Transforming web_search tool to XAI format" ) transformed_tools.append(self._transform_web_search_tool(tool)) - + elif tool_type == "x_search": # Transform x_search to XAI format verbose_logger.debug( "XAI: Transforming x_search tool to XAI format" ) transformed_tools.append(self._transform_x_search_tool(tool)) - + else: # Keep other tools as-is transformed_tools.append(tool) else: transformed_tools.append(tool) - + params["tools"] = transformed_tools - + return params def validate_environment( @@ -207,21 +211,19 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): ) -> dict: """ Validate environment and set up headers for XAI API. - + Uses XAI_API_KEY from environment or litellm_params. """ litellm_params = litellm_params or GenericLiteLLMParams() api_key = ( - litellm_params.api_key - or litellm.api_key - or get_secret_str("XAI_API_KEY") + litellm_params.api_key or litellm.api_key or get_secret_str("XAI_API_KEY") ) - + if not api_key: raise ValueError( "XAI API key is required. Set XAI_API_KEY environment variable or pass api_key parameter." ) - + headers.update( { "Authorization": f"Bearer {api_key}", @@ -236,7 +238,7 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): ) -> str: """ Get the complete URL for XAI Responses API endpoint. - + Returns: str: The full URL for the XAI /responses endpoint """ @@ -246,13 +248,12 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): or get_secret_str("XAI_API_BASE") or XAI_API_BASE ) - + # Remove trailing slashes api_base = api_base.rstrip("/") - + return f"{api_base}/responses" def supports_native_websocket(self) -> bool: """XAI does not support native WebSocket for Responses API""" return False - diff --git a/litellm/llms/zai/chat/transformation.py b/litellm/llms/zai/chat/transformation.py index fb1d67df35..c932dcd2e0 100644 --- a/litellm/llms/zai/chat/transformation.py +++ b/litellm/llms/zai/chat/transformation.py @@ -48,7 +48,9 @@ class ZAIChatConfig(OpenAIGPTConfig): import litellm try: - if litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider): + if litellm.supports_reasoning( + model=model, custom_llm_provider=self.custom_llm_provider + ): base_params.append("thinking") except Exception: pass diff --git a/litellm/main.py b/litellm/main.py index 4e4ce976ac..f2ce894ba3 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -2244,7 +2244,9 @@ def completion( # type: ignore # noqa: PLR0915 client=client, ) elif custom_llm_provider == "bedrock_mantle": - api_base = api_base or litellm.api_base or get_secret("BEDROCK_MANTLE_API_BASE") + api_base = ( + api_base or litellm.api_base or get_secret("BEDROCK_MANTLE_API_BASE") + ) api_key = api_key or litellm.api_key or get_secret("BEDROCK_MANTLE_API_KEY") headers = headers or litellm.headers config = litellm.BedrockMantleChatConfig.get_config() @@ -2272,14 +2274,16 @@ def completion( # type: ignore # noqa: PLR0915 elif custom_llm_provider == "a2a": # A2A (Agent-to-Agent) Protocol # Resolve agent configuration from registry if model format is "a2a/" - api_base, api_key, headers = ( - litellm.A2AConfig.resolve_agent_config_from_registry( - model=model, - api_base=api_base, - api_key=api_key, - headers=headers, - optional_params=optional_params, - ) + ( + api_base, + api_key, + headers, + ) = litellm.A2AConfig.resolve_agent_config_from_registry( + model=model, + api_base=api_base, + api_key=api_key, + headers=headers, + optional_params=optional_params, ) # Fall back to environment variables and defaults @@ -3751,9 +3755,9 @@ def completion( # type: ignore # noqa: PLR0915 "aws_region_name" not in optional_params or optional_params["aws_region_name"] is None ): - optional_params["aws_region_name"] = ( - aws_bedrock_client.meta.region_name - ) + optional_params[ + "aws_region_name" + ] = aws_bedrock_client.meta.region_name bedrock_route = BedrockModelInfo.get_bedrock_route(model) if bedrock_route == "converse": @@ -5193,7 +5197,9 @@ def embedding( # noqa: PLR0915 ) try: - model_info = get_model_info(model=model, custom_llm_provider="vertex_ai") + model_info = get_model_info( + model=model, custom_llm_provider="vertex_ai" + ) uses_embed_content = model_info.get("uses_embed_content", False) except Exception: uses_embed_content = False @@ -6154,9 +6160,9 @@ def adapter_completion( new_kwargs = translation_obj.translate_completion_input_params(kwargs=kwargs) response: Union[ModelResponse, CustomStreamWrapper] = completion(**new_kwargs) # type: ignore - translated_response: Optional[Union[BaseModel, AdapterCompletionStreamWrapper]] = ( - None - ) + translated_response: Optional[ + Union[BaseModel, AdapterCompletionStreamWrapper] + ] = None if isinstance(response, ModelResponse): translated_response = translation_obj.translate_completion_output_params( response=response @@ -6336,7 +6342,9 @@ async def atranscription(*args, **kwargs) -> TranscriptionResponse: if existing_duration is None: calculated_duration = calculate_request_duration(file) if calculated_duration is not None: - response._hidden_params["audio_transcription_duration"] = calculated_duration + response._hidden_params[ + "audio_transcription_duration" + ] = calculated_duration return response except Exception as e: @@ -6559,7 +6567,9 @@ def transcription( if existing_duration is None: calculated_duration = calculate_request_duration(file) if calculated_duration is not None: - response._hidden_params["audio_transcription_duration"] = calculated_duration + response._hidden_params[ + "audio_transcription_duration" + ] = calculated_duration if response is None: raise ValueError("Unmapped provider passed in. Unable to get the response.") @@ -6863,9 +6873,9 @@ def speech( # noqa: PLR0915 ElevenLabsTextToSpeechConfig.ELEVENLABS_QUERY_PARAMS_KEY ] = query_params - litellm_params_dict[ElevenLabsTextToSpeechConfig.ELEVENLABS_VOICE_ID_KEY] = ( - voice_id - ) + litellm_params_dict[ + ElevenLabsTextToSpeechConfig.ELEVENLABS_VOICE_ID_KEY + ] = voice_id if api_base is not None: litellm_params_dict["api_base"] = api_base @@ -7444,9 +7454,9 @@ def stream_chunk_builder( # noqa: PLR0915 ] if len(content_chunks) > 0: - response["choices"][0]["message"]["content"] = ( - processor.get_combined_content(content_chunks) - ) + response["choices"][0]["message"][ + "content" + ] = processor.get_combined_content(content_chunks) thinking_blocks = [ chunk @@ -7457,9 +7467,9 @@ def stream_chunk_builder( # noqa: PLR0915 ] if len(thinking_blocks) > 0: - response["choices"][0]["message"]["thinking_blocks"] = ( - processor.get_combined_thinking_content(thinking_blocks) - ) + response["choices"][0]["message"][ + "thinking_blocks" + ] = processor.get_combined_thinking_content(thinking_blocks) reasoning_chunks = [ chunk @@ -7470,9 +7480,9 @@ def stream_chunk_builder( # noqa: PLR0915 ] if len(reasoning_chunks) > 0: - response["choices"][0]["message"]["reasoning_content"] = ( - processor.get_combined_reasoning_content(reasoning_chunks) - ) + response["choices"][0]["message"][ + "reasoning_content" + ] = processor.get_combined_reasoning_content(reasoning_chunks) annotation_chunks = [ chunk @@ -7626,12 +7636,15 @@ async def acount_tokens( from litellm.utils import ProviderConfigManager # Determine provider from model string - resolved_model, custom_llm_provider, dynamic_api_key, dynamic_api_base = ( - get_llm_provider( - model=model, - api_base=api_base, - api_key=api_key, - ) + ( + resolved_model, + custom_llm_provider, + dynamic_api_key, + dynamic_api_base, + ) = get_llm_provider( + model=model, + api_base=api_base, + api_key=api_key, ) # Use dynamic key/base if not explicitly provided diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d0c250fb0d..9b1d81fee4 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -2565,32 +2565,6 @@ "supports_parallel_function_calling": true, "supports_tool_choice": true }, - "azure/gpt-35-turbo-0301": { - "deprecation_date": "2025-02-13", - "input_cost_per_token": 2e-07, - "litellm_provider": "azure", - "max_input_tokens": 4097, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 2e-06, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "azure/gpt-35-turbo-0613": { - "deprecation_date": "2025-02-13", - "input_cost_per_token": 1.5e-06, - "litellm_provider": "azure", - "max_input_tokens": 4097, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 2e-06, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, "azure/gpt-35-turbo-1106": { "deprecation_date": "2025-03-31", "input_cost_per_token": 1e-06, @@ -8185,72 +8159,6 @@ "supports_reasoning": true, "supports_tool_choice": true }, - "chat-bison": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-chat-models", - "max_input_tokens": 8192, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_tool_choice": true - }, - "chat-bison-32k": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-chat-models", - "max_input_tokens": 32000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_tool_choice": true - }, - "chat-bison-32k@002": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-chat-models", - "max_input_tokens": 32000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_tool_choice": true - }, - "chat-bison@001": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-chat-models", - "max_input_tokens": 8192, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_tool_choice": true - }, - "chat-bison@002": { - "deprecation_date": "2025-04-09", - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-chat-models", - "max_input_tokens": 8192, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_tool_choice": true - }, "chatdolphin": { "input_cost_per_token": 5e-07, "litellm_provider": "nlp_cloud", @@ -8288,60 +8196,6 @@ "/v1/audio/transcriptions" ] }, - "claude-3-5-haiku-20241022": { - "cache_creation_input_token_cost": 1e-06, - "cache_creation_input_token_cost_above_1hr": 6e-06, - "cache_read_input_token_cost": 8e-08, - "deprecation_date": "2025-10-01", - "input_cost_per_token": 8e-07, - "litellm_provider": "anthropic", - "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 4e-06, - "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 - }, - "supports_assistant_prefill": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tool_use_system_prompt_tokens": 264 - }, - "claude-3-5-haiku-latest": { - "cache_creation_input_token_cost": 1.25e-06, - "cache_creation_input_token_cost_above_1hr": 6e-06, - "cache_read_input_token_cost": 1e-07, - "deprecation_date": "2025-10-01", - "input_cost_per_token": 1e-06, - "litellm_provider": "anthropic", - "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 5e-06, - "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 - }, - "supports_assistant_prefill": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tool_use_system_prompt_tokens": 264 - }, "claude-haiku-4-5-20251001": { "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, @@ -8384,83 +8238,6 @@ "supports_tool_choice": true, "supports_vision": true }, - "claude-3-5-sonnet-20240620": { - "cache_creation_input_token_cost": 3.75e-06, - "cache_creation_input_token_cost_above_1hr": 6e-06, - "cache_read_input_token_cost": 3e-07, - "deprecation_date": "2025-06-01", - "input_cost_per_token": 3e-06, - "litellm_provider": "anthropic", - "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.5e-05, - "supports_assistant_prefill": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 - }, - "claude-3-5-sonnet-20241022": { - "cache_creation_input_token_cost": 3.75e-06, - "cache_creation_input_token_cost_above_1hr": 6e-06, - "cache_read_input_token_cost": 3e-07, - "deprecation_date": "2025-10-01", - "input_cost_per_token": 3e-06, - "litellm_provider": "anthropic", - "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.5e-05, - "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 - }, - "supports_assistant_prefill": true, - "supports_computer_use": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tool_use_system_prompt_tokens": 159 - }, - "claude-3-5-sonnet-latest": { - "cache_creation_input_token_cost": 3.75e-06, - "cache_creation_input_token_cost_above_1hr": 6e-06, - "cache_read_input_token_cost": 3e-07, - "deprecation_date": "2025-06-01", - "input_cost_per_token": 3e-06, - "litellm_provider": "anthropic", - "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.5e-05, - "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 - }, - "supports_assistant_prefill": true, - "supports_computer_use": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tool_use_system_prompt_tokens": 159 - }, "claude-3-7-sonnet-20250219": { "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, @@ -8490,34 +8267,6 @@ "supports_web_search": true, "tool_use_system_prompt_tokens": 159 }, - "claude-3-7-sonnet-latest": { - "cache_creation_input_token_cost": 3.75e-06, - "cache_creation_input_token_cost_above_1hr": 6e-06, - "cache_read_input_token_cost": 3e-07, - "deprecation_date": "2025-06-01", - "input_cost_per_token": 3e-06, - "litellm_provider": "anthropic", - "max_input_tokens": 200000, - "max_output_tokens": 64000, - "max_tokens": 64000, - "mode": "chat", - "output_cost_per_token": 1.5e-05, - "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 - }, - "supports_assistant_prefill": true, - "supports_computer_use": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 - }, "claude-3-haiku-20240307": { "cache_creation_input_token_cost": 3e-07, "cache_creation_input_token_cost_above_1hr": 6e-06, @@ -8557,26 +8306,6 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 395 }, - "claude-3-opus-latest": { - "cache_creation_input_token_cost": 1.875e-05, - "cache_creation_input_token_cost_above_1hr": 6e-06, - "cache_read_input_token_cost": 1.5e-06, - "deprecation_date": "2025-03-01", - "input_cost_per_token": 1.5e-05, - "litellm_provider": "anthropic", - "max_input_tokens": 200000, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 7.5e-05, - "supports_assistant_prefill": true, - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 395 - }, "claude-4-opus-20250514": { "cache_creation_input_token_cost": 1.875e-05, "cache_read_input_token_cost": 1.5e-06, @@ -9025,185 +8754,6 @@ "mode": "chat", "output_cost_per_token": 1.923e-06 }, - "code-bison": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-text-models", - "max_input_tokens": 6144, - "max_output_tokens": 1024, - "max_tokens": 1024, - "mode": "chat", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_tool_choice": true - }, - "code-bison-32k@002": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-text-models", - "max_input_tokens": 6144, - "max_output_tokens": 1024, - "max_tokens": 1024, - "mode": "completion", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "code-bison32k": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-text-models", - "max_input_tokens": 6144, - "max_output_tokens": 1024, - "max_tokens": 1024, - "mode": "completion", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "code-bison@001": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-text-models", - "max_input_tokens": 6144, - "max_output_tokens": 1024, - "max_tokens": 1024, - "mode": "completion", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "code-bison@002": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-text-models", - "max_input_tokens": 6144, - "max_output_tokens": 1024, - "max_tokens": 1024, - "mode": "completion", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "code-gecko": { - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-text-models", - "max_input_tokens": 2048, - "max_output_tokens": 64, - "max_tokens": 64, - "mode": "completion", - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "code-gecko-latest": { - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-text-models", - "max_input_tokens": 2048, - "max_output_tokens": 64, - "max_tokens": 64, - "mode": "completion", - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "code-gecko@001": { - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-text-models", - "max_input_tokens": 2048, - "max_output_tokens": 64, - "max_tokens": 64, - "mode": "completion", - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "code-gecko@002": { - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-text-models", - "max_input_tokens": 2048, - "max_output_tokens": 64, - "max_tokens": 64, - "mode": "completion", - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "codechat-bison": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-chat-models", - "max_input_tokens": 6144, - "max_output_tokens": 1024, - "max_tokens": 1024, - "mode": "chat", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_tool_choice": true - }, - "codechat-bison-32k": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-chat-models", - "max_input_tokens": 32000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_tool_choice": true - }, - "codechat-bison-32k@002": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-chat-models", - "max_input_tokens": 32000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_tool_choice": true - }, - "codechat-bison@001": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-chat-models", - "max_input_tokens": 6144, - "max_output_tokens": 1024, - "max_tokens": 1024, - "mode": "chat", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_tool_choice": true - }, - "codechat-bison@002": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-chat-models", - "max_input_tokens": 6144, - "max_output_tokens": 1024, - "max_tokens": 1024, - "mode": "chat", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_tool_choice": true - }, - "codechat-bison@latest": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-chat-models", - "max_input_tokens": 6144, - "max_output_tokens": 1024, - "max_tokens": 1024, - "mode": "chat", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_tool_choice": true - }, "codestral/codestral-2405": { "input_cost_per_token": 0.0, "litellm_provider": "codestral", @@ -13718,475 +13268,6 @@ "supports_response_schema": true, "supports_tool_choice": true }, - "gemini-1.0-pro": { - "input_cost_per_character": 1.25e-07, - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "input_cost_per_video_per_second": 0.002, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 32760, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 3.75e-07, - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-1.0-pro-001": { - "deprecation_date": "2025-04-09", - "input_cost_per_character": 1.25e-07, - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "input_cost_per_video_per_second": 0.002, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 32760, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 3.75e-07, - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-1.0-pro-002": { - "deprecation_date": "2025-04-09", - "input_cost_per_character": 1.25e-07, - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "input_cost_per_video_per_second": 0.002, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 32760, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 3.75e-07, - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-1.0-pro-vision": { - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "litellm_provider": "vertex_ai-vision-models", - "max_images_per_prompt": 16, - "max_input_tokens": 16384, - "max_output_tokens": 2048, - "max_tokens": 2048, - "max_video_length": 2, - "max_videos_per_prompt": 1, - "mode": "chat", - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.0-pro-vision-001": { - "deprecation_date": "2025-04-09", - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "litellm_provider": "vertex_ai-vision-models", - "max_images_per_prompt": 16, - "max_input_tokens": 16384, - "max_output_tokens": 2048, - "max_tokens": 2048, - "max_video_length": 2, - "max_videos_per_prompt": 1, - "mode": "chat", - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.0-ultra": { - "input_cost_per_character": 1.25e-07, - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "input_cost_per_video_per_second": 0.002, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 8192, - "max_output_tokens": 2048, - "max_tokens": 2048, - "mode": "chat", - "output_cost_per_character": 3.75e-07, - "output_cost_per_token": 1.5e-06, - "source": "As of Jun, 2024. There is no available doc on vertex ai pricing gemini-1.0-ultra-001. Using gemini-1.0-pro pricing. Got max_tokens info here: https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-1.0-ultra-001": { - "input_cost_per_character": 1.25e-07, - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "input_cost_per_video_per_second": 0.002, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 8192, - "max_output_tokens": 2048, - "max_tokens": 2048, - "mode": "chat", - "output_cost_per_character": 3.75e-07, - "output_cost_per_token": 1.5e-06, - "source": "As of Jun, 2024. There is no available doc on vertex ai pricing gemini-1.0-ultra-001. Using gemini-1.0-pro pricing. Got max_tokens info here: https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-1.5-flash": { - "deprecation_date": "2025-09-29", - "input_cost_per_audio_per_second": 2e-06, - "input_cost_per_audio_per_second_above_128k_tokens": 4e-06, - "input_cost_per_character": 1.875e-08, - "input_cost_per_character_above_128k_tokens": 2.5e-07, - "input_cost_per_image": 2e-05, - "input_cost_per_image_above_128k_tokens": 4e-05, - "input_cost_per_token": 7.5e-08, - "input_cost_per_token_above_128k_tokens": 1e-06, - "input_cost_per_video_per_second": 2e-05, - "input_cost_per_video_per_second_above_128k_tokens": 4e-05, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_character": 7.5e-08, - "output_cost_per_character_above_128k_tokens": 1.5e-07, - "output_cost_per_token": 3e-07, - "output_cost_per_token_above_128k_tokens": 6e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.5-flash-001": { - "deprecation_date": "2025-05-24", - "input_cost_per_audio_per_second": 2e-06, - "input_cost_per_audio_per_second_above_128k_tokens": 4e-06, - "input_cost_per_character": 1.875e-08, - "input_cost_per_character_above_128k_tokens": 2.5e-07, - "input_cost_per_image": 2e-05, - "input_cost_per_image_above_128k_tokens": 4e-05, - "input_cost_per_token": 7.5e-08, - "input_cost_per_token_above_128k_tokens": 1e-06, - "input_cost_per_video_per_second": 2e-05, - "input_cost_per_video_per_second_above_128k_tokens": 4e-05, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_character": 7.5e-08, - "output_cost_per_character_above_128k_tokens": 1.5e-07, - "output_cost_per_token": 3e-07, - "output_cost_per_token_above_128k_tokens": 6e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.5-flash-002": { - "deprecation_date": "2025-09-24", - "input_cost_per_audio_per_second": 2e-06, - "input_cost_per_audio_per_second_above_128k_tokens": 4e-06, - "input_cost_per_character": 1.875e-08, - "input_cost_per_character_above_128k_tokens": 2.5e-07, - "input_cost_per_image": 2e-05, - "input_cost_per_image_above_128k_tokens": 4e-05, - "input_cost_per_token": 7.5e-08, - "input_cost_per_token_above_128k_tokens": 1e-06, - "input_cost_per_video_per_second": 2e-05, - "input_cost_per_video_per_second_above_128k_tokens": 4e-05, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_character": 7.5e-08, - "output_cost_per_character_above_128k_tokens": 1.5e-07, - "output_cost_per_token": 3e-07, - "output_cost_per_token_above_128k_tokens": 6e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-1.5-flash", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.5-flash-exp-0827": { - "deprecation_date": "2025-09-29", - "input_cost_per_audio_per_second": 2e-06, - "input_cost_per_audio_per_second_above_128k_tokens": 4e-06, - "input_cost_per_character": 1.875e-08, - "input_cost_per_character_above_128k_tokens": 2.5e-07, - "input_cost_per_image": 2e-05, - "input_cost_per_image_above_128k_tokens": 4e-05, - "input_cost_per_token": 4.688e-09, - "input_cost_per_token_above_128k_tokens": 1e-06, - "input_cost_per_video_per_second": 2e-05, - "input_cost_per_video_per_second_above_128k_tokens": 4e-05, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_character": 1.875e-08, - "output_cost_per_character_above_128k_tokens": 3.75e-08, - "output_cost_per_token": 4.6875e-09, - "output_cost_per_token_above_128k_tokens": 9.375e-09, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.5-flash-preview-0514": { - "deprecation_date": "2025-09-29", - "input_cost_per_audio_per_second": 2e-06, - "input_cost_per_audio_per_second_above_128k_tokens": 4e-06, - "input_cost_per_character": 1.875e-08, - "input_cost_per_character_above_128k_tokens": 2.5e-07, - "input_cost_per_image": 2e-05, - "input_cost_per_image_above_128k_tokens": 4e-05, - "input_cost_per_token": 7.5e-08, - "input_cost_per_token_above_128k_tokens": 1e-06, - "input_cost_per_video_per_second": 2e-05, - "input_cost_per_video_per_second_above_128k_tokens": 4e-05, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_character": 1.875e-08, - "output_cost_per_character_above_128k_tokens": 3.75e-08, - "output_cost_per_token": 4.6875e-09, - "output_cost_per_token_above_128k_tokens": 9.375e-09, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.5-pro": { - "deprecation_date": "2025-09-29", - "input_cost_per_audio_per_second": 3.125e-05, - "input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05, - "input_cost_per_character": 3.125e-07, - "input_cost_per_character_above_128k_tokens": 6.25e-07, - "input_cost_per_image": 0.00032875, - "input_cost_per_image_above_128k_tokens": 0.0006575, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_128k_tokens": 2.5e-06, - "input_cost_per_video_per_second": 0.00032875, - "input_cost_per_video_per_second_above_128k_tokens": 0.0006575, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 2097152, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 1.25e-06, - "output_cost_per_character_above_128k_tokens": 2.5e-06, - "output_cost_per_token": 5e-06, - "output_cost_per_token_above_128k_tokens": 1e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.5-pro-001": { - "deprecation_date": "2025-05-24", - "input_cost_per_audio_per_second": 3.125e-05, - "input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05, - "input_cost_per_character": 3.125e-07, - "input_cost_per_character_above_128k_tokens": 6.25e-07, - "input_cost_per_image": 0.00032875, - "input_cost_per_image_above_128k_tokens": 0.0006575, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_128k_tokens": 2.5e-06, - "input_cost_per_video_per_second": 0.00032875, - "input_cost_per_video_per_second_above_128k_tokens": 0.0006575, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 1.25e-06, - "output_cost_per_character_above_128k_tokens": 2.5e-06, - "output_cost_per_token": 5e-06, - "output_cost_per_token_above_128k_tokens": 1e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.5-pro-002": { - "deprecation_date": "2025-09-24", - "input_cost_per_audio_per_second": 3.125e-05, - "input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05, - "input_cost_per_character": 3.125e-07, - "input_cost_per_character_above_128k_tokens": 6.25e-07, - "input_cost_per_image": 0.00032875, - "input_cost_per_image_above_128k_tokens": 0.0006575, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_128k_tokens": 2.5e-06, - "input_cost_per_video_per_second": 0.00032875, - "input_cost_per_video_per_second_above_128k_tokens": 0.0006575, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 2097152, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 1.25e-06, - "output_cost_per_character_above_128k_tokens": 2.5e-06, - "output_cost_per_token": 5e-06, - "output_cost_per_token_above_128k_tokens": 1e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-1.5-pro", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.5-pro-preview-0215": { - "deprecation_date": "2025-09-29", - "input_cost_per_audio_per_second": 3.125e-05, - "input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05, - "input_cost_per_character": 3.125e-07, - "input_cost_per_character_above_128k_tokens": 6.25e-07, - "input_cost_per_image": 0.00032875, - "input_cost_per_image_above_128k_tokens": 0.0006575, - "input_cost_per_token": 7.8125e-08, - "input_cost_per_token_above_128k_tokens": 1.5625e-07, - "input_cost_per_video_per_second": 0.00032875, - "input_cost_per_video_per_second_above_128k_tokens": 0.0006575, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 1.25e-06, - "output_cost_per_character_above_128k_tokens": 2.5e-06, - "output_cost_per_token": 3.125e-07, - "output_cost_per_token_above_128k_tokens": 6.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, - "gemini-1.5-pro-preview-0409": { - "deprecation_date": "2025-09-29", - "input_cost_per_audio_per_second": 3.125e-05, - "input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05, - "input_cost_per_character": 3.125e-07, - "input_cost_per_character_above_128k_tokens": 6.25e-07, - "input_cost_per_image": 0.00032875, - "input_cost_per_image_above_128k_tokens": 0.0006575, - "input_cost_per_token": 7.8125e-08, - "input_cost_per_token_above_128k_tokens": 1.5625e-07, - "input_cost_per_video_per_second": 0.00032875, - "input_cost_per_video_per_second_above_128k_tokens": 0.0006575, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 1.25e-06, - "output_cost_per_character_above_128k_tokens": 2.5e-06, - "output_cost_per_token": 3.125e-07, - "output_cost_per_token_above_128k_tokens": 6.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_tool_choice": true - }, - "gemini-1.5-pro-preview-0514": { - "deprecation_date": "2025-09-29", - "input_cost_per_audio_per_second": 3.125e-05, - "input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05, - "input_cost_per_character": 3.125e-07, - "input_cost_per_character_above_128k_tokens": 6.25e-07, - "input_cost_per_image": 0.00032875, - "input_cost_per_image_above_128k_tokens": 0.0006575, - "input_cost_per_token": 7.8125e-08, - "input_cost_per_token_above_128k_tokens": 1.5625e-07, - "input_cost_per_video_per_second": 0.00032875, - "input_cost_per_video_per_second_above_128k_tokens": 0.0006575, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 1.25e-06, - "output_cost_per_character_above_128k_tokens": 2.5e-06, - "output_cost_per_token": 3.125e-07, - "output_cost_per_token_above_128k_tokens": 6.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, "gemini-2.0-flash": { "cache_read_input_token_cost": 2.5e-08, "deprecation_date": "2026-06-01", @@ -14265,54 +13346,6 @@ "supports_vision": true, "supports_web_search": true }, - "gemini-2.0-flash-exp": { - "cache_read_input_token_cost": 3.75e-08, - "input_cost_per_audio_per_second": 0, - "input_cost_per_audio_per_second_above_128k_tokens": 0, - "input_cost_per_character": 0, - "input_cost_per_character_above_128k_tokens": 0, - "input_cost_per_image": 0, - "input_cost_per_image_above_128k_tokens": 0, - "input_cost_per_token": 1.5e-07, - "input_cost_per_token_above_128k_tokens": 0, - "input_cost_per_video_per_second": 0, - "input_cost_per_video_per_second_above_128k_tokens": 0, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_character": 0, - "output_cost_per_character_above_128k_tokens": 0, - "output_cost_per_token": 6e-07, - "output_cost_per_token_above_128k_tokens": 0, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_audio_output": true, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true - }, "gemini-2.0-flash-lite": { "cache_read_input_token_cost": 1.875e-08, "deprecation_date": "2026-06-01", @@ -14385,235 +13418,6 @@ "supports_vision": true, "supports_web_search": true }, - "gemini-2.0-flash-live-preview-04-09": { - "cache_read_input_token_cost": 7.5e-08, - "input_cost_per_audio_token": 3e-06, - "input_cost_per_image": 3e-06, - "input_cost_per_token": 5e-07, - "input_cost_per_video_per_second": 3e-06, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_audio_token": 1.2e-05, - "output_cost_per_token": 2e-06, - "rpm": 10, - "source": "https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/gemini#gemini-2-0-flash-live-preview-04-09", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "audio" - ], - "supports_audio_output": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 250000 - }, - "gemini-2.0-flash-preview-image-generation": { - "deprecation_date": "2025-11-14", - "cache_read_input_token_cost": 2.5e-08, - "input_cost_per_audio_token": 7e-07, - "input_cost_per_token": 1e-07, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 4e-07, - "source": "https://ai.google.dev/pricing#2_0flash", - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_audio_input": true, - "supports_audio_output": true, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true - }, - "gemini-2.0-flash-thinking-exp": { - "deprecation_date": "2025-12-02", - "cache_read_input_token_cost": 0.0, - "input_cost_per_audio_per_second": 0, - "input_cost_per_audio_per_second_above_128k_tokens": 0, - "input_cost_per_character": 0, - "input_cost_per_character_above_128k_tokens": 0, - "input_cost_per_image": 0, - "input_cost_per_image_above_128k_tokens": 0, - "input_cost_per_token": 0, - "input_cost_per_token_above_128k_tokens": 0, - "input_cost_per_video_per_second": 0, - "input_cost_per_video_per_second_above_128k_tokens": 0, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_character": 0, - "output_cost_per_character_above_128k_tokens": 0, - "output_cost_per_token": 0, - "output_cost_per_token_above_128k_tokens": 0, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-2.0-flash", - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_audio_output": true, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true - }, - "gemini-2.0-flash-thinking-exp-01-21": { - "deprecation_date": "2025-12-02", - "cache_read_input_token_cost": 0.0, - "input_cost_per_audio_per_second": 0, - "input_cost_per_audio_per_second_above_128k_tokens": 0, - "input_cost_per_character": 0, - "input_cost_per_character_above_128k_tokens": 0, - "input_cost_per_image": 0, - "input_cost_per_image_above_128k_tokens": 0, - "input_cost_per_token": 0, - "input_cost_per_token_above_128k_tokens": 0, - "input_cost_per_video_per_second": 0, - "input_cost_per_video_per_second_above_128k_tokens": 0, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65536, - "max_pdf_size_mb": 30, - "max_tokens": 65536, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_character": 0, - "output_cost_per_character_above_128k_tokens": 0, - "output_cost_per_token": 0, - "output_cost_per_token_above_128k_tokens": 0, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-2.0-flash", - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_audio_output": false, - "supports_function_calling": false, - "supports_parallel_function_calling": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": false, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true - }, - "gemini-2.0-pro-exp-02-05": { - "cache_read_input_token_cost": 3.125e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 2097152, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_input": true, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_video_input": true, - "supports_vision": true, - "supports_web_search": true - }, "gemini-2.5-flash": { "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, @@ -14708,57 +13512,6 @@ "supports_web_search": false, "tpm": 8000000 }, - "gemini-2.5-flash-image-preview": { - "deprecation_date": "2026-01-15", - "cache_read_input_token_cost": 7.5e-08, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_image_token": 3e-07, - "input_cost_per_token": 3e-07, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "image_generation", - "output_cost_per_image": 0.039, - "output_cost_per_image_token": 3e-05, - "output_cost_per_reasoning_token": 3e-05, - "output_cost_per_token": 3e-05, - "rpm": 100000, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 8000000 - }, "gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -15181,96 +13934,6 @@ "supports_vision": true, "supports_web_search": true }, - "gemini-2.5-flash-preview-04-17": { - "cache_read_input_token_cost": 3.75e-08, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 1.5e-07, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_reasoning_token": 3.5e-06, - "output_cost_per_token": 6e-07, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true - }, - "gemini-2.5-flash-preview-05-20": { - "deprecation_date": "2025-11-18", - "cache_read_input_token_cost": 7.5e-08, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_reasoning_token": 2.5e-06, - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true - }, "gemini-2.5-pro": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, @@ -15703,193 +14366,6 @@ "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, "supports_service_tier": true }, - "gemini-2.5-pro-exp-03-25": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_input": true, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_video_input": true, - "supports_vision": true, - "supports_web_search": true - }, - "gemini-2.5-pro-preview-03-25": { - "deprecation_date": "2025-12-02", - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_audio_token": 1.25e-06, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true - }, - "gemini-2.5-pro-preview-05-06": { - "deprecation_date": "2025-12-02", - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_audio_token": 1.25e-06, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supported_regions": [ - "global" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true - }, - "gemini-2.5-pro-preview-06-05": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_audio_token": 1.25e-06, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true - }, "gemini-2.5-pro-preview-tts": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, @@ -16072,63 +14548,23 @@ "max_input_tokens": 1000000, "max_output_tokens": 8192, "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 0, + "mode": "embedding", "output_cost_per_token": 0, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/gemini-experimental", - "supports_function_calling": false, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-pro": { - "input_cost_per_character": 1.25e-07, - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "input_cost_per_video_per_second": 0.002, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 32760, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 3.75e-07, - "output_cost_per_token": 1.5e-06, + "output_vector_size": 3072, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true + "uses_embed_content": true }, - "gemini-pro-experimental": { - "input_cost_per_character": 0, - "input_cost_per_token": 0, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1000000, - "max_output_tokens": 8192, + "vertex_ai/gemini-embedding-2-preview": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "vertex_ai", + "max_input_tokens": 8192, "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 0, + "mode": "embedding", "output_cost_per_token": 0, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/gemini-experimental", - "supports_function_calling": false, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-pro-vision": { - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "litellm_provider": "vertex_ai-vision-models", - "max_images_per_prompt": 16, - "max_input_tokens": 16384, - "max_output_tokens": 2048, - "max_tokens": 2048, - "max_video_length": 2, - "max_videos_per_prompt": 1, - "mode": "chat", - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true, - "supports_vision": true + "output_vector_size": 3072, + "source": "https://ai.google.dev/gemini-api/docs/embeddings#multimodal", + "supports_multimodal": true, + "uses_embed_content": true }, "gemini/gemini-embedding-001": { "input_cost_per_token": 1.5e-07, @@ -16163,339 +14599,15 @@ "input_cost_per_token": 7.5e-08, "input_cost_per_token_above_128k_tokens": 1.5e-07, "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, + "max_input_tokens": 8192, "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 3e-07, - "output_cost_per_token_above_128k_tokens": 6e-07, - "rpm": 2000, - "source": "https://ai.google.dev/pricing", - "supports_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 4000000 - }, - "gemini/gemini-1.5-flash-001": { - "cache_creation_input_token_cost": 1e-06, - "cache_read_input_token_cost": 1.875e-08, - "deprecation_date": "2025-05-24", - "input_cost_per_token": 7.5e-08, - "input_cost_per_token_above_128k_tokens": 1.5e-07, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 3e-07, - "output_cost_per_token_above_128k_tokens": 6e-07, - "rpm": 2000, - "source": "https://ai.google.dev/pricing", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 4000000 - }, - "gemini/gemini-1.5-flash-002": { - "cache_creation_input_token_cost": 1e-06, - "cache_read_input_token_cost": 1.875e-08, - "deprecation_date": "2025-09-24", - "input_cost_per_token": 7.5e-08, - "input_cost_per_token_above_128k_tokens": 1.5e-07, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 3e-07, - "output_cost_per_token_above_128k_tokens": 6e-07, - "rpm": 2000, - "source": "https://ai.google.dev/pricing", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 4000000 - }, - "gemini/gemini-1.5-flash-8b": { - "deprecation_date": "2025-09-29", - "input_cost_per_token": 0, - "input_cost_per_token_above_128k_tokens": 0, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", + "mode": "embedding", "output_cost_per_token": 0, - "output_cost_per_token_above_128k_tokens": 0, - "rpm": 4000, - "source": "https://ai.google.dev/pricing", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 4000000 - }, - "gemini/gemini-1.5-flash-8b-exp-0827": { - "deprecation_date": "2025-09-29", - "input_cost_per_token": 0, - "input_cost_per_token_above_128k_tokens": 0, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 0, - "output_cost_per_token_above_128k_tokens": 0, - "rpm": 4000, - "source": "https://ai.google.dev/pricing", - "supports_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 4000000 - }, - "gemini/gemini-1.5-flash-8b-exp-0924": { - "deprecation_date": "2025-09-29", - "input_cost_per_token": 0, - "input_cost_per_token_above_128k_tokens": 0, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 0, - "output_cost_per_token_above_128k_tokens": 0, - "rpm": 4000, - "source": "https://ai.google.dev/pricing", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 4000000 - }, - "gemini/gemini-1.5-flash-exp-0827": { - "deprecation_date": "2025-09-29", - "input_cost_per_token": 0, - "input_cost_per_token_above_128k_tokens": 0, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 0, - "output_cost_per_token_above_128k_tokens": 0, - "rpm": 2000, - "source": "https://ai.google.dev/pricing", - "supports_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 4000000 - }, - "gemini/gemini-1.5-flash-latest": { - "deprecation_date": "2025-09-29", - "input_cost_per_token": 7.5e-08, - "input_cost_per_token_above_128k_tokens": 1.5e-07, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 3e-07, - "output_cost_per_token_above_128k_tokens": 6e-07, - "rpm": 2000, - "source": "https://ai.google.dev/pricing", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 4000000 - }, - "gemini/gemini-1.5-pro": { - "deprecation_date": "2025-09-29", - "input_cost_per_token": 3.5e-06, - "input_cost_per_token_above_128k_tokens": 7e-06, - "litellm_provider": "gemini", - "max_input_tokens": 2097152, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.05e-05, - "output_cost_per_token_above_128k_tokens": 2.1e-05, - "rpm": 1000, - "source": "https://ai.google.dev/pricing", - "supports_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 4000000 - }, - "gemini/gemini-1.5-pro-001": { - "deprecation_date": "2025-05-24", - "input_cost_per_token": 3.5e-06, - "input_cost_per_token_above_128k_tokens": 7e-06, - "litellm_provider": "gemini", - "max_input_tokens": 2097152, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.05e-05, - "output_cost_per_token_above_128k_tokens": 2.1e-05, - "rpm": 1000, - "source": "https://ai.google.dev/pricing", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 4000000 - }, - "gemini/gemini-1.5-pro-002": { - "deprecation_date": "2025-09-24", - "input_cost_per_token": 3.5e-06, - "input_cost_per_token_above_128k_tokens": 7e-06, - "litellm_provider": "gemini", - "max_input_tokens": 2097152, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.05e-05, - "output_cost_per_token_above_128k_tokens": 2.1e-05, - "rpm": 1000, - "source": "https://ai.google.dev/pricing", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 4000000 - }, - "gemini/gemini-1.5-pro-exp-0801": { - "deprecation_date": "2025-09-29", - "input_cost_per_token": 3.5e-06, - "input_cost_per_token_above_128k_tokens": 7e-06, - "litellm_provider": "gemini", - "max_input_tokens": 2097152, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.05e-05, - "output_cost_per_token_above_128k_tokens": 2.1e-05, - "rpm": 1000, - "source": "https://ai.google.dev/pricing", - "supports_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 4000000 - }, - "gemini/gemini-1.5-pro-exp-0827": { - "deprecation_date": "2025-09-29", - "input_cost_per_token": 0, - "input_cost_per_token_above_128k_tokens": 0, - "litellm_provider": "gemini", - "max_input_tokens": 2097152, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 0, - "output_cost_per_token_above_128k_tokens": 0, - "rpm": 1000, - "source": "https://ai.google.dev/pricing", - "supports_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 4000000 - }, - "gemini/gemini-1.5-pro-latest": { - "deprecation_date": "2025-09-29", - "input_cost_per_token": 3.5e-06, - "input_cost_per_token_above_128k_tokens": 7e-06, - "litellm_provider": "gemini", - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.05e-06, - "output_cost_per_token_above_128k_tokens": 2.1e-05, - "rpm": 1000, - "source": "https://ai.google.dev/pricing", - "supports_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 4000000 + "output_vector_size": 3072, + "rpm": 10000, + "source": "https://ai.google.dev/gemini-api/docs/embeddings#multimodal", + "supports_multimodal": true, + "tpm": 10000000 }, "gemini/gemini-2.0-flash": { "cache_read_input_token_cost": 2.5e-08, @@ -16577,55 +14689,6 @@ "supports_web_search": true, "tpm": 10000000 }, - "gemini/gemini-2.0-flash-exp": { - "cache_read_input_token_cost": 0.0, - "input_cost_per_audio_per_second": 0, - "input_cost_per_audio_per_second_above_128k_tokens": 0, - "input_cost_per_character": 0, - "input_cost_per_character_above_128k_tokens": 0, - "input_cost_per_image": 0, - "input_cost_per_image_above_128k_tokens": 0, - "input_cost_per_token": 0, - "input_cost_per_token_above_128k_tokens": 0, - "input_cost_per_video_per_second": 0, - "input_cost_per_video_per_second_above_128k_tokens": 0, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_character": 0, - "output_cost_per_character_above_128k_tokens": 0, - "output_cost_per_token": 0, - "output_cost_per_token_above_128k_tokens": 0, - "rpm": 10, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-2.0-flash", - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_audio_output": true, - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 4000000 - }, "gemini/gemini-2.0-flash-lite": { "cache_read_input_token_cost": 1.875e-08, "deprecation_date": "2026-06-01", @@ -16663,275 +14726,6 @@ "supports_web_search": true, "tpm": 4000000 }, - "gemini/gemini-2.0-flash-lite-preview-02-05": { - "deprecation_date": "2025-12-09", - "cache_read_input_token_cost": 1.875e-08, - "input_cost_per_audio_token": 7.5e-08, - "input_cost_per_token": 7.5e-08, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 3e-07, - "rpm": 60000, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-2.0-flash-lite", - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 10000000 - }, - "gemini/gemini-2.0-flash-live-001": { - "deprecation_date": "2025-12-09", - "cache_read_input_token_cost": 7.5e-08, - "input_cost_per_audio_token": 2.1e-06, - "input_cost_per_image": 2.1e-06, - "input_cost_per_token": 3.5e-07, - "input_cost_per_video_per_second": 2.1e-06, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_audio_token": 8.5e-06, - "output_cost_per_token": 1.5e-06, - "rpm": 10, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2-0-flash-live-001", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "audio" - ], - "supports_audio_output": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 250000 - }, - "gemini/gemini-2.0-flash-preview-image-generation": { - "deprecation_date": "2025-11-14", - "cache_read_input_token_cost": 2.5e-08, - "input_cost_per_audio_token": 7e-07, - "input_cost_per_token": 1e-07, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 4e-07, - "rpm": 10000, - "source": "https://ai.google.dev/pricing#2_0flash", - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_audio_input": true, - "supports_audio_output": true, - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 10000000 - }, - "gemini/gemini-2.0-flash-thinking-exp": { - "deprecation_date": "2025-12-02", - "cache_read_input_token_cost": 0.0, - "input_cost_per_audio_per_second": 0, - "input_cost_per_audio_per_second_above_128k_tokens": 0, - "input_cost_per_character": 0, - "input_cost_per_character_above_128k_tokens": 0, - "input_cost_per_image": 0, - "input_cost_per_image_above_128k_tokens": 0, - "input_cost_per_token": 0, - "input_cost_per_token_above_128k_tokens": 0, - "input_cost_per_video_per_second": 0, - "input_cost_per_video_per_second_above_128k_tokens": 0, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65536, - "max_pdf_size_mb": 30, - "max_tokens": 65536, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_character": 0, - "output_cost_per_character_above_128k_tokens": 0, - "output_cost_per_token": 0, - "output_cost_per_token_above_128k_tokens": 0, - "rpm": 10, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-2.0-flash", - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_audio_output": true, - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 4000000 - }, - "gemini/gemini-2.0-flash-thinking-exp-01-21": { - "deprecation_date": "2025-12-02", - "cache_read_input_token_cost": 0.0, - "input_cost_per_audio_per_second": 0, - "input_cost_per_audio_per_second_above_128k_tokens": 0, - "input_cost_per_character": 0, - "input_cost_per_character_above_128k_tokens": 0, - "input_cost_per_image": 0, - "input_cost_per_image_above_128k_tokens": 0, - "input_cost_per_token": 0, - "input_cost_per_token_above_128k_tokens": 0, - "input_cost_per_video_per_second": 0, - "input_cost_per_video_per_second_above_128k_tokens": 0, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65536, - "max_pdf_size_mb": 30, - "max_tokens": 65536, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_character": 0, - "output_cost_per_character_above_128k_tokens": 0, - "output_cost_per_token": 0, - "output_cost_per_token_above_128k_tokens": 0, - "rpm": 10, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-2.0-flash", - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_audio_output": true, - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 4000000 - }, - "gemini/gemini-2.0-pro-exp-02-05": { - "cache_read_input_token_cost": 0.0, - "input_cost_per_audio_per_second": 0, - "input_cost_per_audio_per_second_above_128k_tokens": 0, - "input_cost_per_character": 0, - "input_cost_per_character_above_128k_tokens": 0, - "input_cost_per_image": 0, - "input_cost_per_image_above_128k_tokens": 0, - "input_cost_per_token": 0, - "input_cost_per_token_above_128k_tokens": 0, - "input_cost_per_video_per_second": 0, - "input_cost_per_video_per_second_above_128k_tokens": 0, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 2097152, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_character": 0, - "output_cost_per_character_above_128k_tokens": 0, - "output_cost_per_token": 0, - "output_cost_per_token_above_128k_tokens": 0, - "rpm": 2, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supports_audio_input": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_video_input": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 1000000 - }, "gemini/gemini-2.5-flash": { "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, @@ -17029,56 +14823,6 @@ "supports_web_search": true, "tpm": 8000000 }, - "gemini/gemini-2.5-flash-image-preview": { - "deprecation_date": "2026-01-15", - "cache_read_input_token_cost": 7.5e-08, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "image_generation", - "output_cost_per_image": 0.039, - "output_cost_per_image_token": 3e-05, - "output_cost_per_reasoning_token": 3e-05, - "output_cost_per_token": 3e-05, - "rpm": 100000, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 8000000 - }, "gemini/gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -17470,96 +15214,6 @@ "supports_web_search": true, "tpm": 250000 }, - "gemini/gemini-2.5-flash-preview-04-17": { - "cache_read_input_token_cost": 3.75e-08, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 1.5e-07, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_reasoning_token": 3.5e-06, - "output_cost_per_token": 6e-07, - "rpm": 10, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 250000 - }, - "gemini/gemini-2.5-flash-preview-05-20": { - "deprecation_date": "2025-11-18", - "cache_read_input_token_cost": 7.5e-08, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_reasoning_token": 2.5e-06, - "output_cost_per_token": 2.5e-06, - "rpm": 10, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 250000 - }, "gemini/gemini-2.5-flash-preview-tts": { "input_cost_per_token": 3e-07, "litellm_provider": "gemini", @@ -17984,177 +15638,6 @@ "cache_read_input_token_cost_priority": 9e-08, "supports_service_tier": true }, - "gemini/gemini-2.5-pro-exp-03-25": { - "cache_read_input_token_cost": 0.0, - "input_cost_per_token": 0.0, - "input_cost_per_token_above_200k_tokens": 0.0, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 0.0, - "output_cost_per_token_above_200k_tokens": 0.0, - "rpm": 5, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_input": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_video_input": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 250000 - }, - "gemini/gemini-2.5-pro-preview-03-25": { - "deprecation_date": "2025-12-02", - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_audio_token": 7e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "rpm": 10000, - "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-pro-preview", - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 10000000 - }, - "gemini/gemini-2.5-pro-preview-05-06": { - "deprecation_date": "2025-12-02", - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_audio_token": 7e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "rpm": 10000, - "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-pro-preview", - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 10000000 - }, - "gemini/gemini-2.5-pro-preview-06-05": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_audio_token": 7e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "rpm": 10000, - "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-pro-preview", - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 10000000 - }, "gemini/gemini-2.5-pro-preview-tts": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, @@ -18278,41 +15761,6 @@ "tpm": 250000, "rpm": 10 }, - "gemini/gemini-pro": { - "input_cost_per_token": 3.5e-07, - "input_cost_per_token_above_128k_tokens": 7e-07, - "litellm_provider": "gemini", - "max_input_tokens": 32760, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.05e-06, - "output_cost_per_token_above_128k_tokens": 2.1e-06, - "rpd": 30000, - "rpm": 360, - "source": "https://ai.google.dev/gemini-api/docs/models/gemini", - "supports_function_calling": true, - "supports_tool_choice": true, - "tpm": 120000 - }, - "gemini/gemini-pro-vision": { - "input_cost_per_token": 3.5e-07, - "input_cost_per_token_above_128k_tokens": 7e-07, - "litellm_provider": "gemini", - "max_input_tokens": 30720, - "max_output_tokens": 2048, - "max_tokens": 2048, - "mode": "chat", - "output_cost_per_token": 1.05e-06, - "output_cost_per_token_above_128k_tokens": 2.1e-06, - "rpd": 30000, - "rpm": 360, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 120000 - }, "gemini/gemma-3-27b-it": { "input_cost_per_audio_per_second": 0, "input_cost_per_audio_per_second_above_128k_tokens": 0, @@ -18420,36 +15868,6 @@ "video" ] }, - "gemini/veo-3.0-fast-generate-preview": { - "deprecation_date": "2025-11-12", - "litellm_provider": "gemini", - "max_input_tokens": 1024, - "max_tokens": 1024, - "mode": "video_generation", - "output_cost_per_second": 0.4, - "source": "https://ai.google.dev/gemini-api/docs/video", - "supported_modalities": [ - "text" - ], - "supported_output_modalities": [ - "video" - ] - }, - "gemini/veo-3.0-generate-preview": { - "deprecation_date": "2025-11-12", - "litellm_provider": "gemini", - "max_input_tokens": 1024, - "max_tokens": 1024, - "mode": "video_generation", - "output_cost_per_second": 0.75, - "source": "https://ai.google.dev/gemini-api/docs/video", - "supported_modalities": [ - "text" - ], - "supported_output_modalities": [ - "video" - ] - }, "gemini/veo-3.1-fast-generate-preview": { "litellm_provider": "gemini", "max_input_tokens": 1024, @@ -19373,31 +16791,6 @@ "supports_system_messages": true, "supports_tool_choice": true }, - "gpt-3.5-turbo-0301": { - "input_cost_per_token": 1.5e-06, - "litellm_provider": "openai", - "max_input_tokens": 4097, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 2e-06, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, - "gpt-3.5-turbo-0613": { - "input_cost_per_token": 1.5e-06, - "litellm_provider": "openai", - "max_input_tokens": 4097, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 2e-06, - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, "gpt-3.5-turbo-1106": { "deprecation_date": "2026-09-28", "input_cost_per_token": 1e-06, @@ -19425,18 +16818,6 @@ "supports_system_messages": true, "supports_tool_choice": true }, - "gpt-3.5-turbo-16k-0613": { - "input_cost_per_token": 3e-06, - "litellm_provider": "openai", - "max_input_tokens": 16385, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 4e-06, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, "gpt-3.5-turbo-instruct": { "input_cost_per_token": 1.5e-06, "litellm_provider": "text-completion-openai", @@ -19483,18 +16864,6 @@ "supports_system_messages": true, "supports_tool_choice": true }, - "gpt-4-0314": { - "input_cost_per_token": 3e-05, - "litellm_provider": "openai", - "max_input_tokens": 8192, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 6e-05, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, "gpt-4-0613": { "deprecation_date": "2025-06-06", "input_cost_per_token": 3e-05, @@ -19524,57 +16893,6 @@ "supports_system_messages": true, "supports_tool_choice": true }, - "gpt-4-1106-vision-preview": { - "deprecation_date": "2024-12-06", - "input_cost_per_token": 1e-05, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 3e-05, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gpt-4-32k": { - "input_cost_per_token": 6e-05, - "litellm_provider": "openai", - "max_input_tokens": 32768, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 0.00012, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, - "gpt-4-32k-0314": { - "input_cost_per_token": 6e-05, - "litellm_provider": "openai", - "max_input_tokens": 32768, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 0.00012, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, - "gpt-4-32k-0613": { - "input_cost_per_token": 6e-05, - "litellm_provider": "openai", - "max_input_tokens": 32768, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 0.00012, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, "gpt-4-turbo": { "input_cost_per_token": 1e-05, "litellm_provider": "openai", @@ -19622,21 +16940,6 @@ "supports_system_messages": true, "supports_tool_choice": true }, - "gpt-4-vision-preview": { - "deprecation_date": "2024-12-06", - "input_cost_per_token": 1e-05, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 3e-05, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, "gpt-4.1": { "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_priority": 8.75e-07, @@ -19854,47 +17157,6 @@ "supports_service_tier": true, "supports_vision": true }, - "gpt-4.5-preview": { - "cache_read_input_token_cost": 3.75e-05, - "input_cost_per_token": 7.5e-05, - "input_cost_per_token_batches": 3.75e-05, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 16384, - "max_tokens": 16384, - "mode": "chat", - "output_cost_per_token": 0.00015, - "output_cost_per_token_batches": 7.5e-05, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gpt-4.5-preview-2025-02-27": { - "cache_read_input_token_cost": 3.75e-05, - "deprecation_date": "2025-07-14", - "input_cost_per_token": 7.5e-05, - "input_cost_per_token_batches": 3.75e-05, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 16384, - "max_tokens": 16384, - "mode": "chat", - "output_cost_per_token": 0.00015, - "output_cost_per_token_batches": 7.5e-05, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, "gpt-4o": { "cache_read_input_token_cost": 1.25e-06, "cache_read_input_token_cost_priority": 2.125e-06, @@ -19998,23 +17260,6 @@ "supports_system_messages": true, "supports_tool_choice": true }, - "gpt-4o-audio-preview-2024-10-01": { - "input_cost_per_audio_token": 4e-05, - "input_cost_per_token": 2.5e-06, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 16384, - "max_tokens": 16384, - "mode": "chat", - "output_cost_per_audio_token": 8e-05, - "output_cost_per_token": 1e-05, - "supports_audio_input": true, - "supports_audio_output": true, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, "gpt-4o-audio-preview-2024-12-17": { "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 2.5e-06, @@ -20478,25 +17723,6 @@ "supports_system_messages": true, "supports_tool_choice": true }, - "gpt-4o-realtime-preview-2024-10-01": { - "cache_creation_input_audio_token_cost": 2e-05, - "cache_read_input_token_cost": 2.5e-06, - "input_cost_per_audio_token": 0.0001, - "input_cost_per_token": 5e-06, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_audio_token": 0.0002, - "output_cost_per_token": 2e-05, - "supports_audio_input": true, - "supports_audio_output": true, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, "gpt-4o-realtime-preview-2024-12-17": { "cache_read_input_token_cost": 2.5e-06, "input_cost_per_audio_token": 4e-05, @@ -25700,62 +22926,6 @@ "supports_tool_choice": true, "supports_vision": true }, - "o1-mini": { - "cache_read_input_token_cost": 5.5e-07, - "input_cost_per_token": 1.1e-06, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 65536, - "max_tokens": 65536, - "mode": "chat", - "output_cost_per_token": 4.4e-06, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_vision": true - }, - "o1-mini-2024-09-12": { - "deprecation_date": "2025-10-27", - "cache_read_input_token_cost": 1.5e-06, - "input_cost_per_token": 3e-06, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 65536, - "max_tokens": 65536, - "mode": "chat", - "output_cost_per_token": 1.2e-05, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_vision": true - }, - "o1-preview": { - "cache_read_input_token_cost": 7.5e-06, - "input_cost_per_token": 1.5e-05, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 32768, - "max_tokens": 32768, - "mode": "chat", - "output_cost_per_token": 6e-05, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_vision": true - }, - "o1-preview-2024-09-12": { - "cache_read_input_token_cost": 7.5e-06, - "input_cost_per_token": 1.5e-05, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 32768, - "max_tokens": 32768, - "mode": "chat", - "output_cost_per_token": 6e-05, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_vision": true - }, "o1-pro": { "input_cost_per_token": 0.00015, "input_cost_per_token_batches": 7.5e-05, @@ -26622,15 +23792,6 @@ "mode": "moderation", "output_cost_per_token": 0.0 }, - "omni-moderation-latest-intents": { - "input_cost_per_token": 0.0, - "litellm_provider": "openai", - "max_input_tokens": 32768, - "max_output_tokens": 0, - "max_tokens": 0, - "mode": "moderation", - "output_cost_per_token": 0.0 - }, "openai.gpt-oss-120b-1:0": { "input_cost_per_token": 1.5e-07, "litellm_provider": "bedrock_converse", @@ -28380,56 +25541,6 @@ "mode": "chat", "output_cost_per_token": 2e-07 }, - "perplexity/llama-3.1-sonar-huge-128k-online": { - "deprecation_date": "2025-02-22", - "input_cost_per_token": 5e-06, - "litellm_provider": "perplexity", - "max_input_tokens": 127072, - "max_output_tokens": 127072, - "max_tokens": 127072, - "mode": "chat", - "output_cost_per_token": 5e-06 - }, - "perplexity/llama-3.1-sonar-large-128k-chat": { - "deprecation_date": "2025-02-22", - "input_cost_per_token": 1e-06, - "litellm_provider": "perplexity", - "max_input_tokens": 131072, - "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", - "output_cost_per_token": 1e-06 - }, - "perplexity/llama-3.1-sonar-large-128k-online": { - "deprecation_date": "2025-02-22", - "input_cost_per_token": 1e-06, - "litellm_provider": "perplexity", - "max_input_tokens": 127072, - "max_output_tokens": 127072, - "max_tokens": 127072, - "mode": "chat", - "output_cost_per_token": 1e-06 - }, - "perplexity/llama-3.1-sonar-small-128k-chat": { - "deprecation_date": "2025-02-22", - "input_cost_per_token": 2e-07, - "litellm_provider": "perplexity", - "max_input_tokens": 131072, - "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", - "output_cost_per_token": 2e-07 - }, - "perplexity/llama-3.1-sonar-small-128k-online": { - "deprecation_date": "2025-02-22", - "input_cost_per_token": 2e-07, - "litellm_provider": "perplexity", - "max_input_tokens": 127072, - "max_output_tokens": 127072, - "max_tokens": 127072, - "mode": "chat", - "output_cost_per_token": 2e-07 - }, "perplexity/mistral-7b-instruct": { "input_cost_per_token": 7e-08, "litellm_provider": "perplexity", @@ -30212,60 +27323,6 @@ "litellm_provider": "tavily", "mode": "search" }, - "text-bison": { - "input_cost_per_character": 2.5e-07, - "litellm_provider": "vertex_ai-text-models", - "max_input_tokens": 8192, - "max_output_tokens": 2048, - "max_tokens": 2048, - "mode": "completion", - "output_cost_per_character": 5e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "text-bison32k": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-text-models", - "max_input_tokens": 8192, - "max_output_tokens": 1024, - "max_tokens": 1024, - "mode": "completion", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "text-bison32k@002": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-text-models", - "max_input_tokens": 8192, - "max_output_tokens": 1024, - "max_tokens": 1024, - "mode": "completion", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "text-bison@001": { - "input_cost_per_character": 2.5e-07, - "litellm_provider": "vertex_ai-text-models", - "max_input_tokens": 8192, - "max_output_tokens": 1024, - "max_tokens": 1024, - "mode": "completion", - "output_cost_per_character": 5e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "text-bison@002": { - "input_cost_per_character": 2.5e-07, - "litellm_provider": "vertex_ai-text-models", - "max_input_tokens": 8192, - "max_output_tokens": 1024, - "max_tokens": 1024, - "mode": "completion", - "output_cost_per_character": 5e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, "text-completion-codestral/codestral-2405": { "input_cost_per_token": 0.0, "litellm_provider": "text-completion-codestral", @@ -30410,16 +27467,6 @@ "output_vector_size": 768, "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models" }, - "text-multilingual-embedding-preview-0409": { - "input_cost_per_token": 6.25e-09, - "litellm_provider": "vertex_ai-embedding-models", - "max_input_tokens": 3072, - "max_tokens": 3072, - "mode": "embedding", - "output_cost_per_token": 0, - "output_vector_size": 768, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, "text-unicorn": { "input_cost_per_token": 1e-05, "litellm_provider": "vertex_ai-text-models", @@ -30440,61 +27487,6 @@ "output_cost_per_token": 2.8e-05, "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" }, - "textembedding-gecko": { - "input_cost_per_character": 2.5e-08, - "input_cost_per_token": 1e-07, - "litellm_provider": "vertex_ai-embedding-models", - "max_input_tokens": 3072, - "max_tokens": 3072, - "mode": "embedding", - "output_cost_per_token": 0, - "output_vector_size": 768, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "textembedding-gecko-multilingual": { - "input_cost_per_character": 2.5e-08, - "input_cost_per_token": 1e-07, - "litellm_provider": "vertex_ai-embedding-models", - "max_input_tokens": 3072, - "max_tokens": 3072, - "mode": "embedding", - "output_cost_per_token": 0, - "output_vector_size": 768, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "textembedding-gecko-multilingual@001": { - "input_cost_per_character": 2.5e-08, - "input_cost_per_token": 1e-07, - "litellm_provider": "vertex_ai-embedding-models", - "max_input_tokens": 3072, - "max_tokens": 3072, - "mode": "embedding", - "output_cost_per_token": 0, - "output_vector_size": 768, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "textembedding-gecko@001": { - "input_cost_per_character": 2.5e-08, - "input_cost_per_token": 1e-07, - "litellm_provider": "vertex_ai-embedding-models", - "max_input_tokens": 3072, - "max_tokens": 3072, - "mode": "embedding", - "output_cost_per_token": 0, - "output_vector_size": 768, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "textembedding-gecko@003": { - "input_cost_per_character": 2.5e-08, - "input_cost_per_token": 1e-07, - "litellm_provider": "vertex_ai-embedding-models", - "max_input_tokens": 3072, - "max_tokens": 3072, - "mode": "embedding", - "output_cost_per_token": 0, - "output_vector_size": 768, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, "together-ai-21.1b-41b": { "input_cost_per_token": 8e-07, "litellm_provider": "together_ai", @@ -32896,36 +29888,6 @@ "supports_tool_choice": true, "supports_vision": true }, - "vertex_ai/claude-3-5-sonnet-v2": { - "input_cost_per_token": 3e-06, - "litellm_provider": "vertex_ai-anthropic_models", - "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.5e-05, - "supports_assistant_prefill": true, - "supports_computer_use": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "vertex_ai/claude-3-5-sonnet-v2@20241022": { - "input_cost_per_token": 3e-06, - "litellm_provider": "vertex_ai-anthropic_models", - "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.5e-05, - "supports_assistant_prefill": true, - "supports_computer_use": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_tool_choice": true, - "supports_vision": true - }, "vertex_ai/claude-3-5-sonnet@20240620": { "input_cost_per_token": 3e-06, "litellm_provider": "vertex_ai-anthropic_models", @@ -32943,7 +29905,7 @@ "vertex_ai/claude-3-7-sonnet@20250219": { "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, - "deprecation_date": "2025-06-01", + "deprecation_date": "2026-05-11", "input_cost_per_token": 3e-06, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, @@ -33959,6 +30921,9 @@ "mode": "chat", "output_cost_per_token": 2.2e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#partner-models", + "supported_regions": [ + "global" + ], "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true @@ -33973,6 +30938,7 @@ "mode": "chat", "output_cost_per_token": 3.2e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#glm-models", + "supported_regions": ["global"], "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -34224,36 +31190,6 @@ "video" ] }, - "vertex_ai/veo-3.0-fast-generate-preview": { - "deprecation_date": "2025-11-12", - "litellm_provider": "vertex_ai-video-models", - "max_input_tokens": 1024, - "max_tokens": 1024, - "mode": "video_generation", - "output_cost_per_second": 0.15, - "source": "https://ai.google.dev/gemini-api/docs/video", - "supported_modalities": [ - "text" - ], - "supported_output_modalities": [ - "video" - ] - }, - "vertex_ai/veo-3.0-generate-preview": { - "deprecation_date": "2025-11-12", - "litellm_provider": "vertex_ai-video-models", - "max_input_tokens": 1024, - "max_tokens": 1024, - "mode": "video_generation", - "output_cost_per_second": 0.4, - "source": "https://ai.google.dev/gemini-api/docs/video", - "supported_modalities": [ - "text" - ], - "supported_output_modalities": [ - "video" - ] - }, "vertex_ai/veo-3.0-fast-generate-001": { "litellm_provider": "vertex_ai-video-models", "max_input_tokens": 1024, diff --git a/litellm/ocr/__init__.py b/litellm/ocr/__init__.py index 53f455619d..a20b0ef6ca 100644 --- a/litellm/ocr/__init__.py +++ b/litellm/ocr/__init__.py @@ -2,4 +2,3 @@ from .main import aocr, ocr __all__ = ["ocr", "aocr"] - diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index e76a222b2e..edee50bdfc 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -113,7 +113,7 @@ async def allm_passthrough_route( # Only call raise_for_status if it's a Response object (not a generator) if isinstance(response, httpx.Response): response.raise_for_status() - + return response else: # This shouldn't happen when allm_passthrough_route=True, but handle it for type safety @@ -216,11 +216,11 @@ def llm_passthrough_route( ) litellm_params_dict = get_litellm_params(**kwargs) - + # Add model_id to litellm_params if present in kwargs (for Bedrock Application Inference Profiles) if "model_id" in kwargs: litellm_params_dict["model_id"] = kwargs["model_id"] - + litellm_logging_obj.update_environment_variables( model=model, litellm_params=litellm_params_dict, @@ -363,7 +363,7 @@ async def _async_passthrough_request( """ # client.client.send returns a coroutine for async clients response_result = client.client.send(request=request, stream=is_streaming_request) - + # Check if it's a coroutine and await it if asyncio.iscoroutine(response_result): if is_streaming_request: @@ -416,7 +416,6 @@ async def _async_streaming( raw_bytes: List[bytes] = [] async for chunk in iter_response.aiter_bytes(): # type: ignore - raw_bytes.append(chunk) yield chunk diff --git a/litellm/passthrough/utils.py b/litellm/passthrough/utils.py index fe1ecad96c..ef4357d1ca 100644 --- a/litellm/passthrough/utils.py +++ b/litellm/passthrough/utils.py @@ -11,7 +11,7 @@ class BasePassthroughUtils: def get_merged_query_parameters( existing_url: httpx.URL, request_query_params: Mapping[str, Union[str, list]], - default_query_params: Optional[Dict[str, Union[str, list]]] = None + default_query_params: Optional[Dict[str, Union[str, list]]] = None, ) -> Dict[str, Union[str, List[str]]]: # Get the existing query params from the target URL existing_query_string = existing_url.query.decode("utf-8") @@ -65,6 +65,7 @@ class BasePassthroughUtils: return headers + class CommonUtils: @staticmethod def encode_bedrock_runtime_modelid_arn(endpoint: str) -> str: @@ -77,37 +78,36 @@ class CommonUtils: arn:aws:bedrock:ap-southeast-1:123456789012:application-inference-profile%2Fabdefg12334 so that it is treated as one part of the path. Otherwise, the encoded endpoint will return 500 error when passed to Bedrock endpoint. - + See the apis in https://docs.aws.amazon.com/bedrock/latest/APIReference/API_Operations_Amazon_Bedrock_Runtime.html for more details on the regex patterns of modelId which we use in the regex logic below. - + Args: endpoint (str): The original endpoint string which may contain ARNs that contain slashes. - + Returns: str: The endpoint with properly encoded ARN slashes """ import re # Early exit: if no ARN detected, return unchanged - if 'arn:aws:' not in endpoint: + if "arn:aws:" not in endpoint: return endpoint # Handle all patterns in one go - more efficient and cleaner patterns = [ # Custom model with 2 slashes (order matters - do this first) - (r'(custom-model)/([a-z0-9.-]+)/([a-z0-9]+)', r'\1%2F\2%2F\3'), - + (r"(custom-model)/([a-z0-9.-]+)/([a-z0-9]+)", r"\1%2F\2%2F\3"), # All other resource types with 1 slash - (r'(:application-inference-profile)/', r'\1%2F'), - (r'(:inference-profile)/', r'\1%2F'), - (r'(:foundation-model)/', r'\1%2F'), - (r'(:imported-model)/', r'\1%2F'), - (r'(:provisioned-model)/', r'\1%2F'), - (r'(:prompt)/', r'\1%2F'), - (r'(:endpoint)/', r'\1%2F'), - (r'(:prompt-router)/', r'\1%2F'), - (r'(:default-prompt-router)/', r'\1%2F'), + (r"(:application-inference-profile)/", r"\1%2F"), + (r"(:inference-profile)/", r"\1%2F"), + (r"(:foundation-model)/", r"\1%2F"), + (r"(:imported-model)/", r"\1%2F"), + (r"(:provisioned-model)/", r"\1%2F"), + (r"(:prompt)/", r"\1%2F"), + (r"(:endpoint)/", r"\1%2F"), + (r"(:prompt-router)/", r"\1%2F"), + (r"(:default-prompt-router)/", r"\1%2F"), ] for pattern, replacement in patterns: @@ -116,4 +116,4 @@ class CommonUtils: endpoint = re.sub(pattern, replacement, endpoint) break # Exit after first match since each ARN has only one resource type - return endpoint \ No newline at end of file + return endpoint diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index c670146be3..357d21eb09 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -388,7 +388,6 @@ class MCPRequestHandler: ) ) - # If end_user has explicit MCP server permissions, apply intersection if len(allowed_mcp_servers_for_end_user) > 0: verbose_logger.debug( @@ -547,16 +546,16 @@ class MCPRequestHandler: agent_obj_perm = await MCPRequestHandler._get_agent_object_permission( user_api_key_auth ) - agent_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server( - server_id=server_id, - user_api_key_auth=user_api_key_auth, - agent_object_permission=agent_obj_perm, + agent_tools = ( + await MCPRequestHandler._get_agent_tool_permissions_for_server( + server_id=server_id, + user_api_key_auth=user_api_key_auth, + agent_object_permission=agent_obj_perm, + ) ) if agent_tools is not None: if allowed_tools is not None: - allowed_tools = list( - set(allowed_tools) & set(agent_tools) - ) + allowed_tools = list(set(allowed_tools) & set(agent_tools)) else: allowed_tools = agent_tools return allowed_tools @@ -621,13 +620,18 @@ class MCPRequestHandler: key_object_permission = MCPRequestHandler._get_key_object_permission( user_api_key_auth ) - if key_object_permission is None and user_api_key_auth and user_api_key_auth.object_permission_id: + if ( + key_object_permission is None + and user_api_key_auth + and user_api_key_auth.object_permission_id + ): from litellm.proxy.auth.auth_checks import get_object_permission from litellm.proxy.proxy_server import ( prisma_client, proxy_logging_obj, user_api_key_cache, ) + if prisma_client is not None: key_object_permission = await get_object_permission( object_permission_id=user_api_key_auth.object_permission_id, @@ -725,7 +729,6 @@ class MCPRequestHandler: return [] if prisma_client is None: - verbose_logger.debug("prisma_client is None") return [] @@ -740,7 +743,6 @@ class MCPRequestHandler: route="/mcp", ) - if end_user_obj is None or end_user_obj.object_permission is None: return [] @@ -796,9 +798,7 @@ class MCPRequestHandler: return agent_row.object_permission except Exception as e: - verbose_logger.warning( - f"Failed to get agent object permission: {str(e)}" - ) + verbose_logger.warning(f"Failed to get agent object permission: {str(e)}") return None @staticmethod @@ -877,9 +877,7 @@ class MCPRequestHandler: if obj_perm is None: return None - mcp_tool_permissions = getattr( - obj_perm, "mcp_tool_permissions", None - ) + mcp_tool_permissions = getattr(obj_perm, "mcp_tool_permissions", None) if not mcp_tool_permissions: return None if isinstance(mcp_tool_permissions, dict): diff --git a/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py b/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py index db18885721..48884d8227 100644 --- a/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py @@ -653,7 +653,9 @@ async def byok_authorize_post( # Reject new codes if the store is at capacity (prevents memory exhaustion # from a burst of abandoned OAuth flows). if len(_byok_auth_codes) >= _AUTH_CODES_MAX_SIZE: - raise HTTPException(status_code=503, detail="Too many pending authorization flows") + raise HTTPException( + status_code=503, detail="Too many pending authorization flows" + ) if code_challenge_method != "S256": raise HTTPException( @@ -745,6 +747,7 @@ async def byok_token( from litellm.proxy._experimental.mcp_server.server import ( _invalidate_byok_cred_cache, ) + _invalidate_byok_cred_cache(user_id, server_id) except Exception as exc: verbose_proxy_logger.error( diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 119e8171a1..45ec1bcebd 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -68,9 +68,13 @@ def _prepare_mcp_server_data( # Handle tool name override serialization if data.tool_name_to_display_name is not None: - data_dict["tool_name_to_display_name"] = safe_dumps(data.tool_name_to_display_name) + data_dict["tool_name_to_display_name"] = safe_dumps( + data.tool_name_to_display_name + ) if data.tool_name_to_description is not None: - data_dict["tool_name_to_description"] = safe_dumps(data.tool_name_to_description) + data_dict["tool_name_to_description"] = safe_dumps( + data.tool_name_to_description + ) # mcp_access_groups is already List[str], no serialization needed @@ -405,7 +409,9 @@ async def update_mcp_server( # Pre-fetch existing record once if we need it for auth_type or credential logic existing = None - has_credentials = "credentials" in data_dict and data_dict["credentials"] is not None + has_credentials = ( + "credentials" in data_dict and data_dict["credentials"] is not None + ) if data.auth_type or has_credentials: existing = await prisma_client.db.litellm_mcpservertable.find_unique( where={"server_id": data.server_id} @@ -746,7 +752,9 @@ async def get_mcp_submissions( ) items = [LiteLLM_MCPServerTable(**r.model_dump()) for r in rows] - pending = sum(1 for i in items if i.approval_status == MCPApprovalStatus.pending_review) + pending = sum( + 1 for i in items if i.approval_status == MCPApprovalStatus.pending_review + ) active = sum(1 for i in items if i.approval_status == MCPApprovalStatus.active) rejected = sum(1 for i in items if i.approval_status == MCPApprovalStatus.rejected) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index ad1dadb122..af3a715051 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -141,9 +141,7 @@ def _resolve_oauth2_server_for_root_endpoints( ) registry = global_mcp_server_manager.get_filtered_registry(client_ip=client_ip) - oauth2_servers = [ - s for s in registry.values() if s.auth_type == MCPAuth.oauth2 - ] + oauth2_servers = [s for s in registry.values() if s.auth_type == MCPAuth.oauth2] if len(oauth2_servers) == 1: return oauth2_servers[0] return None @@ -197,9 +195,7 @@ async def authorize_with_server( parsed_auth_url = urlparse(mcp_server.authorization_url) existing_params = dict(parse_qsl(parsed_auth_url.query)) existing_params.update(params) - final_url = urlunparse( - parsed_auth_url._replace(query=urlencode(existing_params)) - ) + final_url = urlunparse(parsed_auth_url._replace(query=urlencode(existing_params))) return RedirectResponse(final_url) @@ -333,7 +329,9 @@ async def authorize( lookup_name: Optional[str] = mcp_server_name or client_id client_ip = IPAddressUtils.get_mcp_client_ip(request) mcp_server = ( - global_mcp_server_manager.get_mcp_server_by_name(lookup_name, client_ip=client_ip) + global_mcp_server_manager.get_mcp_server_by_name( + lookup_name, client_ip=client_ip + ) if lookup_name else None ) @@ -513,16 +511,18 @@ def _build_oauth_protected_resource_response( ) ], "resource": resource_url, - "scopes_supported": mcp_server.scopes if mcp_server and mcp_server.scopes else [], + "scopes_supported": mcp_server.scopes + if mcp_server and mcp_server.scopes + else [], } # Standard MCP pattern: /.well-known/oauth-protected-resource/mcp/{server_name} # This is the pattern expected by standard MCP clients (mcp-inspector, VSCode Copilot) -@router.get(f"/.well-known/oauth-protected-resource{'' if get_server_root_path() == '/' else get_server_root_path()}/mcp/{{mcp_server_name}}") -async def oauth_protected_resource_mcp_standard( - request: Request, mcp_server_name: str -): +@router.get( + f"/.well-known/oauth-protected-resource{'' if get_server_root_path() == '/' else get_server_root_path()}/mcp/{{mcp_server_name}}" +) +async def oauth_protected_resource_mcp_standard(request: Request, mcp_server_name: str): """ OAuth protected resource discovery endpoint using standard MCP URL pattern. @@ -541,7 +541,9 @@ async def oauth_protected_resource_mcp_standard( # LiteLLM legacy pattern: /.well-known/oauth-protected-resource/{server_name}/mcp # Kept for backward compatibility with existing deployments -@router.get(f"/.well-known/oauth-protected-resource{'' if get_server_root_path() == '/' else get_server_root_path()}/{{mcp_server_name}}/mcp") +@router.get( + f"/.well-known/oauth-protected-resource{'' if get_server_root_path() == '/' else get_server_root_path()}/{{mcp_server_name}}/mcp" +) @router.get("/.well-known/oauth-protected-resource") async def oauth_protected_resource_mcp( request: Request, mcp_server_name: Optional[str] = None @@ -561,6 +563,7 @@ async def oauth_protected_resource_mcp( use_standard_pattern=False, ) + """ https://datatracker.ietf.org/doc/html/rfc8414#section-3.1 RFC 8414: Path-aware OAuth discovery @@ -620,17 +623,23 @@ def _build_oauth_authorization_server_response( "authorization_endpoint": authorization_endpoint, "token_endpoint": token_endpoint, "response_types_supported": ["code"], - "scopes_supported": mcp_server.scopes if mcp_server and mcp_server.scopes else [], + "scopes_supported": mcp_server.scopes + if mcp_server and mcp_server.scopes + else [], "grant_types_supported": ["authorization_code", "refresh_token"], "code_challenge_methods_supported": ["S256"], "token_endpoint_auth_methods_supported": ["client_secret_post"], # Claude expects a registration endpoint, even if we just fake it - "registration_endpoint": f"{request_base_url}/{mcp_server_name}/register" if mcp_server_name else f"{request_base_url}/register", + "registration_endpoint": f"{request_base_url}/{mcp_server_name}/register" + if mcp_server_name + else f"{request_base_url}/register", } # Standard MCP pattern: /.well-known/oauth-authorization-server/mcp/{server_name} -@router.get(f"/.well-known/oauth-authorization-server{'' if get_server_root_path() == '/' else get_server_root_path()}/mcp/{{mcp_server_name}}") +@router.get( + f"/.well-known/oauth-authorization-server{'' if get_server_root_path() == '/' else get_server_root_path()}/mcp/{{mcp_server_name}}" +) async def oauth_authorization_server_mcp_standard( request: Request, mcp_server_name: str ): @@ -647,7 +656,9 @@ async def oauth_authorization_server_mcp_standard( # LiteLLM legacy pattern and root endpoint -@router.get(f"/.well-known/oauth-authorization-server{'' if get_server_root_path() == '/' else get_server_root_path()}/{{mcp_server_name}}") +@router.get( + f"/.well-known/oauth-authorization-server{'' if get_server_root_path() == '/' else get_server_root_path()}/{{mcp_server_name}}" +) @router.get("/.well-known/oauth-authorization-server") async def oauth_authorization_server_mcp( request: Request, mcp_server_name: Optional[str] = None @@ -671,9 +682,7 @@ async def openid_configuration(request: Request): # Additional legacy pattern support @router.get("/.well-known/oauth-authorization-server/{mcp_server_name}/mcp") -async def oauth_authorization_server_legacy( - request: Request, mcp_server_name: str -): +async def oauth_authorization_server_legacy(request: Request, mcp_server_name: str): """ OAuth authorization server discovery for legacy /{server_name}/mcp pattern. """ @@ -710,9 +719,7 @@ async def register_client(request: Request, mcp_server_name: Optional[str] = Non client_name=data.get("client_name", ""), grant_types=data.get("grant_types", []), response_types=data.get("response_types", []), - token_endpoint_auth_method=data.get( - "token_endpoint_auth_method", "" - ), + token_endpoint_auth_method=data.get("token_endpoint_auth_method", ""), fallback_client_id=resolved.server_name or resolved.name, ) return dummy_return diff --git a/litellm/proxy/_experimental/mcp_server/mcp_debug.py b/litellm/proxy/_experimental/mcp_server/mcp_debug.py index 46741a9df9..254f208e23 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_debug.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_debug.py @@ -254,9 +254,7 @@ class MCPDebug: return debug @staticmethod - def wrap_send_with_debug_headers( - send: Send, debug_headers: Dict[str, str] - ) -> Send: + def wrap_send_with_debug_headers(send: Send, debug_headers: Dict[str, str]) -> Send: """ Return a new ASGI ``send`` callable that injects *debug_headers* into the ``http.response.start`` message. @@ -315,9 +313,7 @@ class MCPDebug: break scope_headers = MCPRequestHandler._safe_get_headers_from_scope(scope) - litellm_key = MCPRequestHandler.get_litellm_api_key_from_headers( - scope_headers - ) + litellm_key = MCPRequestHandler.get_litellm_api_key_from_headers(scope_headers) return MCPDebug.build_debug_headers( inbound_headers=raw_headers, diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index b10bfde491..43fe54fdfb 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -501,12 +501,12 @@ class MCPServerManager: ) # Update tool name to server name mapping (for both prefixed and base names) - self.tool_name_to_mcp_server_name_mapping[base_tool_name] = ( - server_prefix - ) - self.tool_name_to_mcp_server_name_mapping[prefixed_tool_name] = ( - server_prefix - ) + self.tool_name_to_mcp_server_name_mapping[ + base_tool_name + ] = server_prefix + self.tool_name_to_mcp_server_name_mapping[ + prefixed_tool_name + ] = server_prefix registered_count += 1 verbose_logger.debug( @@ -970,7 +970,9 @@ class MCPServerManager: # Handle stdio transport if transport == MCPTransport.stdio: - resolved_env = stdio_env if stdio_env is not None else dict(server.env or {}) + resolved_env = ( + stdio_env if stdio_env is not None else dict(server.env or {}) + ) # Ensure npm-based STDIO MCP servers have a writable cache dir. # In containers the default (~/.npm or /app/.npm) may not exist @@ -2355,7 +2357,9 @@ class MCPServerManager: prisma_client = get_prisma_client_or_throw( "Database not connected. Connect a database to your proxy" ) - db_mcp_servers = await get_all_mcp_servers(prisma_client, approval_status="active") + db_mcp_servers = await get_all_mcp_servers( + prisma_client, approval_status="active" + ) verbose_logger.info(f"Found {len(db_mcp_servers)} MCP servers in database") previous_registry = self.registry diff --git a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py index 0de381ee1d..84a2e94467 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py +++ b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py @@ -124,11 +124,18 @@ class MCPOAuth2TokenCache(InMemoryCache): # Safely parse expires_in — providers may return null or non-numeric values raw_expires_in = body.get("expires_in") try: - expires_in = int(raw_expires_in) if raw_expires_in is not None else MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL + expires_in = ( + int(raw_expires_in) + if raw_expires_in is not None + else MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL + ) except (TypeError, ValueError): expires_in = MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL - ttl = max(expires_in - MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS, MCP_OAUTH2_TOKEN_CACHE_MIN_TTL) + ttl = max( + expires_in - MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS, + MCP_OAUTH2_TOKEN_CACHE_MIN_TTL, + ) verbose_logger.info( "Fetched OAuth2 token for MCP server %s (expires in %ds)", diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index 5f6cb87b26..4b4818892b 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -71,6 +71,7 @@ def load_openapi_spec(filepath: str) -> Dict[str, Any]: raise return asyncio.run(load_openapi_spec_async(filepath)) + async def load_openapi_spec_async(filepath: str) -> Dict[str, Any]: if filepath.startswith("http://") or filepath.startswith("https://"): client = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) @@ -92,26 +93,55 @@ def get_base_url(spec: Dict[str, Any], spec_path: Optional[str] = None) -> str: """Extract base URL from OpenAPI spec.""" # OpenAPI 3.x if "servers" in spec and spec["servers"]: - return spec["servers"][0]["url"] + server_url = spec["servers"][0]["url"] + + # If the server URL is relative (starts with /), derive base from spec_path + if server_url.startswith("/") and spec_path: + if spec_path.startswith("http://") or spec_path.startswith("https://"): + # Extract base URL from spec_path (e.g., https://petstore3.swagger.io/api/v3/openapi.json) + # Combine domain with the relative server URL + from urllib.parse import urlparse + + parsed = urlparse(spec_path) + base_domain = f"{parsed.scheme}://{parsed.netloc}" + full_base_url = base_domain + server_url + verbose_logger.info( + f"OpenAPI spec has relative server URL '{server_url}'. " + f"Deriving base from spec_path: {full_base_url}" + ) + return full_base_url + + return server_url # OpenAPI 2.x (Swagger) elif "host" in spec: scheme = spec.get("schemes", ["https"])[0] base_path = spec.get("basePath", "") return f"{scheme}://{spec['host']}{base_path}" - + # Fallback: derive base URL from spec_path if it's a URL - if spec_path and (spec_path.startswith("http://") or spec_path.startswith("https://")): - for suffix in ["/openapi.json", "/openapi.yaml", "/swagger.json", "/swagger.yaml"]: + if spec_path and ( + spec_path.startswith("http://") or spec_path.startswith("https://") + ): + for suffix in [ + "/openapi.json", + "/openapi.yaml", + "/swagger.json", + "/swagger.yaml", + ]: if spec_path.endswith(suffix): - base_url = spec_path[:-len(suffix)] - verbose_logger.info(f"No server info in OpenAPI spec. Using derived base URL: {base_url}") + base_url = spec_path[: -len(suffix)] + verbose_logger.info( + f"No server info in OpenAPI spec. Using derived base URL: {base_url}" + ) return base_url - + if spec_path.split("/")[-1].endswith((".json", ".yaml", ".yml")): base_url = "/".join(spec_path.split("/")[:-1]) - verbose_logger.info(f"No server info in OpenAPI spec. Using derived base URL: {base_url}") + verbose_logger.info( + f"No server info in OpenAPI spec. Using derived base URL: {base_url}" + ) return base_url - + return "" @@ -165,7 +195,9 @@ def resolve_operation_params( path_level = _resolve_param_list(path_item.get("parameters", []), component_params) op_level = _resolve_param_list(operation.get("parameters", []), component_params) op_keys = {(p["name"], p.get("in")) for p in op_level} - merged = [p for p in path_level if (p["name"], p.get("in")) not in op_keys] + op_level + merged = [ + p for p in path_level if (p["name"], p.get("in")) not in op_keys + ] + op_level result = dict(operation) result["parameters"] = merged return result @@ -350,7 +382,9 @@ def create_tool_function( url, params=params, json=json_body, headers=effective_headers ) elif original_method == "delete": - response = await client.delete(url, params=params, headers=effective_headers) + response = await client.delete( + url, params=params, headers=effective_headers + ) elif original_method == "patch": response = await client.patch( url, params=params, json=json_body, headers=effective_headers diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index f10263ba57..948d16ceff 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -119,7 +119,9 @@ if MCP_AVAILABLE: prisma_client = get_prisma_client_or_throw( "Database not connected. Connect a database to use OAuth2 MCP tools." ) - cred = await get_user_oauth_credential(prisma_client, user_id, server_id) + cred = await get_user_oauth_credential( + prisma_client, user_id, server_id + ) if cred and cred.get("access_token"): if is_oauth_credential_expired(cred): verbose_logger.debug( @@ -192,7 +194,9 @@ if MCP_AVAILABLE: if c.get("access_token") and c.get("server_id") } except Exception: - verbose_logger.debug("Failed to bulk-fetch OAuth credentials", exc_info=True) + verbose_logger.debug( + "Failed to bulk-fetch OAuth credentials", exc_info=True + ) return {} def _create_tool_response_objects(tools, server_mcp_info): @@ -411,10 +415,11 @@ if MCP_AVAILABLE: ) allowed_server_ids_set.update(servers) - allowed_server_ids, _ip_blocked_count = ( - global_mcp_server_manager.filter_server_ids_by_ip_with_info( - list(allowed_server_ids_set), _rest_client_ip - ) + ( + allowed_server_ids, + _ip_blocked_count, + ) = global_mcp_server_manager.filter_server_ids_by_ip_with_info( + list(allowed_server_ids_set), _rest_client_ip ) list_tools_result = [] @@ -428,8 +433,12 @@ if MCP_AVAILABLE: # IP-filter error reporting if the resolved UUID is not in allowed_server_ids. _name_resolved = None if server_id not in allowed_server_ids: - _name_resolved = global_mcp_server_manager.get_mcp_server_by_name(server_id) - if _name_resolved is not None and _name_resolved.server_id in set(allowed_server_ids): + _name_resolved = global_mcp_server_manager.get_mcp_server_by_name( + server_id + ) + if _name_resolved is not None and _name_resolved.server_id in set( + allowed_server_ids + ): server_id = _name_resolved.server_id if server_id not in allowed_server_ids: @@ -477,7 +486,9 @@ if MCP_AVAILABLE: server, mcp_server_auth_headers, mcp_auth_header ) # Single-server request: targeted lookup is more efficient than a bulk fetch. - user_oauth_extra_headers = await _get_user_oauth_extra_headers(server, user_api_key_dict) + user_oauth_extra_headers = await _get_user_oauth_extra_headers( + server, user_api_key_dict + ) try: list_tools_result = await _get_tools_for_single_server( @@ -540,7 +551,9 @@ if MCP_AVAILABLE: server, mcp_server_auth_headers, mcp_auth_header ) user_oauth_extra_headers = await _get_user_oauth_extra_headers( - server, user_api_key_dict, prefetched_creds=prefetched_oauth_creds + server, + user_api_key_dict, + prefetched_creds=prefetched_oauth_creds, ) try: @@ -632,21 +645,24 @@ if MCP_AVAILABLE: tool_arguments = data.get("arguments") proxy_base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) - data, logging_obj = ( - await proxy_base_llm_response_processor.common_processing_pre_call_logic( - request=request, - user_api_key_dict=user_api_key_dict, - proxy_config=proxy_config, - route_type=CallTypes.call_mcp_tool.value, - proxy_logging_obj=proxy_logging_obj, - general_settings=general_settings, - ) + ( + data, + logging_obj, + ) = await proxy_base_llm_response_processor.common_processing_pre_call_logic( + request=request, + user_api_key_dict=user_api_key_dict, + proxy_config=proxy_config, + route_type=CallTypes.call_mcp_tool.value, + proxy_logging_obj=proxy_logging_obj, + general_settings=general_settings, ) # Extract MCP auth headers from request and add to data dict - mcp_auth_header, mcp_server_auth_headers, raw_headers_from_request = ( - _extract_mcp_headers_from_request(request, MCPRequestHandler) - ) + ( + mcp_auth_header, + mcp_server_auth_headers, + raw_headers_from_request, + ) = _extract_mcp_headers_from_request(request, MCPRequestHandler) if mcp_auth_header: data["mcp_auth_header"] = mcp_auth_header if mcp_server_auth_headers: @@ -745,7 +761,9 @@ if MCP_AVAILABLE: client_id: Optional[str] = creds.get("client_id") client_secret: Optional[str] = creds.get("client_secret") scopes_raw = creds.get("scopes") - scopes: Optional[List[str]] = scopes_raw if isinstance(scopes_raw, list) else None + scopes: Optional[List[str]] = ( + scopes_raw if isinstance(scopes_raw, list) else None + ) return client_id, client_secret, scopes async def _execute_with_mcp_client( @@ -848,7 +866,9 @@ if MCP_AVAILABLE: if operation is None: continue - resolved_op = resolve_operation_params(operation, path_item, components) + resolved_op = resolve_operation_params( + operation, path_item, components + ) op_id = operation.get("operationId", f"{method}_{path}") summary = operation.get("summary", "") @@ -857,7 +877,9 @@ if MCP_AVAILABLE: tools.append( { "name": op_id, - "description": description or summary or f"{method.upper()} {path}", + "description": description + or summary + or f"{method.upper()} {path}", "inputSchema": input_schema, } ) diff --git a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py index e5cb6a0098..0bafd7da26 100644 --- a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py +++ b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py @@ -60,10 +60,14 @@ class SemanticMCPToolFilter: all_tools = [] for server_id, server in registry.items(): try: - tools = await global_mcp_server_manager.get_tools_for_server(server_id) + tools = await global_mcp_server_manager.get_tools_for_server( + server_id + ) all_tools.extend(tools) except Exception as e: - verbose_logger.warning(f"Failed to fetch tools from server {server_id}: {e}") + verbose_logger.warning( + f"Failed to fetch tools from server {server_id}: {e}" + ) continue if not all_tools: @@ -71,7 +75,9 @@ class SemanticMCPToolFilter: self.tool_router = None return - verbose_logger.info(f"Fetched {len(all_tools)} tools from {len(registry)} MCP servers") + verbose_logger.info( + f"Fetched {len(all_tools)} tools from {len(registry)} MCP servers" + ) self._build_router(all_tools) except Exception as e: @@ -83,7 +89,7 @@ class SemanticMCPToolFilter: """Extract name and description from MCP tool or OpenAI function dict.""" name: str description: str - + if isinstance(tool, dict): # OpenAI function format name = tool.get("name", "") @@ -92,7 +98,7 @@ class SemanticMCPToolFilter: # MCPTool object name = str(tool.name) description = str(tool.description) if tool.description else str(tool.name) - + return name, description def _build_router(self, tools: List) -> None: @@ -136,9 +142,7 @@ class SemanticMCPToolFilter: auto_sync="local", ) - verbose_logger.info( - f"Built semantic router with {len(routes)} tools" - ) + verbose_logger.info(f"Built semantic router with {len(routes)} tools") except Exception as e: verbose_logger.error(f"Failed to build semantic router: {e}") @@ -165,16 +169,18 @@ class SemanticMCPToolFilter: # Early returns for cases where we can't/shouldn't filter if not self.enabled: return available_tools - + if not available_tools: return available_tools - + if not query or not query.strip(): return available_tools # Router should be built on startup - if not, something went wrong if self.tool_router is None: - verbose_logger.warning("Router not initialized - was build_router_from_mcp_registry() called on startup?") + verbose_logger.warning( + "Router not initialized - was build_router_from_mcp_registry() called on startup?" + ) return available_tools # Run semantic filtering @@ -182,10 +188,10 @@ class SemanticMCPToolFilter: limit = top_k or self.top_k matches = self.tool_router(text=query, limit=limit) matched_tool_names = self._extract_tool_names_from_matches(matches) - + if not matched_tool_names: return available_tools - + return self._get_tools_by_names(matched_tool_names, available_tools) except Exception as e: @@ -196,15 +202,15 @@ class SemanticMCPToolFilter: """Extract tool names from semantic router match results.""" if not matches: return [] - + # Handle single match if hasattr(matches, "name") and matches.name: return [matches.name] - + # Handle list of matches if isinstance(matches, list): return [m.name for m in matches if hasattr(m, "name") and m.name] - + return [] def _get_tools_by_names( @@ -217,7 +223,7 @@ class SemanticMCPToolFilter: tool_name, _ = self._extract_tool_info(tool) if tool_name in tool_names: matched_tools.append(tool) - + # Reorder to match semantic router's ordering tool_map = {self._extract_tool_info(t)[0]: t for t in matched_tools} return [tool_map[name] for name in tool_names if name in tool_map] diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index d6d44042ff..da5f18d82c 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -81,6 +81,7 @@ def _write_byok_cred_cache( _byok_cred_cache.clear() _byok_cred_cache[(user_id, server_id)] = (credential, time.monotonic()) + # Check if MCP is available # "mcp" requires python 3.10 or higher, but several litellm users use python 3.8 # We're making this conditional import to avoid breaking users who use python 3.8. @@ -182,7 +183,7 @@ if MCP_AVAILABLE: session_manager = StreamableHTTPSessionManager( app=server, event_store=None, - json_response=False, # enables SSE streaming + json_response=False, # enables SSE streaming stateless=True, ) @@ -341,9 +342,9 @@ if MCP_AVAILABLE: host_progress_callback = None try: host_ctx = server.request_context - if host_ctx and hasattr(host_ctx, 'meta') and host_ctx.meta: - host_token = getattr(host_ctx.meta, 'progressToken', None) - if host_token and hasattr(host_ctx, 'session') and host_ctx.session: + if host_ctx and hasattr(host_ctx, "meta") and host_ctx.meta: + host_token = getattr(host_ctx.meta, "progressToken", None) + if host_token and hasattr(host_ctx, "session") and host_ctx.session: host_session = host_ctx.session async def forward_progress(progress: float, total: float | None): @@ -352,14 +353,20 @@ if MCP_AVAILABLE: await host_session.send_progress_notification( progress_token=host_token, progress=progress, - total=total + total=total, + ) + verbose_logger.debug( + f"Forwarded progress {progress}/{total} to Host" ) - verbose_logger.debug(f"Forwarded progress {progress}/{total} to Host") except Exception as e: - verbose_logger.error(f"Failed to forward progress to Host: {e}") + verbose_logger.error( + f"Failed to forward progress to Host: {e}" + ) host_progress_callback = forward_progress - verbose_logger.debug(f"Host progressToken captured: {host_token[:8]}...") + verbose_logger.debug( + f"Host progressToken captured: {host_token[:8]}..." + ) except Exception as e: verbose_logger.warning(f"Could not capture host progress context: {e}") try: @@ -711,6 +718,7 @@ if MCP_AVAILABLE: Checks both the full tool name and unprefixed version (without server prefix). This allows users to configure simple tool names regardless of prefixing. + Comparison is case-insensitive to handle OpenAPI operationIds that may be in camelCase. Args: tool_name: The tool name to check (may be prefixed like "server-tool_name") @@ -723,13 +731,15 @@ if MCP_AVAILABLE: split_server_prefix_from_name, ) - # Check if the full name is in the list - if tool_name in filter_list: + # Normalize filter list to lowercase for case-insensitive comparison + filter_list_lower = [f.lower() for f in filter_list] + + if tool_name.lower() in filter_list_lower: return True - # Check if the unprefixed name is in the list + # Check if the unprefixed name is in the list (case-insensitive) unprefixed_name, _ = split_server_prefix_from_name(tool_name) - return unprefixed_name in filter_list + return unprefixed_name.lower() in filter_list_lower def filter_tools_by_allowed_tools( tools: List[MCPTool], @@ -831,18 +841,18 @@ if MCP_AVAILABLE: ) allowed_mcp_server_ids = ( - await global_mcp_server_manager.get_allowed_mcp_servers( - user_api_key_auth - ) + await global_mcp_server_manager.get_allowed_mcp_servers(user_api_key_auth) ) - allowed_mcp_server_ids, _ip_blocked = ( - global_mcp_server_manager.filter_server_ids_by_ip_with_info( - allowed_mcp_server_ids, client_ip - ) + ( + allowed_mcp_server_ids, + _ip_blocked, + ) = global_mcp_server_manager.filter_server_ids_by_ip_with_info( + allowed_mcp_server_ids, client_ip ) verbose_logger.debug( "MCP IP filter: client_ip=%s, allowed_server_ids=%s", - client_ip, allowed_mcp_server_ids, + client_ip, + allowed_mcp_server_ids, ) if _ip_blocked > 0: verbose_logger.debug( @@ -867,7 +877,6 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, allowed_mcp_servers=allowed_mcp_servers, ) - return allowed_mcp_servers @@ -906,7 +915,9 @@ if MCP_AVAILABLE: prisma_client = get_prisma_client_or_throw( "Database not connected. Connect a database to use OAuth2 MCP tools." ) - cred = await get_user_oauth_credential(prisma_client, user_id, server_id) + cred = await get_user_oauth_credential( + prisma_client, user_id, server_id + ) if cred and cred.get("access_token"): if is_oauth_credential_expired(cred): verbose_logger.debug( @@ -929,7 +940,9 @@ if MCP_AVAILABLE: Returns a dict keyed by server_id to avoid N+1 queries in asyncio.gather loops. """ - user_id = getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None + user_id = ( + getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None + ) if not user_id: return {} try: @@ -1058,7 +1071,6 @@ if MCP_AVAILABLE: # Attach user identifiers using the standard helper if user_api_key_auth is not None: - LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( data=list_tools_request_data, user_api_key_dict=user_api_key_auth, @@ -1123,7 +1135,9 @@ if MCP_AVAILABLE: # If no OAuth2 token came from request headers, fall back to pre-fetched creds if extra_headers is None and server.auth_type == MCPAuth.oauth2: extra_headers = await _get_user_oauth_extra_headers_from_db( - server, user_api_key_auth, prefetched_creds=_prefetched_oauth_creds + server, + user_api_key_auth, + prefetched_creds=_prefetched_oauth_creds, ) try: @@ -1253,7 +1267,6 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, ) - # Get prompts from each allowed server all_prompts = [] for server in allowed_mcp_servers: @@ -1312,7 +1325,6 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, ) - all_resources: List[Resource] = [] for server in allowed_mcp_servers: if server is None: @@ -1368,7 +1380,6 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, ) - all_resource_templates: List[ResourceTemplate] = [] for server in allowed_mcp_servers: if server is None: @@ -1866,7 +1877,9 @@ if MCP_AVAILABLE: # configured auth_type so the generator doesn't need to know the prefix. auth_header_value: Optional[str] = None if mcp_auth_header: - server_auth_type = getattr(mcp_server, "auth_type", None) if mcp_server else None + server_auth_type = ( + getattr(mcp_server, "auth_type", None) if mcp_server else None + ) if server_auth_type == MCPAuth.api_key: auth_header_value = f"ApiKey {mcp_auth_header}" elif server_auth_type == MCPAuth.basic: @@ -1902,12 +1915,8 @@ if MCP_AVAILABLE: # Deprecated: Local MCP Server Tool ######################################################### else: - local_content = await _handle_local_mcp_tool( - original_tool_name, arguments - ) - response = CallToolResult( - content=cast(Any, local_content), isError=False - ) + local_content = await _handle_local_mcp_tool(original_tool_name, arguments) + response = CallToolResult(content=cast(Any, local_content), isError=False) return response @@ -2028,7 +2037,6 @@ if MCP_AVAILABLE: detail="User not allowed to get this prompt.", ) - # Extract server name from prefixed prompt name original_prompt_name, server_name = split_server_prefix_from_name(name) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 6e2e6c7e25..1174740948 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1,40 +1,59 @@ import enum import json from datetime import datetime -from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Literal, - Optional, Union) +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union import httpx -from pydantic import (BaseModel, ConfigDict, Field, Json, field_validator, - model_validator) +from pydantic import ( + BaseModel, + ConfigDict, + Field, + Json, + field_validator, + model_validator, +) from typing_extensions import Required, TypedDict from litellm._uuid import uuid from litellm.types.integrations.slack_alerting import AlertType -from litellm.types.llms.openai import (AllMessageValues, OpenAIFileObject, - ResponsesAPIResponse) -from litellm.types.mcp import (MCPAuthType, MCPCredentials, MCPTransport, - MCPTransportType) +from litellm.types.llms.openai import ( + AllMessageValues, + OpenAIFileObject, + ResponsesAPIResponse, +) +from litellm.types.mcp import ( + MCPAuthType, + MCPCredentials, + MCPTransport, + MCPTransportType, +) from litellm.types.mcp_server.mcp_server_manager import MCPInfo from litellm.types.router import RouterErrors, UpdateRouterConfig from litellm.types.secret_managers.main import KeyManagementSystem -from litellm.types.utils import (CallTypes, CostBreakdown, EmbeddingResponse, - GenericBudgetConfigType, ImageResponse, - LiteLLMBatch, LiteLLMFineTuningJob, - LiteLLMPydanticObjectBase, ModelResponse, - ProviderField, StandardCallbackDynamicParams, - StandardLoggingGuardrailInformation, - StandardLoggingMCPToolCall, - StandardLoggingModelInformation, - StandardLoggingPayloadErrorInformation, - StandardLoggingPayloadStatus, - StandardLoggingVectorStoreRequest, - StandardPassThroughResponseObject, - TextCompletionResponse) +from litellm.types.utils import ( + CallTypes, + CostBreakdown, + EmbeddingResponse, + GenericBudgetConfigType, + ImageResponse, + LiteLLMBatch, + LiteLLMFineTuningJob, + LiteLLMPydanticObjectBase, + ModelResponse, + ProviderField, + StandardCallbackDynamicParams, + StandardLoggingGuardrailInformation, + StandardLoggingMCPToolCall, + StandardLoggingModelInformation, + StandardLoggingPayloadErrorInformation, + StandardLoggingPayloadStatus, + StandardLoggingVectorStoreRequest, + StandardPassThroughResponseObject, + TextCompletionResponse, +) from litellm.types.videos.main import VideoObject -from .types_utils.utils import (get_instance_fn, - validate_custom_validate_return_type) +from .types_utils.utils import get_instance_fn, validate_custom_validate_return_type if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -1111,13 +1130,16 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase): # BYOM submission fields — set by the endpoint, not by the caller. # Any caller-provided values are silently overridden before persistence. approval_status: Optional[str] = Field( - None, description="Server-managed: set by the endpoint; caller values are overridden." + None, + description="Server-managed: set by the endpoint; caller values are overridden.", ) submitted_by: Optional[str] = Field( - None, description="Server-managed: set by the endpoint; caller values are overridden." + None, + description="Server-managed: set by the endpoint; caller values are overridden.", ) submitted_at: Optional[datetime] = Field( - None, description="Server-managed: set by the endpoint; caller values are overridden." + None, + description="Server-managed: set by the endpoint; caller values are overridden.", ) @model_validator(mode="before") @@ -2442,7 +2464,9 @@ class UserAPIKeyAuth( user_max_budget: Optional[float] = None request_route: Optional[str] = None user: Optional[Any] = None # Expanded user object when expand=user is used - created_by_user: Optional[Any] = None # Expanded created_by user when expand=user is used + created_by_user: Optional[ + Any + ] = None # Expanded created_by user when expand=user is used end_user_object_permission: Optional[LiteLLM_ObjectPermissionTable] = None model_config = ConfigDict(arbitrary_types_allowed=True) @@ -2486,8 +2510,7 @@ class UserAPIKeyAuth( This is used to track number of requests/spend for health check calls. """ - from litellm.constants import \ - LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME + from litellm.constants import LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME return cls( api_key=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, @@ -2519,8 +2542,7 @@ class UserAPIKeyAuth( This is used to track actions performed by automated system jobs. """ - from litellm.constants import \ - LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME + from litellm.constants import LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME return cls( api_key=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, @@ -2926,8 +2948,7 @@ class LiteLLM_AuditLogs(LiteLLMPydanticObjectBase): @model_validator(mode="after") def mask_api_keys(self): - from litellm.litellm_core_utils.sensitive_data_masker import \ - SensitiveDataMasker + from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker masker = SensitiveDataMasker(sensitive_patterns={"key"}) diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 63e0dad332..9b18c76629 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -39,8 +39,7 @@ def _jsonrpc_error( def _get_agent(agent_id: str): """Look up an agent by ID or name. Returns None if not found.""" - from litellm.proxy.agent_endpoints.agent_registry import \ - global_agent_registry + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry agent = global_agent_registry.get_agent_by_id(agent_id=agent_id) if agent is None: @@ -137,8 +136,9 @@ async def _handle_stream_message( and request_data is not None and proxy_logging_obj is not None ): - from litellm.proxy.common_request_processing import \ - ProxyBaseLLMRequestProcessing + from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + ) def _ndjson_chunk(chunk: Any) -> str: if hasattr(chunk, "model_dump"): @@ -238,8 +238,9 @@ async def get_agent_card( The URL in the agent card is rewritten to point to the LiteLLM proxy, so all subsequent A2A calls go through LiteLLM for logging and cost tracking. """ - from litellm.proxy.agent_endpoints.auth.agent_permission_handler import \ - AgentRequestHandler + from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( + AgentRequestHandler, + ) try: agent = _get_agent(agent_id) @@ -303,10 +304,15 @@ async def invoke_agent_a2a( # noqa: PLR0915 """ from litellm.a2a_protocol import asend_message from litellm.a2a_protocol.main import A2A_SDK_AVAILABLE - from litellm.proxy.agent_endpoints.auth.agent_permission_handler import \ - AgentRequestHandler - from litellm.proxy.proxy_server import (general_settings, proxy_config, - proxy_logging_obj, version) + from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( + AgentRequestHandler, + ) + from litellm.proxy.proxy_server import ( + general_settings, + proxy_config, + proxy_logging_obj, + version, + ) body = {} try: @@ -393,8 +399,9 @@ async def invoke_agent_a2a( # noqa: PLR0915 ) # Add litellm data (user_api_key, user_id, team_id, etc.) - from litellm.proxy.common_request_processing import \ - ProxyBaseLLMRequestProcessing + from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + ) processor = ProxyBaseLLMRequestProcessing(data=body) data, logging_obj = await processor.common_processing_pre_call_logic( diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py index cb277d44ee..8f95141499 100644 --- a/litellm/proxy/agent_endpoints/a2a_routing.py +++ b/litellm/proxy/agent_endpoints/a2a_routing.py @@ -14,7 +14,7 @@ from litellm._logging import verbose_proxy_logger def route_a2a_agent_request(data: dict, route_type: str) -> Optional[Any]: """ Route A2A agent requests directly to litellm with injected API base. - + Returns None if not an A2A request (allows normal routing to continue). """ # Import here to avoid circular imports @@ -23,31 +23,31 @@ def route_a2a_agent_request(data: dict, route_type: str) -> Optional[Any]: ROUTE_ENDPOINT_MAPPING, ProxyModelNotFoundError, ) - + model_name = data.get("model", "") - + # Check if this is an A2A agent request if not isinstance(model_name, str) or not model_name.startswith("a2a/"): return None - + # Extract agent name (e.g., "a2a/my-agent" -> "my-agent") agent_name = model_name[4:] - + # Look up agent in registry agent = global_agent_registry.get_agent_by_name(agent_name) if agent is None: verbose_proxy_logger.error(f"[A2A] Agent '{agent_name}' not found in registry") route_name = ROUTE_ENDPOINT_MAPPING.get(route_type, route_type) raise ProxyModelNotFoundError(route=route_name, model_name=model_name) - + # Get API base URL from agent config if not agent.agent_card_params or "url" not in agent.agent_card_params: verbose_proxy_logger.error(f"[A2A] Agent '{agent_name}' has no URL configured") route_name = ROUTE_ENDPOINT_MAPPING.get(route_type, route_type) raise ProxyModelNotFoundError(route=route_name, model_name=model_name) - + # Inject API base and route to litellm data["api_base"] = agent.agent_card_params["url"] verbose_proxy_logger.debug(f"[A2A] Routing {model_name} to {data['api_base']}") - + return getattr(litellm, f"{route_type}")(**data) diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index ce6b1055ee..436de8b0af 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -5,8 +5,9 @@ from typing import Any, Dict, List, Optional import litellm from litellm.litellm_core_utils.safe_json_dumps import safe_dumps -from litellm.proxy.management_helpers.object_permission_utils import \ - handle_update_object_permission_common +from litellm.proxy.management_helpers.object_permission_utils import ( + handle_update_object_permission_common, +) from litellm.proxy.utils import PrismaClient from litellm.types.agents import AgentConfig, AgentResponse, PatchAgentRequest @@ -151,7 +152,12 @@ class AgentRegistry: if object_permission_id is not None: create_data["object_permission_id"] = object_permission_id - for rate_field in ("tpm_limit", "rpm_limit", "session_tpm_limit", "session_rpm_limit"): + for rate_field in ( + "tpm_limit", + "rpm_limit", + "session_tpm_limit", + "session_rpm_limit", + ): _val = agent.get(rate_field) if _val is not None: create_data[rate_field] = _val @@ -165,9 +171,13 @@ class AgentRegistry: created_agent_dict = created_agent.model_dump() if created_agent.object_permission is not None: try: - created_agent_dict["object_permission"] = created_agent.object_permission.model_dump() + created_agent_dict[ + "object_permission" + ] = created_agent.object_permission.model_dump() except Exception: - created_agent_dict["object_permission"] = created_agent.object_permission.dict() + created_agent_dict[ + "object_permission" + ] = created_agent.object_permission.dict() return AgentResponse(**created_agent_dict) # type: ignore except Exception as e: raise Exception(f"Error adding agent to DB: {str(e)}") @@ -208,7 +218,6 @@ class AgentRegistry: The patched agent """ try: - existing_agent = await prisma_client.db.litellm_agentstable.find_unique( where={"agent_id": agent_id} ) @@ -231,7 +240,12 @@ class AgentRegistry: augment_agent.get("agent_card_params") ) - for rate_field in ("tpm_limit", "rpm_limit", "session_tpm_limit", "session_rpm_limit"): + for rate_field in ( + "tpm_limit", + "rpm_limit", + "session_tpm_limit", + "session_rpm_limit", + ): if rate_field in agent: update_data[rate_field] = agent.get(rate_field) if "static_headers" in agent: @@ -249,12 +263,10 @@ class AgentRegistry: existing_object_permission_id = existing_agent.get( "object_permission_id" ) - object_permission_id = ( - await handle_update_object_permission_common( - agent_copy, - existing_object_permission_id, - prisma_client, - ) + object_permission_id = await handle_update_object_permission_common( + agent_copy, + existing_object_permission_id, + prisma_client, ) if object_permission_id is not None: update_data["object_permission_id"] = object_permission_id @@ -271,9 +283,13 @@ class AgentRegistry: patched_agent_dict = patched_agent.model_dump() if patched_agent.object_permission is not None: try: - patched_agent_dict["object_permission"] = patched_agent.object_permission.model_dump() + patched_agent_dict[ + "object_permission" + ] = patched_agent.object_permission.model_dump() except Exception: - patched_agent_dict["object_permission"] = patched_agent.object_permission.dict() + patched_agent_dict[ + "object_permission" + ] = patched_agent.object_permission.dict() return AgentResponse(**patched_agent_dict) # type: ignore except Exception as e: raise Exception(f"Error patching agent in DB: {str(e)}") @@ -330,7 +346,12 @@ class AgentRegistry: "updated_at": datetime.now(timezone.utc), } - for rate_field in ("tpm_limit", "rpm_limit", "session_tpm_limit", "session_rpm_limit"): + for rate_field in ( + "tpm_limit", + "rpm_limit", + "session_tpm_limit", + "session_rpm_limit", + ): _val = agent.get(rate_field) if _val is not None: update_data[rate_field] = _val @@ -345,12 +366,10 @@ class AgentRegistry: else None ) agent_copy = dict(agent) - object_permission_id = ( - await handle_update_object_permission_common( - agent_copy, - existing_object_permission_id, - prisma_client, - ) + object_permission_id = await handle_update_object_permission_common( + agent_copy, + existing_object_permission_id, + prisma_client, ) if object_permission_id is not None: update_data["object_permission_id"] = object_permission_id @@ -365,9 +384,13 @@ class AgentRegistry: updated_agent_dict = updated_agent.model_dump() if updated_agent.object_permission is not None: try: - updated_agent_dict["object_permission"] = updated_agent.object_permission.model_dump() + updated_agent_dict[ + "object_permission" + ] = updated_agent.object_permission.model_dump() except Exception: - updated_agent_dict["object_permission"] = updated_agent.object_permission.dict() + updated_agent_dict[ + "object_permission" + ] = updated_agent.object_permission.dict() return AgentResponse(**updated_agent_dict) # type: ignore except Exception as e: raise Exception(f"Error updating agent in DB: {str(e)}") @@ -391,7 +414,9 @@ class AgentRegistry: # object_permission is eagerly loaded via include above if agent.object_permission is not None: try: - agent_dict["object_permission"] = agent.object_permission.model_dump() + agent_dict[ + "object_permission" + ] = agent.object_permission.model_dump() except Exception: agent_dict["object_permission"] = agent.object_permission.dict() agents.append(agent_dict) diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 646e6d59c3..6e5d4562b5 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -177,9 +177,10 @@ async def get_agents( for agent in returned_agents: if agent.litellm_params is None: agent.litellm_params = {} - agent.litellm_params["is_public"] = ( - litellm.public_agent_groups is not None - and (agent.agent_id in litellm.public_agent_groups) + agent.litellm_params[ + "is_public" + ] = litellm.public_agent_groups is not None and ( + agent.agent_id in litellm.public_agent_groups ) if health_check: @@ -206,18 +207,18 @@ async def get_agents( AGENT_HEALTH_CHECK_GATHER_TIMEOUT_SECONDS, ) health_results = [ - {"agent_id": agent.agent_id, "healthy": False, "error": "Health check timed out"} + { + "agent_id": agent.agent_id, + "healthy": False, + "error": "Health check timed out", + } for agent in agents_with_url ] healthy_ids = { - result["agent_id"] - for result in health_results - if result["healthy"] + result["agent_id"] for result in health_results if result["healthy"] } returned_agents = [ - agent - for agent in agents_with_url - if agent.agent_id in healthy_ids + agent for agent in agents_with_url if agent.agent_id in healthy_ids ] + agents_without_url return returned_agents @@ -236,8 +237,9 @@ async def get_agents( #### CRUD ENDPOINTS FOR AGENTS #### -from litellm.proxy.agent_endpoints.agent_registry import \ - global_agent_registry as AGENT_REGISTRY +from litellm.proxy.agent_endpoints.agent_registry import ( + global_agent_registry as AGENT_REGISTRY, +) @router.post( @@ -376,13 +378,13 @@ async def get_agent_by_id( agent_dict = agent_row.model_dump() if agent_row.object_permission is not None: try: - agent_dict["object_permission"] = ( - agent_row.object_permission.model_dump() - ) + agent_dict[ + "object_permission" + ] = agent_row.object_permission.model_dump() except Exception: - agent_dict["object_permission"] = ( - agent_row.object_permission.dict() - ) + agent_dict[ + "object_permission" + ] = agent_row.object_permission.dict() agent = AgentResponse(**agent_dict) # type: ignore else: # Agent found in memory — refresh spend from DB @@ -698,8 +700,9 @@ async def make_agent_public( try: # Update the public model groups import litellm - from litellm.proxy.agent_endpoints.agent_registry import \ - global_agent_registry as AGENT_REGISTRY + from litellm.proxy.agent_endpoints.agent_registry import ( + global_agent_registry as AGENT_REGISTRY, + ) from litellm.proxy.proxy_server import proxy_config # Check if user has admin permissions @@ -814,8 +817,9 @@ async def make_agents_public( try: # Update the public model groups import litellm - from litellm.proxy.agent_endpoints.agent_registry import \ - global_agent_registry as AGENT_REGISTRY + from litellm.proxy.agent_endpoints.agent_registry import ( + global_agent_registry as AGENT_REGISTRY, + ) from litellm.proxy.proxy_server import proxy_config # Load existing config diff --git a/litellm/proxy/agent_endpoints/model_list_helpers.py b/litellm/proxy/agent_endpoints/model_list_helpers.py index c640300bb8..37308b92f7 100644 --- a/litellm/proxy/agent_endpoints/model_list_helpers.py +++ b/litellm/proxy/agent_endpoints/model_list_helpers.py @@ -18,7 +18,7 @@ async def append_agents_to_model_group( ) -> List[ModelGroupInfoProxy]: """ Append A2A agents to model groups list for UI display. - + Converts agents to model format with "a2a/" naming so they appear in playground and work with LiteLLM routing. """ @@ -31,7 +31,7 @@ async def append_agents_to_model_group( allowed_agent_ids = await AgentRequestHandler.get_allowed_agents( user_api_key_auth=user_api_key_dict ) - + for agent_id in allowed_agent_ids: agent = global_agent_registry.get_agent_by_id(agent_id) if agent is not None: @@ -43,10 +43,8 @@ async def append_agents_to_model_group( ) ) except Exception as e: - verbose_proxy_logger.debug( - f"Error appending agents to model_group/info: {e}" - ) - + verbose_proxy_logger.debug(f"Error appending agents to model_group/info: {e}") + return model_groups @@ -56,7 +54,7 @@ async def append_agents_to_model_info( ) -> List[dict]: """ Append A2A agents to model info list for UI display. - + Converts agents to model format with "a2a/" naming so they appear in models page and work with LiteLLM routing. """ @@ -69,28 +67,28 @@ async def append_agents_to_model_info( allowed_agent_ids = await AgentRequestHandler.get_allowed_agents( user_api_key_auth=user_api_key_dict ) - + for agent_id in allowed_agent_ids: agent = global_agent_registry.get_agent_by_id(agent_id) if agent is not None: - models.append({ - "model_name": f"a2a/{agent.agent_name}", - "litellm_params": { - "model": f"a2a/{agent.agent_name}", - "custom_llm_provider": "a2a", - }, - "model_info": { - "id": agent.agent_id, - "mode": "chat", - "db_model": True, - "created_by": agent.created_by, - "created_at": agent.created_at, - "updated_at": agent.updated_at, - }, - }) + models.append( + { + "model_name": f"a2a/{agent.agent_name}", + "litellm_params": { + "model": f"a2a/{agent.agent_name}", + "custom_llm_provider": "a2a", + }, + "model_info": { + "id": agent.agent_id, + "mode": "chat", + "db_model": True, + "created_by": agent.created_by, + "created_at": agent.created_at, + "updated_at": agent.updated_at, + }, + } + ) except Exception as e: - verbose_proxy_logger.debug( - f"Error appending agents to v2/model/info: {e}" - ) - + verbose_proxy_logger.debug(f"Error appending agents to v2/model/info: {e}") + return models diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 5b23b47923..69d69354fd 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -30,7 +30,7 @@ async def anthropic_response( # noqa: PLR0915 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ - Use `{PROXY_BASE_URL}/anthropic/v1/messages` instead - [Docs](https://docs.litellm.ai/docs/anthropic_completion). + Use `{PROXY_BASE_URL}/anthropic/v1/messages` instead - [Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion). This was a BETA endpoint that calls 100+ LLMs in the anthropic format. """ @@ -257,7 +257,7 @@ async def event_logging_batch( ): """ Stubbed endpoint for Anthropic event logging batch requests. - + This endpoint accepts event logging requests but does nothing with them. It exists to prevent 404 errors from Claude Code clients that send telemetry. """ diff --git a/litellm/proxy/anthropic_endpoints/skills_endpoints.py b/litellm/proxy/anthropic_endpoints/skills_endpoints.py index 69509e1f53..cd19e7731f 100644 --- a/litellm/proxy/anthropic_endpoints/skills_endpoints.py +++ b/litellm/proxy/anthropic_endpoints/skills_endpoints.py @@ -82,7 +82,7 @@ async def create_skill( # Read form data and convert UploadFile objects to file data tuples form_data = await get_form_data(request) data = await convert_upload_files_to_file_data(form_data) - + # Extract model for routing (header > query > body) model = ( data.get("model") @@ -91,10 +91,10 @@ async def create_skill( ) if model: data["model"] = model - + if "custom_llm_provider" not in data: data["custom_llm_provider"] = custom_llm_provider - + # Process request using ProxyBaseLLMRequestProcessing processor = ProxyBaseLLMRequestProcessing(data=data) try: @@ -181,7 +181,7 @@ async def list_skills( # Read request body body = await request.body() data = orjson.loads(body) if body else {} - + # Use query params if not in body if "limit" not in data and limit is not None: data["limit"] = limit @@ -189,7 +189,7 @@ async def list_skills( data["after_id"] = after_id if "before_id" not in data and before_id is not None: data["before_id"] = before_id - + # Extract model for routing (header > query > body) model = ( data.get("model") @@ -198,11 +198,11 @@ async def list_skills( ) if model: data["model"] = model - + # Set custom_llm_provider: body > query param > default if "custom_llm_provider" not in data: data["custom_llm_provider"] = custom_llm_provider - + # Process request using ProxyBaseLLMRequestProcessing processor = ProxyBaseLLMRequestProcessing(data=data) try: @@ -287,10 +287,10 @@ async def get_skill( # Read request body body = await request.body() data = orjson.loads(body) if body else {} - + # Set skill_id from path parameter data["skill_id"] = skill_id - + # Extract model for routing (header > query > body) model = ( data.get("model") @@ -299,11 +299,11 @@ async def get_skill( ) if model: data["model"] = model - + # Set custom_llm_provider: body > query param > default if "custom_llm_provider" not in data: data["custom_llm_provider"] = custom_llm_provider - + # Process request using ProxyBaseLLMRequestProcessing processor = ProxyBaseLLMRequestProcessing(data=data) try: @@ -390,10 +390,10 @@ async def delete_skill( # Read request body body = await request.body() data = orjson.loads(body) if body else {} - + # Set skill_id from path parameter data["skill_id"] = skill_id - + # Extract model for routing (header > query > body) model = ( data.get("model") @@ -402,11 +402,11 @@ async def delete_skill( ) if model: data["model"] = model - + # Set custom_llm_provider: body > query param > default if "custom_llm_provider" not in data: data["custom_llm_provider"] = custom_llm_provider - + # Process request using ProxyBaseLLMRequestProcessing processor = ProxyBaseLLMRequestProcessing(data=data) try: @@ -435,4 +435,3 @@ async def delete_skill( proxy_logging_obj=proxy_logging_obj, version=version, ) - diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index db794c5ac3..b4c123b361 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -11,8 +11,7 @@ Run checks for: import asyncio import re import time -from typing import (TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, - cast) +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast from fastapi import HTTPException, Request, status from pydantic import BaseModel @@ -21,33 +20,48 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.caching.dual_cache import LimitedSizeOrderedDict -from litellm.constants import (CLI_JWT_EXPIRATION_HOURS, CLI_JWT_TOKEN_NAME, - DEFAULT_ACCESS_GROUP_CACHE_TTL, - DEFAULT_IN_MEMORY_TTL, - DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, - DEFAULT_MAX_RECURSE_DEPTH, - EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE) +from litellm.constants import ( + CLI_JWT_EXPIRATION_HOURS, + CLI_JWT_TOKEN_NAME, + DEFAULT_ACCESS_GROUP_CACHE_TTL, + DEFAULT_IN_MEMORY_TTL, + DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, + DEFAULT_MAX_RECURSE_DEPTH, + EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE, +) from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider -from litellm.proxy._types import (RBAC_ROLES, CallInfo, - LiteLLM_AccessGroupTable, - LiteLLM_BudgetTable, LiteLLM_EndUserTable, - Litellm_EntityType, LiteLLM_JWTAuth, - LiteLLM_ObjectPermissionTable, - LiteLLM_OrganizationMembershipTable, - LiteLLM_OrganizationTable, - LiteLLM_ProjectTableCachedObj, - LiteLLM_TagTable, LiteLLM_TeamMembership, - LiteLLM_TeamTable, - LiteLLM_TeamTableCachedObj, - LiteLLM_UserTable, LiteLLMRoutes, - LitellmUserRoles, NewTeamRequest, - ProxyErrorTypes, ProxyException, - RoleBasedPermissions, SpecialModelNames, - UserAPIKeyAuth) +from litellm.proxy._types import ( + RBAC_ROLES, + CallInfo, + LiteLLM_AccessGroupTable, + LiteLLM_BudgetTable, + LiteLLM_EndUserTable, + Litellm_EntityType, + LiteLLM_JWTAuth, + LiteLLM_ObjectPermissionTable, + LiteLLM_OrganizationMembershipTable, + LiteLLM_OrganizationTable, + LiteLLM_ProjectTableCachedObj, + LiteLLM_TagTable, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LiteLLM_TeamTableCachedObj, + LiteLLM_UserTable, + LiteLLMRoutes, + LitellmUserRoles, + NewTeamRequest, + ProxyErrorTypes, + ProxyException, + RoleBasedPermissions, + SpecialModelNames, + UserAPIKeyAuth, +) from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.guardrails.tool_name_extraction import ( - TOOL_CAPABLE_CALL_TYPES, extract_request_tool_names) + TOOL_CAPABLE_CALL_TYPES, + extract_request_tool_names, +) from litellm.proxy.route_llm_request import route_request from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics from litellm.router import Router @@ -281,8 +295,7 @@ def _guardrail_modification_check( if not _request_metadata.get("guardrails"): return - from litellm.proxy.guardrails.guardrail_helpers import \ - can_modify_guardrails + from litellm.proxy.guardrails.guardrail_helpers import can_modify_guardrails if not can_modify_guardrails(team_object): raise HTTPException( @@ -304,8 +317,9 @@ async def check_tools_allowlist( effective allowlist is read from valid_token.metadata and valid_token.team_metadata. Raises ProxyException with tool_access_denied if a tool is not allowed. """ - from litellm.litellm_core_utils.api_route_to_call_types import \ - get_call_types_for_route + from litellm.litellm_core_utils.api_route_to_call_types import ( + get_call_types_for_route, + ) if valid_token is None: return @@ -408,10 +422,8 @@ async def common_checks( # noqa: PLR0915 # Require trace id for agent keys when agent has require_trace_id_on_calls_by_agent if valid_token is not None and valid_token.agent_id: - from litellm.proxy.agent_endpoints.agent_registry import \ - global_agent_registry - from litellm.proxy.litellm_pre_call_utils import \ - get_chain_id_from_headers + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers agent = global_agent_registry.get_agent_by_id(agent_id=valid_token.agent_id) if agent is not None: @@ -1962,8 +1974,9 @@ class ExperimentalUIJWTToken: def get_experimental_ui_login_jwt_auth_token(user_info: LiteLLM_UserTable) -> str: from datetime import timedelta - from litellm.proxy.common_utils.encrypt_decrypt_utils import \ - encrypt_value_helper + from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + encrypt_value_helper, + ) if user_info.user_role is None: raise Exception("User role is required for experimental UI login") @@ -2009,8 +2022,9 @@ class ExperimentalUIJWTToken: """ from datetime import timedelta - from litellm.proxy.common_utils.encrypt_decrypt_utils import \ - encrypt_value_helper + from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + encrypt_value_helper, + ) if user_info.user_role is None: raise Exception("User role is required for CLI JWT login") @@ -2049,8 +2063,9 @@ class ExperimentalUIJWTToken: import json from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth - from litellm.proxy.common_utils.encrypt_decrypt_utils import \ - decrypt_value_helper + from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + decrypt_value_helper, + ) decrypted_token = decrypt_value_helper( hashed_token, key="ui_hash_key", exception_type="debug" @@ -2366,10 +2381,8 @@ async def _get_resources_from_access_groups( # Lazy import to avoid circular imports if prisma_client is None or user_api_key_cache is None: from litellm.proxy.proxy_server import prisma_client as _prisma_client - from litellm.proxy.proxy_server import \ - proxy_logging_obj as _proxy_logging_obj - from litellm.proxy.proxy_server import \ - user_api_key_cache as _user_api_key_cache + from litellm.proxy.proxy_server import proxy_logging_obj as _proxy_logging_obj + from litellm.proxy.proxy_server import user_api_key_cache as _user_api_key_cache prisma_client = prisma_client or _prisma_client user_api_key_cache = user_api_key_cache or _user_api_key_cache @@ -3325,8 +3338,7 @@ async def _tag_max_budget_check( BudgetExceededError if any tag is over its max budget. Triggers a budget alert if any tag is over its max budget. """ - from litellm.proxy.common_utils.http_parsing_utils import \ - get_tags_from_request_body + from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body if prisma_client is None: return diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 59bc4190fd..9a24041faa 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -320,13 +320,13 @@ def normalize_request_route(route: str) -> str: This prevents high cardinality in Prometheus metrics by collapsing routes like: - /v1/responses/1234567890 -> /v1/responses/{response_id} - /v1/threads/thread_123 -> /v1/threads/{thread_id} - + Args: route: The request route path - + Returns: Normalized route with dynamic parameters replaced by placeholders - + Examples: >>> normalize_request_route("/v1/responses/abc123") '/v1/responses/{response_id}' @@ -339,58 +339,90 @@ def normalize_request_route(route: str) -> str: # Format: (regex_pattern, replacement_template) patterns = [ # Responses API - must come before generic patterns - (r'^(/(?:openai/)?v1/responses)/([^/]+)(/input_items)$', r'\1/{response_id}\3'), - (r'^(/(?:openai/)?v1/responses)/([^/]+)(/cancel)$', r'\1/{response_id}\3'), - (r'^(/(?:openai/)?v1/responses)/([^/]+)$', r'\1/{response_id}'), - (r'^(/responses)/([^/]+)(/input_items)$', r'\1/{response_id}\3'), - (r'^(/responses)/([^/]+)(/cancel)$', r'\1/{response_id}\3'), - (r'^(/responses)/([^/]+)$', r'\1/{response_id}'), - + (r"^(/(?:openai/)?v1/responses)/([^/]+)(/input_items)$", r"\1/{response_id}\3"), + (r"^(/(?:openai/)?v1/responses)/([^/]+)(/cancel)$", r"\1/{response_id}\3"), + (r"^(/(?:openai/)?v1/responses)/([^/]+)$", r"\1/{response_id}"), + (r"^(/responses)/([^/]+)(/input_items)$", r"\1/{response_id}\3"), + (r"^(/responses)/([^/]+)(/cancel)$", r"\1/{response_id}\3"), + (r"^(/responses)/([^/]+)$", r"\1/{response_id}"), # Threads API - (r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/steps)/([^/]+)$', r'\1/{thread_id}\3/{run_id}\5/{step_id}'), - (r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/steps)$', r'\1/{thread_id}\3/{run_id}\5'), - (r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/cancel)$', r'\1/{thread_id}\3/{run_id}\5'), - (r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/submit_tool_outputs)$', r'\1/{thread_id}\3/{run_id}\5'), - (r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)$', r'\1/{thread_id}\3/{run_id}'), - (r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)$', r'\1/{thread_id}\3'), - (r'^(/(?:openai/)?v1/threads)/([^/]+)(/messages)/([^/]+)$', r'\1/{thread_id}\3/{message_id}'), - (r'^(/(?:openai/)?v1/threads)/([^/]+)(/messages)$', r'\1/{thread_id}\3'), - (r'^(/(?:openai/)?v1/threads)/([^/]+)$', r'\1/{thread_id}'), - + ( + r"^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/steps)/([^/]+)$", + r"\1/{thread_id}\3/{run_id}\5/{step_id}", + ), + ( + r"^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/steps)$", + r"\1/{thread_id}\3/{run_id}\5", + ), + ( + r"^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/cancel)$", + r"\1/{thread_id}\3/{run_id}\5", + ), + ( + r"^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/submit_tool_outputs)$", + r"\1/{thread_id}\3/{run_id}\5", + ), + ( + r"^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)$", + r"\1/{thread_id}\3/{run_id}", + ), + (r"^(/(?:openai/)?v1/threads)/([^/]+)(/runs)$", r"\1/{thread_id}\3"), + ( + r"^(/(?:openai/)?v1/threads)/([^/]+)(/messages)/([^/]+)$", + r"\1/{thread_id}\3/{message_id}", + ), + (r"^(/(?:openai/)?v1/threads)/([^/]+)(/messages)$", r"\1/{thread_id}\3"), + (r"^(/(?:openai/)?v1/threads)/([^/]+)$", r"\1/{thread_id}"), # Vector Stores API - (r'^(/(?:openai/)?v1/vector_stores)/([^/]+)(/files)/([^/]+)$', r'\1/{vector_store_id}\3/{file_id}'), - (r'^(/(?:openai/)?v1/vector_stores)/([^/]+)(/files)$', r'\1/{vector_store_id}\3'), - (r'^(/(?:openai/)?v1/vector_stores)/([^/]+)(/file_batches)/([^/]+)$', r'\1/{vector_store_id}\3/{batch_id}'), - (r'^(/(?:openai/)?v1/vector_stores)/([^/]+)(/file_batches)$', r'\1/{vector_store_id}\3'), - (r'^(/(?:openai/)?v1/vector_stores)/([^/]+)$', r'\1/{vector_store_id}'), - + ( + r"^(/(?:openai/)?v1/vector_stores)/([^/]+)(/files)/([^/]+)$", + r"\1/{vector_store_id}\3/{file_id}", + ), + ( + r"^(/(?:openai/)?v1/vector_stores)/([^/]+)(/files)$", + r"\1/{vector_store_id}\3", + ), + ( + r"^(/(?:openai/)?v1/vector_stores)/([^/]+)(/file_batches)/([^/]+)$", + r"\1/{vector_store_id}\3/{batch_id}", + ), + ( + r"^(/(?:openai/)?v1/vector_stores)/([^/]+)(/file_batches)$", + r"\1/{vector_store_id}\3", + ), + (r"^(/(?:openai/)?v1/vector_stores)/([^/]+)$", r"\1/{vector_store_id}"), # Assistants API - (r'^(/(?:openai/)?v1/assistants)/([^/]+)$', r'\1/{assistant_id}'), - + (r"^(/(?:openai/)?v1/assistants)/([^/]+)$", r"\1/{assistant_id}"), # Files API - (r'^(/(?:openai/)?v1/files)/([^/]+)(/content)$', r'\1/{file_id}\3'), - (r'^(/(?:openai/)?v1/files)/([^/]+)$', r'\1/{file_id}'), - + (r"^(/(?:openai/)?v1/files)/([^/]+)(/content)$", r"\1/{file_id}\3"), + (r"^(/(?:openai/)?v1/files)/([^/]+)$", r"\1/{file_id}"), # Batches API - (r'^(/(?:openai/)?v1/batches)/([^/]+)(/cancel)$', r'\1/{batch_id}\3'), - (r'^(/(?:openai/)?v1/batches)/([^/]+)$', r'\1/{batch_id}'), - + (r"^(/(?:openai/)?v1/batches)/([^/]+)(/cancel)$", r"\1/{batch_id}\3"), + (r"^(/(?:openai/)?v1/batches)/([^/]+)$", r"\1/{batch_id}"), # Fine-tuning API - (r'^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)(/events)$', r'\1/{fine_tuning_job_id}\3'), - (r'^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)(/cancel)$', r'\1/{fine_tuning_job_id}\3'), - (r'^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)(/checkpoints)$', r'\1/{fine_tuning_job_id}\3'), - (r'^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)$', r'\1/{fine_tuning_job_id}'), - + ( + r"^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)(/events)$", + r"\1/{fine_tuning_job_id}\3", + ), + ( + r"^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)(/cancel)$", + r"\1/{fine_tuning_job_id}\3", + ), + ( + r"^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)(/checkpoints)$", + r"\1/{fine_tuning_job_id}\3", + ), + (r"^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)$", r"\1/{fine_tuning_job_id}"), # Models API - (r'^(/(?:openai/)?v1/models)/([^/]+)$', r'\1/{model}'), + (r"^(/(?:openai/)?v1/models)/([^/]+)$", r"\1/{model}"), ] - + # Apply patterns in order for pattern, replacement in patterns: normalized = re.sub(pattern, replacement, route) if normalized != route: return normalized - + # Return original route if no pattern matched return route @@ -646,6 +678,7 @@ def get_customer_user_header_from_mapping(user_id_mapping) -> Optional[str]: return header_name return None + def _get_customer_id_from_standard_headers( request_headers: Optional[dict], ) -> Optional[str]: @@ -681,7 +714,9 @@ def get_end_user_id_from_request_body( from litellm.proxy.proxy_server import general_settings # Check 1: Standard customer ID headers (always checked, no configuration required) - customer_id = _get_customer_id_from_standard_headers(request_headers=request_headers) + customer_id = _get_customer_id_from_standard_headers( + request_headers=request_headers + ) if customer_id is not None: return customer_id @@ -736,8 +771,7 @@ def get_end_user_id_from_request_body( user_id_from_metadata_field = metadata_dict.get("user_id") if user_id_from_metadata_field is not None: return str(user_id_from_metadata_field) - - + # Check 6: 'safety_identifier' in request body (OpenAI Responses API parameter) # SECURITY NOTE: safety_identifier can be set by any caller in the request body. # Only use this for end-user identification in trusted environments where you control diff --git a/litellm/proxy/auth/ip_address_utils.py b/litellm/proxy/auth/ip_address_utils.py index 651d678533..34fab4849e 100644 --- a/litellm/proxy/auth/ip_address_utils.py +++ b/litellm/proxy/auth/ip_address_utils.py @@ -127,6 +127,7 @@ class IPAddressUtils: from litellm.proxy.proxy_server import ( general_settings as proxy_general_settings, ) + general_settings = proxy_general_settings except ImportError: general_settings = {} diff --git a/litellm/proxy/auth/litellm_license.py b/litellm/proxy/auth/litellm_license.py index b81109b77c..ec2c1eb8e1 100644 --- a/litellm/proxy/auth/litellm_license.py +++ b/litellm/proxy/auth/litellm_license.py @@ -144,7 +144,7 @@ class LicenseCheck: ): return False return total_users > self.airgapped_license_data["max_users"] - + def is_team_count_over_limit(self, team_count: int) -> bool: """ Check if the license is over the limit @@ -152,7 +152,9 @@ class LicenseCheck: if self.airgapped_license_data is None: return False - _max_teams_in_license: Optional[int] = self.airgapped_license_data.get("max_teams") + _max_teams_in_license: Optional[int] = self.airgapped_license_data.get( + "max_teams" + ) if "max_teams" not in self.airgapped_license_data or not isinstance( _max_teams_in_license, int ): @@ -171,7 +173,7 @@ class LicenseCheck: padding_needed = len(license_key) % 4 if padding_needed: license_key += "=" * (4 - padding_needed) - + decoded = base64.b64decode(license_key) message, signature = decoded.split(b".", 1) diff --git a/litellm/proxy/auth/login_utils.py b/litellm/proxy/auth/login_utils.py index c7e22516fe..702f975150 100644 --- a/litellm/proxy/auth/login_utils.py +++ b/litellm/proxy/auth/login_utils.py @@ -258,7 +258,9 @@ async def authenticate_user( # noqa: PLR0915 hash_password = hash_token(token=password) if secrets.compare_digest( password.encode("utf-8"), _password.encode("utf-8") - ) or secrets.compare_digest(hash_password.encode("utf-8"), _password.encode("utf-8")): + ) or secrets.compare_digest( + hash_password.encode("utf-8"), _password.encode("utf-8") + ): if os.getenv("DATABASE_URL") is not None: response = await generate_key_helper_fn( request_type="key", @@ -340,4 +342,3 @@ def create_ui_token_object( disabled_non_admin_personal_key_creation=disabled_non_admin_personal_key_creation, server_root_path=get_server_root_path(), ) - diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index 32f209a763..bf76f99db6 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -80,7 +80,6 @@ async def get_mcp_server_ids( # Make a direct SQL query to get just the mcp_servers try: - result = await prisma_client.db.litellm_objectpermissiontable.find_unique( where={"object_permission_id": user_api_key_dict.object_permission_id}, ) @@ -108,16 +107,27 @@ def get_key_models( """ all_models: List[str] = [] if len(user_api_key_dict.models) > 0: - all_models = user_api_key_dict.models + all_models = list( + user_api_key_dict.models + ) # copy to avoid mutating cached objects if SpecialModelNames.all_team_models.value in all_models: - all_models = user_api_key_dict.team_models + all_models = list( + user_api_key_dict.team_models + ) # copy to avoid mutating cached objects if SpecialModelNames.all_proxy_models.value in all_models: - all_models = proxy_model_list + all_models = list(proxy_model_list) # copy to avoid mutating caller's list + if include_model_access_groups: + all_models.extend(model_access_groups.keys()) all_models = _get_models_from_access_groups( - model_access_groups=model_access_groups, all_models=all_models + model_access_groups=model_access_groups, + all_models=all_models, + include_model_access_groups=include_model_access_groups, ) + # deduplicate while preserving order + all_models = list(dict.fromkeys(all_models)) + verbose_proxy_logger.debug("ALL KEY MODELS - {}".format(len(all_models))) return all_models @@ -141,8 +151,8 @@ def get_team_models( all_models_set.update(team_models) if SpecialModelNames.all_proxy_models.value in all_models_set: all_models_set.update(proxy_model_list) - - all_models = list(all_models_set) + if include_model_access_groups: + all_models_set.update(model_access_groups.keys()) all_models = _get_models_from_access_groups( model_access_groups=model_access_groups, @@ -150,6 +160,9 @@ def get_team_models( include_model_access_groups=include_model_access_groups, ) + # deduplicate while preserving order + all_models = list(dict.fromkeys(all_models)) + verbose_proxy_logger.debug("ALL TEAM MODELS - {}".format(len(all_models))) return all_models @@ -176,6 +189,7 @@ def get_complete_model_list( """ unique_models = [] + def append_unique(models): for model in models: if model not in unique_models: @@ -188,7 +202,7 @@ def get_complete_model_list( else: append_unique(proxy_model_list) if include_model_access_groups: - append_unique(list(model_access_groups.keys())) # TODO: keys order + append_unique(list(model_access_groups.keys())) # TODO: keys order if user_model: append_unique([user_model]) diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 12edb74af3..2c56f5d8bc 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -94,7 +94,7 @@ class RouteChecks: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, - detail=f"Virtual key is not allowed to call this route. Only allowed to call routes: {valid_token.allowed_routes}. Tried to call route: {route}" + detail=f"Virtual key is not allowed to call this route. Only allowed to call routes: {valid_token.allowed_routes}. Tried to call route: {route}", ) @staticmethod @@ -292,7 +292,7 @@ class RouteChecks: if route in LiteLLMRoutes.anthropic_routes.value: return True - + if route in LiteLLMRoutes.google_routes.value: return True @@ -300,7 +300,7 @@ class RouteChecks: route=route, allowed_routes=LiteLLMRoutes.mcp_routes.value ): return True - + if RouteChecks.check_route_access( route=route, allowed_routes=LiteLLMRoutes.agent_routes.value ): diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 70ed7ad3c8..d683015213 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -25,38 +25,53 @@ from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import ( - ExperimentalUIJWTToken, _cache_key_object, _delete_cache_key_object, - _get_user_role, _is_user_proxy_admin, _virtual_key_max_budget_alert_check, - _virtual_key_max_budget_check, _virtual_key_soft_budget_check, - can_key_call_model, common_checks, get_end_user_object, - get_jwt_key_mapping_object, get_key_object, get_project_object, - get_team_object, get_user_object, is_valid_fallback_model) -from litellm.proxy.auth.auth_exception_handler import \ - UserAPIKeyAuthExceptionHandler -from litellm.proxy.auth.auth_utils import (abbreviate_api_key, - get_end_user_id_from_request_body, - get_model_from_request, - get_request_route, - normalize_request_route, - pre_db_read_auth_checks, - route_in_additonal_public_routes) + ExperimentalUIJWTToken, + _cache_key_object, + _delete_cache_key_object, + _get_user_role, + _is_user_proxy_admin, + _virtual_key_max_budget_alert_check, + _virtual_key_max_budget_check, + _virtual_key_soft_budget_check, + can_key_call_model, + common_checks, + get_end_user_object, + get_jwt_key_mapping_object, + get_key_object, + get_project_object, + get_team_object, + get_user_object, + is_valid_fallback_model, +) +from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler +from litellm.proxy.auth.auth_utils import ( + abbreviate_api_key, + get_end_user_id_from_request_body, + get_model_from_request, + get_request_route, + normalize_request_route, + pre_db_read_auth_checks, + route_in_additonal_public_routes, +) from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler from litellm.proxy.auth.oauth2_check import Oauth2Handler from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request from litellm.proxy.auth.route_checks import RouteChecks -from litellm.proxy.common_utils.cache_coordinator import \ - EventDrivenCacheCoordinator +from litellm.proxy.common_utils.cache_coordinator import EventDrivenCacheCoordinator from litellm.proxy.common_utils.http_parsing_utils import ( - _read_request_body, _safe_get_request_headers, - populate_request_with_path_params) + _read_request_body, + _safe_get_request_headers, + populate_request_with_path_params, +) from litellm.proxy.common_utils.realtime_utils import _realtime_request_body from litellm.proxy.utils import PrismaClient, ProxyLogging, normalize_route_for_root_path from litellm.secret_managers.main import get_secret_bool from litellm.types.services import ServiceTypes try: - from litellm_enterprise.proxy.auth.user_api_key_auth import \ - enterprise_custom_auth as _enterprise_custom_auth + from litellm_enterprise.proxy.auth.user_api_key_auth import ( + enterprise_custom_auth as _enterprise_custom_auth, + ) enterprise_custom_auth: Optional[Callable] = _enterprise_custom_auth except ImportError as e: @@ -336,8 +351,9 @@ def get_api_key( Tuple[Optional[str], Optional[str]]: Tuple of the api_key and the passed_in_key """ from litellm.proxy.auth.route_checks import RouteChecks - from litellm.proxy.common_utils.http_parsing_utils import \ - _safe_get_request_query_params + from litellm.proxy.common_utils.http_parsing_utils import ( + _safe_get_request_query_params, + ) api_key = api_key passed_in_key: Optional[str] = None @@ -506,15 +522,20 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 request_data: dict, custom_litellm_key_header: Optional[str] = None, ) -> UserAPIKeyAuth: - from litellm.proxy.proxy_server import (general_settings, jwt_handler, - litellm_proxy_admin_name, - llm_model_list, llm_router, - master_key, - model_max_budget_limiter, - open_telemetry_logger, - prisma_client, proxy_logging_obj, - user_api_key_cache, - user_custom_auth) + from litellm.proxy.proxy_server import ( + general_settings, + jwt_handler, + litellm_proxy_admin_name, + llm_model_list, + llm_router, + master_key, + model_max_budget_limiter, + open_telemetry_logger, + prisma_client, + proxy_logging_obj, + user_api_key_cache, + user_custom_auth, + ) parent_otel_span: Optional[Span] = None start_time = datetime.now() @@ -617,12 +638,18 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 # Only apply OAuth2 M2M authentication to LLM API routes and info routes, not UI/management routes # This allows UI SSO to work separately from API M2M authentication # Note: Info routes are already scoped to the user - if RouteChecks.is_llm_api_route(route=route) or RouteChecks.is_info_route(route=route): + if RouteChecks.is_llm_api_route(route=route) or RouteChecks.is_info_route( + route=route + ): # When both OAuth2 and JWT auth are enabled, use token format to decide: # - JWT tokens (3 dot-separated parts) -> skip OAuth2, fall through to JWT handler # - Opaque tokens -> use OAuth2 handler # This allows JWT for users and OAuth2 for M2M on the same instance - is_jwt_token = jwt_handler.is_jwt(token=api_key) if general_settings.get("enable_jwt_auth", False) is True else False + is_jwt_token = ( + jwt_handler.is_jwt(token=api_key) + if general_settings.get("enable_jwt_auth", False) is True + else False + ) if not is_jwt_token: # return UserAPIKeyAuth object # helper to check if the api_key is a valid oauth2 token @@ -778,8 +805,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 model = get_model_from_request(request_data, route) skip_budget_checks = False if model is not None and llm_router is not None: - from litellm.proxy.auth.auth_checks import \ - _is_model_cost_zero + from litellm.proxy.auth.auth_checks import _is_model_cost_zero skip_budget_checks = _is_model_cost_zero( model=model, llm_router=llm_router @@ -884,9 +910,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 route=route, ) if _end_user_object is not None: - end_user_params["allowed_model_region"] = ( - _end_user_object.allowed_model_region - ) + end_user_params[ + "allowed_model_region" + ] = _end_user_object.allowed_model_region if _end_user_object.litellm_budget_table is not None: _apply_budget_limits_to_end_user_params( end_user_params=end_user_params, @@ -895,8 +921,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 ) elif litellm.max_end_user_budget_id is not None: # End user doesn't exist yet, but apply default budget limits if configured - from litellm.proxy.auth.auth_checks import \ - get_default_end_user_budget + from litellm.proxy.auth.auth_checks import ( + get_default_end_user_budget, + ) default_budget = await get_default_end_user_budget( prisma_client=prisma_client, @@ -1453,9 +1480,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 if _end_user_object is not None: valid_token_dict.update(end_user_params) - valid_token_dict["end_user_object_permission"] = ( - _end_user_object.object_permission - ) + valid_token_dict[ + "end_user_object_permission" + ] = _end_user_object.object_permission # check if token is from litellm-ui, litellm ui makes keys to allow users to login with sso. These keys can only be used for LiteLLM UI functions # sso/login, ui/login, /key functions and /user functions @@ -1677,8 +1704,7 @@ async def _lookup_end_user_and_apply_budget( valid_token=valid_token, end_user_params=end_user_params ) elif litellm.max_end_user_budget_id is not None: - from litellm.proxy.auth.auth_checks import \ - get_default_end_user_budget + from litellm.proxy.auth.auth_checks import get_default_end_user_budget default_budget = await get_default_end_user_budget( prisma_client=prisma_client, @@ -1709,10 +1735,14 @@ async def _run_post_custom_auth_checks( route: str, parent_otel_span: Optional[Span], ) -> UserAPIKeyAuth: - from litellm.proxy.proxy_server import (general_settings, llm_router, - model_max_budget_limiter, - prisma_client, proxy_logging_obj, - user_api_key_cache) + from litellm.proxy.proxy_server import ( + general_settings, + llm_router, + model_max_budget_limiter, + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) # 1. Look up end_user object from DB if end_user_id is set end_user_object = None diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 3fdebd423e..32501fdc54 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -123,11 +123,12 @@ async def create_batch( # noqa: PLR0915 # Apply team-level batch output expiry enforcement team_metadata = user_api_key_dict.team_metadata or {} - enforced_batch_expiry = team_metadata.get( - "enforced_batch_output_expires_after" - ) + enforced_batch_expiry = team_metadata.get("enforced_batch_output_expires_after") if enforced_batch_expiry is not None: - if "anchor" not in enforced_batch_expiry or "seconds" not in enforced_batch_expiry: + if ( + "anchor" not in enforced_batch_expiry + or "seconds" not in enforced_batch_expiry + ): raise HTTPException( status_code=500, detail={ @@ -148,12 +149,12 @@ async def create_batch( # noqa: PLR0915 input_file_id = _create_batch_data.get("input_file_id", None) unified_file_id: Union[str, Literal[False]] = False - + model_from_file_id = None if input_file_id: model_from_file_id = decode_model_from_file_id(input_file_id) unified_file_id = _is_base64_encoded_unified_file_id(input_file_id) - + # SCENARIO 1: File ID is encoded with model info if model_from_file_id is not None and input_file_id: credentials = get_credentials_for_model( @@ -161,20 +162,20 @@ async def create_batch( # noqa: PLR0915 model_id=model_from_file_id, operation_context="batch creation (file created with model)", ) - + original_file_id = get_original_file_id(input_file_id) _create_batch_data["input_file_id"] = original_file_id prepare_data_with_credentials( data=_create_batch_data, # type: ignore credentials=credentials, ) - + # Create batch using model credentials response = await litellm.acreate_batch( custom_llm_provider=credentials["custom_llm_provider"], - **_create_batch_data # type: ignore + **_create_batch_data, # type: ignore ) - + # Encode the batch ID and related file IDs with model information if response and hasattr(response, "id") and response.id: original_batch_id = response.id @@ -184,24 +185,24 @@ async def create_batch( # noqa: PLR0915 id_type="batch", ) response.id = encoded_batch_id - + if hasattr(response, "output_file_id") and response.output_file_id: response.output_file_id = encode_file_id_with_model( file_id=response.output_file_id, model=model_from_file_id ) - + if hasattr(response, "error_file_id") and response.error_file_id: response.error_file_id = encode_file_id_with_model( file_id=response.error_file_id, model=model_from_file_id ) - + verbose_proxy_logger.debug( f"Created batch using model: {model_from_file_id}, " f"original_batch_id: {original_batch_id}, encoded: {encoded_batch_id}" ) - + response.input_file_id = input_file_id - + elif ( litellm.enable_loadbalancing_on_batch_endpoints is True and is_router_model @@ -250,7 +251,7 @@ async def create_batch( # noqa: PLR0915 or request.query_params.get("model") or request.headers.get("x-litellm-model") ) - + # SCENARIO 2 & 3: Model from header/query OR custom_llm_provider fallback if model_param: # SCENARIO 2: Use model-based routing from header/query/body @@ -259,16 +260,16 @@ async def create_batch( # noqa: PLR0915 model_id=model_param, operation_context="batch creation", ) - + prepare_data_with_credentials( data=_create_batch_data, # type: ignore credentials=credentials, ) - + # Create batch using model credentials response = await litellm.acreate_batch( custom_llm_provider=credentials["custom_llm_provider"], - **_create_batch_data # type: ignore + **_create_batch_data, # type: ignore ) encode_batch_response_ids(response, model=model_param) @@ -338,7 +339,7 @@ async def create_batch( # noqa: PLR0915 dependencies=[Depends(user_api_key_auth)], tags=["batch"], ) -async def retrieve_batch( # noqa: PLR0915 +async def retrieve_batch( # noqa: PLR0915 request: Request, fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -395,7 +396,7 @@ async def retrieve_batch( # noqa: PLR0915 # FIX: First, try to read from ManagedObjectTable for consistent state managed_files_obj = proxy_logging_obj.get_proxy_hook("managed_files") from litellm.proxy.proxy_server import prisma_client - + db_batch_object, response = await get_batch_from_database( batch_id=batch_id, unified_batch_id=unified_batch_id, @@ -403,9 +404,14 @@ async def retrieve_batch( # noqa: PLR0915 prisma_client=prisma_client, verbose_proxy_logger=verbose_proxy_logger, ) - + # If batch is in a terminal state, return immediately - if response is not None and response.status in ["completed", "failed", "cancelled", "expired"]: + if response is not None and response.status in [ + "completed", + "failed", + "cancelled", + "expired", + ]: # Call hooks and return response = await proxy_logging_obj.post_call_success_hook( data=data, user_api_key_dict=user_api_key_dict, response=response @@ -415,18 +421,18 @@ async def retrieve_batch( # noqa: PLR0915 # but not input_file_id. Resolve raw provider ID to unified ID. if unified_batch_id: await resolve_input_file_id_to_unified(response, prisma_client) - + asyncio.create_task( proxy_logging_obj.update_request_status( litellm_call_id=data.get("litellm_call_id", ""), status="success" ) ) - + hidden_params = getattr(response, "_hidden_params", {}) or {} model_id = hidden_params.get("model_id", None) or "" cache_key = hidden_params.get("cache_key", None) or "" api_base = hidden_params.get("api_base", None) or "" - + fastapi_response.headers.update( ProxyBaseLLMRequestProcessing.get_custom_headers( user_api_key_dict=user_api_key_dict, @@ -438,9 +444,9 @@ async def retrieve_batch( # noqa: PLR0915 request_data=data, ) ) - + return response - + # If batch is still processing, sync with provider to get latest state if response is not None: verbose_proxy_logger.debug( @@ -455,7 +461,7 @@ async def retrieve_batch( # noqa: PLR0915 model_id=model_from_id, operation_context="batch retrieval (batch created with model)", ) - + original_batch_id = get_original_file_id(batch_id) prepare_data_with_credentials( data=data, @@ -464,11 +470,11 @@ async def retrieve_batch( # noqa: PLR0915 ) # Fix: The helper sets "file_id" but we need "batch_id" data["batch_id"] = data.pop("file_id", original_batch_id) - + # Retrieve batch using model credentials response = await litellm.aretrieve_batch( custom_llm_provider=credentials["custom_llm_provider"], - **data # type: ignore + **data, # type: ignore ) encode_batch_response_ids(response, model=model_from_id) @@ -476,8 +482,10 @@ async def retrieve_batch( # noqa: PLR0915 verbose_proxy_logger.debug( f"Retrieved batch using model: {model_from_id}, original_id: {original_batch_id}" ) - - elif litellm.enable_loadbalancing_on_batch_endpoints is True or unified_batch_id: + + elif ( + litellm.enable_loadbalancing_on_batch_endpoints is True or unified_batch_id + ): if llm_router is None: raise HTTPException( status_code=500, @@ -489,10 +497,12 @@ async def retrieve_batch( # noqa: PLR0915 response = await llm_router.aretrieve_batch(**data) # type: ignore response._hidden_params["unified_batch_id"] = unified_batch_id if unified_batch_id: - model_id_from_batch = get_model_id_from_unified_batch_id(unified_batch_id) + model_id_from_batch = get_model_id_from_unified_batch_id( + unified_batch_id + ) if model_id_from_batch: response._hidden_params["model_id"] = model_id_from_batch - + # SCENARIO 3: Fallback to custom_llm_provider (uses env variables) else: custom_llm_provider = ( @@ -504,7 +514,7 @@ async def retrieve_batch( # noqa: PLR0915 response = await litellm.aretrieve_batch( custom_llm_provider=custom_llm_provider, **data # type: ignore ) - + # FIX: Update the database with the latest state from provider await update_batch_in_database( batch_id=batch_id, @@ -636,10 +646,10 @@ async def list_batches( # Try to use managed objects table for listing batches (returns encoded IDs) managed_files_obj = proxy_logging_obj.get_proxy_hook("managed_files") - if managed_files_obj is not None and hasattr(managed_files_obj, "list_user_batches"): - verbose_proxy_logger.debug( - "Using managed objects table for batch listing" - ) + if managed_files_obj is not None and hasattr( + managed_files_obj, "list_user_batches" + ): + verbose_proxy_logger.debug("Using managed objects table for batch listing") response = await managed_files_obj.list_user_batches( user_api_key_dict=user_api_key_dict, limit=limit, @@ -648,25 +658,25 @@ async def list_batches( target_model_names=target_model_names, llm_router=llm_router, ) - elif (model_param := ( + elif model_param := ( data.get("model") or request.query_params.get("model") or request.headers.get("x-litellm-model") - )): + ): # SCENARIO 2: Use model-based routing from header/query/body credentials = get_credentials_for_model( llm_router=llm_router, model_id=model_param, operation_context="batch listing", ) - + data.update(credentials) - + response = await litellm.alist_batches( custom_llm_provider=credentials["custom_llm_provider"], after=after, limit=limit, - **data # type: ignore + **data, # type: ignore ) # Encode batch IDs in the list response so clients can use @@ -676,12 +686,16 @@ async def list_batches( encode_batch_response_ids(batch, model=model_param) verbose_proxy_logger.debug(f"Listed batches using model: {model_param}") - + # SCENARIO 2 (alternative): target_model_names based routing elif target_model_names or data.get("target_model_names", None): - target_model_names = target_model_names or data.get("target_model_names", None) + target_model_names = target_model_names or data.get( + "target_model_names", None + ) if target_model_names is None: - raise ValueError("target_model_names is required for this routing scenario") + raise ValueError( + "target_model_names is required for this routing scenario" + ) model = target_model_names.split(",")[0] data.pop("model", None) response = await llm_router.alist_batches( @@ -690,7 +704,7 @@ async def list_batches( limit=limit, **data, ) - + # SCENARIO 3: Fallback to custom_llm_provider (uses env variables) else: custom_llm_provider = ( @@ -795,13 +809,13 @@ async def cancel_batch( try: # Check for encoded batch ID with model info model_from_id = decode_model_from_file_id(batch_id) - + # Create CancelBatchRequest with batch_id to enable ownership checking _cancel_batch_request = CancelBatchRequest( batch_id=batch_id, ) data = cast(dict, _cancel_batch_request) - + unified_batch_id = _is_base64_encoded_unified_file_id(batch_id) base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) @@ -835,7 +849,7 @@ async def cancel_batch( model_id=model_from_id, operation_context="batch cancellation (batch created with model)", ) - + original_batch_id = get_original_file_id(batch_id) prepare_data_with_credentials( data=data, @@ -844,11 +858,11 @@ async def cancel_batch( ) # Fix: The helper sets "file_id" but we need "batch_id" data["batch_id"] = data.pop("file_id", original_batch_id) - + # Cancel batch using model credentials response = await litellm.acancel_batch( custom_llm_provider=credentials["custom_llm_provider"], - **data # type: ignore + **data, # type: ignore ) encode_batch_response_ids(response, model=model_from_id) @@ -856,7 +870,7 @@ async def cancel_batch( verbose_proxy_logger.debug( f"Cancelled batch using model: {model_from_id}, original_id: {original_batch_id}" ) - + # SCENARIO 2: target_model_names based routing elif unified_batch_id: if llm_router is None: @@ -870,14 +884,13 @@ async def cancel_batch( # Hook has already extracted model and unwrapped batch_id into data dict response = await llm_router.acancel_batch(**data) # type: ignore response._hidden_params["unified_batch_id"] = unified_batch_id - + # Ensure model_id is set for the post_call_success_hook to re-encode IDs if not response._hidden_params.get("model_id") and data.get("model"): response._hidden_params["model_id"] = data["model"] - + # SCENARIO 3: Fallback to custom_llm_provider (uses env variables) else: - custom_llm_provider = ( provider or data.pop("custom_llm_provider", None) or "openai" ) @@ -893,7 +906,7 @@ async def cancel_batch( # FIX: Update the database with the new cancelled state managed_files_obj = proxy_logging_obj.get_proxy_hook("managed_files") from litellm.proxy.proxy_server import prisma_client - + await update_batch_in_database( batch_id=batch_id, unified_batch_id=unified_batch_id, diff --git a/litellm/proxy/client/__init__.py b/litellm/proxy/client/__init__.py index 89574bfd24..370585728b 100644 --- a/litellm/proxy/client/__init__.py +++ b/litellm/proxy/client/__init__.py @@ -6,4 +6,12 @@ from .exceptions import UnauthorizedError from .users import UsersManagementClient from .health import HealthManagementClient -__all__ = ["Client", "ChatClient", "ModelsManagementClient", "ModelGroupsManagementClient", "UsersManagementClient", "UnauthorizedError", "HealthManagementClient"] +__all__ = [ + "Client", + "ChatClient", + "ModelsManagementClient", + "ModelGroupsManagementClient", + "UsersManagementClient", + "UnauthorizedError", + "HealthManagementClient", +] diff --git a/litellm/proxy/client/chat.py b/litellm/proxy/client/chat.py index 91fc33002b..064c6162b0 100644 --- a/litellm/proxy/client/chat.py +++ b/litellm/proxy/client/chat.py @@ -139,11 +139,7 @@ class ChatClient: url = f"{self._base_url}/chat/completions" # Build request data with required fields - data: Dict[str, Any] = { - "model": model, - "messages": messages, - "stream": True - } + data: Dict[str, Any] = {"model": model, "messages": messages, "stream": True} # Add optional parameters if provided if temperature is not None: @@ -165,27 +161,24 @@ class ChatClient: session = requests.Session() try: response = session.post( - url, - headers=self._get_headers(), - json=data, - stream=True + url, headers=self._get_headers(), json=data, stream=True ) response.raise_for_status() - + # Parse SSE stream for line in response.iter_lines(): if line: - line = line.decode('utf-8') - if line.startswith('data: '): + line = line.decode("utf-8") + if line.startswith("data: "): data_str = line[6:] # Remove 'data: ' prefix - if data_str.strip() == '[DONE]': + if data_str.strip() == "[DONE]": break try: chunk = json.loads(data_str) yield chunk except json.JSONDecodeError: continue - + except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 64b3233536..aeb59e78a5 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -334,7 +334,10 @@ def _normalize_teams(teams, team_details): """ if isinstance(team_details, list) and team_details: return [ - {"team_id": i.get("team_id") or i.get("id"), "team_alias": i.get("team_alias")} + { + "team_id": i.get("team_id") or i.get("id"), + "team_alias": i.get("team_alias"), + } for i in team_details if isinstance(i, dict) and (i.get("team_id") or i.get("id")) ] @@ -608,7 +611,9 @@ def whoami(): click.echo(f"Token age: {age_hours:.1f} hours") if age_hours > CLI_JWT_EXPIRATION_HOURS: - click.echo(f"⚠️ Warning: Token is more than {CLI_JWT_EXPIRATION_HOURS} hours old and may have expired.") + click.echo( + f"⚠️ Warning: Token is more than {CLI_JWT_EXPIRATION_HOURS} hours old and may have expired." + ) # Export functions for use by other CLI commands diff --git a/litellm/proxy/client/cli/commands/chat.py b/litellm/proxy/client/cli/commands/chat.py index 41ded68ed0..a078b76610 100644 --- a/litellm/proxy/client/cli/commands/chat.py +++ b/litellm/proxy/client/cli/commands/chat.py @@ -28,47 +28,53 @@ def _get_available_models(ctx: click.Context) -> List[Dict[str, Any]]: return [] -def _select_model(console: Console, available_models: List[Dict[str, Any]]) -> Optional[str]: +def _select_model( + console: Console, available_models: List[Dict[str, Any]] +) -> Optional[str]: """Interactive model selection""" if not available_models: - console.print("[yellow]No models available or could not fetch models list.[/yellow]") + console.print( + "[yellow]No models available or could not fetch models list.[/yellow]" + ) model_name = Prompt.ask("Please enter a model name") return model_name if model_name.strip() else None - + # Display available models in a table table = Table(title="Available Models") table.add_column("Index", style="cyan", no_wrap=True) table.add_column("Model ID", style="green") table.add_column("Owned By", style="yellow") MAX_MODELS_TO_DISPLAY = 200 - + models_to_display: List[Dict[str, Any]] = available_models[:MAX_MODELS_TO_DISPLAY] for i, model in enumerate(models_to_display): # Limit to first 200 models table.add_row( - str(i + 1), - str(model.get("id", "")), - str(model.get("owned_by", "")) + str(i + 1), str(model.get("id", "")), str(model.get("owned_by", "")) ) - + if len(available_models) > MAX_MODELS_TO_DISPLAY: - console.print(f"\n[dim]... and {len(available_models) - MAX_MODELS_TO_DISPLAY} more models[/dim]") - + console.print( + f"\n[dim]... and {len(available_models) - MAX_MODELS_TO_DISPLAY} more models[/dim]" + ) + console.print(table) - + while True: try: choice = Prompt.ask( "\nSelect a model by entering the index number (or type a model name directly)", - default="1" + default="1", ).strip() - + # Try to parse as index try: index = int(choice) - 1 if 0 <= index < len(available_models): return available_models[index]["id"] else: - console.print(f"[red]Invalid index. Please enter a number between 1 and {len(available_models)}[/red]") + console.print( + f"[red]Invalid index. Please enter a number between 1 and {len(available_models)}[/red]" + ) continue except ValueError: # Not a number, treat as model name @@ -77,7 +83,7 @@ def _select_model(console: Console, available_models: List[Dict[str, Any]]) -> O else: console.print("[red]Please enter a valid model name or index[/red]") continue - + except KeyboardInterrupt: console.print("\n[yellow]Model selection cancelled.[/yellow]") return None @@ -112,20 +118,20 @@ def chat( system: Optional[str] = None, ): """Interactive chat with streaming responses - + Examples: - + # Chat with a specific model litellm-proxy chat gpt-4 - + # Chat without specifying model (will show model selection) litellm-proxy chat - + # Chat with custom settings litellm-proxy chat gpt-4 --temperature 0.9 --system "You are a helpful coding assistant" """ console = Console() - + # If no model specified, show model selection if not model: available_models = _get_available_models(ctx) @@ -133,27 +139,29 @@ def chat( if not model: console.print("[red]No model selected. Exiting.[/red]") return - + client = ChatClient(ctx.obj["base_url"], ctx.obj["api_key"]) - + # Initialize conversation history messages: List[Dict[str, Any]] = [] - + # Add system message if provided if system: messages.append({"role": "system", "content": system}) - + # Display welcome message - console.print(Panel.fit( - f"[bold blue]LiteLLM Interactive Chat[/bold blue]\n" - f"Model: [green]{model}[/green]\n" - f"Temperature: [yellow]{temperature}[/yellow]\n" - f"Max Tokens: [yellow]{max_tokens or 'unlimited'}[/yellow]\n\n" - f"Type your messages and press Enter. Type '/quit' or '/exit' to end the session.\n" - f"Type '/help' for more commands.", - title="🤖 Chat Session" - )) - + console.print( + Panel.fit( + f"[bold blue]LiteLLM Interactive Chat[/bold blue]\n" + f"Model: [green]{model}[/green]\n" + f"Temperature: [yellow]{temperature}[/yellow]\n" + f"Max Tokens: [yellow]{max_tokens or 'unlimited'}[/yellow]\n\n" + f"Type your messages and press Enter. Type '/quit' or '/exit' to end the session.\n" + f"Type '/help' for more commands.", + title="🤖 Chat Session", + ) + ) + try: while True: # Get user input @@ -162,27 +170,42 @@ def chat( except (EOFError, KeyboardInterrupt): console.print("\n[yellow]Chat session ended.[/yellow]") break - + # Handle special commands should_exit, messages, new_model = _handle_special_commands( console, user_input, messages, system, ctx ) - + if should_exit: break if new_model: model = new_model - + # Check if this was a special command that was handled (not a normal message) - if user_input.lower().startswith(('/quit', '/exit', '/q', '/help', '/clear', '/history', '/save', '/load', '/model')) or not user_input: + if ( + user_input.lower().startswith( + ( + "/quit", + "/exit", + "/q", + "/help", + "/clear", + "/history", + "/save", + "/load", + "/model", + ) + ) + or not user_input + ): continue - + # Add user message to conversation messages.append({"role": "user", "content": user_input}) - + # Display assistant label console.print("\n[bold green]Assistant:[/bold green]") - + # Stream the response assistant_content = _stream_response( console=console, @@ -192,13 +215,13 @@ def chat( temperature=temperature, max_tokens=max_tokens, ) - + # Add assistant message to conversation history if assistant_content: messages.append({"role": "assistant", "content": assistant_content}) else: console.print("[red]Error: No content received from the model[/red]") - + except KeyboardInterrupt: console.print("\n[yellow]Chat session interrupted.[/yellow]") @@ -230,19 +253,23 @@ def _show_history(console: Console, messages: List[Dict[str, Any]]): if not messages: console.print("[yellow]No conversation history.[/yellow]") return - + console.print(Panel.fit("[bold]Conversation History[/bold]", title="History")) - + for i, message in enumerate(messages, 1): role = message["role"] content = message["content"] - + if role == "system": - console.print(f"[dim]{i}. [bold magenta]System:[/bold magenta] {content}[/dim]") + console.print( + f"[dim]{i}. [bold magenta]System:[/bold magenta] {content}[/dim]" + ) elif role == "user": console.print(f"{i}. [bold cyan]You:[/bold cyan] {content}") elif role == "assistant": - console.print(f"{i}. [bold green]Assistant:[/bold green] {content[:100]}{'...' if len(content) > 100 else ''}") + console.print( + f"{i}. [bold green]Assistant:[/bold green] {content[:100]}{'...' if len(content) > 100 else ''}" + ) def _save_conversation(console: Console, messages: List[Dict[str, Any]], command: str): @@ -251,32 +278,34 @@ def _save_conversation(console: Console, messages: List[Dict[str, Any]], command if len(parts) < 2: console.print("[red]Usage: /save [/red]") return - + filename = parts[1] - if not filename.endswith('.json'): - filename += '.json' - + if not filename.endswith(".json"): + filename += ".json" + try: - with open(filename, 'w') as f: + with open(filename, "w") as f: json.dump(messages, f, indent=2) console.print(f"[green]Conversation saved to {filename}[/green]") except Exception as e: console.print(f"[red]Error saving conversation: {e}[/red]") -def _load_conversation(console: Console, command: str, system: Optional[str]) -> List[Dict[str, Any]]: +def _load_conversation( + console: Console, command: str, system: Optional[str] +) -> List[Dict[str, Any]]: """Load conversation from a file""" parts = command.split() if len(parts) < 2: console.print("[red]Usage: /load [/red]") return [] - + filename = parts[1] - if not filename.endswith('.json'): - filename += '.json' - + if not filename.endswith(".json"): + filename += ".json" + try: - with open(filename, 'r') as f: + with open(filename, "r") as f: messages = json.load(f) console.print(f"[green]Conversation loaded from {filename}[/green]") return messages @@ -284,7 +313,7 @@ def _load_conversation(console: Console, command: str, system: Optional[str]) -> console.print(f"[red]File not found: {filename}[/red]") except Exception as e: console.print(f"[red]Error loading conversation: {e}[/red]") - + # Return empty list or just system message if load failed if system: return [{"role": "system", "content": system}] @@ -292,35 +321,35 @@ def _load_conversation(console: Console, command: str, system: Optional[str]) -> def _handle_special_commands( - console: Console, - user_input: str, - messages: List[Dict[str, Any]], + console: Console, + user_input: str, + messages: List[Dict[str, Any]], system: Optional[str], - ctx: click.Context + ctx: click.Context, ) -> tuple[bool, List[Dict[str, Any]], Optional[str]]: """Handle special chat commands. Returns (should_exit, updated_messages, updated_model)""" - if user_input.lower() in ['/quit', '/exit', '/q']: + if user_input.lower() in ["/quit", "/exit", "/q"]: console.print("[yellow]Chat session ended.[/yellow]") return True, messages, None - elif user_input.lower() == '/help': + elif user_input.lower() == "/help": _show_help(console) return False, messages, None - elif user_input.lower() == '/clear': + elif user_input.lower() == "/clear": new_messages = [] if system: new_messages.append({"role": "system", "content": system}) console.print("[green]Conversation history cleared.[/green]") return False, new_messages, None - elif user_input.lower() == '/history': + elif user_input.lower() == "/history": _show_history(console, messages) return False, messages, None - elif user_input.lower().startswith('/save'): + elif user_input.lower().startswith("/save"): _save_conversation(console, messages, user_input) return False, messages, None - elif user_input.lower().startswith('/load'): + elif user_input.lower().startswith("/load"): new_messages = _load_conversation(console, user_input, system) return False, new_messages, None - elif user_input.lower() == '/model': + elif user_input.lower() == "/model": available_models = _get_available_models(ctx) new_model = _select_model(console, available_models) if new_model: @@ -329,12 +358,19 @@ def _handle_special_commands( return False, messages, None elif not user_input: return False, messages, None - + # Not a special command return False, messages, None -def _stream_response(console: Console, client: ChatClient, model: str, messages: List[Dict[str, Any]], temperature: float, max_tokens: Optional[int]) -> Optional[str]: +def _stream_response( + console: Console, + client: ChatClient, + model: str, + messages: List[Dict[str, Any]], + temperature: float, + max_tokens: Optional[int], +) -> Optional[str]: """Stream the model response and return the complete content""" try: assistant_content = "" @@ -351,18 +387,20 @@ def _stream_response(console: Console, client: ChatClient, model: str, messages: assistant_content += content console.print(content, end="") sys.stdout.flush() - + console.print() # Add newline after streaming return assistant_content if assistant_content else None - + except requests.exceptions.HTTPError as e: console.print(f"\n[red]Error: HTTP {e.response.status_code}[/red]") try: error_body = e.response.json() - console.print(f"[red]{error_body.get('error', {}).get('message', 'Unknown error')}[/red]") + console.print( + f"[red]{error_body.get('error', {}).get('message', 'Unknown error')}[/red]" + ) except json.JSONDecodeError: console.print(f"[red]{e.response.text}[/red]") return None except Exception as e: console.print(f"\n[red]Error: {str(e)}[/red]") - return None \ No newline at end of file + return None diff --git a/litellm/proxy/client/cli/commands/http.py b/litellm/proxy/client/cli/commands/http.py index dba36f9d92..b724f3cf2c 100644 --- a/litellm/proxy/client/cli/commands/http.py +++ b/litellm/proxy/client/cli/commands/http.py @@ -61,7 +61,9 @@ def request( key, value = h.split(":", 1) headers[key.strip()] = value.strip() except ValueError: - raise click.BadParameter(f"Invalid header format: {h}. Expected format: 'key:value'") + raise click.BadParameter( + f"Invalid header format: {h}. Expected format: 'key:value'" + ) # Parse JSON data if provided json_data = None diff --git a/litellm/proxy/client/cli/commands/keys.py b/litellm/proxy/client/cli/commands/keys.py index 2e10304b3f..a007d26019 100644 --- a/litellm/proxy/client/cli/commands/keys.py +++ b/litellm/proxy/client/cli/commands/keys.py @@ -24,8 +24,15 @@ def keys(): @click.option("--organization-id", type=str, help="Filter keys by organization ID") @click.option("--key-hash", type=str, help="Filter by specific key hash") @click.option("--key-alias", type=str, help="Filter by key alias") -@click.option("--return-full-object", is_flag=True, default=True, help="Return the full key object") -@click.option("--include-team-keys", is_flag=True, help="Include team keys in the response") +@click.option( + "--return-full-object", + is_flag=True, + default=True, + help="Return the full key object", +) +@click.option( + "--include-team-keys", is_flag=True, help="Include team keys in the response" +) @click.option( "--format", "output_format", @@ -65,7 +72,9 @@ def list( if output_format == "json": rich.print_json(data=response) else: - rich.print(f"Showing {len(response.get('keys', []))} keys out of {response.get('total_count', 0)}") + rich.print( + f"Showing {len(response.get('keys', []))} keys out of {response.get('total_count', 0)}" + ) table = Table(title="API Keys") table.add_column("Key Hash", style="cyan") table.add_column("Alias", style="green") @@ -87,12 +96,18 @@ def list( @click.option("--models", type=str, help="Comma-separated list of allowed models") @click.option("--aliases", type=str, help="JSON string of model alias mappings") @click.option("--spend", type=float, help="Maximum spend limit for this key") -@click.option("--duration", type=str, help="Duration for which the key is valid (e.g. '24h', '7d')") +@click.option( + "--duration", + type=str, + help="Duration for which the key is valid (e.g. '24h', '7d')", +) @click.option("--key-alias", type=str, help="Alias/name for the key") @click.option("--team-id", type=str, help="Team ID to associate the key with") @click.option("--user-id", type=str, help="User ID to associate the key with") @click.option("--budget-id", type=str, help="Budget ID to associate the key with") -@click.option("--config", type=str, help="JSON string of additional configuration parameters") +@click.option( + "--config", type=str, help="JSON string of additional configuration parameters" +) @click.pass_context def generate( ctx: click.Context, @@ -139,7 +154,9 @@ def generate( @keys.command() @click.option("--keys", type=str, help="Comma-separated list of API keys to delete") -@click.option("--key-aliases", type=str, help="Comma-separated list of key aliases to delete") +@click.option( + "--key-aliases", type=str, help="Comma-separated list of key aliases to delete" +) @click.pass_context def delete(ctx: click.Context, keys: Optional[str], key_aliases: Optional[str]): """Delete API keys by key or alias""" @@ -171,11 +188,16 @@ def _parse_created_since_filter(created_since: Optional[str]) -> Optional[dateti else: return datetime.strptime(created_since, "%Y-%m-%d") except ValueError: - click.echo(f"Error: Invalid date format '{created_since}'. Use YYYY-MM-DD_HH:MM or YYYY-MM-DD", err=True) + click.echo( + f"Error: Invalid date format '{created_since}'. Use YYYY-MM-DD_HH:MM or YYYY-MM-DD", + err=True, + ) raise click.Abort() -def _fetch_all_keys_with_pagination(source_client: KeysManagementClient, source_base_url: str) -> List[Dict[str, Any]]: +def _fetch_all_keys_with_pagination( + source_client: KeysManagementClient, source_base_url: str +) -> List[Dict[str, Any]]: """Fetch all keys from source instance using pagination.""" click.echo(f"Fetching keys from source server: {source_base_url}") source_keys = [] @@ -183,7 +205,9 @@ def _fetch_all_keys_with_pagination(source_client: KeysManagementClient, source_ page_size = 100 # Use a larger page size to minimize API calls while True: - source_response = source_client.list(return_full_object=True, page=page, size=page_size) + source_response = source_client.list( + return_full_object=True, page=page, size=page_size + ) # source_client.list() returns Dict[str, Any] when return_request is False (default) assert isinstance(source_response, dict), "Expected dict response from list API" page_keys = source_response.get("keys", []) @@ -204,7 +228,9 @@ def _fetch_all_keys_with_pagination(source_client: KeysManagementClient, source_ def _filter_keys_by_created_since( - source_keys: List[Dict[str, Any]], created_since_dt: Optional[datetime], created_since: str + source_keys: List[Dict[str, Any]], + created_since_dt: Optional[datetime], + created_since: str, ) -> List[Dict[str, Any]]: """Filter keys by created_since date if specified.""" if not created_since_dt: @@ -217,7 +243,9 @@ def _filter_keys_by_created_since( # Parse the key's created_at timestamp if isinstance(key_created_at, str): if "T" in key_created_at: - key_dt = datetime.fromisoformat(key_created_at.replace("Z", "+00:00")) + key_dt = datetime.fromisoformat( + key_created_at.replace("Z", "+00:00") + ) else: key_dt = datetime.fromisoformat(key_created_at) @@ -228,7 +256,9 @@ def _filter_keys_by_created_since( if key_dt >= created_since_dt: filtered_keys.append(key) - click.echo(f"Filtered {len(source_keys)} keys to {len(filtered_keys)} keys created since {created_since}") + click.echo( + f"Filtered {len(source_keys)} keys to {len(filtered_keys)} keys created since {created_since}" + ) return filtered_keys @@ -251,7 +281,9 @@ def _display_dry_run_table(source_keys: List[Dict[str, Any]]) -> None: dt = datetime.fromisoformat(created_at.replace("Z", "+00:00")) created_at = dt.strftime("%Y-%m-%d %H:%M") - table.add_row(str(key.get("key_alias", "")), str(key.get("user_id", "")), str(created_at)) + table.add_row( + str(key.get("key_alias", "")), str(key.get("user_id", "")), str(created_at) + ) rich.print(table) @@ -260,7 +292,16 @@ def _prepare_key_import_data(key: Dict[str, Any]) -> Dict[str, Any]: import_data = {} # Copy relevant fields if they exist - for field in ["models", "aliases", "spend", "key_alias", "team_id", "user_id", "budget_id", "config"]: + for field in [ + "models", + "aliases", + "spend", + "key_alias", + "team_id", + "user_id", + "budget_id", + "config", + ]: if key.get(field): import_data[field] = key[field] @@ -298,16 +339,29 @@ def _import_keys_to_destination( @keys.command(name="import") @click.option( - "--source-base-url", required=True, help="Base URL of the source LiteLLM proxy server to import keys from" + "--source-base-url", + required=True, + help="Base URL of the source LiteLLM proxy server to import keys from", ) -@click.option("--source-api-key", help="API key for authentication to the source server") -@click.option("--dry-run", is_flag=True, help="Show what would be imported without actually importing") @click.option( - "--created-since", help="Only import keys created after this date/time (format: YYYY-MM-DD_HH:MM or YYYY-MM-DD)" + "--source-api-key", help="API key for authentication to the source server" +) +@click.option( + "--dry-run", + is_flag=True, + help="Show what would be imported without actually importing", +) +@click.option( + "--created-since", + help="Only import keys created after this date/time (format: YYYY-MM-DD_HH:MM or YYYY-MM-DD)", ) @click.pass_context def import_keys( - ctx: click.Context, source_base_url: str, source_api_key: Optional[str], dry_run: bool, created_since: Optional[str] + ctx: click.Context, + source_base_url: str, + source_api_key: Optional[str], + dry_run: bool, + created_since: Optional[str], ): """Import API keys from another LiteLLM instance""" # Parse created_since filter if provided @@ -323,7 +377,9 @@ def import_keys( # Filter keys by created_since if specified if created_since: - source_keys = _filter_keys_by_created_since(source_keys, created_since_dt, created_since) + source_keys = _filter_keys_by_created_since( + source_keys, created_since_dt, created_since + ) if not source_keys: click.echo("No keys found in source instance.") @@ -336,7 +392,9 @@ def import_keys( return # Import each key - imported_count, failed_count = _import_keys_to_destination(source_keys, dest_client) + imported_count, failed_count = _import_keys_to_destination( + source_keys, dest_client + ) # Summary click.echo("\nImport completed:") diff --git a/litellm/proxy/client/cli/commands/models.py b/litellm/proxy/client/cli/commands/models.py index 4ff59e6be8..8acafbd88a 100644 --- a/litellm/proxy/client/cli/commands/models.py +++ b/litellm/proxy/client/cli/commands/models.py @@ -129,7 +129,9 @@ def list_models(ctx: click.Context, output_format: Literal["table", "json"]) -> table.add_row( str(model.get("id", "")), str(model.get("object", "model")), - format_timestamp(created) if isinstance(created, int) else format_iso_datetime_str(created), + format_timestamp(created) + if isinstance(created, int) + else format_iso_datetime_str(created), str(model.get("owned_by", "")), ) @@ -151,7 +153,9 @@ def list_models(ctx: click.Context, output_format: Literal["table", "json"]) -> help="Model info in key=value format (can be specified multiple times)", ) @click.pass_context -def add_model(ctx: click.Context, model_name: str, param: tuple[str, ...], info: tuple[str, ...]) -> None: +def add_model( + ctx: click.Context, model_name: str, param: tuple[str, ...], info: tuple[str, ...] +) -> None: """Add a new model to the proxy""" # Convert parameters from key=value format to dict model_params = dict(p.split("=", 1) for p in param) @@ -180,7 +184,9 @@ def delete_model(ctx: click.Context, model_id: str) -> None: @click.option("--id", "model_id", help="ID of the model to retrieve") @click.option("--name", "model_name", help="Name of the model to retrieve") @click.pass_context -def get_model(ctx: click.Context, model_id: Optional[str], model_name: Optional[str]) -> None: +def get_model( + ctx: click.Context, model_id: Optional[str], model_name: Optional[str] +) -> None: """Get information about a specific model""" if not model_id and not model_name: raise click.UsageError("Either --id or --name must be provided") @@ -205,7 +211,9 @@ def get_model(ctx: click.Context, model_id: Optional[str], model_name: Optional[ help="Comma-separated list of columns to display. Valid columns: public_model, upstream_model, credential_name, created_at, updated_at, id, input_cost, output_cost. Default: public_model,upstream_model,updated_at", ) @click.pass_context -def get_models_info(ctx: click.Context, output_format: Literal["table", "json"], columns: str) -> None: +def get_models_info( + ctx: click.Context, output_format: Literal["table", "json"], columns: str +) -> None: """Get detailed information about all models""" client = create_client(ctx) models_info = client.models.info() @@ -226,22 +234,30 @@ def get_models_info(ctx: click.Context, output_format: Literal["table", "json"], "upstream_model": { "header": "Upstream Model", "style": "green", - "get_value": lambda m: str(m.get("litellm_params", {}).get("model", "")), + "get_value": lambda m: str( + m.get("litellm_params", {}).get("model", "") + ), }, "credential_name": { "header": "Credential Name", "style": "yellow", - "get_value": lambda m: str(m.get("litellm_params", {}).get("litellm_credential_name", "")), + "get_value": lambda m: str( + m.get("litellm_params", {}).get("litellm_credential_name", "") + ), }, "created_at": { "header": "Created At", "style": "magenta", - "get_value": lambda m: format_iso_datetime_str(m.get("model_info", {}).get("created_at")), + "get_value": lambda m: format_iso_datetime_str( + m.get("model_info", {}).get("created_at") + ), }, "updated_at": { "header": "Updated At", "style": "magenta", - "get_value": lambda m: format_iso_datetime_str(m.get("model_info", {}).get("updated_at")), + "get_value": lambda m: format_iso_datetime_str( + m.get("model_info", {}).get("updated_at") + ), }, "id": { "header": "ID", @@ -252,13 +268,17 @@ def get_models_info(ctx: click.Context, output_format: Literal["table", "json"], "header": "Input Cost", "style": "green", "justify": "right", - "get_value": lambda m: format_cost_per_1k_tokens(m.get("model_info", {}).get("input_cost_per_token")), + "get_value": lambda m: format_cost_per_1k_tokens( + m.get("model_info", {}).get("input_cost_per_token") + ), }, "output_cost": { "header": "Output Cost", "style": "green", "justify": "right", - "get_value": lambda m: format_cost_per_1k_tokens(m.get("model_info", {}).get("output_cost_per_token")), + "get_value": lambda m: format_cost_per_1k_tokens( + m.get("model_info", {}).get("output_cost_per_token") + ), }, } @@ -267,7 +287,11 @@ def get_models_info(ctx: click.Context, output_format: Literal["table", "json"], for col_name in requested_columns: if col_name in column_configs: config = column_configs[col_name] - table.add_column(config["header"], style=config["style"], justify=config.get("justify", "left")) + table.add_column( + config["header"], + style=config["style"], + justify=config.get("justify", "left"), + ) else: click.echo(f"Warning: Unknown column '{col_name}'", err=True) @@ -298,7 +322,9 @@ def get_models_info(ctx: click.Context, output_format: Literal["table", "json"], help="Model info in key=value format (can be specified multiple times)", ) @click.pass_context -def update_model(ctx: click.Context, model_id: str, param: tuple[str, ...], info: tuple[str, ...]) -> None: +def update_model( + ctx: click.Context, model_id: str, param: tuple[str, ...], info: tuple[str, ...] +) -> None: """Update an existing model's configuration""" # Convert parameters from key=value format to dict model_params = dict(p.split("=", 1) for p in param) @@ -328,7 +354,10 @@ def _filter_model(model, model_regex, access_group_regex): if access_group_regex: if not isinstance(access_groups, list): return False - if not any(isinstance(group, str) and access_group_regex.search(group) for group in access_groups): + if not any( + isinstance(group, str) and access_group_regex.search(group) + for group in access_groups + ): return False return True @@ -364,18 +393,32 @@ def get_model_list_from_yaml_file(yaml_file: str) -> list[dict[str, Any]]: with open(yaml_file, "r") as f: data = yaml.safe_load(f) if not data or "model_list" not in data: - raise click.ClickException("YAML file must contain a 'model_list' key with a list of models.") + raise click.ClickException( + "YAML file must contain a 'model_list' key with a list of models." + ) model_list = data["model_list"] if not isinstance(model_list, list): raise click.ClickException("'model_list' must be a list of model definitions.") return model_list -def _get_filtered_model_list(model_list, only_models_matching_regex, only_access_groups_matching_regex): +def _get_filtered_model_list( + model_list, only_models_matching_regex, only_access_groups_matching_regex +): """Return a list of models that pass the filter criteria.""" - model_regex = re.compile(only_models_matching_regex) if only_models_matching_regex else None - access_group_regex = re.compile(only_access_groups_matching_regex) if only_access_groups_matching_regex else None - return [model for model in model_list if _filter_model(model, model_regex, access_group_regex)] + model_regex = ( + re.compile(only_models_matching_regex) if only_models_matching_regex else None + ) + access_group_regex = ( + re.compile(only_access_groups_matching_regex) + if only_access_groups_matching_regex + else None + ) + return [ + model + for model in model_list + if _filter_model(model, model_regex, access_group_regex) + ] def _import_models_get_table_title(dry_run: bool) -> str: @@ -386,8 +429,14 @@ def _import_models_get_table_title(dry_run: bool) -> str: @models.command("import") -@click.argument("yaml_file", type=click.Path(exists=True, dir_okay=False, readable=True)) -@click.option("--dry-run", is_flag=True, help="Show what would be imported without making any changes.") +@click.argument( + "yaml_file", type=click.Path(exists=True, dir_okay=False, readable=True) +) +@click.option( + "--dry-run", + is_flag=True, + help="Show what would be imported without making any changes.", +) @click.option( "--only-models-matching-regex", default=None, diff --git a/litellm/proxy/client/cli/commands/teams.py b/litellm/proxy/client/cli/commands/teams.py index 57397ca01a..51a3250162 100644 --- a/litellm/proxy/client/cli/commands/teams.py +++ b/litellm/proxy/client/cli/commands/teams.py @@ -19,11 +19,11 @@ def teams(): def display_teams_table(teams: List[Dict[str, Any]]) -> None: """Display teams in a formatted table""" console = Console() - + if not teams: console.print("❌ No teams found for your user.") return - + table = Table(title="Available Teams") table.add_column("Index", style="cyan", no_wrap=True) table.add_column("Team Alias", style="magenta") @@ -31,13 +31,13 @@ def display_teams_table(teams: List[Dict[str, Any]]) -> None: table.add_column("Models", style="yellow") table.add_column("Max Budget", style="blue") table.add_column("Role", style="red") - + for i, team in enumerate(teams): team_alias = team.get("team_alias") or "N/A" team_id = team.get("team_id", "N/A") models = team.get("models", []) max_budget = team.get("max_budget") - + # Format models list if models: if len(models) > 3: @@ -46,25 +46,22 @@ def display_teams_table(teams: List[Dict[str, Any]]) -> None: models_str = ", ".join(models) else: models_str = "All models" - + # Format budget budget_str = f"${max_budget}" if max_budget else "Unlimited" - + # Try to determine role (this might vary based on API response structure) role = "Member" # Default role - if isinstance(team, dict) and 'members_with_roles' in team and team['members_with_roles']: + if ( + isinstance(team, dict) + and "members_with_roles" in team + and team["members_with_roles"] + ): # This would need to be implemented based on actual API response structure pass - - table.add_row( - str(i + 1), - team_alias, - team_id, - models_str, - budget_str, - role - ) - + + table.add_row(str(i + 1), team_alias, team_id, models_str, budget_str, role) + console.print(table) @@ -73,7 +70,7 @@ def display_teams_table(teams: List[Dict[str, Any]]) -> None: def list(ctx: click.Context): """List teams that you belong to""" client = Client(ctx.obj["base_url"], ctx.obj["api_key"]) - + try: # Use list() for simpler response structure (returns array directly) teams = client.teams.list() @@ -93,7 +90,7 @@ def list(ctx: click.Context): def available(ctx: click.Context): """List teams that are available to join""" client = Client(ctx.obj["base_url"], ctx.obj["api_key"]) - + try: teams = client.teams.get_available() if teams: @@ -118,47 +115,48 @@ def assign_key(ctx: click.Context, team_id: Optional[str]): """Assign your current CLI key to a team""" client = Client(ctx.obj["base_url"], ctx.obj["api_key"]) api_key = ctx.obj["api_key"] - + if not api_key: click.echo("❌ No API key found. Please login first using 'litellm login'") raise click.Abort() - + try: # If no team_id provided, show teams and let user select if not team_id: teams = client.teams.list() - + if not teams: click.echo("❌ No teams found for your user.") return - + # Use interactive selection from auth module from .auth import prompt_team_selection + selected_team = prompt_team_selection(teams) - + if selected_team: - team_id = selected_team.get('team_id') + team_id = selected_team.get("team_id") else: click.echo("❌ Operation cancelled.") return - + # Update the key with the selected team if team_id: click.echo(f"\n🔄 Assigning your key to team: {team_id}") client.keys.update(key=api_key, team_id=team_id) click.echo(f"✅ Successfully assigned key to team: {team_id}") - + # Show team details if available teams = client.teams.list() for team in teams: - if team.get('team_id') == team_id: - models = team.get('models', []) + if team.get("team_id") == team_id: + models = team.get("models", []) if models: click.echo(f"🎯 You can now access models: {', '.join(models)}") else: click.echo("🎯 You can now access all available models") break - + except requests.exceptions.HTTPError as e: click.echo(f"Error: HTTP {e.response.status_code}", err=True) error_body = e.response.json() diff --git a/litellm/proxy/client/cli/commands/users.py b/litellm/proxy/client/cli/commands/users.py index 9887a8d0df..36b29b8fe6 100644 --- a/litellm/proxy/client/cli/commands/users.py +++ b/litellm/proxy/client/cli/commands/users.py @@ -2,16 +2,20 @@ import click import rich from ... import UsersManagementClient + @click.group() def users(): """Manage users on your LiteLLM proxy server""" pass + @users.command("list") @click.pass_context def list_users(ctx: click.Context): """List all users""" - client = UsersManagementClient(base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"]) + client = UsersManagementClient( + base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"] + ) users = client.list_users() if isinstance(users, dict) and "users" in users: users = users["users"] @@ -20,6 +24,7 @@ def list_users(ctx: click.Context): return from rich.table import Table from rich.console import Console + table = Table(title="Users") table.add_column("User ID", style="cyan") table.add_column("Email", style="green") @@ -30,20 +35,24 @@ def list_users(ctx: click.Context): str(user.get("user_id", "")), str(user.get("user_email", "")), str(user.get("user_role", "")), - ", ".join(user.get("teams", []) or []) + ", ".join(user.get("teams", []) or []), ) console = Console() console.print(table) + @users.command("get") @click.option("--id", "user_id", help="ID of the user to retrieve") @click.pass_context def get_user(ctx: click.Context, user_id: str): """Get information about a specific user""" - client = UsersManagementClient(base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"]) + client = UsersManagementClient( + base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"] + ) result = client.get_user(user_id=user_id) rich.print_json(data=result) + @users.command("create") @click.option("--email", required=True, help="User email") @click.option("--role", default="internal_user", help="User role") @@ -53,7 +62,9 @@ def get_user(ctx: click.Context, user_id: str): @click.pass_context def create_user(ctx: click.Context, email, role, alias, team, max_budget): """Create a new user""" - client = UsersManagementClient(base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"]) + client = UsersManagementClient( + base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"] + ) user_data = { "user_email": email, "user_role": role, @@ -67,11 +78,14 @@ def create_user(ctx: click.Context, email, role, alias, team, max_budget): result = client.create_user(user_data) rich.print_json(data=result) + @users.command("delete") @click.argument("user_ids", nargs=-1) @click.pass_context def delete_user(ctx: click.Context, user_ids): """Delete one or more users by user_id""" - client = UsersManagementClient(base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"]) + client = UsersManagementClient( + base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"] + ) result = client.delete_user(list(user_ids)) - rich.print_json(data=result) \ No newline at end of file + rich.print_json(data=result) diff --git a/litellm/proxy/client/cli/interface.py b/litellm/proxy/client/cli/interface.py index 2c6f5f10b4..eba693dc18 100644 --- a/litellm/proxy/client/cli/interface.py +++ b/litellm/proxy/client/cli/interface.py @@ -26,7 +26,7 @@ def styled_prompt(): # Fallback if we can't get terminal size verbose_logger.debug(f"Error getting terminal size: {e}") click.echo("\n" * 3) - + # Unicode box drawing characters top_left = "┌" top_right = "┐" @@ -34,45 +34,47 @@ def styled_prompt(): bottom_right = "┘" horizontal = "─" vertical = "│" - + # Create the box with increased width width = 80 top_line = top_left + horizontal * (width - 2) + top_right bottom_line = bottom_left + horizontal * (width - 2) + bottom_right - + # Create styled elements left_border = click.style(vertical, fg="blue", bold=True) right_border = click.style(vertical, fg="blue", bold=True) prompt_text = click.style("> ", fg="cyan", bold=True) - + # Display the complete box structure first to reserve space click.echo(click.style(top_line, fg="blue", bold=True)) - + # Create empty space in the box for input empty_space = " " * (width - 4) click.echo(f"{left_border} {empty_space} {right_border}") - + # Display bottom border to complete the box click.echo(click.style(bottom_line, fg="blue", bold=True)) - + # Now move cursor up to the input line and get input click.echo("\033[2A", nl=False) # Move cursor up 2 lines - click.echo(f"\r{left_border} {prompt_text}", nl=False) # Position at start of input line - + click.echo( + f"\r{left_border} {prompt_text}", nl=False + ) # Position at start of input line + try: # Get user input user_input = input().strip() - + # Move cursor down to after the box click.echo("\033[1B") # Move cursor down 1 line click.echo("") # Add some space after - + except (KeyboardInterrupt, EOFError): # Move cursor down and add space click.echo("\033[1B") click.echo("") raise - + return user_input @@ -93,7 +95,7 @@ def show_commands(): ("help", "Show this help message"), ("quit", "Exit the interactive session"), ] - + click.echo("Available commands:") for cmd, description in commands: click.echo(f" {cmd:<20} {description}") @@ -103,13 +105,13 @@ def show_commands(): def setup_shell(ctx: click.Context): """Set up the interactive shell with banner and initial info.""" from litellm.proxy.common_utils.banner import show_banner - + show_banner() - + # Show server connection info base_url = ctx.obj.get("base_url") click.secho(f"Connected to LiteLLM server: {base_url}\n", fg="green") - + show_commands() @@ -125,10 +127,11 @@ def handle_special_commands(user_input: str) -> bool: elif user_input.lower() == "clear": click.clear() from litellm.proxy.common_utils.banner import show_banner + show_banner() show_commands() return True - + return False @@ -138,33 +141,30 @@ def execute_command(user_input: str, ctx: click.Context): parts = user_input.split() command = parts[0] args = parts[1:] if len(parts) > 1 else [] - + # Import cli here to avoid circular import from . import main + cli = main.cli - + # Check if command exists if command not in cli.commands: click.echo(f"Unknown command: {command}") click.echo("Type 'help' to see available commands.") return - + # Execute the command try: # Create a new argument list for click to parse sys.argv = ["litellm-proxy"] + [command] + args - + # Get the command object and invoke it cmd = cli.commands[command] - + # Create a new context for the subcommand with ctx.scope(): - cmd.main( - args, - parent=ctx, - standalone_mode=False - ) - + cmd.main(args, parent=ctx, standalone_mode=False) + except click.ClickException as e: e.show() except click.Abort: @@ -179,29 +179,29 @@ def execute_command(user_input: str, ctx: click.Context): def interactive_shell(ctx: click.Context): """Run the interactive shell.""" setup_shell(ctx) - + while True: try: # Add some space before the input box to ensure it's positioned well click.echo("\n") # Extra spacing - + # Show styled prompt user_input = styled_prompt() - + if not user_input: continue - + # Handle special commands if handle_special_commands(user_input): if user_input.lower() in ["exit", "quit"]: break continue - + # Execute regular commands execute_command(user_input, ctx) - + except (KeyboardInterrupt, EOFError): click.echo("\nGoodbye!") break except Exception as e: - click.echo(f"Error: {e}") \ No newline at end of file + click.echo(f"Error: {e}") diff --git a/litellm/proxy/client/cli/main.py b/litellm/proxy/client/cli/main.py index eab9b31482..744acf3838 100644 --- a/litellm/proxy/client/cli/main.py +++ b/litellm/proxy/client/cli/main.py @@ -38,15 +38,21 @@ def print_version(base_url: str, api_key: Optional[str]): @click.group(invoke_without_command=True) @click.option( - "--version", "-v", is_flag=True, is_eager=True, expose_value=False, + "--version", + "-v", + is_flag=True, + is_eager=True, + expose_value=False, help="Show the LiteLLM Proxy CLI and server version and exit.", callback=lambda ctx, param, value: ( print_version( ctx.params.get("base_url") or "http://localhost:4000", - ctx.params.get("api_key") + ctx.params.get("api_key"), ) or ctx.exit() - ) if value and not ctx.resilient_parsing else None, + ) + if value and not ctx.resilient_parsing + else None, ) @click.option( "--base-url", @@ -72,7 +78,7 @@ def cli(ctx: click.Context, base_url: str, api_key: Optional[str]) -> None: ctx.obj["base_url"] = base_url ctx.obj["api_key"] = api_key - + # If no subcommand was invoked, start interactive mode if ctx.invoked_subcommand is None: interactive_shell(ctx) diff --git a/litellm/proxy/client/client.py b/litellm/proxy/client/client.py index c9066f70de..12b5cd79f7 100644 --- a/litellm/proxy/client/client.py +++ b/litellm/proxy/client/client.py @@ -34,9 +34,17 @@ class Client: # Initialize resource clients self.http = HTTPClient(base_url=base_url, api_key=api_key, timeout=timeout) - self.models = ModelsManagementClient(base_url=self._base_url, api_key=self._api_key) - self.model_groups = ModelGroupsManagementClient(base_url=self._base_url, api_key=self._api_key) + self.models = ModelsManagementClient( + base_url=self._base_url, api_key=self._api_key + ) + self.model_groups = ModelGroupsManagementClient( + base_url=self._base_url, api_key=self._api_key + ) self.chat = ChatClient(base_url=self._base_url, api_key=self._api_key) self.keys = KeysManagementClient(base_url=self._base_url, api_key=self._api_key) - self.credentials = CredentialsManagementClient(base_url=self._base_url, api_key=self._api_key) - self.teams = TeamsManagementClient(base_url=self._base_url, api_key=self._api_key) + self.credentials = CredentialsManagementClient( + base_url=self._base_url, api_key=self._api_key + ) + self.teams = TeamsManagementClient( + base_url=self._base_url, api_key=self._api_key + ) diff --git a/litellm/proxy/client/health.py b/litellm/proxy/client/health.py index b9da8d9c38..3cfcd151d6 100644 --- a/litellm/proxy/client/health.py +++ b/litellm/proxy/client/health.py @@ -1,10 +1,12 @@ from typing import Optional, Dict, Any from .http_client import HTTPClient + class HealthManagementClient: """ Client for interacting with the health endpoints of the LiteLLM proxy server. """ + def __init__(self, base_url: str, api_key: Optional[str] = None, timeout: int = 30): """ Initialize the HealthManagementClient. @@ -37,4 +39,4 @@ class HealthManagementClient: Optional[str]: The server version if available, otherwise None. """ readiness = self.get_readiness() - return readiness.get("litellm_version") \ No newline at end of file + return readiness.get("litellm_version") diff --git a/litellm/proxy/client/keys.py b/litellm/proxy/client/keys.py index 50fd7b9d9c..d8687cbad1 100644 --- a/litellm/proxy/client/keys.py +++ b/litellm/proxy/client/keys.py @@ -89,7 +89,9 @@ class KeysManagementClient: if include_team_keys is not None: params["include_team_keys"] = str(include_team_keys).lower() - request = requests.Request("GET", url, headers=self._get_headers(), params=params) + request = requests.Request( + "GET", url, headers=self._get_headers(), params=params + ) if return_request: return request @@ -257,7 +259,7 @@ class KeysManagementClient: url = f"{self._base_url}/key/update" data: Dict[str, Any] = {"key": key} - + if key_alias is not None: data["key_alias"] = key_alias if user_id is not None: @@ -283,8 +285,9 @@ class KeysManagementClient: except Exception: raise Exception(f"Error updating key: {response_text}") - - def info(self, key: str, return_request: bool = False) -> Union[Dict[str, Any], requests.Request]: + def info( + self, key: str, return_request: bool = False + ) -> Union[Dict[str, Any], requests.Request]: """ Get information about API keys. diff --git a/litellm/proxy/client/model_groups.py b/litellm/proxy/client/model_groups.py index 2be6e10e54..03bc3eae46 100644 --- a/litellm/proxy/client/model_groups.py +++ b/litellm/proxy/client/model_groups.py @@ -27,7 +27,9 @@ class ModelGroupsManagementClient: headers["Authorization"] = f"Bearer {self._api_key}" return headers - def info(self, return_request: bool = False) -> Union[List[Dict[str, Any]], requests.Request]: + def info( + self, return_request: bool = False + ) -> Union[List[Dict[str, Any]], requests.Request]: """ Get detailed information about all model groups from the server. diff --git a/litellm/proxy/client/models.py b/litellm/proxy/client/models.py index 7943d25b99..d2f5eead28 100644 --- a/litellm/proxy/client/models.py +++ b/litellm/proxy/client/models.py @@ -27,7 +27,9 @@ class ModelsManagementClient: headers["Authorization"] = f"Bearer {self._api_key}" return headers - def list(self, return_request: bool = False) -> Union[List[Dict[str, Any]], requests.Request]: + def list( + self, return_request: bool = False + ) -> Union[List[Dict[str, Any]], requests.Request]: """ Get the list of models supported by the server. @@ -109,7 +111,9 @@ class ModelsManagementClient: raise UnauthorizedError(e) raise - def delete(self, model_id: str, return_request: bool = False) -> Union[Dict[str, Any], requests.Request]: + def delete( + self, model_id: str, return_request: bool = False + ) -> Union[Dict[str, Any], requests.Request]: """ Delete a model from the proxy. @@ -148,7 +152,10 @@ class ModelsManagementClient: raise def get( - self, model_id: Optional[str] = None, model_name: Optional[str] = None, return_request: bool = False + self, + model_id: Optional[str] = None, + model_name: Optional[str] = None, + return_request: bool = False, ) -> Union[Dict[str, Any], requests.Request]: """ Get information about a specific model by its ID or name. @@ -168,7 +175,9 @@ class ModelsManagementClient: NotFoundError: If the model is not found requests.exceptions.RequestException: If the request fails with any other error """ - if (model_id is None and model_name is None) or (model_id is not None and model_name is not None): + if (model_id is None and model_name is None) or ( + model_id is not None and model_name is not None + ): raise ValueError("Exactly one of model_id or model_name must be provided") # If return_request is True, delegate to info @@ -202,7 +211,9 @@ class ModelsManagementClient: ) ) - def info(self, return_request: bool = False) -> Union[List[Dict[str, Any]], requests.Request]: + def info( + self, return_request: bool = False + ) -> Union[List[Dict[str, Any]], requests.Request]: """ Get detailed information about all models from the server. diff --git a/litellm/proxy/client/teams.py b/litellm/proxy/client/teams.py index 4f54b6bbd0..017d074485 100644 --- a/litellm/proxy/client/teams.py +++ b/litellm/proxy/client/teams.py @@ -60,10 +60,10 @@ class TeamsManagementClient: params["organization_id"] = organization_id response = requests.get(url, headers=self._get_headers(), params=params) - + if response.status_code == 401: raise UnauthorizedError("Authentication failed. Check your API key.") - + response.raise_for_status() return response.json() @@ -104,7 +104,7 @@ class TeamsManagementClient: "page_size": page_size, "sort_order": sort_order, } - + if user_id: params["user_id"] = user_id if organization_id: @@ -117,10 +117,10 @@ class TeamsManagementClient: params["sort_by"] = sort_by response = requests.get(url, headers=self._get_headers(), params=params) - + if response.status_code == 401: raise UnauthorizedError("Authentication failed. Check your API key.") - + response.raise_for_status() return response.json() @@ -136,11 +136,11 @@ class TeamsManagementClient: UnauthorizedError: If authentication fails """ url = f"{self._base_url}/team/available" - + response = requests.get(url, headers=self._get_headers()) - + if response.status_code == 401: raise UnauthorizedError("Authentication failed. Check your API key.") - + response.raise_for_status() return response.json() diff --git a/litellm/proxy/client/users.py b/litellm/proxy/client/users.py index 66e2d76bee..df7aa228bc 100644 --- a/litellm/proxy/client/users.py +++ b/litellm/proxy/client/users.py @@ -14,7 +14,9 @@ class UsersManagementClient: headers["Authorization"] = f"Bearer {self.api_key}" return headers - def list_users(self, params: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]: + def list_users( + self, params: Optional[Dict[str, Any]] = None + ) -> List[Dict[str, Any]]: """List users (GET /user/list)""" url = f"{self.base_url}/user/list" response = requests.get(url, headers=self._get_headers(), params=params) @@ -47,7 +49,9 @@ class UsersManagementClient: def delete_user(self, user_ids: List[str]) -> Dict[str, Any]: """Delete users (POST /user/delete)""" url = f"{self.base_url}/user/delete" - response = requests.post(url, headers=self._get_headers(), json={"user_ids": user_ids}) + response = requests.post( + url, headers=self._get_headers(), json={"user_ids": user_ids} + ) if response.status_code == 401: raise UnauthorizedError(response.text) response.raise_for_status() diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 5a3f3a984b..7b9e0a4373 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -249,24 +249,21 @@ async def create_response( def _is_azure_model_router_request(model: str) -> bool: """ Check if the requested model is an Azure Model Router. - + Azure Model Router models follow the pattern: - azure_ai/model_router/ - azure_ai/model-router - model_router/ - model-router - + Args: model: The requested model name - + Returns: bool: True if this is an Azure Model Router request """ model_lower = model.lower() - return ( - "model-router" in model_lower - or "model_router" in model_lower - ) + return "model-router" in model_lower or "model_router" in model_lower def _override_openai_response_model( @@ -1233,7 +1230,9 @@ class ProxyBaseLLMRequestProcessing: data=self.data, user_api_key_dict=user_api_key_dict, response=None, - request_headers=(self.data.get("proxy_server_request") or {}).get("headers", {}), + request_headers=(self.data.get("proxy_server_request") or {}).get( + "headers", {} + ), ) if callback_headers: headers.update(callback_headers) diff --git a/litellm/proxy/common_utils/banner.py b/litellm/proxy/common_utils/banner.py index 1deca22373..9983f12837 100644 --- a/litellm/proxy/common_utils/banner.py +++ b/litellm/proxy/common_utils/banner.py @@ -1,4 +1,3 @@ - # LiteLLM ASCII banner LITELLM_BANNER = """ ██╗ ██╗████████╗███████╗██╗ ██╗ ███╗ ███╗ ██║ ██║╚══██╔══╝██╔════╝██║ ██║ ████╗ ████║ @@ -11,7 +10,8 @@ LITELLM_BANNER = """ ██╗ ██╗████████╗██ def show_banner(): """Display the LiteLLM CLI banner.""" try: - import click - click.echo(f"\n{LITELLM_BANNER}\n") + import click + + click.echo(f"\n{LITELLM_BANNER}\n") except ImportError: - print("\n") # noqa: T201 \ No newline at end of file + print("\n") # noqa: T201 diff --git a/litellm/proxy/common_utils/cache_coordinator.py b/litellm/proxy/common_utils/cache_coordinator.py index 4eceb83af5..ccc73c5e6d 100644 --- a/litellm/proxy/common_utils/cache_coordinator.py +++ b/litellm/proxy/common_utils/cache_coordinator.py @@ -181,9 +181,7 @@ class EventDrivenCacheCoordinator: event_to_wait = await self._claim_role() if event_to_wait is not None: - return await self._wait_for_signal_and_get( - event_to_wait, cache_key, cache - ) + return await self._wait_for_signal_and_get(event_to_wait, cache_key, cache) try: result = await self._load_and_cache(cache_key, cache, load_fn) diff --git a/litellm/proxy/common_utils/custom_openapi_spec.py b/litellm/proxy/common_utils/custom_openapi_spec.py index 69472c2cda..a93749c395 100644 --- a/litellm/proxy/common_utils/custom_openapi_spec.py +++ b/litellm/proxy/common_utils/custom_openapi_spec.py @@ -8,34 +8,31 @@ class CustomOpenAPISpec: Handler for customizing OpenAPI specifications with Pydantic models for documentation purposes without runtime validation. """ - + CHAT_COMPLETION_PATHS = [ "/v1/chat/completions", - "/chat/completions", + "/chat/completions", "/engines/{model}/chat/completions", - "/openai/deployments/{model}/chat/completions" + "/openai/deployments/{model}/chat/completions", ] - + EMBEDDING_PATHS = [ "/v1/embeddings", "/embeddings", - "/engines/{model}/embeddings", - "/openai/deployments/{model}/embeddings" + "/engines/{model}/embeddings", + "/openai/deployments/{model}/embeddings", ] - - RESPONSES_API_PATHS = [ - "/v1/responses", - "/responses" - ] - + + RESPONSES_API_PATHS = ["/v1/responses", "/responses"] + @staticmethod def get_pydantic_schema(model_class) -> Optional[Dict[str, Any]]: """ Get JSON schema from a Pydantic model, handling both v1 and v2 APIs. - + Args: model_class: Pydantic model class - + Returns: JSON schema dict or None if failed """ @@ -52,14 +49,18 @@ class CustomOpenAPISpec: except Exception as e: # FastAPI 0.120+ may fail schema generation for certain types (e.g., openai.Timeout) # Log the error and return None to skip schema generation for this model - verbose_proxy_logger.debug(f"Failed to generate schema for {model_class}: {e}") + verbose_proxy_logger.debug( + f"Failed to generate schema for {model_class}: {e}" + ) return None - + @staticmethod - def add_schema_to_components(openapi_schema: Dict[str, Any], schema_name: str, schema_def: Dict[str, Any]) -> None: + def add_schema_to_components( + openapi_schema: Dict[str, Any], schema_name: str, schema_def: Dict[str, Any] + ) -> None: """ Add a schema definition to the OpenAPI components/schemas section. - + Args: openapi_schema: The OpenAPI schema dict to modify schema_name: Name for the schema component @@ -70,122 +71,148 @@ class CustomOpenAPISpec: openapi_schema["components"] = {} if "schemas" not in openapi_schema["components"]: openapi_schema["components"]["schemas"] = {} - + # Add the schema - CustomOpenAPISpec._move_defs_to_components(openapi_schema, {schema_name: schema_def}) - + CustomOpenAPISpec._move_defs_to_components( + openapi_schema, {schema_name: schema_def} + ) + @staticmethod - def add_request_body_to_paths(openapi_schema: Dict[str, Any], paths: List[str], schema_ref: str) -> None: + def add_request_body_to_paths( + openapi_schema: Dict[str, Any], paths: List[str], schema_ref: str + ) -> None: """ Add request body with expanded form fields for better Swagger UI display. This keeps the request body but expands it to show individual fields in the UI. - + Args: openapi_schema: The OpenAPI schema dict to modify paths: List of paths to update schema_ref: Reference to the schema component (e.g., "#/components/schemas/ModelName") """ for path in paths: - if path in openapi_schema.get("paths", {}) and "post" in openapi_schema["paths"][path]: + if ( + path in openapi_schema.get("paths", {}) + and "post" in openapi_schema["paths"][path] + ): # Get the actual schema to extract ALL field definitions - schema_name = schema_ref.split("/")[-1] # Extract "ProxyChatCompletionRequest" from the ref - actual_schema = openapi_schema.get("components", {}).get("schemas", {}).get(schema_name, {}) + schema_name = schema_ref.split("/")[ + -1 + ] # Extract "ProxyChatCompletionRequest" from the ref + actual_schema = ( + openapi_schema.get("components", {}) + .get("schemas", {}) + .get(schema_name, {}) + ) schema_properties = actual_schema.get("properties", {}) required_fields = actual_schema.get("required", []) - + # Extract $defs and add them to components/schemas # This fixes Pydantic v2 $defs not being resolvable in Swagger/OpenAPI if "$defs" in actual_schema: - CustomOpenAPISpec._move_defs_to_components(openapi_schema, actual_schema["$defs"]) - + CustomOpenAPISpec._move_defs_to_components( + openapi_schema, actual_schema["$defs"] + ) + # Create an expanded inline schema instead of just a $ref # This makes Swagger UI show all individual fields in the request body editor expanded_schema = { "type": "object", "required": required_fields, - "properties": {} + "properties": {}, } - + # Add all properties with their full definitions for field_name, field_def in schema_properties.items(): - expanded_field = CustomOpenAPISpec._expand_field_definition(field_def) - + expanded_field = CustomOpenAPISpec._expand_field_definition( + field_def + ) + # Rewrite $defs references to use components/schemas instead - expanded_field = CustomOpenAPISpec._rewrite_defs_refs(expanded_field) - + expanded_field = CustomOpenAPISpec._rewrite_defs_refs( + expanded_field + ) + # Add a simple example for the messages field if field_name == "messages": expanded_field["example"] = [ {"role": "user", "content": "Hello, how are you?"} ] - + expanded_schema["properties"][field_name] = expanded_field - + # Set the request body with the expanded schema openapi_schema["paths"][path]["post"]["requestBody"] = { "required": True, - "content": { - "application/json": { - "schema": expanded_schema - } - } + "content": {"application/json": {"schema": expanded_schema}}, } - + # Keep any existing parameters (like path parameters) but remove conflicting query params if "parameters" in openapi_schema["paths"][path]["post"]: - existing_params = openapi_schema["paths"][path]["post"]["parameters"] + existing_params = openapi_schema["paths"][path]["post"][ + "parameters" + ] # Only keep path parameters, remove query params that conflict with request body filtered_params = [ - param for param in existing_params - if param.get("in") == "path" + param for param in existing_params if param.get("in") == "path" ] - openapi_schema["paths"][path]["post"]["parameters"] = filtered_params - + openapi_schema["paths"][path]["post"][ + "parameters" + ] = filtered_params + @staticmethod - def _move_defs_to_components(openapi_schema: Dict[str, Any], defs: Dict[str, Any]) -> None: + def _move_defs_to_components( + openapi_schema: Dict[str, Any], defs: Dict[str, Any] + ) -> None: """ Move $defs from Pydantic v2 schema to OpenAPI components/schemas. This makes the definitions resolvable in Swagger/OpenAPI viewers. - + Args: openapi_schema: The OpenAPI schema dict to modify defs: The $defs dictionary from Pydantic schema """ if not defs: return - + # Ensure components/schemas exists if "components" not in openapi_schema: openapi_schema["components"] = {} if "schemas" not in openapi_schema["components"]: openapi_schema["components"]["schemas"] = {} - + # Add each definition to components/schemas for def_name, def_schema in defs.items(): # Recursively rewrite any nested $defs references within this definition rewritten_def = CustomOpenAPISpec._rewrite_defs_refs(def_schema) openapi_schema["components"]["schemas"][def_name] = rewritten_def - + # If this definition also has $defs, process them recursively if "$defs" in def_schema: - CustomOpenAPISpec._move_defs_to_components(openapi_schema, def_schema["$defs"]) - + CustomOpenAPISpec._move_defs_to_components( + openapi_schema, def_schema["$defs"] + ) + @staticmethod def _rewrite_defs_refs(schema: Any) -> Any: """ Recursively rewrite $ref values from #/$defs/... to #/components/schemas/... This converts Pydantic v2 references to OpenAPI-compatible references. - + Args: schema: Schema object to process (can be dict, list, or primitive) - + Returns: Schema with rewritten references """ if isinstance(schema, dict): result = {} for key, value in schema.items(): - if key == "$ref" and isinstance(value, str) and value.startswith("#/$defs/"): + if ( + key == "$ref" + and isinstance(value, str) + and value.startswith("#/$defs/") + ): # Rewrite the reference to use components/schemas def_name = value.replace("#/$defs/", "") result[key] = f"#/components/schemas/{def_name}" @@ -200,22 +227,22 @@ class CustomOpenAPISpec: return [CustomOpenAPISpec._rewrite_defs_refs(item) for item in schema] else: return schema - + @staticmethod def _extract_field_schema(field_def: Dict[str, Any]) -> Dict[str, Any]: """ Extract a simple schema from a Pydantic field definition for parameter display. - + Args: field_def: Pydantic field definition - + Returns: Simplified schema for OpenAPI parameter """ # Handle simple types if "type" in field_def: return {"type": field_def["type"]} - + # Handle anyOf (Optional fields in Pydantic v2) if "anyOf" in field_def: any_of = field_def["anyOf"] @@ -225,168 +252,186 @@ class CustomOpenAPISpec: return option # Fallback to string if all else fails return {"type": "string"} - + # Default fallback return {"type": "string"} - + @staticmethod def _expand_field_definition(field_def: Dict[str, Any]) -> Dict[str, Any]: """ Expand a Pydantic field definition for inline use in OpenAPI schema. This creates a full field definition that Swagger UI can render as individual form fields. - + Args: field_def: Pydantic field definition - + Returns: Expanded field definition for OpenAPI schema """ # Return the field definition as-is since Pydantic already provides proper schemas return field_def.copy() - + @staticmethod def add_request_schema( - openapi_schema: Dict[str, Any], - model_class: Type, - schema_name: str, + openapi_schema: Dict[str, Any], + model_class: Type, + schema_name: str, paths: List[str], - operation_name: str + operation_name: str, ) -> Dict[str, Any]: """ Generic method to add a request schema to OpenAPI specification. - + Args: openapi_schema: The OpenAPI schema dict to modify model_class: The Pydantic model class to get schema from schema_name: Name for the schema component paths: List of paths to add the request body to operation_name: Name of the operation for logging (e.g., "chat completion", "embedding") - + Returns: Modified OpenAPI schema """ try: # Get the schema for the model class request_schema = CustomOpenAPISpec.get_pydantic_schema(model_class) - + # Only proceed if we successfully got the schema if request_schema is not None: # Add schema to components - CustomOpenAPISpec.add_schema_to_components(openapi_schema, schema_name, request_schema) - + CustomOpenAPISpec.add_schema_to_components( + openapi_schema, schema_name, request_schema + ) + # Add request body to specified endpoints CustomOpenAPISpec.add_request_body_to_paths( - openapi_schema, - paths, - f"#/components/schemas/{schema_name}" + openapi_schema, paths, f"#/components/schemas/{schema_name}" + ) + + verbose_proxy_logger.debug( + f"Successfully added {schema_name} schema to OpenAPI spec" ) - - verbose_proxy_logger.debug(f"Successfully added {schema_name} schema to OpenAPI spec") else: verbose_proxy_logger.debug(f"Could not get schema for {schema_name}") - + except Exception as e: # If schema addition fails, continue without it - verbose_proxy_logger.debug(f"Failed to add {operation_name} request schema: {str(e)}") - + verbose_proxy_logger.debug( + f"Failed to add {operation_name} request schema: {str(e)}" + ) + return openapi_schema - + @staticmethod - def add_chat_completion_request_schema(openapi_schema: Dict[str, Any]) -> Dict[str, Any]: + def add_chat_completion_request_schema( + openapi_schema: Dict[str, Any] + ) -> Dict[str, Any]: """ Add ProxyChatCompletionRequest schema to chat completion endpoints for documentation. This shows the request body in Swagger without runtime validation. - + Args: openapi_schema: The OpenAPI schema dict to modify - + Returns: Modified OpenAPI schema """ try: from litellm.proxy._types import ProxyChatCompletionRequest - + return CustomOpenAPISpec.add_request_schema( openapi_schema=openapi_schema, model_class=ProxyChatCompletionRequest, schema_name="ProxyChatCompletionRequest", paths=CustomOpenAPISpec.CHAT_COMPLETION_PATHS, - operation_name="chat completion" + operation_name="chat completion", ) except ImportError as e: - verbose_proxy_logger.debug(f"Failed to import ProxyChatCompletionRequest: {str(e)}") + verbose_proxy_logger.debug( + f"Failed to import ProxyChatCompletionRequest: {str(e)}" + ) return openapi_schema - + @staticmethod def add_embedding_request_schema(openapi_schema: Dict[str, Any]) -> Dict[str, Any]: """ Add EmbeddingRequest schema to embedding endpoints for documentation. This shows the request body in Swagger without runtime validation. - + Args: openapi_schema: The OpenAPI schema dict to modify - + Returns: Modified OpenAPI schema """ try: from litellm.types.embedding import EmbeddingRequest - + return CustomOpenAPISpec.add_request_schema( openapi_schema=openapi_schema, model_class=EmbeddingRequest, schema_name="EmbeddingRequest", paths=CustomOpenAPISpec.EMBEDDING_PATHS, - operation_name="embedding" + operation_name="embedding", ) except ImportError as e: verbose_proxy_logger.debug(f"Failed to import EmbeddingRequest: {str(e)}") return openapi_schema - + @staticmethod - def add_responses_api_request_schema(openapi_schema: Dict[str, Any]) -> Dict[str, Any]: + def add_responses_api_request_schema( + openapi_schema: Dict[str, Any] + ) -> Dict[str, Any]: """ Add ResponsesAPIRequestParams schema to responses API endpoints for documentation. This shows the request body in Swagger without runtime validation. - + Args: openapi_schema: The OpenAPI schema dict to modify - + Returns: Modified OpenAPI schema """ try: from litellm.types.llms.openai import ResponsesAPIRequestParams - + return CustomOpenAPISpec.add_request_schema( openapi_schema=openapi_schema, model_class=ResponsesAPIRequestParams, schema_name="ResponsesAPIRequestParams", paths=CustomOpenAPISpec.RESPONSES_API_PATHS, - operation_name="responses API" + operation_name="responses API", ) except ImportError as e: - verbose_proxy_logger.debug(f"Failed to import ResponsesAPIRequestParams: {str(e)}") + verbose_proxy_logger.debug( + f"Failed to import ResponsesAPIRequestParams: {str(e)}" + ) return openapi_schema - + @staticmethod - def add_llm_api_request_schema_body(openapi_schema: Dict[str, Any]) -> Dict[str, Any]: + def add_llm_api_request_schema_body( + openapi_schema: Dict[str, Any] + ) -> Dict[str, Any]: """ Add LLM API request schema bodies to OpenAPI specification for documentation. - + Args: openapi_schema: The base OpenAPI schema - + Returns: OpenAPI schema with added request body schemas """ # Add chat completion request schema - openapi_schema = CustomOpenAPISpec.add_chat_completion_request_schema(openapi_schema) - + openapi_schema = CustomOpenAPISpec.add_chat_completion_request_schema( + openapi_schema + ) + # Add embedding request schema openapi_schema = CustomOpenAPISpec.add_embedding_request_schema(openapi_schema) - + # Add responses API request schema - openapi_schema = CustomOpenAPISpec.add_responses_api_request_schema(openapi_schema) - - return openapi_schema \ No newline at end of file + openapi_schema = CustomOpenAPISpec.add_responses_api_request_schema( + openapi_schema + ) + + return openapi_schema diff --git a/litellm/proxy/common_utils/debug_utils.py b/litellm/proxy/common_utils/debug_utils.py index 0cb7f0058f..6f7038377b 100644 --- a/litellm/proxy/common_utils/debug_utils.py +++ b/litellm/proxy/common_utils/debug_utils.py @@ -18,6 +18,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth router = APIRouter() + # Configure garbage collection thresholds from environment variables def configure_gc_thresholds(): """Configure Python garbage collection thresholds from environment variables.""" @@ -30,13 +31,20 @@ def configure_gc_thresholds(): gc.set_threshold(*thresholds) verbose_proxy_logger.info(f"GC thresholds set to: {thresholds}") else: - verbose_proxy_logger.warning(f"GC threshold not set: {gc_threshold_env}. Expected format: 'gen0,gen1,gen2'") + verbose_proxy_logger.warning( + f"GC threshold not set: {gc_threshold_env}. Expected format: 'gen0,gen1,gen2'" + ) except ValueError as e: - verbose_proxy_logger.warning(f"Failed to parse GC threshold: {gc_threshold_env}. Error: {e}") - + verbose_proxy_logger.warning( + f"Failed to parse GC threshold: {gc_threshold_env}. Error: {e}" + ) + # Log current thresholds current_thresholds = gc.get_threshold() - verbose_proxy_logger.info(f"Current GC thresholds: gen0={current_thresholds[0]}, gen1={current_thresholds[1]}, gen2={current_thresholds[2]}") + verbose_proxy_logger.info( + f"Current GC thresholds: gen0={current_thresholds[0]}, gen1={current_thresholds[1]}, gen2={current_thresholds[2]}" + ) + # Initialize GC configuration configure_gc_thresholds() @@ -59,7 +67,6 @@ async def get_active_tasks_stats(): # Count how many active tasks exist, grouped by coroutine function name. counter = Counter() for idx, task in enumerate(active_tasks): - # reasonable max circuit breaker if idx >= MAX_TASKS_TO_CHECK: break @@ -191,17 +198,17 @@ async def get_memory_summary( ) -> Dict[str, Any]: """ Get simplified memory usage summary for the proxy. - + Returns: - worker_pid: Process ID - status: Overall health based on memory usage - memory: Process memory usage and RAM info - caches: Cache item counts and descriptions - garbage_collector: GC status and pending object counts - + Example usage: curl http://localhost:4000/debug/memory/summary -H "Authorization: Bearer sk-1234" - + For detailed analysis, call GET /debug/memory/details For cache management, use the cache management endpoints """ @@ -210,25 +217,25 @@ async def get_memory_summary( proxy_logging_obj, user_api_key_cache, ) - + # Get process memory info process_memory = {} health_status = "healthy" - + try: import psutil - + process = psutil.Process() memory_info = process.memory_info() memory_mb = memory_info.rss / (1024 * 1024) memory_percent = process.memory_percent() - + process_memory = { "summary": f"{memory_mb:.1f} MB ({memory_percent:.1f}% of system memory)", "ram_usage_mb": round(memory_mb, 2), "system_memory_percent": round(memory_percent, 2), } - + # Check memory health status if memory_percent > 80: health_status = "critical" @@ -236,16 +243,18 @@ async def get_memory_summary( health_status = "warning" else: health_status = "healthy" - + except ImportError: - process_memory["error"] = "Install psutil for memory monitoring: pip install psutil" + process_memory[ + "error" + ] = "Install psutil for memory monitoring: pip install psutil" except Exception as e: process_memory["error"] = str(e) - + # Get cache information caches: Dict[str, Any] = {} total_cache_items = 0 - + try: # User API key cache user_cache_items = len(user_api_key_cache.in_memory_cache.cache_dict) @@ -253,9 +262,9 @@ async def get_memory_summary( caches["user_api_keys"] = { "count": user_cache_items, "count_readable": f"{user_cache_items:,}", - "what_it_stores": "Validated API keys for faster authentication" + "what_it_stores": "Validated API keys for faster authentication", } - + # Router cache if llm_router is not None: router_cache_items = len(llm_router.cache.in_memory_cache.cache_dict) @@ -263,9 +272,9 @@ async def get_memory_summary( caches["llm_responses"] = { "count": router_cache_items, "count_readable": f"{router_cache_items:,}", - "what_it_stores": "LLM responses for identical requests" + "what_it_stores": "LLM responses for identical requests", } - + # Proxy logging cache logging_cache_items = len( proxy_logging_obj.internal_usage_cache.dual_cache.in_memory_cache.cache_dict @@ -274,26 +283,28 @@ async def get_memory_summary( caches["usage_tracking"] = { "count": logging_cache_items, "count_readable": f"{logging_cache_items:,}", - "what_it_stores": "Usage metrics before database write" + "what_it_stores": "Usage metrics before database write", } - + except Exception as e: caches["error"] = str(e) - + # Get garbage collector stats gc_enabled = gc.isenabled() objects_pending = gc.get_count()[0] uncollectable = len(gc.garbage) - + gc_info = { "status": "enabled" if gc_enabled else "disabled", "objects_awaiting_collection": objects_pending, } - + # Add warning if garbage collection issues detected if uncollectable > 0: - gc_info["warning"] = f"{uncollectable} uncollectable objects (possible memory leak)" - + gc_info[ + "warning" + ] = f"{uncollectable} uncollectable objects (possible memory leak)" + return { "worker_pid": os.getpid(), "status": health_status, @@ -314,13 +325,13 @@ def _get_gc_statistics() -> Dict[str, Any]: "generation_0": gc.get_threshold()[0], "generation_1": gc.get_threshold()[1], "generation_2": gc.get_threshold()[2], - "explanation": "Number of allocations before automatic collection for each generation" + "explanation": "Number of allocations before automatic collection for each generation", }, "current_counts": { "generation_0": gc.get_count()[0], "generation_1": gc.get_count()[1], "generation_2": gc.get_count()[2], - "explanation": "Current number of allocated objects in each generation" + "explanation": "Current number of allocated objects in each generation", }, "collection_history": [ { @@ -338,21 +349,17 @@ def _get_object_type_counts(top_n: int) -> Tuple[int, List[Dict[str, Any]]]: """Count objects by type and return total count and top N types.""" type_counts: Counter = Counter() total_objects = 0 - + for obj in gc.get_objects(): total_objects += 1 obj_type = type(obj).__name__ type_counts[obj_type] += 1 - + top_object_types = [ - { - "type": obj_type, - "count": count, - "count_readable": f"{count:,}" - } + {"type": obj_type, "count": count, "count_readable": f"{count:,}"} for obj_type, count in type_counts.most_common(top_n) ] - + return total_objects, top_object_types @@ -362,11 +369,15 @@ def _get_uncollectable_objects_info() -> Dict[str, Any]: return { "count": len(uncollectable), "sample_types": [type(obj).__name__ for obj in uncollectable[:10]], - "warning": "If count > 0, you may have reference cycles preventing garbage collection" if len(uncollectable) > 0 else None, + "warning": "If count > 0, you may have reference cycles preventing garbage collection" + if len(uncollectable) > 0 + else None, } -def _get_cache_memory_stats(user_api_key_cache, llm_router, proxy_logging_obj, redis_usage_cache) -> Dict[str, Any]: +def _get_cache_memory_stats( + user_api_key_cache, llm_router, proxy_logging_obj, redis_usage_cache +) -> Dict[str, Any]: """Calculate memory usage for all caches.""" cache_stats: Dict[str, Any] = {} try: @@ -377,20 +388,26 @@ def _get_cache_memory_stats(user_api_key_cache, llm_router, proxy_logging_obj, r "num_items": len(user_api_key_cache.in_memory_cache.cache_dict), "cache_dict_size_bytes": user_cache_size, "ttl_dict_size_bytes": user_ttl_size, - "total_size_mb": round((user_cache_size + user_ttl_size) / (1024 * 1024), 2), + "total_size_mb": round( + (user_cache_size + user_ttl_size) / (1024 * 1024), 2 + ), } - + # Router cache if llm_router is not None: - router_cache_size = sys.getsizeof(llm_router.cache.in_memory_cache.cache_dict) + router_cache_size = sys.getsizeof( + llm_router.cache.in_memory_cache.cache_dict + ) router_ttl_size = sys.getsizeof(llm_router.cache.in_memory_cache.ttl_dict) cache_stats["llm_router_cache"] = { "num_items": len(llm_router.cache.in_memory_cache.cache_dict), "cache_dict_size_bytes": router_cache_size, "ttl_dict_size_bytes": router_ttl_size, - "total_size_mb": round((router_cache_size + router_ttl_size) / (1024 * 1024), 2), + "total_size_mb": round( + (router_cache_size + router_ttl_size) / (1024 * 1024), 2 + ), } - + # Proxy logging cache logging_cache_size = sys.getsizeof( proxy_logging_obj.internal_usage_cache.dual_cache.in_memory_cache.cache_dict @@ -404,9 +421,11 @@ def _get_cache_memory_stats(user_api_key_cache, llm_router, proxy_logging_obj, r ), "cache_dict_size_bytes": logging_cache_size, "ttl_dict_size_bytes": logging_ttl_size, - "total_size_mb": round((logging_cache_size + logging_ttl_size) / (1024 * 1024), 2), + "total_size_mb": round( + (logging_cache_size + logging_ttl_size) / (1024 * 1024), 2 + ), } - + # Redis cache info if redis_usage_cache is not None: cache_stats["redis_usage_cache"] = { @@ -415,22 +434,29 @@ def _get_cache_memory_stats(user_api_key_cache, llm_router, proxy_logging_obj, r } # Try to get Redis connection pool info if available try: - if hasattr(redis_usage_cache, 'redis_client') and redis_usage_cache.redis_client: - if hasattr(redis_usage_cache.redis_client, 'connection_pool'): + if ( + hasattr(redis_usage_cache, "redis_client") + and redis_usage_cache.redis_client + ): + if hasattr(redis_usage_cache.redis_client, "connection_pool"): pool_info = redis_usage_cache.redis_client.connection_pool # type: ignore cache_stats["redis_usage_cache"]["connection_pool"] = { - "max_connections": pool_info.max_connections if hasattr(pool_info, 'max_connections') else None, - "connection_class": pool_info.connection_class.__name__ if hasattr(pool_info, 'connection_class') else None, + "max_connections": pool_info.max_connections + if hasattr(pool_info, "max_connections") + else None, + "connection_class": pool_info.connection_class.__name__ + if hasattr(pool_info, "connection_class") + else None, } except Exception as e: verbose_proxy_logger.debug(f"Error getting Redis pool info: {e}") else: cache_stats["redis_usage_cache"] = {"enabled": False} - + except Exception as e: verbose_proxy_logger.debug(f"Error calculating cache stats: {e}") cache_stats["error"] = str(e) - + return cache_stats @@ -440,108 +466,115 @@ def _get_router_memory_stats(llm_router) -> Dict[str, Any]: try: if llm_router is not None: # Model list memory size - if hasattr(llm_router, 'model_list') and llm_router.model_list: + if hasattr(llm_router, "model_list") and llm_router.model_list: model_list_size = sys.getsizeof(llm_router.model_list) litellm_router_memory["model_list"] = { "num_models": len(llm_router.model_list), "size_bytes": model_list_size, "size_mb": round(model_list_size / (1024 * 1024), 4), } - + # Model names set - if hasattr(llm_router, 'model_names') and llm_router.model_names: + if hasattr(llm_router, "model_names") and llm_router.model_names: model_names_size = sys.getsizeof(llm_router.model_names) litellm_router_memory["model_names_set"] = { "num_model_groups": len(llm_router.model_names), "size_bytes": model_names_size, "size_mb": round(model_names_size / (1024 * 1024), 4), } - + # Deployment names list - if hasattr(llm_router, 'deployment_names') and llm_router.deployment_names: + if hasattr(llm_router, "deployment_names") and llm_router.deployment_names: deployment_names_size = sys.getsizeof(llm_router.deployment_names) litellm_router_memory["deployment_names"] = { "num_deployments": len(llm_router.deployment_names), "size_bytes": deployment_names_size, "size_mb": round(deployment_names_size / (1024 * 1024), 4), } - + # Deployment latency map - if hasattr(llm_router, 'deployment_latency_map') and llm_router.deployment_latency_map: + if ( + hasattr(llm_router, "deployment_latency_map") + and llm_router.deployment_latency_map + ): latency_map_size = sys.getsizeof(llm_router.deployment_latency_map) litellm_router_memory["deployment_latency_map"] = { "num_tracked_deployments": len(llm_router.deployment_latency_map), "size_bytes": latency_map_size, "size_mb": round(latency_map_size / (1024 * 1024), 4), } - + # Fallback configuration - if hasattr(llm_router, 'fallbacks') and llm_router.fallbacks: + if hasattr(llm_router, "fallbacks") and llm_router.fallbacks: fallbacks_size = sys.getsizeof(llm_router.fallbacks) litellm_router_memory["fallbacks"] = { "num_fallback_configs": len(llm_router.fallbacks), "size_bytes": fallbacks_size, "size_mb": round(fallbacks_size / (1024 * 1024), 4), } - + # Total router object size router_obj_size = sys.getsizeof(llm_router) litellm_router_memory["router_object"] = { "size_bytes": router_obj_size, "size_mb": round(router_obj_size / (1024 * 1024), 4), } - + else: litellm_router_memory = {"note": "Router not initialized"} except Exception as e: verbose_proxy_logger.debug(f"Error getting router memory info: {e}") litellm_router_memory = {"error": str(e)} - + return litellm_router_memory -def _get_process_memory_info(worker_pid: int, include_process_info: bool) -> Optional[Dict[str, Any]]: +def _get_process_memory_info( + worker_pid: int, include_process_info: bool +) -> Optional[Dict[str, Any]]: """Get process-level memory information using psutil.""" if not include_process_info: return None - + try: import psutil - + process = psutil.Process() memory_info = process.memory_info() ram_usage_mb = round(memory_info.rss / (1024 * 1024), 2) virtual_memory_mb = round(memory_info.vms / (1024 * 1024), 2) memory_percent = round(process.memory_percent(), 2) - + return { "pid": worker_pid, "summary": f"Worker PID {worker_pid} using {ram_usage_mb:.1f} MB of RAM ({memory_percent:.1f}% of system memory)", "ram_usage": { "megabytes": ram_usage_mb, - "description": "Actual physical RAM used by this process" + "description": "Actual physical RAM used by this process", }, "virtual_memory": { "megabytes": virtual_memory_mb, - "description": "Total virtual memory allocated (includes swapped memory)" + "description": "Total virtual memory allocated (includes swapped memory)", }, "system_memory_percent": { "percent": memory_percent, - "description": "Percentage of total system RAM being used" + "description": "Percentage of total system RAM being used", }, "open_file_handles": { - "count": process.num_fds() if hasattr(process, "num_fds") else "N/A (Windows)", - "description": "Number of open file descriptors/handles" + "count": process.num_fds() + if hasattr(process, "num_fds") + else "N/A (Windows)", + "description": "Number of open file descriptors/handles", }, "threads": { "count": process.num_threads(), - "description": "Number of active threads in this process" - } + "description": "Number of active threads in this process", + }, } except ImportError: return { "pid": worker_pid, - "error": "psutil not installed. Install with: pip install psutil" + "error": "psutil not installed. Install with: pip install psutil", } except Exception as e: verbose_proxy_logger.debug(f"Error getting process info: {e}") @@ -556,7 +589,7 @@ async def get_memory_details( ) -> Dict[str, Any]: """ Get detailed memory diagnostics for deep debugging. - + Returns: - worker_pid: Process ID - process_memory: RAM usage, virtual memory, file handles, threads @@ -565,14 +598,14 @@ async def get_memory_details( - uncollectable: Objects that can't be garbage collected (potential leaks) - cache_memory: Memory usage of user_api_key, router, and logging caches - router_memory: Memory usage of router components (model_list, deployment_names, etc.) - + Query Parameters: - top_n: Number of top object types to return (default: 20) - include_process_info: Include process-level memory info using psutil (default: true) - + Example usage: curl "http://localhost:4000/debug/memory/details?top_n=30" -H "Authorization: Bearer sk-1234" - + All memory sizes are reported in both bytes and MB. """ from litellm.proxy.proxy_server import ( @@ -581,17 +614,19 @@ async def get_memory_details( user_api_key_cache, redis_usage_cache, ) - + worker_pid = os.getpid() - + # Collect all diagnostics using helper functions gc_stats = _get_gc_statistics() total_objects, top_object_types = _get_object_type_counts(top_n) uncollectable_info = _get_uncollectable_objects_info() - cache_stats = _get_cache_memory_stats(user_api_key_cache, llm_router, proxy_logging_obj, redis_usage_cache) + cache_stats = _get_cache_memory_stats( + user_api_key_cache, llm_router, proxy_logging_obj, redis_usage_cache + ) litellm_router_memory = _get_router_memory_stats(llm_router) process_info = _get_process_memory_info(worker_pid, include_process_info) - + return { "worker_pid": worker_pid, "process_memory": process_info, @@ -616,33 +651,33 @@ async def configure_gc_thresholds_endpoint( ) -> Dict[str, Any]: """ Configure Python garbage collection thresholds. - + Lower thresholds mean more frequent GC cycles (less memory, more CPU overhead). Higher thresholds mean less frequent GC cycles (more memory, less CPU overhead). - + Returns: - message: Confirmation message - previous_thresholds: Old threshold values - new_thresholds: New threshold values - objects_awaiting_collection: Current object count in gen-0 - tip: Hint about when next collection will occur - + Query Parameters: - generation_0: Number of allocations before gen-0 collection (default: 700) - - generation_1: Number of gen-0 collections before gen-1 collection (default: 10) + - generation_1: Number of gen-0 collections before gen-1 collection (default: 10) - generation_2: Number of gen-1 collections before gen-2 collection (default: 10) - + Example for more aggressive collection: curl -X POST "http://localhost:4000/debug/memory/gc/configure?generation_0=500" -H "Authorization: Bearer sk-1234" - + Example for less aggressive collection: curl -X POST "http://localhost:4000/debug/memory/gc/configure?generation_0=1000" -H "Authorization: Bearer sk-1234" - + Monitor memory usage with GET /debug/memory/summary after changes. """ # Get current thresholds for logging old_thresholds = gc.get_threshold() - + # Set new thresholds with error handling try: gc.set_threshold(generation_0, generation_1, generation_2) @@ -653,19 +688,18 @@ async def configure_gc_thresholds_endpoint( except Exception as e: verbose_proxy_logger.error(f"Failed to set GC thresholds: {e}") raise HTTPException( - status_code=500, - detail=f"Failed to set GC thresholds: {str(e)}" + status_code=500, detail=f"Failed to set GC thresholds: {str(e)}" ) - + # Get current object count to show immediate impact current_count = gc.get_count()[0] - + return { "message": "GC thresholds updated", "previous_thresholds": f"{old_thresholds[0]}, {old_thresholds[1]}, {old_thresholds[2]}", "new_thresholds": f"{generation_0}, {generation_1}, {generation_2}", "objects_awaiting_collection": current_count, - "tip": f"Next collection will run after {generation_0 - current_count} more allocations" + "tip": f"Next collection will run after {generation_0 - current_count} more allocations", } diff --git a/litellm/proxy/common_utils/encrypt_decrypt_utils.py b/litellm/proxy/common_utils/encrypt_decrypt_utils.py index 6c47d220c4..a5da5798f4 100644 --- a/litellm/proxy/common_utils/encrypt_decrypt_utils.py +++ b/litellm/proxy/common_utils/encrypt_decrypt_utils.py @@ -60,7 +60,6 @@ def decrypt_value_helper( # if it's not str - do not decrypt it, return the value return value except Exception as e: - error_message = f"Error decrypting value for key: {key}, Did your master_key/salt key change recently? \nError: {str(e)}\nSet permanent salt key - https://docs.litellm.ai/docs/proxy/prod#5-set-litellm-salt-key" if exception_type == "debug": verbose_proxy_logger.debug(error_message) diff --git a/litellm/proxy/common_utils/get_routes.py b/litellm/proxy/common_utils/get_routes.py index bf3773037e..743c3b6e9d 100644 --- a/litellm/proxy/common_utils/get_routes.py +++ b/litellm/proxy/common_utils/get_routes.py @@ -24,14 +24,12 @@ class GetRoutes: "methods": getattr(route, "methods", None), "name": getattr(route, "name", None), "endpoint": ( - endpoint_route.__name__ - if getattr(route, "endpoint", None) - else None + endpoint_route.__name__ if getattr(route, "endpoint", None) else None ), } routes.append(route_info) return routes - + @staticmethod def get_routes_for_mounted_app( route: BaseRoute, @@ -40,17 +38,19 @@ class GetRoutes: Get routes for a mounted sub-application. """ routes: List[Dict[str, Any]] = [] - mount_path = getattr(route, 'path', '') - sub_app = getattr(route, 'app', None) - if sub_app and hasattr(sub_app, 'routes'): + mount_path = getattr(route, "path", "") + sub_app = getattr(route, "app", None) + if sub_app and hasattr(sub_app, "routes"): for sub_route in sub_app.routes: # Get endpoint - either from endpoint attribute or app attribute - endpoint_func = getattr(sub_route, "endpoint", None) or getattr(sub_route, "app", None) - + endpoint_func = getattr(sub_route, "endpoint", None) or getattr( + sub_route, "app", None + ) + if endpoint_func is not None: sub_route_path = getattr(sub_route, "path", "") - full_path = mount_path.rstrip('/') + sub_route_path - + full_path = mount_path.rstrip("/") + sub_route_path + route_info = { "path": full_path, "methods": getattr(sub_route, "methods", ["GET", "POST"]), @@ -60,7 +60,6 @@ class GetRoutes: } routes.append(route_info) return routes - @staticmethod def _safe_get_endpoint_name(endpoint_function: Any) -> Optional[str]: @@ -68,12 +67,16 @@ class GetRoutes: Safely get the name of the endpoint function. """ try: - if hasattr(endpoint_function, '__name__'): - return getattr(endpoint_function, '__name__') - elif hasattr(endpoint_function, '__class__') and hasattr(endpoint_function.__class__, '__name__'): - return getattr(endpoint_function.__class__, '__name__') + if hasattr(endpoint_function, "__name__"): + return getattr(endpoint_function, "__name__") + elif hasattr(endpoint_function, "__class__") and hasattr( + endpoint_function.__class__, "__name__" + ): + return getattr(endpoint_function.__class__, "__name__") else: return None except Exception: - verbose_logger.exception(f"Error getting endpoint name for route: {endpoint_function}") - return None \ No newline at end of file + verbose_logger.exception( + f"Error getting endpoint name for route: {endpoint_function}" + ) + return None diff --git a/litellm/proxy/common_utils/html_forms/cli_sso_success.py b/litellm/proxy/common_utils/html_forms/cli_sso_success.py index 7da140505a..51f0775d90 100644 --- a/litellm/proxy/common_utils/html_forms/cli_sso_success.py +++ b/litellm/proxy/common_utils/html_forms/cli_sso_success.py @@ -4,11 +4,11 @@ from litellm.proxy.common_utils.banner import LITELLM_BANNER def render_cli_sso_success_page() -> str: """ Renders the CLI SSO authentication success page with minimal styling - + Returns: str: HTML content for the success page """ - + html_content = f""" @@ -204,4 +204,4 @@ def render_cli_sso_success_page() -> str: """ - return html_content \ No newline at end of file + return html_content diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index dc7b25ea09..1dd2526212 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -70,7 +70,9 @@ async def _read_request_body(request: Optional[Request]) -> Dict: parsed_body = json.loads(body_str) except json.JSONDecodeError: # If both orjson and json.loads fail, throw a proper error - verbose_proxy_logger.error(f"Invalid JSON payload received: {str(e)}") + verbose_proxy_logger.error( + f"Invalid JSON payload received: {str(e)}" + ) raise ProxyException( message=f"Invalid JSON payload: {str(e)}", type="invalid_request_error", @@ -106,6 +108,7 @@ def _safe_get_request_parsed_body(request: Optional[Request]) -> Optional[dict]: return {key: parsed_body[key] for key in accepted_keys} return None + def _safe_get_request_query_params(request: Optional[Request]) -> Dict: if request is None: return {} @@ -119,6 +122,7 @@ def _safe_get_request_query_params(request: Optional[Request]) -> Dict: ) return {} + def _safe_set_request_parsed_body( request: Optional[Request], parsed_body: dict, @@ -257,16 +261,16 @@ async def convert_upload_files_to_file_data( ) -> Dict[str, Any]: """ Convert FastAPI UploadFile objects to file data tuples for litellm. - + Converts UploadFile objects to tuples of (filename, content, content_type) which is the format expected by httpx and litellm's HTTP handlers. - + Args: form_data: Dictionary containing form data with potential UploadFile objects - + Returns: Dictionary with UploadFile objects converted to file data tuples - + Example: ```python form_data = await get_form_data(request) @@ -304,9 +308,10 @@ async def get_request_body(request: Request) -> Dict[str, Any]: if request.method == "POST": if request.headers.get("content-type", "") == "application/json": return await _read_request_body(request) - elif ( - "multipart/form-data" in request.headers.get("content-type", "") - or "application/x-www-form-urlencoded" in request.headers.get("content-type", "") + elif "multipart/form-data" in request.headers.get( + "content-type", "" + ) or "application/x-www-form-urlencoded" in request.headers.get( + "content-type", "" ): return await get_form_data(request) else: @@ -317,25 +322,24 @@ async def get_request_body(request: Request) -> Dict[str, Any]: def extract_nested_form_metadata( - form_data: Dict[str, Any], - prefix: str = "litellm_metadata[" + form_data: Dict[str, Any], prefix: str = "litellm_metadata[" ) -> Dict[str, Any]: """ Extract nested metadata from form data with bracket notation. - + Handles form data that uses bracket notation to represent nested dictionaries, such as litellm_metadata[spend_logs_metadata][owner] = "value". - + This is commonly encountered when SDKs or clients send form data with nested structures using bracket notation instead of JSON. - + Args: form_data: Dictionary containing form data (from request.form()) prefix: The prefix to look for in form keys (default: "litellm_metadata[") - + Returns: Dictionary with nested structure reconstructed from bracket notation - + Example: Input form_data: { @@ -344,7 +348,7 @@ def extract_nested_form_metadata( "litellm_metadata[tags]": "production", "other_field": "value" } - + Output: { "spend_logs_metadata": { @@ -356,36 +360,36 @@ def extract_nested_form_metadata( """ if not form_data: return {} - + metadata: Dict[str, Any] = {} - + for key, value in form_data.items(): # Skip keys that don't start with the prefix if not isinstance(key, str) or not key.startswith(prefix): continue - + # Skip UploadFile objects - they should not be in metadata if isinstance(value, UploadFile): verbose_proxy_logger.warning( f"Skipping UploadFile in metadata extraction for key: {key}" ) continue - + # Extract the nested path from bracket notation # Example: "litellm_metadata[spend_logs_metadata][owner]" -> ["spend_logs_metadata", "owner"] try: # Remove the prefix and strip trailing ']' path_string = key.replace(prefix, "").rstrip("]") - + # Split by "][" to get individual path parts parts = path_string.split("][") - + if not parts or not parts[0]: verbose_proxy_logger.warning( f"Invalid metadata key format (empty path): {key}" ) continue - + # Navigate/create nested dictionary structure current = metadata for part in parts[:-1]: @@ -403,23 +407,21 @@ def extract_nested_form_metadata( verbose_proxy_logger.warning( f"Cannot set value - parent is not a dict for key: {key}" ) - + except Exception as e: - verbose_proxy_logger.error( - f"Error parsing metadata key '{key}': {str(e)}" - ) + verbose_proxy_logger.error(f"Error parsing metadata key '{key}': {str(e)}") continue - + return metadata def get_tags_from_request_body(request_body: dict) -> List[str]: """ Extract tags from request body metadata. - + Args: request_body: The request body dictionary - + Returns: List of tag names (strings), empty list if no valid tags found """ @@ -440,30 +442,28 @@ def get_tags_from_request_body(request_body: dict) -> List[str]: return [tag for tag in combined_tags if isinstance(tag, str)] -def populate_request_with_path_params( - request_data: dict, request: Request -) -> dict: +def populate_request_with_path_params(request_data: dict, request: Request) -> dict: """ Copy FastAPI path params and query params into the request payload so downstream checks (e.g. vector store RBAC, organization RBAC) see them the same way as body params. - + Since path_params may not be available during dependency injection, we parse the URL path directly for known patterns. - + Args: request_data: The request data dictionary to populate request: The FastAPI Request object - + Returns: dict: Updated request_data with path parameters and query parameters added - """ + """ # Add query parameters to request_data (for GET requests, etc.) query_params = _safe_get_request_query_params(request) if query_params: for key, value in query_params.items(): # Don't overwrite existing values from request body request_data.setdefault(key, value) - + # Try to get path_params if available (sometimes populated by FastAPI) path_params = getattr(request, "path_params", None) if isinstance(path_params, dict) and path_params: @@ -494,7 +494,7 @@ def _add_vector_store_id_from_path(request_data: dict, request: Request) -> None Parse the request path to find /vector_stores/{vector_store_id}/... segments. When found, ensure both vector_store_id and vector_store_ids are populated. - + Args: request_data: The request data dictionary to populate request: The FastAPI Request object diff --git a/litellm/proxy/common_utils/load_config_utils.py b/litellm/proxy/common_utils/load_config_utils.py index b8533ad00d..3c7329c2c0 100644 --- a/litellm/proxy/common_utils/load_config_utils.py +++ b/litellm/proxy/common_utils/load_config_utils.py @@ -29,7 +29,7 @@ def get_file_contents_from_s3(bucket_name, object_key): # Read the file contents and directly parse YAML file_contents = response["Body"].read().decode("utf-8") verbose_proxy_logger.debug("File contents retrieved from S3") - + # Parse YAML directly from string config = yaml.safe_load(file_contents) return config @@ -71,12 +71,12 @@ def download_python_file_from_s3( ) -> bool: """ Download a Python file from S3 and save it to local filesystem. - + Args: bucket_name (str): S3 bucket name object_key (str): S3 object key (file path in bucket) local_file_path (str): Local path where file should be saved - + Returns: bool: True if successful, False otherwise """ @@ -85,6 +85,7 @@ def download_python_file_from_s3( from botocore.credentials import Credentials from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + base_aws_llm = BaseAWSLLM() credentials: Credentials = base_aws_llm.get_credentials() @@ -94,24 +95,26 @@ def download_python_file_from_s3( aws_secret_access_key=credentials.secret_key, aws_session_token=credentials.token, ) - + verbose_proxy_logger.debug( f"Downloading Python file {object_key} from S3 bucket: {bucket_name}" ) response = s3_client.get_object(Bucket=bucket_name, Key=object_key) - + # Read the file contents file_contents = response["Body"].read().decode("utf-8") verbose_proxy_logger.debug(f"File contents: {file_contents}") - + # Ensure directory exists os.makedirs(os.path.dirname(local_file_path), exist_ok=True) - + # Write to local file - with open(local_file_path, 'w') as f: + with open(local_file_path, "w") as f: f.write(file_contents) - - verbose_proxy_logger.debug(f"Python file downloaded successfully to {local_file_path}") + + verbose_proxy_logger.debug( + f"Python file downloaded successfully to {local_file_path}" + ) return True except ImportError as e: @@ -129,12 +132,12 @@ async def download_python_file_from_gcs( ) -> bool: """ Download a Python file from GCS and save it to local filesystem. - + Args: bucket_name (str): GCS bucket name object_key (str): GCS object key (file path in bucket) local_file_path (str): Local path where file should be saved - + Returns: bool: True if successful, False otherwise """ @@ -147,22 +150,26 @@ async def download_python_file_from_gcs( file_contents = await gcs_bucket.download_gcs_object(object_key) if file_contents is None: raise Exception(f"File contents are None for {object_key}") - + # file_contents is a bytes object, decode it file_contents = file_contents.decode("utf-8") - + # Ensure directory exists os.makedirs(os.path.dirname(local_file_path), exist_ok=True) - + # Write to local file - with open(local_file_path, 'w') as f: + with open(local_file_path, "w") as f: f.write(file_contents) - - verbose_proxy_logger.debug(f"Python file downloaded successfully to {local_file_path}") + + verbose_proxy_logger.debug( + f"Python file downloaded successfully to {local_file_path}" + ) return True except Exception as e: - verbose_proxy_logger.exception(f"Error downloading Python file from GCS: {str(e)}") + verbose_proxy_logger.exception( + f"Error downloading Python file from GCS: {str(e)}" + ) return False diff --git a/litellm/proxy/common_utils/openai_endpoint_utils.py b/litellm/proxy/common_utils/openai_endpoint_utils.py index bedaf31e75..6df5491f37 100644 --- a/litellm/proxy/common_utils/openai_endpoint_utils.py +++ b/litellm/proxy/common_utils/openai_endpoint_utils.py @@ -61,6 +61,7 @@ def get_custom_llm_provider_from_request_query(request: Request) -> Optional[str return request.query_params["custom_llm_provider"] return None + def get_custom_llm_provider_from_request_headers(request: Request) -> Optional[str]: """ Get the `custom_llm_provider` from the request header `custom-llm-provider` diff --git a/litellm/proxy/common_utils/openapi_schema_compat.py b/litellm/proxy/common_utils/openapi_schema_compat.py index f97cee9415..83b1751c94 100644 --- a/litellm/proxy/common_utils/openapi_schema_compat.py +++ b/litellm/proxy/common_utils/openapi_schema_compat.py @@ -19,17 +19,17 @@ def get_openapi_schema_with_compat( ) -> Dict[str, Any]: """ Generate OpenAPI schema with compatibility handling for FastAPI 0.120+. - + This function patches Pydantic's schema generation to handle non-serializable types like openai.Timeout that cause PydanticSchemaGenerationError in FastAPI 0.120+. - + Args: get_openapi_func: The FastAPI get_openapi function title: API title version: API version description: API description routes: List of routes - + Returns: OpenAPI schema dictionary """ @@ -41,18 +41,21 @@ def get_openapi_schema_with_compat( # Store original method original_unknown_type_schema = GenerateSchema._unknown_type_schema - + def patched_unknown_type_schema(self, obj): """Patch to handle openai.Timeout and other non-serializable types""" # Check if it's openai.Timeout or similar types obj_str = str(obj) - obj_module = getattr(obj, '__module__', '') - - if (obj_module == 'openai' and 'Timeout' in obj_str) or \ - (hasattr(obj, '__name__') and obj.__name__ == 'Timeout' and obj_module == 'openai'): + obj_module = getattr(obj, "__module__", "") + + if (obj_module == "openai" and "Timeout" in obj_str) or ( + hasattr(obj, "__name__") + and obj.__name__ == "Timeout" + and obj_module == "openai" + ): # Return a simple string schema for Timeout types return core_schema.str_schema() - + # For other unknown types, try to return a default schema # This prevents the error from propagating try: @@ -60,10 +63,10 @@ def get_openapi_schema_with_compat( except Exception: # Last resort: return string schema return core_schema.str_schema() - + # Apply patch - setattr(GenerateSchema, '_unknown_type_schema', patched_unknown_type_schema) - + setattr(GenerateSchema, "_unknown_type_schema", patched_unknown_type_schema) + try: openapi_schema = get_openapi_func( title=title, @@ -73,13 +76,17 @@ def get_openapi_schema_with_compat( ) finally: # Restore original method - setattr(GenerateSchema, '_unknown_type_schema', original_unknown_type_schema) - + setattr( + GenerateSchema, "_unknown_type_schema", original_unknown_type_schema + ) + return openapi_schema - + except (ImportError, AttributeError) as e: # If patching fails, try normal generation with error handling - verbose_proxy_logger.debug(f"Could not patch Pydantic schema generation: {e}. Trying normal generation.") + verbose_proxy_logger.debug( + f"Could not patch Pydantic schema generation: {e}. Trying normal generation." + ) try: return get_openapi_func( title=title, @@ -91,16 +98,24 @@ def get_openapi_schema_with_compat( # Check if it's a PydanticSchemaGenerationError by checking the error type name # This avoids import issues if PydanticSchemaGenerationError is not available error_type_name = type(pydantic_error).__name__ - if error_type_name == "PydanticSchemaGenerationError" or "PydanticSchemaGenerationError" in str(type(pydantic_error)): + if ( + error_type_name == "PydanticSchemaGenerationError" + or "PydanticSchemaGenerationError" in str(type(pydantic_error)) + ): # If we still get the error, log it and return minimal schema - verbose_proxy_logger.warning(f"PydanticSchemaGenerationError during schema generation: {pydantic_error}") + verbose_proxy_logger.warning( + f"PydanticSchemaGenerationError during schema generation: {pydantic_error}" + ) return { "openapi": "3.0.0", - "info": {"title": title, "version": version, "description": description or ""}, + "info": { + "title": title, + "version": version, + "description": description or "", + }, "paths": {}, "components": {"schemas": {}}, } else: # Re-raise if it's a different error raise - diff --git a/litellm/proxy/common_utils/performance_utils.py b/litellm/proxy/common_utils/performance_utils.py index 5bfa6f31c7..6853a86d1d 100644 --- a/litellm/proxy/common_utils/performance_utils.py +++ b/litellm/proxy/common_utils/performance_utils.py @@ -37,7 +37,7 @@ def _should_sample(profile_sampling_rate: float) -> bool: return True # Always sample elif profile_sampling_rate <= 0.0: return False # Never sample - + # Use deterministic sampling based on counter for consistent rate global _sample_counter with _sample_counter_lock: @@ -54,7 +54,9 @@ def _start_profiling(profile_sampling_rate: float) -> None: if _profiler is None: _profiler = cProfile.Profile() _profiler.enable() - verbose_proxy_logger.info(f"Profiling started with sampling rate: {profile_sampling_rate}") + verbose_proxy_logger.info( + f"Profiling started with sampling rate: {profile_sampling_rate}" + ) def _start_profiling_for_request(profile_sampling_rate: float) -> bool: @@ -88,19 +90,21 @@ def _save_stats(profile_file: PathLib) -> None: def profile_endpoint(sampling_rate: float = 1.0): """Decorator to sample endpoint hits and save to a profile file. - + Args: sampling_rate: Rate of requests to profile (0.0 to 1.0) - 1.0: Profile all requests (100%) - 0.1: Profile 1 in 10 requests (10%) - 0.0: Profile no requests (0%) """ + def decorator(func): def set_last_profile_path(path: PathLib) -> None: global _last_profile_file_path _last_profile_file_path = path if inspect.iscoroutinefunction(func): + @functools.wraps(func) async def async_wrapper(*args, **kwargs): is_sampling = _start_profiling_for_request(sampling_rate) @@ -115,8 +119,10 @@ def profile_endpoint(sampling_rate: float = 1.0): if is_sampling: _save_stats(file_path_obj) raise + return async_wrapper else: + @functools.wraps(func) def sync_wrapper(*args, **kwargs): is_sampling = _start_profiling_for_request(sampling_rate) @@ -131,19 +137,21 @@ def profile_endpoint(sampling_rate: float = 1.0): if is_sampling: _save_stats(file_path_obj) raise + return sync_wrapper + return decorator def enable_line_profiler() -> None: """Enable line_profiler for dynamic function wrapping. - + Raises: ImportError: If line_profiler is not available """ global _line_profiler from line_profiler import LineProfiler # Will raise ImportError if not available - + with _line_profiler_lock: if _line_profiler is None: _line_profiler = LineProfiler() @@ -152,11 +160,11 @@ def enable_line_profiler() -> None: def wrap_function_with_line_profiler(module: Any, function_name: str) -> bool: """Dynamically wrap a function with line_profiler. - + Args: module: The module containing the function function_name: Name of the function to wrap - + Returns: True if wrapping was successful, False otherwise """ @@ -164,10 +172,10 @@ def wrap_function_with_line_profiler(module: Any, function_name: str) -> bool: enable_line_profiler() # May raise ImportError if not available except ImportError: return False - + if _line_profiler is None: return False - + try: original_function = getattr(module, function_name, None) if original_function is None: @@ -175,15 +183,15 @@ def wrap_function_with_line_profiler(module: Any, function_name: str) -> bool: f"Function {function_name} not found in module {module.__name__}" ) return False - + # Store original function if not already wrapped if function_name not in _wrapped_functions: _wrapped_functions[function_name] = original_function - + # Wrap with line_profiler profiled_function = _line_profiler(original_function) setattr(module, function_name, profiled_function) - + verbose_proxy_logger.info( f"Wrapped {module.__name__}.{function_name} with line_profiler" ) @@ -197,68 +205,66 @@ def wrap_function_with_line_profiler(module: Any, function_name: str) -> bool: def wrap_function_directly(func: Callable) -> Callable: """Wrap a function directly with line_profiler. - + This is the recommended way to profile functions, especially closures or functions created dynamically (like wrapper_async in litellm/utils.py). - + Args: func: The function to wrap - + Returns: The wrapped function that will be profiled when called - + Raises: ImportError: If line_profiler is not available RuntimeError: If line_profiler cannot be enabled or function cannot be wrapped """ import warnings - + enable_line_profiler() # Will raise ImportError if not available - + if _line_profiler is None: raise RuntimeError("Line profiler was not initialized") - + # Suppress warnings about __wrapped__ - we intentionally want to profile the wrapper with warnings.catch_warnings(): - warnings.filterwarnings('ignore', message='.*__wrapped__.*', category=UserWarning) + warnings.filterwarnings( + "ignore", message=".*__wrapped__.*", category=UserWarning + ) # Add function to line_profiler and wrap it _line_profiler.add_function(func) profiled_function = _line_profiler(func) - - verbose_proxy_logger.info( - f"Wrapped function {func.__name__} with line_profiler" - ) + + verbose_proxy_logger.info(f"Wrapped function {func.__name__} with line_profiler") return profiled_function def collect_line_profiler_stats(output_file: Optional[str] = None) -> None: """Collect and save line_profiler statistics. - + This can be called manually to collect stats at any time, or it's automatically called on shutdown if register_shutdown_handler() was used. - + Args: output_file: Optional path to save stats. If None, prints to stdout. """ global _line_profiler - + with _line_profiler_lock: if _line_profiler is None: verbose_proxy_logger.debug("Line profiler not enabled, nothing to collect") return - + try: if output_file: # Save to file output_path = PathLib(output_file) _line_profiler.dump_stats(str(output_path)) - verbose_proxy_logger.info( - f"Line profiler stats saved to {output_path}" - ) + verbose_proxy_logger.info(f"Line profiler stats saved to {output_path}") else: # Print to stdout from io import StringIO - + stream = StringIO() _line_profiler.print_stats(stream=stream) stats_output = stream.getvalue() @@ -269,20 +275,22 @@ def collect_line_profiler_stats(output_file: Optional[str] = None) -> None: def register_shutdown_handler(output_file: Optional[str] = None) -> None: """Register a shutdown handler to collect line_profiler stats. - + This registers an atexit handler that will automatically save profiling statistics when the Python process exits. Safe to call multiple times (only registers once). - + Args: output_file: Optional path to save stats on shutdown. Defaults to 'line_profile_stats.lprof' """ if output_file is None: output_file = "line_profile_stats.lprof" - + def shutdown_handler(): collect_line_profiler_stats(output_file=output_file) - + atexit.register(shutdown_handler) - verbose_proxy_logger.debug(f"Registered line_profiler shutdown handler for {output_file}") + verbose_proxy_logger.debug( + f"Registered line_profiler shutdown handler for {output_file}" + ) diff --git a/litellm/proxy/common_utils/rbac_utils.py b/litellm/proxy/common_utils/rbac_utils.py index 5ce77ec836..b54b5e8f45 100644 --- a/litellm/proxy/common_utils/rbac_utils.py +++ b/litellm/proxy/common_utils/rbac_utils.py @@ -35,7 +35,11 @@ async def check_feature_access_for_user( ): return - from litellm.proxy.proxy_server import general_settings, prisma_client, user_api_key_cache + from litellm.proxy.proxy_server import ( + general_settings, + prisma_client, + user_api_key_cache, + ) disable_flag = f"disable_{feature_name}_for_internal_users" allow_team_admins_flag = f"allow_{feature_name}_for_team_admins" @@ -46,7 +50,9 @@ async def check_feature_access_for_user( # Feature is disabled. Check if team/org admins are exempted. if general_settings.get(allow_team_admins_flag, False): - from litellm.proxy.management_endpoints.common_utils import _user_has_admin_privileges + from litellm.proxy.management_endpoints.common_utils import ( + _user_has_admin_privileges, + ) is_admin = await _user_has_admin_privileges( user_api_key_dict=user_api_key_dict, diff --git a/litellm/proxy/common_utils/realtime_utils.py b/litellm/proxy/common_utils/realtime_utils.py index 4af7ad2514..ee31a902ed 100644 --- a/litellm/proxy/common_utils/realtime_utils.py +++ b/litellm/proxy/common_utils/realtime_utils.py @@ -11,5 +11,3 @@ def _realtime_request_body(model: Optional[str]) -> bytes: string formatting work while keeping memory usage bounded. """ return f'{{"model": "{model or ""}"}}'.encode() - - diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 8ce73d29c8..674214b19e 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -616,4 +616,4 @@ class ResetBudgetJob: await ResetBudgetJob._reset_budget_common( item=key, current_time=current_time, item_type="key" ) - return key \ No newline at end of file + return key diff --git a/litellm/proxy/container_endpoints/endpoints.py b/litellm/proxy/container_endpoints/endpoints.py index 8581603eea..45870d73d4 100644 --- a/litellm/proxy/container_endpoints/endpoints.py +++ b/litellm/proxy/container_endpoints/endpoints.py @@ -91,10 +91,10 @@ async def create_container( or await get_custom_llm_provider_from_request_body(request=request) or "openai" ) - + # Add custom_llm_provider to data data["custom_llm_provider"] = custom_llm_provider - + # Process request using ProxyBaseLLMRequestProcessing processor = ProxyBaseLLMRequestProcessing(data=data) try: @@ -184,7 +184,7 @@ async def list_containers( or get_custom_llm_provider_from_request_query(request=request) or "openai" ) - + # Add custom_llm_provider to data data["custom_llm_provider"] = custom_llm_provider @@ -278,7 +278,7 @@ async def retrieve_container( or get_custom_llm_provider_from_request_query(request=request) or "openai" ) - + # Add custom_llm_provider to data data["custom_llm_provider"] = custom_llm_provider @@ -372,7 +372,7 @@ async def delete_container( or get_custom_llm_provider_from_request_query(request=request) or "openai" ) - + # Add custom_llm_provider to data data["custom_llm_provider"] = custom_llm_provider diff --git a/litellm/proxy/container_endpoints/handler_factory.py b/litellm/proxy/container_endpoints/handler_factory.py index dc10e39bc9..078f0c9bc4 100644 --- a/litellm/proxy/container_endpoints/handler_factory.py +++ b/litellm/proxy/container_endpoints/handler_factory.py @@ -38,17 +38,26 @@ def _get_container_provider_config(custom_llm_provider: str): """Get the container provider config for the given provider.""" if custom_llm_provider == "openai": from litellm.llms.openai.containers.transformation import OpenAIContainerConfig + return OpenAIContainerConfig() else: - raise ValueError(f"Container API not supported for provider: {custom_llm_provider}") + raise ValueError( + f"Container API not supported for provider: {custom_llm_provider}" + ) -def _create_handler_for_path_params(path_params: List[str], route_type: str, returns_binary: bool = False, is_multipart: bool = False): +def _create_handler_for_path_params( + path_params: List[str], + route_type: str, + returns_binary: bool = False, + is_multipart: bool = False, +): """ Dynamically create a handler with the correct path parameter signature. """ # For binary content endpoints, use a different handler if returns_binary and path_params == ["container_id", "file_id"]: + async def handler_binary_content( request: Request, container_id: str, @@ -61,10 +70,12 @@ def _create_handler_for_path_params(path_params: List[str], route_type: str, ret file_id=file_id, user_api_key_dict=user_api_key_dict, ) + return handler_binary_content - + # For multipart file upload endpoints if is_multipart: + async def handler_multipart_upload( request: Request, container_id: str, @@ -78,10 +89,12 @@ def _create_handler_for_path_params(path_params: List[str], route_type: str, ret route_type=route_type, container_id=container_id, ) + return handler_multipart_upload - + # Create handlers for different path parameter combinations if path_params == ["container_id"]: + async def handler_container_id( request: Request, container_id: str, @@ -95,9 +108,11 @@ def _create_handler_for_path_params(path_params: List[str], route_type: str, ret route_type=route_type, path_params={"container_id": container_id}, ) + return handler_container_id - + elif path_params == ["container_id", "file_id"]: + async def handler_container_file( request: Request, container_id: str, @@ -112,8 +127,9 @@ def _create_handler_for_path_params(path_params: List[str], route_type: str, ret route_type=route_type, path_params={"container_id": container_id, "file_id": file_id}, ) + return handler_container_file - + else: # Fallback for no path params async def handler_no_params( @@ -128,6 +144,7 @@ def _create_handler_for_path_params(path_params: List[str], route_type: str, ret route_type=route_type, path_params={}, ) + return handler_no_params @@ -139,7 +156,7 @@ async def _process_binary_request( ): """ Process binary content requests using the proper transformation pattern. - + This uses the provider config transformations and llm_http_handler to maintain consistency with the established pattern. """ @@ -153,13 +170,13 @@ async def _process_binary_request( or get_custom_llm_provider_from_request_query(request=request) or "openai" ) - + # Get the provider config container_provider_config = _get_container_provider_config(custom_llm_provider) - + # Build litellm_params - credentials are resolved by provider config from env litellm_params = GenericLiteLLMParams() - + # Create logging object logging_obj = Logging( model="container-file-content", @@ -170,10 +187,10 @@ async def _process_binary_request( litellm_call_id="", function_id="", ) - + # Use the HTTP handler to make the request handler = BaseLLMHTTPHandler() - + try: content = await handler.async_container_file_content_handler( container_id=container_id, @@ -182,7 +199,7 @@ async def _process_binary_request( litellm_params=litellm_params, logging_obj=logging_obj, ) - + # Determine content type based on common file extensions in the file_id content_type = "application/octet-stream" file_id_lower = file_id.lower() @@ -200,12 +217,12 @@ async def _process_binary_request( content_type = "text/plain" elif ".pdf" in file_id_lower: content_type = "application/pdf" - + return Response( content=content, media_type=content_type, ) - + except Exception as e: raise e @@ -239,16 +256,17 @@ async def _process_multipart_upload_request( # Parse multipart form data and convert files form_data = await get_form_data(request) data = await convert_upload_files_to_file_data(form_data) - + if "file" not in data: from fastapi import HTTPException + raise HTTPException(status_code=400, detail="Missing required 'file' field") - + # convert_upload_files_to_file_data returns list of tuples, extract single file file_list = data["file"] if isinstance(file_list, list) and len(file_list) > 0: data["file"] = file_list[0] - + data["container_id"] = container_id custom_llm_provider = ( @@ -354,12 +372,12 @@ async def _process_request( def register_container_file_endpoints(router: APIRouter) -> None: """ Register ALL container file endpoints from JSON config to the router. - + This single function registers all endpoints defined in endpoints.json, eliminating the need for manual endpoint definitions. """ config = _load_endpoints_config() - + for endpoint_config in config["endpoints"]: path = endpoint_config["path"] method = endpoint_config["method"].lower() @@ -367,13 +385,15 @@ def register_container_file_endpoints(router: APIRouter) -> None: route_type = endpoint_config["async_name"] returns_binary = endpoint_config.get("returns_binary", False) is_multipart = endpoint_config.get("is_multipart", False) - + # Create handler with correct signature for path params - handler = _create_handler_for_path_params(path_params, route_type, returns_binary, is_multipart) - + handler = _create_handler_for_path_params( + path_params, route_type, returns_binary, is_multipart + ) + # Register routes route_method = getattr(router, method) - + # For binary endpoints, don't use ORJSONResponse if returns_binary: # Register both /v1/... and /... paths without JSON response class @@ -382,7 +402,7 @@ def register_container_file_endpoints(router: APIRouter) -> None: dependencies=[Depends(user_api_key_auth)], tags=["containers"], )(handler) - + route_method( path, dependencies=[Depends(user_api_key_auth)], @@ -396,7 +416,7 @@ def register_container_file_endpoints(router: APIRouter) -> None: response_class=ORJSONResponse, tags=["containers"], )(handler) - + route_method( path, dependencies=[Depends(user_api_key_auth)], diff --git a/litellm/proxy/credential_endpoints/endpoints.py b/litellm/proxy/credential_endpoints/endpoints.py index 9f228bb118..64f860fc4f 100644 --- a/litellm/proxy/credential_endpoints/endpoints.py +++ b/litellm/proxy/credential_endpoints/endpoints.py @@ -21,11 +21,15 @@ router = APIRouter() class CredentialHelperUtils: @staticmethod - def encrypt_credential_values(credential: CredentialItem, new_encryption_key: Optional[str] = None) -> CredentialItem: + def encrypt_credential_values( + credential: CredentialItem, new_encryption_key: Optional[str] = None + ) -> CredentialItem: """Encrypt values in credential.credential_values and add to DB""" encrypted_credential_values = {} for key, value in (credential.credential_values or {}).items(): - encrypted_credential_values[key] = encrypt_value_helper(value, new_encryption_key) + encrypted_credential_values[key] = encrypt_value_helper( + value, new_encryption_key + ) # Return a new object to avoid mutating the caller's credential, which # is kept in memory and should remain unencrypted. @@ -142,17 +146,49 @@ async def get_credentials( tags=["credential management"], response_model=CredentialItem, ) +async def get_credential_by_name( + request: Request, + fastapi_response: Response, + credential_name: str = Path( + ..., description="The credential name, percent-decoded; may contain slashes" + ), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + [BETA] endpoint. This might change unexpectedly. + """ + try: + for credential in litellm.credential_list: + if credential.credential_name == credential_name: + masked_credential = CredentialItem( + credential_name=credential.credential_name, + credential_values=_get_masked_values( + credential.credential_values, + unmasked_length=4, + number_of_asterisks=4, + ), + credential_info=credential.credential_info, + ) + return masked_credential + raise HTTPException( + status_code=404, + detail="Credential not found. Got credential name: " + credential_name, + ) + except Exception as e: + verbose_proxy_logger.exception(e) + raise handle_exception_on_proxy(e) + + @router.get( "/credentials/by_model/{model_id}", dependencies=[Depends(user_api_key_auth)], tags=["credential management"], response_model=CredentialItem, ) -async def get_credential( +async def get_credential_by_model( request: Request, fastapi_response: Response, - credential_name: str = Path(..., description="The credential name, percent-decoded; may contain slashes"), - model_id: Optional[str] = None, + model_id: str = Path(..., description="The model ID to look up credentials for"), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -161,48 +197,25 @@ async def get_credential( from litellm.proxy.proxy_server import llm_router try: - if model_id: - if llm_router is None: - raise HTTPException(status_code=500, detail="LLM router not found") - model = llm_router.get_deployment(model_id) - if model is None: - raise HTTPException(status_code=404, detail="Model not found") - credential_values = llm_router.get_deployment_credentials(model_id) - if credential_values is None: - raise HTTPException(status_code=404, detail="Model not found") - masked_credential_values = _get_masked_values( - credential_values, - unmasked_length=4, - number_of_asterisks=4, - ) - credential = CredentialItem( - credential_name="{}-credential-{}".format(model.model_name, model_id), - credential_values=masked_credential_values, - credential_info={}, - ) - # return credential object - return credential - elif credential_name: - for credential in litellm.credential_list: - if credential.credential_name == credential_name: - masked_credential = CredentialItem( - credential_name=credential.credential_name, - credential_values=_get_masked_values( - credential.credential_values, - unmasked_length=4, - number_of_asterisks=4, - ), - credential_info=credential.credential_info, - ) - return masked_credential - raise HTTPException( - status_code=404, - detail="Credential not found. Got credential name: " + credential_name, - ) - else: - raise HTTPException( - status_code=404, detail="Credential name or model ID required" - ) + if llm_router is None: + raise HTTPException(status_code=500, detail="LLM router not found") + model = llm_router.get_deployment(model_id) + if model is None: + raise HTTPException(status_code=404, detail="Model not found") + credential_values = llm_router.get_deployment_credentials(model_id) + if credential_values is None: + raise HTTPException(status_code=404, detail="Model not found") + masked_credential_values = _get_masked_values( + credential_values, + unmasked_length=4, + number_of_asterisks=4, + ) + credential = CredentialItem( + credential_name="{}-credential-{}".format(model.model_name, model_id), + credential_values=masked_credential_values, + credential_info={}, + ) + return credential except Exception as e: verbose_proxy_logger.exception(e) raise handle_exception_on_proxy(e) @@ -216,7 +229,9 @@ async def get_credential( async def delete_credential( request: Request, fastapi_response: Response, - credential_name: str = Path(..., description="The credential name, percent-decoded; may contain slashes"), + credential_name: str = Path( + ..., description="The credential name, percent-decoded; may contain slashes" + ), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -246,7 +261,9 @@ async def delete_credential( def update_db_credential( - db_credential: CredentialItem, updated_patch: CredentialItem, new_encryption_key: Optional[str] = None + db_credential: CredentialItem, + updated_patch: CredentialItem, + new_encryption_key: Optional[str] = None, ) -> CredentialItem: """ Update a credential in the DB. @@ -293,7 +310,9 @@ async def update_credential( request: Request, fastapi_response: Response, credential: CredentialItem, - credential_name: str = Path(..., description="The credential name, percent-decoded; may contain slashes"), + credential_name: str = Path( + ..., description="The credential name, percent-decoded; may contain slashes" + ), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ diff --git a/litellm/proxy/custom_hooks/custom_ui_sso_hook.py b/litellm/proxy/custom_hooks/custom_ui_sso_hook.py index 8bb6b27409..ebf0437671 100644 --- a/litellm/proxy/custom_hooks/custom_ui_sso_hook.py +++ b/litellm/proxy/custom_hooks/custom_ui_sso_hook.py @@ -15,6 +15,7 @@ class CustomSSOLoginHandler(CustomLogger): Useful when you have an OAuth proxy in front of LiteLLM and you want to use the headers from the proxy to sign in the user """ + async def handle_custom_ui_sso_sign_in( self, request: Request, @@ -30,4 +31,6 @@ class CustomSSOLoginHandler(CustomLogger): picture="https://test.com/test.png", provider="test", ) -custom_ui_sso_sign_in_handler = CustomSSOLoginHandler() \ No newline at end of file + + +custom_ui_sso_sign_in_handler = CustomSSOLoginHandler() diff --git a/litellm/proxy/custom_sso.py b/litellm/proxy/custom_sso.py index 651aed4c56..43fb9f97ce 100644 --- a/litellm/proxy/custom_sso.py +++ b/litellm/proxy/custom_sso.py @@ -20,9 +20,10 @@ from litellm.proxy import proxy_server async def custom_sso_handler(userIDPInfo: OpenID) -> SSOUserDefinedValues: try: - if userIDPInfo.id is None: - raise ValueError(f"No ID found for user. userIDPInfo.id is None {userIDPInfo}") + raise ValueError( + f"No ID found for user. userIDPInfo.id is None {userIDPInfo}" + ) # Access extra fields from the IDP response (requires GENERIC_USER_EXTRA_ATTRIBUTES env var) # Example: Set GENERIC_USER_EXTRA_ATTRIBUTES="group,NTID,domain" to capture these fields @@ -31,7 +32,9 @@ async def custom_sso_handler(userIDPInfo: OpenID) -> SSOUserDefinedValues: # check if user exists in litellm proxy DB if proxy_server.prisma_client is not None: - _user_info = await proxy_server.prisma_client.get_data(user_id=userIDPInfo.id) + _user_info = await proxy_server.prisma_client.get_data( + user_id=userIDPInfo.id + ) return SSOUserDefinedValues( models=[], diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 28b1e6601b..a305d5be1e 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -856,9 +856,7 @@ class DBSpendUpdateWriter: or {} ), len( - db_spend_update_transactions.get( - "agent_list_transactions" - ) + db_spend_update_transactions.get("agent_list_transactions") or {} ), ) @@ -1345,7 +1343,9 @@ class DBSpendUpdateWriter: ) ### UPDATE AGENT TABLE ### - agent_list_transactions = db_spend_update_transactions["agent_list_transactions"] + agent_list_transactions = db_spend_update_transactions[ + "agent_list_transactions" + ] await DBSpendUpdateWriter._update_entity_spend_in_db( entity_name="Agent", transactions=agent_list_transactions, @@ -1615,14 +1615,14 @@ class DBSpendUpdateWriter: # Add cache-related fields if they exist if "cache_read_input_tokens" in transaction: - common_data["cache_read_input_tokens"] = ( - transaction.get("cache_read_input_tokens", 0) - ) + common_data[ + "cache_read_input_tokens" + ] = transaction.get("cache_read_input_tokens", 0) if "cache_creation_input_tokens" in transaction: - common_data["cache_creation_input_tokens"] = ( - transaction.get( - "cache_creation_input_tokens", 0 - ) + common_data[ + "cache_creation_input_tokens" + ] = transaction.get( + "cache_creation_input_tokens", 0 ) if entity_type == "tag" and "request_id" in transaction: diff --git a/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py index f47b694d44..75e9b9580b 100644 --- a/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py @@ -54,9 +54,9 @@ class DailySpendUpdateQueue(BaseUpdateQueue): def __init__(self): super().__init__() - self.update_queue: asyncio.Queue[Dict[str, BaseDailySpendTransaction]] = ( - asyncio.Queue(maxsize=LITELLM_ASYNCIO_QUEUE_MAXSIZE) - ) + self.update_queue: asyncio.Queue[ + Dict[str, BaseDailySpendTransaction] + ] = asyncio.Queue(maxsize=LITELLM_ASYNCIO_QUEUE_MAXSIZE) async def add_update(self, update: Dict[str, BaseDailySpendTransaction]): """Enqueue an update.""" @@ -73,9 +73,9 @@ class DailySpendUpdateQueue(BaseUpdateQueue): Combine all updates in the queue into a single update. This is used to reduce the size of the in-memory queue. """ - updates: List[Dict[str, BaseDailySpendTransaction]] = ( - await self.flush_all_updates_from_in_memory_queue() - ) + updates: List[ + Dict[str, BaseDailySpendTransaction] + ] = await self.flush_all_updates_from_in_memory_queue() aggregated_updates = self.get_aggregated_daily_spend_update_transactions( updates ) diff --git a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py index 6f86e82cf2..546ea05998 100644 --- a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py +++ b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py @@ -36,7 +36,7 @@ class PodLockManager: """ Attempt to acquire the lock for a specific cron job using Redis. Uses the SET command with NX and EX options to ensure atomicity. - + Args: cronjob_id: The ID of the cron job to lock """ diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 4f38e71bbf..c51c06df2f 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -10,31 +10,36 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast from litellm._logging import verbose_proxy_logger from litellm.caching import RedisCache -from litellm.constants import (MAX_REDIS_BUFFER_DEQUEUE_COUNT, - REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, - REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, - REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, - REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, - REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, - REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, - REDIS_UPDATE_BUFFER_KEY) +from litellm.constants import ( + MAX_REDIS_BUFFER_DEQUEUE_COUNT, + REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, + REDIS_UPDATE_BUFFER_KEY, +) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps -from litellm.proxy._types import (DailyAgentSpendTransaction, - DailyEndUserSpendTransaction, - DailyOrganizationSpendTransaction, - DailyTagSpendTransaction, - DailyTeamSpendTransaction, - DailyUserSpendTransaction, - DBSpendUpdateTransactions) -from litellm.proxy.db.db_transaction_queue.base_update_queue import \ - service_logger_obj -from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import \ - DailySpendUpdateQueue -from litellm.proxy.db.db_transaction_queue.spend_update_queue import \ - SpendUpdateQueue +from litellm.proxy._types import ( + DailyAgentSpendTransaction, + DailyEndUserSpendTransaction, + DailyOrganizationSpendTransaction, + DailyTagSpendTransaction, + DailyTeamSpendTransaction, + DailyUserSpendTransaction, + DBSpendUpdateTransactions, +) +from litellm.proxy.db.db_transaction_queue.base_update_queue import service_logger_obj +from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( + DailySpendUpdateQueue, +) +from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue from litellm.secret_managers.main import str_to_bool -from litellm.types.caching import (RedisPipelineLpopOperation, - RedisPipelineRpushOperation) +from litellm.types.caching import ( + RedisPipelineLpopOperation, + RedisPipelineRpushOperation, +) from litellm.types.services import ServiceTypes if TYPE_CHECKING: @@ -66,9 +71,9 @@ class RedisUpdateBuffer: """ from litellm.proxy.proxy_server import general_settings - _use_redis_transaction_buffer: Optional[Union[bool, str]] = ( - general_settings.get("use_redis_transaction_buffer", False) - ) + _use_redis_transaction_buffer: Optional[ + Union[bool, str] + ] = general_settings.get("use_redis_transaction_buffer", False) if isinstance(_use_redis_transaction_buffer, str): _use_redis_transaction_buffer = str_to_bool(_use_redis_transaction_buffer) if _use_redis_transaction_buffer is None: @@ -210,13 +215,41 @@ class RedisUpdateBuffer: # Build a list of rpush operations, skipping empty/None transaction sets _queue_configs: List[Tuple[Any, str, ServiceTypes]] = [ - (db_spend_update_transactions, REDIS_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_SPEND_UPDATE_QUEUE), - (daily_spend_update_transactions, REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE), - (daily_team_spend_update_transactions, REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_TEAM_SPEND_UPDATE_QUEUE), - (daily_org_spend_update_transactions, REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_ORG_SPEND_UPDATE_QUEUE), - (daily_end_user_spend_update_transactions, REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_END_USER_SPEND_UPDATE_QUEUE), - (daily_agent_spend_update_transactions, REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE), - (daily_tag_spend_update_transactions, REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_TAG_SPEND_UPDATE_QUEUE), + ( + db_spend_update_transactions, + REDIS_UPDATE_BUFFER_KEY, + ServiceTypes.REDIS_SPEND_UPDATE_QUEUE, + ), + ( + daily_spend_update_transactions, + REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, + ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE, + ), + ( + daily_team_spend_update_transactions, + REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, + ServiceTypes.REDIS_DAILY_TEAM_SPEND_UPDATE_QUEUE, + ), + ( + daily_org_spend_update_transactions, + REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, + ServiceTypes.REDIS_DAILY_ORG_SPEND_UPDATE_QUEUE, + ), + ( + daily_end_user_spend_update_transactions, + REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, + ServiceTypes.REDIS_DAILY_END_USER_SPEND_UPDATE_QUEUE, + ), + ( + daily_agent_spend_update_transactions, + REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, + ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE, + ), + ( + daily_tag_spend_update_transactions, + REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, + ServiceTypes.REDIS_DAILY_TAG_SPEND_UPDATE_QUEUE, + ), ] rpush_list: List[RedisPipelineRpushOperation] = [] @@ -361,13 +394,33 @@ class RedisUpdateBuffer: return None, None, None, None, None, None, None lpop_list: List[RedisPipelineLpopOperation] = [ - RedisPipelineLpopOperation(key=REDIS_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), - RedisPipelineLpopOperation(key=REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), - RedisPipelineLpopOperation(key=REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), - RedisPipelineLpopOperation(key=REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), - RedisPipelineLpopOperation(key=REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), - RedisPipelineLpopOperation(key=REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), - RedisPipelineLpopOperation(key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), + RedisPipelineLpopOperation( + key=REDIS_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT + ), + RedisPipelineLpopOperation( + key=REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, + count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, + ), + RedisPipelineLpopOperation( + key=REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, + count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, + ), + RedisPipelineLpopOperation( + key=REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, + count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, + ), + RedisPipelineLpopOperation( + key=REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, + count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, + ), + RedisPipelineLpopOperation( + key=REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, + count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, + ), + RedisPipelineLpopOperation( + key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, + count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, + ), ] raw_results = await self.redis_cache.async_lpop_pipeline(lpop_list=lpop_list) @@ -399,7 +452,9 @@ class RedisUpdateBuffer: db_spend, cast(Optional[Dict[str, DailyUserSpendTransaction]], daily_results[0]), cast(Optional[Dict[str, DailyTeamSpendTransaction]], daily_results[1]), - cast(Optional[Dict[str, DailyOrganizationSpendTransaction]], daily_results[2]), + cast( + Optional[Dict[str, DailyOrganizationSpendTransaction]], daily_results[2] + ), cast(Optional[Dict[str, DailyEndUserSpendTransaction]], daily_results[3]), cast(Optional[Dict[str, DailyAgentSpendTransaction]], daily_results[4]), cast(Optional[Dict[str, DailyTagSpendTransaction]], daily_results[5]), @@ -455,7 +510,7 @@ class RedisUpdateBuffer: async def get_all_daily_org_spend_update_transactions_from_redis_buffer( self, - ) -> Optional[Dict[str, DailyOrganizationSpendTransaction]]: + ) -> Optional[Dict[str, DailyOrganizationSpendTransaction]]: """ Gets all the daily organization spend update transactions from Redis """ @@ -471,7 +526,7 @@ class RedisUpdateBuffer: json.loads(transaction) for transaction in list_of_transactions ] return cast( - Dict[str, DailyOrganizationSpendTransaction], + Dict[str, DailyOrganizationSpendTransaction], DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( list_of_daily_spend_update_transactions ), diff --git a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py index 8c04bae259..ba9423c6ef 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -131,9 +131,12 @@ class SpendLogCleanup: # If we have a pod lock manager, try to acquire the lock if self.pod_lock_manager and self.pod_lock_manager.redis_cache: - lock_acquired = await self.pod_lock_manager.acquire_lock( - cronjob_id=SPEND_LOG_CLEANUP_JOB_NAME, - ) or False + lock_acquired = ( + await self.pod_lock_manager.acquire_lock( + cronjob_id=SPEND_LOG_CLEANUP_JOB_NAME, + ) + or False + ) verbose_proxy_logger.info( f"Lock acquisition attempt: {'successful' if lock_acquired else 'failed'} at {datetime.now()}" ) @@ -158,7 +161,11 @@ class SpendLogCleanup: return # Return after error handling finally: # Only release the lock if it was actually acquired - if lock_acquired and self.pod_lock_manager and self.pod_lock_manager.redis_cache: + if ( + lock_acquired + and self.pod_lock_manager + and self.pod_lock_manager.redis_cache + ): await self.pod_lock_manager.release_lock( cronjob_id=SPEND_LOG_CLEANUP_JOB_NAME ) diff --git a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py index b7cd06a64f..727e8dc1d5 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py @@ -3,10 +3,15 @@ from typing import Dict, List, Optional from litellm._logging import verbose_proxy_logger from litellm.constants import LITELLM_ASYNCIO_QUEUE_MAXSIZE -from litellm.proxy._types import (DBSpendUpdateTransactions, - Litellm_EntityType, SpendUpdateQueueItem) +from litellm.proxy._types import ( + DBSpendUpdateTransactions, + Litellm_EntityType, + SpendUpdateQueueItem, +) from litellm.proxy.db.db_transaction_queue.base_update_queue import ( - BaseUpdateQueue, service_logger_obj) + BaseUpdateQueue, + service_logger_obj, +) from litellm.types.services import ServiceTypes diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index b2efbf9d07..213dd39adc 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -60,7 +60,11 @@ class PrismaDBExceptionHandler: if isinstance(e, DB_CONNECTION_ERROR_TYPES): return True if isinstance( - e, (prisma.errors.ClientNotConnectedError, prisma.errors.HTTPClientClosedError) + e, + ( + prisma.errors.ClientNotConnectedError, + prisma.errors.HTTPClientClosedError, + ), ): return True if isinstance(e, prisma.errors.PrismaError): diff --git a/litellm/proxy/db/spend_log_tool_index.py b/litellm/proxy/db/spend_log_tool_index.py index 6e8c63675e..835d76e0ee 100644 --- a/litellm/proxy/db/spend_log_tool_index.py +++ b/litellm/proxy/db/spend_log_tool_index.py @@ -50,7 +50,9 @@ def _parse_tool_names_from_payload(payload: Dict[str, Any]) -> Set[str]: _add_tool_calls_to_set(response_obj.get("tool_calls"), tool_names) choices = response_obj.get("choices") if isinstance(choices, list) and choices: - msg = choices[0].get("message") if isinstance(choices[0], dict) else None + msg = ( + choices[0].get("message") if isinstance(choices[0], dict) else None + ) if isinstance(msg, dict): _add_tool_calls_to_set(msg.get("tool_calls"), tool_names) @@ -101,9 +103,7 @@ async def process_spend_logs_tool_usage( continue if isinstance(start_time, str): try: - start_time = datetime.fromisoformat( - start_time.replace("Z", "+00:00") - ) + start_time = datetime.fromisoformat(start_time.replace("Z", "+00:00")) except (ValueError, TypeError): continue if start_time.tzinfo is None: @@ -111,11 +111,13 @@ async def process_spend_logs_tool_usage( tool_names = _parse_tool_names_from_payload(payload) for tool_name in tool_names: - index_rows.append({ - "request_id": request_id, - "tool_name": tool_name, - "start_time": start_time, - }) + index_rows.append( + { + "request_id": request_id, + "tool_name": tool_name, + "start_time": start_time, + } + ) if not index_rows: return @@ -131,11 +133,13 @@ async def process_spend_logs_tool_usage( continue if st.tzinfo is None: st = st.replace(tzinfo=timezone.utc) - index_data.append({ - "request_id": r["request_id"], - "tool_name": r["tool_name"], - "start_time": st, - }) + index_data.append( + { + "request_id": r["request_id"], + "tool_name": r["tool_name"], + "start_time": st, + } + ) if index_data: await prisma_client.db.litellm_spendlogtoolindex.create_many( data=index_data, diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py index 0eda012d51..6b34c974cf 100644 --- a/litellm/proxy/db/tool_registry_writer.py +++ b/litellm/proxy/db/tool_registry_writer.py @@ -358,9 +358,7 @@ class ToolPolicyRegistry: blocked: set = set() for op_id in (object_permission_id, team_object_permission_id): if op_id and op_id.strip(): - blocked.update( - self._blocked_tools_by_op_id.get(op_id.strip(), []) - ) + blocked.update(self._blocked_tools_by_op_id.get(op_id.strip(), [])) result: Dict[str, str] = {} for name in tool_names: if name in blocked: diff --git a/litellm/proxy/dd_span_tagger.py b/litellm/proxy/dd_span_tagger.py index 08b7d928d0..7bbfe50a01 100644 --- a/litellm/proxy/dd_span_tagger.py +++ b/litellm/proxy/dd_span_tagger.py @@ -48,7 +48,9 @@ class DDSpanTagger: """ try: if user_api_key_dict.key_alias: - set_active_span_tag("litellm.key_alias", str(user_api_key_dict.key_alias)) + set_active_span_tag( + "litellm.key_alias", str(user_api_key_dict.key_alias) + ) if user_api_key_dict.token: set_active_span_tag("litellm.key_hash", str(user_api_key_dict.token)) if requested_model: diff --git a/litellm/proxy/fine_tuning_endpoints/endpoints.py b/litellm/proxy/fine_tuning_endpoints/endpoints.py index 46f719a7c9..ff6300a4fa 100644 --- a/litellm/proxy/fine_tuning_endpoints/endpoints.py +++ b/litellm/proxy/fine_tuning_endpoints/endpoints.py @@ -281,7 +281,9 @@ async def retrieve_fine_tuning_job( except Exception: request_body = {} - custom_llm_provider = request_body.get("custom_llm_provider", None) or custom_llm_provider + custom_llm_provider = ( + request_body.get("custom_llm_provider", None) or custom_llm_provider + ) ## CHECK IF MANAGED FILE ID unified_finetuning_job_id: Union[str, Literal[False]] = False diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 4c866a2499..2b20876ba2 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -566,9 +566,7 @@ class GuardrailSubmissionItem(BaseModel): guardrail_name: str status: str # pending_review | active | rejected team_id: Optional[str] = None - team_guardrail: bool = ( - False # True when submitted via team (team_id set); use to distinguish team vs regular guardrails - ) + team_guardrail: bool = False # True when submitted via team (team_id set); use to distinguish team vs regular guardrails litellm_params: Optional[Dict[str, Any]] = None guardrail_info: Optional[Dict[str, Any]] = None submitted_by_user_id: Optional[str] = None @@ -663,9 +661,9 @@ async def register_guardrail( guardrail_info = dict(request.guardrail_info or {}) guardrail_info["submitted_by_user_id"] = user_api_key_dict.user_id guardrail_info["submitted_by_email"] = user_api_key_dict.user_email - guardrail_info["team_guardrail"] = ( - True # Mark as team submission for filtering/display - ) + guardrail_info[ + "team_guardrail" + ] = True # Mark as team submission for filtering/display guardrail_info_str = safe_dumps(guardrail_info) try: @@ -769,9 +767,7 @@ async def list_guardrail_submissions( active_count = sum( 1 for r in all_team_rows if (r.status or "active") == "active" ) - rejected = sum( - 1 for r in all_team_rows if (r.status or "active") == "rejected" - ) + rejected = sum(1 for r in all_team_rows if (r.status or "active") == "rejected") # Apply filters to get the submissions list rows = all_team_rows @@ -1810,9 +1806,9 @@ async def get_provider_specific_params(): lakera_v2_fields = _get_fields_from_model(LakeraV2GuardrailConfigModel) tool_permission_fields = _get_fields_from_model(ToolPermissionGuardrailConfigModel) - tool_permission_fields["ui_friendly_name"] = ( - ToolPermissionGuardrailConfigModel.ui_friendly_name() - ) + tool_permission_fields[ + "ui_friendly_name" + ] = ToolPermissionGuardrailConfigModel.ui_friendly_name() # Return the provider-specific parameters provider_params = { @@ -2085,10 +2081,10 @@ async def apply_guardrail( from litellm.proxy.utils import handle_exception_on_proxy try: - active_guardrail: Optional[CustomGuardrail] = ( - GUARDRAIL_REGISTRY.get_initialized_guardrail_callback( - guardrail_name=request.guardrail_name - ) + active_guardrail: Optional[ + CustomGuardrail + ] = GUARDRAIL_REGISTRY.get_initialized_guardrail_callback( + guardrail_name=request.guardrail_name ) if active_guardrail is None: raise HTTPException( diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py b/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py index 744329f85f..a465c5428d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py @@ -116,9 +116,7 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr AzureTextModerationGuardrailResponse, ) - chunks = self.split_text_by_words( - text, AZURE_CONTENT_SAFETY_MAX_TEXT_LENGTH - ) + chunks = self.split_text_by_words(text, AZURE_CONTENT_SAFETY_MAX_TEXT_LENGTH) last_response: Optional[AzureTextModerationGuardrailResponse] = None diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 6800dff55a..8ef188bb23 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -470,9 +470,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): response = getattr(e, "response", None) if isinstance(response, httpx.Response): try: - status_code, detail_message = ( - self._parse_bedrock_guardrail_error_response(response) - ) + ( + status_code, + detail_message, + ) = self._parse_bedrock_guardrail_error_response(response) self.add_standard_logging_guardrail_information_to_request_data( guardrail_provider=self.guardrail_provider, guardrail_json_response={"error": detail_message}, @@ -795,9 +796,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ######################################################### ########## 1. Make the Bedrock API request ########## ######################################################### - bedrock_guardrail_response: Optional[Union[BedrockGuardrailResponse, str]] = ( - None - ) + bedrock_guardrail_response: Optional[ + Union[BedrockGuardrailResponse, str] + ] = None try: bedrock_guardrail_response = await self.make_bedrock_api_request( source="INPUT", messages=filtered_messages, request_data=data @@ -867,9 +868,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ######################################################### ########## 1. Make the Bedrock API request ########## ######################################################### - bedrock_guardrail_response: Optional[Union[BedrockGuardrailResponse, str]] = ( - None - ) + bedrock_guardrail_response: Optional[ + Union[BedrockGuardrailResponse, str] + ] = None try: bedrock_guardrail_response = await self.make_bedrock_api_request( source="INPUT", messages=filtered_messages, request_data=data diff --git a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py index 51bc6d08ac..b8a2111c01 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py @@ -41,9 +41,7 @@ def initialize_guardrail( guardrail_name = guardrail.get("guardrail_name") if not guardrail_name: - raise ValueError( - "Block Code Execution guardrail requires a guardrail_name" - ) + raise ValueError("Block Code Execution guardrail requires a guardrail_name") blocked_languages: Optional[List[str]] = cast( Optional[List[str]], diff --git a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py index e76a02a6e4..efd781681a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py +++ b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py @@ -347,9 +347,9 @@ class BlockCodeExecutionGuardrail(CustomGuardrail): **kwargs: Any, ) -> None: # Normalize to type expected by CustomGuardrail - _event_hook: Optional[Union[GuardrailEventHooks, List[GuardrailEventHooks]]] = ( - None - ) + _event_hook: Optional[ + Union[GuardrailEventHooks, List[GuardrailEventHooks]] + ] = None if event_hook is not None: if isinstance(event_hook, list): _event_hook = [ @@ -483,9 +483,7 @@ class BlockCodeExecutionGuardrail(CustomGuardrail): # For responses, always enforce the block action (no intent check needed). # For requests with detect_execution_intent, require execution intent. effective_block = action_taken == "block" and ( - is_response - or not self.detect_execution_intent - or has_execution_intent + is_response or not self.detect_execution_intent or has_execution_intent ) if detections is not None: detections.append( diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py index d166e66dba..a956688fd4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py @@ -12,8 +12,10 @@ from typing import TYPE_CHECKING from litellm.types.guardrails import SupportedGuardrailIntegrations from .custom_code_guardrail import CustomCodeGuardrail -from .response_rejection_code import (DEFAULT_REJECTION_PHRASES, - RESPONSE_REJECTION_GUARDRAIL_CODE) +from .response_rejection_code import ( + DEFAULT_REJECTION_PHRASES, + RESPONSE_REJECTION_GUARDRAIL_CODE, +) if TYPE_CHECKING: from litellm.types.guardrails import Guardrail, LitellmParams diff --git a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/__init__.py index c1ffa33764..79f1992da4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/__init__.py @@ -1,4 +1,3 @@ from .dynamoai import DynamoAIGuardrails __all__ = ["DynamoAIGuardrails"] - diff --git a/litellm/proxy/guardrails/guardrail_hooks/enkryptai/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/enkryptai/__init__.py index ab62679a60..74d07d3f71 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/enkryptai/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/enkryptai/__init__.py @@ -38,4 +38,3 @@ guardrail_initializer_registry = { guardrail_class_registry = { SupportedGuardrailIntegrations.ENKRYPTAI.value: EnkryptAIGuardrails, } - diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index 2d0ce040a6..1872084508 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -219,9 +219,9 @@ class GenericGuardrailAPI(CustomGuardrail): additional_provider_specific_params or {} ) - self.unreachable_fallback: Literal["fail_closed", "fail_open"] = ( - unreachable_fallback - ) + self.unreachable_fallback: Literal[ + "fail_closed", "fail_open" + ] = unreachable_fallback # Set supported event hooks if "supported_event_hooks" not in kwargs: @@ -295,7 +295,9 @@ class GenericGuardrailAPI(CustomGuardrail): error: Exception, http_status_code: Optional[int] = None, ) -> GenericGuardrailAPIInputs: - status_suffix = f" http_status_code={http_status_code}" if http_status_code else "" + status_suffix = ( + f" http_status_code={http_status_code}" if http_status_code else "" + ) verbose_proxy_logger.critical( "Generic Guardrail API unreachable (fail-open). Proceeding without guardrail.%s " "guardrail_name=%s api_base=%s input_type=%s litellm_call_id=%s litellm_trace_id=%s", @@ -496,9 +498,7 @@ class GenericGuardrailAPI(CustomGuardrail): e, inputs, input_type, logging_obj ) except httpx.HTTPStatusError as e: - status_code = getattr( - getattr(e, "response", None), "status_code", None - ) + status_code = getattr(getattr(e, "response", None), "status_code", None) is_unreachable = status_code in (502, 503, 504) return self._handle_guardrail_request_error( e, inputs, input_type, logging_obj, is_unreachable=is_unreachable diff --git a/litellm/proxy/guardrails/guardrail_hooks/grayswan/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/grayswan/__init__.py index 3aca9078c0..c6dee3f841 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/grayswan/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/grayswan/__init__.py @@ -42,10 +42,12 @@ def initialize_guardrail( policy_id=_get_config_value(litellm_params, optional_params, "policy_id"), streaming_end_of_stream_only=_get_config_value( litellm_params, optional_params, "streaming_end_of_stream_only" - ) or False, + ) + or False, streaming_sampling_rate=_get_config_value( litellm_params, optional_params, "streaming_sampling_rate" - ) or 5, + ) + or 5, fail_open=_get_config_value(litellm_params, optional_params, "fail_open"), guardrail_timeout=_get_config_value( litellm_params, optional_params, "guardrail_timeout" diff --git a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/__init__.py index e397d8098a..2f22e4c33d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/__init__.py @@ -26,14 +26,14 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" verify_ssl = getattr(litellm_params, "verify_ssl", True) # Get optional params - optional_params = getattr(litellm_params, "optional_params", IBMDetectorOptionalParams()) + optional_params = getattr( + litellm_params, "optional_params", IBMDetectorOptionalParams() + ) detector_params = getattr(optional_params, "detector_params", {}) extra_headers = getattr(optional_params, "extra_headers", {}) score_threshold = getattr(optional_params, "score_threshold", None) block_on_detection = getattr(optional_params, "block_on_detection", True) - - is_detector_server = litellm_params.is_detector_server if is_detector_server is None: is_detector_server = True diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index dbda524ca0..6b917bc794 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -358,9 +358,7 @@ class LakeraAIGuardrail(CustomGuardrail): # when some choices have null content (e.g. tool-call-only). response_messages: List[AllMessageValues] = [] choice_indices: List[int] = [] - response_dict = ( - response.model_dump() if hasattr(response, "model_dump") else {} - ) + response_dict = response.model_dump() if hasattr(response, "model_dump") else {} for i, choice in enumerate(response_dict.get("choices", [])): msg = choice.get("message") if not msg: @@ -395,7 +393,9 @@ class LakeraAIGuardrail(CustomGuardrail): for idx, msg in enumerate(assistant_messages): if idx < len(choice_indices): choice_idx = choice_indices[idx] - response_dict["choices"][choice_idx]["message"]["content"] = msg.get("content", "") + response_dict["choices"][choice_idx]["message"][ + "content" + ] = msg.get("content", "") add_guardrail_to_applied_guardrails_header( request_data=data, guardrail_name=self.guardrail_name ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py index 5850103132..a72f3e4c3f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py @@ -7,7 +7,17 @@ import os import uuid -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Type, Union, TypedDict +from typing import ( + TYPE_CHECKING, + Any, + Dict, + List, + Literal, + Optional, + Type, + Union, + TypedDict, +) try: import ulid @@ -92,7 +102,9 @@ class LassoGuardrail(CustomGuardrail): ) self.lasso_api_key = lasso_api_key or api_key or os.environ.get("LASSO_API_KEY") self.user_id = user_id or os.environ.get("LASSO_USER_ID") - self.conversation_id = conversation_id or os.environ.get("LASSO_CONVERSATION_ID") + self.conversation_id = conversation_id or os.environ.get( + "LASSO_CONVERSATION_ID" + ) self.mask = mask or False if self.lasso_api_key is None: @@ -102,7 +114,9 @@ class LassoGuardrail(CustomGuardrail): ) self.api_base = ( - api_base or os.getenv("LASSO_API_BASE") or "https://server.lasso.security/gateway/v3" + api_base + or os.getenv("LASSO_API_BASE") + or "https://server.lasso.security/gateway/v3" ) verbose_proxy_logger.debug( @@ -127,7 +141,7 @@ class LassoGuardrail(CustomGuardrail): async def async_pre_call_hook( self, user_api_key_dict: UserAPIKeyAuth, - cache: DualCache, # Deprecated, use global_cache instead (kept to align with CustomGuardrail interface) + cache: DualCache, # Deprecated, use global_cache instead (kept to align with CustomGuardrail interface) data: dict, call_type: Literal[ "completion", @@ -155,7 +169,9 @@ class LassoGuardrail(CustomGuardrail): # The conversation_id is being stored in the cache so it can be used by the post_call hook self._get_or_generate_conversation_id(data, global_cache) - return await self._run_lasso_guardrail(data, global_cache, message_type="PROMPT") + return await self._run_lasso_guardrail( + data, global_cache, message_type="PROMPT" + ) @log_guardrail_information async def async_moderation_hook( @@ -206,7 +222,9 @@ class LassoGuardrail(CustomGuardrail): response_messages = [] for choice in response.choices: if hasattr(choice, "message") and choice.message.content: - response_messages.append({"role": "assistant", "content": choice.message.content}) + response_messages.append( + {"role": "assistant", "content": choice.message.content} + ) if response_messages: # Include litellm_call_id from original data for conversation_id consistency @@ -215,35 +233,53 @@ class LassoGuardrail(CustomGuardrail): "litellm_call_id": data.get("litellm_call_id"), } - # Handle masking for post-call if self.mask: headers = self._prepare_headers(response_data, global_cache) - payload = self._prepare_payload(response_messages, response_data, global_cache, "COMPLETION") + payload = self._prepare_payload( + response_messages, response_data, global_cache, "COMPLETION" + ) api_url = f"{self.api_base}/classifix" try: - lasso_response = await self._call_lasso_api(headers=headers, payload=payload, api_url=api_url) + lasso_response = await self._call_lasso_api( + headers=headers, payload=payload, api_url=api_url + ) self._process_lasso_response(lasso_response) # Apply masking to the actual response if masked content is available masked_messages = lasso_response.get("messages") - if lasso_response.get("violations_detected") and masked_messages: - self._apply_masking_to_model_response(response, masked_messages) - verbose_proxy_logger.debug("Applied Lasso masking to model response") + if ( + lasso_response.get("violations_detected") + and masked_messages + ): + self._apply_masking_to_model_response( + response, masked_messages + ) + verbose_proxy_logger.debug( + "Applied Lasso masking to model response" + ) except Exception as e: if isinstance(e, HTTPException): raise e - verbose_proxy_logger.error(f"Error in post-call Lasso masking: {str(e)}") - raise LassoGuardrailAPIError(f"Failed to apply post-call masking: {str(e)}") + verbose_proxy_logger.error( + f"Error in post-call Lasso masking: {str(e)}" + ) + raise LassoGuardrailAPIError( + f"Failed to apply post-call masking: {str(e)}" + ) else: # Use the same data for conversation_id consistency (no cache access needed) - await self._run_lasso_guardrail(response_data, cache=global_cache, message_type="COMPLETION") + await self._run_lasso_guardrail( + response_data, cache=global_cache, message_type="COMPLETION" + ) verbose_proxy_logger.debug("Post-call Lasso validation completed") else: verbose_proxy_logger.warning("No response messages found to validate") else: - verbose_proxy_logger.warning(f"Unexpected response type for post-call hook: {type(response)}") + verbose_proxy_logger.warning( + f"Unexpected response type for post-call hook: {type(response)}" + ) return response @@ -337,7 +373,9 @@ class LassoGuardrail(CustomGuardrail): if self.mask: return await self._handle_masking(data, cache, message_type, messages) else: - return await self._handle_classification(data, cache, message_type, messages) + return await self._handle_classification( + data, cache, message_type, messages + ) async def _handle_classification( self, @@ -369,7 +407,9 @@ class LassoGuardrail(CustomGuardrail): headers = self._prepare_headers(data, cache) payload = self._prepare_payload(messages, data, cache, message_type) api_url = f"{self.api_base}/classifix" - response = await self._call_lasso_api(headers=headers, payload=payload, api_url=api_url) + response = await self._call_lasso_api( + headers=headers, payload=payload, api_url=api_url + ) self._process_lasso_response(response) # Apply masking to messages if violations detected and masked messages are available @@ -411,10 +451,14 @@ class LassoGuardrail(CustomGuardrail): elif error.response.status_code == 429: raise LassoGuardrailAPIError("Lasso API rate limit exceeded") else: - raise LassoGuardrailAPIError(f"API error: {error.response.status_code}") + raise LassoGuardrailAPIError( + f"API error: {error.response.status_code}" + ) # Generic error handling - raise LassoGuardrailAPIError(f"Failed to verify request safety with Lasso API: {str(error)}") + raise LassoGuardrailAPIError( + f"Failed to verify request safety with Lasso API: {str(error)}" + ) def _log_masking_applied( self, @@ -452,7 +496,7 @@ class LassoGuardrail(CustomGuardrail): headers["lasso-user-id"] = self.user_id # Always include conversation_id (generated or provided) - conversation_id = self._get_or_generate_conversation_id(data, cache) + conversation_id = self._get_or_generate_conversation_id(data, cache) headers["lasso-conversation-id"] = conversation_id @@ -495,7 +539,9 @@ class LassoGuardrail(CustomGuardrail): ) -> LassoResponse: """Call the Lasso API and return the response.""" url = api_url or f"{self.api_base}/classify" - verbose_proxy_logger.debug(f"Calling Lasso API with messageType: {payload.get('messageType')}") + verbose_proxy_logger.debug( + f"Calling Lasso API with messageType: {payload.get('messageType')}" + ) response = await self.async_handler.post( url=url, headers=headers, @@ -533,7 +579,9 @@ class LassoGuardrail(CustomGuardrail): """ if response and response.get("violations_detected") is True: violated_deputies = self._parse_violated_deputies(response) - verbose_proxy_logger.warning(f"Lasso guardrail detected violations: {violated_deputies}") + verbose_proxy_logger.warning( + f"Lasso guardrail detected violations: {violated_deputies}" + ) # Check if any findings have "BLOCK" action blocking_violations = self._check_for_blocking_actions(response) @@ -609,11 +657,17 @@ class LassoGuardrail(CustomGuardrail): """Apply masking to the actual model response when mask=True and masked content is available.""" masked_index = 0 for choice in model_response.choices: - if hasattr(choice, "message") and choice.message.content and masked_index < len(masked_messages): + if ( + hasattr(choice, "message") + and choice.message.content + and masked_index < len(masked_messages) + ): # Replace the content with the masked version from Lasso choice.message.content = masked_messages[masked_index]["content"] masked_index += 1 - verbose_proxy_logger.debug(f"Applied masked content to choice {masked_index}") + verbose_proxy_logger.debug( + f"Applied masked content to choice {masked_index}" + ) @staticmethod def get_config_model() -> Optional[Type["GuardrailConfigModel"]]: diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/__init__.py index 111f8dc783..8eb4960264 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/__init__.py @@ -47,9 +47,13 @@ def initialize_guardrail( competitor_intent_config=getattr( litellm_params, "competitor_intent_config", None ), - end_session_after_n_fails=getattr(litellm_params, "end_session_after_n_fails", None), + end_session_after_n_fails=getattr( + litellm_params, "end_session_after_n_fails", None + ), on_violation=getattr(litellm_params, "on_violation", None), - realtime_violation_message=getattr(litellm_params, "realtime_violation_message", None), + realtime_violation_message=getattr( + litellm_params, "realtime_violation_message", None + ), ) litellm.logging_callback_manager.add_litellm_callback(content_filter_guardrail) diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/competitor_intent/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/competitor_intent/__init__.py index 85b92cb16b..d373fc2481 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/competitor_intent/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/competitor_intent/__init__.py @@ -4,10 +4,14 @@ Competitor intent: entity + intent disambiguation with safe (non-competitor) def Base logic in base.py; industry-specific checkers in submodules (e.g. airline.py). """ -from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.competitor_intent.airline import \ - AirlineCompetitorIntentChecker +from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.competitor_intent.airline import ( + AirlineCompetitorIntentChecker, +) from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.competitor_intent.base import ( - BaseCompetitorIntentChecker, normalize, text_for_entity_matching) + BaseCompetitorIntentChecker, + normalize, + text_for_entity_matching, +) __all__ = [ "BaseCompetitorIntentChecker", diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/competitor_intent/airline.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/competitor_intent/airline.py index 5da0bd25fc..90b45262c4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/competitor_intent/airline.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/competitor_intent/airline.py @@ -106,13 +106,9 @@ AIRLINE_COMPARISON_SIGNALS = [ # Explicit markers: strong override when present. AIRLINE_EXPLICIT_COMPETITOR_MARKER = r"\b(airways?|airline|carrier)\b" -AIRLINE_EXPLICIT_OTHER_MEANING_MARKER = ( - r"\b(fly|travel|going|visit|layover|stopover|transit)\b.{0,12}\b(to|in|via|from)\b.{0,8}\b" -) +AIRLINE_EXPLICIT_OTHER_MEANING_MARKER = r"\b(fly|travel|going|visit|layover|stopover|transit)\b.{0,12}\b(to|in|via|from)\b.{0,8}\b" -_MAJOR_AIRLINES_PATH = ( - Path(__file__).resolve().parent / "major_airlines.json" -) +_MAJOR_AIRLINES_PATH = Path(__file__).resolve().parent / "major_airlines.json" def _load_competitors_excluding_brand(brand_self: List[str]) -> List[str]: @@ -163,7 +159,9 @@ class AirlineCompetitorIntentChecker(BaseCompetitorIntentChecker): if not merged.get("explicit_competitor_marker"): merged["explicit_competitor_marker"] = AIRLINE_EXPLICIT_COMPETITOR_MARKER if not merged.get("explicit_other_meaning_marker"): - merged["explicit_other_meaning_marker"] = AIRLINE_EXPLICIT_OTHER_MEANING_MARKER + merged[ + "explicit_other_meaning_marker" + ] = AIRLINE_EXPLICIT_OTHER_MEANING_MARKER if not merged.get("domain_words"): merged["domain_words"] = ["airline", "airlines", "carrier"] if not merged.get("competitors"): @@ -184,12 +182,15 @@ class AirlineCompetitorIntentChecker(BaseCompetitorIntentChecker): def _classify_ambiguous(self, text: str, token: str) -> Tuple[str, float]: """Other meaning vs competitor using airline signals and explicit markers.""" text_lower = text.lower() - if self._explicit_competitor_marker and self._explicit_competitor_marker.search( - text_lower - ) and _word_boundary_match(text_lower, token.lower()): + if ( + self._explicit_competitor_marker + and self._explicit_competitor_marker.search(text_lower) + and _word_boundary_match(text_lower, token.lower()) + ): return "COMPETITOR", 0.85 - if self._explicit_other_meaning_marker and self._explicit_other_meaning_marker.search( - text_lower + if ( + self._explicit_other_meaning_marker + and self._explicit_other_meaning_marker.search(text_lower) ): return "OTHER_MEANING", 0.85 # Operational-only: baggage/lounge/check-in/refund with no comparison → product query diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index 6c83833c66..e4da1c1ae7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -212,17 +212,17 @@ class ContentFilterGuardrail(CustomGuardrail): self.image_model = image_model # Store loaded categories self.loaded_categories: Dict[str, CategoryConfig] = {} - self.category_keywords: Dict[str, Tuple[str, str, ContentFilterAction]] = ( - {} - ) # keyword -> (category, severity, action) + self.category_keywords: Dict[ + str, Tuple[str, str, ContentFilterAction] + ] = {} # keyword -> (category, severity, action) # Always-block keywords are checked after exceptions (exceptions take precedence) self.always_block_category_keywords: Dict[ str, Tuple[str, str, ContentFilterAction] ] = {} # Store conditional categories (identifier_words + block_words) - self.conditional_categories: Dict[str, Dict[str, Any]] = ( - {} - ) # category_name -> {identifier_words, block_words, action, severity} + self.conditional_categories: Dict[ + str, Dict[str, Any] + ] = {} # category_name -> {identifier_words, block_words, action, severity} # Competitor intent checker (optional; airline uses major_airlines.json, generic requires competitors) self._competitor_intent_checker: Optional[BaseCompetitorIntentChecker] = None @@ -1078,7 +1078,11 @@ class ContentFilterGuardrail(CustomGuardrail): return None # Always-block keywords are checked after exceptions. - for keyword, (category, severity, action) in self.always_block_category_keywords.items(): + for keyword, ( + category, + severity, + action, + ) in self.always_block_category_keywords.items(): keyword_pattern_str = self._keyword_to_regex_pattern(keyword) if " " in keyword: keyword_found = bool(re.search(keyword_pattern_str, text_lower)) diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/guardrail_benchmarks/test_eval.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/guardrail_benchmarks/test_eval.py index ca66b4da65..56398739b9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/guardrail_benchmarks/test_eval.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/guardrail_benchmarks/test_eval.py @@ -199,9 +199,7 @@ def _confusion_matrix(checker, cases: List[dict], label: str): precision = tp / (tp + fp) if (tp + fp) > 0 else 0 recall = tp / (tp + fn) if (tp + fn) > 0 else 0 f1 = ( - 2 * precision * recall / (precision + recall) - if (precision + recall) > 0 - else 0 + 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0 ) accuracy = (tp + tn) / total if total > 0 else 0 @@ -212,9 +210,18 @@ def _confusion_matrix(checker, cases: List[dict], label: str): avg_lat = sum(latencies) / len(latencies) if latencies else 0 metrics = { - "total": total, "tp": tp, "tn": tn, "fp": fp, "fn": fn, - "precision": precision, "recall": recall, "f1": f1, "accuracy": accuracy, - "p50": p50, "p95": p95, "avg_lat": avg_lat, + "total": total, + "tp": tp, + "tn": tn, + "fp": fp, + "fn": fn, + "precision": precision, + "recall": recall, + "f1": f1, + "accuracy": accuracy, + "p50": p50, + "p95": p95, + "avg_lat": avg_lat, } _print_confusion_report(label, metrics, wrong) result = _save_confusion_results(label, metrics, wrong, rows) @@ -586,4 +593,6 @@ class TestInvestmentLlmJudgeClaude: return _load_jsonl("block_investment.jsonl") def test_confusion_matrix(self, blocker, cases): - _confusion_matrix(blocker, cases, "Block Investment — LLM Judge (claude-haiku-4.5)") + _confusion_matrix( + blocker, cases, "Block Investment — LLM Judge (claude-haiku-4.5)" + ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/__init__.py index 237364f971..5060fade8c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/__init__.py @@ -14,7 +14,9 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" # Default to always-on. Only disable if the user explicitly sets default_on: false. # We check the raw guardrail dict because LitellmParams normalizes None → False, # making it impossible to distinguish "not set" from "explicitly false" via litellm_params. - _raw_default_on = cast(Dict[str, Any], guardrail).get("litellm_params", {}).get("default_on") + _raw_default_on = ( + cast(Dict[str, Any], guardrail).get("litellm_params", {}).get("default_on") + ) _default_on = False if _raw_default_on is False else True _callback = MCPEndUserPermissionGuardrail( diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py index 385143abc0..794edf092d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py @@ -86,7 +86,7 @@ class MCPSecurityGuardrail(CustomGuardrail): if not isinstance(server_url, str): continue if server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL_PREFIX): - name = server_url[len(LITELLM_PROXY_MCP_SERVER_URL_PREFIX):] + name = server_url[len(LITELLM_PROXY_MCP_SERVER_URL_PREFIX) :] if name: server_names.add(name) return server_names @@ -98,8 +98,8 @@ class MCPSecurityGuardrail(CustomGuardrail): if not tools or not isinstance(tools, list): return set() - requested_servers = ( - MCPSecurityGuardrail._extract_mcp_server_names_from_tools(tools) + requested_servers = MCPSecurityGuardrail._extract_mcp_server_names_from_tools( + tools ) if not requested_servers: return set() diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py index 2c42917294..1a119ec56b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py @@ -13,7 +13,10 @@ from typing import TYPE_CHECKING, Any, Literal, Optional, Type, cast from urllib.parse import urlparse from litellm._logging import verbose_proxy_logger -from litellm.integrations.custom_guardrail import CustomGuardrail, log_guardrail_information +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.llms.custom_httpx.http_handler import ( @@ -51,23 +54,33 @@ class NomaV2Guardrail(CustomGuardrail): block_failures: Optional[bool] = None, **kwargs: Any, ) -> None: - self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) + self.async_handler = get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback + ) self.api_key = api_key or os.environ.get("NOMA_API_KEY") - self.api_base = (api_base or os.environ.get("NOMA_API_BASE") or _DEFAULT_API_BASE).rstrip("/") + self.api_base = ( + api_base or os.environ.get("NOMA_API_BASE") or _DEFAULT_API_BASE + ).rstrip("/") self.application_id = application_id or os.environ.get("NOMA_APPLICATION_ID") if monitor_mode is None: - self.monitor_mode = os.environ.get("NOMA_MONITOR_MODE", "false").lower() == "true" + self.monitor_mode = ( + os.environ.get("NOMA_MONITOR_MODE", "false").lower() == "true" + ) else: self.monitor_mode = monitor_mode if block_failures is None: - self.block_failures = os.environ.get("NOMA_BLOCK_FAILURES", "true").lower() == "true" + self.block_failures = ( + os.environ.get("NOMA_BLOCK_FAILURES", "true").lower() == "true" + ) else: self.block_failures = block_failures if self._requires_api_key(api_base=self.api_base) and not self.api_key: - raise ValueError("Noma v2 guardrail requires api_key when using Noma SaaS endpoint") + raise ValueError( + "Noma v2 guardrail requires api_key when using Noma SaaS endpoint" + ) if "supported_event_hooks" not in kwargs: kwargs["supported_event_hooks"] = [ @@ -128,7 +141,9 @@ class NomaV2Guardrail(CustomGuardrail): ) -> dict: payload_request_data = deepcopy(request_data) if logging_obj is not None: - payload_request_data["litellm_logging_obj"] = getattr(logging_obj, "model_call_details", None) + payload_request_data["litellm_logging_obj"] = getattr( + logging_obj, "model_call_details", None + ) payload: dict[str, Any] = { "inputs": inputs, @@ -285,13 +300,17 @@ class NomaV2Guardrail(CustomGuardrail): action=action, ) - guardrail_status = "success" if action == _Action.NONE else "guardrail_intervened" + guardrail_status = ( + "success" if action == _Action.NONE else "guardrail_intervened" + ) return processed_inputs except NomaBlockedMessage as e: guardrail_status = "guardrail_intervened" guardrail_json_response = ( - response_json if isinstance(response_json, dict) else getattr(e, "detail", {"error": "blocked"}) + response_json + if isinstance(response_json, dict) + else getattr(e, "detail", {"error": "blocked"}) ) raise except Exception as e: diff --git a/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py b/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py index 1cfc805dbf..a07f537135 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py +++ b/litellm/proxy/guardrails/guardrail_hooks/onyx/onyx.py @@ -94,7 +94,6 @@ class OnyxGuardrail(CustomGuardrail): input_type: Literal["request", "response"], logging_obj: Optional["LiteLLMLoggingObj"] = None, ) -> GenericGuardrailAPIInputs: - conversation_id = ( logging_obj.litellm_call_id if logging_obj else str(uuid.uuid4()) ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/openai/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/openai/__init__.py index 8ca708fdcc..678d611fdc 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/openai/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/openai/__init__.py @@ -14,7 +14,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" guardrail_name = guardrail.get("guardrail_name") if not guardrail_name: raise ValueError("OpenAI Moderation: guardrail_name is required") - + openai_moderation_guardrail = OpenAIModerationGuardrail( guardrail_name=guardrail_name, **{ @@ -27,14 +27,11 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" }, ) - litellm.logging_callback_manager.add_litellm_callback( - openai_moderation_guardrail - ) + litellm.logging_callback_manager.add_litellm_callback(openai_moderation_guardrail) return openai_moderation_guardrail - guardrail_initializer_registry = { SupportedGuardrailIntegrations.OPENAI_MODERATION.value: initialize_guardrail, } diff --git a/litellm/proxy/guardrails/guardrail_hooks/openai/base.py b/litellm/proxy/guardrails/guardrail_hooks/openai/base.py index d93e05168a..872d09cd88 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/openai/base.py +++ b/litellm/proxy/guardrails/guardrail_hooks/openai/base.py @@ -49,4 +49,4 @@ class OpenAIGuardrailBase: user_prompt += text_content + "\n" result = user_prompt.strip() - return result if result else None \ No newline at end of file + return result if result else None diff --git a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py index 6160fb4143..4bd9434572 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py +++ b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py @@ -15,7 +15,7 @@ from fastapi import HTTPException from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( CustomGuardrail, - log_guardrail_information + log_guardrail_information, ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.custom_httpx.http_handler import ( diff --git a/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py b/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py index 55c7e72c36..2974febe02 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py +++ b/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py @@ -104,7 +104,7 @@ class PangeaHandler(CustomGuardrail): super().__init__( guardrail_name=guardrail_name, supported_event_hooks=supported_event_hooks, - **kwargs + **kwargs, ) verbose_proxy_logger.debug( f"Initialized Pangea Guardrail: name={guardrail_name}, recipe={pangea_input_recipe}, api_base={self.api_base}" @@ -172,7 +172,7 @@ class PangeaHandler(CustomGuardrail): user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, - call_type: str + call_type: str, ): transformer = None messages: Any = None @@ -184,10 +184,7 @@ class PangeaHandler(CustomGuardrail): ai_guard_payload = { "debug": False, - "input": { - "messages": messages, # type: ignore - "tools": data.get("tools") - }, + "input": {"messages": messages, "tools": data.get("tools")}, # type: ignore "event_type": "input", } if self.pangea_input_recipe: @@ -205,12 +202,11 @@ class PangeaHandler(CustomGuardrail): output = ai_guard_response.get("result", {}).get("output", {}) if call_type == "text_completion" or call_type == "atext_completion": - data = transformer.update_original_body(output["messages"]) # type: ignore + data = transformer.update_original_body(output["messages"]) # type: ignore else: data["messages"] = output["messages"] return data - @log_guardrail_information async def async_pre_call_hook( self, @@ -227,7 +223,9 @@ class PangeaHandler(CustomGuardrail): return data try: - return await self._async_pre_call_hook(user_api_key_dict, cache, data, call_type) + return await self._async_pre_call_hook( + user_api_key_dict, cache, data, call_type + ) except HTTPException: raise except Exception as e: @@ -237,7 +235,7 @@ class PangeaHandler(CustomGuardrail): "error": "Error in Pangea Guardrail", "guardrail_name": self.guardrail_name, "exceptions": str(e), - } + }, ) from e async def _async_post_call_success_hook( @@ -321,7 +319,9 @@ class PangeaHandler(CustomGuardrail): ) return data try: - return await self._async_post_call_success_hook(data, user_api_key_dict, response) + return await self._async_post_call_success_hook( + data, user_api_key_dict, response + ) except HTTPException: raise except Exception as e: @@ -331,7 +331,7 @@ class PangeaHandler(CustomGuardrail): "error": "Error in Pangea Guardrail", "guardrail_name": self.guardrail_name, "exceptions": str(e), - } + }, ) from e @staticmethod diff --git a/litellm/proxy/guardrails/guardrail_hooks/pillar/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/pillar/__init__.py index 2e4213a34d..5ef6f32ead 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/pillar/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/pillar/__init__.py @@ -32,9 +32,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" on_flagged_action=getattr(litellm_params, "on_flagged_action", "monitor"), event_hook=litellm_params.mode, default_on=litellm_params.default_on, - async_mode=_get_config_value( - litellm_params, optional_params, "async_mode" - ), + async_mode=_get_config_value(litellm_params, optional_params, "async_mode"), persist_session=_get_config_value( litellm_params, optional_params, "persist_session" ), diff --git a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py index ef22b09930..1b3f11e56f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py +++ b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py @@ -91,11 +91,15 @@ def _truncate_evidence_payload( step = max(1, len(evidence_text) // 2) while len(encoded.encode("utf-8")) > max_bytes and evidence_text: evidence_text = ( - evidence_text[:-step] if len(evidence_text) > step else evidence_text[:-1] + evidence_text[:-step] + if len(evidence_text) > step + else evidence_text[:-1] ) step = max(1, step // 2) truncated_text = ( - f"{evidence_text}...[truncated]" if evidence_text else "[truncated]" + f"{evidence_text}...[truncated]" + if evidence_text + else "[truncated]" ) working_entry["evidence"] = truncated_text working_entry["evidence_truncated"] = True @@ -120,7 +124,9 @@ def build_pillar_response_headers(metadata_store: Dict[str, Any]) -> Dict[str, s headers["x-pillar-flagged"] = str(metadata_store["pillar_flagged"]).lower() if "pillar_scanners" in metadata_store: - headers["x-pillar-scanners"] = _encode_json_for_header(metadata_store["pillar_scanners"]) + headers["x-pillar-scanners"] = _encode_json_for_header( + metadata_store["pillar_scanners"] + ) if "pillar_evidence" in metadata_store: truncated_evidence, encoded_value, truncated_flag = _truncate_evidence_payload( @@ -169,7 +175,9 @@ class PillarGuardrail(CustomGuardrail): SUPPORTED_FALLBACK_ACTIONS = ["allow", "block"] DEFAULT_FALLBACK_ACTION = "allow" BASE_API_URL = "https://api.pillar.security" - DEFAULT_TIMEOUT = 5.0 # 5 seconds - fast failure detection with graceful degradation + DEFAULT_TIMEOUT = ( + 5.0 # 5 seconds - fast failure detection with graceful degradation + ) def __init__( self, @@ -201,7 +209,9 @@ class PillarGuardrail(CustomGuardrail): LiteLLM virtual key context (user_id, team_id, key_alias, etc.) is always automatically passed as X-LiteLLM-* headers to enable application/user tracking. """ - self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) + self.async_handler = get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback + ) self.api_key = api_key or os.environ.get("PILLAR_API_KEY") if self.api_key is None: @@ -219,10 +229,14 @@ class PillarGuardrail(CustomGuardrail): self.on_flagged_action = action else: if action: - verbose_proxy_logger.warning(f"Invalid action '{action}', using default") + verbose_proxy_logger.warning( + f"Invalid action '{action}', using default" + ) self.on_flagged_action = self.DEFAULT_ON_FLAGGED_ACTION - verbose_proxy_logger.debug(f"Pillar Guardrail: Initialized with on_flagged_action: {self.on_flagged_action}") + verbose_proxy_logger.debug( + f"Pillar Guardrail: Initialized with on_flagged_action: {self.on_flagged_action}" + ) self.async_mode = self._resolve_bool_config( provided_value=async_mode, @@ -260,14 +274,18 @@ class PillarGuardrail(CustomGuardrail): ) self.fallback_on_error = self.DEFAULT_FALLBACK_ACTION - verbose_proxy_logger.debug(f"Pillar Guardrail: Initialized with fallback_on_error: {self.fallback_on_error}") + verbose_proxy_logger.debug( + f"Pillar Guardrail: Initialized with fallback_on_error: {self.fallback_on_error}" + ) # Set timeout with graceful fallback on invalid configuration if timeout is not None: self.timeout = timeout else: try: - self.timeout = float(os.environ.get("PILLAR_TIMEOUT", str(self.DEFAULT_TIMEOUT))) + self.timeout = float( + os.environ.get("PILLAR_TIMEOUT", str(self.DEFAULT_TIMEOUT)) + ) except (ValueError, TypeError): verbose_proxy_logger.warning( f"Pillar Guardrail: Invalid PILLAR_TIMEOUT value '{os.environ.get('PILLAR_TIMEOUT')}', " @@ -330,14 +348,18 @@ class PillarGuardrail(CustomGuardrail): """ event_type = GuardrailEventHooks.pre_call if self.should_run_guardrail(data=data, event_type=event_type) is not True: - verbose_proxy_logger.debug(f"Pillar Guardrail: Pre-call scanning disabled for {self.guardrail_name}") + verbose_proxy_logger.debug( + f"Pillar Guardrail: Pre-call scanning disabled for {self.guardrail_name}" + ) return data verbose_proxy_logger.debug("Pillar Guardrail: Pre-call hook") result = await self.run_pillar_guardrail(data, user_api_key_dict) # Add guardrail name to response headers - add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) + add_guardrail_to_applied_guardrails_header( + request_data=data, guardrail_name=self.guardrail_name + ) return result @@ -373,14 +395,18 @@ class PillarGuardrail(CustomGuardrail): """ event_type = GuardrailEventHooks.during_call if self.should_run_guardrail(data=data, event_type=event_type) is not True: - verbose_proxy_logger.debug(f"Pillar Guardrail: During-call scanning disabled for {self.guardrail_name}") + verbose_proxy_logger.debug( + f"Pillar Guardrail: During-call scanning disabled for {self.guardrail_name}" + ) return data verbose_proxy_logger.debug("Pillar Guardrail: During-call moderation hook") result = await self.run_pillar_guardrail(data, user_api_key_dict) # Add guardrail name to response headers - add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) + add_guardrail_to_applied_guardrails_header( + request_data=data, guardrail_name=self.guardrail_name + ) return result @@ -407,7 +433,9 @@ class PillarGuardrail(CustomGuardrail): """ event_type = GuardrailEventHooks.post_call if self.should_run_guardrail(data=data, event_type=event_type) is not True: - verbose_proxy_logger.debug(f"Pillar Guardrail: Post-call scanning disabled for {self.guardrail_name}") + verbose_proxy_logger.debug( + f"Pillar Guardrail: Post-call scanning disabled for {self.guardrail_name}" + ) return response verbose_proxy_logger.debug("Pillar Guardrail: Post-call hook") @@ -415,11 +443,15 @@ class PillarGuardrail(CustomGuardrail): # Extract response messages in the format Pillar expects response_dict = response.model_dump() if hasattr(response, "model_dump") else {} # type: ignore[union-attr] response_messages = [ - choice.get("message") for choice in response_dict.get("choices", []) if choice.get("message") + choice.get("message") + for choice in response_dict.get("choices", []) + if choice.get("message") ] if not response_messages: - verbose_proxy_logger.debug("Pillar Guardrail: No response content to scan, skipping post-call analysis") + verbose_proxy_logger.debug( + "Pillar Guardrail: No response content to scan, skipping post-call analysis" + ) return response # Create complete conversation: original messages + response messages @@ -430,7 +462,9 @@ class PillarGuardrail(CustomGuardrail): await self.run_pillar_guardrail(post_call_data, user_api_key_dict) # Add guardrail name to response headers - add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) + add_guardrail_to_applied_guardrails_header( + request_data=data, guardrail_name=self.guardrail_name + ) return response @@ -438,7 +472,9 @@ class PillarGuardrail(CustomGuardrail): # CORE LOGIC METHOD # ========================================================================= - async def run_pillar_guardrail(self, data: dict, user_api_key_dict: UserAPIKeyAuth) -> dict: + async def run_pillar_guardrail( + self, data: dict, user_api_key_dict: UserAPIKeyAuth + ) -> dict: """ Core method to run the Pillar guardrail scan. @@ -454,7 +490,9 @@ class PillarGuardrail(CustomGuardrail): """ # Check if messages are present if not data.get("messages"): - verbose_proxy_logger.debug("Pillar Guardrail: No messages detected, bypassing security scan") + verbose_proxy_logger.debug( + "Pillar Guardrail: No messages detected, bypassing security scan" + ) return data try: @@ -476,7 +514,9 @@ class PillarGuardrail(CustomGuardrail): raise e # Handle API communication errors based on fallback_on_error setting - verbose_proxy_logger.error(f"Pillar Guardrail: API communication failed - {str(e)}") + verbose_proxy_logger.error( + f"Pillar Guardrail: API communication failed - {str(e)}" + ) return self._handle_api_error(e, data) @@ -536,7 +576,7 @@ class PillarGuardrail(CustomGuardrail): headers: Dict[str, str] = { "Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json", - } + } # Add Pillar-specific headers based on configuration self._set_bool_header(headers, "plr_scanners", self.include_scanners) @@ -560,7 +600,9 @@ class PillarGuardrail(CustomGuardrail): return headers - def _set_bool_header(self, headers: Dict[str, str], header_name: str, value: Optional[bool]) -> None: + def _set_bool_header( + self, headers: Dict[str, str], header_name: str, value: Optional[bool] + ) -> None: """Apply a boolean value as a lowercase string HTTP header when provided.""" if value is None: @@ -701,7 +743,9 @@ class PillarGuardrail(CustomGuardrail): ) return payload - async def _call_pillar_api(self, headers: Dict[str, str], payload: Dict[str, Any]) -> Dict[str, Any]: + async def _call_pillar_api( + self, headers: Dict[str, str], payload: Dict[str, Any] + ) -> Dict[str, Any]: """ Call the Pillar API and return the response. @@ -726,10 +770,14 @@ class PillarGuardrail(CustomGuardrail): flagged = res.get("flagged") session_id = res.get("session_id") - verbose_proxy_logger.debug(f"Pillar Guardrail: Analysis complete - flagged={flagged}, session={session_id}") + verbose_proxy_logger.debug( + f"Pillar Guardrail: Analysis complete - flagged={flagged}, session={session_id}" + ) return res - def _process_pillar_response(self, pillar_response: Dict[str, Any], original_data: dict) -> None: + def _process_pillar_response( + self, pillar_response: Dict[str, Any], original_data: dict + ) -> None: """ Process the Pillar API response and handle detections based on configuration. @@ -746,19 +794,25 @@ class PillarGuardrail(CustomGuardrail): flagged = pillar_response.get("flagged", False) metadata_field = get_metadata_variable_name_from_kwargs(original_data) - if metadata_field not in original_data or not isinstance(original_data.get(metadata_field), dict): + if metadata_field not in original_data or not isinstance( + original_data.get(metadata_field), dict + ): original_data[metadata_field] = {} metadata_store = original_data[metadata_field] # Backwards compatibility - ensure metadata alias exists when different key used if metadata_field != "metadata": - if "metadata" not in original_data or not isinstance(original_data.get("metadata"), dict): + if "metadata" not in original_data or not isinstance( + original_data.get("metadata"), dict + ): original_data["metadata"] = metadata_store # Store session_id from Pillar response for potential reuse pillar_session_id = pillar_response.get("session_id") if pillar_session_id: - verbose_proxy_logger.debug(f"Pillar Guardrail: Received session_id from server: {pillar_session_id}") + verbose_proxy_logger.debug( + f"Pillar Guardrail: Received session_id from server: {pillar_session_id}" + ) # Store in request metadata for use in subsequent hooks if "pillar_session_id" not in metadata_store: metadata_store["pillar_session_id"] = pillar_session_id @@ -776,7 +830,9 @@ class PillarGuardrail(CustomGuardrail): if self.on_flagged_action == "block": self._raise_pillar_detection_exception(pillar_response) elif self.on_flagged_action == "mask": - verbose_proxy_logger.info("Pillar Guardrail: Masking mode - masking flagged content") + verbose_proxy_logger.info( + "Pillar Guardrail: Masking mode - masking flagged content" + ) masked_messages = pillar_response.get("masked_session_messages", []) if masked_messages: original_data["messages"] = masked_messages @@ -785,11 +841,15 @@ class PillarGuardrail(CustomGuardrail): "Pillar Guardrail: Masking requested but no masked_session_messages in response" ) elif self.on_flagged_action == "monitor": - verbose_proxy_logger.info("Pillar Guardrail: Monitoring mode - allowing flagged content to proceed") + verbose_proxy_logger.info( + "Pillar Guardrail: Monitoring mode - allowing flagged content to proceed" + ) build_pillar_response_headers(metadata_store) - def _raise_pillar_detection_exception(self, pillar_response: Dict[str, Any]) -> None: + def _raise_pillar_detection_exception( + self, pillar_response: Dict[str, Any] + ) -> None: """ Raise an HTTPException for Pillar security detections. @@ -802,7 +862,7 @@ class PillarGuardrail(CustomGuardrail): pillar_response_dict = { "session_id": pillar_response.get("session_id"), } - + # Conditionally include scanners and evidence based on config if self.include_scanners: pillar_response_dict["scanners"] = pillar_response.get("scanners", {}) @@ -815,7 +875,9 @@ class PillarGuardrail(CustomGuardrail): "pillar_response": pillar_response_dict, } - verbose_proxy_logger.warning("Pillar Guardrail: Request blocked - Security threats detected") + verbose_proxy_logger.warning( + "Pillar Guardrail: Request blocked - Security threats detected" + ) raise HTTPException(status_code=400, detail=error_detail) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index b84c74bee4..8d94a7051a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -1172,9 +1172,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): yield mock_response_stream except Exception as e: - verbose_proxy_logger.error( - f"Error masking streaming PII output: {str(e)}" - ) + verbose_proxy_logger.error(f"Error masking streaming PII output: {str(e)}") for chunk in all_chunks: yield chunk diff --git a/litellm/proxy/guardrails/guardrail_hooks/qualifire/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/qualifire/__init__.py index 8c29cfcd30..b9f7aed4a2 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/qualifire/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/qualifire/__init__.py @@ -19,8 +19,12 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" hallucinations_check=getattr(litellm_params, "hallucinations_check", None), grounding_check=getattr(litellm_params, "grounding_check", None), pii_check=getattr(litellm_params, "pii_check", None), - content_moderation_check=getattr(litellm_params, "content_moderation_check", None), - tool_selection_quality_check=getattr(litellm_params, "tool_selection_quality_check", None), + content_moderation_check=getattr( + litellm_params, "content_moderation_check", None + ), + tool_selection_quality_check=getattr( + litellm_params, "tool_selection_quality_check", None + ), assertions=getattr(litellm_params, "assertions", None), on_flagged=getattr(litellm_params, "on_flagged", "block"), guardrail_name=guardrail.get("guardrail_name", ""), diff --git a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py index ad05e7656c..10a50c39e3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py +++ b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/route_loader.py @@ -52,14 +52,18 @@ class SemanticGuardRouteLoader: def load_custom_routes_file(file_path: str) -> List[Dict[str, Any]]: """Load custom routes from a YAML file.""" if not os.path.exists(file_path): - raise ValueError(f"SemanticGuard: custom routes file not found: {file_path}") + raise ValueError( + f"SemanticGuard: custom routes file not found: {file_path}" + ) with open(file_path, "r") as f: data = yaml.safe_load(f) if isinstance(data, list): return data if isinstance(data, dict): return [data] - raise ValueError(f"SemanticGuard: invalid custom routes file format in {file_path}") + raise ValueError( + f"SemanticGuard: invalid custom routes file format in {file_path}" + ) @classmethod def build_routes( diff --git a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py index 465c4a86c2..be48991500 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/semantic_guard/semantic_guard.py @@ -52,7 +52,9 @@ class SemanticGuardrail(CustomGuardrail): custom_routes_file: Optional[str] = None, custom_routes: Optional[List[Dict[str, Any]]] = None, on_flagged_action: str = "block", - event_hook: Optional[Union[GuardrailEventHooks, List[GuardrailEventHooks], Mode]] = None, + event_hook: Optional[ + Union[GuardrailEventHooks, List[GuardrailEventHooks], Mode] + ] = None, default_on: bool = False, **kwargs, ): @@ -86,11 +88,13 @@ class SemanticGuardrail(CustomGuardrail): "Provide route_templates or custom_routes." ) - self.semantic_router: "SemanticRouter" = SemanticGuardRouteLoader.build_semantic_router( - routes=routes, - litellm_router=llm_router, - embedding_model=embedding_model, - global_threshold=similarity_threshold, + self.semantic_router: "SemanticRouter" = ( + SemanticGuardRouteLoader.build_semantic_router( + routes=routes, + litellm_router=llm_router, + embedding_model=embedding_model, + global_threshold=similarity_threshold, + ) ) self.route_count = len(routes) diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index bec76acc50..6dd0288cb0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -109,8 +109,16 @@ class ToolPermissionGuardrail(CustomGuardrail): self._compiled_rule_patterns[rule.id] = compiled_patterns # Normalize to lowercase for case-insensitive handling - self.default_action = default_action.lower() if isinstance(default_action, str) else default_action - self.on_disallowed_action = on_disallowed_action.lower() if isinstance(on_disallowed_action, str) else on_disallowed_action + self.default_action = ( + default_action.lower() + if isinstance(default_action, str) + else default_action + ) + self.on_disallowed_action = ( + on_disallowed_action.lower() + if isinstance(on_disallowed_action, str) + else on_disallowed_action + ) verbose_proxy_logger.debug( "Tool Permission Guardrail initialized with %d rules, default_action: %s", @@ -246,7 +254,11 @@ class ToolPermissionGuardrail(CustomGuardrail): return {} def _collect_argument_paths( - self, value: Any, current_path: str, collected: Dict[str, List[Any]], depth: int = 0 + self, + value: Any, + current_path: str, + collected: Dict[str, List[Any]], + depth: int = 0, ) -> None: from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py index 368948414e..12510d051d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py @@ -154,9 +154,10 @@ class ToolPolicyGuardrail(CustomGuardrail): if not tool_names: return inputs - object_permission_id, team_object_permission_id = ( - _get_request_object_permission_ids(request_data) - ) + ( + object_permission_id, + team_object_permission_id, + ) = _get_request_object_permission_ids(request_data) from litellm.proxy.db.tool_registry_writer import get_tool_policy_registry registry = get_tool_policy_registry() @@ -172,7 +173,8 @@ class ToolPolicyGuardrail(CustomGuardrail): blocked = [name for name in tool_names if policy_map.get(name) == "blocked"] if blocked: verbose_proxy_logger.warning( - "ToolPolicyGuardrail: blocking tool(s) %s (input_policy=blocked)", blocked + "ToolPolicyGuardrail: blocking tool(s) %s (input_policy=blocked)", + blocked, ) raise HTTPException( status_code=400, @@ -199,7 +201,9 @@ class ToolPolicyGuardrail(CustomGuardrail): if msg.get("role") != "tool": continue tool_call_id = msg.get("tool_call_id") - source_tool = tc_id_to_name.get(tool_call_id, "") if tool_call_id else "" + source_tool = ( + tc_id_to_name.get(tool_call_id, "") if tool_call_id else "" + ) if not source_tool: continue if registry.get_output_policy(source_tool) == "untrusted": diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index c35eadcb6f..84bbf6d20e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -85,7 +85,6 @@ class UnifiedLLMGuardrails(CustomLogger): add_guardrail_to_applied_guardrails_header, ) - verbose_proxy_logger.debug("Running UnifiedLLMGuardrails pre-call hook") guardrail_to_apply: CustomGuardrail = data.pop("guardrail_to_apply", None) diff --git a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py index fbf9e1f29f..e0e59bfef3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py @@ -42,17 +42,29 @@ class ZscalerAIGuard(CustomGuardrail): "ZSCALER_AI_GUARD_URL", "https://api.us1.zseclipse.net/v1/detection/execute-policy", ) - self.policy_id = policy_id if policy_id is not None else int(os.getenv("ZSCALER_AI_GUARD_POLICY_ID", -1)) + self.policy_id = ( + policy_id + if policy_id is not None + else int(os.getenv("ZSCALER_AI_GUARD_POLICY_ID", -1)) + ) self.api_key = api_key or os.getenv("ZSCALER_AI_GUARD_API_KEY") - self.send_user_api_key_alias = send_user_api_key_alias if send_user_api_key_alias is not None else os.getenv( - "SEND_USER_API_KEY_ALIAS", "False" - ).lower() in ("true", "1") - self.send_user_api_key_user_id = send_user_api_key_user_id if send_user_api_key_user_id is not None else os.getenv( - "SEND_USER_API_KEY_USER_ID", "False" - ).lower() in ("true", "1") - self.send_user_api_key_team_id = send_user_api_key_team_id if send_user_api_key_team_id is not None else os.getenv( - "SEND_USER_API_KEY_TEAM_ID", "False" - ).lower() in ("true", "1") + self.send_user_api_key_alias = ( + send_user_api_key_alias + if send_user_api_key_alias is not None + else os.getenv("SEND_USER_API_KEY_ALIAS", "False").lower() in ("true", "1") + ) + self.send_user_api_key_user_id = ( + send_user_api_key_user_id + if send_user_api_key_user_id is not None + else os.getenv("SEND_USER_API_KEY_USER_ID", "False").lower() + in ("true", "1") + ) + self.send_user_api_key_team_id = ( + send_user_api_key_team_id + if send_user_api_key_team_id is not None + else os.getenv("SEND_USER_API_KEY_TEAM_ID", "False").lower() + in ("true", "1") + ) verbose_proxy_logger.debug( f"""send_user_api_key_alias: {self.send_user_api_key_alias}, @@ -65,7 +77,9 @@ class ZscalerAIGuard(CustomGuardrail): verbose_proxy_logger.debug("ZscalerAIGuard Initializing ...") @staticmethod - def _resolve_metadata_value(request_data: Optional[dict], key: str) -> Optional[str]: + def _resolve_metadata_value( + request_data: Optional[dict], key: str + ) -> Optional[str]: """ Resolve metadata value from request_data, checking both metadata locations. @@ -157,17 +171,20 @@ class ZscalerAIGuard(CustomGuardrail): kwargs = {} if self.send_user_api_key_alias: - kwargs["user_api_key_alias"] = self._resolve_metadata_value( - request_data, "user_api_key_alias" - ) or "N/A" + kwargs["user_api_key_alias"] = ( + self._resolve_metadata_value(request_data, "user_api_key_alias") + or "N/A" + ) if self.send_user_api_key_team_id: - kwargs["user_api_key_team_id"] = self._resolve_metadata_value( - request_data, "user_api_key_team_id" - ) or "N/A" + kwargs["user_api_key_team_id"] = ( + self._resolve_metadata_value(request_data, "user_api_key_team_id") + or "N/A" + ) if self.send_user_api_key_user_id: - kwargs["user_api_key_user_id"] = self._resolve_metadata_value( - request_data, "user_api_key_user_id" - ) or "N/A" + kwargs["user_api_key_user_id"] = ( + self._resolve_metadata_value(request_data, "user_api_key_user_id") + or "N/A" + ) verbose_proxy_logger.debug(f"inside apply_guardrail kwargs: {kwargs}") zscaler_ai_guard_result = None @@ -184,7 +201,9 @@ class ZscalerAIGuard(CustomGuardrail): content=concatenated_text, **kwargs, ) - verbose_proxy_logger.debug(f"response from zscaler ai guards: {zscaler_ai_guard_result}") + verbose_proxy_logger.debug( + f"response from zscaler ai guards: {zscaler_ai_guard_result}" + ) if ( zscaler_ai_guard_result and zscaler_ai_guard_result.get("action") == "BLOCK" diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index 46ea667f46..d41be370f7 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -506,7 +506,9 @@ class InMemoryGuardrailHandler: guardrail_type, ) - _guardrail_class = get_instance_fn(guardrail_type, config_file_path=config_file_path) + _guardrail_class = get_instance_fn( + guardrail_type, config_file_path=config_file_path + ) mode = litellm_params.mode if mode is None: diff --git a/litellm/proxy/guardrails/tool_name_extraction.py b/litellm/proxy/guardrails/tool_name_extraction.py index db24fa2277..c554c4fc9a 100644 --- a/litellm/proxy/guardrails/tool_name_extraction.py +++ b/litellm/proxy/guardrails/tool_name_extraction.py @@ -40,22 +40,28 @@ def _extract_mcp_tool_names(data: dict) -> List[str]: def _register_standalone_extractors() -> None: if STANDALONE_EXTRACTORS: return - STANDALONE_EXTRACTORS[CallTypes.generate_content.value] = _extract_generate_content_tool_names - STANDALONE_EXTRACTORS[CallTypes.agenerate_content.value] = _extract_generate_content_tool_names + STANDALONE_EXTRACTORS[ + CallTypes.generate_content.value + ] = _extract_generate_content_tool_names + STANDALONE_EXTRACTORS[ + CallTypes.agenerate_content.value + ] = _extract_generate_content_tool_names STANDALONE_EXTRACTORS[CallTypes.call_mcp_tool.value] = _extract_mcp_tool_names # Tool-capable call types (routes that can send tools in the request) -TOOL_CAPABLE_CALL_TYPES = frozenset({ - CallTypes.completion.value, - CallTypes.acompletion.value, - CallTypes.responses.value, - CallTypes.aresponses.value, - CallTypes.anthropic_messages.value, - CallTypes.generate_content.value, - CallTypes.agenerate_content.value, - CallTypes.call_mcp_tool.value, -}) +TOOL_CAPABLE_CALL_TYPES = frozenset( + { + CallTypes.completion.value, + CallTypes.acompletion.value, + CallTypes.responses.value, + CallTypes.aresponses.value, + CallTypes.anthropic_messages.value, + CallTypes.generate_content.value, + CallTypes.agenerate_content.value, + CallTypes.call_mcp_tool.value, + } +) def extract_request_tool_names(route: str, data: dict) -> List[str]: diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index 3314c5ca2e..529949c6dd 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -386,19 +386,11 @@ async def guardrails_usage_detail( _litellm_params = getattr(guardrail, "litellm_params", None) or ( guardrail.get("litellm_params") if isinstance(guardrail, dict) else None ) - litellm_params = ( - _litellm_params - if isinstance(_litellm_params, dict) - else {} - ) + litellm_params = _litellm_params if isinstance(_litellm_params, dict) else {} _guardrail_info = getattr(guardrail, "guardrail_info", None) or ( guardrail.get("guardrail_info") if isinstance(guardrail, dict) else None ) - guardrail_info = ( - _guardrail_info - if isinstance(_guardrail_info, dict) - else {} - ) + guardrail_info = _guardrail_info if isinstance(_guardrail_info, dict) else {} _guardrail_name = getattr(guardrail, "guardrail_name", None) or ( guardrail.get("guardrail_name") if isinstance(guardrail, dict) else None ) diff --git a/litellm/proxy/guardrails/usage_tracking.py b/litellm/proxy/guardrails/usage_tracking.py index 248f3a1987..8907c9201a 100644 --- a/litellm/proxy/guardrails/usage_tracking.py +++ b/litellm/proxy/guardrails/usage_tracking.py @@ -85,7 +85,9 @@ async def process_spend_logs_guardrail_usage( date_key = _date_str(start_time) for entry in _parse_guardrail_info_from_payload(payload): - guardrail_id = entry.get("guardrail_id") or entry.get("guardrail_name") or "" + guardrail_id = ( + entry.get("guardrail_id") or entry.get("guardrail_name") or "" + ) if not guardrail_id: continue key = (guardrail_id, date_key) @@ -98,12 +100,14 @@ async def process_spend_logs_guardrail_usage( else: daily_guardrail[key]["flagged_count"] += 1 policy_id = entry.get("policy_id") - index_rows.append({ - "request_id": request_id, - "guardrail_id": guardrail_id, - "policy_id": policy_id, - "start_time": start_time, - }) + index_rows.append( + { + "request_id": request_id, + "guardrail_id": guardrail_id, + "policy_id": policy_id, + "start_time": start_time, + } + ) if not daily_guardrail and not index_rows: return @@ -119,12 +123,14 @@ async def process_spend_logs_guardrail_usage( st = datetime.fromisoformat(st.replace("Z", "+00:00")) except (ValueError, TypeError): continue - index_data.append({ - "request_id": r["request_id"], - "guardrail_id": r["guardrail_id"], - "policy_id": r.get("policy_id"), - "start_time": st, - }) + index_data.append( + { + "request_id": r["request_id"], + "guardrail_id": r["guardrail_id"], + "policy_id": r.get("policy_id"), + "start_time": st, + } + ) try: await prisma_client.db.litellm_spendlogguardrailindex.create_many( data=index_data, diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index b401528f64..ef9436f2d8 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -1143,7 +1143,9 @@ async def _db_health_readiness_check(): try: time_diff = datetime.now() - db_health_cache["last_updated"] - if db_health_cache["status"] == "connected" and time_diff < timedelta(seconds=15): + if db_health_cache["status"] == "connected" and time_diff < timedelta( + seconds=15 + ): return db_health_cache if prisma_client is None: @@ -1167,7 +1169,10 @@ async def _db_health_readiness_check(): verbose_proxy_logger.info( "_db_health_readiness_check: reconnect succeeded" ) - db_health_cache = {"status": "connected", "last_updated": datetime.now()} + db_health_cache = { + "status": "connected", + "last_updated": datetime.now(), + } return db_health_cache except Exception: verbose_proxy_logger.error( diff --git a/litellm/proxy/health_endpoints/health_app_factory.py b/litellm/proxy/health_endpoints/health_app_factory.py index 7737969318..c4fe383365 100644 --- a/litellm/proxy/health_endpoints/health_app_factory.py +++ b/litellm/proxy/health_endpoints/health_app_factory.py @@ -1,7 +1,8 @@ from fastapi import FastAPI from litellm.proxy.health_endpoints._health_endpoints import router as health_router + def build_health_app(): health_app = FastAPI(title="LiteLLM Health Endpoints") health_app.include_router(health_router) - return health_app \ No newline at end of file + return health_app diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 5bebcc9207..ba8a4672ca 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -60,17 +60,20 @@ else: RateLimitStatus = Dict[str, Any] RateLimitDescriptor = Dict[str, Any] + class BatchFileUsage(BaseModel): """ Internal model for batch file usage tracking, used for batch rate limiting """ + total_tokens: int request_count: int + class _PROXY_BatchRateLimiter(CustomLogger): """ Rate limiter for batch API requests. - + Handles rate limiting at two points: 1. Batch submission - reads input file and reserves capacity 2. Batch completion - reads output file and adjusts for actual usage @@ -83,7 +86,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): ): """ Initialize the batch rate limiter. - + Note: These dependencies are automatically injected by ProxyLogging._add_proxy_hooks() when this hook is registered in PROXY_HOOKS. See BATCH_RATE_LIMITER_INTEGRATION.md. @@ -106,22 +109,29 @@ class _PROXY_BatchRateLimiter(CustomLogger): # Find the descriptor for this status descriptor_index = next( - (i for i, d in enumerate(descriptors) - if d.get("key") == status.get("descriptor_key")), - 0 + ( + i + for i, d in enumerate(descriptors) + if d.get("key") == status.get("descriptor_key") + ), + 0, ) - descriptor: RateLimitDescriptor = descriptors[descriptor_index] if descriptors else {"key": "", "value": "", "rate_limit": None} - + descriptor: RateLimitDescriptor = ( + descriptors[descriptor_index] + if descriptors + else {"key": "", "value": "", "rate_limit": None} + ) + now = datetime.now().timestamp() window_size = self.parallel_request_limiter.window_size reset_time = now + window_size reset_time_formatted = datetime.fromtimestamp(reset_time).strftime( "%Y-%m-%d %H:%M:%S UTC" ) - + remaining_display = max(0, status["limit_remaining"]) current_limit = status["current_limit"] - + if limit_type == "requests": detail = ( f"Batch rate limit exceeded for {descriptor.get('key', 'unknown')}: {descriptor.get('value', 'unknown')}. " @@ -136,7 +146,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): f"out of {current_limit} TPM limit. " f"Limit resets at: {reset_time_formatted}" ) - + raise HTTPException( status_code=429, detail=detail, @@ -155,7 +165,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): ) -> None: """ Check rate limits and increment counters by the batch amounts. - + Raises HTTPException if any limit would be exceeded. """ from litellm.types.caching import RedisPipelineIncrementOperation @@ -168,30 +178,32 @@ class _PROXY_BatchRateLimiter(CustomLogger): tpm_limit_type=None, model_has_failures=False, ) - + # Check current usage without incrementing rate_limit_response = await self.parallel_request_limiter.should_rate_limit( descriptors=descriptors, parent_otel_span=user_api_key_dict.parent_otel_span, read_only=True, ) - + # Verify batch won't exceed any limits for status in rate_limit_response["statuses"]: rate_limit_type = status["rate_limit_type"] limit_remaining = status["limit_remaining"] - + required_capacity = ( - batch_usage.request_count if rate_limit_type == "requests" - else batch_usage.total_tokens if rate_limit_type == "tokens" + batch_usage.request_count + if rate_limit_type == "requests" + else batch_usage.total_tokens + if rate_limit_type == "tokens" else 0 ) - + if required_capacity > limit_remaining: self._raise_rate_limit_error( status, descriptors, batch_usage, rate_limit_type ) - + # Build pipeline operations for batch increments # Reuse the same keys that descriptors check pipeline_operations: List[RedisPipelineIncrementOperation] = [] @@ -229,7 +241,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): ttl=self.parallel_request_limiter.window_size, ) ) - + # Execute increments if pipeline_operations: await self.parallel_request_limiter.async_increment_tokens_with_ttl_preservation( @@ -245,12 +257,12 @@ class _PROXY_BatchRateLimiter(CustomLogger): ) -> BatchFileUsage: """ Count number of requests and tokens in a batch input file. - + Args: file_id: The file ID to read custom_llm_provider: The custom LLM provider to use for token encoding user_api_key_dict: User authentication information for file access (required for managed files) - + Returns: BatchFileUsage with total_tokens and request_count """ @@ -259,6 +271,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, ) + # Managed files require bypassing the HTTP endpoint (which runs access-check hooks) # and calling the managed files hook directly with the user's credentials. is_managed_file = _is_base64_encoded_unified_file_id(file_id) @@ -275,9 +288,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): user_api_key_dict=user_api_key_dict, ) - file_content_as_dict = _get_file_content_as_dictionary( - file_content.content - ) + file_content_as_dict = _get_file_content_as_dictionary(file_content.content) input_file_usage = _get_batch_job_input_file_usage( file_content_dictionary=file_content_as_dict, @@ -288,7 +299,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): total_tokens=input_file_usage.total_tokens, request_count=request_count, ) - + except Exception as e: verbose_proxy_logger.error( f"Error counting input file usage for {file_id}: {str(e)}" @@ -302,14 +313,14 @@ class _PROXY_BatchRateLimiter(CustomLogger): ) -> Any: """ Fetch file content from managed files hook. - + This is needed for managed files because they require proper user context to verify file ownership and access permissions. - + Args: file_id: The managed file ID (base64 encoded) user_api_key_dict: User authentication information - + Returns: HttpxBinaryResponseContent with the file content """ @@ -323,29 +334,25 @@ class _PROXY_BatchRateLimiter(CustomLogger): f"Cannot import proxy_server dependencies: {str(e)}. " "Managed files require proxy_server to be initialized." ) - + # Get the managed files hook if proxy_logging_obj is None: raise ValueError( "proxy_logging_obj not available. Cannot access managed files hook." ) - + managed_files_obj = proxy_logging_obj.get_proxy_hook("managed_files") if managed_files_obj is None: raise ValueError( "Managed files hook not found. Cannot access managed file." ) - + if not isinstance(managed_files_obj, BaseFileEndpoints): - raise ValueError( - "Managed files hook is not a BaseFileEndpoints instance." - ) - + raise ValueError("Managed files hook is not a BaseFileEndpoints instance.") + if llm_router is None: - raise ValueError( - "llm_router not available. Cannot access managed files." - ) - + raise ValueError("llm_router not available. Cannot access managed files.") + # Use the managed files hook to get file content # This properly handles user permissions and file ownership file_content = await managed_files_obj.afile_content( @@ -353,7 +360,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): litellm_parent_otel_span=user_api_key_dict.parent_otel_span, llm_router=llm_router, ) - + return file_content async def async_pre_call_hook( @@ -365,7 +372,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): ) -> Union[Exception, str, Dict, None]: """ Pre-call hook for batch operations. - + Only handles batch creation (acreate_batch): - Reads input file - Counts tokens and requests @@ -433,7 +440,9 @@ class _PROXY_BatchRateLimiter(CustomLogger): batch_usage=batch_usage, ) - verbose_proxy_logger.debug("Batch rate limit check passed, counters incremented") + verbose_proxy_logger.debug( + "Batch rate limit check passed, counters incremented" + ) return data except HTTPException: @@ -445,10 +454,3 @@ class _PROXY_BatchRateLimiter(CustomLogger): ) # Don't block the request if rate limiting fails return data - - - - - - - diff --git a/litellm/proxy/hooks/dynamic_rate_limiter.py b/litellm/proxy/hooks/dynamic_rate_limiter.py index f1c1d487cc..14fde51210 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter.py @@ -103,9 +103,9 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger): """ try: # Get model info first for conversion - model_group_info: Optional[ModelGroupInfo] = ( - self.llm_router.get_model_group_info(model_group=model) - ) + model_group_info: Optional[ + ModelGroupInfo + ] = self.llm_router.get_model_group_info(model_group=model) weight: float = 1 if ( @@ -277,16 +277,16 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger): ) = await self.check_available_usage( model=model_info["model_name"], priority=key_priority ) - response._hidden_params["additional_headers"] = ( - { # Add additional response headers - easier debugging - "x-litellm-model_group": model_info["model_name"], - "x-ratelimit-remaining-litellm-project-tokens": available_tpm, - "x-ratelimit-remaining-litellm-project-requests": available_rpm, - "x-ratelimit-remaining-model-tokens": model_tpm, - "x-ratelimit-remaining-model-requests": model_rpm, - "x-ratelimit-current-active-projects": active_projects, - } - ) + response._hidden_params[ + "additional_headers" + ] = { # Add additional response headers - easier debugging + "x-litellm-model_group": model_info["model_name"], + "x-ratelimit-remaining-litellm-project-tokens": available_tpm, + "x-ratelimit-remaining-litellm-project-requests": available_rpm, + "x-ratelimit-remaining-model-tokens": model_tpm, + "x-ratelimit-remaining-model-requests": model_rpm, + "x-ratelimit-current-active-projects": active_projects, + } return response return await super().async_post_call_success_hook( diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py index 95c9c80612..2d61203ad5 100644 --- a/litellm/proxy/hooks/key_management_event_hooks.py +++ b/litellm/proxy/hooks/key_management_event_hooks.py @@ -151,9 +151,7 @@ class KeyManagementEventHooks: or f"virtual-key-{existing_key_row.token}" ) new_secret_name = ( - response.key_alias - or data.key_alias - or initial_secret_name + response.key_alias or data.key_alias or initial_secret_name ) verbose_proxy_logger.info( "Updating secret in secret manager: secret_name=%s", diff --git a/litellm/proxy/hooks/litellm_skills/main.py b/litellm/proxy/hooks/litellm_skills/main.py index c2ad1e2944..83e419bc23 100644 --- a/litellm/proxy/hooks/litellm_skills/main.py +++ b/litellm/proxy/hooks/litellm_skills/main.py @@ -46,12 +46,12 @@ class SkillsInjectionHook(CustomLogger): - Skills with 'litellm:' prefix are fetched from LiteLLM DB - For Anthropic models: native skills pass through, LiteLLM skills converted to tools - For non-Anthropic models: LiteLLM skills are converted to tools + execute_code tool - + Post-call (async_post_call_success_deployment_hook): - If response has litellm_code_execution tool call, automatically execute code - Continue conversation loop until model gives final response - Return response with generated files inline - + This hook is called automatically by litellm during completion calls. """ @@ -60,7 +60,7 @@ class SkillsInjectionHook(CustomLogger): DEFAULT_MAX_ITERATIONS, DEFAULT_SANDBOX_TIMEOUT, ) - + self.optional_params = kwargs self.prompt_handler = SkillPromptInjectionHandler() self.max_iterations = kwargs.get("max_iterations", DEFAULT_MAX_ITERATIONS) @@ -95,7 +95,9 @@ class SkillsInjectionHook(CustomLogger): if not skills or not isinstance(skills, list): return data - verbose_proxy_logger.debug(f"SkillsInjectionHook: Processing {len(skills)} skills") + verbose_proxy_logger.debug( + f"SkillsInjectionHook: Processing {len(skills)} skills" + ) litellm_skills: List[LiteLLM_SkillsTable] = [] anthropic_skills: List[Dict[str, Any]] = [] @@ -132,7 +134,6 @@ class SkillsInjectionHook(CustomLogger): return data - def _process_for_messages_api( self, data: dict, @@ -141,7 +142,7 @@ class SkillsInjectionHook(CustomLogger): ) -> dict: """ Process skills for messages API (Anthropic format tools). - + - Converts skills to Anthropic-style tools (name, description, input_schema) - Extracts and injects SKILL.md content into system prompt - Adds litellm_code_execution tool for code execution @@ -150,7 +151,7 @@ class SkillsInjectionHook(CustomLogger): from litellm.llms.litellm_proxy.skills.code_execution import ( get_litellm_code_execution_tool_anthropic, ) - + tools = data.get("tools", []) skill_contents: List[str] = [] all_skill_files: Dict[str, Dict[str, bytes]] = {} @@ -159,12 +160,12 @@ class SkillsInjectionHook(CustomLogger): for skill in litellm_skills: # Convert skill to Anthropic-style tool tools.append(self.prompt_handler.convert_skill_to_anthropic_tool(skill)) - + # Extract skill content from file if available content = self.prompt_handler.extract_skill_content(skill) if content: skill_contents.append(content) - + # Extract all files for code execution skill_files = self.prompt_handler.extract_all_files(skill) if skill_files: @@ -187,7 +188,7 @@ class SkillsInjectionHook(CustomLogger): if all_skill_files: code_exec_tool = get_litellm_code_execution_tool_anthropic() data["tools"] = data.get("tools", []) + [code_exec_tool] - + # Store skill files in litellm_metadata for automatic code execution data["litellm_metadata"] = data.get("litellm_metadata", {}) data["litellm_metadata"]["_skill_files"] = all_skill_files @@ -211,7 +212,7 @@ class SkillsInjectionHook(CustomLogger): ) -> dict: """ Process skills for non-Anthropic models (OpenAI format tools). - + - Converts skills to OpenAI-style tools - Extracts and injects SKILL.md content - Adds execute_code tool for code execution @@ -225,12 +226,12 @@ class SkillsInjectionHook(CustomLogger): for skill in litellm_skills: # Convert skill to OpenAI-style tool tools.append(self.prompt_handler.convert_skill_to_tool(skill)) - + # Extract skill content from file if available content = self.prompt_handler.extract_skill_content(skill) if content: skill_contents.append(content) - + # Extract all files for code execution skill_files = self.prompt_handler.extract_all_files(skill) if skill_files: @@ -245,15 +246,18 @@ class SkillsInjectionHook(CustomLogger): # Inject skill content into system prompt if skill_contents: - data = self.prompt_handler.inject_skill_content_to_messages(data, skill_contents) + data = self.prompt_handler.inject_skill_content_to_messages( + data, skill_contents + ) # Add litellm_code_execution tool if we have skill files if all_skill_files: from litellm.llms.litellm_proxy.skills.code_execution import ( get_litellm_code_execution_tool, ) + data["tools"] = data.get("tools", []) + [get_litellm_code_execution_tool()] - + # Store skill files in litellm_metadata for automatic code execution # Using litellm_metadata instead of metadata to avoid conflicts with user metadata data["litellm_metadata"] = data.get("litellm_metadata", {}) @@ -271,7 +275,9 @@ class SkillsInjectionHook(CustomLogger): return data - async def _fetch_skill_from_db(self, skill_id: str) -> Optional[LiteLLM_SkillsTable]: + async def _fetch_skill_from_db( + self, skill_id: str + ) -> Optional[LiteLLM_SkillsTable]: """ Fetch a skill from the LiteLLM database. @@ -320,10 +326,10 @@ class SkillsInjectionHook(CustomLogger): ) -> Optional[Any]: """ Post-call hook to handle automatic code execution. - - Handles both OpenAI format (response.choices) and Anthropic/messages API + + Handles both OpenAI format (response.choices) and Anthropic/messages API format (response["content"]). - + If the response contains a tool call (litellm_code_execution or skill tool): 1. Execute the code in sandbox 2. Add result to messages @@ -338,95 +344,107 @@ class SkillsInjectionHook(CustomLogger): # Check if code execution is enabled for this request litellm_metadata = request_data.get("litellm_metadata") or {} metadata = request_data.get("metadata") or {} - - code_exec_enabled = ( - litellm_metadata.get("_litellm_code_execution_enabled") or - metadata.get("_litellm_code_execution_enabled") - ) + + code_exec_enabled = litellm_metadata.get( + "_litellm_code_execution_enabled" + ) or metadata.get("_litellm_code_execution_enabled") if not code_exec_enabled: return None - + # Get skill files - skill_files_by_id = ( - litellm_metadata.get("_skill_files") or - metadata.get("_skill_files", {}) + skill_files_by_id = litellm_metadata.get("_skill_files") or metadata.get( + "_skill_files", {} ) all_skill_files: Dict[str, bytes] = {} for files_dict in skill_files_by_id.values(): all_skill_files.update(files_dict) - + if not all_skill_files: verbose_proxy_logger.warning( "SkillsInjectionHook: No skill files found, cannot execute code" ) return None - + # Check for tool calls - handle both Anthropic and OpenAI formats tool_calls = self._extract_tool_calls(response) if not tool_calls: return None - + # Check if any tool call needs execution (litellm_code_execution or skill tool) has_executable_tool = False for tc in tool_calls: tool_name = tc.get("name", "") # Execute if it's litellm_code_execution OR a skill tool (skill_xxx) - if tool_name == LiteLLMInternalTools.CODE_EXECUTION.value or tool_name.startswith("skill_"): + if ( + tool_name == LiteLLMInternalTools.CODE_EXECUTION.value + or tool_name.startswith("skill_") + ): has_executable_tool = True break - + if not has_executable_tool: return None - + verbose_proxy_logger.debug( "SkillsInjectionHook: Detected tool call, starting execution loop" ) - + # Start the agentic loop return await self._execute_code_loop_messages_api( data=request_data, response=response, skill_files=all_skill_files, ) - + def _extract_tool_calls(self, response: Any) -> List[Dict[str, Any]]: """Extract tool calls from response, handling both formats.""" tool_calls = [] - + # Get content - handle both dict and object responses content = None if isinstance(response, dict): content = response.get("content", []) elif hasattr(response, "content"): content = response.content - + # Anthropic/messages API format: response has "content" list with tool_use blocks if content: for block in content: if isinstance(block, dict) and block.get("type") == "tool_use": - tool_calls.append({ - "id": block.get("id"), - "name": block.get("name"), - "input": block.get("input", {}), - }) - elif hasattr(block, "type") and getattr(block, "type", None) == "tool_use": - tool_calls.append({ - "id": getattr(block, "id", None), - "name": getattr(block, "name", None), - "input": getattr(block, "input", {}), - }) - + tool_calls.append( + { + "id": block.get("id"), + "name": block.get("name"), + "input": block.get("input", {}), + } + ) + elif ( + hasattr(block, "type") + and getattr(block, "type", None) == "tool_use" + ): + tool_calls.append( + { + "id": getattr(block, "id", None), + "name": getattr(block, "name", None), + "input": getattr(block, "input", {}), + } + ) + # OpenAI format: response has choices[0].message.tool_calls if not tool_calls and hasattr(response, "choices") and response.choices: # type: ignore[union-attr] msg = response.choices[0].message # type: ignore[union-attr] if hasattr(msg, "tool_calls") and msg.tool_calls: for tc in msg.tool_calls: - tool_calls.append({ - "id": tc.id, - "name": tc.function.name, - "input": json.loads(tc.function.arguments) if tc.function.arguments else {}, - }) - + tool_calls.append( + { + "id": tc.id, + "name": tc.function.name, + "input": json.loads(tc.function.arguments) + if tc.function.arguments + else {}, + } + ) + return tool_calls async def _execute_code_loop_messages_api( @@ -437,7 +455,7 @@ class SkillsInjectionHook(CustomLogger): ) -> Any: """ Execute the code execution loop for messages API (Anthropic format). - + Returns the final response with generated files inline. """ import litellm @@ -454,23 +472,31 @@ class SkillsInjectionHook(CustomLogger): "SkillsInjectionHook: Response is None, cannot execute code loop" ) return None - + model = data.get("model", "") messages = list(data.get("messages", [])) tools = data.get("tools", []) max_tokens = data.get("max_tokens", 4096) - + executor = SkillsSandboxExecutor(timeout=self.sandbox_timeout) generated_files: List[Dict[str, Any]] = [] current_response = response - + for iteration in range(self.max_iterations): # Extract tool calls from current response tool_calls = self._extract_tool_calls(current_response) - stop_reason = current_response.get("stop_reason") if isinstance(current_response, dict) else getattr(current_response, "stop_reason", None) - + stop_reason = ( + current_response.get("stop_reason") + if isinstance(current_response, dict) + else getattr(current_response, "stop_reason", None) + ) + # Get content for assistant message - convert to plain dicts - raw_content = current_response.get("content", []) if isinstance(current_response, dict) else getattr(current_response, "content", []) + raw_content = ( + current_response.get("content", []) + if isinstance(current_response, dict) + else getattr(current_response, "content", []) + ) content_blocks = [] for block in raw_content or []: if isinstance(block, dict): @@ -481,11 +507,11 @@ class SkillsInjectionHook(CustomLogger): content_blocks.append(dict(block.__dict__)) else: content_blocks.append({"type": "text", "text": str(block)}) - + # Build assistant message for conversation history (Anthropic format) assistant_msg = {"role": "assistant", "content": content_blocks} messages.append(assistant_msg) - + # Check if we're done (no tool calls) if stop_reason != "tool_use" or not tool_calls: verbose_proxy_logger.debug( @@ -493,33 +519,39 @@ class SkillsInjectionHook(CustomLogger): f"{len(generated_files)} files generated" ) return self._attach_files_to_response(current_response, generated_files) - + # Process tool calls tool_results = [] for tc in tool_calls: tool_name = tc.get("name", "") tool_id = tc.get("id", "") tool_input = tc.get("input", {}) - + # Execute if it's litellm_code_execution OR a skill tool if tool_name == LiteLLMInternalTools.CODE_EXECUTION.value: code = tool_input.get("code", "") - result = await self._execute_code(code, skill_files, executor, generated_files) + result = await self._execute_code( + code, skill_files, executor, generated_files + ) elif tool_name.startswith("skill_"): # Skill tool - execute the skill's code - result = await self._execute_skill_tool(tool_name, tool_input, skill_files, executor, generated_files) + result = await self._execute_skill_tool( + tool_name, tool_input, skill_files, executor, generated_files + ) else: result = f"Tool '{tool_name}' not handled" - - tool_results.append({ - "type": "tool_result", - "tool_use_id": tool_id, - "content": result, - }) - + + tool_results.append( + { + "type": "tool_result", + "tool_use_id": tool_id, + "content": result, + } + ) + # Add tool results to messages (Anthropic format) messages.append({"role": "user", "content": tool_results}) - + # Make next LLM call verbose_proxy_logger.debug( f"SkillsInjectionHook: Making LLM call iteration {iteration + 2}" @@ -537,11 +569,9 @@ class SkillsInjectionHook(CustomLogger): ) return self._attach_files_to_response(response, generated_files) except Exception as e: - verbose_proxy_logger.error( - f"SkillsInjectionHook: LLM call failed: {e}" - ) + verbose_proxy_logger.error(f"SkillsInjectionHook: LLM call failed: {e}") return self._attach_files_to_response(response, generated_files) - + verbose_proxy_logger.warning( f"SkillsInjectionHook: Max iterations ({self.max_iterations}) reached" ) @@ -556,26 +586,30 @@ class SkillsInjectionHook(CustomLogger): ) -> str: """Execute code in sandbox and return result string.""" try: - verbose_proxy_logger.debug(f"SkillsInjectionHook: Executing code ({len(code)} chars)") - + verbose_proxy_logger.debug( + f"SkillsInjectionHook: Executing code ({len(code)} chars)" + ) + exec_result = executor.execute(code=code, skill_files=skill_files) - + result = exec_result.get("output", "") or "" - + # Collect generated files if exec_result.get("files"): for f in exec_result["files"]: - generated_files.append({ - "name": f["name"], - "mime_type": f["mime_type"], - "content_base64": f["content_base64"], - "size": len(base64.b64decode(f["content_base64"])), - }) + generated_files.append( + { + "name": f["name"], + "mime_type": f["mime_type"], + "content_base64": f["content_base64"], + "size": len(base64.b64decode(f["content_base64"])), + } + ) result += f"\n\nGenerated file: {f['name']}" - + if exec_result.get("error"): result += f"\n\nError: {exec_result['error']}" - + return result or "Code executed successfully" except Exception as e: return f"Code execution failed: {str(e)}" @@ -591,23 +625,31 @@ class SkillsInjectionHook(CustomLogger): """Execute a skill tool by generating and running code based on skill content.""" # Generate code based on available skill modules # Look for Python modules in the skill - python_modules = [p for p in skill_files.keys() if p.endswith(".py") and not p.endswith("__init__.py")] - + python_modules = [ + p + for p in skill_files.keys() + if p.endswith(".py") and not p.endswith("__init__.py") + ] + # Try to find the main builder/creator module main_module = None for mod in python_modules: - if "builder" in mod.lower() or "creator" in mod.lower() or "generator" in mod.lower(): + if ( + "builder" in mod.lower() + or "creator" in mod.lower() + or "generator" in mod.lower() + ): main_module = mod break - + if not main_module and python_modules: # Use first non-init module main_module = python_modules[0] - + if main_module: # Convert path to import: "core/gif_builder.py" -> "core.gif_builder" import_path = main_module.replace("/", ".").replace(".py", "") - + # Generate code that imports and uses the module code = f""" # Auto-generated code to execute skill @@ -650,7 +692,7 @@ for f in os.listdir('.'): code = """ print('No executable skill module found') """ - + return await self._execute_code(code, skill_files, executor, generated_files) async def _execute_code_loop( @@ -661,7 +703,7 @@ print('No executable skill module found') ) -> Any: """ Execute the code execution loop until model gives final response. - + Returns the final response with generated files inline. """ import litellm @@ -671,36 +713,35 @@ print('No executable skill module found') from litellm.llms.litellm_proxy.skills.sandbox_executor import ( SkillsSandboxExecutor, ) - + model = data.get("model", "") messages = list(data.get("messages", [])) tools = data.get("tools", []) - + # Keys to exclude when passing through to acompletion # These are either handled explicitly or are internal LiteLLM fields - _EXCLUDED_ACOMPLETION_KEYS = frozenset({ - "messages", - "model", - "tools", - "metadata", - "litellm_metadata", - "container", - }) - - kwargs = { - k: v for k, v in data.items() - if k not in _EXCLUDED_ACOMPLETION_KEYS - } - + _EXCLUDED_ACOMPLETION_KEYS = frozenset( + { + "messages", + "model", + "tools", + "metadata", + "litellm_metadata", + "container", + } + ) + + kwargs = {k: v for k, v in data.items() if k not in _EXCLUDED_ACOMPLETION_KEYS} + executor = SkillsSandboxExecutor(timeout=self.sandbox_timeout) generated_files: List[Dict[str, Any]] = [] current_response: Any = response - + for iteration in range(self.max_iterations): # OpenAI format response has choices[0].message assistant_message = current_response.choices[0].message # type: ignore[union-attr] stop_reason = current_response.choices[0].finish_reason # type: ignore[union-attr] - + # Build assistant message for conversation history assistant_msg_dict: Dict[str, Any] = { "role": "assistant", @@ -713,13 +754,13 @@ print('No executable skill module found') "type": "function", "function": { "name": tc.function.name, - "arguments": tc.function.arguments - } + "arguments": tc.function.arguments, + }, } for tc in assistant_message.tool_calls ] messages.append(assistant_msg_dict) - + # Check if we're done (no tool calls) if stop_reason != "tool_calls" or not assistant_message.tool_calls: verbose_proxy_logger.debug( @@ -728,11 +769,11 @@ print('No executable skill module found') ) # Attach generated files to response return self._attach_files_to_response(current_response, generated_files) - + # Process tool calls for tool_call in assistant_message.tool_calls: tool_name = tool_call.function.name - + if tool_name == LiteLLMInternalTools.CODE_EXECUTION.value: tool_result = await self._execute_code_tool( tool_call=tool_call, @@ -743,13 +784,15 @@ print('No executable skill module found') else: # Non-code-execution tool - cannot handle tool_result = f"Tool '{tool_name}' not handled automatically" - - messages.append({ - "role": "tool", - "tool_call_id": tool_call.id, - "content": tool_result, - }) - + + messages.append( + { + "role": "tool", + "tool_call_id": tool_call.id, + "content": tool_result, + } + ) + # Make next LLM call using the messages API verbose_proxy_logger.debug( f"SkillsInjectionHook: Making LLM call iteration {iteration + 2}" @@ -760,13 +803,13 @@ print('No executable skill module found') tools=tools, max_tokens=kwargs.get("max_tokens", 4096), ) - + # Max iterations reached verbose_proxy_logger.warning( f"SkillsInjectionHook: Max iterations ({self.max_iterations}) reached" ) return self._attach_files_to_response(current_response, generated_files) - + async def _execute_code_tool( self, tool_call: Any, @@ -778,48 +821,50 @@ print('No executable skill module found') try: args = json.loads(tool_call.function.arguments) code = args.get("code", "") - + verbose_proxy_logger.debug( f"SkillsInjectionHook: Executing code ({len(code)} chars)" ) - + exec_result = executor.execute( code=code, skill_files=skill_files, ) - + # Build tool result content tool_result = exec_result.get("output", "") or "" - + # Collect generated files if exec_result.get("files"): tool_result += "\n\nGenerated files:" for f in exec_result["files"]: file_content = base64.b64decode(f["content_base64"]) - generated_files.append({ - "name": f["name"], - "mime_type": f["mime_type"], - "content_base64": f["content_base64"], - "size": len(file_content), - }) + generated_files.append( + { + "name": f["name"], + "mime_type": f["mime_type"], + "content_base64": f["content_base64"], + "size": len(file_content), + } + ) tool_result += f"\n- {f['name']} ({len(file_content)} bytes)" - + verbose_proxy_logger.debug( f"SkillsInjectionHook: Generated file {f['name']} " f"({len(file_content)} bytes)" ) - + if exec_result.get("error"): tool_result += f"\n\nError:\n{exec_result['error']}" - + return tool_result - + except Exception as e: verbose_proxy_logger.error( f"SkillsInjectionHook: Code execution failed: {e}" ) return f"Code execution failed: {str(e)}" - + def _attach_files_to_response( self, response: Any, @@ -827,13 +872,13 @@ print('No executable skill module found') ) -> Any: """ Attach generated files to the response object. - + Files are added to response._litellm_generated_files for easy access. For dict responses, files are added as a key. """ if not generated_files: return response - + # Handle dict response (Anthropic/messages API format) if isinstance(response, dict): response["_litellm_generated_files"] = generated_files @@ -841,23 +886,23 @@ print('No executable skill module found') f"SkillsInjectionHook: Attached {len(generated_files)} files to dict response" ) return response - + # Handle object response (OpenAI format) try: response._litellm_generated_files = generated_files except AttributeError: pass - + # Also add to model_extra if available (for serialization) if hasattr(response, "model_extra"): if response.model_extra is None: response.model_extra = {} response.model_extra["_litellm_generated_files"] = generated_files - + verbose_proxy_logger.debug( f"SkillsInjectionHook: Attached {len(generated_files)} files to response" ) - + return response diff --git a/litellm/proxy/hooks/max_budget_per_session_limiter.py b/litellm/proxy/hooks/max_budget_per_session_limiter.py index a981207f00..59fb101f55 100644 --- a/litellm/proxy/hooks/max_budget_per_session_limiter.py +++ b/litellm/proxy/hooks/max_budget_per_session_limiter.py @@ -208,9 +208,7 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): async def _get_current_spend(self, cache_key: str) -> float: """Read current accumulated spend for a session.""" - if ( - self.internal_usage_cache.dual_cache.redis_cache is not None - ): + if self.internal_usage_cache.dual_cache.redis_cache is not None: try: result = await self.internal_usage_cache.dual_cache.redis_cache.async_get_cache( key=cache_key @@ -252,9 +250,7 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): return await self._in_memory_increment_spend(cache_key, amount) - async def _in_memory_increment_spend( - self, cache_key: str, amount: float - ) -> float: + async def _in_memory_increment_spend(self, cache_key: str, amount: float) -> float: current = await self.internal_usage_cache.async_get_cache( key=cache_key, litellm_parent_otel_span=None, diff --git a/litellm/proxy/hooks/max_iterations_limiter.py b/litellm/proxy/hooks/max_iterations_limiter.py index b6fde2b178..df9a298ca0 100644 --- a/litellm/proxy/hooks/max_iterations_limiter.py +++ b/litellm/proxy/hooks/max_iterations_limiter.py @@ -148,9 +148,7 @@ class _PROXY_MaxIterationsHandler(CustomLogger): return None - def _get_max_iterations( - self, user_api_key_dict: UserAPIKeyAuth - ) -> Optional[int]: + def _get_max_iterations(self, user_api_key_dict: UserAPIKeyAuth) -> Optional[int]: """Extract max_iterations from agent litellm_params, with fallback to key metadata.""" # Try agent litellm_params first agent_id = user_api_key_dict.agent_id diff --git a/litellm/proxy/hooks/mcp_semantic_filter/hook.py b/litellm/proxy/hooks/mcp_semantic_filter/hook.py index fc9349c2a4..725d085edb 100644 --- a/litellm/proxy/hooks/mcp_semantic_filter/hook.py +++ b/litellm/proxy/hooks/mcp_semantic_filter/hook.py @@ -26,40 +26,40 @@ if TYPE_CHECKING: class SemanticToolFilterHook(CustomLogger): """ Pre-call hook that filters MCP tools semantically. - + This hook: 1. Extracts the user query from messages 2. Filters tools based on semantic similarity to the query 3. Returns only the top-k most relevant tools to the LLM """ - + def __init__(self, semantic_filter: "SemanticMCPToolFilter"): """ Initialize the hook. - + Args: semantic_filter: SemanticMCPToolFilter instance """ super().__init__() self.filter = semantic_filter - + verbose_proxy_logger.debug( f"Initialized SemanticToolFilterHook with filter: " f"enabled={semantic_filter.enabled}, top_k={semantic_filter.top_k}" ) - + def _should_expand_mcp_tools(self, tools: List[Any]) -> bool: """ Check if tools contain MCP references with server_url="litellm_proxy". - + Only expands MCP tools pointing to litellm proxy, not external MCP servers. """ from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) - + return LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools) - + async def _expand_mcp_tools( self, tools: List[Any], @@ -67,7 +67,7 @@ class SemanticToolFilterHook(CustomLogger): ) -> List[Dict[str, Any]]: """ Expand MCP references to actual tool definitions. - + Reuses LiteLLM_Proxy_MCP_Handler._process_mcp_tools_to_openai_format which internally does: parse -> fetch -> filter -> deduplicate -> transform """ @@ -77,46 +77,56 @@ class SemanticToolFilterHook(CustomLogger): # Parse to separate MCP tools from other tools mcp_tools, _ = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools) - + if not mcp_tools: return [] - + # Use single combined method instead of 3 separate calls # This already handles: fetch -> filter by allowed_tools -> deduplicate -> transform - openai_tools, _ = await LiteLLM_Proxy_MCP_Handler._process_mcp_tools_to_openai_format( - user_api_key_auth=user_api_key_dict, - mcp_tools_with_litellm_proxy=mcp_tools + ( + openai_tools, + _, + ) = await LiteLLM_Proxy_MCP_Handler._process_mcp_tools_to_openai_format( + user_api_key_auth=user_api_key_dict, mcp_tools_with_litellm_proxy=mcp_tools ) - + # Convert Pydantic models to dicts for compatibility openai_tools_as_dicts = [] for tool in openai_tools: if hasattr(tool, "model_dump"): tool_dict = tool.model_dump(exclude_none=True) - verbose_proxy_logger.debug(f"Converted Pydantic tool to dict: {type(tool).__name__} -> dict with keys: {list(tool_dict.keys())}") + verbose_proxy_logger.debug( + f"Converted Pydantic tool to dict: {type(tool).__name__} -> dict with keys: {list(tool_dict.keys())}" + ) openai_tools_as_dicts.append(tool_dict) elif hasattr(tool, "dict"): tool_dict = tool.dict(exclude_none=True) - verbose_proxy_logger.debug(f"Converted Pydantic tool (v1) to dict: {type(tool).__name__} -> dict") + verbose_proxy_logger.debug( + f"Converted Pydantic tool (v1) to dict: {type(tool).__name__} -> dict" + ) openai_tools_as_dicts.append(tool_dict) elif isinstance(tool, dict): - verbose_proxy_logger.debug(f"Tool is already a dict with keys: {list(tool.keys())}") + verbose_proxy_logger.debug( + f"Tool is already a dict with keys: {list(tool.keys())}" + ) openai_tools_as_dicts.append(tool) else: - verbose_proxy_logger.warning(f"Tool is unknown type: {type(tool)}, passing as-is") + verbose_proxy_logger.warning( + f"Tool is unknown type: {type(tool)}, passing as-is" + ) openai_tools_as_dicts.append(tool) - + verbose_proxy_logger.debug( f"Expanded {len(mcp_tools)} MCP reference(s) to {len(openai_tools_as_dicts)} tools (all as dicts)" ) - + return openai_tools_as_dicts - + def _get_metadata_variable_name(self, data: dict) -> str: if "litellm_metadata" in data: return "litellm_metadata" return "metadata" - + async def async_pre_call_hook( self, user_api_key_dict: "UserAPIKeyAuth", @@ -126,16 +136,16 @@ class SemanticToolFilterHook(CustomLogger): ) -> Optional[Union[Exception, str, dict]]: """ Filter tools before LLM call based on user query. - + This hook is called before the LLM request is made. It filters the tools list to only include semantically relevant tools. - + Args: user_api_key_dict: User authentication cache: Cache instance data: Request data containing messages and tools call_type: Type of call (completion, acompletion, etc.) - + Returns: Modified data dict with filtered tools, or None if no changes """ @@ -145,100 +155,104 @@ class SemanticToolFilterHook(CustomLogger): f"Skipping semantic filter for call_type={call_type}" ) return None - + # Check if tools are present tools = data.get("tools") if not tools: verbose_proxy_logger.debug("No tools in request, skipping semantic filter") return None - + original_tool_count = len(tools) - + # Check for MCP references (server_url="litellm_proxy") and expand them if self._should_expand_mcp_tools(tools): verbose_proxy_logger.debug( "Detected litellm_proxy MCP references, expanding before semantic filtering" ) - + try: - expanded_tools = await self._expand_mcp_tools( - tools, user_api_key_dict - ) - + expanded_tools = await self._expand_mcp_tools(tools, user_api_key_dict) + if not expanded_tools: verbose_proxy_logger.warning( "No tools expanded from MCP references" ) return None - + verbose_proxy_logger.info( f"Expanded {len(tools)} MCP reference(s) to {len(expanded_tools)} tools" ) - + # Update tools for filtering tools = expanded_tools original_tool_count = len(tools) - + except Exception as e: verbose_proxy_logger.error( f"Failed to expand MCP references: {e}", exc_info=True ) return None - + # Check if messages are present (try both "messages" and "input" for responses API) messages = data.get("messages", []) if not messages: messages = data.get("input", []) if not messages: - verbose_proxy_logger.debug("No messages in request, skipping semantic filter") + verbose_proxy_logger.debug( + "No messages in request, skipping semantic filter" + ) return None - + # Check if filter is enabled if not self.filter.enabled: verbose_proxy_logger.debug("Semantic filter disabled, skipping") return None - + try: # Extract user query from messages user_query = self.filter.extract_user_query(messages) if not user_query: - verbose_proxy_logger.debug("No user query found, skipping semantic filter") + verbose_proxy_logger.debug( + "No user query found, skipping semantic filter" + ) return None - + verbose_proxy_logger.debug( f"Applying semantic filter to {len(tools)} tools " f"with query: '{user_query[:50]}...'" ) - + # Filter tools semantically filtered_tools = await self.filter.filter_tools( query=user_query, available_tools=tools, # type: ignore ) - + # Always update tools and emit header (even if count unchanged) data["tools"] = filtered_tools - + # Store filter stats and tool names for response header filter_stats = f"{original_tool_count}->{len(filtered_tools)}" tool_names_csv = self._get_tool_names_csv(filtered_tools) - + _metadata_variable_name = self._get_metadata_variable_name(data) - data[_metadata_variable_name]["litellm_semantic_filter_stats"] = filter_stats - data[_metadata_variable_name]["litellm_semantic_filter_tools"] = tool_names_csv - - verbose_proxy_logger.info( - f"Semantic tool filter: {filter_stats} tools" - ) - + data[_metadata_variable_name][ + "litellm_semantic_filter_stats" + ] = filter_stats + data[_metadata_variable_name][ + "litellm_semantic_filter_tools" + ] = tool_names_csv + + verbose_proxy_logger.info(f"Semantic tool filter: {filter_stats} tools") + return data - + except Exception as e: verbose_proxy_logger.warning( f"Semantic tool filter hook failed: {e}. Proceeding with all tools." ) return None - + async def async_post_call_response_headers_hook( self, data: dict, @@ -248,39 +262,46 @@ class SemanticToolFilterHook(CustomLogger): ) -> Optional[Dict[str, str]]: """Add semantic filter stats and tool names to response headers.""" from litellm.constants import MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH - + _metadata_variable_name = self._get_metadata_variable_name(data) metadata = data[_metadata_variable_name] - + filter_stats = metadata.get("litellm_semantic_filter_stats") if not filter_stats: return None - + headers = {"x-litellm-semantic-filter": filter_stats} - + # Add CSV of filtered tool names (nginx-safe length) tool_names_csv = metadata.get("litellm_semantic_filter_tools", "") if tool_names_csv: if len(tool_names_csv) > MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH: - tool_names_csv = tool_names_csv[:MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH - 3] + "..." - + tool_names_csv = ( + tool_names_csv[: MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH - 3] + + "..." + ) + headers["x-litellm-semantic-filter-tools"] = tool_names_csv - + return headers - + def _get_tool_names_csv(self, tools: List[Any]) -> str: """Extract tool names and return as CSV string.""" if not tools: return "" - + tool_names = [] for tool in tools: - name = tool.get("name", "") if isinstance(tool, dict) else getattr(tool, "name", "") + name = ( + tool.get("name", "") + if isinstance(tool, dict) + else getattr(tool, "name", "") + ) if name: tool_names.append(name) - + return ",".join(tool_names) - + @staticmethod async def initialize_from_config( config: Optional[Dict[str, Any]], @@ -288,29 +309,29 @@ class SemanticToolFilterHook(CustomLogger): ) -> Optional["SemanticToolFilterHook"]: """ Initialize semantic tool filter from proxy config. - + Args: config: Proxy configuration dict (litellm_settings.mcp_semantic_tool_filter) llm_router: LiteLLM router instance for embeddings - + Returns: SemanticToolFilterHook instance if enabled, None otherwise """ from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( SemanticMCPToolFilter, ) + if not config or not config.get("enabled", False): verbose_proxy_logger.debug("Semantic tool filter not enabled in config") return None - + if llm_router is None: verbose_proxy_logger.warning( "Cannot initialize semantic filter: llm_router is None" ) return None - + try: - embedding_model = config.get( "embedding_model", DEFAULT_MCP_SEMANTIC_FILTER_EMBEDDING_MODEL ) @@ -318,7 +339,7 @@ class SemanticToolFilterHook(CustomLogger): similarity_threshold = config.get( "similarity_threshold", DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD ) - + semantic_filter = SemanticMCPToolFilter( embedding_model=embedding_model, litellm_router_instance=llm_router, @@ -326,20 +347,20 @@ class SemanticToolFilterHook(CustomLogger): similarity_threshold=similarity_threshold, enabled=True, ) - + # Build router from MCP registry on startup await semantic_filter.build_router_from_mcp_registry() - + hook = SemanticToolFilterHook(semantic_filter) - + verbose_proxy_logger.info( f"✅ MCP Semantic Tool Filter enabled: " f"embedding_model={embedding_model}, top_k={top_k}, " f"similarity_threshold={similarity_threshold}" ) - + return hook - + except ImportError as e: verbose_proxy_logger.warning( f"semantic-router not installed. Install with: " diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 856975ea09..19c8c484b4 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -7,8 +7,18 @@ This is currently in development and not yet ready for production. import binascii import os from datetime import datetime -from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Literal, - Optional, TypedDict, Union, cast) +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + List, + Literal, + Optional, + TypedDict, + Union, + cast, +) from fastapi import HTTPException @@ -165,8 +175,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """Get or lazy-load the batch rate limiter.""" if self._batch_rate_limiter is None: try: - from litellm.proxy.hooks.batch_rate_limiter import \ - _PROXY_BatchRateLimiter + from litellm.proxy.hooks.batch_rate_limiter import ( + _PROXY_BatchRateLimiter, + ) self._batch_rate_limiter = _PROXY_BatchRateLimiter( internal_usage_cache=self.internal_usage_cache, @@ -668,8 +679,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): requested_model: The model being requested descriptors: List of rate limit descriptors to append to """ - from litellm.proxy.auth.auth_utils import (get_key_model_rpm_limit, - get_key_model_tpm_limit) + from litellm.proxy.auth.auth_utils import ( + get_key_model_rpm_limit, + get_key_model_tpm_limit, + ) if not requested_model: return @@ -780,8 +793,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def _get_agent_from_registry(self, agent_id: str) -> Optional[Any]: """Look up an agent from the in-memory registry by ID.""" - from litellm.proxy.agent_endpoints.agent_registry import \ - global_agent_registry + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry return global_agent_registry.get_agent_by_id(agent_id=agent_id) @@ -878,8 +890,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): Returns list of descriptors for API key, user, team, team member, end user, model-specific, agent, and agent-session limits. """ - from litellm.proxy.auth.auth_utils import (get_team_model_rpm_limit, - get_team_model_tpm_limit) + from litellm.proxy.auth.auth_utils import ( + get_team_model_rpm_limit, + get_team_model_tpm_limit, + ) descriptors = [] @@ -1053,8 +1067,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): Returns True if any deployment has failures in the current minute. """ from litellm.proxy.proxy_server import llm_router - from litellm.router_utils.router_callbacks.track_deployment_metrics import \ - get_deployment_failures_for_current_minute + from litellm.router_utils.router_callbacks.track_deployment_metrics import ( + get_deployment_failures_for_current_minute, + ) if llm_router is None: return False @@ -1468,10 +1483,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ Update TPM usage on successful API calls by incrementing counters using pipeline """ - from litellm.litellm_core_utils.core_helpers import \ - _get_parent_otel_span_from_kwargs - from litellm.proxy.common_utils.callback_utils import \ - get_model_group_from_litellm_kwargs + from litellm.litellm_core_utils.core_helpers import ( + _get_parent_otel_span_from_kwargs, + ) + from litellm.proxy.common_utils.callback_utils import ( + get_model_group_from_litellm_kwargs, + ) from litellm.types.caching import RedisPipelineIncrementOperation rate_limit_type = self.get_rate_limit_type() @@ -1655,14 +1672,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ Decrement max parallel requests counter for the API Key """ - from litellm.litellm_core_utils.core_helpers import \ - _get_parent_otel_span_from_kwargs + from litellm.litellm_core_utils.core_helpers import ( + _get_parent_otel_span_from_kwargs, + ) from litellm.types.caching import RedisPipelineIncrementOperation try: - litellm_parent_otel_span: Union[Span, None] = ( - _get_parent_otel_span_from_kwargs(kwargs) - ) + litellm_parent_otel_span: Union[ + Span, None + ] = _get_parent_otel_span_from_kwargs(kwargs) # Get metadata from standard_logging_object - this correctly handles both # 'metadata' and 'litellm_metadata' fields from litellm_params standard_logging_object = kwargs.get("standard_logging_object") or {} diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 8abec22e60..49ff94d0a9 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -12,7 +12,11 @@ from litellm.litellm_core_utils.core_helpers import ( ) from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.auth.auth_checks import get_key_object, get_team_object, log_db_metrics +from litellm.proxy.auth.auth_checks import ( + get_key_object, + get_team_object, + log_db_metrics, +) from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.utils import ProxyUpdateSpend from litellm.types.utils import ( @@ -39,8 +43,8 @@ class _ProxyDBLogger(CustomLogger): if _ProxyDBLogger._should_track_errors_in_db() is False: return elif request_route is not None and not ( - RouteChecks.is_llm_api_route(route=request_route) or - RouteChecks.is_info_route(route=request_route) + RouteChecks.is_llm_api_route(route=request_route) + or RouteChecks.is_info_route(route=request_route) ): return diff --git a/litellm/proxy/hooks/rate_limiter_utils.py b/litellm/proxy/hooks/rate_limiter_utils.py index c6b0ef65bc..927bac0de5 100644 --- a/litellm/proxy/hooks/rate_limiter_utils.py +++ b/litellm/proxy/hooks/rate_limiter_utils.py @@ -36,19 +36,9 @@ def convert_priority_to_percent( if val_type == "percent": return float(val_num) - elif ( - val_type == "rpm" - and model_info - and model_info.rpm - and model_info.rpm > 0 - ): + elif val_type == "rpm" and model_info and model_info.rpm and model_info.rpm > 0: return float(val_num) / model_info.rpm - elif ( - val_type == "tpm" - and model_info - and model_info.tpm - and model_info.tpm > 0 - ): + elif val_type == "tpm" and model_info and model_info.tpm and model_info.tpm > 0: return float(val_num) / model_info.tpm # Fallback: treat as percent diff --git a/litellm/proxy/hooks/responses_id_security.py b/litellm/proxy/hooks/responses_id_security.py index 39f33ade38..3a23347f35 100644 --- a/litellm/proxy/hooks/responses_id_security.py +++ b/litellm/proxy/hooks/responses_id_security.py @@ -269,7 +269,9 @@ class ResponsesIDSecurity(CustomLogger): if isinstance(response, ResponsesAPIResponse): response = cast( ResponsesAPIResponse, - self._encrypt_response_id(response, user_api_key_dict, request_cache=None), + self._encrypt_response_id( + response, user_api_key_dict, request_cache=None + ), ) return response @@ -288,5 +290,7 @@ class ResponsesIDSecurity(CustomLogger): == "/v1/responses" # only encrypt the response id for the responses api and not general_settings.get("disable_responses_id_security", False) ): - chunk = self._encrypt_response_id(chunk, user_api_key_dict, request_encryption_cache) + chunk = self._encrypt_response_id( + chunk, user_api_key_dict, request_encryption_cache + ) yield chunk diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index 4a8eb8e741..4f994b87f5 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -242,9 +242,13 @@ async def image_edit_api( ``` """ if image is not None and image_array is not None: - raise HTTPException(status_code=422, detail="Cannot specify both 'image' and 'image[]'") + raise HTTPException( + status_code=422, detail="Cannot specify both 'image' and 'image[]'" + ) if mask is not None and mask_array is not None: - raise HTTPException(status_code=422, detail="Cannot specify both 'mask' and 'mask[]'") + raise HTTPException( + status_code=422, detail="Cannot specify both 'mask' and 'mask[]'" + ) if image is None and image_array is not None: image = image_array if mask is None and mask_array is not None: @@ -280,7 +284,7 @@ async def image_edit_api( data["image"] = image_files if mask_files: data["mask"] = mask_files - + # Ensure prompt exists in data (default to None for models that don't require it) if "prompt" not in data: data["prompt"] = None diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index cf4729db94..daf2867699 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -196,12 +196,12 @@ def _get_dynamic_logging_metadata( user_api_key_dict: UserAPIKeyAuth, proxy_config: ProxyConfig ) -> Optional[TeamCallbackMetadata]: callback_settings_obj: Optional[TeamCallbackMetadata] = None - key_dynamic_logging_settings: Optional[dict] = ( - KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict) - ) - team_dynamic_logging_settings: Optional[dict] = ( - KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(user_api_key_dict) - ) + key_dynamic_logging_settings: Optional[ + dict + ] = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict) + team_dynamic_logging_settings: Optional[ + dict + ] = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(user_api_key_dict) ######################################################################################### # Key-based callbacks ######################################################################################### @@ -262,14 +262,14 @@ def clean_headers( ) -> dict: """ Removes litellm api key from headers - + Args: headers: Request headers litellm_key_header_name: Custom header name for LiteLLM API key forward_llm_provider_auth_headers: Whether to forward provider auth headers authenticated_with_header: Which header was used for LiteLLM authentication (e.g., "x-litellm-api-key", "authorization", "x-api-key") - + Returns: Cleaned headers dict """ @@ -283,14 +283,17 @@ def clean_headers( header_lower = header.lower() if header_lower == "authorization" and is_anthropic_oauth_key(value): - if authenticated_with_header is None or authenticated_with_header.lower() != "authorization": + if ( + authenticated_with_header is None + or authenticated_with_header.lower() != "authorization" + ): clean_headers[header] = value continue # Special handling for x-api-key: forward it based on authenticated_with_header elif header_lower == "x-api-key": - if ( - forward_llm_provider_auth_headers - and (authenticated_with_header is None or authenticated_with_header.lower() != "x-api-key") + if forward_llm_provider_auth_headers and ( + authenticated_with_header is None + or authenticated_with_header.lower() != "x-api-key" ): clean_headers[header] = value elif ( @@ -625,7 +628,6 @@ class LiteLLMProxyRequestSetup: "x-litellm-session-id" ) - if agent_id_from_header: metadata_from_headers["agent_id"] = agent_id_from_header verbose_proxy_logger.debug( @@ -751,11 +753,11 @@ class LiteLLMProxyRequestSetup: ## KEY-LEVEL SPEND LOGS / TAGS if "tags" in key_metadata and key_metadata["tags"] is not None: - data[_metadata_variable_name]["tags"] = ( - LiteLLMProxyRequestSetup._merge_tags( - request_tags=data[_metadata_variable_name].get("tags"), - tags_to_add=key_metadata["tags"], - ) + data[_metadata_variable_name][ + "tags" + ] = LiteLLMProxyRequestSetup._merge_tags( + request_tags=data[_metadata_variable_name].get("tags"), + tags_to_add=key_metadata["tags"], ) if "disable_global_guardrails" in key_metadata and isinstance( key_metadata["disable_global_guardrails"], bool @@ -925,7 +927,6 @@ async def add_litellm_data_to_request( # noqa: PLR0915 verbose_proxy_logger.debug(f"Request Headers: {_headers}") verbose_proxy_logger.debug(f"Raw Headers: {_raw_headers}") - if forward_llm_auth and "x-api-key" in _headers: data["api_key"] = _headers["x-api-key"] verbose_proxy_logger.debug( @@ -1052,9 +1053,9 @@ async def add_litellm_data_to_request( # noqa: PLR0915 data[_metadata_variable_name]["litellm_api_version"] = version if general_settings is not None: - data[_metadata_variable_name]["global_max_parallel_requests"] = ( - general_settings.get("global_max_parallel_requests", None) - ) + data[_metadata_variable_name][ + "global_max_parallel_requests" + ] = general_settings.get("global_max_parallel_requests", None) ### KEY-LEVEL Controls key_metadata = user_api_key_dict.metadata @@ -1142,14 +1143,14 @@ async def add_litellm_data_to_request( # noqa: PLR0915 ] = user_api_key_dict.user_max_budget data[_metadata_variable_name]["user_api_key_metadata"] = user_api_key_dict.metadata - data[_metadata_variable_name]["user_api_key_team_metadata"] = ( - user_api_key_dict.team_metadata + data[_metadata_variable_name][ + "user_api_key_team_metadata" + ] = user_api_key_dict.team_metadata + data[_metadata_variable_name]["user_api_key_object_permission_id"] = getattr( + user_api_key_dict, "object_permission_id", None ) - data[_metadata_variable_name]["user_api_key_object_permission_id"] = ( - getattr(user_api_key_dict, "object_permission_id", None) - ) - data[_metadata_variable_name]["user_api_key_team_object_permission_id"] = ( - getattr(user_api_key_dict, "team_object_permission_id", None) + data[_metadata_variable_name]["user_api_key_team_object_permission_id"] = getattr( + user_api_key_dict, "team_object_permission_id", None ) data[_metadata_variable_name]["headers"] = _headers data[_metadata_variable_name]["endpoint"] = str(request.url) diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index d58dca5aec..caaec12f7a 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -108,7 +108,10 @@ async def _sync_add_access_group_to_teams( if team is not None and access_group_id not in (team.access_group_ids or []): await tx.litellm_teamtable.update( where={"team_id": team_id}, - data={"access_group_ids": list(team.access_group_ids or []) + [access_group_id]}, + data={ + "access_group_ids": list(team.access_group_ids or []) + + [access_group_id] + }, ) @@ -121,7 +124,11 @@ async def _sync_remove_access_group_from_teams( if team is not None and access_group_id in (team.access_group_ids or []): await tx.litellm_teamtable.update( where={"team_id": team_id}, - data={"access_group_ids": [ag for ag in team.access_group_ids if ag != access_group_id]}, + data={ + "access_group_ids": [ + ag for ag in team.access_group_ids if ag != access_group_id + ] + }, ) @@ -134,7 +141,10 @@ async def _sync_add_access_group_to_keys( if key is not None and access_group_id not in (key.access_group_ids or []): await tx.litellm_verificationtoken.update( where={"token": token}, - data={"access_group_ids": list(key.access_group_ids or []) + [access_group_id]}, + data={ + "access_group_ids": list(key.access_group_ids or []) + + [access_group_id] + }, ) @@ -147,7 +157,11 @@ async def _sync_remove_access_group_from_keys( if key is not None and access_group_id in (key.access_group_ids or []): await tx.litellm_verificationtoken.update( where={"token": token}, - data={"access_group_ids": [ag for ag in key.access_group_ids if ag != access_group_id]}, + data={ + "access_group_ids": [ + ag for ag in key.access_group_ids if ag != access_group_id + ] + }, ) @@ -175,7 +189,9 @@ async def _patch_team_caches_add_access_group( if cached_team.access_group_ids is None: cached_team.access_group_ids = [access_group_id] elif access_group_id not in cached_team.access_group_ids: - cached_team.access_group_ids = list(cached_team.access_group_ids) + [access_group_id] + cached_team.access_group_ids = list(cached_team.access_group_ids) + [ + access_group_id + ] else: continue await _cache_team_object( @@ -230,7 +246,9 @@ async def _patch_key_caches_add_access_group( if cached_key.access_group_ids is None: cached_key.access_group_ids = [access_group_id] elif access_group_id not in cached_key.access_group_ids: - cached_key.access_group_ids = list(cached_key.access_group_ids) + [access_group_id] + cached_key.access_group_ids = list(cached_key.access_group_ids) + [ + access_group_id + ] else: continue await _cache_key_object( @@ -281,7 +299,9 @@ async def create_access_group( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ) -> AccessGroupResponse: _require_proxy_admin(user_api_key_dict) - prisma_client = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value) + prisma_client = get_prisma_client_or_throw( + CommonProxyErrors.db_not_connected_error.value + ) try: async with prisma_client.db.tx() as tx: @@ -330,10 +350,16 @@ async def create_access_group( await _cache_access_group_record(record) await _patch_team_caches_add_access_group( - data.assigned_team_ids or [], record.access_group_id, user_api_key_cache, proxy_logging_obj + data.assigned_team_ids or [], + record.access_group_id, + user_api_key_cache, + proxy_logging_obj, ) await _patch_key_caches_add_access_group( - data.assigned_key_ids or [], record.access_group_id, user_api_key_cache, proxy_logging_obj + data.assigned_key_ids or [], + record.access_group_id, + user_api_key_cache, + proxy_logging_obj, ) return _record_to_response(record) @@ -347,7 +373,9 @@ async def list_access_groups( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ) -> List[AccessGroupResponse]: _require_proxy_admin(user_api_key_dict) - prisma_client = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value) + prisma_client = get_prisma_client_or_throw( + CommonProxyErrors.db_not_connected_error.value + ) records = await prisma_client.db.litellm_accessgrouptable.find_many( order={"created_at": "desc"} @@ -364,7 +392,9 @@ async def get_access_group( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ) -> AccessGroupResponse: _require_proxy_admin(user_api_key_dict) - prisma_client = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value) + prisma_client = get_prisma_client_or_throw( + CommonProxyErrors.db_not_connected_error.value + ) record = await prisma_client.db.litellm_accessgrouptable.find_unique( where={"access_group_id": access_group_id} @@ -387,12 +417,24 @@ async def update_access_group( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ) -> AccessGroupResponse: _require_proxy_admin(user_api_key_dict) - prisma_client = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value) + prisma_client = get_prisma_client_or_throw( + CommonProxyErrors.db_not_connected_error.value + ) update_fields = data.model_dump(exclude_unset=True) update_data: dict = {"updated_by": user_api_key_dict.user_id} for field, value in update_fields.items(): - if field in ("assigned_team_ids", "assigned_key_ids", "access_model_names", "access_mcp_server_ids", "access_agent_ids") and value is None: + if ( + field + in ( + "assigned_team_ids", + "assigned_key_ids", + "access_model_names", + "access_mcp_server_ids", + "access_agent_ids", + ) + and value is None + ): value = [] update_data[field] = value @@ -418,8 +460,16 @@ async def update_access_group( old_team_ids: Set[str] = set(existing.assigned_team_ids or []) old_key_ids: Set[str] = set(existing.assigned_key_ids or []) - new_team_ids: Set[str] = set(update_fields["assigned_team_ids"] or []) if "assigned_team_ids" in update_fields else old_team_ids - new_key_ids: Set[str] = set(update_fields["assigned_key_ids"] or []) if "assigned_key_ids" in update_fields else old_key_ids + new_team_ids: Set[str] = ( + set(update_fields["assigned_team_ids"] or []) + if "assigned_team_ids" in update_fields + else old_team_ids + ) + new_key_ids: Set[str] = ( + set(update_fields["assigned_key_ids"] or []) + if "assigned_key_ids" in update_fields + else old_key_ids + ) teams_to_add = list(new_team_ids - old_team_ids) teams_to_remove = list(old_team_ids - new_team_ids) @@ -432,9 +482,13 @@ async def update_access_group( ) await _sync_add_access_group_to_teams(tx, teams_to_add, access_group_id) - await _sync_remove_access_group_from_teams(tx, teams_to_remove, access_group_id) + await _sync_remove_access_group_from_teams( + tx, teams_to_remove, access_group_id + ) await _sync_add_access_group_to_keys(tx, keys_to_add, access_group_id) - await _sync_remove_access_group_from_keys(tx, keys_to_remove, access_group_id) + await _sync_remove_access_group_from_keys( + tx, keys_to_remove, access_group_id + ) except HTTPException: raise except Exception as e: @@ -449,10 +503,18 @@ async def update_access_group( from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache await _cache_access_group_record(record) - await _patch_team_caches_add_access_group(teams_to_add, access_group_id, user_api_key_cache, proxy_logging_obj) - await _patch_team_caches_remove_access_group(teams_to_remove, access_group_id, user_api_key_cache, proxy_logging_obj) - await _patch_key_caches_add_access_group(keys_to_add, access_group_id, user_api_key_cache, proxy_logging_obj) - await _patch_key_caches_remove_access_group(keys_to_remove, access_group_id, user_api_key_cache, proxy_logging_obj) + await _patch_team_caches_add_access_group( + teams_to_add, access_group_id, user_api_key_cache, proxy_logging_obj + ) + await _patch_team_caches_remove_access_group( + teams_to_remove, access_group_id, user_api_key_cache, proxy_logging_obj + ) + await _patch_key_caches_add_access_group( + keys_to_add, access_group_id, user_api_key_cache, proxy_logging_obj + ) + await _patch_key_caches_remove_access_group( + keys_to_remove, access_group_id, user_api_key_cache, proxy_logging_obj + ) return _record_to_response(record) @@ -466,7 +528,9 @@ async def delete_access_group( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ) -> None: _require_proxy_admin(user_api_key_dict) - prisma_client = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value) + prisma_client = get_prisma_client_or_throw( + CommonProxyErrors.db_not_connected_error.value + ) try: affected_team_ids: List[str] = [] @@ -487,10 +551,9 @@ async def delete_access_group( teams_with_group = await tx.litellm_teamtable.find_many( where={"access_group_ids": {"hasSome": [access_group_id]}} ) - all_affected_team_ids: Set[str] = ( - {team.team_id for team in teams_with_group} - | set(existing.assigned_team_ids or []) - ) + all_affected_team_ids: Set[str] = { + team.team_id for team in teams_with_group + } | set(existing.assigned_team_ids or []) affected_team_ids = list(all_affected_team_ids) # Union of: keys that have this access_group_id in their own access_group_ids @@ -498,31 +561,50 @@ async def delete_access_group( keys_with_group = await tx.litellm_verificationtoken.find_many( where={"access_group_ids": {"hasSome": [access_group_id]}} ) - all_affected_key_tokens: Set[str] = ( - {key.token for key in keys_with_group} - | set(existing.assigned_key_ids or []) - ) + all_affected_key_tokens: Set[str] = { + key.token for key in keys_with_group + } | set(existing.assigned_key_ids or []) affected_key_tokens = list(all_affected_key_tokens) # Update teams returned by find_many directly — we already have their data. for team in teams_with_group: await tx.litellm_teamtable.update( where={"team_id": team.team_id}, - data={"access_group_ids": [ag for ag in (team.access_group_ids or []) if ag != access_group_id]}, + data={ + "access_group_ids": [ + ag + for ag in (team.access_group_ids or []) + if ag != access_group_id + ] + }, ) # Use _sync_remove only for out-of-sync teams not found by the hasSome query. - out_of_sync_team_ids = set(existing.assigned_team_ids or []) - {t.team_id for t in teams_with_group} - await _sync_remove_access_group_from_teams(tx, list(out_of_sync_team_ids), access_group_id) + out_of_sync_team_ids = set(existing.assigned_team_ids or []) - { + t.team_id for t in teams_with_group + } + await _sync_remove_access_group_from_teams( + tx, list(out_of_sync_team_ids), access_group_id + ) # Update keys returned by find_many directly — we already have their data. for key in keys_with_group: await tx.litellm_verificationtoken.update( where={"token": key.token}, - data={"access_group_ids": [ag for ag in (key.access_group_ids or []) if ag != access_group_id]}, + data={ + "access_group_ids": [ + ag + for ag in (key.access_group_ids or []) + if ag != access_group_id + ] + }, ) # Use _sync_remove only for out-of-sync keys not found by the hasSome query. - out_of_sync_key_tokens = set(existing.assigned_key_ids or []) - {k.token for k in keys_with_group} - await _sync_remove_access_group_from_keys(tx, list(out_of_sync_key_tokens), access_group_id) + out_of_sync_key_tokens = set(existing.assigned_key_ids or []) - { + k.token for k in keys_with_group + } + await _sync_remove_access_group_from_keys( + tx, list(out_of_sync_key_tokens), access_group_id + ) await tx.litellm_accessgrouptable.delete( where={"access_group_id": access_group_id} @@ -551,7 +633,9 @@ async def delete_access_group( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail=CommonProxyErrors.db_not_connected_error.value, ) - if "P2025" in str(e) or ("record" in str(e).lower() and "not found" in str(e).lower()): + if "P2025" in str(e) or ( + "record" in str(e).lower() and "not found" in str(e).lower() + ): raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail=f"Access group '{access_group_id}' not found", diff --git a/litellm/proxy/management_endpoints/cache_settings_endpoints.py b/litellm/proxy/management_endpoints/cache_settings_endpoints.py index c9eeea15e2..8d6ec2cec1 100644 --- a/litellm/proxy/management_endpoints/cache_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/cache_settings_endpoints.py @@ -31,31 +31,32 @@ class CacheSettingsManager: Manages cache settings initialization and updates. Tracks last cache params to avoid unnecessary reinitialization. """ - + _last_cache_params: Optional[Dict[str, Any]] = None - + @staticmethod def _cache_params_equal(params1: Dict[str, Any], params2: Dict[str, Any]) -> bool: """ Compare two cache parameter dictionaries for equality. Normalizes values and filters out UI-only fields. """ + # Normalize by removing None values and UI-only fields def normalize(params: Dict[str, Any]) -> Dict[str, Any]: normalized = {} for k, v in params.items(): - if k == 'redis_type': # Skip UI-only field + if k == "redis_type": # Skip UI-only field continue if v is not None: # Convert to string for comparison to handle different types normalized[k] = str(v) if not isinstance(v, (list, dict)) else v return normalized - + normalized1 = normalize(params1) normalized2 = normalize(params2) - + return normalized1 == normalized2 - + @staticmethod async def init_cache_settings_in_db(prisma_client, proxy_config): """ @@ -63,7 +64,7 @@ class CacheSettingsManager: Only reinitializes if cache params have changed. """ import json - + try: cache_config = await prisma_client.db.litellm_cacheconfig.find_unique( where={"id": "cache_config"} @@ -75,44 +76,47 @@ class CacheSettingsManager: cache_settings_dict = json.loads(cache_settings_json) else: cache_settings_dict = cache_settings_json - + # Decrypt cache settings decrypted_settings = proxy_config._decrypt_db_variables( variables_dict=cache_settings_dict ) - + # Remove redis_type if present (UI-only field, not a Cache parameter) # We derive it for UI in get_cache_settings endpoint - cache_params = {k: v for k, v in decrypted_settings.items() if k != "redis_type"} - + cache_params = { + k: v for k, v in decrypted_settings.items() if k != "redis_type" + } + # Check if cache params have changed - if CacheSettingsManager._last_cache_params is not None and CacheSettingsManager._cache_params_equal( - CacheSettingsManager._last_cache_params, cache_params + if ( + CacheSettingsManager._last_cache_params is not None + and CacheSettingsManager._cache_params_equal( + CacheSettingsManager._last_cache_params, cache_params + ) ): verbose_proxy_logger.debug( "Cache settings unchanged, skipping reinitialization" ) return - + # Initialize cache only if params changed or cache not initialized proxy_config._init_cache(cache_params=cache_params) - + # Store the params we just initialized CacheSettingsManager._last_cache_params = cache_params.copy() - + # Switch on LLM response caching proxy_config.switch_on_llm_response_caching() - - verbose_proxy_logger.info( - "Cache settings initialized from database" - ) + + verbose_proxy_logger.info("Cache settings initialized from database") except Exception as e: verbose_proxy_logger.exception( "litellm.proxy.management_endpoints.cache_settings_endpoints.py::CacheSettingsManager::init_cache_settings_in_db - {}".format( str(e) ) ) - + @staticmethod def update_cache_params(cache_params: Dict[str, Any]): """ @@ -143,13 +147,13 @@ class CacheTestRequest(BaseModel): class CacheTestResponse(BaseModel): status: str = Field(description="Connection status: 'success' or 'failed'") message: str = Field(description="Connection result message") - error: Optional[str] = Field(default=None, description="Error message if connection failed") + error: Optional[str] = Field( + default=None, description="Error message if connection failed" + ) class CacheSettingsUpdateRequest(BaseModel): - cache_settings: Dict[str, Any] = Field( - description="Cache settings to save" - ) + cache_settings: Dict[str, Any] = Field(description="Cache settings to save") @router.get( @@ -163,17 +167,17 @@ async def get_cache_settings( ): """ Get cache configuration and available settings. - + Returns: - fields: List of all configurable cache settings with their metadata (type, description, default, options) - current_values: Current values of cache settings from database """ from litellm.proxy.proxy_server import prisma_client, proxy_config - + try: # Get cache settings fields from types file cache_fields = [field.model_copy(deep=True) for field in CACHE_SETTINGS_FIELDS] - + # Try to get cache settings from database current_values = {} if prisma_client is not None: @@ -187,12 +191,12 @@ async def get_cache_settings( cache_settings_dict = json.loads(cache_settings_json) else: cache_settings_dict = cache_settings_json - + # Decrypt environment variables decrypted_settings = proxy_config._decrypt_db_variables( variables_dict=cache_settings_dict ) - + # Derive redis_type for UI based on settings # UI uses redis_type to show/hide fields, backend only stores 'type' if decrypted_settings.get("type") == "redis": @@ -202,26 +206,23 @@ async def get_cache_settings( decrypted_settings["redis_type"] = "sentinel" else: decrypted_settings["redis_type"] = "node" - + current_values = decrypted_settings - + # Update field values with current values for field in cache_fields: if field.field_name in current_values: field.field_value = current_values[field.field_name] - + return CacheSettingsResponse( fields=cache_fields, current_values=current_values, redis_type_descriptions=REDIS_TYPE_DESCRIPTIONS, ) except Exception as e: - verbose_proxy_logger.error( - f"Error fetching cache settings: {str(e)}" - ) + verbose_proxy_logger.error(f"Error fetching cache settings: {str(e)}") raise HTTPException( - status_code=500, - detail=f"Error fetching cache settings: {str(e)}" + status_code=500, detail=f"Error fetching cache settings: {str(e)}" ) @@ -237,35 +238,35 @@ async def test_cache_connection( ): """ Test cache connection with provided credentials. - + Creates a temporary cache instance and uses its test_connection method to verify the credentials work without affecting global state. """ from litellm import Cache - + try: cache_settings = request.cache_settings.copy() - verbose_proxy_logger.debug("Testing cache connection with settings: %s", cache_settings) - + verbose_proxy_logger.debug( + "Testing cache connection with settings: %s", cache_settings + ) + # Only support Redis for now if cache_settings.get("type") != "redis": return CacheTestResponse( status="failed", message="Only Redis cache type is currently supported for testing", ) - + # Create temporary cache instance temp_cache = Cache(**cache_settings) - + # Use the cache's test_connection method result = await temp_cache.cache.test_connection() - + return CacheTestResponse(**result) - + except Exception as e: - verbose_proxy_logger.error( - f"Error testing cache connection: {str(e)}" - ) + verbose_proxy_logger.error(f"Error testing cache connection: {str(e)}") return CacheTestResponse( status="failed", message=f"Cache connection test failed: {str(e)}", @@ -284,7 +285,7 @@ async def update_cache_settings( ): """ Save cache settings to database and initialize cache. - + This endpoint: 1. Encrypts sensitive fields (passwords, etc.) 2. Saves to LiteLLM_CacheConfig table @@ -295,13 +296,13 @@ async def update_cache_settings( proxy_config, store_model_in_db, ) - + if prisma_client is None: raise HTTPException( status_code=500, detail={"error": "Database not connected. Please connect a database."}, ) - + if store_model_in_db is not True: raise HTTPException( status_code=500, @@ -309,15 +310,15 @@ async def update_cache_settings( "error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature." }, ) - + try: cache_settings = request.cache_settings.copy() - + # Encrypt sensitive fields (keep redis_type for storage) encrypted_settings = proxy_config._encrypt_env_variables( environment_variables=cache_settings ) - + # Save to database await prisma_client.db.litellm_cacheconfig.upsert( where={"id": "cache_config"}, @@ -331,36 +332,34 @@ async def update_cache_settings( }, }, ) - + # Reinitialize cache with new settings # Decrypt for initialization decrypted_settings = proxy_config._decrypt_db_variables( variables_dict=encrypted_settings ) - + # Remove redis_type if present (UI-only field, not a Cache parameter) - cache_params = {k: v for k, v in decrypted_settings.items() if k != "redis_type"} - + cache_params = { + k: v for k, v in decrypted_settings.items() if k != "redis_type" + } + # Initialize cache (frontend sends type="redis", not redis_type) proxy_config._init_cache(cache_params=cache_params) - + # Update the last cache params to avoid reinitializing unnecessarily CacheSettingsManager.update_cache_params(cache_params) - + # Switch on LLM response caching proxy_config.switch_on_llm_response_caching() - + return { "message": "Cache settings updated successfully", "status": "success", "settings": cache_settings, } except Exception as e: - verbose_proxy_logger.error( - f"Error updating cache settings: {str(e)}" - ) + verbose_proxy_logger.error(f"Error updating cache settings: {str(e)}") raise HTTPException( - status_code=500, - detail=f"Error updating cache settings: {str(e)}" + status_code=500, detail=f"Error updating cache settings: {str(e)}" ) - diff --git a/litellm/proxy/management_endpoints/callback_management_endpoints.py b/litellm/proxy/management_endpoints/callback_management_endpoints.py index 3bb7511fef..9132d3fe1d 100644 --- a/litellm/proxy/management_endpoints/callback_management_endpoints.py +++ b/litellm/proxy/management_endpoints/callback_management_endpoints.py @@ -26,7 +26,7 @@ async def list_callbacks(): # Get callbacks organized by type using the callback manager utility callbacks_by_type = logging_callback_manager.get_callbacks_by_type() - + return callbacks_by_type @@ -38,17 +38,17 @@ async def list_callbacks(): async def get_callback_configs(): """ Get Available Callback Configurations - + Returns the configuration details for all available logging callbacks, including supported parameters, field types, and descriptions. """ config_path = os.path.join( os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "integrations", - "callback_configs.json" + "callback_configs.json", ) - + with open(config_path, "r") as f: configs = json.load(f) - - return configs \ No newline at end of file + + return configs diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 02961748e7..011d2f7485 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -48,7 +48,9 @@ def _is_user_agent_tag(tag: Optional[str]) -> bool: if not tag: return False normalized_tag = tag.strip().lower() - return normalized_tag.startswith("user-agent:") or normalized_tag.startswith("user agent:") + return normalized_tag.startswith("user-agent:") or normalized_tag.startswith( + "user agent:" + ) def compute_tag_metadata_totals(records: List[Any]) -> SpendMetrics: @@ -103,26 +105,24 @@ def update_breakdown_metrics( # Update API key breakdown for this model if record.api_key not in breakdown.models[record.model].api_key_breakdown: - breakdown.models[record.model].api_key_breakdown[record.api_key] = ( - KeyMetricWithMetadata( - metrics=SpendMetrics(), - metadata=KeyMetadata( - key_alias=api_key_metadata.get(record.api_key, {}).get( - "key_alias", None - ), - team_id=api_key_metadata.get(record.api_key, {}).get( - "team_id", None - ), + breakdown.models[record.model].api_key_breakdown[ + record.api_key + ] = KeyMetricWithMetadata( + metrics=SpendMetrics(), + metadata=KeyMetadata( + key_alias=api_key_metadata.get(record.api_key, {}).get( + "key_alias", None ), - ) - ) - breakdown.models[record.model].api_key_breakdown[record.api_key].metrics = ( - update_metrics( - breakdown.models[record.model] - .api_key_breakdown[record.api_key] - .metrics, - record, + team_id=api_key_metadata.get(record.api_key, {}).get( + "team_id", None + ), + ), ) + breakdown.models[record.model].api_key_breakdown[ + record.api_key + ].metrics = update_metrics( + breakdown.models[record.model].api_key_breakdown[record.api_key].metrics, + record, ) # Update model group breakdown @@ -218,24 +218,22 @@ def update_breakdown_metrics( # Update API key breakdown for this provider if record.api_key not in breakdown.providers[provider].api_key_breakdown: - breakdown.providers[provider].api_key_breakdown[record.api_key] = ( - KeyMetricWithMetadata( - metrics=SpendMetrics(), - metadata=KeyMetadata( - key_alias=api_key_metadata.get(record.api_key, {}).get( - "key_alias", None - ), - team_id=api_key_metadata.get(record.api_key, {}).get( - "team_id", None - ), + breakdown.providers[provider].api_key_breakdown[ + record.api_key + ] = KeyMetricWithMetadata( + metrics=SpendMetrics(), + metadata=KeyMetadata( + key_alias=api_key_metadata.get(record.api_key, {}).get( + "key_alias", None ), - ) - ) - breakdown.providers[provider].api_key_breakdown[record.api_key].metrics = ( - update_metrics( - breakdown.providers[provider].api_key_breakdown[record.api_key].metrics, - record, + team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None), + ), ) + breakdown.providers[provider].api_key_breakdown[ + record.api_key + ].metrics = update_metrics( + breakdown.providers[provider].api_key_breakdown[record.api_key].metrics, + record, ) # Update endpoint breakdown @@ -251,26 +249,26 @@ def update_breakdown_metrics( # Update API key breakdown for this endpoint if record.api_key not in breakdown.endpoints[record.endpoint].api_key_breakdown: - breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key] = ( - KeyMetricWithMetadata( - metrics=SpendMetrics(), - metadata=KeyMetadata( - key_alias=api_key_metadata.get(record.api_key, {}).get( - "key_alias", None - ), - team_id=api_key_metadata.get(record.api_key, {}).get( - "team_id", None - ), + breakdown.endpoints[record.endpoint].api_key_breakdown[ + record.api_key + ] = KeyMetricWithMetadata( + metrics=SpendMetrics(), + metadata=KeyMetadata( + key_alias=api_key_metadata.get(record.api_key, {}).get( + "key_alias", None ), - ) - ) - breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key].metrics = ( - update_metrics( - breakdown.endpoints[record.endpoint] - .api_key_breakdown[record.api_key] - .metrics, - record, + team_id=api_key_metadata.get(record.api_key, {}).get( + "team_id", None + ), + ), ) + breakdown.endpoints[record.endpoint].api_key_breakdown[ + record.api_key + ].metrics = update_metrics( + breakdown.endpoints[record.endpoint] + .api_key_breakdown[record.api_key] + .metrics, + record, ) # Update api key breakdown @@ -309,26 +307,24 @@ def update_breakdown_metrics( # Update API key breakdown for this entity if record.api_key not in breakdown.entities[entity_value].api_key_breakdown: - breakdown.entities[entity_value].api_key_breakdown[record.api_key] = ( - KeyMetricWithMetadata( - metrics=SpendMetrics(), - metadata=KeyMetadata( - key_alias=api_key_metadata.get(record.api_key, {}).get( - "key_alias", None - ), - team_id=api_key_metadata.get(record.api_key, {}).get( - "team_id", None - ), + breakdown.entities[entity_value].api_key_breakdown[ + record.api_key + ] = KeyMetricWithMetadata( + metrics=SpendMetrics(), + metadata=KeyMetadata( + key_alias=api_key_metadata.get(record.api_key, {}).get( + "key_alias", None ), - ) - ) - breakdown.entities[entity_value].api_key_breakdown[record.api_key].metrics = ( - update_metrics( - breakdown.entities[entity_value] - .api_key_breakdown[record.api_key] - .metrics, - record, + team_id=api_key_metadata.get(record.api_key, {}).get( + "team_id", None + ), + ), ) + breakdown.entities[entity_value].api_key_breakdown[ + record.api_key + ].metrics = update_metrics( + breakdown.entities[entity_value].api_key_breakdown[record.api_key].metrics, + record, ) return breakdown @@ -347,8 +343,7 @@ async def get_api_key_metadata( where={"token": {"in": list(api_keys)}} ) result = { - k.token: {"key_alias": k.key_alias, "team_id": k.team_id} - for k in key_records + k.token: {"key_alias": k.key_alias, "team_id": k.team_id} for k in key_records } # For any keys not found in the active table, check the deleted keys table @@ -523,9 +518,7 @@ def _build_aggregated_sql_query( # Exclude specific entities if exclude_entity_ids: - placeholders = ", ".join( - f"${p + i}" for i in range(len(exclude_entity_ids)) - ) + placeholders = ", ".join(f"${p + i}" for i in range(len(exclude_entity_ids))) sql_conditions.append(f'"{entity_id_field}" NOT IN ({placeholders})') sql_params.extend(exclude_entity_ids) p += len(exclude_entity_ids) @@ -799,8 +792,12 @@ async def get_daily_activity_aggregated( total_api_requests=aggregated["totals"].api_requests, total_successful_requests=aggregated["totals"].successful_requests, total_failed_requests=aggregated["totals"].failed_requests, - total_cache_read_input_tokens=aggregated["totals"].cache_read_input_tokens, - total_cache_creation_input_tokens=aggregated["totals"].cache_creation_input_tokens, + total_cache_read_input_tokens=aggregated[ + "totals" + ].cache_read_input_tokens, + total_cache_creation_input_tokens=aggregated[ + "totals" + ].cache_creation_input_tokens, page=1, total_pages=1, has_more=False, diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index e22f4e1b67..efc42d3355 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -92,10 +92,7 @@ def _team_member_has_permission( if permission not in team_obj.team_member_permissions: return False for member in team_obj.members_with_roles: - if ( - member.user_id is not None - and member.user_id == user_api_key_dict.user_id - ): + if member.user_id is not None and member.user_id == user_api_key_dict.user_id: return True return False diff --git a/litellm/proxy/management_endpoints/config_override_endpoints.py b/litellm/proxy/management_endpoints/config_override_endpoints.py index 0e364e9bd1..d78c5526e6 100644 --- a/litellm/proxy/management_endpoints/config_override_endpoints.py +++ b/litellm/proxy/management_endpoints/config_override_endpoints.py @@ -218,9 +218,7 @@ async def update_hashicorp_vault_config( _set_env_vars(config_data) try: - proxy_config.initialize_secret_manager( - key_management_system="hashicorp_vault" - ) + proxy_config.initialize_secret_manager(key_management_system="hashicorp_vault") except Exception as e: _set_env_vars(previous_env) verbose_proxy_logger.exception( @@ -295,9 +293,7 @@ async def get_hashicorp_vault_config( # Decrypt then mask sensitive fields so plaintext secrets are never sent to the UI decrypted_data = proxy_config._decrypt_db_variables(config_data) - masked_data = _mask_sensitive_fields( - decrypted_data, HASHICORP_SENSITIVE_FIELDS - ) + masked_data = _mask_sensitive_fields(decrypted_data, HASHICORP_SENSITIVE_FIELDS) return ConfigOverrideSettingsResponse( config_type="hashicorp_vault", @@ -307,9 +303,7 @@ async def get_hashicorp_vault_config( # Fallback to env vars — also mask sensitive values env_values = _get_current_env_values(HASHICORP_ENV_VAR_MAPPING) - masked_env_values = _mask_sensitive_fields( - env_values, HASHICORP_SENSITIVE_FIELDS - ) + masked_env_values = _mask_sensitive_fields(env_values, HASHICORP_SENSITIVE_FIELDS) return ConfigOverrideSettingsResponse( config_type="hashicorp_vault", @@ -399,7 +393,9 @@ async def test_hashicorp_vault_connection( # Step 2: Verify the token is valid via token/lookup-self try: - async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.SecretManager) + async_client = get_async_httpx_client( + llm_provider=httpxSpecialProvider.SecretManager + ) lookup_url = f"{client.vault_addr}/v1/auth/token/lookup-self" if client.vault_namespace: headers["X-Vault-Namespace"] = client.vault_namespace diff --git a/litellm/proxy/management_endpoints/cost_tracking_settings.py b/litellm/proxy/management_endpoints/cost_tracking_settings.py index 4418d934c8..bf24d8924d 100644 --- a/litellm/proxy/management_endpoints/cost_tracking_settings.py +++ b/litellm/proxy/management_endpoints/cost_tracking_settings.py @@ -67,7 +67,12 @@ def _resolve_model_for_cost_lookup(model: str) -> Tuple[str, Optional[str]]: f"Resolved model '{model}' to base_model '{base_model}' from router" ) custom_llm_provider = litellm_params.get("custom_llm_provider") - return str(base_model), str(custom_llm_provider) if custom_llm_provider is not None else None + return ( + str(base_model), + str(custom_llm_provider) + if custom_llm_provider is not None + else None, + ) resolved_model = litellm_params.get("model") @@ -76,7 +81,12 @@ def _resolve_model_for_cost_lookup(model: str) -> Tuple[str, Optional[str]]: f"Resolved model '{model}' to '{resolved_model}' from router" ) custom_llm_provider = litellm_params.get("custom_llm_provider") - return str(resolved_model), str(custom_llm_provider) if custom_llm_provider is not None else None + return ( + str(resolved_model), + str(custom_llm_provider) + if custom_llm_provider is not None + else None, + ) except Exception as e: verbose_proxy_logger.debug( f"Could not resolve model '{model}' from router: {e}" @@ -114,30 +124,28 @@ async def get_cost_discount_config( ): """ Get current cost discount configuration. - + Returns the cost_discount_config from litellm_settings. """ from litellm.proxy.proxy_server import prisma_client, proxy_config - + if prisma_client is None: raise HTTPException( status_code=500, detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - + try: # Load config from DB config = await proxy_config.get_config() - + # Get cost_discount_config from litellm_settings litellm_settings = config.get("litellm_settings", {}) cost_discount_config = litellm_settings.get("cost_discount_config", {}) - + return {"values": cost_discount_config} except Exception as e: - verbose_proxy_logger.error( - f"Error fetching cost discount config: {str(e)}" - ) + verbose_proxy_logger.error(f"Error fetching cost discount config: {str(e)}") return {"values": {}} @@ -152,10 +160,10 @@ async def update_cost_discount_config( ): """ Update cost discount configuration. - + Updates the cost_discount_config in litellm_settings. Discounts should be between 0 and 1 (e.g., 0.05 = 5% discount). - + Example: ```json { @@ -170,13 +178,13 @@ async def update_cost_discount_config( proxy_config, store_model_in_db, ) - + if prisma_client is None: raise HTTPException( status_code=500, detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - + if store_model_in_db is not True: raise HTTPException( status_code=500, @@ -184,13 +192,13 @@ async def update_cost_discount_config( "error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature." }, ) - + # Validate that all providers are valid LiteLLM providers invalid_providers = [] for provider in cost_discount_config.keys(): if provider not in LlmProvidersSet: invalid_providers.append(provider) - + if invalid_providers: raise HTTPException( status_code=400, @@ -198,53 +206,50 @@ async def update_cost_discount_config( "error": f"Invalid provider(s): {', '.join(invalid_providers)}. Must be valid LiteLLM providers. See https://docs.litellm.ai/docs/providers for the full list." }, ) - + # Validate discount values are between 0 and 1 for provider, discount in cost_discount_config.items(): if not isinstance(discount, (int, float)): raise HTTPException( - status_code=400, - detail=f"Discount for {provider} must be a number" + status_code=400, detail=f"Discount for {provider} must be a number" ) if not (0 <= discount <= 1): raise HTTPException( status_code=400, - detail=f"Discount for {provider} must be between 0 and 1 (0% to 100%)" + detail=f"Discount for {provider} must be between 0 and 1 (0% to 100%)", ) - + try: # Load existing config config = await proxy_config.get_config() - + # Ensure litellm_settings exists if "litellm_settings" not in config: config["litellm_settings"] = {} - + # Update cost_discount_config config["litellm_settings"]["cost_discount_config"] = cost_discount_config - + # Save the updated config to DB await proxy_config.save_config(new_config=config) - + # Update in-memory litellm.cost_discount_config litellm.cost_discount_config = cost_discount_config - + verbose_proxy_logger.info( f"Updated cost_discount_config: {cost_discount_config}" ) - + return { "message": "Cost discount configuration updated successfully", "status": "success", - "values": cost_discount_config + "values": cost_discount_config, } except Exception as e: - verbose_proxy_logger.error( - f"Error updating cost discount config: {str(e)}" - ) + verbose_proxy_logger.error(f"Error updating cost discount config: {str(e)}") raise HTTPException( status_code=500, - detail={"error": f"Failed to update cost discount config: {str(e)}"} + detail={"error": f"Failed to update cost discount config: {str(e)}"}, ) @@ -258,30 +263,28 @@ async def get_cost_margin_config( ): """ Get current cost margin configuration. - + Returns the cost_margin_config from litellm_settings. """ from litellm.proxy.proxy_server import prisma_client, proxy_config - + if prisma_client is None: raise HTTPException( status_code=500, detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - + try: # Load config from DB config = await proxy_config.get_config() - + # Get cost_margin_config from litellm_settings litellm_settings = config.get("litellm_settings", {}) cost_margin_config = litellm_settings.get("cost_margin_config", {}) - + return {"values": cost_margin_config} except Exception as e: - verbose_proxy_logger.error( - f"Error fetching cost margin config: {str(e)}" - ) + verbose_proxy_logger.error(f"Error fetching cost margin config: {str(e)}") return {"values": {}} @@ -296,14 +299,14 @@ async def update_cost_margin_config( ): """ Update cost margin configuration. - + Updates the cost_margin_config in litellm_settings. Margins can be: - Percentage: {"openai": 0.10} = 10% margin - Fixed amount: {"openai": {"fixed_amount": 0.001}} = $0.001 per request - Combined: {"vertex_ai": {"percentage": 0.08, "fixed_amount": 0.0005}} - Global: {"global": 0.05} = 5% global margin on all providers - + Example: ```json { @@ -319,13 +322,13 @@ async def update_cost_margin_config( proxy_config, store_model_in_db, ) - + if prisma_client is None: raise HTTPException( status_code=500, detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - + if store_model_in_db is not True: raise HTTPException( status_code=500, @@ -333,13 +336,13 @@ async def update_cost_margin_config( "error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature." }, ) - + # Validate that all providers are valid LiteLLM providers (except "global") invalid_providers = [] for provider in cost_margin_config.keys(): if provider != "global" and provider not in LlmProvidersSet: invalid_providers.append(provider) - + if invalid_providers: raise HTTPException( status_code=400, @@ -347,7 +350,7 @@ async def update_cost_margin_config( "error": f"Invalid provider(s): {', '.join(invalid_providers)}. Must be valid LiteLLM providers or 'global'. See https://docs.litellm.ai/docs/providers for the full list." }, ) - + # Validate margin values for provider, margin_value in cost_margin_config.items(): if isinstance(margin_value, (int, float)): @@ -355,7 +358,7 @@ async def update_cost_margin_config( if not (0 <= margin_value <= 10): # Allow up to 1000% margin raise HTTPException( status_code=400, - detail=f"Margin percentage for {provider} must be between 0 and 10 (0% to 1000%)" + detail=f"Margin percentage for {provider} must be between 0 and 10 (0% to 1000%)", ) elif isinstance(margin_value, dict): # Complex format: {"percentage": 0.08, "fixed_amount": 0.0005} @@ -364,69 +367,65 @@ async def update_cost_margin_config( if not isinstance(percentage, (int, float)): raise HTTPException( status_code=400, - detail=f"Margin percentage for {provider} must be a number" + detail=f"Margin percentage for {provider} must be a number", ) if not (0 <= percentage <= 10): raise HTTPException( status_code=400, - detail=f"Margin percentage for {provider} must be between 0 and 10 (0% to 1000%)" + detail=f"Margin percentage for {provider} must be between 0 and 10 (0% to 1000%)", ) if "fixed_amount" in margin_value: fixed_amount = margin_value["fixed_amount"] if not isinstance(fixed_amount, (int, float)): raise HTTPException( status_code=400, - detail=f"Fixed margin amount for {provider} must be a number" + detail=f"Fixed margin amount for {provider} must be a number", ) if fixed_amount < 0: raise HTTPException( status_code=400, - detail=f"Fixed margin amount for {provider} must be non-negative" + detail=f"Fixed margin amount for {provider} must be non-negative", ) if not margin_value: # Empty dict raise HTTPException( status_code=400, - detail=f"Margin config for {provider} cannot be empty. Must include 'percentage' and/or 'fixed_amount'" + detail=f"Margin config for {provider} cannot be empty. Must include 'percentage' and/or 'fixed_amount'", ) else: raise HTTPException( status_code=400, - detail=f"Margin for {provider} must be a number (percentage) or dict with 'percentage' and/or 'fixed_amount'" + detail=f"Margin for {provider} must be a number (percentage) or dict with 'percentage' and/or 'fixed_amount'", ) - + try: # Load existing config config = await proxy_config.get_config() - + # Ensure litellm_settings exists if "litellm_settings" not in config: config["litellm_settings"] = {} - + # Update cost_margin_config config["litellm_settings"]["cost_margin_config"] = cost_margin_config - + # Save the updated config to DB await proxy_config.save_config(new_config=config) - + # Update in-memory litellm.cost_margin_config litellm.cost_margin_config = cost_margin_config - - verbose_proxy_logger.info( - f"Updated cost_margin_config: {cost_margin_config}" - ) - + + verbose_proxy_logger.info(f"Updated cost_margin_config: {cost_margin_config}") + return { "message": "Cost margin configuration updated successfully", "status": "success", - "values": cost_margin_config + "values": cost_margin_config, } except Exception as e: - verbose_proxy_logger.error( - f"Error updating cost margin config: {str(e)}" - ) + verbose_proxy_logger.error(f"Error updating cost margin config: {str(e)}") raise HTTPException( status_code=500, - detail={"error": f"Failed to update cost margin config: {str(e)}"} + detail={"error": f"Failed to update cost margin config: {str(e)}"}, ) @@ -520,7 +519,9 @@ async def estimate_cost( input_cost = cost_breakdown.get("input_cost", 0.0) if cost_breakdown else 0.0 output_cost = cost_breakdown.get("output_cost", 0.0) if cost_breakdown else 0.0 - margin_cost = cost_breakdown.get("margin_total_amount", 0.0) if cost_breakdown else 0.0 + margin_cost = ( + cost_breakdown.get("margin_total_amount", 0.0) if cost_breakdown else 0.0 + ) # Get model info for per-token pricing display try: @@ -538,23 +539,29 @@ async def estimate_cost( custom_llm_provider = resolved_provider # Calculate daily and monthly costs - daily_cost, daily_input_cost, daily_output_cost, daily_margin_cost = ( - _calculate_period_costs( - num_requests=request.num_requests_per_day, - cost_per_request=cost_per_request, - input_cost=input_cost, - output_cost=output_cost, - margin_cost=margin_cost, - ) + ( + daily_cost, + daily_input_cost, + daily_output_cost, + daily_margin_cost, + ) = _calculate_period_costs( + num_requests=request.num_requests_per_day, + cost_per_request=cost_per_request, + input_cost=input_cost, + output_cost=output_cost, + margin_cost=margin_cost, ) - monthly_cost, monthly_input_cost, monthly_output_cost, monthly_margin_cost = ( - _calculate_period_costs( - num_requests=request.num_requests_per_month, - cost_per_request=cost_per_request, - input_cost=input_cost, - output_cost=output_cost, - margin_cost=margin_cost, - ) + ( + monthly_cost, + monthly_input_cost, + monthly_output_cost, + monthly_margin_cost, + ) = _calculate_period_costs( + num_requests=request.num_requests_per_month, + cost_per_request=cost_per_request, + input_cost=input_cost, + output_cost=output_cost, + margin_cost=margin_cost, ) return CostEstimateResponse( @@ -579,4 +586,3 @@ async def estimate_cost( output_cost_per_token=output_cost_per_token, provider=custom_llm_provider, ) - diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index bce5f6cda7..084c2f47d0 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -21,13 +21,15 @@ from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm._logging import verbose_proxy_logger from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.management_endpoints.common_daily_activity import \ - get_daily_activity +from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity from litellm.proxy.management_helpers.object_permission_utils import ( - _set_object_permission, handle_update_object_permission_common) + _set_object_permission, + handle_update_object_permission_common, +) from litellm.proxy.utils import handle_exception_on_proxy -from litellm.types.proxy.management_endpoints.common_daily_activity import \ - SpendAnalyticsPaginatedResponse +from litellm.types.proxy.management_endpoints.common_daily_activity import ( + SpendAnalyticsPaginatedResponse, +) router = APIRouter() @@ -111,8 +113,9 @@ async def unblock_user(data: BlockUsers): ``` """ try: - from enterprise.enterprise_hooks.blocked_user_list import \ - _ENTERPRISE_BlockedUserList + from enterprise.enterprise_hooks.blocked_user_list import ( + _ENTERPRISE_BlockedUserList, + ) except ImportError: raise HTTPException( status_code=400, @@ -164,7 +167,10 @@ def new_budget_request(data: NewCustomerRequest) -> Optional[BudgetNewRequest]: if budget_kv_pairs: budget_request = BudgetNewRequest(**budget_kv_pairs) - if budget_request.budget_reset_at is None and budget_request.budget_duration is not None: + if ( + budget_request.budget_reset_at is None + and budget_request.budget_duration is not None + ): budget_request.budget_reset_at = datetime.utcnow() + timedelta( seconds=duration_in_seconds(duration=budget_request.budget_duration) ) @@ -296,8 +302,11 @@ async def new_end_user( - end-user object - currently allowed models """ - from litellm.proxy.proxy_server import (litellm_proxy_admin_name, - llm_router, prisma_client) + from litellm.proxy.proxy_server import ( + litellm_proxy_admin_name, + llm_router, + prisma_client, + ) if prisma_client is None: raise HTTPException( @@ -373,7 +382,13 @@ async def new_end_user( response_dict = end_user_record.model_dump() if response_dict.get("object_permission"): # Remove reverse relations from object_permission - for field in ["teams", "verification_tokens", "organizations", "users", "end_users"]: + for field in [ + "teams", + "verification_tokens", + "organizations", + "users", + "end_users", + ]: response_dict["object_permission"].pop(field, None) return response_dict @@ -432,7 +447,8 @@ async def end_user_info( ) user_info = await prisma_client.db.litellm_endusertable.find_first( - where={"user_id": end_user_id}, include={"litellm_budget_table": True, "object_permission": True} + where={"user_id": end_user_id}, + include={"litellm_budget_table": True, "object_permission": True}, ) if user_info is None: @@ -447,11 +463,17 @@ async def end_user_info( response_dict = user_info.model_dump(exclude_none=True) if response_dict.get("object_permission"): # Remove reverse relations from object_permission - for field in ["teams", "verification_tokens", "organizations", "users", "end_users"]: + for field in [ + "teams", + "verification_tokens", + "organizations", + "users", + "end_users", + ]: response_dict["object_permission"].pop(field, None) return response_dict - + except Exception as e: verbose_proxy_logger.exception( "litellm.proxy.management_endpoints.customer_endpoints.end_user_info(): Exception occured - {}".format( @@ -460,6 +482,7 @@ async def end_user_info( ) raise handle_exception_on_proxy(e) + @router.post( "/customer/update", tags=["Customer Management"], @@ -527,8 +550,7 @@ async def update_end_user( ``` """ - from litellm.proxy.proxy_server import (litellm_proxy_admin_name, - prisma_client) + from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client try: data_json: dict = data.json() @@ -645,7 +667,13 @@ async def update_end_user( response_dict = response.model_dump() if response_dict.get("object_permission"): # Remove reverse relations from object_permission - for field in ["teams", "verification_tokens", "organizations", "users", "end_users"]: + for field in [ + "teams", + "verification_tokens", + "organizations", + "users", + "end_users", + ]: response_dict["object_permission"].pop(field, None) return response_dict @@ -751,6 +779,7 @@ async def delete_end_user( ) raise handle_exception_on_proxy(e) + @router.get( "/customer/list", tags=["Customer Management"], @@ -808,11 +837,17 @@ async def list_end_user( item_dict = item.model_dump() # Remove reverse relations from object_permission if item_dict.get("object_permission"): - for field in ["teams", "verification_tokens", "organizations", "users", "end_users"]: + for field in [ + "teams", + "verification_tokens", + "organizations", + "users", + "end_users", + ]: item_dict["object_permission"].pop(field, None) returned_response.append(LiteLLM_EndUserTable(**item_dict)) return returned_response - + except Exception as e: verbose_proxy_logger.exception( "litellm.proxy.management_endpoints.customer_endpoints.list_end_user(): Exception occured - {}".format( @@ -821,6 +856,7 @@ async def list_end_user( ) raise handle_exception_on_proxy(e) + @router.get( "/customer/daily/activity", tags=["Customer Management"], @@ -844,7 +880,6 @@ async def get_customer_daily_activity( exclude_end_user_ids: Optional[str] = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): - """ Get daily activity for specific organizations or all accessible organizations. """ @@ -864,7 +899,6 @@ async def get_customer_daily_activity( exclude_end_user_ids.split(",") if exclude_end_user_ids else None ) - # Fetch organization aliases for metadata where_condition = {} if end_user_ids_list: @@ -872,10 +906,7 @@ async def get_customer_daily_activity( end_user_aliases = await prisma_client.db.litellm_endusertable.find_many( where=where_condition ) - end_user_alias_metadata = { - e.user_id: {"alias": e.alias} - for e in end_user_aliases - } + end_user_alias_metadata = {e.user_id: {"alias": e.alias} for e in end_user_aliases} # Query daily activity for organizations return await get_daily_activity( @@ -891,4 +922,4 @@ async def get_customer_daily_activity( api_key=api_key, page=page, page_size=page_size, - ) \ No newline at end of file + ) diff --git a/litellm/proxy/management_endpoints/fallback_management_endpoints.py b/litellm/proxy/management_endpoints/fallback_management_endpoints.py index 7e5e871efc..f91b95acd6 100644 --- a/litellm/proxy/management_endpoints/fallback_management_endpoints.py +++ b/litellm/proxy/management_endpoints/fallback_management_endpoints.py @@ -110,9 +110,7 @@ async def create_fallback( if data.model in data.fallback_models: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, - detail={ - "error": f"Model '{data.model}' cannot be its own fallback" - }, + detail={"error": f"Model '{data.model}' cannot be its own fallback"}, ) # Check if we need to store in DB @@ -165,9 +163,7 @@ async def create_fallback( "param_name": "router_settings", "param_value": router_settings_json, }, - "update": { - "param_value": router_settings_json - }, + "update": {"param_value": router_settings_json}, }, ) @@ -346,18 +342,14 @@ async def delete_fallback( "param_name": "router_settings", "param_value": router_settings_json, }, - "update": { - "param_value": router_settings_json - }, + "update": {"param_value": router_settings_json}, }, ) # Update the in-memory router configuration setattr(llm_router, fallback_key, updated_fallbacks) - verbose_proxy_logger.info( - f"Fallback deleted: {model} (type: {fallback_type})" - ) + verbose_proxy_logger.info(f"Fallback deleted: {model} (type: {fallback_type})") return FallbackDeleteResponse( model=model, diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 80094c9abd..eb22a0ddc4 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -61,9 +61,9 @@ def _update_internal_new_user_params(data_json: dict, data: NewUserRequest) -> d auto_create_key = data_json.pop("auto_create_key", True) if auto_create_key is False: - data_json["table_name"] = ( - "user" # only create a user, don't create key if 'auto_create_key' set to False - ) + data_json[ + "table_name" + ] = "user" # only create a user, don't create key if 'auto_create_key' set to False if litellm.default_internal_user_params and ( data.user_role != LitellmUserRoles.PROXY_ADMIN.value @@ -142,7 +142,9 @@ async def _check_duplicate_user_field( error_label = label or field_name raise HTTPException( status_code=409, - detail={"error": f"User with {error_label} {existing_value} already exists"}, + detail={ + "error": f"User with {error_label} {existing_value} already exists" + }, ) @@ -415,18 +417,19 @@ async def new_user( status_code=403, detail="License is over limit. Please contact support@berri.ai to upgrade your license.", ) - + # Only proxy admins can create administrative users # Check if user_api_key_dict is actually a UserAPIKeyAuth instance (not a Depends object) # This can happen when the function is called directly in tests if ( - data.user_role in [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY] + data.user_role + in [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY] and isinstance(user_api_key_dict, UserAPIKeyAuth) and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN ): raise HTTPException( status_code=403, - detail=f"Only proxy admins can create administrative users (proxy_admin, proxy_admin_viewer). Attempted to create user with role: {data.user_role}. Your role: {user_api_key_dict.user_role}" + detail=f"Only proxy admins can create administrative users (proxy_admin, proxy_admin_viewer). Attempted to create user with role: {data.user_role}. Your role: {user_api_key_dict.user_role}", ) data_json = data.json() # type: ignore @@ -615,7 +618,9 @@ async def user_info( user_id is None and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN ): - return await _get_user_info_for_proxy_admin(user_api_key_dict=user_api_key_dict) + return await _get_user_info_for_proxy_admin( + user_api_key_dict=user_api_key_dict + ) elif user_id is None: user_id = user_api_key_dict.user_id ## GET USER ROW ## @@ -623,7 +628,7 @@ async def user_info( user_info = None if user_id is not None: user_info = await prisma_client.get_data(user_id=user_id) - + if user_info is None: raise HTTPException( status_code=404, @@ -755,11 +760,11 @@ async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth): _teams_in_db = [LiteLLM_TeamTable(**team) for team in _teams_in_db] _teams_in_db.sort(key=lambda x: (getattr(x, "team_alias", "") or "")) returned_keys = _process_keys_for_user_info(keys=keys_in_db, all_teams=_teams_in_db) - + # Get admin's own user_id and user_info admin_user_id = user_api_key_dict.user_id admin_user_info = None - + if admin_user_id is not None: admin_user_info = await prisma_client.get_data(user_id=admin_user_id) if admin_user_info is not None: @@ -768,7 +773,7 @@ async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth): if isinstance(admin_user_info, BaseModel) else admin_user_info ) - + return UserInfoResponse( user_id=admin_user_id, user_info=admin_user_info, @@ -801,11 +806,11 @@ def _process_keys_for_user_info( except Exception: # if using pydantic v1 _key = key.dict() - + # Filter out UI session tokens (team_id="litellm-dashboard") if _key.get("team_id") == UI_SESSION_TOKEN_TEAM_ID: continue - + if ( "team_id" in _key and _key["team_id"] is not None @@ -829,8 +834,8 @@ def _update_internal_user_params( data_json: dict, data: Union[UpdateUserRequest, UpdateUserRequestNoUserIDorEmail] ) -> dict: non_default_values = {} - fields_set = data.fields_set() if hasattr(data, 'fields_set') else set() - + fields_set = data.fields_set() if hasattr(data, "fields_set") else set() + for k, v in data_json.items(): if k == "max_budget": if "max_budget" in fields_set: @@ -867,9 +872,9 @@ def _update_internal_user_params( "budget_duration" not in non_default_values ): # applies internal user limits, if user role updated if is_internal_user and litellm.internal_user_budget_duration is not None: - non_default_values["budget_duration"] = ( - litellm.internal_user_budget_duration - ) + non_default_values[ + "budget_duration" + ] = litellm.internal_user_budget_duration from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time non_default_values["budget_reset_at"] = get_budget_reset_time( @@ -1489,7 +1494,9 @@ async def _authorize_user_list_request( if user_api_key_dict.user_id is None: raise HTTPException( status_code=403, - detail={"error": "Only proxy admins and organization admins can list users."}, + detail={ + "error": "Only proxy admins and organization admins can list users." + }, ) try: caller_user = await get_user_object( @@ -1502,12 +1509,16 @@ async def _authorize_user_list_request( except ValueError: raise HTTPException( status_code=403, - detail={"error": "Only proxy admins and organization admins can list users."}, + detail={ + "error": "Only proxy admins and organization admins can list users." + }, ) if caller_user is None: raise HTTPException( status_code=403, - detail={"error": "Only proxy admins and organization admins can list users."}, + detail={ + "error": "Only proxy admins and organization admins can list users." + }, ) allowed_org_ids = [ @@ -1518,17 +1529,23 @@ async def _authorize_user_list_request( if not allowed_org_ids: raise HTTPException( status_code=403, - detail={"error": "Only proxy admins and organization admins can list users."}, + detail={ + "error": "Only proxy admins and organization admins can list users." + }, ) # If client also sent organization_ids, intersect with allowed orgs if organization_ids: - requested = set(oid.strip() for oid in organization_ids.split(",") if oid.strip()) + requested = set( + oid.strip() for oid in organization_ids.split(",") if oid.strip() + ) intersection = list(requested & set(allowed_org_ids)) if not intersection: raise HTTPException( status_code=403, - detail={"error": "You do not have org_admin access to the requested organization(s)."}, + detail={ + "error": "You do not have org_admin access to the requested organization(s)." + }, ) allowed_org_ids = intersection @@ -1661,7 +1678,9 @@ async def get_users( } if organization_ids: - org_id_list = [oid.strip() for oid in organization_ids.split(",") if oid.strip()] + org_id_list = [ + oid.strip() for oid in organization_ids.split(",") if oid.strip() + ] if org_id_list: where_conditions["organization_memberships"] = { "some": {"organization_id": {"in": org_id_list}} @@ -1968,8 +1987,11 @@ async def _resolve_org_filter_for_user_search( if team_id is not None: return await _resolve_team_org_filter( - user_api_key_dict, team_id, prisma_client, - user_api_key_cache, proxy_logging_obj, + user_api_key_dict, + team_id, + prisma_client, + user_api_key_cache, + proxy_logging_obj, ) raise HTTPException( @@ -2110,13 +2132,13 @@ async def ui_view_users( } # Query users with pagination and filters - users: Optional[List[BaseModel]] = ( - await prisma_client.db.litellm_usertable.find_many( - where=where_conditions, - skip=skip, - take=page_size, - order={"created_at": "desc"}, - ) + users: Optional[ + List[BaseModel] + ] = await prisma_client.db.litellm_usertable.find_many( + where=where_conditions, + skip=skip, + take=page_size, + order={"created_at": "desc"}, ) if not users: diff --git a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py index a2a38cad14..e474cb7d15 100644 --- a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py +++ b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py @@ -1,4 +1,3 @@ - from fastapi import APIRouter, Depends, HTTPException, Query from litellm.proxy._types import ( diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 941b2c276d..eb4bb7f884 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -127,7 +127,10 @@ def _calculate_key_rotation_time(rotation_interval: str) -> datetime: def _set_key_rotation_fields( - data: dict, auto_rotate: bool, rotation_interval: Optional[str], existing_key_alias: Optional[str] = None + data: dict, + auto_rotate: bool, + rotation_interval: Optional[str], + existing_key_alias: Optional[str] = None, ) -> None: """ Helper function to set rotation fields in key data if auto_rotate is enabled. @@ -3103,7 +3106,10 @@ async def delete_verification_tokens( hashed_token = hash_token(cast(str, key)) user_api_key_cache.delete_cache(hashed_token) - return {"deleted_keys": deleted_tokens, "failed_tokens": failed_tokens}, _keys_being_deleted + return { + "deleted_keys": deleted_tokens, + "failed_tokens": failed_tokens, + }, _keys_being_deleted def _transform_verification_tokens_to_deleted_records( @@ -3206,7 +3212,7 @@ async def delete_key_aliases( ) -async def _rotate_master_key( # noqa: PLR0915 +async def _rotate_master_key( # noqa: PLR0915 prisma_client: PrismaClient, user_api_key_dict: UserAPIKeyAuth, current_master_key: str, @@ -3421,6 +3427,8 @@ async def _insert_deprecated_key( "Failed to insert deprecated key for grace period: %s", deprecated_err, ) + + async def _execute_virtual_key_regeneration( *, prisma_client: PrismaClient, @@ -3984,8 +3992,7 @@ def _get_member_team_ids_from_objects( team.team_id for team in team_objects if any( - member.user_id is not None - and member.user_id == user_api_key_dict.user_id + member.user_id is not None and member.user_id == user_api_key_dict.user_id for member in team.members_with_roles ) ] @@ -4268,9 +4275,7 @@ async def key_aliases( where_sql = " AND ".join(where_parts) - count_sql = ( - f'SELECT COUNT(*) AS count FROM "LiteLLM_VerificationToken" WHERE {where_sql}' - ) + count_sql = f'SELECT COUNT(*) AS count FROM "LiteLLM_VerificationToken" WHERE {where_sql}' count_rows = await prisma_client.db.query_raw(count_sql, *query_params) total_count = int(count_rows[0]["count"]) if count_rows else 0 @@ -4285,7 +4290,9 @@ async def key_aliases( f" LIMIT ${limit_idx} OFFSET ${offset_idx}" ) alias_rows = await prisma_client.db.query_raw(aliases_sql, *aliases_params) - aliases: List[str] = [row["key_alias"] for row in alias_rows if row.get("key_alias")] + aliases: List[str] = [ + row["key_alias"] for row in alias_rows if row.get("key_alias") + ] total_pages = -(-total_count // size) if total_count > 0 else 0 verbose_proxy_logger.debug( diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index a9ff61dab5..3e5b729cea 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -175,9 +175,7 @@ if MCP_AVAILABLE: _base_validate_and_normalize_mcp_server_payload(payload) _validate_mcp_server_name_fields(payload) - _VALID_MCP_REQUIRED_FIELDS: frozenset = frozenset( - NewMCPServerRequest.model_fields - ) + _VALID_MCP_REQUIRED_FIELDS: frozenset = frozenset(NewMCPServerRequest.model_fields) def _validate_mcp_required_fields(payload: Any) -> None: """Validate submission payload against admin-configured mcp_required_fields.""" @@ -426,11 +424,17 @@ if MCP_AVAILABLE: inherited_credentials["scopes"] = existing_server.scopes # AWS SigV4 fields if existing_server.aws_access_key_id: - inherited_credentials["aws_access_key_id"] = existing_server.aws_access_key_id + inherited_credentials[ + "aws_access_key_id" + ] = existing_server.aws_access_key_id if existing_server.aws_secret_access_key: - inherited_credentials["aws_secret_access_key"] = existing_server.aws_secret_access_key + inherited_credentials[ + "aws_secret_access_key" + ] = existing_server.aws_secret_access_key if existing_server.aws_session_token: - inherited_credentials["aws_session_token"] = existing_server.aws_session_token + inherited_credentials[ + "aws_session_token" + ] = existing_server.aws_session_token if existing_server.aws_region_name: inherited_credentials["aws_region_name"] = existing_server.aws_region_name if existing_server.aws_service_name: @@ -736,8 +740,7 @@ if MCP_AVAILABLE: check_db_only=True, ) user_in_team = any( - m.user_id is not None - and m.user_id == user_api_key_dict.user_id + m.user_id is not None and m.user_id == user_api_key_dict.user_id for m in team_obj.members_with_roles ) if not user_in_team: @@ -746,20 +749,26 @@ if MCP_AVAILABLE: detail="You do not have permission to view MCP servers for this team.", ) - redacted_mcp_servers = await _get_team_scoped_mcp_server_list(sanitized_team_id) + redacted_mcp_servers = await _get_team_scoped_mcp_server_list( + sanitized_team_id + ) else: user_mcp_management_mode = _get_user_mcp_management_mode() if user_mcp_management_mode == "view_all" and not is_restricted_virtual_key: - servers = await global_mcp_server_manager.get_all_mcp_servers_unfiltered() + servers = ( + await global_mcp_server_manager.get_all_mcp_servers_unfiltered() + ) redacted_mcp_servers = _redact_mcp_credentials_list(servers) else: auth_contexts = await build_effective_auth_contexts(user_api_key_dict) aggregated_servers: Dict[str, LiteLLM_MCPServerTable] = {} for auth_context in auth_contexts: - servers = await global_mcp_server_manager.get_all_allowed_mcp_servers( - user_api_key_auth=auth_context + servers = ( + await global_mcp_server_manager.get_all_allowed_mcp_servers( + user_api_key_auth=auth_context + ) ) for server in servers: if server.server_id not in aggregated_servers: @@ -788,8 +797,10 @@ if MCP_AVAILABLE: if getattr(s, "is_byok", False) ] if byok_server_ids: - cred_rows = await _byok_prisma_client.db.litellm_mcpusercredentials.find_many( - where={"user_id": user_id, "server_id": {"in": byok_server_ids}} + cred_rows = ( + await _byok_prisma_client.db.litellm_mcpusercredentials.find_many( + where={"user_id": user_id, "server_id": {"in": byok_server_ids}} + ) ) cred_set = {r.server_id for r in cred_rows} for server in redacted_mcp_servers: @@ -941,7 +952,9 @@ if MCP_AVAILABLE: ): raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, - detail={"error": "Admin access required to view MCP server submissions."}, + detail={ + "error": "Admin access required to view MCP server submissions." + }, ) prisma_client = get_prisma_client_or_throw( @@ -967,7 +980,9 @@ if MCP_AVAILABLE: if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, - detail={"error": "Admin access required to approve MCP server submissions."}, + detail={ + "error": "Admin access required to approve MCP server submissions." + }, ) prisma_client = get_prisma_client_or_throw( @@ -1013,7 +1028,9 @@ if MCP_AVAILABLE: if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, - detail={"error": "Admin access required to reject MCP server submissions."}, + detail={ + "error": "Admin access required to reject MCP server submissions." + }, ) prisma_client = get_prisma_client_or_throw( @@ -1078,8 +1095,11 @@ if MCP_AVAILABLE: client_ip = IPAddressUtils.get_mcp_client_ip(request) registry_server = global_mcp_server_manager.get_mcp_server_by_id(server_id) - if registry_server is not None and not global_mcp_server_manager._is_server_accessible_from_ip( - registry_server, client_ip + if ( + registry_server is not None + and not global_mcp_server_manager._is_server_accessible_from_ip( + registry_server, client_ip + ) ): registry_server = None if registry_server is None: @@ -1114,8 +1134,10 @@ if MCP_AVAILABLE: exists = does_mcp_server_exist(mcp_server_records, server_id) else: # Registry/config server: use same access logic as list endpoint - allowed_server_ids = await global_mcp_server_manager.get_allowed_mcp_servers( - user_api_key_dict + allowed_server_ids = ( + await global_mcp_server_manager.get_allowed_mcp_servers( + user_api_key_dict + ) ) exists = mcp_server.server_id in allowed_server_ids @@ -1313,10 +1335,9 @@ if MCP_AVAILABLE: global_mcp_server_manager, ) - server = ( - global_mcp_server_manager.get_mcp_server_by_id(server_id) - or global_mcp_server_manager.get_mcp_server_by_name(server_id) - ) + server = global_mcp_server_manager.get_mcp_server_by_id( + server_id + ) or global_mcp_server_manager.get_mcp_server_by_name(server_id) if server is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -1522,10 +1543,13 @@ if MCP_AVAILABLE: detail={"error": "User ID not found in token"}, ) if payload.save: - await store_user_credential(prisma_client, user_id, server_id, payload.credential) + await store_user_credential( + prisma_client, user_id, server_id, payload.credential + ) from litellm.proxy._experimental.mcp_server.server import ( _invalidate_byok_cred_cache, ) + _invalidate_byok_cred_cache(user_id, server_id) return MCPUserCredentialResponse(server_id=server_id, has_credential=True) # save=False: credential not persisted @@ -1559,6 +1583,7 @@ if MCP_AVAILABLE: from litellm.proxy._experimental.mcp_server.server import ( _invalidate_byok_cred_cache, ) + _invalidate_byok_cred_cache(user_id, server_id) return MCPUserCredentialResponse(server_id=server_id, has_credential=False) @@ -1637,7 +1662,9 @@ if MCP_AVAILABLE: # Only delete if the stored credential is actually an OAuth2 token. # This prevents accidentally deleting a BYOK credential if one exists # for the same (user_id, server_id) pair. - cred_to_delete = await get_user_oauth_credential(prisma_client, user_id, server_id) + cred_to_delete = await get_user_oauth_credential( + prisma_client, user_id, server_id + ) if cred_to_delete is not None: try: await delete_user_credential(prisma_client, user_id, server_id) @@ -1916,9 +1943,7 @@ if MCP_AVAILABLE: query: Optional[str] = Query( None, description="Search filter for server names and descriptions" ), - category: Optional[str] = Query( - None, description="Filter by category" - ), + category: Optional[str] = Query(None, description="Filter by category"), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -1952,15 +1977,11 @@ if MCP_AVAILABLE: # Apply category filter if category: - servers = [ - s for s in servers if s.get("category", "") == category - ] + servers = [s for s in servers if s.get("category", "") == category] # Extract unique categories from the full list (before filtering) all_servers = registry.get("servers", []) - categories = sorted( - set(s.get("category", "Other") for s in all_servers) - ) + categories = sorted(set(s.get("category", "Other") for s in all_servers)) return { "servers": servers, diff --git a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py index 000682bbf8..b05cfef576 100644 --- a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py @@ -31,19 +31,17 @@ from litellm.types.proxy.management_endpoints.model_management_endpoints import router = APIRouter() -def validate_models_exist( - model_names: List[str], llm_router -) -> Tuple[bool, List[str]]: +def validate_models_exist(model_names: List[str], llm_router) -> Tuple[bool, List[str]]: """ Validate that all requested model names exist in the router. Checks only exact model name matches. - + Returns: Tuple[bool, List[str]]: (all_valid, missing_models) """ if llm_router is None: return False, model_names - + router_model_names = set(llm_router.get_model_names()) missing = [m for m in model_names if m not in router_model_names] return (len(missing) == 0, missing) @@ -54,24 +52,24 @@ def add_access_group_to_deployment( ) -> Tuple[Dict[str, Any], bool]: """ Add an access group to a deployment's model_info. - + Args: model_info: The model_info dictionary from the deployment access_group: The access group name to add - + Returns: Tuple[Dict[str, Any], bool]: (updated_model_info, was_modified) """ access_groups = model_info.get("access_groups", []) - + # Check if access group already exists if access_group in access_groups: return model_info, False - + # Add the access group access_groups.append(access_group) model_info["access_groups"] = access_groups - + return model_info, True @@ -82,31 +80,29 @@ async def update_deployments_with_access_group( ) -> int: """ Update all deployments for the given model names to include the access group. - + Args: model_names: List of model names whose deployments should be updated access_group: The access group name to add prisma_client: Database client - + Returns: int: Number of deployments updated """ models_updated = 0 - + for model_name in model_names: - verbose_proxy_logger.debug( - f"Updating deployments for model_name: {model_name}" - ) - + verbose_proxy_logger.debug(f"Updating deployments for model_name: {model_name}") + # Get all deployments with this model_name deployments = await prisma_client.db.litellm_proxymodeltable.find_many( where={"model_name": model_name} ) - + verbose_proxy_logger.debug( f"Found {len(deployments)} deployments for model_name: {model_name}" ) - + # If no deployments found, this is a config model (not in DB) if len(deployments) == 0: raise HTTPException( @@ -115,29 +111,29 @@ async def update_deployments_with_access_group( "error": f"Can't find model '{model_name}' in Database. Access group management is only supported for database models." }, ) - + # Update each deployment for deployment in deployments: model_info = deployment.model_info or {} - + # Add access group using helper updated_model_info, was_modified = add_access_group_to_deployment( model_info=model_info, access_group=access_group, ) - + # Only update in DB if modified if was_modified: await prisma_client.db.litellm_proxymodeltable.update( where={"model_id": deployment.model_id}, data={"model_info": json.dumps(updated_model_info)}, ) - + models_updated += 1 verbose_proxy_logger.debug( f"Updated deployment {deployment.model_id} with access group: {access_group}" ) - + return models_updated @@ -155,9 +151,7 @@ async def update_specific_deployments_with_access_group( """ models_updated = 0 for model_id in model_ids: - verbose_proxy_logger.debug( - f"Updating specific deployment model_id: {model_id}" - ) + verbose_proxy_logger.debug(f"Updating specific deployment model_id: {model_id}") deployment = await prisma_client.db.litellm_proxymodeltable.find_unique( where={"model_id": model_id} ) @@ -190,24 +184,24 @@ def remove_access_group_from_deployment( ) -> Tuple[Dict[str, Any], bool]: """ Remove an access group from a deployment's model_info. - + Args: model_info: The model_info dictionary from the deployment access_group: The access group name to remove - + Returns: Tuple[Dict[str, Any], bool]: (updated_model_info, was_modified) """ access_groups = model_info.get("access_groups", []) - + # Check if access group exists if access_group not in access_groups: return model_info, False - + # Remove the access group access_groups.remove(access_group) model_info["access_groups"] = access_groups - + return model_info, True @@ -216,31 +210,31 @@ async def get_all_access_groups_from_db( ) -> Dict[str, AccessGroupInfo]: """ Get all access groups from the database. - + Returns: Dict[str, AccessGroupInfo]: Dictionary mapping access_group name to info """ # Get all deployments deployments = await prisma_client.db.litellm_proxymodeltable.find_many() - + # Build access group map access_group_map: Dict[str, Dict[str, Any]] = {} - + for deployment in deployments: model_info = deployment.model_info or {} access_groups = model_info.get("access_groups", []) model_name = deployment.model_name - + for access_group in access_groups: if access_group not in access_group_map: access_group_map[access_group] = { "model_names": set(), "deployment_count": 0, } - + access_group_map[access_group]["model_names"].add(model_name) access_group_map[access_group]["deployment_count"] += 1 - + # Convert to AccessGroupInfo objects result = {} for access_group, data in access_group_map.items(): @@ -249,7 +243,7 @@ async def get_all_access_groups_from_db( model_names=sorted(list(data["model_names"])), deployment_count=data["deployment_count"], ) - + return result @@ -295,18 +289,18 @@ async def create_model_group( llm_router, prisma_client, ) - + verbose_proxy_logger.debug( f"Creating access group: {data.access_group} with models: {data.model_names}" ) - + # Validation: Check if access_group is provided if not data.access_group or not data.access_group.strip(): raise HTTPException( status_code=400, detail={"error": "access_group is required and cannot be empty"}, ) - + # Validation: Check that at least one of model_names or model_ids is provided has_model_names = data.model_names and len(data.model_names) > 0 has_model_ids = data.model_ids and len(data.model_ids) > 0 @@ -314,7 +308,9 @@ async def create_model_group( if not has_model_names and not has_model_ids: raise HTTPException( status_code=400, - detail={"error": "Either model_names or model_ids must be provided and non-empty"}, + detail={ + "error": "Either model_names or model_ids must be provided and non-empty" + }, ) # If model_ids is provided, use it (more precise targeting) @@ -333,26 +329,28 @@ async def create_model_group( status_code=400, detail={"error": f"Model(s) not found: {', '.join(missing_models)}"}, ) - + # Check if database is connected if prisma_client is None: raise HTTPException( status_code=500, detail={"error": "Database not connected. Cannot create access group."}, ) - + try: # Check if access group already exists existing_access_groups = await get_all_access_groups_from_db( prisma_client=prisma_client ) - + if data.access_group in existing_access_groups: raise HTTPException( status_code=409, - detail={"error": f"Access group '{data.access_group}' already exists. Use PUT /access_group/{data.access_group}/update to modify it."}, + detail={ + "error": f"Access group '{data.access_group}' already exists. Use PUT /access_group/{data.access_group}/update to modify it." + }, ) - + # Update deployments using the appropriate method if use_model_ids: assert data.model_ids is not None @@ -368,20 +366,20 @@ async def create_model_group( access_group=data.access_group, prisma_client=prisma_client, ) - + await clear_cache() - + verbose_proxy_logger.info( f"Successfully created access group '{data.access_group}' with {models_updated} models updated" ) - + return NewModelGroupResponse( access_group=data.access_group, model_names=data.model_names, model_ids=data.model_ids, models_updated=models_updated, ) - + except HTTPException: raise except Exception as e: @@ -418,26 +416,26 @@ async def list_access_groups( - ListAccessGroupsResponse with all access groups """ from litellm.proxy.proxy_server import prisma_client - + if prisma_client is None: raise HTTPException( status_code=500, detail={"error": "Database not connected."}, ) - + try: access_groups_map = await get_all_access_groups_from_db( prisma_client=prisma_client ) - + # Sort by access group name access_groups_list = sorted( access_groups_map.values(), key=lambda x: x.access_group, ) - + return ListAccessGroupsResponse(access_groups=access_groups_list) - + except Exception as e: verbose_proxy_logger.exception(f"Error listing access groups: {str(e)}") raise HTTPException( @@ -475,26 +473,26 @@ async def get_access_group_info( - HTTPException 404: If access group not found """ from litellm.proxy.proxy_server import prisma_client - + if prisma_client is None: raise HTTPException( status_code=500, detail={"error": "Database not connected."}, ) - + try: access_groups_map = await get_all_access_groups_from_db( prisma_client=prisma_client ) - + if access_group not in access_groups_map: raise HTTPException( status_code=404, detail={"error": f"Access group '{access_group}' not found"}, ) - + return access_groups_map[access_group] - + except HTTPException: raise except Exception as e: @@ -547,17 +545,17 @@ async def update_access_group( - HTTPException 404: If access group not found """ from litellm.proxy.proxy_server import llm_router, prisma_client - + if prisma_client is None: raise HTTPException( status_code=500, detail={"error": "Database not connected."}, ) - + verbose_proxy_logger.debug( f"Updating access group: {access_group} with models: {data.model_names}" ) - + # Validation: Check that at least one of model_names or model_ids is provided has_model_names = data.model_names and len(data.model_names) > 0 has_model_ids = data.model_ids and len(data.model_ids) > 0 @@ -565,11 +563,13 @@ async def update_access_group( if not has_model_names and not has_model_ids: raise HTTPException( status_code=400, - detail={"error": "Either model_names or model_ids must be provided and non-empty"}, + detail={ + "error": "Either model_names or model_ids must be provided and non-empty" + }, ) use_model_ids = has_model_ids - + # Validation: Check if access group exists try: access_groups_map = await get_all_access_groups_from_db( @@ -587,7 +587,7 @@ async def update_access_group( status_code=500, detail={"error": f"Failed to check access group existence: {str(e)}"}, ) - + # Validation: Check if all new models exist (only if using model_names path) if not use_model_ids and has_model_names: assert data.model_names is not None @@ -601,26 +601,25 @@ async def update_access_group( status_code=400, detail={"error": f"Model(s) not found: {', '.join(missing_models)}"}, ) - + try: # Step 1: Remove access group from ALL DB deployments (skip config models) all_deployments = await prisma_client.db.litellm_proxymodeltable.find_many() - + for deployment in all_deployments: model_info = deployment.model_info or {} - updated_model_info, was_modified = remove_access_group_from_deployment( model_info=model_info, access_group=access_group, ) - + if was_modified: await prisma_client.db.litellm_proxymodeltable.update( where={"model_id": deployment.model_id}, data={"model_info": json.dumps(updated_model_info)}, ) - + # Step 2: Add access group using the appropriate method if use_model_ids: assert data.model_ids is not None @@ -636,21 +635,21 @@ async def update_access_group( access_group=access_group, prisma_client=prisma_client, ) - + # Clear cache and reload models to pick up the access group changes await clear_cache() - + verbose_proxy_logger.info( f"Successfully updated access group '{access_group}' with {models_updated} models updated" ) - + return NewModelGroupResponse( access_group=access_group, model_names=data.model_names, model_ids=data.model_ids, models_updated=models_updated, ) - + except HTTPException: raise except Exception as e: @@ -694,15 +693,15 @@ async def delete_access_group( - HTTPException 404: If access group not found """ from litellm.proxy.proxy_server import prisma_client - + if prisma_client is None: raise HTTPException( status_code=500, detail={"error": "Database not connected."}, ) - + verbose_proxy_logger.debug(f"Deleting access group: {access_group}") - + # Validation: Check if access group exists try: access_groups_map = await get_all_access_groups_from_db( @@ -720,40 +719,40 @@ async def delete_access_group( status_code=500, detail={"error": f"Failed to check access group existence: {str(e)}"}, ) - + try: # Remove access group from all DB deployments (skip config models) all_deployments = await prisma_client.db.litellm_proxymodeltable.find_many() models_updated = 0 - + for deployment in all_deployments: model_info = deployment.model_info or {} - + updated_model_info, was_modified = remove_access_group_from_deployment( model_info=model_info, access_group=access_group, ) - + if was_modified: await prisma_client.db.litellm_proxymodeltable.update( where={"model_id": deployment.model_id}, data={"model_info": json.dumps(updated_model_info)}, ) models_updated += 1 - + # Clear cache and reload models to pick up the access group changes await clear_cache() - + verbose_proxy_logger.info( f"Successfully deleted access group '{access_group}' from {models_updated} deployments" ) - + return DeleteModelGroupResponse( access_group=access_group, models_updated=models_updated, message=f"Access group '{access_group}' deleted successfully", ) - + except HTTPException: raise except Exception as e: @@ -764,4 +763,3 @@ async def delete_access_group( status_code=500, detail={"error": f"Failed to delete access group: {str(e)}"}, ) - diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 248b34c3df..e5af8e8755 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -375,7 +375,7 @@ async def _update_team_model_in_db( ) -> PrismaCompatibleUpdateDBModel: """ Handle team model updates with proper alias management. - + If patch_data contains a team_id: - Creates unique internal model_name and team alias - Adds model to team object @@ -383,36 +383,37 @@ async def _update_team_model_in_db( """ # Validate team_id if present in patch_data from litellm.proxy.proxy_server import premium_user - + await ModelManagementAuthChecks.allow_team_model_action( model_params=patch_data, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, premium_user=premium_user, ) - + patch_team_id = patch_data.model_info.team_id if patch_data.model_info else None - + # No team_id in patch, proceed with standard update if patch_team_id is None: return update_db_model(db_model=db_model, updated_patch=patch_data) - + # Determine public model name public_model_name = _get_public_model_name( patch_data=patch_data, db_model=db_model, ) - + # Ensure model_info exists and set team_public_model_name if patch_data.model_info is None: from litellm.types.router import ModelInfo + patch_data.model_info = ModelInfo() patch_data.model_info.team_public_model_name = public_model_name - + # Check if team assignment is new or changed db_team_id = db_model.model_info.team_id if db_model.model_info else None is_new_team_assignment = db_team_id != patch_team_id - + if is_new_team_assignment: await _setup_new_team_model_assignment( team_id=patch_team_id, @@ -428,7 +429,7 @@ async def _update_team_model_in_db( patch_data=patch_data, user_api_key_dict=user_api_key_dict, ) - + return update_db_model(db_model=db_model, updated_patch=patch_data) @@ -439,10 +440,10 @@ def _get_public_model_name( """Determine the public model name from patch or existing model.""" if patch_data.model_name: return patch_data.model_name - + if db_model.model_info and db_model.model_info.team_public_model_name: return db_model.model_info.team_public_model_name - + return db_model.model_name @@ -455,7 +456,7 @@ async def _setup_new_team_model_assignment( """Set up a new team model with unique name, alias, and team membership.""" unique_model_name = f"model_name_{team_id}_{uuid.uuid4()}" patch_data.model_name = unique_model_name - + await update_team( data=UpdateTeamRequest( team_id=team_id, @@ -464,7 +465,7 @@ async def _setup_new_team_model_assignment( user_api_key_dict=user_api_key_dict, http_request=Request(scope={"type": "http"}), ) - + await team_model_add( data=TeamModelAddRequest( team_id=team_id, @@ -484,11 +485,9 @@ async def _update_existing_team_model_assignment( ) -> None: """Update an existing team model if the public name changed.""" old_public_name = ( - db_model.model_info.team_public_model_name - if db_model.model_info - else None + db_model.model_info.team_public_model_name if db_model.model_info else None ) - + # Update alias only if public name changed if old_public_name and public_model_name != old_public_name: await update_team( @@ -499,7 +498,7 @@ async def _update_existing_team_model_assignment( user_api_key_dict=user_api_key_dict, http_request=Request(scope={"type": "http"}), ) - + # Keep existing unique model_name patch_data.model_name = None @@ -1330,16 +1329,15 @@ async def clear_cache(): ) return - try: # Only clear DB models, preserve config models verbose_proxy_logger.debug("Clearing only DB models, preserving config models") - + # Get current models and filter out DB models current_models = llm_router.model_list.copy() config_models = [] db_model_ids = [] - + for model in current_models: model_info = model.get("model_info", {}) if model_info.get("db_model", False): @@ -1348,20 +1346,22 @@ async def clear_cache(): else: # This is a config model, preserve it config_models.append(model) - + # Clear only DB models for model_id in db_model_ids: llm_router.delete_deployment(id=model_id) - + # Clear auto routers llm_router.auto_routers.clear() - + # Reload only DB models await proxy_config.add_deployment( prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj ) - - verbose_proxy_logger.debug(f"Cleared {len(db_model_ids)} DB models, preserved {len(config_models)} config models") + + verbose_proxy_logger.debug( + f"Cleared {len(db_model_ids)} DB models, preserved {len(config_models)} config models" + ) except Exception as e: verbose_proxy_logger.exception( f"Failed to clear cache and reload models. Due to error - {str(e)}" diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 103b2efcdd..edea0c79c9 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -175,12 +175,16 @@ async def new_organization( if data.max_budget is not None and data.max_budget < 0: raise HTTPException( status_code=400, - detail={"error": f"max_budget cannot be negative. Received: {data.max_budget}"} + detail={ + "error": f"max_budget cannot be negative. Received: {data.max_budget}" + }, ) if data.soft_budget is not None and data.soft_budget < 0: raise HTTPException( status_code=400, - detail={"error": f"soft_budget cannot be negative. Received: {data.soft_budget}"} + detail={ + "error": f"soft_budget cannot be negative. Received: {data.soft_budget}" + }, ) user_object_correct_type: Optional[LiteLLM_UserTable] = None @@ -297,7 +301,7 @@ async def get_organization_daily_activity( from litellm.proxy.proxy_server import ( prisma_client, ) - + if prisma_client is None: raise HTTPException( status_code=500, @@ -433,12 +437,16 @@ async def update_organization( if data.max_budget is not None and data.max_budget < 0: raise HTTPException( status_code=400, - detail={"error": f"max_budget cannot be negative. Received: {data.max_budget}"} + detail={ + "error": f"max_budget cannot be negative. Received: {data.max_budget}" + }, ) if data.soft_budget is not None and data.soft_budget < 0: raise HTTPException( status_code=400, - detail={"error": f"soft_budget cannot be negative. Received: {data.soft_budget}"} + detail={ + "error": f"soft_budget cannot be negative. Received: {data.soft_budget}" + }, ) if data.updated_by is None: @@ -675,13 +683,15 @@ async def list_organization( response = [] else: where_conditions["organization_id"] = org_id - response = await prisma_client.db.litellm_organizationtable.find_many( - where=where_conditions, - include={ - "litellm_budget_table": True, - "members": True, - "teams": True, - }, + response = ( + await prisma_client.db.litellm_organizationtable.find_many( + where=where_conditions, + include={ + "litellm_budget_table": True, + "members": True, + "teams": True, + }, + ) ) else: # Filter by membership and any additional filters @@ -716,20 +726,20 @@ async def info_organization(organization_id: str): if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) - response: Optional[LiteLLM_OrganizationTableWithMembers] = ( - await prisma_client.db.litellm_organizationtable.find_unique( - where={"organization_id": organization_id}, - include={ - "litellm_budget_table": True, - "members": { - "include": { - "user": True, - } - }, - "teams": True, - "object_permission": True, + response: Optional[ + LiteLLM_OrganizationTableWithMembers + ] = await prisma_client.db.litellm_organizationtable.find_unique( + where={"organization_id": organization_id}, + include={ + "litellm_budget_table": True, + "members": { + "include": { + "user": True, + } }, - ) + "teams": True, + "object_permission": True, + }, ) if response is None: @@ -1025,16 +1035,16 @@ async def organization_member_update( }, data={"budget_id": budget_id}, ) - final_organization_membership: Optional[BaseModel] = ( - await prisma_client.db.litellm_organizationmembership.find_unique( - where={ - "user_id_organization_id": { - "user_id": data.user_id, - "organization_id": data.organization_id, - } - }, - include={"litellm_budget_table": True}, - ) + final_organization_membership: Optional[ + BaseModel + ] = await prisma_client.db.litellm_organizationmembership.find_unique( + where={ + "user_id_organization_id": { + "user_id": data.user_id, + "organization_id": data.organization_id, + } + }, + include={"litellm_budget_table": True}, ) if final_organization_membership is None: diff --git a/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py b/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py index 040acd9222..0ae1143512 100644 --- a/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py +++ b/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py @@ -101,9 +101,7 @@ class AiPolicySuggester: template_descriptions = [] for t in templates: examples = t.get("example_sentences", []) - examples_str = ( - ", ".join(f'"{e}"' for e in examples) if examples else "none" - ) + examples_str = ", ".join(f'"{e}"' for e in examples) if examples else "none" entry = ( f"- ID: {t['id']}\n" f" Title: {t['title']}\n" @@ -121,9 +119,7 @@ class AiPolicySuggester: "Available templates:\n\n" + "\n\n".join(template_descriptions) ) - def _build_user_prompt( - self, attack_examples: List[str], description: str - ) -> str: + def _build_user_prompt(self, attack_examples: List[str], description: str) -> str: parts = [] filtered_examples = [e for e in attack_examples if e.strip()] if filtered_examples: diff --git a/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py index 4a42e493d3..57578d98b7 100644 --- a/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py +++ b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py @@ -12,8 +12,7 @@ All /policy management endpoints import copy import json import os -from typing import (TYPE_CHECKING, Any, AsyncIterator, List, Literal, Optional, - cast) +from typing import TYPE_CHECKING, Any, AsyncIterator, List, Literal, Optional, cast from fastapi import APIRouter, Depends, HTTPException, Request from fastapi.responses import Response, StreamingResponse @@ -21,34 +20,40 @@ from pydantic import BaseModel, Field from typing_extensions import TypedDict from litellm._logging import verbose_proxy_logger -from litellm.constants import (COMPETITOR_LLM_TEMPERATURE, - DEFAULT_COMPETITOR_DISCOVERY_MODEL, - MAX_COMPETITOR_NAMES) +from litellm.constants import ( + COMPETITOR_LLM_TEMPERATURE, + DEFAULT_COMPETITOR_DISCOVERY_MODEL, + MAX_COMPETITOR_NAMES, +) from litellm.integrations.custom_guardrail import CustomGuardrail -from litellm.llms.openai.chat.guardrail_translation.handler import \ - OpenAIChatCompletionsHandler +from litellm.llms.openai.chat.guardrail_translation.handler import ( + OpenAIChatCompletionsHandler, +) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.guardrails.guardrail_hooks.custom_code import ( - RESPONSE_REJECTION_GUARDRAIL_CODE, CustomCodeGuardrail) + RESPONSE_REJECTION_GUARDRAIL_CODE, + CustomCodeGuardrail, +) from litellm.proxy.guardrails.guardrail_registry import GuardrailRegistry from litellm.proxy.management_helpers.utils import management_endpoint_wrapper from litellm.proxy.policy_engine.policy_registry import get_policy_registry from litellm.proxy.policy_engine.policy_resolver import PolicyResolver -from litellm.types.proxy.policy_engine import (PolicyGuardrailsResponse, - PolicyInfoResponse, - PolicyListResponse, - PolicyMatchContext, - PolicyScopeResponse, - PolicySummaryItem, - PolicyTestResponse, - PolicyValidateRequest, - PolicyValidationResponse) +from litellm.types.proxy.policy_engine import ( + PolicyGuardrailsResponse, + PolicyInfoResponse, + PolicyListResponse, + PolicyMatchContext, + PolicyScopeResponse, + PolicySummaryItem, + PolicyTestResponse, + PolicyValidateRequest, + PolicyValidationResponse, +) from litellm.types.utils import GenericGuardrailAPIInputs, ModelResponse if TYPE_CHECKING: - from litellm.litellm_core_utils.litellm_logging import \ - Logging as LiteLLMLoggingObj + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj router = APIRouter() @@ -295,8 +300,7 @@ async def test_policies_and_guardrails( Use inputs for a single call (legacy). """ - from litellm.litellm_core_utils.litellm_logging import \ - Logging as LiteLLMLoggingObj + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.proxy_server import chat_completion, proxy_logging_obj from litellm.proxy.utils import handle_exception_on_proxy @@ -658,8 +662,7 @@ async def get_policy_templates( return _load_policy_templates_from_local_backup() try: - from litellm.llms.custom_httpx.http_handler import \ - get_async_httpx_client + from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.llms.custom_http import httpxSpecialProvider async_client = get_async_httpx_client( @@ -1151,8 +1154,9 @@ async def suggest_policy_templates( Calls an LLM with tool calling to match user requirements to available templates. """ - from litellm.proxy.management_endpoints.policy_endpoints.ai_policy_suggester import \ - AiPolicySuggester + from litellm.proxy.management_endpoints.policy_endpoints.ai_policy_suggester import ( + AiPolicySuggester, + ) templates = _load_policy_templates_from_local_backup() suggester = AiPolicySuggester() @@ -1222,8 +1226,9 @@ async def _test_guardrail_definitions( text: str, ) -> List[GuardrailTestResultEntry]: """Instantiate and run each guardrail definition against the text.""" - from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import \ - ContentFilterGuardrail + from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( + ContentFilterGuardrail, + ) results: List[GuardrailTestResultEntry] = [] diff --git a/litellm/proxy/management_endpoints/router_settings_endpoints.py b/litellm/proxy/management_endpoints/router_settings_endpoints.py index 4d4c41a3dc..c98c4620d9 100644 --- a/litellm/proxy/management_endpoints/router_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/router_settings_endpoints.py @@ -53,16 +53,16 @@ def _get_routing_strategies_from_router_class() -> List[str]: """ # Get the __init__ signature sig = inspect.signature(Router.__init__) - + # Get the routing_strategy parameter routing_strategy_param = sig.parameters.get("routing_strategy") - + if routing_strategy_param and routing_strategy_param.annotation: # Extract Literal values using get_args literal_values = get_args(routing_strategy_param.annotation) if literal_values: return list(literal_values) - + raise ValueError("Unable to extract routing strategies from Router class") @@ -77,31 +77,33 @@ async def get_router_settings( ): """ Get router configuration and available settings. - + Returns: - fields: List of all configurable router settings with their metadata (type, description, default, options) The routing_strategy field includes available options extracted from the Router class - current_values: Current values of router settings from config """ from litellm.proxy.proxy_server import llm_router, proxy_config - + try: # Get available routing strategies dynamically from Router class available_routing_strategies = _get_routing_strategies_from_router_class() - + # Get router settings fields from types file - router_fields = [field.model_copy(deep=True) for field in ROUTER_SETTINGS_FIELDS] - + router_fields = [ + field.model_copy(deep=True) for field in ROUTER_SETTINGS_FIELDS + ] + # Populate routing_strategy field with available options and descriptions for field in router_fields: if field.field_name == "routing_strategy": field.options = available_routing_strategies break - + # Try to get router settings from config config = await proxy_config.get_config() router_settings_from_config = config.get("router_settings", {}) - + # Get current values from llm_router if initialized current_values = {} if llm_router is not None: @@ -110,24 +112,22 @@ async def get_router_settings( if hasattr(llm_router, field.field_name): value = getattr(llm_router, field.field_name) current_values[field.field_name] = value - + # Merge with config values (config takes precedence) current_values.update(router_settings_from_config) - + # Update field values with current values for field in router_fields: if field.field_name in current_values: field.field_value = current_values[field.field_name] - + return RouterSettingsResponse( fields=router_fields, current_values=current_values, routing_strategy_descriptions=ROUTING_STRATEGY_DESCRIPTIONS, ) except Exception as e: - verbose_proxy_logger.error( - f"Error fetching router settings: {str(e)}" - ) + verbose_proxy_logger.error(f"Error fetching router settings: {str(e)}") raise @@ -142,11 +142,11 @@ async def get_router_fields( ): """ Get router settings field definitions without values. - + Returns only the field metadata (type, description, default, options) without populating field_value. This is useful for UI components that need to know what fields to render, but will get the actual values from a different endpoint. - + Returns: - fields: List of all configurable router settings with their metadata (type, description, default, options) The routing_strategy field includes available options extracted from the Router class @@ -156,27 +156,26 @@ async def get_router_fields( try: # Get available routing strategies dynamically from Router class available_routing_strategies = _get_routing_strategies_from_router_class() - + # Get router settings fields from types file - router_fields = [field.model_copy(deep=True) for field in ROUTER_SETTINGS_FIELDS] - + router_fields = [ + field.model_copy(deep=True) for field in ROUTER_SETTINGS_FIELDS + ] + # Populate routing_strategy field with available options for field in router_fields: if field.field_name == "routing_strategy": field.options = available_routing_strategies break - + # Ensure field_value is None for all fields (don't populate values) for field in router_fields: field.field_value = None - + return RouterFieldsResponse( fields=router_fields, routing_strategy_descriptions=ROUTING_STRATEGY_DESCRIPTIONS, ) except Exception as e: - verbose_proxy_logger.error( - f"Error fetching router fields: {str(e)}" - ) + verbose_proxy_logger.error(f"Error fetching router fields: {str(e)}") raise - diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index 73fcce72c3..2d657d96c1 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -210,18 +210,18 @@ def _build_scim_metadata( async def _get_scim_upsert_user_setting() -> bool: """ Get the scim_upsert_user setting from litellm_settings. - + Returns: True if scim_upsert_user is not set or is True (default behavior), False if scim_upsert_user is explicitly set to False (SCIM 2.0 strict mode) """ try: from litellm.proxy.proxy_server import proxy_config - + config = await proxy_config.get_config() litellm_settings = config.get("litellm_settings", {}) or {} scim_upsert_user = litellm_settings.get("scim_upsert_user", True) - + # Default to True if not set (backward compatibility) return bool(scim_upsert_user) except Exception as e: @@ -250,7 +250,7 @@ async def _extract_group_member_ids(group: SCIMGroup) -> GroupMemberExtractionRe existing_member_ids = [] created_users = [] all_member_ids = [] - + # Check the feature flag scim_upsert_user = await _get_scim_upsert_user_setting() @@ -262,9 +262,7 @@ async def _extract_group_member_ids(group: SCIMGroup) -> GroupMemberExtractionRe if not user_id or not user_id.strip(): raise HTTPException( status_code=400, - detail={ - "error": "Invalid member: user ID cannot be empty." - }, + detail={"error": "Invalid member: user ID cannot be empty."}, ) # Check if user exists @@ -293,7 +291,7 @@ async def _extract_group_member_ids(group: SCIMGroup) -> GroupMemberExtractionRe status_code=400, detail={ "error": f"User with ID '{user_id}' does not exist. " - "Please create the user first via POST /Users before adding to group." + "Please create the user first via POST /Users before adding to group." }, ) @@ -652,9 +650,7 @@ async def get_resource_type( """ Get a single ResourceType by ID per RFC 7644. """ - verbose_proxy_logger.debug( - "SCIM ResourceType request for id=%s", resource_type_id - ) + verbose_proxy_logger.debug("SCIM ResourceType request for id=%s", resource_type_id) base_url = str(request.base_url).rstrip("/") + "/scim/v2" resource_types = _get_resource_types(base_url) for rt in resource_types: @@ -769,13 +765,13 @@ async def get_users( where_conditions["user_email"] = email # Get users from database - users: List[LiteLLM_UserTable] = ( - await prisma_client.db.litellm_usertable.find_many( - where=where_conditions, - skip=(startIndex - 1), - take=count, - order={"created_at": "desc"}, - ) + users: List[ + LiteLLM_UserTable + ] = await prisma_client.db.litellm_usertable.find_many( + where=where_conditions, + skip=(startIndex - 1), + take=count, + order={"created_at": "desc"}, ) # Get total count for pagination @@ -1143,7 +1139,12 @@ def _apply_patch_ops( for name_key, name_val in val.items(): name_key_lower = name_key.lower() if name_key_lower in ("givenname", "familyname"): - _handle_name_update(f"name.{name_key_lower}", op_type, name_val, scim_metadata) + _handle_name_update( + f"name.{name_key_lower}", + op_type, + name_val, + scim_metadata, + ) continue if path == "displayname": @@ -1174,7 +1175,7 @@ async def patch_team_membership( ) -> bool: """ Add or remove user from teams - + Handles duplicate membership gracefully (idempotent operation). If a user is already in a team, that's fine - we don't treat it as an error. """ @@ -1588,9 +1589,7 @@ async def _process_group_patch_operations( if not member_id or not member_id.strip(): raise HTTPException( status_code=400, - detail={ - "error": "Invalid member: user ID cannot be empty." - }, + detail={"error": "Invalid member: user ID cannot be empty."}, ) user = await prisma_client.db.litellm_usertable.find_unique( @@ -1613,7 +1612,7 @@ async def _process_group_patch_operations( status_code=400, detail={ "error": f"User with ID '{member_id}' does not exist. " - "Please create the user first via POST /Users before adding to group." + "Please create the user first via POST /Users before adding to group." }, ) diff --git a/litellm/proxy/management_endpoints/sso/__init__.py b/litellm/proxy/management_endpoints/sso/__init__.py index 8144e7c53f..0f77e84cef 100644 --- a/litellm/proxy/management_endpoints/sso/__init__.py +++ b/litellm/proxy/management_endpoints/sso/__init__.py @@ -9,4 +9,3 @@ from litellm.proxy.management_endpoints.sso.custom_microsoft_sso import ( ) __all__ = ["CustomMicrosoftSSO"] - diff --git a/litellm/proxy/management_endpoints/sso/custom_microsoft_sso.py b/litellm/proxy/management_endpoints/sso/custom_microsoft_sso.py index 61b3a8231a..191212d6f0 100644 --- a/litellm/proxy/management_endpoints/sso/custom_microsoft_sso.py +++ b/litellm/proxy/management_endpoints/sso/custom_microsoft_sso.py @@ -75,7 +75,11 @@ class CustomMicrosoftSSO(MicrosoftSSO): custom_userinfo_endpoint or f"https://graph.microsoft.com/{self.version}/me" ) - if custom_authorization_endpoint or custom_token_endpoint or custom_userinfo_endpoint: + if ( + custom_authorization_endpoint + or custom_token_endpoint + or custom_userinfo_endpoint + ): verbose_proxy_logger.debug( f"Using custom Microsoft SSO endpoints - " f"authorization: {authorization_endpoint}, " @@ -88,4 +92,3 @@ class CustomMicrosoftSSO(MicrosoftSSO): token_endpoint=token_endpoint, userinfo_endpoint=userinfo_endpoint, ) - diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index b7714d3f86..8e3061c203 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -25,7 +25,6 @@ from litellm.proxy.management_endpoints.common_daily_activity import ( ) from litellm.proxy.management_helpers.utils import handle_budget_for_entity from litellm.types.tag_management import ( - LiteLLM_DailyTagSpendTable, TagConfig, TagDeleteRequest, TagInfoRequest, @@ -96,7 +95,7 @@ async def new_tag( - description: Optional[str] - Description of what this tag represents - models: List[str] - List of either 'model_id' or 'model_name' allowed for this tag - budget_id: Optional[str] - The id for a budget (tpm/rpm/max budget) for the tag - + ### IF NO BUDGET ID - CREATE ONE WITH THESE PARAMS ### - max_budget: Optional[float] - Max budget for tag - tpm_limit: Optional[int] - Max tpm limit for tag @@ -208,7 +207,7 @@ async def _add_tag_to_deployment(deployment: "Deployment", tag: str): if db_model is None: raise HTTPException( status_code=404, - detail=f"Model {deployment.model_info.id} not found in database" + detail=f"Model {deployment.model_info.id} not found in database", ) # Prisma returns litellm_params as dict (already parsed from JSON) @@ -252,7 +251,7 @@ async def update_tag( - description: Optional[str] - Updated description - models: List[str] - Updated list of allowed LLM models - budget_id: Optional[str] - The id for a budget to associate with the tag - + ### BUDGET UPDATE PARAMS ### - max_budget: Optional[float] - Max budget for tag - tpm_limit: Optional[int] - Max tpm limit for tag @@ -295,7 +294,7 @@ async def update_tag( "models": tag.models or [], "model_info": json.dumps(model_info), } - + # Add budget_id if it changed if budget_id != existing_tag.budget_id: update_data["budget_id"] = budget_id @@ -383,7 +382,10 @@ async def info_tag( } # Add budget info if available - if hasattr(tag_record, "litellm_budget_table") and tag_record.litellm_budget_table: + if ( + hasattr(tag_record, "litellm_budget_table") + and tag_record.litellm_budget_table + ): tag_dict["litellm_budget_table"] = tag_record.litellm_budget_table requested_tags[tag_record.tag_name] = tag_dict @@ -438,31 +440,39 @@ async def list_tags( } # Add budget info if available - if hasattr(tag_record, "litellm_budget_table") and tag_record.litellm_budget_table: + if ( + hasattr(tag_record, "litellm_budget_table") + and tag_record.litellm_budget_table + ): tag_dict["litellm_budget_table"] = tag_record.litellm_budget_table list_of_tags.append(tag_dict) ## QUERY DYNAMIC TAGS ## - dynamic_tags = await prisma_client.db.litellm_dailytagspend.find_many( - distinct=["tag"], + # Use group_by instead of find_many(distinct=["tag"]). + # Prisma's distinct fetches all columns for all rows and deduplicates + # in application code, which is extremely slow on large tables. + # See: https://www.prisma.io/docs/orm/prisma-client/queries/aggregation-grouping-summarizing#distinct-under-the-hood + dynamic_tag_rows = await prisma_client.db.litellm_dailytagspend.group_by( + by=["tag"], + where={"tag": {"not": None}}, + # The old find_many(distinct=...) returned arbitrary timestamps from + # whichever row Prisma happened to pick. MIN/MAX give more meaningful + # values: earliest appearance and most recent activity. + _min={"created_at": True}, + _max={"updated_at": True}, ) - dynamic_tags_list = [ - LiteLLM_DailyTagSpendTable(**dynamic_tag.model_dump()) - for dynamic_tag in dynamic_tags - ] - dynamic_tag_config = [ { - "name": tag.tag, + "name": row["tag"], "description": "This is just a spend tag that was passed dynamically in a request. It does not control any LLM models.", "models": None, - "created_at": tag.created_at.isoformat(), - "updated_at": tag.updated_at.isoformat(), + "created_at": row["_min"]["created_at"].isoformat(), + "updated_at": row["_max"]["updated_at"].isoformat(), } - for tag in dynamic_tags_list - if tag.tag not in stored_tag_names + for row in dynamic_tag_rows + if row["tag"] not in stored_tag_names ] return list_of_tags + dynamic_tag_config diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 633de86aa6..9c8e6f7282 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -759,19 +759,25 @@ async def new_team( # noqa: PLR0915 if data.max_budget is not None and data.max_budget < 0: raise HTTPException( status_code=400, - detail={"error": f"max_budget cannot be negative. Received: {data.max_budget}"} + detail={ + "error": f"max_budget cannot be negative. Received: {data.max_budget}" + }, ) if data.team_member_budget is not None and data.team_member_budget < 0: raise HTTPException( status_code=400, - detail={"error": f"team_member_budget cannot be negative. Received: {data.team_member_budget}"} + detail={ + "error": f"team_member_budget cannot be negative. Received: {data.team_member_budget}" + }, ) if data.soft_budget is not None and data.soft_budget < 0: raise HTTPException( status_code=400, - detail={"error": f"soft_budget cannot be negative. Received: {data.soft_budget}"} + detail={ + "error": f"soft_budget cannot be negative. Received: {data.soft_budget}" + }, ) - + if data.soft_budget is not None: if data.max_budget is not None: # If max_budget is set, soft_budget must be strictly lower than max_budget @@ -780,7 +786,7 @@ async def new_team( # noqa: PLR0915 status_code=400, detail={ "error": f"soft_budget ({data.soft_budget}) must be strictly lower than max_budget ({data.max_budget})" - } + }, ) # Check if license is over limit @@ -940,12 +946,16 @@ async def new_team( # noqa: PLR0915 complete_team_data.members_with_roles = [] complete_team_data_dict = complete_team_data.model_dump(exclude_none=True) - + # Serialize router_settings to JSON (matching key creation pattern) router_settings_value = getattr(data, "router_settings", None) - router_settings_json = safe_dumps(router_settings_value) if router_settings_value is not None else safe_dumps({}) + router_settings_json = ( + safe_dumps(router_settings_value) + if router_settings_value is not None + else safe_dumps({}) + ) complete_team_data_dict["router_settings"] = router_settings_json - + complete_team_data_dict = prisma_client.jsonify_team_object( db_data=complete_team_data_dict ) @@ -1121,7 +1131,9 @@ async def fetch_and_validate_organization( validate_team_org_change( team=LiteLLM_TeamTable(**existing_team_row.model_dump()), - organization=LiteLLM_OrganizationTableWithMembers(**organization_row.model_dump()), + organization=LiteLLM_OrganizationTableWithMembers( + **organization_row.model_dump() + ), llm_router=llm_router, ) @@ -1129,7 +1141,9 @@ async def fetch_and_validate_organization( def validate_team_org_change( - team: LiteLLM_TeamTable, organization: LiteLLM_OrganizationTableWithMembers, llm_router: Router + team: LiteLLM_TeamTable, + organization: LiteLLM_OrganizationTableWithMembers, + llm_router: Router, ) -> bool: """ Validate that a team can be moved to an organization. @@ -1180,7 +1194,9 @@ def validate_team_org_change( # Check if the team's user_id is a member of the org team_members = [m.user_id for m in team.members_with_roles] - org_members = [m.user_id for m in organization.members] if organization.members else [] + org_members = ( + [m.user_id for m in organization.members] if organization.members else [] + ) not_in_org = [ m for m in team_members @@ -1226,7 +1242,7 @@ def validate_team_org_change( "/team/update", tags=["team management"], dependencies=[Depends(user_api_key_auth)] ) @management_endpoint_wrapper -async def update_team( # noqa: PLR0915 +async def update_team( # noqa: PLR0915 data: UpdateTeamRequest, http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -1314,24 +1330,32 @@ async def update_team( # noqa: PLR0915 ) if data.team_id is None: - raise HTTPException(status_code=400, detail={"error": "No team id passed in"}) + raise HTTPException( + status_code=400, detail={"error": "No team id passed in"} + ) verbose_proxy_logger.debug("/team/update - %s", data) # Validate budget values are not negative if data.max_budget is not None and data.max_budget < 0: raise HTTPException( status_code=400, - detail={"error": f"max_budget cannot be negative. Received: {data.max_budget}"} + detail={ + "error": f"max_budget cannot be negative. Received: {data.max_budget}" + }, ) if data.team_member_budget is not None and data.team_member_budget < 0: raise HTTPException( status_code=400, - detail={"error": f"team_member_budget cannot be negative. Received: {data.team_member_budget}"} + detail={ + "error": f"team_member_budget cannot be negative. Received: {data.team_member_budget}" + }, ) if data.soft_budget is not None and data.soft_budget < 0: raise HTTPException( status_code=400, - detail={"error": f"soft_budget cannot be negative. Received: {data.soft_budget}"} + detail={ + "error": f"soft_budget cannot be negative. Received: {data.soft_budget}" + }, ) existing_team_row = await prisma_client.db.litellm_teamtable.find_unique( @@ -1343,28 +1367,38 @@ async def update_team( # noqa: PLR0915 status_code=404, detail={"error": f"Team not found, passed team_id={data.team_id}"}, ) - + if data.soft_budget is not None: - max_budget_to_check = data.max_budget if data.max_budget is not None else existing_team_row.max_budget + max_budget_to_check = ( + data.max_budget + if data.max_budget is not None + else existing_team_row.max_budget + ) if max_budget_to_check is not None: if data.soft_budget >= max_budget_to_check: raise HTTPException( status_code=400, detail={ "error": f"soft_budget ({data.soft_budget}) must be strictly lower than max_budget ({max_budget_to_check})" - } + }, ) - + if data.max_budget is not None: - existing_soft_budget = getattr(existing_team_row, 'soft_budget', None) - soft_budget_to_check = data.soft_budget if data.soft_budget is not None else existing_soft_budget - if soft_budget_to_check is not None and isinstance(soft_budget_to_check, (int, float)): + existing_soft_budget = getattr(existing_team_row, "soft_budget", None) + soft_budget_to_check = ( + data.soft_budget + if data.soft_budget is not None + else existing_soft_budget + ) + if soft_budget_to_check is not None and isinstance( + soft_budget_to_check, (int, float) + ): if data.max_budget <= soft_budget_to_check: raise HTTPException( status_code=400, detail={ "error": f"max_budget ({data.max_budget}) must be strictly greater than soft_budget ({soft_budget_to_check})" - } + }, ) if ( @@ -1465,16 +1499,19 @@ async def update_team( # noqa: PLR0915 updated_kv["model_id"] = _model_id # Serialize router_settings to JSON if present (matching key update pattern) - if "router_settings" in updated_kv and updated_kv["router_settings"] is not None: + if ( + "router_settings" in updated_kv + and updated_kv["router_settings"] is not None + ): updated_kv["router_settings"] = safe_dumps(updated_kv["router_settings"]) updated_kv = prisma_client.jsonify_team_object(db_data=updated_kv) - team_row: Optional[LiteLLM_TeamTable] = ( - await prisma_client.db.litellm_teamtable.update( - where={"team_id": data.team_id}, - data=updated_kv, - include={"litellm_model_table": True}, # type: ignore - ) + team_row: Optional[ + LiteLLM_TeamTable + ] = await prisma_client.db.litellm_teamtable.update( + where={"team_id": data.team_id}, + data=updated_kv, + include={"litellm_model_table": True}, # type: ignore ) if team_row is None or team_row.team_id is None: @@ -1483,7 +1520,9 @@ async def update_team( # noqa: PLR0915 detail={"error": "Team doesn't exist. Got={}".format(team_row)}, ) - verbose_proxy_logger.info("Successfully updated team - %s, info", team_row.team_id) + verbose_proxy_logger.info( + "Successfully updated team - %s, info", team_row.team_id + ) await _cache_team_object( team_id=team_row.team_id, team_table=LiteLLM_TeamTableCachedObj(**team_row.model_dump()), @@ -1834,14 +1873,14 @@ async def _validate_and_populate_member_user_info( ) -> Member: """ Validate and populate user_email/user_id for a member. - + Logic: 1. If both user_email and user_id are provided, verify they belong to the same user (use user_email as source of truth) 2. If only user_email is provided, populate user_id from DB 3. If only user_id is provided, populate user_email from DB (if user exists) 4. If only user_id is provided and doesn't exist, allow it to pass with user_email as None (will be upserted later) 5. If user_email and user_id mismatch, throw error - + Returns a Member with user_email and user_id populated (user_email may be None if only user_id provided and user doesn't exist). """ if member.user_email is None and member.user_id is None: @@ -1849,7 +1888,7 @@ async def _validate_and_populate_member_user_info( status_code=400, detail={"error": "Either user_id or user_email must be provided"}, ) - + # Case 1: Both user_email and user_id provided - verify they match if member.user_email is not None and member.user_id is not None: # Use user_email as source of truth @@ -1859,13 +1898,13 @@ async def _validate_and_populate_member_user_info( table_name="user", query_type="find_all", ) - + if users_by_email is None or ( isinstance(users_by_email, list) and len(users_by_email) == 0 ): # User doesn't exist yet - this is fine, will be created later return member - + if isinstance(users_by_email, list) and len(users_by_email) > 1: raise HTTPException( status_code=400, @@ -1873,10 +1912,10 @@ async def _validate_and_populate_member_user_info( "error": f"Multiple users found with email '{member.user_email}'. Please use 'user_id' instead." }, ) - + # Get the single user user_by_email = users_by_email[0] - + # Verify the user_id matches if user_by_email.user_id != member.user_id: raise HTTPException( @@ -1885,56 +1924,61 @@ async def _validate_and_populate_member_user_info( "error": f"user_email '{member.user_email}' and user_id '{member.user_id}' do not belong to the same user." }, ) - + # Both match, return as is return member - + # Case 2: Only user_email provided - populate user_id from DB if member.user_email is not None and member.user_id is None: user_by_email = await prisma_client.db.litellm_usertable.find_first( where={"user_email": {"equals": member.user_email, "mode": "insensitive"}} ) - + if user_by_email is None: # User doesn't exist yet - this is fine, will be created later return member - + # Check for multiple users with same email users_by_email = await prisma_client.get_data( key_val={"user_email": member.user_email}, table_name="user", query_type="find_all", ) - - if users_by_email and isinstance(users_by_email, list) and len(users_by_email) > 1: + + if ( + users_by_email + and isinstance(users_by_email, list) + and len(users_by_email) > 1 + ): raise HTTPException( status_code=400, detail={ "error": f"Multiple users found with email '{member.user_email}'. Please use 'user_id' instead." }, ) - + # Populate user_id member.user_id = user_by_email.user_id return member - + # Case 3: Only user_id provided - populate user_email from DB if user exists if member.user_id is not None and member.user_email is None: user_by_id = await prisma_client.db.litellm_usertable.find_unique( where={"user_id": member.user_id} ) - + if user_by_id is None: # User doesn't exist yet - allow it to pass with user_email as None # Will be upserted later with just user_id and null email return member - + # Populate user_email member.user_email = user_by_id.user_email return member - + return member + @router.post( "/team/member_add", tags=["team management"], @@ -2023,14 +2067,16 @@ async def team_member_add( prisma_client=prisma_client, ) - updated_team, updated_users, updated_team_memberships = ( - await _add_team_members_to_team( - data=data, - complete_team_data=complete_team_data, - prisma_client=prisma_client, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, - ) + ( + updated_team, + updated_users, + updated_team_memberships, + ) = await _add_team_members_to_team( + data=data, + complete_team_data=complete_team_data, + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, ) # Check if updated_team is None @@ -2212,15 +2258,15 @@ async def team_member_delete( ) # Fetch keys before deletion to persist them - keys_to_delete: List[LiteLLM_VerificationToken] = ( - await prisma_client.db.litellm_verificationtoken.find_many( - where={ - "user_id": {"in": list(user_ids_to_delete)}, - "team_id": data.team_id, - } - ) + keys_to_delete: List[ + LiteLLM_VerificationToken + ] = await prisma_client.db.litellm_verificationtoken.find_many( + where={ + "user_id": {"in": list(user_ids_to_delete)}, + "team_id": data.team_id, + } ) - + if keys_to_delete: await _persist_deleted_verification_tokens( keys=keys_to_delete, @@ -2602,10 +2648,10 @@ async def delete_team( team_rows: List[LiteLLM_TeamTable] = [] for team_id in data.team_ids: try: - team_row_base: Optional[BaseModel] = ( - await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id} - ) + team_row_base: Optional[ + BaseModel + ] = await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id} ) if team_row_base is None: raise Exception @@ -2664,10 +2710,10 @@ async def delete_team( _persist_deleted_verification_tokens, ) - keys_to_delete: List[LiteLLM_VerificationToken] = ( - await prisma_client.db.litellm_verificationtoken.find_many( - where={"team_id": {"in": data.team_ids}} - ) + keys_to_delete: List[ + LiteLLM_VerificationToken + ] = await prisma_client.db.litellm_verificationtoken.find_many( + where={"team_id": {"in": data.team_ids}} ) if keys_to_delete: @@ -2706,7 +2752,6 @@ async def delete_team( return deleted_teams - def _transform_teams_to_deleted_records( teams: List[LiteLLM_TeamTable], user_api_key_dict: UserAPIKeyAuth, @@ -2729,7 +2774,13 @@ def _transform_teams_to_deleted_records( ) record = deleted_record.model_dump() - for json_field in ["members_with_roles", "metadata", "model_spend", "model_max_budget", "router_settings"]: + for json_field in [ + "members_with_roles", + "metadata", + "model_spend", + "model_max_budget", + "router_settings", + ]: if json_field in record and record[json_field] is not None: record[json_field] = json.dumps(record[json_field]) @@ -2748,9 +2799,7 @@ async def _save_deleted_team_records( """Save deleted team records to the database.""" if not records: return - await prisma_client.db.litellm_deletedteamtable.create_many( - data=records - ) + await prisma_client.db.litellm_deletedteamtable.create_many(data=records) async def _persist_deleted_team_records( @@ -2770,6 +2819,7 @@ async def _persist_deleted_team_records( prisma_client=prisma_client, ) + async def validate_membership( user_api_key_dict: UserAPIKeyAuth, team_table: LiteLLM_TeamTable ): @@ -2806,9 +2856,7 @@ async def validate_membership( ) # Check direct team membership - if user_api_key_dict.user_id in [ - m.user_id for m in team_table.members_with_roles - ]: + if user_api_key_dict.user_id in [m.user_id for m in team_table.members_with_roles]: return # Check if user is an org admin for the team's organization @@ -2827,23 +2875,6 @@ async def validate_membership( ) -def _unfurl_all_proxy_models( - team_info: LiteLLM_TeamTable, llm_router: Router -) -> LiteLLM_TeamTable: - if ( - SpecialModelNames.all_proxy_models.value in team_info.models - and llm_router is not None - ): - team_models: set[str] = set() # make set to avoid duplicates - for model in team_info.models: - if model != SpecialModelNames.all_proxy_models.value: - team_models.add(model) - for model in llm_router.get_model_names(): - team_models.add(model) - team_info.models = list(team_models) - return team_info - - async def _add_team_member_budget_table( team_member_budget_id: str, prisma_client: PrismaClient, @@ -2902,11 +2933,11 @@ async def team_info( ) try: - team_info: Optional[BaseModel] = ( - await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id}, - include={"object_permission": True}, - ) + team_info: Optional[ + BaseModel + ] = await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id}, + include={"object_permission": True}, ) if team_info is None: raise Exception @@ -2972,9 +3003,6 @@ async def team_info( team_info_response_object=_team_info, ) - # ## UNFURL 'all-proxy-models' into the team_info.models list ## - # if llm_router is not None: - # _team_info = _unfurl_all_proxy_models(_team_info, llm_router) response_object = TeamInfoResponseObject( team_id=team_id, team_info=_team_info, @@ -3364,7 +3392,9 @@ async def list_team_v2( order=order_by if order_by else {"created_at": "desc"}, # Default sort ) # Get total count for pagination - total_count = await prisma_client.db.litellm_teamtable.count(where=where_conditions) + total_count = await prisma_client.db.litellm_teamtable.count( + where=where_conditions + ) # Calculate total pages total_pages = -(-total_count // page_size) # Ceiling division diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 84e945e888..084a528e8a 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -679,9 +679,9 @@ async def _setup_role_mappings() -> Optional["RoleMappings"]: import ast try: - generic_user_role_mappings_data: Dict[LitellmUserRoles, List[str]] = ( - ast.literal_eval(generic_role_mappings) - ) + generic_user_role_mappings_data: Dict[ + LitellmUserRoles, List[str] + ] = ast.literal_eval(generic_role_mappings) if isinstance(generic_user_role_mappings_data, dict): from litellm.types.proxy.management_endpoints.ui_sso import RoleMappings @@ -1017,9 +1017,9 @@ def apply_user_info_values_to_sso_user_defined_values( else: # SSO didn't provide a valid role, fall back to DB role or default if user_info is None or user_info.user_role is None: - user_defined_values["user_role"] = ( - LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value - ) + user_defined_values[ + "user_role" + ] = LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value verbose_proxy_logger.debug( "No SSO or DB role found, using default: INTERNAL_USER_VIEW_ONLY" ) @@ -1447,9 +1447,9 @@ async def insert_sso_user( if user_defined_values.get("max_budget") is None: user_defined_values["max_budget"] = litellm.max_internal_user_budget if user_defined_values.get("budget_duration") is None: - user_defined_values["budget_duration"] = ( - litellm.internal_user_budget_duration - ) + user_defined_values[ + "budget_duration" + ] = litellm.internal_user_budget_duration if user_defined_values["user_role"] is None: user_defined_values["user_role"] = LitellmUserRoles.INTERNAL_USER_VIEW_ONLY @@ -2550,9 +2550,9 @@ class MicrosoftSSOHandler: # if user is trying to get the raw sso response for debugging, return the raw sso response if return_raw_sso_response: - original_msft_result[MicrosoftSSOHandler.GRAPH_API_RESPONSE_KEY] = ( - user_team_ids - ) + original_msft_result[ + MicrosoftSSOHandler.GRAPH_API_RESPONSE_KEY + ] = user_team_ids original_msft_result["app_roles"] = app_roles return original_msft_result or {} @@ -2671,9 +2671,9 @@ class MicrosoftSSOHandler: # Fetch user membership from Microsoft Graph API all_group_ids = [] - next_link: Optional[str] = ( - MicrosoftSSOHandler.graph_api_user_groups_endpoint - ) + next_link: Optional[ + str + ] = MicrosoftSSOHandler.graph_api_user_groups_endpoint auth_headers = {"Authorization": f"Bearer {access_token}"} page_count = 0 diff --git a/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py b/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py index 55263dcd6b..9d3ecdba92 100644 --- a/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py +++ b/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py @@ -31,20 +31,25 @@ router = APIRouter() class TagActiveUsersResponse(BaseModel): """Response for tag active users metrics""" + tag: str active_users: int date: str # The specific date or period identifier - period_start: Optional[str] = None # For WAU/MAU, this will be the start of the period + period_start: Optional[ + str + ] = None # For WAU/MAU, this will be the start of the period period_end: Optional[str] = None # For WAU/MAU, this will be the end of the period class ActiveUsersAnalyticsResponse(BaseModel): """Response for active users analytics""" + results: List[TagActiveUsersResponse] class TagSummaryMetrics(BaseModel): """Summary metrics for a tag""" + tag: str unique_users: int total_requests: int @@ -56,22 +61,25 @@ class TagSummaryMetrics(BaseModel): class TagSummaryResponse(BaseModel): """Response for tag summary analytics""" + results: List[TagSummaryMetrics] class DistinctTagResponse(BaseModel): """Response for distinct user agent tags""" + tag: str class DistinctTagsResponse(BaseModel): """Response for all distinct user agent tags""" - results: List[DistinctTagResponse] + results: List[DistinctTagResponse] class PerUserMetrics(BaseModel): """Metrics for individual user""" + user_id: str user_email: Optional[str] = None user_agent: Optional[str] = None @@ -84,6 +92,7 @@ class PerUserMetrics(BaseModel): class PerUserAnalyticsResponse(BaseModel): """Response for per-user analytics""" + results: List[PerUserMetrics] total_count: int page: int @@ -102,21 +111,21 @@ async def get_distinct_user_agent_tags( ): """ Get all distinct user agent tags up to a maximum of {MAX_TAGS} tags. - + This endpoint returns all unique user agent tags found in the database, sorted by frequency of usage. - + Returns: DistinctTagsResponse: List of distinct user agent tags """ from litellm.proxy.proxy_server import prisma_client - + if prisma_client is None: raise HTTPException( status_code=500, detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - + try: sql_query = f""" SELECT @@ -128,16 +137,13 @@ async def get_distinct_user_agent_tags( ORDER BY usage_count DESC LIMIT {MAX_TAGS} """ - + db_response = await prisma_client.db.query_raw(sql_query) - - results = [ - DistinctTagResponse(tag=row["tag"]) - for row in db_response - ] - + + results = [DistinctTagResponse(tag=row["tag"]) for row in db_response] + return DistinctTagsResponse(results=results) - + except Exception as e: raise HTTPException( status_code=500, @@ -164,39 +170,44 @@ async def get_daily_active_users( ): """ Get Daily Active Users (DAU) by tags for the last {MAX_DAYS} days ending on UTC today + 1 day. - + This endpoint efficiently calculates unique users per tag for each of the last {MAX_DAYS} days using a single optimized SQL query, perfect for dashboard time series visualization. - + Args: tag_filter: Optional filter to specific tag (legacy) tag_filters: Optional filter to multiple specific tags (takes precedence over tag_filter) - + Returns: ActiveUsersAnalyticsResponse: DAU data by tag for each of the last {MAX_DAYS} days """ from litellm.proxy.proxy_server import prisma_client - + if prisma_client is None: raise HTTPException( status_code=500, detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - + try: # Calculate end_date as UTC today + 1 day from datetime import timezone - end_dt = datetime.now(timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=1) + + end_dt = datetime.now(timezone.utc).replace( + hour=0, minute=0, second=0, microsecond=0 + ) + timedelta(days=1) end_date = end_dt.strftime("%Y-%m-%d") - + # Calculate date range (last MAX_DAYS days) start_dt = end_dt - timedelta(days=MAX_DAYS) start_date = start_dt.strftime("%Y-%m-%d") - + # Build SQL query with optional tag filter(s) - where_clause = "WHERE dts.date >= $1 AND dts.date <= $2 AND vt.user_id IS NOT NULL" + where_clause = ( + "WHERE dts.date >= $1 AND dts.date <= $2 AND vt.user_id IS NOT NULL" + ) params = [start_date, end_date] - + # Handle multiple tag filters (takes precedence over single tag filter) if tag_filters and len(tag_filters) > 0: tag_conditions = [] @@ -208,7 +219,7 @@ async def get_daily_active_users( elif tag_filter: where_clause += " AND dts.tag ILIKE $3" params.append(f"%{tag_filter}%") - + sql_query = f""" SELECT dts.tag, @@ -220,20 +231,18 @@ async def get_daily_active_users( GROUP BY dts.tag, dts.date ORDER BY dts.date DESC, active_users DESC """ - + db_response = await prisma_client.db.query_raw(sql_query, *params) - + results = [ TagActiveUsersResponse( - tag=row["tag"], - active_users=row["active_users"], - date=row["date"] + tag=row["tag"], active_users=row["active_users"], date=row["date"] ) for row in db_response ] - + return ActiveUsersAnalyticsResponse(results=results) - + except Exception as e: raise HTTPException( status_code=500, @@ -260,44 +269,51 @@ async def get_weekly_active_users( ): """ Get Weekly Active Users (WAU) by tags for the last {MAX_WEEKS} weeks ending on UTC today + 1 day. - + Shows week-by-week breakdown: - Week 1 (Jan 1): Earliest week (7 weeks ago) - Week 2 (Jan 8): Next week (6 weeks ago) - Week 3 (Jan 15): Next week (5 weeks ago) - ... and so on for {MAX_WEEKS} weeks total - Week 7: Most recent week ending on UTC today + 1 day - + Args: tag_filter: Optional filter to specific tag (legacy) tag_filters: Optional filter to multiple specific tags (takes precedence over tag_filter) - + Returns: ActiveUsersAnalyticsResponse: WAU data by tag for each of the last {MAX_WEEKS} weeks with descriptive week labels (e.g., "Week 1 (Jan 1)") """ from litellm.proxy.proxy_server import prisma_client - + if prisma_client is None: raise HTTPException( status_code=500, detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - + try: # Calculate end_date as UTC today + 1 day from datetime import timezone - end_dt = datetime.now(timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=1) + + end_dt = datetime.now(timezone.utc).replace( + hour=0, minute=0, second=0, microsecond=0 + ) + timedelta(days=1) end_date = end_dt.strftime("%Y-%m-%d") - + # Calculate date range for all weeks (49 days total) # Start from 48 days before end_date to cover exactly MAX_WEEKS complete weeks - start_dt = end_dt - timedelta(days=(MAX_WEEKS * 7 - 1)) # MAX_WEEKS weeks * 7 days - 1 + start_dt = end_dt - timedelta( + days=(MAX_WEEKS * 7 - 1) + ) # MAX_WEEKS weeks * 7 days - 1 start_date = start_dt.strftime("%Y-%m-%d") - + # Build SQL query with optional tag filter(s) - where_clause = "WHERE dts.date >= $1 AND dts.date <= $2 AND vt.user_id IS NOT NULL" + where_clause = ( + "WHERE dts.date >= $1 AND dts.date <= $2 AND vt.user_id IS NOT NULL" + ) params = [start_date, end_date] - + # Handle multiple tag filters (takes precedence over single tag filter) if tag_filters and len(tag_filters) > 0: tag_conditions = [] @@ -309,7 +325,7 @@ async def get_weekly_active_users( elif tag_filter: where_clause += " AND dts.tag ILIKE $3" params.append(f"%{tag_filter}%") - + # Use window function to group by weeks with clear week numbering sql_query = f""" WITH weekly_data AS ( @@ -338,22 +354,24 @@ async def get_weekly_active_users( GROUP BY tag, week_offset ORDER BY week_offset DESC, active_users DESC """ - + db_response = await prisma_client.db.query_raw(sql_query, *params) - + results = [ TagActiveUsersResponse( tag=row["tag"], active_users=row["active_users"], - date=row["date"], # This will be "Week 1 (Jan 15)", "Week 2 (Jan 8)", etc. + date=row[ + "date" + ], # This will be "Week 1 (Jan 15)", "Week 2 (Jan 8)", etc. period_start=row["period_start"], - period_end=row["period_end"] + period_end=row["period_end"], ) for row in db_response ] - + return ActiveUsersAnalyticsResponse(results=results) - + except Exception as e: raise HTTPException( status_code=500, @@ -380,44 +398,51 @@ async def get_monthly_active_users( ): """ Get Monthly Active Users (MAU) by tags for the last {MAX_MONTHS} months ending on UTC today + 1 day. - + Shows month-by-month breakdown: - Month 1 (Nov): Earliest month (7 months ago, 30-day period) - Month 2 (Dec): Next month (6 months ago) - Month 3 (Jan): Next month (5 months ago) - ... and so on for {MAX_MONTHS} months total - Month 7: Most recent month ending on UTC today + 1 day - + Args: tag_filter: Optional filter to specific tag (legacy) tag_filters: Optional filter to multiple specific tags (takes precedence over tag_filter) - + Returns: ActiveUsersAnalyticsResponse: MAU data by tag for each of the last {MAX_MONTHS} months with descriptive month labels (e.g., "Month 1 (Nov)") """ from litellm.proxy.proxy_server import prisma_client - + if prisma_client is None: raise HTTPException( status_code=500, detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - + try: # Calculate end_date as UTC today + 1 day from datetime import timezone - end_dt = datetime.now(timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=1) + + end_dt = datetime.now(timezone.utc).replace( + hour=0, minute=0, second=0, microsecond=0 + ) + timedelta(days=1) end_date = end_dt.strftime("%Y-%m-%d") - + # Calculate date range for all months (210 days total) # Start from 209 days before end_date to cover exactly MAX_MONTHS complete months - start_dt = end_dt - timedelta(days=(MAX_MONTHS * 30 - 1)) # MAX_MONTHS months * 30 days - 1 + start_dt = end_dt - timedelta( + days=(MAX_MONTHS * 30 - 1) + ) # MAX_MONTHS months * 30 days - 1 start_date = start_dt.strftime("%Y-%m-%d") - + # Build SQL query with optional tag filter(s) - where_clause = "WHERE dts.date >= $1 AND dts.date <= $2 AND vt.user_id IS NOT NULL" + where_clause = ( + "WHERE dts.date >= $1 AND dts.date <= $2 AND vt.user_id IS NOT NULL" + ) params = [start_date, end_date] - + # Handle multiple tag filters (takes precedence over single tag filter) if tag_filters and len(tag_filters) > 0: tag_conditions = [] @@ -429,7 +454,7 @@ async def get_monthly_active_users( elif tag_filter: where_clause += " AND dts.tag ILIKE $3" params.append(f"%{tag_filter}%") - + # Use window function to group by months (30-day periods) with clear month numbering sql_query = f""" WITH monthly_data AS ( @@ -458,22 +483,22 @@ async def get_monthly_active_users( GROUP BY tag, month_offset ORDER BY month_offset DESC, active_users DESC """ - + db_response = await prisma_client.db.query_raw(sql_query, *params) - + results = [ TagActiveUsersResponse( tag=row["tag"], active_users=row["active_users"], date=row["date"], # This will be "Month 1 (Jan)", "Month 2 (Dec)", etc. period_start=row["period_start"], - period_end=row["period_end"] + period_end=row["period_end"], ) for row in db_response ] - + return ActiveUsersAnalyticsResponse(results=results) - + except Exception as e: raise HTTPException( status_code=500, @@ -488,12 +513,8 @@ async def get_monthly_active_users( dependencies=[Depends(user_api_key_auth)], ) async def get_tag_summary( - start_date: str = Query( - description="Start date in YYYY-MM-DD format" - ), - end_date: str = Query( - description="End date in YYYY-MM-DD format" - ), + start_date: str = Query(description="Start date in YYYY-MM-DD format"), + end_date: str = Query(description="End date in YYYY-MM-DD format"), tag_filter: Optional[str] = Query( default=None, description="Filter by specific tag (optional)", @@ -506,33 +527,33 @@ async def get_tag_summary( ): """ Get summary analytics for tags including unique users, requests, tokens, and spend. - + Args: start_date: Start date for the analytics period (YYYY-MM-DD) end_date: End date for the analytics period (YYYY-MM-DD) tag_filter: Optional filter to specific tag (legacy) tag_filters: Optional filter to multiple specific tags (takes precedence over tag_filter) - + Returns: TagSummaryResponse: Summary analytics data by tag """ from litellm.proxy.proxy_server import prisma_client - + if prisma_client is None: raise HTTPException( status_code=500, detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - + try: # Validate date format datetime.strptime(start_date, "%Y-%m-%d") datetime.strptime(end_date, "%Y-%m-%d") - + # Build SQL query with optional tag filter(s) where_clause = "WHERE dts.date >= $1 AND dts.date <= $2" params = [start_date, end_date] - + # Handle multiple tag filters (takes precedence over single tag filter) if tag_filters and len(tag_filters) > 0: tag_conditions = [] @@ -544,7 +565,7 @@ async def get_tag_summary( elif tag_filter: where_clause += " AND dts.tag ILIKE $3" params.append(f"%{tag_filter}%") - + sql_query = f""" SELECT dts.tag, @@ -560,9 +581,9 @@ async def get_tag_summary( GROUP BY dts.tag ORDER BY total_requests DESC """ - + db_response = await prisma_client.db.query_raw(sql_query, *params) - + results = [ TagSummaryMetrics( tag=row["tag"], @@ -571,13 +592,13 @@ async def get_tag_summary( successful_requests=int(row["successful_requests"] or 0), failed_requests=int(row["failed_requests"] or 0), total_tokens=int(row["total_tokens"] or 0), - total_spend=float(row["total_spend"] or 0.0) + total_spend=float(row["total_spend"] or 0.0), ) for row in db_response ] - + return TagSummaryResponse(results=results) - + except ValueError as e: raise HTTPException( status_code=400, @@ -606,63 +627,62 @@ async def get_per_user_analytics( description="Filter by multiple specific tags (optional, takes precedence over tag_filter)", ), page: int = Query(default=1, description="Page number for pagination", ge=1), - page_size: int = Query( - default=50, description="Items per page", ge=1, le=1000 - ), + page_size: int = Query(default=50, description="Items per page", ge=1, le=1000), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ Get per-user analytics including successful requests, tokens, and spend by individual users. - + This endpoint provides usage metrics broken down by individual users based on their tag activity during the last 30 days ending on UTC today + 1 day. - + Args: tag_filter: Optional filter to specific tag (legacy) tag_filters: Optional filter to multiple specific tags (takes precedence over tag_filter) page: Page number for pagination page_size: Number of items per page - + Returns: PerUserAnalyticsResponse: Analytics data broken down by individual users for the last 30 days """ from litellm.proxy.proxy_server import prisma_client - + if prisma_client is None: raise HTTPException( status_code=500, detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - + try: # Calculate end_date as UTC today + 1 day from datetime import timezone - end_dt = datetime.now(timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=1) + + end_dt = datetime.now(timezone.utc).replace( + hour=0, minute=0, second=0, microsecond=0 + ) + timedelta(days=1) end_date = end_dt.strftime("%Y-%m-%d") - + # Calculate date range (last 30 days) start_dt = end_dt - timedelta(days=30) start_date = start_dt.strftime("%Y-%m-%d") - + # Build where clause with date range - where_clause: Dict[str, Any] = { - "date": {"gte": start_date, "lte": end_date} - } - + where_clause: Dict[str, Any] = {"date": {"gte": start_date, "lte": end_date}} + # Add tag filtering if provided if tag_filters and len(tag_filters) > 0: where_clause["tag"] = {"in": tag_filters} elif tag_filter: where_clause["tag"] = {"contains": tag_filter} - + # Get all tag records in the date range with optional tag filtering tag_records = await prisma_client.db.litellm_dailytagspend.find_many( where=where_clause ) - + # Get unique api_keys api_keys = set(record.api_key for record in tag_records if record.api_key) - + if not api_keys: return PerUserAnalyticsResponse( results=[], @@ -671,31 +691,28 @@ async def get_per_user_analytics( page_size=page_size, total_pages=0, ) - + # Lookup user_id for each api_key api_key_records = await prisma_client.db.litellm_verificationtoken.find_many( where={"token": {"in": list(api_keys)}} ) - + # Create mapping from api_key to user_id api_key_to_user_id = { - record.token: record.user_id - for record in api_key_records - if record.user_id + record.token: record.user_id for record in api_key_records if record.user_id } - + # Get user emails for the user_ids user_ids = list(set(api_key_to_user_id.values())) user_records = await prisma_client.db.litellm_usertable.find_many( where={"user_id": {"in": user_ids}} ) - + # Create mapping from user_id to user_email user_id_to_email = { - record.user_id: record.user_email - for record in user_records + record.user_id: record.user_email for record in user_records } - + # Aggregate metrics by user user_metrics: Dict[str, PerUserMetrics] = {} @@ -703,42 +720,46 @@ async def get_per_user_analytics( if record.api_key in api_key_to_user_id: user_id = api_key_to_user_id[record.api_key] tag = record.tag # Use the full tag as user_agent - + if user_id not in user_metrics: user_metrics[user_id] = PerUserMetrics( user_id=user_id, user_email=user_id_to_email.get(user_id), - user_agent=tag + user_agent=tag, ) else: # If tag is different, keep the first one or prioritize certain ones if tag and not user_metrics[user_id].user_agent: user_metrics[user_id].user_agent = tag - + # Aggregate metrics - user_metrics[user_id].successful_requests += record.successful_requests or 0 + user_metrics[user_id].successful_requests += ( + record.successful_requests or 0 + ) user_metrics[user_id].failed_requests += record.failed_requests or 0 user_metrics[user_id].total_requests += record.api_requests or 0 # Calculate total_tokens from prompt_tokens + completion_tokens prompt_tokens = record.prompt_tokens or 0 completion_tokens = record.completion_tokens or 0 - user_metrics[user_id].total_tokens += int(prompt_tokens + completion_tokens) + user_metrics[user_id].total_tokens += int( + prompt_tokens + completion_tokens + ) user_metrics[user_id].spend += record.spend or 0.0 - + # Convert to list and sort by successful requests (descending) results = sorted( list(user_metrics.values()), key=lambda x: x.successful_requests, - reverse=True + reverse=True, ) - + # Apply pagination total_count = len(results) total_pages = (total_count + page_size - 1) // page_size start_idx = (page - 1) * page_size end_idx = start_idx + page_size paginated_results = results[start_idx:end_idx] - + return PerUserAnalyticsResponse( results=paginated_results, total_count=total_count, @@ -746,9 +767,9 @@ async def get_per_user_analytics( page_size=page_size, total_pages=total_pages, ) - + except Exception as e: raise HTTPException( status_code=500, detail=f"Failed to fetch per-user analytics: {str(e)}", - ) \ No newline at end of file + ) diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 0f426bf604..8aba8307b9 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -18,7 +18,6 @@ if TYPE_CHECKING: LiteLLM_ObjectPermissionTable, LiteLLM_TeamTableCachedObj, ) - async def attach_object_permission_to_dict( @@ -27,30 +26,32 @@ async def attach_object_permission_to_dict( ) -> Dict: """ Helper method to attach object_permission to a dictionary if object_permission_id is set. - + This function: 1. Checks if the dictionary has an object_permission_id 2. If found, queries the database for the corresponding object permission 3. Converts the object permission to a dictionary format 4. Attaches it to the input dictionary under the 'object_permission' key - + Args: data_dict: The dictionary to attach object_permission to prisma_client: The database client - + Returns: Dict: The input dictionary with object_permission attached if found - + Raises: ValueError: If prisma_client is None """ if prisma_client is None: raise ValueError("Prisma client not found") - + object_permission_id = data_dict.get("object_permission_id") if object_permission_id: - object_permission = await prisma_client.db.litellm_objectpermissiontable.find_unique( - where={"object_permission_id": object_permission_id}, + object_permission = ( + await prisma_client.db.litellm_objectpermissiontable.find_unique( + where={"object_permission_id": object_permission_id}, + ) ) if object_permission: # Convert to dict if needed @@ -168,21 +169,24 @@ async def _set_object_permission( if not isinstance(permission_data, dict): data_json.pop("object_permission") return data_json - + # Clean data: exclude None values and object_permission_id clean_data = { - k: v for k, v in permission_data.items() + k: v + for k, v in permission_data.items() if v is not None and k != "object_permission_id" } - + # Serialize mcp_tool_permissions to JSON string for GraphQL compatibility if "mcp_tool_permissions" in clean_data: - clean_data["mcp_tool_permissions"] = safe_dumps(clean_data["mcp_tool_permissions"]) - + clean_data["mcp_tool_permissions"] = safe_dumps( + clean_data["mcp_tool_permissions"] + ) + created_permission = await prisma_client.db.litellm_objectpermissiontable.create( data=clean_data ) - + data_json["object_permission_id"] = created_permission.object_permission_id data_json.pop("object_permission") return data_json @@ -204,10 +208,10 @@ async def _resolve_team_allowed_mcp_servers( ) direct_servers: List[str] = team_object_permission.mcp_servers or [] - access_group_servers: List[str] = ( - await MCPRequestHandler._get_mcp_servers_from_access_groups( - team_object_permission.mcp_access_groups or [] - ) + access_group_servers: List[ + str + ] = await MCPRequestHandler._get_mcp_servers_from_access_groups( + team_object_permission.mcp_access_groups or [] ) raw_tool_perms = team_object_permission.mcp_tool_permissions or {} if isinstance(raw_tool_perms, str): @@ -359,4 +363,4 @@ async def validate_key_mcp_servers_against_team( raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail={"error": detail}, - ) \ No newline at end of file + ) diff --git a/litellm/proxy/management_helpers/user_invitation.py b/litellm/proxy/management_helpers/user_invitation.py index 80f57fef2a..d2d800aa77 100644 --- a/litellm/proxy/management_helpers/user_invitation.py +++ b/litellm/proxy/management_helpers/user_invitation.py @@ -14,7 +14,7 @@ async def create_invitation_for_user( Create an invitation for the user to onboard to LiteLLM Admin UI. """ from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client - + if prisma_client is None: raise HTTPException( status_code=400, @@ -44,4 +44,4 @@ async def create_invitation_for_user( "error": "User id does not exist in 'LiteLLM_UserTable'. Fix this by creating user via `/user/new`." }, ) - raise HTTPException(status_code=500, detail={"error": str(e)}) \ No newline at end of file + raise HTTPException(status_code=500, detail={"error": str(e)}) diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index 67a1ea659f..7d485fdebd 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -64,19 +64,19 @@ async def handle_budget_for_entity( ) -> Optional[str]: """ Common helper to handle budget creation/updates for entities (organizations, tags, etc). - + This function: 1. Creates a new budget if budget_id is None but budget fields are provided 2. Updates an existing budget if budget fields are provided and budget_id exists 3. Returns the budget_id to use (existing or newly created) - + Args: data: The request object (e.g., TagNewRequest, NewOrganizationRequest, etc.) containing budget fields existing_budget_id: The existing budget_id if updating an entity, None if creating new user_api_key_dict: User authentication info prisma_client: Database client litellm_proxy_admin_name: Admin name for audit trail - + Returns: Optional[str]: The budget_id to use, or None if no budget was created/updated """ @@ -88,7 +88,9 @@ async def handle_budget_for_entity( budget_params = LiteLLM_BudgetTable.model_fields.keys() # Extract budget fields from data - _json_data = data.model_dump(exclude_none=True) if hasattr(data, "model_dump") else data + _json_data = ( + data.model_dump(exclude_none=True) if hasattr(data, "model_dump") else data + ) _budget_data = {k: v for k, v in _json_data.items() if k in budget_params} # Check if budget_id is explicitly provided in the data diff --git a/litellm/proxy/ocr_endpoints/__init__.py b/litellm/proxy/ocr_endpoints/__init__.py index 3488912f66..1a5b0ecbd7 100644 --- a/litellm/proxy/ocr_endpoints/__init__.py +++ b/litellm/proxy/ocr_endpoints/__init__.py @@ -1,2 +1 @@ # OCR Endpoints - diff --git a/litellm/proxy/openai_evals_endpoints/endpoints.py b/litellm/proxy/openai_evals_endpoints/endpoints.py index d87fe01add..98d409adf5 100644 --- a/litellm/proxy/openai_evals_endpoints/endpoints.py +++ b/litellm/proxy/openai_evals_endpoints/endpoints.py @@ -593,6 +593,7 @@ async def cancel_eval( version=version, ) + # =================================== # Run API Endpoints # =================================== @@ -659,7 +660,11 @@ async def create_run( request.headers.get("x-litellm-model") or request.query_params.get("model") or data.get("model") - or (data.get("completion", {}).get("model") if isinstance(data.get("completion"), dict) else None) + or ( + data.get("completion", {}).get("model") + if isinstance(data.get("completion"), dict) + else None + ) ) if model: data["model"] = model @@ -753,9 +758,7 @@ async def list_runs( } # Extract model for routing (header > query) - model = request.headers.get("x-litellm-model") or request.query_params.get( - "model" - ) + model = request.headers.get("x-litellm-model") or request.query_params.get("model") if model: data["model"] = model @@ -842,9 +845,7 @@ async def get_run( } # Extract model for routing (header > query) - model = request.headers.get("x-litellm-model") or request.query_params.get( - "model" - ) + model = request.headers.get("x-litellm-model") or request.query_params.get("model") if model: data["model"] = model diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 343ea11967..5d546733b7 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -153,7 +153,7 @@ def decode_model_from_file_id(encoded_id: str) -> Optional[str]: try: if not isinstance(encoded_id, str): return None - + # Remove prefix if present (file-, batch_, etc.) if encoded_id.startswith("file-"): b64_part = encoded_id[5:] # Remove "file-" @@ -161,14 +161,14 @@ def decode_model_from_file_id(encoded_id: str) -> Optional[str]: b64_part = encoded_id[6:] # Remove "batch_" else: b64_part = encoded_id - + padded = b64_part + "=" * (-len(b64_part) % 4) - decoded = base64.urlsafe_b64decode(padded).decode() + decoded = base64.urlsafe_b64decode(padded).decode() if decoded.startswith("litellm:") and ";model," in decoded: match = re.search(r";model,([^;]+)", decoded) if match: return match.group(1).strip() - + return None except Exception: return None @@ -182,7 +182,7 @@ def get_original_file_id(encoded_id: str) -> str: try: if not isinstance(encoded_id, str): return encoded_id - + # Remove prefix if present (file-, batch_, etc.) if encoded_id.startswith("file-"): b64_part = encoded_id[5:] # Remove "file-" @@ -190,15 +190,15 @@ def get_original_file_id(encoded_id: str) -> str: b64_part = encoded_id[6:] # Remove "batch_" else: b64_part = encoded_id - + padded = b64_part + "=" * (-len(b64_part) % 4) decoded = base64.urlsafe_b64decode(padded).decode() - + if decoded.startswith("litellm:") and ";model," in decoded: match = re.search(r"litellm:([^;]+);model,", decoded) if match: return match.group(1) - + return encoded_id except Exception: return encoded_id @@ -227,12 +227,12 @@ def extract_model_from_sources( 2. Request headers (x-litellm-model) 3. Query parameters (?model=) 4. Request body/data dict - + Args: file_id: File ID that may contain embedded model info request: FastAPI request object data: Optional request data dictionary - + Returns: Tuple of (model_from_id, model_from_param) - model_from_id: Model decoded from file ID (if embedded) @@ -240,17 +240,17 @@ def extract_model_from_sources( """ if data is None: data = {} - + # Check if file_id has embedded model info model_from_id = decode_model_from_file_id(file_id) - + # Check other sources for model parameter model_from_param = ( data.get("model") or request.query_params.get("model") or request.headers.get("x-litellm-model") ) - + return model_from_id, model_from_param @@ -261,28 +261,28 @@ def get_credentials_for_model( ): """ Retrieve API credentials for a model from the LLM Router. - + Args: llm_router: LiteLLM Router instance model_id: Model name or deployment ID operation_context: Description for error messages (e.g., "file upload", "batch creation") - + Returns: Dictionary with credentials (api_key, api_base, custom_llm_provider, etc.) - + Raises: HTTPException: If router not initialized or model not found """ from fastapi import HTTPException - + if llm_router is None: raise HTTPException( status_code=500, detail={"error": "Router not initialized. Cannot use model-based routing."}, ) - + credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id) - + if credentials is None: raise HTTPException( status_code=400, @@ -290,7 +290,7 @@ def get_credentials_for_model( "error": f"Model '{model_id}' not found in model_list. Please check your config.yaml." }, ) - + return credentials @@ -301,7 +301,7 @@ def prepare_data_with_credentials( ) -> None: """ Update data dictionary with model credentials (in-place). - + Args: data: Data dictionary to update credentials: Credentials from router @@ -309,7 +309,7 @@ def prepare_data_with_credentials( """ data.update(credentials) data.pop("custom_llm_provider", None) - + if file_id is not None: data["file_id"] = file_id @@ -323,21 +323,21 @@ def handle_model_based_routing( ) -> tuple[bool, Optional[str], Optional[str], Optional[dict]]: """ Orchestrate model-based credential routing for file operations. - + Args: file_id: File ID (may contain embedded model info) request: FastAPI request object llm_router: LiteLLM Router instance data: Request data dictionary check_file_id_encoding: Whether to check for embedded model in file_id - + Returns: Tuple of (should_use_model_routing, model_used, original_file_id, credentials) - should_use_model_routing: True if model-based routing should be used - model_used: The model name being used - original_file_id: Decoded file ID (if it was encoded) - credentials: Model credentials dict - + Raises: HTTPException: If router unavailable or model not found """ @@ -346,7 +346,7 @@ def handle_model_based_routing( request=request, data=data, ) - + # Priority 1: Model embedded in file_id if check_file_id_encoding and model_from_id is not None: credentials = get_credentials_for_model( @@ -356,7 +356,7 @@ def handle_model_based_routing( ) original_file_id = get_original_file_id(file_id) return True, model_from_id, original_file_id, credentials - + # Priority 2: Model from header/query/body elif model_from_param is not None: credentials = get_credentials_for_model( @@ -365,7 +365,7 @@ def handle_model_based_routing( operation_context="file operation", ) return True, model_from_param, None, credentials - + # No model-based routing needed return False, None, None, None @@ -433,24 +433,24 @@ EXTENSION_TO_MIME_TYPE = { def detect_content_type_from_filename(filename: str) -> str: """ Detect content type from filename using extension. - + Uses Python's mimetypes module with custom overrides for common cases. Normalizes jpg to jpeg for consistency. """ if not filename: return "application/octet-stream" - + # Try custom mapping first filename_lower = filename.lower() for ext, mime_type in EXTENSION_TO_MIME_TYPE.items(): if filename_lower.endswith(ext): return mime_type - + # Fall back to Python's mimetypes mime_type_guess, _ = mimetypes.guess_type(filename) if mime_type_guess is not None: return mime_type_guess - + return "application/octet-stream" @@ -459,44 +459,44 @@ def normalize_mime_type_for_provider( ) -> str: """ Normalize MIME type for specific provider requirements. - + Currently handles: - Gemini: Normalizes image/jpg to image/jpeg - + Args: mime_type: Original MIME type provider: Provider name (e.g., "gemini", "vertex_ai") - + Returns: str: Normalized MIME type """ normalized = mime_type.lower().strip() - + # Gemini/Vertex AI requires image/jpeg, not image/jpg if provider and ("gemini" in provider.lower() or "vertex_ai" in provider.lower()): if normalized == "image/jpg": normalized = "image/jpeg" - + # General normalization: always normalize jpg to jpeg if normalized == "image/jpg": normalized = "image/jpeg" - + return normalized def is_gemini_supported_mime_type(mime_type: str) -> bool: """ Check if a MIME type is supported by Gemini multimodal models. - + Supported categories: - Images: image/png, image/jpeg, image/webp - Video: 3gpp, wmv, webm, mp4, mpg, mpegps, mpeg, quicktime, x-flv - Audio: webm, wav, pcm, opus, mp4, mpga, mpeg, m4a, mp3, flac, aac - Documents: text/plain, application/pdf - + Args: mime_type: MIME type to check - + Returns: bool: True if supported, False otherwise """ @@ -512,35 +512,36 @@ def is_gemini_supported_mime_type(mime_type: str) -> bool: def get_content_type_from_file_object(file_object: Optional[dict]) -> str: """ Determine content type from file object (from database or API response). - + Extracts filename from file object and uses detect_content_type_from_filename. Falls back to default if file object is invalid or filename not found. - + Args: file_object: File object dictionary (can be None) - + Returns: str: MIME type (defaults to "application/octet-stream" if cannot be determined) """ if not file_object: return "application/octet-stream" - + # Handle JSON string if isinstance(file_object, str): import json + try: file_object = json.loads(file_object) except json.JSONDecodeError: return "application/octet-stream" - + if not isinstance(file_object, dict): return "application/octet-stream" - + # Try to get filename filename = file_object.get("filename", "") if filename: return detect_content_type_from_filename(filename) - + return "application/octet-stream" @@ -553,28 +554,30 @@ def get_content_type_from_file_object(file_object: Optional[dict]) -> str: class FileCreationParams: """ Structured parameters extracted from file creation requests. - + Attributes: target_storage: Storage backend name (e.g., "azure_storage", "default") target_model_names: List of model names for managed files model: Model parameter for multi-account routing """ - + target_storage: str = "default" target_model_names: List[str] = field(default_factory=list) model: Optional[str] = None - + def __post_init__(self): """Normalize and validate parameters after initialization.""" if self.target_model_names is None: self.target_model_names = [] - + # Normalize target_storage if not self.target_storage: self.target_storage = "default" - + # Strip whitespace from model names - self.target_model_names = [name.strip() for name in self.target_model_names if name.strip()] + self.target_model_names = [ + name.strip() for name in self.target_model_names if name.strip() + ] async def extract_file_creation_params( @@ -585,30 +588,30 @@ async def extract_file_creation_params( ) -> FileCreationParams: """ Extract file creation parameters from request. - + Args: request: FastAPI request object request_body: Optional pre-parsed request body target_model_names_form: target_model_names from form field (comma-separated string) target_storage_form: target_storage from form field (defaults to "default") - + Returns: FileCreationParams: Structured parameters extracted from the request """ from litellm.proxy.common_utils.http_parsing_utils import _read_request_body - + if request_body is None: request_body = await _read_request_body(request=request) or {} - + # Extract target_storage (simplified - just use form parameter) target_storage = _extract_target_storage_simple(target_storage_form) - + # Extract target_model_names (simplified - just use form parameter) target_model_names = _extract_target_model_names_simple(target_model_names_form) - + # Extract model parameter model = _extract_model_param(request, request_body) - + return FileCreationParams( target_storage=target_storage, target_model_names=target_model_names, @@ -619,10 +622,10 @@ async def extract_file_creation_params( def _extract_target_storage_simple(target_storage_form: Optional[str] = None) -> str: """ Extract target_storage parameter from form field. - + Args: target_storage_form: target_storage from form field - + Returns: str: Target storage backend name, or "default" """ @@ -631,26 +634,30 @@ def _extract_target_storage_simple(target_storage_form: Optional[str] = None) -> return "default" -def _extract_target_model_names_simple(target_model_names_form: Optional[str] = None) -> List[str]: +def _extract_target_model_names_simple( + target_model_names_form: Optional[str] = None, +) -> List[str]: """ Extract target_model_names parameter from form field. """ if not target_model_names_form: return [] - + # Parse comma-separated string into list if isinstance(target_model_names_form, str): - return [name.strip() for name in target_model_names_form.split(",") if name.strip()] + return [ + name.strip() for name in target_model_names_form.split(",") if name.strip() + ] elif isinstance(target_model_names_form, list): return [str(name).strip() for name in target_model_names_form if name] - + return [] def _extract_model_param(request: "Request", request_body: dict) -> Optional[str]: """ Extract model parameter from request. - + Priority: 1. request_body.model 2. Query parameter (?model=) @@ -699,14 +706,14 @@ async def get_batch_from_database( ): """ Try to retrieve batch object from ManagedObjectTable for consistent state. - + Args: batch_id: The batch ID (may be unified/encoded) unified_batch_id: Result from _is_base64_encoded_unified_file_id() managed_files_obj: The managed_files proxy hook object prisma_client: Prisma database client verbose_proxy_logger: Logger instance - + Returns: Tuple of (db_batch_object, response_batch) - db_batch_object: Raw database object (or None) @@ -714,35 +721,39 @@ async def get_batch_from_database( """ import json from litellm.types.utils import LiteLLMBatch - + if managed_files_obj is None or not unified_batch_id: return None, None - + try: if not prisma_client: return None, None - + db_batch_object = await prisma_client.db.litellm_managedobjecttable.find_first( where={"unified_object_id": batch_id} ) - + if not db_batch_object or not db_batch_object.file_object: return None, None - + # Parse the batch object from database - batch_data = json.loads(db_batch_object.file_object) if isinstance(db_batch_object.file_object, str) else db_batch_object.file_object + batch_data = ( + json.loads(db_batch_object.file_object) + if isinstance(db_batch_object.file_object, str) + else db_batch_object.file_object + ) response = LiteLLMBatch(**batch_data) response.id = batch_id # The stored batch object has the raw provider input_file_id. Resolve to unified ID. await resolve_input_file_id_to_unified(response, prisma_client) - + verbose_proxy_logger.debug( f"Retrieved batch {batch_id} from ManagedObjectTable with status={response.status}" ) - + return db_batch_object, response - + except Exception as e: verbose_proxy_logger.warning( f"Failed to retrieve batch from ManagedObjectTable: {e}, falling back to provider" @@ -762,7 +773,7 @@ async def update_batch_in_database( ): """ Update batch status and object in ManagedObjectTable. - + Args: batch_id: The batch ID (unified/encoded) unified_batch_id: Result from _is_base64_encoded_unified_file_id() @@ -774,18 +785,18 @@ async def update_batch_in_database( operation: Description of operation ("update", "cancel", etc.) """ import litellm.utils - + if managed_files_obj is None or not unified_batch_id: return - + try: if not prisma_client: return - + # Only update if status has changed (when db_batch_object is provided) if db_batch_object and response.status == db_batch_object.status: return - + if db_batch_object: verbose_proxy_logger.info( f"Updating batch {batch_id} status from {db_batch_object.status} to {response.status}" @@ -794,10 +805,10 @@ async def update_batch_in_database( verbose_proxy_logger.info( f"Updating batch {batch_id} status to {response.status} after {operation}" ) - + # Normalize status for database storage db_status = response.status if response.status != "completed" else "complete" - + await prisma_client.db.litellm_managedobjecttable.update( where={"unified_object_id": batch_id}, data={ diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 8a02f96926..973836b13d 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -145,7 +145,7 @@ async def route_create_file( ) -> OpenAIFileObject: """ Route file creation request to the appropriate provider. - + Priority: 1. If target_storage is specified and not "default" -> use storage backend 2. If model parameter provided -> use model credentials and encode ID @@ -153,7 +153,7 @@ async def route_create_file( 4. If enable_loadbalancing_on_batch_endpoints -> deprecated loadbalancing 5. Else -> use custom_llm_provider with files_settings """ - + # Handle custom storage backend if target_storage and target_storage != "default": from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -162,7 +162,7 @@ async def route_create_file( # Extract file data file_data = extract_file_data(cast(Any, _create_file_request.get("file"))) - + # Use storage backend service to handle upload file_object = await StorageBackendFileService.upload_file_to_storage_backend( file_data=file_data, @@ -172,9 +172,9 @@ async def route_create_file( proxy_logging_obj=proxy_logging_obj, user_api_key_dict=user_api_key_dict, ) - + return file_object - + # NEW: Handle model-based routing (no DB required) if model is not None: # Get credentials from model_list via router @@ -183,19 +183,19 @@ async def route_create_file( model_id=model, operation_context="file upload", ) - + # Merge credentials into the request prepare_data_with_credentials( data=_create_file_request, # type: ignore credentials=credentials, ) - + # Create the file with model credentials response = await litellm.acreate_file( - **_create_file_request, - custom_llm_provider=credentials["custom_llm_provider"] + **_create_file_request, + custom_llm_provider=credentials["custom_llm_provider"], ) # type: ignore - + # Encode the file ID with model information if response and hasattr(response, "id") and response.id: original_id = response.id @@ -204,9 +204,9 @@ async def route_create_file( verbose_proxy_logger.debug( f"Encoded file ID: {original_id} -> {encoded_id} (model: {model})" ) - + return response - + # Handle managed files (supports loadbalancing via llm_router.acreate_file) # Priority: Check for managed files BEFORE deprecated loadbalancing if target_model_names_list: @@ -339,7 +339,7 @@ async def create_file( # noqa: PLR0915 target_model_names_form=target_model_names, target_storage_form=target_storage, ) - + target_storage = file_params.target_storage target_model_names_list = file_params.target_model_names model_param = file_params.model @@ -358,18 +358,19 @@ async def create_file( # noqa: PLR0915 purpose = cast(OpenAIFilesPurpose, purpose) data = {} - + # Parse expires_after if provided expires_after: Optional[FileExpiresAfter] = None form_data_raw = await request.form() form_data_dict: Dict[str, Any] = dict(form_data_raw) - extracted_litellm_metadata: Optional[Dict[str, Any]] = extract_nested_form_metadata( - form_data=form_data_dict, - prefix="litellm_metadata[" + extracted_litellm_metadata: Optional[ + Dict[str, Any] + ] = extract_nested_form_metadata( + form_data=form_data_dict, prefix="litellm_metadata[" ) expires_after_anchor = form_data_raw.get("expires_after[anchor]") expires_after_seconds_str = form_data_raw.get("expires_after[seconds]") - + # Add litellm_metadata to data if provided (from form field) if extracted_litellm_metadata is not None: data["litellm_metadata"] = extracted_litellm_metadata @@ -382,7 +383,7 @@ async def create_file( # noqa: PLR0915 "error": "Both expires_after[anchor] and expires_after[seconds] must be provided if expires_after is specified", }, ) - + # Validate expires_after[anchor] is a string (not UploadFile) if isinstance(expires_after_anchor, UploadFile): raise HTTPException( @@ -391,7 +392,7 @@ async def create_file( # noqa: PLR0915 "error": "expires_after[anchor] must be a string, not a file upload", }, ) - + # Validate expires_after[seconds] is a string (not UploadFile) # Use positive isinstance check for proper type narrowing (matches codebase pattern) if not isinstance(expires_after_seconds_str, str): @@ -403,7 +404,7 @@ async def create_file( # noqa: PLR0915 ) # After this check, mypy knows expires_after_seconds_str is str expires_after_seconds_str_validated: str = expires_after_seconds_str - + # Validate anchor is "created_at" if expires_after_anchor != "created_at": raise HTTPException( @@ -412,7 +413,7 @@ async def create_file( # noqa: PLR0915 "error": f"expires_after[anchor] must be 'created_at', got '{expires_after_anchor}'", }, ) - + # Convert seconds to int try: expires_after_seconds = int(expires_after_seconds_str_validated) @@ -423,7 +424,7 @@ async def create_file( # noqa: PLR0915 "error": f"expires_after[seconds] must be a valid integer, got '{expires_after_seconds_str}': {e}", }, ) - + # Use literal "created_at" (not variable) for TypedDict to satisfy Literal type expires_after = FileExpiresAfter( anchor="created_at", # Literal, not expires_after_anchor variable @@ -458,7 +459,10 @@ async def create_file( # noqa: PLR0915 team_metadata = user_api_key_dict.team_metadata or {} enforced_file_expiry = team_metadata.get("enforced_file_expires_after") if enforced_file_expiry is not None: - if "anchor" not in enforced_file_expiry or "seconds" not in enforced_file_expiry: + if ( + "anchor" not in enforced_file_expiry + or "seconds" not in enforced_file_expiry + ): raise HTTPException( status_code=500, detail={ @@ -477,15 +481,13 @@ async def create_file( # noqa: PLR0915 seconds=int(enforced_file_expiry["seconds"]), ) - verbose_proxy_logger.debug( - "create_file expires_after: %s", expires_after - ) + verbose_proxy_logger.debug("create_file expires_after: %s", expires_after) _create_file_request = CreateFileRequest( file=file_data, purpose=cast(CREATE_FILE_REQUESTS_PURPOSE, purpose), expires_after=expires_after, - **data + **data, ) response = await route_create_file( @@ -657,9 +659,11 @@ async def get_file_content( # noqa: PLR0915 param="None", code=500, ) - + # Check if file is stored in a storage backend (check DB) - if hasattr(managed_files_obj, "prisma_client") and getattr(managed_files_obj, "prisma_client", None): + if hasattr(managed_files_obj, "prisma_client") and getattr( + managed_files_obj, "prisma_client", None + ): prisma_client = getattr(managed_files_obj, "prisma_client") db_file = await prisma_client.db.litellm_managedfiletable.find_first( where={"unified_file_id": file_id} @@ -669,17 +673,18 @@ async def get_file_content( # noqa: PLR0915 from litellm.llms.base_llm.files.storage_backend_factory import ( get_storage_backend, ) - + storage_backend_name = db_file.storage_backend storage_url = db_file.storage_url - + try: # Get storage backend (uses same env vars as callback) storage_backend = get_storage_backend(storage_backend_name) file_content = await storage_backend.download_file(storage_url) - + # Return file content from fastapi.responses import Response as FastAPIResponse + return FastAPIResponse( content=file_content, media_type="application/octet-stream", @@ -691,7 +696,7 @@ async def get_file_content( # noqa: PLR0915 param="file_id", code=400, ) - + model = cast(Optional[str], data.get("model")) if model: response = await llm_router.afile_content( @@ -713,14 +718,19 @@ async def get_file_content( # noqa: PLR0915 ) else: # Check for model-based credential routing - should_route, model_used, original_file_id, credentials = handle_model_based_routing( + ( + should_route, + model_used, + original_file_id, + credentials, + ) = handle_model_based_routing( file_id=file_id, request=request, llm_router=llm_router, data=data, check_file_id_encoding=True, ) - + if should_route: # Use model-based routing with credentials from config prepare_data_with_credentials( @@ -728,15 +738,19 @@ async def get_file_content( # noqa: PLR0915 credentials=credentials, # type: ignore file_id=original_file_id, # Use decoded file ID if from encoded ID ) - + response = await litellm.afile_content( custom_llm_provider=credentials["custom_llm_provider"], # type: ignore - **data + **data, ) # type: ignore - + verbose_proxy_logger.debug( f"Retrieved file content using model: {model_used}" - + (f", file_id: {file_id} -> {original_file_id}" if original_file_id else "") + + ( + f", file_id: {file_id} -> {original_file_id}" + if original_file_id + else "" + ) ) else: # Fallback to default behavior (uses env variables or provider-based routing) @@ -854,7 +868,6 @@ async def get_file( data: Dict = {"file_id": file_id} try: - custom_llm_provider = ( provider or get_custom_llm_provider_from_request_headers(request=request) @@ -880,15 +893,20 @@ async def get_file( ## Check for model-based credential routing from litellm.proxy.proxy_server import llm_router - - should_route, model_used, original_file_id, credentials = handle_model_based_routing( + + ( + should_route, + model_used, + original_file_id, + credentials, + ) = handle_model_based_routing( file_id=file_id, request=request, llm_router=llm_router, data=data, check_file_id_encoding=True, ) - + if should_route: # Use model-based routing with credentials from config prepare_data_with_credentials( @@ -898,16 +916,21 @@ async def get_file( ) response = await litellm.afile_retrieve(**data) # type: ignore - + # Keep the encoded ID in response if it was originally encoded - if original_file_id and response and hasattr(response, "id") and response.id: + if ( + original_file_id + and response + and hasattr(response, "id") + and response.id + ): response.id = file_id - + verbose_proxy_logger.debug( f"Retrieved file using model: {model_used}" + (f", original_id: {original_file_id}" if original_file_id else "") ) - + ## EXISTING: check if file_id is a litellm managed file elif _is_base64_encoded_unified_file_id(file_id): managed_files_obj = proxy_logging_obj.get_proxy_hook("managed_files") @@ -1044,7 +1067,7 @@ async def delete_file( or await get_custom_llm_provider_from_request_body(request=request) or "openai" ) - + # Call common_processing_pre_call_logic to trigger permission checks base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) ( @@ -1059,7 +1082,7 @@ async def delete_file( proxy_config=proxy_config, route_type="afile_delete", ) - + # Include original request and headers in the data data = await add_litellm_data_to_request( data=data, @@ -1071,14 +1094,19 @@ async def delete_file( ) # Check for model-based credential routing - should_route, model_used, original_file_id, credentials = handle_model_based_routing( + ( + should_route, + model_used, + original_file_id, + credentials, + ) = handle_model_based_routing( file_id=file_id, request=request, llm_router=llm_router, data=data, check_file_id_encoding=True, ) - + if should_route: # Use model-based routing with credentials from config prepare_data_with_credentials( @@ -1086,14 +1114,14 @@ async def delete_file( credentials=credentials, # type: ignore file_id=original_file_id, ) - + response = await litellm.afile_delete(**data) # type: ignore - + verbose_proxy_logger.debug( f"Deleted file using model: {model_used}" + (f", original_id: {original_file_id}" if original_file_id else "") ) - + ## EXISTING: check if file_id is a litellm managed file elif _is_base64_encoded_unified_file_id(file_id): managed_files_obj = proxy_logging_obj.get_proxy_hook("managed_files") @@ -1246,7 +1274,7 @@ async def list_files( ) response: Optional[Any] = None - + # Check for model-based credential routing (no file_id encoding check for list) should_route, model_used, _, credentials = handle_model_based_routing( file_id="", # No file_id for list endpoint @@ -1255,18 +1283,18 @@ async def list_files( data=data, check_file_id_encoding=False, ) - + if should_route: # Use model-based routing with credentials from config data.update(credentials) # type: ignore response = await litellm.afile_list( custom_llm_provider=credentials["custom_llm_provider"], # type: ignore purpose=purpose, - **data # type: ignore + **data, # type: ignore ) - + verbose_proxy_logger.debug(f"Listed files using model: {model_used}") - + elif target_model_names and isinstance(target_model_names, str): target_model_names_list = target_model_names.split(",") if len(target_model_names_list) != 1: diff --git a/litellm/proxy/openai_files_endpoints/storage_backend_service.py b/litellm/proxy/openai_files_endpoints/storage_backend_service.py index 991fff1d3f..9adeeb995a 100644 --- a/litellm/proxy/openai_files_endpoints/storage_backend_service.py +++ b/litellm/proxy/openai_files_endpoints/storage_backend_service.py @@ -22,14 +22,14 @@ from litellm.types.utils import SpecialEnums class StorageBackendFileService: """ Service for handling file uploads to storage backends. - + This service encapsulates the logic for: - Uploading files to storage backends - Creating file objects with storage metadata - Generating unified file IDs for managed files - Storing files in the managed files system """ - + @staticmethod async def upload_file_to_storage_backend( file_data: Mapping[str, Any], @@ -41,7 +41,7 @@ class StorageBackendFileService: ) -> OpenAIFileObject: """ Upload a file to a storage backend and create a file object. - + Args: file_data: File data dictionary from extract_file_data() target_storage: Storage backend name (e.g., "azure_storage") @@ -49,10 +49,10 @@ class StorageBackendFileService: purpose: File purpose (e.g., "user_data", "batch") proxy_logging_obj: Proxy logging object for accessing hooks user_api_key_dict: User API key authentication data - + Returns: OpenAIFileObject: Created file object with storage metadata - + Raises: ProxyException: If storage backend is invalid or upload fails """ @@ -66,12 +66,12 @@ class StorageBackendFileService: param="target_storage", code=400, ) - + # Extract file information file_content = file_data["content"] filename = file_data.get("filename", "file") content_type = file_data.get("content_type", "application/octet-stream") - + # Upload to storage backend storage_url = await storage_backend.upload_file( file_content=file_content, @@ -80,20 +80,22 @@ class StorageBackendFileService: path_prefix="", file_naming_strategy="uuid", ) - + verbose_proxy_logger.debug( f"Storage backend upload complete: backend={target_storage}, url={storage_url}" ) - + # Create file object with storage metadata - file_object = StorageBackendFileService._create_file_object_with_storage_metadata( - file_content=file_content, - filename=filename, - purpose=purpose, - target_storage=target_storage, - storage_url=storage_url, + file_object = ( + StorageBackendFileService._create_file_object_with_storage_metadata( + file_content=file_content, + filename=filename, + purpose=purpose, + target_storage=target_storage, + storage_url=storage_url, + ) ) - + # Store in managed files if target_model_names provided if target_model_names: await StorageBackendFileService._store_in_managed_files( @@ -105,9 +107,9 @@ class StorageBackendFileService: proxy_logging_obj=proxy_logging_obj, user_api_key_dict=user_api_key_dict, ) - + return file_object - + @staticmethod def _create_file_object_with_storage_metadata( file_content: bytes, @@ -118,14 +120,14 @@ class StorageBackendFileService: ) -> OpenAIFileObject: """ Create an OpenAIFileObject with storage backend metadata. - + Args: file_content: File content bytes filename: Original filename purpose: File purpose target_storage: Storage backend name storage_url: URL where file is stored - + Returns: OpenAIFileObject: File object with storage metadata in _hidden_params """ @@ -139,17 +141,22 @@ class StorageBackendFileService: filename=filename, status="uploaded", ) - + # Store storage metadata in hidden params - if not hasattr(file_object, "_hidden_params") or file_object._hidden_params is None: + if ( + not hasattr(file_object, "_hidden_params") + or file_object._hidden_params is None + ): file_object._hidden_params = {} - file_object._hidden_params.update({ - "storage_backend": target_storage, - "storage_url": storage_url, - }) - + file_object._hidden_params.update( + { + "storage_backend": target_storage, + "storage_url": storage_url, + } + ) + return file_object - + @staticmethod def _create_unified_file_id( file_type: str, @@ -158,29 +165,31 @@ class StorageBackendFileService: ) -> str: """ Create a base64-encoded unified file ID for managed files. - + Args: file_type: MIME type of the file target_model_names: List of model names file_id: Original file ID - + Returns: str: Base64-encoded unified file ID """ - unified_file_id_str = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( - file_type, - str(uuid_module.uuid4()), - ",".join(target_model_names), - file_id, - None, + unified_file_id_str = ( + SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + file_type, + str(uuid_module.uuid4()), + ",".join(target_model_names), + file_id, + None, + ) ) - + base64_unified_file_id = ( base64.urlsafe_b64encode(unified_file_id_str.encode()).decode().rstrip("=") ) - + return base64_unified_file_id - + @staticmethod async def _store_in_managed_files( file_object: OpenAIFileObject, @@ -193,7 +202,7 @@ class StorageBackendFileService: ) -> None: """ Store file in managed files system with unified file ID. - + Args: file_object: File object to store file_data: File data dictionary @@ -204,19 +213,18 @@ class StorageBackendFileService: user_api_key_dict: User API key authentication data """ managed_files_obj = proxy_logging_obj.get_proxy_hook("managed_files") - if not managed_files_obj or not isinstance(managed_files_obj, BaseFileEndpoints): + if not managed_files_obj or not isinstance( + managed_files_obj, BaseFileEndpoints + ): verbose_proxy_logger.warning( "Managed files hook not available, skipping managed files storage" ) return managed_files_obj = cast(Any, managed_files_obj) - + # Create model mappings using storage URL - model_mappings = { - model_name: storage_url - for model_name in target_model_names - } - + model_mappings = {model_name: storage_url for model_name in target_model_names} + # Create unified file ID file_type = file_data.get("content_type", "application/octet-stream") base64_unified_file_id = StorageBackendFileService._create_unified_file_id( @@ -224,15 +232,15 @@ class StorageBackendFileService: target_model_names=target_model_names, file_id=file_object.id, ) - + # Update file object ID to unified ID file_object.id = base64_unified_file_id - + verbose_proxy_logger.debug( f"Storing file in managed files: unified_id={base64_unified_file_id}, " f"storage_backend={target_storage}, storage_url={storage_url}" ) - + # Store in managed files await managed_files_obj.store_unified_file_id( file_id=base64_unified_file_id, @@ -241,4 +249,3 @@ class StorageBackendFileService: model_mappings=model_mappings, user_api_key_dict=user_api_key_dict, ) - diff --git a/litellm/proxy/pass_through_endpoints/common_utils.py b/litellm/proxy/pass_through_endpoints/common_utils.py index 804960cdee..3a3783dd57 100644 --- a/litellm/proxy/pass_through_endpoints/common_utils.py +++ b/litellm/proxy/pass_through_endpoints/common_utils.py @@ -14,4 +14,3 @@ def get_litellm_virtual_key(request: Request) -> str: if litellm_api_key: return f"Bearer {litellm_api_key}" return request.headers.get("Authorization", "") - diff --git a/litellm/proxy/pass_through_endpoints/jsonpath_extractor.py b/litellm/proxy/pass_through_endpoints/jsonpath_extractor.py index fde2553be4..5fab1d504b 100644 --- a/litellm/proxy/pass_through_endpoints/jsonpath_extractor.py +++ b/litellm/proxy/pass_through_endpoints/jsonpath_extractor.py @@ -19,16 +19,16 @@ class JsonPathExtractor: ) -> str: """ Extract field values from data using JSONPath-like expressions. - + Supports simple expressions like: - "query" -> data["query"] - "documents[*].text" -> all text fields from documents array - "messages[*].content" -> all content fields from messages array - + Returns concatenated string of all extracted values. """ extracted_values: List[str] = [] - + for expr in jsonpath_expressions: try: value = JsonPathExtractor.evaluate(data, expr) @@ -41,14 +41,14 @@ class JsonPathExtractor: verbose_proxy_logger.debug( "Failed to extract field %s: %s", expr, str(e) ) - + return "\n".join(extracted_values) @staticmethod def evaluate(data: dict, expr: str) -> Union[str, List[str], None]: """ Evaluate a simple JSONPath-like expression. - + Supports: - Simple key: "query" -> data["query"] - Nested key: "foo.bar" -> data["foo"]["bar"] @@ -56,24 +56,24 @@ class JsonPathExtractor: """ if not expr or not data: return None - + parts = expr.replace("[*]", ".[*]").split(".") current: Any = data - + for i, part in enumerate(parts): if current is None: return None - + if part == "[*]": # Wildcard - current should be a list if not isinstance(current, list): return None - + # Get remaining path - remaining_path = ".".join(parts[i + 1:]) + remaining_path = ".".join(parts[i + 1 :]) if not remaining_path: return current - + # Recursively evaluate remaining path for each item results = [] for item in current: @@ -85,11 +85,10 @@ class JsonPathExtractor: else: results.append(result) return results if results else None - + elif isinstance(current, dict): current = current.get(part) else: return None - - return current + return current diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 13f78f30fa..4e3e04a847 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -695,10 +695,10 @@ def _extract_model_from_bedrock_endpoint(endpoint: str) -> str: if action_index is not None and action_index > model_index + 1: # Join all parts between "model" and the action (excluding "model" itself) - return "/".join(endpoint_parts[model_index + 1:action_index]) + return "/".join(endpoint_parts[model_index + 1 : action_index]) # Fallback to taking everything after "model" if no action found - model_parts = [p for p in endpoint_parts[model_index + 1:] if p] + model_parts = [p for p in endpoint_parts[model_index + 1 :] if p] if model_parts: return "/".join(model_parts) @@ -866,10 +866,10 @@ async def handle_bedrock_count_tokens( except BedrockError as e: # Convert BedrockError to HTTPException for FastAPI - verbose_proxy_logger.error(f"BedrockError in handle_bedrock_count_tokens: {str(e)}") - raise HTTPException( - status_code=e.status_code, detail={"error": e.message} + verbose_proxy_logger.error( + f"BedrockError in handle_bedrock_count_tokens: {str(e)}" ) + raise HTTPException(status_code=e.status_code, detail={"error": e.message}) except HTTPException: # Re-raise HTTP exceptions as-is raise @@ -1078,7 +1078,7 @@ async def bedrock_proxy_route( target=str(prepped.url), custom_headers=prepped.headers, # type: ignore is_streaming_request=is_streaming_request, - _forward_headers=True + _forward_headers=True, ) # dynamically construct pass-through endpoint based on incoming path received_value = await endpoint_func( request, @@ -1100,7 +1100,7 @@ def _resolve_vertex_model_from_router( ) -> Tuple[str, str, Optional[str], Optional[str]]: """ Resolve Vertex AI model configuration from router. - + Args: model_id: The model ID extracted from the URL (e.g., "gcp/google/gemini-2.5-flash") llm_router: The LiteLLM router instance @@ -1108,21 +1108,23 @@ def _resolve_vertex_model_from_router( endpoint: The original endpoint path vertex_project: Current vertex project (may be from URL) vertex_location: Current vertex location (may be from URL) - + Returns: Tuple of (encoded_endpoint, endpoint, vertex_project, vertex_location) with resolved values from router config """ if not llm_router: return encoded_endpoint, endpoint, vertex_project, vertex_location - + try: - deployment = llm_router.get_available_deployment_for_pass_through(model=model_id) + deployment = llm_router.get_available_deployment_for_pass_through( + model=model_id + ) if not deployment: return encoded_endpoint, endpoint, vertex_project, vertex_location - + litellm_params = deployment.get("litellm_params", {}) - + # Always override with router config values (they take precedence over URL values) config_vertex_project = litellm_params.get("vertex_project") config_vertex_location = litellm_params.get("vertex_location") @@ -1130,12 +1132,11 @@ def _resolve_vertex_model_from_router( vertex_project = config_vertex_project if config_vertex_location: vertex_location = config_vertex_location - + # Get the actual Vertex AI model name by stripping the provider prefix # e.g., "vertex_ai/gemini-2.0-flash-exp" -> "gemini-2.0-flash-exp" model_from_config = litellm_params.get("model", "") if model_from_config: - # get_llm_provider returns (model, custom_llm_provider, dynamic_api_key, api_base) # For "vertex_ai/gemini-2.0-flash-exp" it returns: # model="gemini-2.0-flash-exp", custom_llm_provider="vertex_ai" @@ -1164,12 +1165,12 @@ def _resolve_vertex_model_from_router( ) encoded_endpoint = encoded_endpoint.replace(model_id, actual_model) endpoint = endpoint.replace(model_id, actual_model) - + except Exception as e: verbose_proxy_logger.debug( f"Error resolving vertex model from router for model {model_id}: {e}" ) - + return encoded_endpoint, endpoint, vertex_project, vertex_location @@ -1634,7 +1635,7 @@ async def _prepare_vertex_auth_headers( vertex_credentials_str = None elif vertex_credentials is not None: # Use credentials from vertex_credentials - # When vertex_credentials are provided (including default credentials), + # When vertex_credentials are provided (including default credentials), # use their project/location values if available if vertex_credentials.vertex_project is not None: vertex_project = vertex_credentials.vertex_project @@ -1740,10 +1741,14 @@ async def _base_vertex_proxy_route( # Check if model is in router config - always do this to resolve custom model names model_id = get_vertex_model_id_from_url(endpoint) if model_id: - if llm_router: # Resolve model configuration from router - encoded_endpoint, endpoint, vertex_project, vertex_location = _resolve_vertex_model_from_router( + ( + encoded_endpoint, + endpoint, + vertex_project, + vertex_location, + ) = _resolve_vertex_model_from_router( model_id=model_id, llm_router=llm_router, encoded_endpoint=encoded_endpoint, @@ -1936,25 +1941,25 @@ async def openai_proxy_route( ): """ Pass-through endpoint for OpenAI API calls. - + Available on both routes: - /openai/{endpoint:path} - Standard OpenAI passthrough route - /openai_passthrough/{endpoint:path} - Dedicated passthrough route (recommended for Responses API) - + Use /openai_passthrough/* when you need guaranteed passthrough to OpenAI without conflicts with LiteLLM's native implementations (e.g., for the Responses API at /v1/responses). - + Examples: Standard route: - /openai/v1/chat/completions - /openai/v1/assistants - /openai/v1/threads - + Dedicated passthrough (for Responses API): - /openai_passthrough/v1/responses - /openai_passthrough/v1/responses/{response_id} - /openai_passthrough/v1/responses/{response_id}/input_items - + [Docs](https://docs.litellm.ai/docs/pass_through/openai_passthrough) """ base_target_url = os.getenv("OPENAI_API_BASE") or "https://api.openai.com/" @@ -2154,9 +2159,7 @@ async def cursor_proxy_route( base_url = httpx.URL(base_target_url) updated_url = base_url.copy_with(path=encoded_endpoint) - auth_value = base64.b64encode( - f"{cursor_api_key}:".encode("utf-8") - ).decode("ascii") + auth_value = base64.b64encode(f"{cursor_api_key}:".encode("utf-8")).decode("ascii") endpoint_func = create_pass_through_route( endpoint=endpoint, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index e70d6cb7fc..20d06b7d53 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -59,7 +59,7 @@ class AnthropicPassthroughLoggingHandler: request_body=request_body, **kwargs, ) - + model = response_body.get("model", "") anthropic_config = get_anthropic_config(url_route) litellm_model_response: ModelResponse = anthropic_config().transform_response( @@ -156,9 +156,9 @@ class AnthropicPassthroughLoggingHandler: litellm_model_response.model = model logging_obj.model_call_details["model"] = model if not logging_obj.model_call_details.get("custom_llm_provider"): - logging_obj.model_call_details["custom_llm_provider"] = ( - litellm.LlmProviders.ANTHROPIC.value - ) + logging_obj.model_call_details[ + "custom_llm_provider" + ] = litellm.LlmProviders.ANTHROPIC.value return kwargs except Exception as e: verbose_proxy_logger.exception( @@ -326,25 +326,26 @@ class AnthropicPassthroughLoggingHandler: try: _json_response = httpx_response.json() - - + # Only handle successful batch job creation (POST requests with 201 status) if httpx_response.status_code == 200 and "id" in _json_response: # Transform Anthropic response to LiteLLM batch format anthropic_batches_config = AnthropicBatchesConfig() - litellm_batch_response = anthropic_batches_config.transform_retrieve_batch_response( - model=None, - raw_response=httpx_response, - logging_obj=logging_obj, - litellm_params={}, + litellm_batch_response = ( + anthropic_batches_config.transform_retrieve_batch_response( + model=None, + raw_response=httpx_response, + logging_obj=logging_obj, + litellm_params={}, + ) ) # Set status to "validating" for newly created batches so polling mechanism picks them up # The polling mechanism only looks for status="validating" jobs litellm_batch_response.status = "validating" - + # Extract batch ID from the response batch_id = _json_response.get("id", "") - + # Get model from request body (batch response doesn't include model) request_body = request_body or {} # Try to extract model from the batch request body, supporting Anthropic's nested structure @@ -363,20 +364,33 @@ class AnthropicPassthroughLoggingHandler: extracted_model = params.get("model") if extracted_model: model_name = extracted_model - - + # Create unified object ID for tracking # Format: base64(litellm_proxy;model_id:{};llm_batch_id:{}) # For Anthropic passthrough, prefix model with "anthropic/" so router can determine provider - actual_model_id = AnthropicPassthroughLoggingHandler.get_actual_model_id_from_router(model_name) - + actual_model_id = ( + AnthropicPassthroughLoggingHandler.get_actual_model_id_from_router( + model_name + ) + ) + # If model not in router, use "anthropic/{model_name}" format so router can determine provider - if actual_model_id == model_name and not actual_model_id.startswith("anthropic/"): + if actual_model_id == model_name and not actual_model_id.startswith( + "anthropic/" + ): actual_model_id = f"anthropic/{model_name}" - unified_id_string = SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format(actual_model_id, batch_id) - unified_object_id = base64.urlsafe_b64encode(unified_id_string.encode()).decode().rstrip("=") - + unified_id_string = ( + SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format( + actual_model_id, batch_id + ) + ) + unified_object_id = ( + base64.urlsafe_b64encode(unified_id_string.encode()) + .decode() + .rstrip("=") + ) + # Store the managed object for cost tracking # This will be picked up by check_batch_cost polling mechanism AnthropicPassthroughLoggingHandler._store_batch_managed_object( @@ -386,31 +400,33 @@ class AnthropicPassthroughLoggingHandler: logging_obj=logging_obj, **kwargs, ) - + # Create a batch job response for logging litellm_model_response = ModelResponse() litellm_model_response.id = str(uuid.uuid4()) litellm_model_response.model = model_name litellm_model_response.object = "batch" litellm_model_response.created = int(start_time.timestamp()) - + # Add batch-specific metadata to indicate this is a pending batch job - litellm_model_response.choices = [Choices( - finish_reason="stop", - index=0, - message={ - "role": "assistant", - "content": f"Batch job {batch_id} created and is pending. Status will be updated when the batch completes.", - "tool_calls": None, - "function_call": None, - "provider_specific_fields": { - "batch_job_id": batch_id, - "batch_job_state": "in_progress", - "unified_object_id": unified_object_id - } - } - )] - + litellm_model_response.choices = [ + Choices( + finish_reason="stop", + index=0, + message={ + "role": "assistant", + "content": f"Batch job {batch_id} created and is pending. Status will be updated when the batch completes.", + "tool_calls": None, + "function_call": None, + "provider_specific_fields": { + "batch_job_id": batch_id, + "batch_job_state": "in_progress", + "unified_object_id": unified_object_id, + }, + }, + ) + ] + # Set response cost to 0 initially (will be updated when batch completes) response_cost = 0.0 kwargs["response_cost"] = response_cost @@ -418,12 +434,12 @@ class AnthropicPassthroughLoggingHandler: kwargs["batch_id"] = batch_id kwargs["unified_object_id"] = unified_object_id kwargs["batch_job_state"] = "in_progress" - + logging_obj.model = model_name logging_obj.model_call_details["model"] = logging_obj.model logging_obj.model_call_details["response_cost"] = response_cost logging_obj.model_call_details["batch_id"] = batch_id - + return { "result": litellm_model_response, "kwargs": kwargs, @@ -435,32 +451,34 @@ class AnthropicPassthroughLoggingHandler: litellm_model_response.model = "anthropic_batch" litellm_model_response.object = "batch" litellm_model_response.created = int(start_time.timestamp()) - + # Add error-specific metadata - litellm_model_response.choices = [Choices( - finish_reason="stop", - index=0, - message={ - "role": "assistant", - "content": f"Batch job creation failed. Status: {httpx_response.status_code}", - "tool_calls": None, - "function_call": None, - "provider_specific_fields": { - "batch_job_state": "failed", - "status_code": httpx_response.status_code - } - } - )] - + litellm_model_response.choices = [ + Choices( + finish_reason="stop", + index=0, + message={ + "role": "assistant", + "content": f"Batch job creation failed. Status: {httpx_response.status_code}", + "tool_calls": None, + "function_call": None, + "provider_specific_fields": { + "batch_job_state": "failed", + "status_code": httpx_response.status_code, + }, + }, + ) + ] + kwargs["response_cost"] = 0.0 kwargs["model"] = "anthropic_batch" kwargs["batch_job_state"] = "failed" - + return { "result": litellm_model_response, "kwargs": kwargs, } - + except Exception as e: verbose_proxy_logger.error(f"Error in batch_creation_handler: {e}") # Return basic response on error @@ -469,27 +487,29 @@ class AnthropicPassthroughLoggingHandler: litellm_model_response.model = "anthropic_batch" litellm_model_response.object = "batch" litellm_model_response.created = int(start_time.timestamp()) - + # Add error-specific metadata - litellm_model_response.choices = [Choices( - finish_reason="stop", - index=0, - message={ - "role": "assistant", - "content": f"Error creating batch job: {str(e)}", - "tool_calls": None, - "function_call": None, - "provider_specific_fields": { - "batch_job_state": "failed", - "error": str(e) - } - } - )] - + litellm_model_response.choices = [ + Choices( + finish_reason="stop", + index=0, + message={ + "role": "assistant", + "content": f"Error creating batch job: {str(e)}", + "tool_calls": None, + "function_call": None, + "provider_specific_fields": { + "batch_job_state": "failed", + "error": str(e), + }, + }, + ) + ] + kwargs["response_cost"] = 0.0 kwargs["model"] = "anthropic_batch" kwargs["batch_job_state"] = "failed" - + return { "result": litellm_model_response, "kwargs": kwargs, @@ -508,15 +528,17 @@ class AnthropicPassthroughLoggingHandler: This will be picked up by the check_batch_cost polling mechanism. """ try: - # Get the managed files hook from the logging object # This is a bit of a hack, but we need access to the proxy logging system from litellm.proxy.proxy_server import proxy_logging_obj - + managed_files_hook = proxy_logging_obj.get_proxy_hook("managed_files") - if managed_files_hook is not None and hasattr(managed_files_hook, 'store_unified_object_id'): + if managed_files_hook is not None and hasattr( + managed_files_hook, "store_unified_object_id" + ): # Create a mock user API key dict for the managed object storage from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + user_api_key_dict = UserAPIKeyAuth( user_id=kwargs.get("user_id", "default-user"), api_key="", @@ -539,9 +561,10 @@ class AnthropicPassthroughLoggingHandler: model_max_budget={}, # Set to empty dict instead of None model_spend={}, # Set to empty dict instead of None ) - + # Store the unified object for batch cost tracking import asyncio + asyncio.create_task( managed_files_hook.store_unified_object_id( # type: ignore unified_object_id=unified_object_id, @@ -552,20 +575,24 @@ class AnthropicPassthroughLoggingHandler: user_api_key_dict=user_api_key_dict, ) ) - + verbose_proxy_logger.info( f"Stored Anthropic batch managed object with unified_object_id={unified_object_id}, batch_id={model_object_id}" ) else: - verbose_proxy_logger.warning("Managed files hook not available, cannot store batch object for cost tracking") - + verbose_proxy_logger.warning( + "Managed files hook not available, cannot store batch object for cost tracking" + ) + except Exception as e: - verbose_proxy_logger.error(f"Error storing Anthropic batch managed object: {e}") + verbose_proxy_logger.error( + f"Error storing Anthropic batch managed object: {e}" + ) @staticmethod def get_actual_model_id_from_router(model_name: str) -> str: from litellm.proxy.proxy_server import llm_router - + if llm_router is not None: # Try to find the model in the router by the model name # Use the existing get_model_ids method from router @@ -573,14 +600,20 @@ class AnthropicPassthroughLoggingHandler: if model_ids and len(model_ids) > 0: # Use the first model ID found actual_model_id = model_ids[0] - verbose_proxy_logger.info(f"Found model ID in router: {actual_model_id}") + verbose_proxy_logger.info( + f"Found model ID in router: {actual_model_id}" + ) return actual_model_id else: # Fallback to model name actual_model_id = model_name - verbose_proxy_logger.warning(f"Model not found in router, using model name: {actual_model_id}") + verbose_proxy_logger.warning( + f"Model not found in router, using model name: {actual_model_id}" + ) return actual_model_id else: # Fallback if router is not available - verbose_proxy_logger.warning(f"Router not available, using model name: {model_name}") + verbose_proxy_logger.warning( + f"Router not available, using model name: {model_name}" + ) return model_name diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cohere_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cohere_passthrough_logging_handler.py index 743f4e4f96..adb1278fee 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cohere_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cohere_passthrough_logging_handler.py @@ -176,7 +176,7 @@ class CoherePassthroughLoggingHandler(BasePassthroughLoggingHandler): request_body=request_body, **kwargs, ) - + # For non-embed routes (e.g., /v2/chat), fall back to chat handler return super().passthrough_chat_handler( httpx_response=httpx_response, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cursor_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cursor_passthrough_logging_handler.py index 2d687928d9..a104f96263 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cursor_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cursor_passthrough_logging_handler.py @@ -126,9 +126,7 @@ class CursorPassthroughLoggingHandler: ) return { - "result": StandardPassThroughResponseObject( - response=response_summary - ), + "result": StandardPassThroughResponseObject(response=response_summary), "kwargs": kwargs, } except Exception as e: diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py index 2bda9ba485..b05cb70f75 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py @@ -42,32 +42,34 @@ class GeminiPassthroughLoggingHandler: ) -> PassThroughEndpointLoggingTypedDict: if "predictLongRunning" in url_route: model = GeminiPassthroughLoggingHandler.extract_model_from_url(url_route) - + gemini_video_config = GeminiVideoConfig() - litellm_video_response = gemini_video_config.transform_video_create_response( - model=model, - raw_response=httpx_response, - logging_obj=logging_obj, - custom_llm_provider="gemini", - request_data=request_body, + litellm_video_response = ( + gemini_video_config.transform_video_create_response( + model=model, + raw_response=httpx_response, + logging_obj=logging_obj, + custom_llm_provider="gemini", + request_data=request_body, + ) ) logging_obj.model = model logging_obj.model_call_details["model"] = model logging_obj.model_call_details["custom_llm_provider"] = "gemini" logging_obj.custom_llm_provider = "gemini" - + response_cost = litellm.completion_cost( completion_response=litellm_video_response, model=model, custom_llm_provider="gemini", call_type="create_video", ) - + # Set response_cost in _hidden_params to prevent recalculation if not hasattr(litellm_video_response, "_hidden_params"): litellm_video_response._hidden_params = {} litellm_video_response._hidden_params["response_cost"] = response_cost - + kwargs["response_cost"] = response_cost kwargs["model"] = model kwargs["custom_llm_provider"] = "gemini" @@ -76,23 +78,27 @@ class GeminiPassthroughLoggingHandler: "result": litellm_video_response, "kwargs": kwargs, } - + if "generateContent" in url_route: model = GeminiPassthroughLoggingHandler.extract_model_from_url(url_route) # Use Gemini config for transformation instance_of_gemini_llm = litellm.GoogleAIStudioGeminiConfig() - litellm_model_response: ModelResponse = instance_of_gemini_llm.transform_response( - model=model, - messages=[{"role": "user", "content": "no-message-pass-through-endpoint"}], - raw_response=httpx_response, - model_response=litellm.ModelResponse(), - logging_obj=logging_obj, - optional_params={}, - litellm_params={}, - api_key="", - request_data={}, - encoding=litellm.encoding, + litellm_model_response: ModelResponse = ( + instance_of_gemini_llm.transform_response( + model=model, + messages=[ + {"role": "user", "content": "no-message-pass-through-endpoint"} + ], + raw_response=httpx_response, + model_response=litellm.ModelResponse(), + logging_obj=logging_obj, + optional_params={}, + litellm_params={}, + api_key="", + request_data={}, + encoding=litellm.encoding, + ) ) kwargs = GeminiPassthroughLoggingHandler._create_gemini_response_logging_payload_for_generate_content( litellm_model_response=litellm_model_response, @@ -134,12 +140,16 @@ class GeminiPassthroughLoggingHandler: - Logs in litellm callbacks """ kwargs: Dict[str, Any] = {} - model = model or GeminiPassthroughLoggingHandler.extract_model_from_url(url_route) - complete_streaming_response = GeminiPassthroughLoggingHandler._build_complete_streaming_response( - all_chunks=all_chunks, - litellm_logging_obj=litellm_logging_obj, - model=model, - url_route=url_route, + model = model or GeminiPassthroughLoggingHandler.extract_model_from_url( + url_route + ) + complete_streaming_response = ( + GeminiPassthroughLoggingHandler._build_complete_streaming_response( + all_chunks=all_chunks, + litellm_logging_obj=litellm_logging_obj, + model=model, + url_route=url_route, + ) ) if complete_streaming_response is None: @@ -194,7 +204,9 @@ class GeminiPassthroughLoggingHandler: continue all_openai_chunks.append(parsed_chunk) - complete_streaming_response = litellm.stream_chunk_builder(chunks=all_openai_chunks) + complete_streaming_response = litellm.stream_chunk_builder( + chunks=all_openai_chunks + ) return complete_streaming_response diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index 6745c559cd..38b2734bc2 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -214,11 +214,16 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): is_image_editing = ( OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(url_route) ) - is_responses = ( - OpenAIPassthroughLoggingHandler.is_openai_responses_route(url_route) + is_responses = OpenAIPassthroughLoggingHandler.is_openai_responses_route( + url_route ) - if not (is_chat_completions or is_image_generation or is_image_editing or is_responses): + if not ( + is_chat_completions + or is_image_generation + or is_image_editing + or is_responses + ): # For unsupported endpoints, return None to let the system fall back to generic behavior return { "result": None, @@ -247,11 +252,13 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): try: response_cost = 0.0 - litellm_model_response: Optional[Union[ModelResponse, TextCompletionResponse, ImageResponse]] = None + litellm_model_response: Optional[ + Union[ModelResponse, TextCompletionResponse, ImageResponse] + ] = None handler_instance = OpenAIPassthroughLoggingHandler() custom_llm_provider = kwargs.get("custom_llm_provider", "openai") - + if is_chat_completions: # Handle chat completions with existing logic provider_config = handler_instance.get_provider_config(model=model) @@ -368,7 +375,9 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): passthrough_logging_payload=passthrough_logging_payload, ) if user: - kwargs["litellm_params"].setdefault("proxy_server_request", {}).setdefault("body", {})["user"] = user + kwargs["litellm_params"].setdefault( + "proxy_server_request", {} + ).setdefault("body", {})["user"] = user # Create standard logging object if litellm_model_response is not None: @@ -527,7 +536,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): custom_llm_provider = litellm_logging_obj.model_call_details.get( "custom_llm_provider", "openai" - ) + ) # Calculate cost using LiteLLM's cost calculator response_cost = litellm.completion_cost( completion_response=complete_response, @@ -536,10 +545,10 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): ) # Preserve existing litellm_params to maintain metadata tags - existing_litellm_params = litellm_logging_obj.model_call_details.get( - "litellm_params", {} - ) or {} - + existing_litellm_params = ( + litellm_logging_obj.model_call_details.get("litellm_params", {}) or {} + ) + # Prepare kwargs for logging kwargs = { "response_cost": response_cost, @@ -559,7 +568,9 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): passthrough_logging_payload=passthrough_logging_payload, ) if user: - kwargs["litellm_params"].setdefault("proxy_server_request", {}).setdefault("body", {})["user"] = user + kwargs["litellm_params"].setdefault( + "proxy_server_request", {} + ).setdefault("body", {})["user"] = user # Create standard logging object get_standard_logging_object_payload( @@ -573,7 +584,9 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): # Update logging object with cost information litellm_logging_obj.model_call_details["model"] = model - litellm_logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider + litellm_logging_obj.model_call_details[ + "custom_llm_provider" + ] = custom_llm_provider litellm_logging_obj.model_call_details["response_cost"] = response_cost verbose_proxy_logger.debug( diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py index f8eb98affc..04fe74bbf2 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py @@ -374,8 +374,13 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): # Safely log the model name: only allow known safe formats, redact otherwise. import re + allowed_pattern = re.compile(r"^[A-Za-z0-9._\-:]+$") - safe_model = model if isinstance(model, str) and allowed_pattern.match(model) else "[REDACTED]" + safe_model = ( + model + if isinstance(model, str) and allowed_pattern.match(model) + else "[REDACTED]" + ) verbose_proxy_logger.debug( f"Vertex AI Live API passthrough cost tracking - " f"Model: {safe_model}, Cost: ${response_cost:.6f}, " diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index 3d5c529a3b..d709956df5 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -40,7 +40,6 @@ EndpointType = Any class VertexPassthroughLoggingHandler: - @staticmethod def vertex_passthrough_handler( httpx_response: httpx.Response, @@ -55,43 +54,45 @@ class VertexPassthroughLoggingHandler: ) -> PassThroughEndpointLoggingTypedDict: if "predictLongRunning" in url_route: model = VertexPassthroughLoggingHandler.extract_model_from_url(url_route) - + vertex_video_config = VertexAIVideoConfig() - litellm_video_response = vertex_video_config.transform_video_create_response( - model=model, - raw_response=httpx_response, - logging_obj=logging_obj, - custom_llm_provider="vertex_ai", - request_data=request_body, + litellm_video_response = ( + vertex_video_config.transform_video_create_response( + model=model, + raw_response=httpx_response, + logging_obj=logging_obj, + custom_llm_provider="vertex_ai", + request_data=request_body, + ) ) - + logging_obj.model = model logging_obj.model_call_details["model"] = model logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai" logging_obj.custom_llm_provider = "vertex_ai" - + response_cost = litellm.completion_cost( completion_response=litellm_video_response, model=model, custom_llm_provider="vertex_ai", call_type="create_video", ) - + # Set response_cost in _hidden_params to prevent recalculation if not hasattr(litellm_video_response, "_hidden_params"): litellm_video_response._hidden_params = {} litellm_video_response._hidden_params["response_cost"] = response_cost - + kwargs["response_cost"] = response_cost kwargs["model"] = model kwargs["custom_llm_provider"] = "vertex_ai" logging_obj.model_call_details["response_cost"] = response_cost - + return { "result": litellm_video_response, "kwargs": kwargs, } - + elif "generateContent" in url_route: model = VertexPassthroughLoggingHandler.extract_model_from_url(url_route) @@ -190,7 +191,6 @@ class VertexPassthroughLoggingHandler: "kwargs": kwargs, } elif "search" in url_route: - litellm_vs_response = ( vertex_search_api_config.transform_search_vector_store_response( response=httpx_response, @@ -262,9 +262,7 @@ class VertexPassthroughLoggingHandler: litellm_prediction_response: Union[ ModelResponse, EmbeddingResponse, ImageResponse ] = ModelResponse() - if vertex_image_generation_class.is_image_generation_response( - _json_response - ): + if vertex_image_generation_class.is_image_generation_response(_json_response): litellm_prediction_response = ( vertex_image_generation_class.process_image_generation_response( _json_response, @@ -294,10 +292,12 @@ class VertexPassthroughLoggingHandler: ) ) else: - litellm_prediction_response = litellm.vertexAITextEmbeddingConfig.transform_vertex_response_to_openai( - response=_json_response, - model=model, - model_response=litellm.EmbeddingResponse(), + litellm_prediction_response = ( + litellm.vertexAITextEmbeddingConfig.transform_vertex_response_to_openai( + response=_json_response, + model=model, + model_response=litellm.EmbeddingResponse(), + ) ) if isinstance(litellm_prediction_response, litellm.EmbeddingResponse): litellm_prediction_response.model = model @@ -440,14 +440,14 @@ class VertexPassthroughLoggingHandler: def extract_model_name_from_vertex_path(vertex_model_path: str) -> str: """ Extract the actual model name from a Vertex AI model path. - + Examples: - publishers/google/models/gemini-2.5-flash -> gemini-2.5-flash - projects/PROJECT_ID/locations/LOCATION/models/MODEL_ID -> MODEL_ID - + Args: vertex_model_path: The full Vertex AI model path - + Returns: The extracted model name for use with LiteLLM """ @@ -457,14 +457,14 @@ class VertexPassthroughLoggingHandler: parts = vertex_model_path.split("models/") if len(parts) > 1: return parts[-1] - + # Handle projects/PROJECT_ID/locations/LOCATION/models/MODEL_ID format elif "projects/" in vertex_model_path and "models/" in vertex_model_path: # Extract everything after the last models/ parts = vertex_model_path.split("models/") if len(parts) > 1: return parts[-1] - + # If no recognized pattern, return the original path return vertex_model_path @@ -581,25 +581,39 @@ class VertexPassthroughLoggingHandler: try: _json_response = httpx_response.json() - + # Only handle successful batch job creation (POST requests) if httpx_response.status_code == 200 and "name" in _json_response: # Transform Vertex AI response to LiteLLM batch format litellm_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( response=_json_response ) - + # Extract batch ID and model from the response - batch_id = VertexAIBatchTransformation._get_batch_id_from_vertex_ai_batch_response(_json_response) + batch_id = VertexAIBatchTransformation._get_batch_id_from_vertex_ai_batch_response( + _json_response + ) model_name = _json_response.get("model", "unknown") - + # Create unified object ID for tracking # Format: base64(litellm_proxy;model_id:{};llm_batch_id:{}) - actual_model_id = VertexPassthroughLoggingHandler.get_actual_model_id_from_router(model_name) + actual_model_id = ( + VertexPassthroughLoggingHandler.get_actual_model_id_from_router( + model_name + ) + ) + + unified_id_string = ( + SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format( + actual_model_id, batch_id + ) + ) + unified_object_id = ( + base64.urlsafe_b64encode(unified_id_string.encode()) + .decode() + .rstrip("=") + ) - unified_id_string = SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format(actual_model_id, batch_id) - unified_object_id = base64.urlsafe_b64encode(unified_id_string.encode()).decode().rstrip("=") - # Store the managed object for cost tracking # This will be picked up by check_batch_cost polling mechanism VertexPassthroughLoggingHandler._store_batch_managed_object( @@ -609,31 +623,33 @@ class VertexPassthroughLoggingHandler: logging_obj=logging_obj, **kwargs, ) - + # Create a batch job response for logging litellm_model_response = ModelResponse() litellm_model_response.id = str(uuid.uuid4()) litellm_model_response.model = model_name litellm_model_response.object = "batch_prediction_job" litellm_model_response.created = int(start_time.timestamp()) - + # Add batch-specific metadata to indicate this is a pending batch job - litellm_model_response.choices = [Choices( - finish_reason="stop", - index=0, - message={ - "role": "assistant", - "content": f"Batch prediction job {batch_id} created and is pending. Status will be updated when the batch completes.", - "tool_calls": None, - "function_call": None, - "provider_specific_fields": { - "batch_job_id": batch_id, - "batch_job_state": "JOB_STATE_PENDING", - "unified_object_id": unified_object_id - } - } - )] - + litellm_model_response.choices = [ + Choices( + finish_reason="stop", + index=0, + message={ + "role": "assistant", + "content": f"Batch prediction job {batch_id} created and is pending. Status will be updated when the batch completes.", + "tool_calls": None, + "function_call": None, + "provider_specific_fields": { + "batch_job_id": batch_id, + "batch_job_state": "JOB_STATE_PENDING", + "unified_object_id": unified_object_id, + }, + }, + ) + ] + # Set response cost to 0 initially (will be updated when batch completes) response_cost = 0.0 kwargs["response_cost"] = response_cost @@ -641,12 +657,12 @@ class VertexPassthroughLoggingHandler: kwargs["batch_id"] = batch_id kwargs["unified_object_id"] = unified_object_id kwargs["batch_job_state"] = "JOB_STATE_PENDING" - + logging_obj.model = model_name logging_obj.model_call_details["model"] = logging_obj.model logging_obj.model_call_details["response_cost"] = response_cost logging_obj.model_call_details["batch_id"] = batch_id - + return { "result": litellm_model_response, "kwargs": kwargs, @@ -658,32 +674,34 @@ class VertexPassthroughLoggingHandler: litellm_model_response.model = "vertex_ai_batch" litellm_model_response.object = "batch_prediction_job" litellm_model_response.created = int(start_time.timestamp()) - + # Add error-specific metadata - litellm_model_response.choices = [Choices( - finish_reason="stop", - index=0, - message={ - "role": "assistant", - "content": f"Batch prediction job creation failed. Status: {httpx_response.status_code}", - "tool_calls": None, - "function_call": None, - "provider_specific_fields": { - "batch_job_state": "JOB_STATE_FAILED", - "status_code": httpx_response.status_code - } - } - )] - + litellm_model_response.choices = [ + Choices( + finish_reason="stop", + index=0, + message={ + "role": "assistant", + "content": f"Batch prediction job creation failed. Status: {httpx_response.status_code}", + "tool_calls": None, + "function_call": None, + "provider_specific_fields": { + "batch_job_state": "JOB_STATE_FAILED", + "status_code": httpx_response.status_code, + }, + }, + ) + ] + kwargs["response_cost"] = 0.0 kwargs["model"] = "vertex_ai_batch" kwargs["batch_job_state"] = "JOB_STATE_FAILED" - + return { "result": litellm_model_response, "kwargs": kwargs, } - + except Exception as e: verbose_proxy_logger.error(f"Error in batch_prediction_jobs_handler: {e}") # Return basic response on error @@ -692,27 +710,29 @@ class VertexPassthroughLoggingHandler: litellm_model_response.model = "vertex_ai_batch" litellm_model_response.object = "batch_prediction_job" litellm_model_response.created = int(start_time.timestamp()) - + # Add error-specific metadata - litellm_model_response.choices = [Choices( - finish_reason="stop", - index=0, - message={ - "role": "assistant", - "content": f"Error creating batch prediction job: {str(e)}", - "tool_calls": None, - "function_call": None, - "provider_specific_fields": { - "batch_job_state": "JOB_STATE_FAILED", - "error": str(e) - } - } - )] - + litellm_model_response.choices = [ + Choices( + finish_reason="stop", + index=0, + message={ + "role": "assistant", + "content": f"Error creating batch prediction job: {str(e)}", + "tool_calls": None, + "function_call": None, + "provider_specific_fields": { + "batch_job_state": "JOB_STATE_FAILED", + "error": str(e), + }, + }, + ) + ] + kwargs["response_cost"] = 0.0 kwargs["model"] = "vertex_ai_batch" kwargs["batch_job_state"] = "JOB_STATE_FAILED" - + return { "result": litellm_model_response, "kwargs": kwargs, @@ -730,15 +750,18 @@ class VertexPassthroughLoggingHandler: Store batch managed object for cost tracking. This will be picked up by the check_batch_cost polling mechanism. """ - try: + try: # Get the managed files hook from the logging object # This is a bit of a hack, but we need access to the proxy logging system from litellm.proxy.proxy_server import proxy_logging_obj - + managed_files_hook = proxy_logging_obj.get_proxy_hook("managed_files") - if managed_files_hook is not None and hasattr(managed_files_hook, 'store_unified_object_id'): + if managed_files_hook is not None and hasattr( + managed_files_hook, "store_unified_object_id" + ): # Create a mock user API key dict for the managed object storage from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + user_api_key_dict = UserAPIKeyAuth( user_id=kwargs.get("user_id", "default-user"), api_key="", @@ -761,9 +784,10 @@ class VertexPassthroughLoggingHandler: model_max_budget={}, # Set to empty dict instead of None model_spend={}, # Set to empty dict instead of None ) - + # Store the unified object for batch cost tracking import asyncio + asyncio.create_task( managed_files_hook.store_unified_object_id( # type: ignore unified_object_id=unified_object_id, @@ -774,39 +798,54 @@ class VertexPassthroughLoggingHandler: user_api_key_dict=user_api_key_dict, ) ) - + verbose_proxy_logger.info( f"Stored batch managed object with unified_object_id={unified_object_id}, batch_id={model_object_id}" ) else: - verbose_proxy_logger.warning("Managed files hook not available, cannot store batch object for cost tracking") - + verbose_proxy_logger.warning( + "Managed files hook not available, cannot store batch object for cost tracking" + ) + except Exception as e: verbose_proxy_logger.error(f"Error storing batch managed object: {e}") @staticmethod def get_actual_model_id_from_router(model_name: str) -> str: from litellm.proxy.proxy_server import llm_router - + if llm_router is not None: # Try to find the model in the router by the extracted model name - extracted_model_name = VertexPassthroughLoggingHandler.extract_model_name_from_vertex_path(model_name) - + extracted_model_name = ( + VertexPassthroughLoggingHandler.extract_model_name_from_vertex_path( + model_name + ) + ) + # Use the existing get_model_ids method from router model_ids = llm_router.get_model_ids(model_name=extracted_model_name) if model_ids and len(model_ids) > 0: # Use the first model ID found actual_model_id = model_ids[0] - verbose_proxy_logger.info(f"Found model ID in router: {actual_model_id}") + verbose_proxy_logger.info( + f"Found model ID in router: {actual_model_id}" + ) return actual_model_id else: # Fallback to constructed model name actual_model_id = extracted_model_name - verbose_proxy_logger.warning(f"Model not found in router, using constructed name: {actual_model_id}") + verbose_proxy_logger.warning( + f"Model not found in router, using constructed name: {actual_model_id}" + ) return actual_model_id else: # Fallback if router is not available - extracted_model_name = VertexPassthroughLoggingHandler.extract_model_name_from_vertex_path(model_name) - verbose_proxy_logger.warning(f"Router not available, using constructed model name: {extracted_model_name}") + extracted_model_name = ( + VertexPassthroughLoggingHandler.extract_model_name_from_vertex_path( + model_name + ) + ) + verbose_proxy_logger.warning( + f"Router not available, using constructed model name: {extracted_model_name}" + ) return extracted_model_name - diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 9173758e2c..9e287c2bec 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -404,7 +404,10 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): headers=headers, params=requested_query_params, ) - elif HttpPassThroughEndpointHelpers.is_multipart(request) is True and not _parsed_body: + elif ( + HttpPassThroughEndpointHelpers.is_multipart(request) is True + and not _parsed_body + ): # Only use multipart handler if we don't have a parsed body # (parsed body means it was JSON despite multipart content-type header) return await HttpPassThroughEndpointHelpers.make_multipart_http_request( @@ -681,8 +684,10 @@ async def pass_through_request( # noqa: PLR0915 # Skip body parsing for multipart requests - make_multipart_http_request will handle it # But if custom_body is provided (e.g., JSON parsed despite multipart content-type), use it - is_multipart = HttpPassThroughEndpointHelpers.is_multipart(request) and not custom_body - + is_multipart = ( + HttpPassThroughEndpointHelpers.is_multipart(request) and not custom_body + ) + if custom_body: _parsed_body = custom_body elif is_multipart: @@ -1133,7 +1138,9 @@ def create_pass_through_route( fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), subpath: str = "", # captures sub-paths when include_subpath=True - custom_body: Optional[dict] = None, # caller-supplied body takes precedence over request-parsed body + custom_body: Optional[ + dict + ] = None, # caller-supplied body takes precedence over request-parsed body ): from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( InitPassThroughEndpointHelpers, @@ -2116,11 +2123,7 @@ class InitPassThroughEndpointHelpers: # If path matches and method filter is provided, check if method is allowed if path_matches: - if ( - method is None - or not route_methods - or method in route_methods - ): + if method is None or not route_methods or method in route_methods: return _registered_pass_through_routes[key] return None diff --git a/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py b/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py index a32659e45b..ae2f8edc74 100644 --- a/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py +++ b/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py @@ -134,9 +134,9 @@ class PassthroughEndpointRouter: vertex_location=location, vertex_credentials=vertex_credentials, ) - self.deployment_key_to_vertex_credentials[deployment_key] = ( - vertex_pass_through_credentials - ) + self.deployment_key_to_vertex_credentials[ + deployment_key + ] = vertex_pass_through_credentials def _get_deployment_key( self, project_id: Optional[str], location: Optional[str] @@ -156,10 +156,10 @@ class PassthroughEndpointRouter: """ if litellm.vector_store_registry is None: return None - vector_store_to_run: Optional[LiteLLM_ManagedVectorStore] = ( - litellm.vector_store_registry.get_litellm_managed_vector_store_from_registry( - vector_store_id=vector_store_id - ) + vector_store_to_run: Optional[ + LiteLLM_ManagedVectorStore + ] = litellm.vector_store_registry.get_litellm_managed_vector_store_from_registry( + vector_store_id=vector_store_id ) return vector_store_to_run diff --git a/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py b/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py index c37703b9df..5683491fed 100644 --- a/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py +++ b/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py @@ -28,12 +28,12 @@ PassThroughGuardrailsConfigInput = Union[ class PassthroughGuardrailHandler: """ Handles guardrail execution for passthrough endpoints. - + Passthrough endpoints use an opt-in model for guardrails: - Guardrails only run when explicitly configured on the endpoint - Supports field-level targeting using JSONPath expressions - Automatically inherits org/team/key level guardrails when enabled - + Guardrails can be specified as: - List format (simple): ["guardrail-1", "guardrail-2"] - Dict format (with settings): {"guardrail-1": {"request_fields": ["query"]}} @@ -45,7 +45,7 @@ class PassthroughGuardrailHandler: ) -> Optional[PassThroughGuardrailsConfig]: """ Normalize guardrails config to dict format. - + Accepts: - List of guardrail names: ["g1", "g2"] -> {"g1": None, "g2": None} - Dict with settings: {"g1": {"request_fields": [...]}} @@ -53,15 +53,15 @@ class PassthroughGuardrailHandler: """ if guardrails_config is None: return None - + # Already a dict - return as-is if isinstance(guardrails_config, dict): return guardrails_config - + # List of guardrail names - convert to dict if isinstance(guardrails_config, list): return {name: None for name in guardrails_config} - + verbose_proxy_logger.debug( "Passthrough guardrails config is not a dict or list, got: %s", type(guardrails_config), @@ -74,8 +74,8 @@ class PassthroughGuardrailHandler: ) -> bool: """ Check if guardrails are enabled for a passthrough endpoint. - - Passthrough endpoints are opt-in only - guardrails only run when + + Passthrough endpoints are opt-in only - guardrails only run when the guardrails config is set with at least one guardrail. """ normalized = PassthroughGuardrailHandler.normalize_config(guardrails_config) @@ -102,14 +102,14 @@ class PassthroughGuardrailHandler: normalized = PassthroughGuardrailHandler.normalize_config(guardrails_config) if normalized is None: return None - + settings = normalized.get(guardrail_name) if settings is None: return None - + if isinstance(settings, dict): return PassThroughGuardrailSettings(**settings) - + return settings @staticmethod @@ -119,14 +119,15 @@ class PassthroughGuardrailHandler: ) -> str: """ Prepare input text for guardrail execution based on field targeting settings. - + If request_fields is specified, extracts only those fields. Otherwise, uses the entire request payload as text. """ if guardrail_settings is None or guardrail_settings.request_fields is None: from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + return safe_dumps(request_data) - + return JsonPathExtractor.extract_fields( data=request_data, jsonpath_expressions=guardrail_settings.request_fields, @@ -139,14 +140,15 @@ class PassthroughGuardrailHandler: ) -> str: """ Prepare output text for guardrail execution based on field targeting settings. - + If response_fields is specified, extracts only those fields. Otherwise, uses the entire response payload as text. """ if guardrail_settings is None or guardrail_settings.response_fields is None: from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + return safe_dumps(response_data) - + return JsonPathExtractor.extract_fields( data=response_data, jsonpath_expressions=guardrail_settings.response_fields, @@ -161,18 +163,18 @@ class PassthroughGuardrailHandler: ) -> dict: """ Execute guardrails for a passthrough endpoint. - + This is the main entry point for passthrough guardrail execution. - + Args: request_data: The request payload user_api_key_dict: User API key authentication info guardrails_config: Passthrough-specific guardrails configuration event_type: "pre_call" for request, "post_call" for response - + Returns: The potentially modified request_data - + Raises: HTTPException if a guardrail blocks the request """ @@ -181,14 +183,14 @@ class PassthroughGuardrailHandler: "Passthrough guardrails not enabled, skipping guardrail execution" ) return request_data - + guardrail_names = PassthroughGuardrailHandler.get_guardrail_names( guardrails_config ) verbose_proxy_logger.debug( "Executing passthrough guardrails: %s", guardrail_names ) - + # Add to request metadata so guardrails know which to run from litellm.proxy.pass_through_endpoints.passthrough_context import ( set_passthrough_guardrails_config, @@ -196,15 +198,15 @@ class PassthroughGuardrailHandler: if "metadata" not in request_data: request_data["metadata"] = {} - + # Set guardrails in metadata using dict format for compatibility request_data["metadata"]["guardrails"] = { name: True for name in guardrail_names } - + # Store passthrough guardrails config in request-scoped context set_passthrough_guardrails_config(guardrails_config) - + return request_data @staticmethod @@ -297,15 +299,15 @@ class PassthroughGuardrailHandler: ) -> Optional[str]: """ Get the text to check for a guardrail, respecting field targeting settings. - + Called by guardrail hooks to get the appropriate text based on passthrough field targeting configuration. - + Args: data: The request/response data dict guardrail_name: Name of the guardrail being executed is_request: True for request (pre_call), False for response (post_call) - + Returns: The text to check, or None to use default behavior """ @@ -316,18 +318,18 @@ class PassthroughGuardrailHandler: passthrough_config = get_passthrough_guardrails_config() if passthrough_config is None: return None - + settings = PassthroughGuardrailHandler.get_settings( passthrough_config, guardrail_name ) if settings is None: return None - + if is_request: if settings.request_fields: return JsonPathExtractor.extract_fields(data, settings.request_fields) else: if settings.response_fields: return JsonPathExtractor.extract_fields(data, settings.response_fields) - + return None diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 1e7118f447..302d7e76ed 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -60,18 +60,14 @@ class PassThroughStreamingHandler: if endpoint_type == EndpointType.VERTEX_AI: # Only handle streamRawPredict (uses Anthropic format) if "streamRawPredict" in url_route or "rawPredict" in url_route: - modified_chunk = ( - ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( - chunk, model_name - ) + modified_chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( + chunk, model_name ) if modified_chunk is not None: chunk = modified_chunk elif endpoint_type == EndpointType.ANTHROPIC: - modified_chunk = ( - ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( - chunk, model_name - ) + modified_chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( + chunk, model_name ) if modified_chunk is not None: chunk = modified_chunk @@ -141,35 +137,31 @@ class PassThroughStreamingHandler: ) kwargs = anthropic_passthrough_logging_handler_result["kwargs"] elif endpoint_type == EndpointType.VERTEX_AI: - vertex_passthrough_logging_handler_result = ( - VertexPassthroughLoggingHandler._handle_logging_vertex_collected_chunks( - litellm_logging_obj=litellm_logging_obj, - passthrough_success_handler_obj=passthrough_success_handler_obj, - url_route=url_route, - request_body=request_body, - endpoint_type=endpoint_type, - start_time=start_time, - all_chunks=all_chunks, - end_time=end_time, - model=model, - ) + vertex_passthrough_logging_handler_result = VertexPassthroughLoggingHandler._handle_logging_vertex_collected_chunks( + litellm_logging_obj=litellm_logging_obj, + passthrough_success_handler_obj=passthrough_success_handler_obj, + url_route=url_route, + request_body=request_body, + endpoint_type=endpoint_type, + start_time=start_time, + all_chunks=all_chunks, + end_time=end_time, + model=model, ) standard_logging_response_object = ( vertex_passthrough_logging_handler_result["result"] ) kwargs = vertex_passthrough_logging_handler_result["kwargs"] elif endpoint_type == EndpointType.OPENAI: - openai_passthrough_logging_handler_result = ( - OpenAIPassthroughLoggingHandler._handle_logging_openai_collected_chunks( - litellm_logging_obj=litellm_logging_obj, - passthrough_success_handler_obj=passthrough_success_handler_obj, - url_route=url_route, - request_body=request_body, - endpoint_type=endpoint_type, - start_time=start_time, - all_chunks=all_chunks, - end_time=end_time, - ) + openai_passthrough_logging_handler_result = OpenAIPassthroughLoggingHandler._handle_logging_openai_collected_chunks( + litellm_logging_obj=litellm_logging_obj, + passthrough_success_handler_obj=passthrough_success_handler_obj, + url_route=url_route, + request_body=request_body, + endpoint_type=endpoint_type, + start_time=start_time, + all_chunks=all_chunks, + end_time=end_time, ) standard_logging_response_object = ( openai_passthrough_logging_handler_result["result"] @@ -187,7 +179,10 @@ class PassThroughStreamingHandler: cache_hit=False, **kwargs, ) - if litellm_logging_obj._should_run_sync_callbacks_for_async_calls() is False: + if ( + litellm_logging_obj._should_run_sync_callbacks_for_async_calls() + is False + ): return executor.submit( diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index 271c2d7a48..33819b888d 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -61,10 +61,19 @@ class PassThroughEndpointLogging: self.TRACKED_LANGFUSE_ROUTES = ["/langfuse/"] # Gemini - self.TRACKED_GEMINI_ROUTES = ["generateContent", "streamGenerateContent", "predictLongRunning"] + self.TRACKED_GEMINI_ROUTES = [ + "generateContent", + "streamGenerateContent", + "predictLongRunning", + ] # Cursor Cloud Agents - self.TRACKED_CURSOR_ROUTES = ["/v0/agents", "/v0/me", "/v0/models", "/v0/repositories"] + self.TRACKED_CURSOR_ROUTES = [ + "/v0/agents", + "/v0/me", + "/v0/models", + "/v0/repositories", + ] # Vertex AI Live API WebSocket self.TRACKED_VERTEX_AI_LIVE_ROUTES = ["/vertex_ai/live"] @@ -274,9 +283,9 @@ class PassThroughEndpointLogging: standard_logging_response_object = vertex_ai_live_handler_result["result"] kwargs = vertex_ai_live_handler_result["kwargs"] - return_dict["standard_logging_response_object"] = ( - standard_logging_response_object - ) + return_dict[ + "standard_logging_response_object" + ] = standard_logging_response_object return_dict["kwargs"] = kwargs return return_dict @@ -299,9 +308,9 @@ class PassThroughEndpointLogging: standard_logging_response_object: Optional[ PassThroughEndpointLoggingResultValues ] = None - logging_obj.model_call_details["passthrough_logging_payload"] = ( - passthrough_logging_payload - ) + logging_obj.model_call_details[ + "passthrough_logging_payload" + ] = passthrough_logging_payload if self.is_assemblyai_route(url_route): if ( AssemblyAIPassthroughLoggingHandler._should_log_request( @@ -478,8 +487,8 @@ class PassThroughEndpointLogging: kwargs["response_cost"] = passthrough_logging_payload.get( "cost_per_request" ) - logging_obj.model_call_details["response_cost"] = ( - passthrough_logging_payload.get("cost_per_request") - ) + logging_obj.model_call_details[ + "response_cost" + ] = passthrough_logging_payload.get("cost_per_request") return kwargs diff --git a/litellm/proxy/policy_engine/attachment_registry.py b/litellm/proxy/policy_engine/attachment_registry.py index 69b3b3599f..530e1fca1f 100644 --- a/litellm/proxy/policy_engine/attachment_registry.py +++ b/litellm/proxy/policy_engine/attachment_registry.py @@ -60,9 +60,7 @@ class AttachmentRegistry: f"Loaded attachment for policy: {attachment.policy}" ) except Exception as e: - verbose_proxy_logger.error( - f"Error loading attachment: {str(e)}" - ) + verbose_proxy_logger.error(f"Error loading attachment: {str(e)}") raise ValueError(f"Invalid attachment: {str(e)}") from e self._initialized = True @@ -97,7 +95,9 @@ class AttachmentRegistry: Returns: List of policy names that are attached to matching scopes """ - return [r["policy_name"] for r in self.get_attached_policies_with_reasons(context)] + return [ + r["policy_name"] for r in self.get_attached_policies_with_reasons(context) + ] def get_attached_policies_with_reasons( self, context: PolicyMatchContext @@ -146,7 +146,8 @@ class AttachmentRegistry: reasons = [] if attachment.tags and context.tags: matching_tags = [ - t for t in context.tags + t + for t in context.tags if PolicyMatcher.matches_pattern(t, attachment.tags) ] if matching_tags: @@ -160,9 +161,7 @@ class AttachmentRegistry: return "+".join(reasons) if reasons else "scope:default" - def is_policy_attached( - self, policy_name: str, context: PolicyMatchContext - ) -> bool: + def is_policy_attached(self, policy_name: str, context: PolicyMatchContext) -> bool: """ Check if a specific policy is attached to the given context. @@ -465,9 +464,13 @@ class AttachmentRegistry: attachment = PolicyAttachment( policy=attachment_response.policy_name, scope=attachment_response.scope, - teams=attachment_response.teams if attachment_response.teams else None, + teams=attachment_response.teams + if attachment_response.teams + else None, keys=attachment_response.keys if attachment_response.keys else None, - models=attachment_response.models if attachment_response.models else None, + models=attachment_response.models + if attachment_response.models + else None, tags=attachment_response.tags if attachment_response.tags else None, ) self._attachments.append(attachment) diff --git a/litellm/proxy/policy_engine/init_policies.py b/litellm/proxy/policy_engine/init_policies.py index b734c0cb5c..3167a0fe8b 100644 --- a/litellm/proxy/policy_engine/init_policies.py +++ b/litellm/proxy/policy_engine/init_policies.py @@ -44,8 +44,12 @@ def _print_policies_on_startup( condition = policy_data.get("condition") description = policy_data.get("description") - guardrails_add = guardrails.get("add", []) if isinstance(guardrails, dict) else [] - guardrails_remove = guardrails.get("remove", []) if isinstance(guardrails, dict) else [] + guardrails_add = ( + guardrails.get("add", []) if isinstance(guardrails, dict) else [] + ) + guardrails_remove = ( + guardrails.get("remove", []) if isinstance(guardrails, dict) else [] + ) inherit_str = f" (inherits: {inherit})" if inherit else "" print( # noqa: T201 @@ -58,7 +62,9 @@ def _print_policies_on_startup( if guardrails_remove: print(f" guardrails.remove: {guardrails_remove}") # noqa: T201 if condition: - model_condition = condition.get("model") if isinstance(condition, dict) else None + model_condition = ( + condition.get("model") if isinstance(condition, dict) else None + ) if model_condition: print(f" condition.model: {model_condition}") # noqa: T201 @@ -258,19 +264,23 @@ def get_policies_summary() -> Dict[str, Any]: "description": policy.description if policy else None, "guardrails_add": policy.guardrails.get_add() if policy else [], "guardrails_remove": policy.guardrails.get_remove() if policy else [], - "condition": policy.condition.model_dump() if policy and policy.condition else None, + "condition": policy.condition.model_dump() + if policy and policy.condition + else None, "resolved_guardrails": resolved_policy.guardrails, "inheritance_chain": resolved_policy.inheritance_chain, } # Add attachment info for attachment in attachment_registry.get_all_attachments(): - summary["attachments"].append({ - "policy": attachment.policy, - "scope": attachment.scope, - "teams": attachment.teams, - "keys": attachment.keys, - "models": attachment.models, - }) + summary["attachments"].append( + { + "policy": attachment.policy, + "scope": attachment.scope, + "teams": attachment.teams, + "keys": attachment.keys, + "models": attachment.models, + } + ) return summary diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index c1e2d76e7c..729b42ce63 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -114,7 +114,8 @@ class PipelineExecutor: return PipelineExecutionResult( terminal_action="modify_response", step_results=step_results, - modify_response_message=step.modify_response_message or error_detail, + modify_response_message=step.modify_response_message + or error_detail, ) # action == "next" → continue to next step diff --git a/litellm/proxy/policy_engine/policy_endpoints.py b/litellm/proxy/policy_engine/policy_endpoints.py index d8de028d6a..3a25e249a4 100644 --- a/litellm/proxy/policy_engine/policy_endpoints.py +++ b/litellm/proxy/policy_engine/policy_endpoints.py @@ -11,17 +11,24 @@ from fastapi import APIRouter, Depends, HTTPException from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.policy_engine.attachment_registry import \ - get_attachment_registry +from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor from litellm.proxy.policy_engine.policy_registry import get_policy_registry from litellm.types.proxy.policy_engine import ( - GuardrailPipeline, PipelineTestRequest, PolicyAttachmentCreateRequest, - PolicyAttachmentDBResponse, PolicyAttachmentListResponse, - PolicyCreateRequest, PolicyDBResponse, PolicyListDBResponse, - PolicyUpdateRequest, PolicyVersionCompareResponse, - PolicyVersionCreateRequest, PolicyVersionListResponse, - PolicyVersionStatusUpdateRequest) + GuardrailPipeline, + PipelineTestRequest, + PolicyAttachmentCreateRequest, + PolicyAttachmentDBResponse, + PolicyAttachmentListResponse, + PolicyCreateRequest, + PolicyDBResponse, + PolicyListDBResponse, + PolicyUpdateRequest, + PolicyVersionCompareResponse, + PolicyVersionCreateRequest, + PolicyVersionListResponse, + PolicyVersionStatusUpdateRequest, +) router = APIRouter() @@ -253,7 +260,11 @@ async def update_policy_version_status( raise except Exception as e: verbose_proxy_logger.exception(f"Error updating version status: {e}") - if "invalid status" in str(e).lower() or "only draft" in str(e).lower() or "cannot promote" in str(e).lower(): + if ( + "invalid status" in str(e).lower() + or "only draft" in str(e).lower() + or "cannot promote" in str(e).lower() + ): raise HTTPException(status_code=400, detail=str(e)) if "not found" in str(e).lower(): raise HTTPException(status_code=404, detail=str(e)) diff --git a/litellm/proxy/policy_engine/policy_matcher.py b/litellm/proxy/policy_engine/policy_matcher.py index 888981f85f..b2788e6355 100644 --- a/litellm/proxy/policy_engine/policy_matcher.py +++ b/litellm/proxy/policy_engine/policy_matcher.py @@ -89,8 +89,7 @@ class PolicyMatcher: return False # Match if ANY context tag matches ANY scope tag pattern if not any( - PolicyMatcher.matches_pattern(tag, scope_tags) - for tag in context.tags + PolicyMatcher.matches_pattern(tag, scope_tags) for tag in context.tags ): return False diff --git a/litellm/proxy/policy_engine/policy_registry.py b/litellm/proxy/policy_engine/policy_registry.py index f8e1ebd7ba..d3df16afde 100644 --- a/litellm/proxy/policy_engine/policy_registry.py +++ b/litellm/proxy/policy_engine/policy_registry.py @@ -12,14 +12,18 @@ from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple from litellm._logging import verbose_proxy_logger -from litellm.types.proxy.policy_engine import (GuardrailPipeline, PipelineStep, - Policy, PolicyCondition, - PolicyCreateRequest, - PolicyDBResponse, - PolicyGuardrails, - PolicyUpdateRequest, - PolicyVersionCompareResponse, - PolicyVersionListResponse) +from litellm.types.proxy.policy_engine import ( + GuardrailPipeline, + PipelineStep, + Policy, + PolicyCondition, + PolicyCreateRequest, + PolicyDBResponse, + PolicyGuardrails, + PolicyUpdateRequest, + PolicyVersionCompareResponse, + PolicyVersionListResponse, +) if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient @@ -468,7 +472,9 @@ class PolicyRegistry: verbose_proxy_logger.exception(f"Error getting policy from DB: {e}") raise Exception(f"Error getting policy from DB: {str(e)}") - def get_policy_by_id_for_request(self, policy_id: str) -> Optional[Tuple[str, Policy]]: + def get_policy_by_id_for_request( + self, policy_id: str + ) -> Optional[Tuple[str, Policy]]: """ Return a policy version by ID from in-memory cache (no DB access). diff --git a/litellm/proxy/policy_engine/policy_resolve_endpoints.py b/litellm/proxy/policy_engine/policy_resolve_endpoints.py index 318e990ff1..54374d90a1 100644 --- a/litellm/proxy/policy_engine/policy_resolve_endpoints.py +++ b/litellm/proxy/policy_engine/policy_resolve_endpoints.py @@ -77,7 +77,9 @@ def _get_tags_from_metadata(metadata: object, json_metadata: object = None) -> l async def _fetch_all_teams(prisma_client: object) -> list: """Fetch teams from DB once. Reuse the result across tag and alias lookups.""" return await prisma_client.db.litellm_teamtable.find_many( # type: ignore - where={}, order={"created_at": "desc"}, take=MAX_POLICY_ESTIMATE_IMPACT_ROWS, + where={}, + order={"created_at": "desc"}, + take=MAX_POLICY_ESTIMATE_IMPACT_ROWS, ) @@ -159,7 +161,8 @@ async def _find_affected_by_team_patterns( if matched_team_ids: keys = await prisma_client.db.litellm_verificationtoken.find_many( # type: ignore where={"team_id": {"in": matched_team_ids}}, - order={"created_at": "desc"}, take=MAX_POLICY_ESTIMATE_IMPACT_ROWS, + order={"created_at": "desc"}, + take=MAX_POLICY_ESTIMATE_IMPACT_ROWS, ) for key in keys: key_alias = key.key_alias or "" @@ -181,7 +184,8 @@ async def _find_affected_keys_by_alias( keys = await prisma_client.db.litellm_verificationtoken.find_many( # type: ignore where=_build_alias_where("key_alias", key_patterns), - order={"created_at": "desc"}, take=MAX_POLICY_ESTIMATE_IMPACT_ROWS, + order={"created_at": "desc"}, + take=MAX_POLICY_ESTIMATE_IMPACT_ROWS, ) for key in keys: key_alias = key.key_alias or "" @@ -364,19 +368,24 @@ async def estimate_attachment_impact( # Tag-based impact if tag_patterns: keys = await prisma_client.db.litellm_verificationtoken.find_many( # type: ignore - where={}, order={"created_at": "desc"}, + where={}, + order={"created_at": "desc"}, take=MAX_POLICY_ESTIMATE_IMPACT_ROWS, ) affected_keys, unnamed_keys = _filter_keys_by_tags(keys, tag_patterns) affected_teams, unnamed_teams = _filter_teams_by_tags( - all_teams, tag_patterns, + all_teams, + tag_patterns, ) # Team-based impact (alias matching + keys belonging to those teams) if team_patterns: new_teams, new_keys, new_unnamed = await _find_affected_by_team_patterns( - prisma_client, all_teams, team_patterns, - affected_teams, affected_keys, + prisma_client, + all_teams, + team_patterns, + affected_teams, + affected_keys, ) affected_teams.extend(new_teams) affected_keys.extend(new_keys) @@ -386,7 +395,9 @@ async def estimate_attachment_impact( key_patterns = request.keys or [] if key_patterns: new_keys = await _find_affected_keys_by_alias( - prisma_client, key_patterns, affected_keys, + prisma_client, + key_patterns, + affected_keys, ) affected_keys.extend(new_keys) diff --git a/litellm/proxy/policy_engine/policy_resolver.py b/litellm/proxy/policy_engine/policy_resolver.py index c802a970a8..2c0a5334b0 100644 --- a/litellm/proxy/policy_engine/policy_resolver.py +++ b/litellm/proxy/policy_engine/policy_resolver.py @@ -11,9 +11,12 @@ Handles: from typing import Dict, List, Optional, Set, Tuple from litellm._logging import verbose_proxy_logger -from litellm.types.proxy.policy_engine import (GuardrailPipeline, Policy, - PolicyMatchContext, - ResolvedPolicy) +from litellm.types.proxy.policy_engine import ( + GuardrailPipeline, + Policy, + PolicyMatchContext, + ResolvedPolicy, +) class PolicyResolver: @@ -87,8 +90,7 @@ class PolicyResolver: Returns: ResolvedPolicy with final guardrails list """ - from litellm.proxy.policy_engine.condition_evaluator import \ - ConditionEvaluator + from litellm.proxy.policy_engine.condition_evaluator import ConditionEvaluator inheritance_chain = PolicyResolver.resolve_inheritance_chain( policy_name=policy_name, policies=policies @@ -152,8 +154,7 @@ class PolicyResolver: List of guardrail names to apply """ from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher - from litellm.proxy.policy_engine.policy_registry import \ - get_policy_registry + from litellm.proxy.policy_engine.policy_registry import get_policy_registry if policies is None: registry = get_policy_registry() @@ -190,9 +191,7 @@ class PolicyResolver: ) result = list(all_guardrails) - verbose_proxy_logger.debug( - f"Final guardrails for context: {result}" - ) + verbose_proxy_logger.debug(f"Final guardrails for context: {result}") return result @@ -218,8 +217,7 @@ class PolicyResolver: List of (policy_name, GuardrailPipeline) tuples """ from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher - from litellm.proxy.policy_engine.policy_registry import \ - get_policy_registry + from litellm.proxy.policy_engine.policy_registry import get_policy_registry if policies is None: registry = get_policy_registry() @@ -281,8 +279,7 @@ class PolicyResolver: Returns: Dictionary mapping policy names to ResolvedPolicy objects """ - from litellm.proxy.policy_engine.policy_registry import \ - get_policy_registry + from litellm.proxy.policy_engine.policy_registry import get_policy_registry if policies is None: registry = get_policy_registry() diff --git a/litellm/proxy/policy_engine/policy_validator.py b/litellm/proxy/policy_engine/policy_validator.py index 89c9b0e2e9..b587e3432b 100644 --- a/litellm/proxy/policy_engine/policy_validator.py +++ b/litellm/proxy/policy_engine/policy_validator.py @@ -70,7 +70,11 @@ class PolicyValidator: ) guardrails = IN_MEMORY_GUARDRAIL_HANDLER.list_in_memory_guardrails() - return {g.get("guardrail_name", "") for g in guardrails if g.get("guardrail_name")} + return { + g.get("guardrail_name", "") + for g in guardrails + if g.get("guardrail_name") + } except Exception as e: verbose_proxy_logger.warning( f"Could not get guardrails from registry: {str(e)}" @@ -145,17 +149,17 @@ class PolicyValidator: # Check if model matches any pattern via pattern router if hasattr(self.llm_router, "pattern_router"): - pattern_deployments = self.llm_router.pattern_router.get_deployments_by_pattern( - model=model + pattern_deployments = ( + self.llm_router.pattern_router.get_deployments_by_pattern( + model=model + ) ) if pattern_deployments: return True return False except Exception as e: - verbose_proxy_logger.warning( - f"Could not check model '{model}': {str(e)}" - ) + verbose_proxy_logger.warning(f"Could not check model '{model}': {str(e)}") return True # Assume valid on error def _validate_inheritance_chain( diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index be2f5ac7c1..153c3c2dba 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -45,7 +45,9 @@ def append_query_params(url: Optional[str], params: dict) -> str: if not isinstance(url, str) or url == "": # Preserve previous startup behavior when DATABASE_URL is absent. # Returning an empty string avoids urlparse type errors in test/dev flows. - verbose_proxy_logger.warning("append_query_params received empty or non-string URL, returning empty string") + verbose_proxy_logger.warning( + "append_query_params received empty or non-string URL, returning empty string" + ) return "" parsed_url = urlparse.urlparse(url) parsed_query = urlparse.parse_qs(parsed_url.query) @@ -347,9 +349,8 @@ class ProxyInitializationHelpers: from litellm.proxy.prometheus_cleanup import wipe_directory - multiproc_dir = ( - os.environ.get("PROMETHEUS_MULTIPROC_DIR") - or os.environ.get("prometheus_multiproc_dir") + multiproc_dir = os.environ.get("PROMETHEUS_MULTIPROC_DIR") or os.environ.get( + "prometheus_multiproc_dir" ) auto_created = not multiproc_dir @@ -853,7 +854,9 @@ def run_server( # noqa: PLR0915 ): check_prisma_schema_diff(db_url=None) else: - if not PrismaManager.setup_database(use_migrate=not use_prisma_db_push): + if not PrismaManager.setup_database( + use_migrate=not use_prisma_db_push + ): print( # noqa "\033[1;31mLiteLLM Proxy: Database setup failed after multiple retries. " "The proxy cannot start safely. Please check your database connection and migration status.\033[0m" diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ff30aa91b2..e01ee8e9d9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4604,9 +4604,7 @@ class ProxyConfig: ) ) - async def _init_hashicorp_vault_config_override( - self, prisma_client: PrismaClient - ): + async def _init_hashicorp_vault_config_override(self, prisma_client: PrismaClient): """ Load Hashicorp Vault config override from DB. Decrypts sensitive fields, sets HCP_VAULT_* env vars, and reinitializes the secret manager. @@ -4645,18 +4643,14 @@ class ProxyConfig: # Reinitialize the secret manager try: - self.initialize_secret_manager( - key_management_system="hashicorp_vault" - ) + self.initialize_secret_manager(key_management_system="hashicorp_vault") except Exception: # Restore previous working env vars instead of wiping all _set_env_vars(previous_env) raise self._last_hashicorp_vault_config = config_data.copy() - verbose_proxy_logger.debug( - "Hashicorp Vault config override loaded from DB" - ) + verbose_proxy_logger.debug("Hashicorp Vault config override loaded from DB") except Exception as e: verbose_proxy_logger.exception( "Error loading Hashicorp Vault config override from DB: %s", @@ -4756,7 +4750,14 @@ class ProxyConfig: } ), }, - "update": {"param_value": safe_dumps({"interval_hours": interval_hours, "force_reload": False})}, + "update": { + "param_value": safe_dumps( + { + "interval_hours": interval_hours, + "force_reload": False, + } + ) + }, }, ) @@ -4857,7 +4858,14 @@ class ProxyConfig: } ), }, - "update": {"param_value": safe_dumps({"interval_hours": interval_hours, "force_reload": False})}, + "update": { + "param_value": safe_dumps( + { + "interval_hours": interval_hours, + "force_reload": False, + } + ) + }, }, ) @@ -5429,9 +5437,7 @@ def _restamp_streaming_chunk_model( return chunk, model_mismatch_logged # For Azure Model Router, preserve the actual model used in each chunk - if _is_azure_model_router_request( - requested_model_from_client - ): + if _is_azure_model_router_request(requested_model_from_client): return chunk, model_mismatch_logged downstream_model = ( @@ -5643,9 +5649,9 @@ class ProxyStartupEvent: """ from litellm.secret_managers.main import str_to_bool - _use_redis_transaction_buffer: Optional[Union[bool, str]] = ( - general_settings.get("use_redis_transaction_buffer", False) - ) + _use_redis_transaction_buffer: Optional[ + Union[bool, str] + ] = general_settings.get("use_redis_transaction_buffer", False) if isinstance(_use_redis_transaction_buffer, str): _use_redis_transaction_buffer = str_to_bool(_use_redis_transaction_buffer) @@ -12523,7 +12529,11 @@ async def reload_model_cost_map( {"interval_hours": None, "force_reload": True} ), }, - "update": {"param_value": safe_dumps({"interval_hours": existing_interval, "force_reload": True})}, + "update": { + "param_value": safe_dumps( + {"interval_hours": existing_interval, "force_reload": True} + ) + }, }, ) @@ -12858,7 +12868,9 @@ async def reload_anthropic_beta_headers( ) existing_beta_interval = None if existing_beta_config and existing_beta_config.param_value: - existing_beta_interval = existing_beta_config.param_value.get("interval_hours") + existing_beta_interval = existing_beta_config.param_value.get( + "interval_hours" + ) await prisma_client.db.litellm_config.upsert( where={"param_name": "anthropic_beta_headers_reload_config"}, @@ -12869,7 +12881,11 @@ async def reload_anthropic_beta_headers( {"interval_hours": None, "force_reload": True} ), }, - "update": {"param_value": safe_dumps({"interval_hours": existing_beta_interval, "force_reload": True})}, + "update": { + "param_value": safe_dumps( + {"interval_hours": existing_beta_interval, "force_reload": True} + ) + }, }, ) diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index a74b9a40a1..eb9ed59055 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -40,8 +40,14 @@ _ENDPOINT_METADATA: Dict[str, Dict[str, str]] = { "messages": {"label": "Messages", "endpoint": "/messages"}, "responses": {"label": "Responses", "endpoint": "/responses"}, "embeddings": {"label": "Embeddings", "endpoint": "/embeddings"}, - "image_generations": {"label": "Image Generations", "endpoint": "/images/generations"}, - "audio_transcriptions": {"label": "Audio Transcriptions", "endpoint": "/audio/transcriptions"}, + "image_generations": { + "label": "Image Generations", + "endpoint": "/images/generations", + }, + "audio_transcriptions": { + "label": "Audio Transcriptions", + "endpoint": "/audio/transcriptions", + }, "audio_speech": {"label": "Audio Speech", "endpoint": "/audio/speech"}, "moderations": {"label": "Moderations", "endpoint": "/moderations"}, "batches": {"label": "Batches", "endpoint": "/batches"}, @@ -52,14 +58,29 @@ _ENDPOINT_METADATA: Dict[str, Dict[str, str]] = { "interactions": {"label": "Interactions", "endpoint": "/interactions"}, "a2a": {"label": "A2A (Agent Gateway)", "endpoint": "/a2a/{agent}/message/send"}, "container": {"label": "Containers", "endpoint": "/containers"}, - "container_files": {"label": "Container Files", "endpoint": "/containers/{id}/files"}, + "container_files": { + "label": "Container Files", + "endpoint": "/containers/{id}/files", + }, "compact": {"label": "Compact", "endpoint": "/responses/compact"}, "files": {"label": "Files", "endpoint": "/files"}, "image_edits": {"label": "Image Edits", "endpoint": "/images/edits"}, - "vector_stores_create": {"label": "Vector Stores (Create)", "endpoint": "/vector_stores"}, - "vector_stores_search": {"label": "Vector Stores (Search)", "endpoint": "/vector_stores/{id}/search"}, - "vector_store_files": {"label": "Vector Store Files", "endpoint": "/vector_stores/{id}/files"}, - "video_generations": {"label": "Video Generations", "endpoint": "/videos/generations"}, + "vector_stores_create": { + "label": "Vector Stores (Create)", + "endpoint": "/vector_stores", + }, + "vector_stores_search": { + "label": "Vector Stores (Search)", + "endpoint": "/vector_stores/{id}/search", + }, + "vector_store_files": { + "label": "Vector Store Files", + "endpoint": "/vector_stores/{id}/files", + }, + "video_generations": { + "label": "Video Generations", + "endpoint": "/videos/generations", + }, "assistants": {"label": "Assistants", "endpoint": "/assistants"}, "fine_tuning": {"label": "Fine Tuning", "endpoint": "/fine_tuning/jobs"}, "text_completion": {"label": "Text Completion", "endpoint": "/completions"}, @@ -110,7 +131,9 @@ def _build_endpoints(raw: Dict[str, Any]) -> List[Dict[str, Any]]: for slug, pd in providers.items() if pd.get("endpoints", {}).get(key) ] - result.append({"key": key, "label": label, "endpoint": path, "providers": supporting}) + result.append( + {"key": key, "label": label, "endpoint": path, "providers": supporting} + ) return result @@ -135,8 +158,14 @@ def _load_endpoints() -> List[Dict[str, Any]]: ) async def public_model_hub(): import litellm - from litellm.proxy.proxy_server import _get_model_group_info, llm_router, prisma_client - from litellm.proxy.health_endpoints._health_endpoints import _convert_health_check_to_dict + from litellm.proxy.proxy_server import ( + _get_model_group_info, + llm_router, + prisma_client, + ) + from litellm.proxy.health_endpoints._health_endpoints import ( + _convert_health_check_to_dict, + ) if llm_router is None: raise HTTPException( @@ -269,7 +298,7 @@ async def get_provider_fields() -> List[ProviderCreateInfo]: os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "proxy", "public_endpoints", - "provider_create_fields.json" + "provider_create_fields.json", ) with open(provider_create_fields_path, "r") as f: @@ -383,7 +412,9 @@ async def get_agent_fields() -> List[AgentCreateInfo]: field_copy["include_in_litellm_params"] = True inherited_fields.append(field_copy) # Append provider credential fields after agent's own fields - agent["credential_fields"] = agent.get("credential_fields", []) + inherited_fields + agent["credential_fields"] = ( + agent.get("credential_fields", []) + inherited_fields + ) # Remove the inherit field from response (not needed by frontend) agent.pop("inherit_credentials_from_provider", None) diff --git a/litellm/proxy/rag_endpoints/__init__.py b/litellm/proxy/rag_endpoints/__init__.py index 4586e4ec72..89c4876069 100644 --- a/litellm/proxy/rag_endpoints/__init__.py +++ b/litellm/proxy/rag_endpoints/__init__.py @@ -3,4 +3,3 @@ from litellm.proxy.rag_endpoints.endpoints import router __all__ = ["router"] - diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 408a017967..76136c12be 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -33,12 +33,12 @@ def _build_file_metadata_entry( ) -> Dict[str, Any]: """ Build a file metadata entry for storing in vector_store_metadata. - + Args: response: The response from litellm.aingest containing file_id file_data: Optional tuple of (filename, content, content_type) file_url: Optional URL if file was ingested from URL - + Returns: Dictionary with file metadata (file_id, filename, file_url, ingested_at, etc.) """ @@ -50,17 +50,17 @@ def _build_file_metadata_entry( file_id = response.get("file_id") elif hasattr(response, "file_id"): file_id = response.file_id - + # Extract file information from file_data tuple filename = None file_size = None content_type = None - + if file_data: filename = file_data[0] file_size = len(file_data[1]) if len(file_data) > 1 else None content_type = file_data[2] if len(file_data) > 2 else None - + # Build file metadata entry file_entry = { "file_id": file_id, @@ -68,13 +68,13 @@ def _build_file_metadata_entry( "file_url": file_url, "ingested_at": datetime.now(timezone.utc).isoformat(), } - + # Add optional fields if available if file_size is not None: file_entry["file_size"] = file_size if content_type is not None: file_entry["content_type"] = content_type - + return file_entry @@ -88,14 +88,14 @@ async def _save_vector_store_to_db_from_rag_ingest( ) -> None: """ Helper function to save a newly created vector store from RAG ingest to the database. - + This function: - Extracts vector store ID and config from the ingest response - Checks if the vector store already exists in the database - Creates a new database entry if it doesn't exist - Adds the vector store to the registry - Tracks team_id and user_id for access control - + Args: response: The response from litellm.aingest() ingest_options: The ingest options containing vector store config @@ -125,12 +125,14 @@ async def _save_vector_store_to_db_from_rag_ingest( vector_store_config = ingest_options.get("vector_store", {}) custom_llm_provider = vector_store_config.get("custom_llm_provider") - + # Extract litellm_vector_store_params for custom name and description litellm_vector_store_params = ingest_options.get("litellm_vector_store_params", {}) custom_vector_store_name = litellm_vector_store_params.get("vector_store_name") - custom_vector_store_description = litellm_vector_store_params.get("vector_store_description") - + custom_vector_store_description = litellm_vector_store_params.get( + "vector_store_description" + ) + # Extract provider-specific params from vector_store_config to save as litellm_params # This ensures params like aws_region_name, embedding_model, etc. are available for search provider_specific_params = {} @@ -161,13 +163,15 @@ async def _save_vector_store_to_db_from_rag_ingest( ) # Initialize metadata with first file - initial_metadata = { - "ingested_files": [file_entry] - } - + initial_metadata = {"ingested_files": [file_entry]} + # Use custom name if provided, otherwise default - vector_store_name = custom_vector_store_name or f"RAG Vector Store - {vector_store_id[:8]}" - vector_store_description = custom_vector_store_description or "Created via RAG ingest endpoint" + vector_store_name = ( + custom_vector_store_name or f"RAG Vector Store - {vector_store_id[:8]}" + ) + vector_store_description = ( + custom_vector_store_description or "Created via RAG ingest endpoint" + ) await create_vector_store_in_db( vector_store_id=vector_store_id, @@ -176,7 +180,9 @@ async def _save_vector_store_to_db_from_rag_ingest( vector_store_name=vector_store_name, vector_store_description=vector_store_description, vector_store_metadata=initial_metadata, - litellm_params=provider_specific_params if provider_specific_params else None, + litellm_params=provider_specific_params + if provider_specific_params + else None, team_id=user_api_key_dict.team_id, user_id=user_api_key_dict.user_id, ) @@ -188,24 +194,26 @@ async def _save_vector_store_to_db_from_rag_ingest( verbose_proxy_logger.info( f"Vector store {vector_store_id} already exists, appending file to metadata" ) - + # Update existing vector store with new file existing_metadata = existing_vector_store.vector_store_metadata or {} if isinstance(existing_metadata, str): import json + existing_metadata = json.loads(existing_metadata) - + ingested_files = existing_metadata.get("ingested_files", []) ingested_files.append(file_entry) existing_metadata["ingested_files"] = ingested_files - + # Update the vector store from litellm.proxy.utils import safe_dumps + await prisma_client.db.litellm_managedvectorstorestable.update( where={"vector_store_id": vector_store_id}, - data={"vector_store_metadata": safe_dumps(existing_metadata)} + data={"vector_store_metadata": safe_dumps(existing_metadata)}, ) - + verbose_proxy_logger.info( f"Added file {file_entry.get('filename') or file_entry.get('file_url', 'Unknown')} to vector store {vector_store_id} metadata" ) @@ -218,7 +226,9 @@ async def _save_vector_store_to_db_from_rag_ingest( async def parse_rag_ingest_request( request: Request, -) -> Tuple[Dict[str, Any], Optional[Tuple[str, bytes, str]], Optional[str], Optional[str]]: +) -> Tuple[ + Dict[str, Any], Optional[Tuple[str, bytes, str]], Optional[str], Optional[str] +]: """ Parse RAG ingest request. @@ -289,7 +299,9 @@ async def parse_rag_ingest_request( if "vector_store" not in ingest_options: raise HTTPException( status_code=400, - detail={"error": "ingest_options must contain 'vector_store' configuration"}, + detail={ + "error": "ingest_options must contain 'vector_store' configuration" + }, ) return ingest_options, file_data, file_url, file_id @@ -355,11 +367,14 @@ async def rag_ingest( try: # Parse request - ingest_options, file_data, file_url, file_id = await parse_rag_ingest_request(request) + ingest_options, file_data, file_url, file_id = await parse_rag_ingest_request( + request + ) # INTERNAL_USER_VIEW_ONLY can ingest to existing vector stores only if ( - user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value + user_api_key_dict.user_role + == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value and not ingest_options.get("vector_store", {}).get("vector_store_id") ): raise HTTPException( @@ -397,7 +412,7 @@ async def rag_ingest( verbose_proxy_logger.debug( f"RAG Ingest - Checking database save conditions: prisma_client={prisma_client is not None}, response={response is not None}, response_type={type(response)}" ) - + if prisma_client is not None and response is not None: await _save_vector_store_to_db_from_rag_ingest( response=response, diff --git a/litellm/proxy/realtime_endpoints/endpoints.py b/litellm/proxy/realtime_endpoints/endpoints.py index bb286d1fd0..d2975ba5fc 100644 --- a/litellm/proxy/realtime_endpoints/endpoints.py +++ b/litellm/proxy/realtime_endpoints/endpoints.py @@ -114,14 +114,14 @@ async def create_realtime_client_secret( ) data = {"model": model} - + # If session is provided, use it; otherwise create one from model if req.session: data["session"] = req.session.model_dump(exclude_none=True) elif req.model: # User provided model at root level, convert to session format data["session"] = {"type": "realtime", "model": model} - + if req.expires_after: data["expires_after"] = req.expires_after.model_dump(exclude_none=True) @@ -275,7 +275,7 @@ async def proxy_realtime_calls( status_code=http_status.HTTP_401_UNAUTHORIZED, media_type="application/json", ) - + openai_ephemeral_key = decoded_payload.get("ephemeral_key", "") model = ( decoded_payload.get("model_id") @@ -328,9 +328,7 @@ async def proxy_realtime_calls( call_type="arealtime_calls", ) - verbose_proxy_logger.debug( - "WebRTC: /v1/realtime/calls (model=%s)", model - ) + verbose_proxy_logger.debug("WebRTC: /v1/realtime/calls (model=%s)", model) llm_call = await route_request( data=data, diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 4253c2ca83..e9c7cce0d7 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -91,12 +91,12 @@ async def responses_api( ) data = await _read_request_body(request=request) - + # Check if polling via cache should be used for this request from litellm.proxy.response_polling.polling_handler import ( should_use_polling_for_request, ) - + should_use_polling = should_use_polling_for_request( background_mode=data.get("background", False), polling_via_cache_enabled=polling_via_cache_enabled, @@ -105,7 +105,7 @@ async def responses_api( llm_router=llm_router, native_background_mode=native_background_mode, ) - + # If polling is enabled, use polling mode if should_use_polling: from litellm.proxy.response_polling.background_streaming import ( @@ -114,26 +114,26 @@ async def responses_api( from litellm.proxy.response_polling.polling_handler import ( ResponsePollingHandler, ) - + verbose_proxy_logger.info( f"Starting background response with polling for model={data.get('model')}" ) - + # Initialize polling handler with configured TTL (from global config) polling_handler = ResponsePollingHandler( redis_cache=redis_usage_cache, - ttl=polling_cache_ttl # Global var set at startup + ttl=polling_cache_ttl, # Global var set at startup ) - + # Generate polling ID polling_id = ResponsePollingHandler.generate_polling_id() - + # Create initial state in Redis initial_state = await polling_handler.create_initial_state( polling_id=polling_id, request_data=data, ) - + # Start background task to stream and update cache asyncio.create_task( background_streaming_task( @@ -156,11 +156,11 @@ async def responses_api( version=version, ) ) - + # Return OpenAI Response object format (initial state) # https://platform.openai.com/docs/api-reference/responses/object return initial_state - + # Normal response flow processor = ProxyBaseLLMRequestProcessing(data=data) try: @@ -182,29 +182,32 @@ async def responses_api( user_api_base=user_api_base, version=version, ) - + # Store in managed objects table if background mode is enabled if data.get("background") and isinstance(response, ResponsesAPIResponse): if response.status in ["queued", "in_progress"]: from litellm_enterprise.proxy.hooks.managed_files import ( # type: ignore _PROXY_LiteLLMManagedFiles, - ) + ) + managed_files_obj = cast( Optional[_PROXY_LiteLLMManagedFiles], proxy_logging_obj.get_proxy_hook("managed_files"), ) - + if managed_files_obj and llm_router: try: # Get the actual deployment model_id from hidden params hidden_params = getattr(response, "_hidden_params", {}) or {} model_id = hidden_params.get("model_id", None) - + if not model_id: verbose_proxy_logger.warning( f"No model_id found in response hidden params for response {response.id}, skipping managed object storage" ) - raise Exception("No model_id found in response hidden params") + raise Exception( + "No model_id found in response hidden params" + ) # Store in managed objects table await managed_files_obj.store_unified_object_id( unified_object_id=response.id, @@ -214,7 +217,7 @@ async def responses_api( file_purpose="response", user_api_key_dict=user_api_key_dict, ) - + verbose_proxy_logger.info( f"Stored background response {response.id} in managed objects table with unified_id={response.id}" ) @@ -222,7 +225,7 @@ async def responses_api( verbose_proxy_logger.error( f"Failed to store background response in managed objects table: {str(e)}" ) - + return response except ModifyResponseException as e: # Guardrail passthrough: return violation message in Responses API format (200) @@ -241,9 +244,7 @@ async def responses_api( model=e.model or data.get("model"), output=cast(Any, [{"content": [{"type": "text", "text": violation_text}]}]), status="completed", - usage=ResponseAPIUsage( - input_tokens=0, output_tokens=0, total_tokens=0 - ), + usage=ResponseAPIUsage(input_tokens=0, output_tokens=0, total_tokens=0), ) return response_obj except Exception as e: @@ -305,26 +306,26 @@ async def cursor_chat_completions( from litellm.types.utils import ModelResponse data = await _read_request_body(request=request) - + # Convert 'messages' to 'input' for Responses API compatibility # Cursor sends 'messages' but Responses API expects 'input' if "messages" in data and "input" not in data: data["input"] = data.pop("messages") - + processor = ProxyBaseLLMRequestProcessing(data=data) def cursor_data_generator(response, user_api_key_dict, request_data): """ Custom generator that transforms Responses API streaming chunks to chat completion chunks. - + This generator is used for the cursor endpoint to convert Responses API format responses to chat completion format that Cursor IDE expects. - + Args: response: The streaming response (BaseResponsesAPIStreamingIterator or other) user_api_key_dict: User API key authentication dict request_data: Request data containing model, logging_obj, etc. - + Returns: Async generator that yields SSE-formatted chat completion chunks """ @@ -332,10 +333,12 @@ async def cursor_chat_completions( if isinstance(response, BaseResponsesAPIStreamingIterator): # Transform Responses API iterator to chat completion iterator # Cast to AsyncIterator[str] since BaseResponsesAPIStreamingIterator implements __aiter__/__anext__ - completion_stream = responses_api_bridge.transformation_handler.get_model_response_iterator( - streaming_response=cast(AsyncIterator[str], response), - sync_stream=False, - json_mode=False, + completion_stream = ( + responses_api_bridge.transformation_handler.get_model_response_iterator( + streaming_response=cast(AsyncIterator[str], response), + sync_stream=False, + json_mode=False, + ) ) # Wrap in CustomStreamWrapper to get the async generator logging_obj = request_data.get("litellm_logging_obj") @@ -381,18 +384,20 @@ async def cursor_chat_completions( # Transform non-streaming Responses API response to chat completions format if isinstance(response, ResponsesAPIResponse): logging_obj = processor.data.get("litellm_logging_obj") - transformed_response = responses_api_bridge.transformation_handler.transform_response( - model=processor.data.get("model", ""), - raw_response=response, - model_response=ModelResponse(), - logging_obj=cast(Any, logging_obj), - request_data=processor.data, - messages=processor.data.get("input", []), - optional_params={}, - litellm_params={}, - encoding=None, - api_key=None, - json_mode=None, + transformed_response = ( + responses_api_bridge.transformation_handler.transform_response( + model=processor.data.get("model", ""), + raw_response=response, + model_response=ModelResponse(), + logging_obj=cast(Any, logging_obj), + request_data=processor.data, + messages=processor.data.get("input", []), + optional_params={}, + litellm_params={}, + encoding=None, + api_key=None, + json_mode=None, + ) ) return transformed_response @@ -470,24 +475,24 @@ async def get_response( if not redis_usage_cache: raise HTTPException( status_code=500, - detail="Redis cache not configured. Polling requires Redis." + detail="Redis cache not configured. Polling requires Redis.", ) - + polling_handler = ResponsePollingHandler(redis_cache=redis_usage_cache) - + # Get current state from cache state = await polling_handler.get_state(response_id) - + if not state: raise HTTPException( status_code=404, - detail=f"Polling response {response_id} not found or expired" + detail=f"Polling response {response_id} not found or expired", ) - + # Return the whole state directly (OpenAI Response object format) # https://platform.openai.com/docs/api-reference/responses/object return state - + # Normal provider response flow data = await _read_request_body(request=request) data["response_id"] = response_id @@ -576,37 +581,28 @@ async def delete_response( if ResponsePollingHandler.is_polling_id(response_id): # Handle polling response deletion if not redis_usage_cache: - raise HTTPException( - status_code=500, - detail="Redis cache not configured." - ) - + raise HTTPException(status_code=500, detail="Redis cache not configured.") + polling_handler = ResponsePollingHandler(redis_cache=redis_usage_cache) - + # Get state to verify access state = await polling_handler.get_state(response_id) - + if not state: raise HTTPException( - status_code=404, - detail=f"Polling response {response_id} not found" + status_code=404, detail=f"Polling response {response_id} not found" ) - + # Delete from cache success = await polling_handler.delete_polling(response_id) - + if success: - return DeleteResponseResult( - id=response_id, - object="response", - deleted=True - ) + return DeleteResponseResult(id=response_id, object="response", deleted=True) else: raise HTTPException( - status_code=500, - detail="Failed to delete polling response" + status_code=500, detail="Failed to delete polling response" ) - + # Normal provider response flow data = await _read_request_body(request=request) data["response_id"] = response_id @@ -850,37 +846,32 @@ async def cancel_response( if ResponsePollingHandler.is_polling_id(response_id): # Handle polling response cancellation if not redis_usage_cache: - raise HTTPException( - status_code=500, - detail="Redis cache not configured." - ) - + raise HTTPException(status_code=500, detail="Redis cache not configured.") + polling_handler = ResponsePollingHandler(redis_cache=redis_usage_cache) - + # Get current state to verify it exists state = await polling_handler.get_state(response_id) - + if not state: raise HTTPException( - status_code=404, - detail=f"Polling response {response_id} not found" + status_code=404, detail=f"Polling response {response_id} not found" ) - + # Cancel the polling response (sets status to "cancelled") success = await polling_handler.cancel_polling(response_id) - + if success: # Fetch the updated state with cancelled status updated_state = await polling_handler.get_state(response_id) - + # Return the whole state directly (now with status="cancelled") return updated_state else: raise HTTPException( - status_code=500, - detail="Failed to cancel polling response" + status_code=500, detail="Failed to cancel polling response" ) - + # Normal provider response flow data = await _read_request_body(request=request) data["response_id"] = response_id diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index 1e37b42f0c..f1b2493976 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -40,30 +40,30 @@ async def background_streaming_task( # noqa: PLR0915 ): """ Background task to stream response and update cache - + Follows OpenAI Response Streaming format: https://platform.openai.com/docs/api-reference/responses-streaming - + Processes streaming events and builds Response object: https://platform.openai.com/docs/api-reference/responses/object """ - + try: verbose_proxy_logger.info(f"Starting background streaming for {polling_id}") - + # Update status to in_progress (OpenAI format) await polling_handler.update_state( polling_id=polling_id, status="in_progress", ) - + # Force streaming mode and remove background flag data["stream"] = True data.pop("background", None) - + # Create processor processor = ProxyBaseLLMRequestProcessing(data=data) - + # Make streaming request response = await processor.base_process_llm_request( request=request, @@ -83,12 +83,14 @@ async def background_streaming_task( # noqa: PLR0915 user_api_base=user_api_base, version=version, ) - + # Process streaming response following OpenAI events format # https://platform.openai.com/docs/api-reference/responses-streaming output_items: dict[str, dict[str, Any]] = {} # Track output items by ID - accumulated_text = {} # Track accumulated text deltas by (item_id, content_index) - + accumulated_text = ( + {} + ) # Track accumulated text deltas by (item_id, content_index) + # ResponsesAPIResponse fields to extract from response.completed usage_data = None reasoning_data = None @@ -106,17 +108,19 @@ async def background_streaming_task( # noqa: PLR0915 user_data = None store_data = None incomplete_details_data = None - + state_dirty = False # Track if state needs to be synced last_update_time = asyncio.get_event_loop().time() UPDATE_INTERVAL = 0.150 # 150ms batching interval - + async def flush_state_if_needed(force: bool = False) -> None: """Flush accumulated state to Redis if interval elapsed or forced""" nonlocal state_dirty, last_update_time - + current_time = asyncio.get_event_loop().time() - if state_dirty and (force or (current_time - last_update_time) >= UPDATE_INTERVAL): + if state_dirty and ( + force or (current_time - last_update_time) >= UPDATE_INTERVAL + ): # Convert output_items dict to list for update output_list = list(output_items.values()) await polling_handler.update_state( @@ -125,23 +129,23 @@ async def background_streaming_task( # noqa: PLR0915 ) state_dirty = False last_update_time = current_time - + # Handle StreamingResponse - if hasattr(response, 'body_iterator'): + if hasattr(response, "body_iterator"): async for chunk in response.body_iterator: # Parse chunk if isinstance(chunk, bytes): - chunk = chunk.decode('utf-8') - + chunk = chunk.decode("utf-8") + if isinstance(chunk, str) and chunk.startswith("data: "): chunk_data = chunk[6:].strip() if chunk_data == "[DONE]": break - + try: event = json.loads(chunk_data) event_type = event.get("type", "") - + # Process different event types based on OpenAI streaming spec if event_type == "response.output_item.added": # New output item added @@ -150,48 +154,52 @@ async def background_streaming_task( # noqa: PLR0915 if item_id: output_items[item_id] = item state_dirty = True - + elif event_type == "response.content_part.added": # Content part added to an output item item_id = event.get("item_id") content_part = event.get("part", {}) - + if item_id and item_id in output_items: # Update the output item with new content if "content" not in output_items[item_id]: output_items[item_id]["content"] = [] output_items[item_id]["content"].append(content_part) state_dirty = True - + elif event_type == "response.output_text.delta": # Text delta - accumulate text content # https://platform.openai.com/docs/api-reference/responses-streaming/response-text-delta item_id = event.get("item_id") content_index = event.get("content_index", 0) delta = event.get("delta", "") - + if item_id and item_id in output_items: # Accumulate text delta key = (item_id, content_index) if key not in accumulated_text: accumulated_text[key] = "" accumulated_text[key] += delta - + # Update the content in output_items if "content" in output_items[item_id]: content_list = output_items[item_id]["content"] if content_index < len(content_list): # Update existing content part with accumulated text - if isinstance(content_list[content_index], dict): - content_list[content_index]["text"] = accumulated_text[key] + if isinstance( + content_list[content_index], dict + ): + content_list[content_index][ + "text" + ] = accumulated_text[key] state_dirty = True - + elif event_type == "response.content_part.done": # Content part completed item_id = event.get("item_id") content_part = event.get("part", {}) content_index = event.get("content_index", 0) - + if item_id and item_id in output_items: # Update with final content from event if "content" in output_items[item_id]: @@ -199,7 +207,7 @@ async def background_streaming_task( # noqa: PLR0915 if content_index < len(content_list): content_list[content_index] = content_part state_dirty = True - + elif event_type == "response.output_item.done": # Output item completed - use final item data item = event.get("item", {}) @@ -207,7 +215,7 @@ async def background_streaming_task( # noqa: PLR0915 if item_id: output_items[item_id] = item state_dirty = True - + elif event_type == "response.in_progress": # Response is now in progress # https://platform.openai.com/docs/api-reference/responses-streaming/response-in-progress @@ -215,32 +223,40 @@ async def background_streaming_task( # noqa: PLR0915 polling_id=polling_id, status="in_progress", ) - + elif event_type == "response.completed": # Response completed - extract all ResponsesAPIResponse fields # https://platform.openai.com/docs/api-reference/responses-streaming/response-completed response_data = event.get("response", {}) - + # Core response fields usage_data = response_data.get("usage") reasoning_data = response_data.get("reasoning") tool_choice_data = response_data.get("tool_choice") tools_data = response_data.get("tools") - + # Additional ResponsesAPIResponse fields model_data = response_data.get("model") instructions_data = response_data.get("instructions") temperature_data = response_data.get("temperature") top_p_data = response_data.get("top_p") - max_output_tokens_data = response_data.get("max_output_tokens") - previous_response_id_data = response_data.get("previous_response_id") + max_output_tokens_data = response_data.get( + "max_output_tokens" + ) + previous_response_id_data = response_data.get( + "previous_response_id" + ) text_data = response_data.get("text") truncation_data = response_data.get("truncation") - parallel_tool_calls_data = response_data.get("parallel_tool_calls") + parallel_tool_calls_data = response_data.get( + "parallel_tool_calls" + ) user_data = response_data.get("user") store_data = response_data.get("store") - incomplete_details_data = response_data.get("incomplete_details") - + incomplete_details_data = response_data.get( + "incomplete_details" + ) + # Also update output from final response if available if "output" in response_data: final_output = response_data.get("output", []) @@ -249,19 +265,19 @@ async def background_streaming_task( # noqa: PLR0915 if item_id: output_items[item_id] = item state_dirty = True - + # Flush state to Redis if interval elapsed await flush_state_if_needed() - + except json.JSONDecodeError as e: verbose_proxy_logger.warning( f"Failed to parse streaming chunk: {e}" ) pass - + # Final flush to ensure all accumulated state is saved await flush_state_if_needed(force=True) - + # Mark as completed with all ResponsesAPIResponse fields await polling_handler.update_state( polling_id=polling_id, @@ -283,25 +299,25 @@ async def background_streaming_task( # noqa: PLR0915 store=store_data, incomplete_details=incomplete_details_data, ) - + verbose_proxy_logger.info( f"Completed background streaming for {polling_id}, output_items={len(output_items)}" ) - + except Exception as e: verbose_proxy_logger.error( f"Error in background streaming task for {polling_id}: {str(e)}" ) import traceback + verbose_proxy_logger.error(traceback.format_exc()) - + await polling_handler.update_state( polling_id=polling_id, status="failed", error={ "type": "internal_error", "message": str(e), - "code": "background_streaming_error" + "code": "background_streaming_error", }, ) - diff --git a/litellm/proxy/response_polling/polling_handler.py b/litellm/proxy/response_polling/polling_handler.py index f0b850049b..71e97a46c6 100644 --- a/litellm/proxy/response_polling/polling_handler.py +++ b/litellm/proxy/response_polling/polling_handler.py @@ -13,29 +13,29 @@ from litellm.types.llms.openai import ResponsesAPIResponse, ResponsesAPIStatus class ResponsePollingHandler: """Handles polling-based responses with Redis cache""" - + CACHE_KEY_PREFIX = "litellm:polling:response:" POLLING_ID_PREFIX = "litellm_poll_" # Clear prefix to identify polling IDs - + def __init__(self, redis_cache: Optional[RedisCache] = None, ttl: int = 3600): self.redis_cache = redis_cache self.ttl = ttl # Time-to-live for cache entries (default: 1 hour) - + @classmethod def generate_polling_id(cls) -> str: """Generate a unique UUID for polling with clear prefix""" return f"{cls.POLLING_ID_PREFIX}{uuid4()}" - + @classmethod def is_polling_id(cls, response_id: str) -> bool: """Check if a response_id is a polling ID""" return response_id.startswith(cls.POLLING_ID_PREFIX) - + @classmethod def get_cache_key(cls, polling_id: str) -> str: """Get Redis cache key for a polling ID""" return f"{cls.CACHE_KEY_PREFIX}{polling_id}" - + async def create_initial_state( self, polling_id: str, @@ -43,19 +43,19 @@ class ResponsePollingHandler: ) -> ResponsesAPIResponse: """ Create initial state in Redis for a polling request - + Uses OpenAI ResponsesAPIResponse object: https://platform.openai.com/docs/api-reference/responses/object - + Args: polling_id: Unique identifier for this polling request request_data: Original request data - + Returns: ResponsesAPIResponse object following OpenAI spec """ created_timestamp = int(datetime.now(timezone.utc).timestamp()) - + # Create OpenAI-compliant response object response = ResponsesAPIResponse( id=polling_id, @@ -66,9 +66,9 @@ class ResponsePollingHandler: metadata=request_data.get("metadata", {}), usage=None, ) - + cache_key = self.get_cache_key(polling_id) - + if self.redis_cache: # Store ResponsesAPIResponse directly in Redis await self.redis_cache.async_set_cache( @@ -79,9 +79,9 @@ class ResponsePollingHandler: verbose_proxy_logger.debug( f"Created initial polling state for {polling_id} with TTL={self.ttl}s" ) - + return response - + async def update_state( self, polling_id: str, @@ -108,10 +108,10 @@ class ResponsePollingHandler: ) -> None: """ Update the polling state in Redis - + Uses OpenAI Response object format with native status types: https://platform.openai.com/docs/api-reference/responses/object - + Args: polling_id: Unique identifier for this polling request status: OpenAI ResponsesAPIStatus value @@ -136,9 +136,9 @@ class ResponsePollingHandler: """ if not self.redis_cache: return - + cache_key = self.get_cache_key(polling_id) - + # Get current state cached_state = await self.redis_cache.async_get_cache(cache_key) if not cached_state: @@ -146,31 +146,31 @@ class ResponsePollingHandler: f"No cached state found for polling_id: {polling_id}" ) return - + # Parse existing ResponsesAPIResponse from cache state = json.loads(cached_state) - + # Update status (using OpenAI native status values) if status: state["status"] = status - + # Replace full output list if provided if output is not None: state["output"] = output - + # Update usage if usage: state["usage"] = usage - + # Handle error (sets status to OpenAI's "failed") if error: state["status"] = "failed" state["error"] = error # Use OpenAI's 'error' field - + # Handle incomplete details if incomplete_details: state["incomplete_details"] = incomplete_details - + # Update reasoning, tool_choice, tools from response.completed if reasoning is not None: state["reasoning"] = reasoning @@ -178,7 +178,7 @@ class ResponsePollingHandler: state["tool_choice"] = tool_choice if tools is not None: state["tools"] = tools - + # Update additional ResponsesAPIResponse fields if model is not None: state["model"] = model @@ -202,36 +202,36 @@ class ResponsePollingHandler: state["user"] = user if store is not None: state["store"] = store - + # Update cache with configured TTL await self.redis_cache.async_set_cache( key=cache_key, value=json.dumps(state), ttl=self.ttl, ) - + output_count = len(state.get("output", [])) verbose_proxy_logger.debug( f"Updated polling state for {polling_id}: status={state['status']}, output_items={output_count}" ) - + async def get_state(self, polling_id: str) -> Optional[Dict[str, Any]]: """Get current polling state from Redis""" if not self.redis_cache: return None - + cache_key = self.get_cache_key(polling_id) cached_state = await self.redis_cache.async_get_cache(cache_key) - + if cached_state: return json.loads(cached_state) - + return None - + async def cancel_polling(self, polling_id: str) -> bool: """ Cancel a polling request - + Following OpenAI Response object format for cancelled status """ await self.update_state( @@ -239,12 +239,12 @@ class ResponsePollingHandler: status="cancelled", ) return True - + async def delete_polling(self, polling_id: str) -> bool: """Delete a polling request from cache""" if not self.redis_cache: return False - + cache_key = self.get_cache_key(polling_id) # Use RedisCache's async_delete_cache method which handles Redis/RedisCluster await self.redis_cache.async_delete_cache(cache_key) @@ -257,38 +257,40 @@ def should_use_polling_for_request( redis_cache, # RedisCache or None model: str, llm_router, # Router instance or None - native_background_mode: Optional[List[str]] = None, # List of models that should use native background mode + native_background_mode: Optional[ + List[str] + ] = None, # List of models that should use native background mode ) -> bool: """ Determine if polling via cache should be used for a request. - + Args: background_mode: Whether background=true was set in the request polling_via_cache_enabled: Config value - False, "all", or list of providers redis_cache: Redis cache instance (required for polling) model: Model name from the request (e.g., "gpt-5" or "openai/gpt-4o") llm_router: LiteLLM router instance for looking up model deployments - native_background_mode: List of model names that should use native provider + native_background_mode: List of model names that should use native provider background mode instead of polling via cache - + Returns: True if polling should be used, False otherwise """ # All conditions must be met if not (background_mode and polling_via_cache_enabled and redis_cache): return False - + # Check if model is in native_background_mode list - these use native provider background mode if native_background_mode and model in native_background_mode: verbose_proxy_logger.debug( f"Model {model} is in native_background_mode list, skipping polling via cache" ) return False - + # "all" enables polling for all providers if polling_via_cache_enabled == "all": return True - + # Check if provider is in the enabled list if isinstance(polling_via_cache_enabled, list): # First, try to get provider from model string format "provider/model" @@ -304,16 +306,16 @@ def should_use_polling_for_request( for idx in indices: deployment_dict = llm_router.model_list[idx] litellm_params = deployment_dict.get("litellm_params", {}) - + # Check custom_llm_provider first dep_provider = litellm_params.get("custom_llm_provider") - + # Then try to extract from model (e.g., "openai/gpt-5") if not dep_provider: dep_model = litellm_params.get("model", "") if "/" in dep_model: dep_provider = dep_model.split("/")[0] - + # If ANY deployment's provider matches, enable polling if dep_provider and dep_provider in polling_via_cache_enabled: verbose_proxy_logger.debug( @@ -324,6 +326,5 @@ def should_use_polling_for_request( verbose_proxy_logger.debug( f"Could not resolve provider for model {model}: {e}" ) - - return False + return False diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 285de6d101..9fb5fe9fee 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -308,12 +308,18 @@ async def route_request( # noqa: PLR0915 - Complex routing function, refactorin if model and llm_router: try: # Try to get deployment credentials for this model - deployment_creds = llm_router.get_deployment_credentials(model_id=model) + deployment_creds = llm_router.get_deployment_credentials( + model_id=model + ) if not deployment_creds: # Try by model group name - deployment = llm_router.get_deployment_by_model_group_name(model_group_name=model) + deployment = llm_router.get_deployment_by_model_group_name( + model_group_name=model + ) if deployment and deployment.litellm_params: - deployment_creds = deployment.litellm_params.model_dump(exclude_none=True) + deployment_creds = deployment.litellm_params.model_dump( + exclude_none=True + ) # If we found credentials, merge them into data (but don't override user-provided values) if deployment_creds: @@ -439,7 +445,7 @@ async def route_request( # noqa: PLR0915 - Complex routing function, refactorin from litellm.proxy.agent_endpoints.a2a_routing import ( route_a2a_agent_request, ) - + result = route_a2a_agent_request(data, route_type) if result is not None: return result diff --git a/litellm/proxy/search_endpoints/__init__.py b/litellm/proxy/search_endpoints/__init__.py index 92b88f783c..085d9446a4 100644 --- a/litellm/proxy/search_endpoints/__init__.py +++ b/litellm/proxy/search_endpoints/__init__.py @@ -5,4 +5,3 @@ from .search_tool_registry import SearchToolRegistry __all__ = [ "SearchToolRegistry", ] - diff --git a/litellm/proxy/search_endpoints/endpoints.py b/litellm/proxy/search_endpoints/endpoints.py index c7a3b88c49..8bed5b5407 100644 --- a/litellm/proxy/search_endpoints/endpoints.py +++ b/litellm/proxy/search_endpoints/endpoints.py @@ -127,35 +127,38 @@ async def search( # Read request body body = await request.body() data = orjson.loads(body) - + # If search_tool_name is provided in URL path, use it (takes precedence over body) if search_tool_name is not None: data["search_tool_name"] = search_tool_name if "search_tool_name" in data and data["search_tool_name"]: data["model"] = data["search_tool_name"] - + if llm_router is not None and hasattr(llm_router, "search_tools"): search_tool_name_value = data["search_tool_name"] - + verbose_proxy_logger.debug( f"Search endpoint - Looking for search_tool_name: {search_tool_name_value}. " f"Available search tools in router: {[tool.get('search_tool_name') for tool in llm_router.search_tools]}. " f"Total search tools: {len(llm_router.search_tools)}" ) - + matching_tools = [ - tool for tool in llm_router.search_tools + tool + for tool in llm_router.search_tools if tool.get("search_tool_name") == search_tool_name_value ] - + if matching_tools: search_tool = matching_tools[0] - search_provider = search_tool.get("litellm_params", {}).get("search_provider") - + search_provider = search_tool.get("litellm_params", {}).get( + "search_provider" + ) + if search_provider: data["custom_llm_provider"] = search_provider - + if "metadata" not in data: data["metadata"] = {} data["metadata"]["model_group"] = search_tool_name_value @@ -189,6 +192,7 @@ async def search( version=version, ) + @router.get( "/v1/search/tools", dependencies=[Depends(user_api_key_auth)], @@ -236,28 +240,27 @@ async def list_search_tools( try: search_tools_list = [] - + if llm_router is not None and hasattr(llm_router, "search_tools"): for tool in llm_router.search_tools: tool_info = { "search_tool_name": tool.get("search_tool_name"), - "search_provider": tool.get("litellm_params", {}).get("search_provider"), + "search_provider": tool.get("litellm_params", {}).get( + "search_provider" + ), } - + # Add description if available if "search_tool_info" in tool and tool["search_tool_info"]: description = tool["search_tool_info"].get("description") if description: tool_info["description"] = description - + search_tools_list.append(tool_info) - - return { - "object": "list", - "data": search_tools_list - } + + return {"object": "list", "data": search_tools_list} except Exception as e: from litellm._logging import verbose_proxy_logger + verbose_proxy_logger.exception(f"Error listing search tools: {e}") raise HTTPException(status_code=500, detail=str(e)) - diff --git a/litellm/proxy/search_endpoints/search_tool_management.py b/litellm/proxy/search_endpoints/search_tool_management.py index 4754316795..c46bbfddca 100644 --- a/litellm/proxy/search_endpoints/search_tool_management.py +++ b/litellm/proxy/search_endpoints/search_tool_management.py @@ -26,10 +26,10 @@ SEARCH_TOOL_REGISTRY = SearchToolRegistry() def _convert_datetime_to_str(value: Union[datetime, str, None]) -> Union[str, None]: """ Convert datetime object to ISO format string. - + Args: value: datetime object, string, or None - + Returns: ISO format string or original value if already string or None """ @@ -97,14 +97,12 @@ async def list_search_tools(): prisma_client=prisma_client ) - db_tool_names = { - tool.get("search_tool_name") for tool in search_tools_from_db - } + db_tool_names = {tool.get("search_tool_name") for tool in search_tools_from_db} search_tool_configs: List[SearchToolInfoResponse] = [] - + config_search_tools = [] - + try: config = await proxy_config.get_config() parsed_tools = proxy_config.parse_search_tools(config) @@ -114,7 +112,7 @@ async def list_search_tools(): verbose_proxy_logger.debug( f"Could not get config-defined search tools: {e}" ) - + for search_tool in config_search_tools: tool_name = search_tool.get("search_tool_name") if tool_name: @@ -138,10 +136,11 @@ async def list_search_tools(): ) search_tool_configs = [ - tool for tool in search_tool_configs + tool + for tool in search_tool_configs if tool.get("search_tool_name") not in db_tool_names ] - + for search_tool in search_tools_from_db: litellm_params_dict = dict(search_tool.get("litellm_params", {})) masked_litellm_params_dict = _get_masked_values( @@ -149,7 +148,7 @@ async def list_search_tools(): unmasked_length=4, number_of_asterisks=4, ) - + search_tool_configs.append( SearchToolInfoResponse( search_tool_id=search_tool.get("search_tool_id"), @@ -508,17 +507,16 @@ async def test_search_tool_connection(request: TestSearchToolConnectionRequest): search_provider = litellm_params.get("search_provider") api_key = litellm_params.get("api_key") api_base = litellm_params.get("api_base") - + if not search_provider: raise HTTPException( - status_code=400, - detail="search_provider is required in litellm_params" + status_code=400, detail="search_provider is required in litellm_params" ) - + verbose_proxy_logger.debug( f"Testing connection to search provider: {search_provider}" ) - + # Make a simple test search query with max_results=1 to minimize cost test_query = "test" response = await asearch( @@ -529,26 +527,28 @@ async def test_search_tool_connection(request: TestSearchToolConnectionRequest): max_results=1, # Minimize results to reduce cost timeout=10.0, # 10 second timeout for test ) - + verbose_proxy_logger.debug( f"Successfully tested connection to {search_provider} search provider" ) - + return { "status": "success", "message": f"Successfully connected to {search_provider} search provider", "test_query": test_query, - "results_count": len(response.results) if response and response.results else 0, + "results_count": len(response.results) + if response and response.results + else 0, } - + except Exception as e: error_message = str(e) error_type = type(e).__name__ - + verbose_proxy_logger.exception( f"Failed to connect to search provider: {error_message}" ) - + # Return error details in a structured format return { "status": "error", @@ -592,31 +592,34 @@ async def get_available_search_providers(): """ try: from litellm.utils import ProviderConfigManager - + available_providers = [] - + # Auto-discover providers from SearchProviders enum for provider in SearchProviders: try: # Get the config class for this provider - config = ProviderConfigManager.get_provider_search_config(provider=provider) - + config = ProviderConfigManager.get_provider_search_config( + provider=provider + ) + if config is not None: # Get the UI-friendly name from the config class ui_name = config.ui_friendly_name() - - available_providers.append({ - "provider_name": provider.value, - "ui_friendly_name": ui_name, - }) + + available_providers.append( + { + "provider_name": provider.value, + "ui_friendly_name": ui_name, + } + ) except Exception as e: verbose_proxy_logger.debug( f"Could not get config for search provider {provider.value}: {e}" ) continue - + return {"providers": available_providers} except Exception as e: verbose_proxy_logger.exception(f"Error getting available search providers: {e}") raise HTTPException(status_code=500, detail=str(e)) - diff --git a/litellm/proxy/search_endpoints/search_tool_registry.py b/litellm/proxy/search_endpoints/search_tool_registry.py index bab92d21de..e9eba1e179 100644 --- a/litellm/proxy/search_endpoints/search_tool_registry.py +++ b/litellm/proxy/search_endpoints/search_tool_registry.py @@ -11,7 +11,7 @@ from litellm.types.search import SearchTool class SearchToolRegistry: - """ + """ Handles adding, removing, and getting search tools in DB + in memory. """ @@ -22,10 +22,10 @@ class SearchToolRegistry: def _convert_prisma_to_dict(prisma_obj) -> dict: """ Convert Prisma result to dict with datetime objects as ISO format strings. - + Args: prisma_obj: Prisma model instance - + Returns: Dict with datetime fields converted to ISO strings """ @@ -40,34 +40,38 @@ class SearchToolRegistry: ########################################################### ########### DB management helpers for search tools ######## ########################################################### - + async def add_search_tool_to_db( self, search_tool: SearchTool, prisma_client: PrismaClient ): """ Add a search tool to the database. - + Args: search_tool: Search tool configuration prisma_client: Prisma client instance - + Returns: Dict with created search tool data """ try: search_tool_name = search_tool.get("search_tool_name") - litellm_params: str = safe_dumps(dict(search_tool.get("litellm_params", {}))) + litellm_params: str = safe_dumps( + dict(search_tool.get("litellm_params", {})) + ) search_tool_info: str = safe_dumps(search_tool.get("search_tool_info", {})) # Create search tool in DB - created_search_tool = await prisma_client.db.litellm_searchtoolstable.create( - data={ - "search_tool_name": search_tool_name, - "litellm_params": litellm_params, - "search_tool_info": search_tool_info, - "created_at": datetime.now(timezone.utc), - "updated_at": datetime.now(timezone.utc), - } + created_search_tool = ( + await prisma_client.db.litellm_searchtoolstable.create( + data={ + "search_tool_name": search_tool_name, + "litellm_params": litellm_params, + "search_tool_info": search_tool_info, + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + } + ) ) # Add search_tool_id to the returned search tool object @@ -86,11 +90,11 @@ class SearchToolRegistry: ): """ Delete a search tool from the database. - + Args: search_tool_id: ID of search tool to delete prisma_client: Prisma client instance - + Returns: Dict with success message """ @@ -99,10 +103,10 @@ class SearchToolRegistry: existing_tool = await prisma_client.db.litellm_searchtoolstable.find_unique( where={"search_tool_id": search_tool_id} ) - + if not existing_tool: raise Exception(f"Search tool with ID {search_tool_id} not found") - + # Delete from DB await prisma_client.db.litellm_searchtoolstable.delete( where={"search_tool_id": search_tool_id} @@ -113,7 +117,9 @@ class SearchToolRegistry: "search_tool_name": existing_tool.search_tool_name, } except Exception as e: - verbose_proxy_logger.exception(f"Error deleting search tool from DB: {str(e)}") + verbose_proxy_logger.exception( + f"Error deleting search tool from DB: {str(e)}" + ) raise Exception(f"Error deleting search tool from DB: {str(e)}") async def update_search_tool_in_db( @@ -121,35 +127,41 @@ class SearchToolRegistry: ): """ Update a search tool in the database. - + Args: search_tool_id: ID of search tool to update search_tool: Updated search tool configuration prisma_client: Prisma client instance - + Returns: Dict with updated search tool data """ try: search_tool_name = search_tool.get("search_tool_name") - litellm_params: str = safe_dumps(dict(search_tool.get("litellm_params", {}))) + litellm_params: str = safe_dumps( + dict(search_tool.get("litellm_params", {})) + ) search_tool_info: str = safe_dumps(search_tool.get("search_tool_info", {})) # Update in DB - updated_search_tool = await prisma_client.db.litellm_searchtoolstable.update( - where={"search_tool_id": search_tool_id}, - data={ - "search_tool_name": search_tool_name, - "litellm_params": litellm_params, - "search_tool_info": search_tool_info, - "updated_at": datetime.now(timezone.utc), - }, + updated_search_tool = ( + await prisma_client.db.litellm_searchtoolstable.update( + where={"search_tool_id": search_tool_id}, + data={ + "search_tool_name": search_tool_name, + "litellm_params": litellm_params, + "search_tool_info": search_tool_info, + "updated_at": datetime.now(timezone.utc), + }, + ) ) # Convert to dict with ISO formatted datetimes return self._convert_prisma_to_dict(updated_search_tool) except Exception as e: - verbose_proxy_logger.exception(f"Error updating search tool in DB: {str(e)}") + verbose_proxy_logger.exception( + f"Error updating search tool in DB: {str(e)}" + ) raise Exception(f"Error updating search tool in DB: {str(e)}") @staticmethod @@ -158,10 +170,10 @@ class SearchToolRegistry: ) -> List[SearchTool]: """ Get all search tools from the database. - + Args: prisma_client: Prisma client instance - + Returns: List of search tool configurations """ @@ -175,12 +187,16 @@ class SearchToolRegistry: search_tools: List[SearchTool] = [] for search_tool in search_tools_from_db: # Convert Prisma result to dict with ISO formatted datetimes - search_tool_dict = SearchToolRegistry._convert_prisma_to_dict(search_tool) + search_tool_dict = SearchToolRegistry._convert_prisma_to_dict( + search_tool + ) search_tools.append(SearchTool(**search_tool_dict)) # type: ignore return search_tools except Exception as e: - verbose_proxy_logger.exception(f"Error getting search tools from DB: {str(e)}") + verbose_proxy_logger.exception( + f"Error getting search tools from DB: {str(e)}" + ) raise Exception(f"Error getting search tools from DB: {str(e)}") async def get_search_tool_by_id_from_db( @@ -188,11 +204,11 @@ class SearchToolRegistry: ) -> Optional[SearchTool]: """ Get a search tool by its ID from the database. - + Args: search_tool_id: ID of search tool to retrieve prisma_client: Prisma client instance - + Returns: Search tool configuration or None if not found """ @@ -208,7 +224,9 @@ class SearchToolRegistry: search_tool_dict = self._convert_prisma_to_dict(search_tool) return SearchTool(**search_tool_dict) # type: ignore except Exception as e: - verbose_proxy_logger.exception(f"Error getting search tool from DB: {str(e)}") + verbose_proxy_logger.exception( + f"Error getting search tool from DB: {str(e)}" + ) raise Exception(f"Error getting search tool from DB: {str(e)}") async def get_search_tool_by_name_from_db( @@ -216,11 +234,11 @@ class SearchToolRegistry: ) -> Optional[SearchTool]: """ Get a search tool by its name from the database. - + Args: search_tool_name: Name of search tool to retrieve prisma_client: Prisma client instance - + Returns: Search tool configuration or None if not found """ @@ -236,6 +254,7 @@ class SearchToolRegistry: search_tool_dict = self._convert_prisma_to_dict(search_tool) return SearchTool(**search_tool_dict) # type: ignore except Exception as e: - verbose_proxy_logger.exception(f"Error getting search tool from DB: {str(e)}") + verbose_proxy_logger.exception( + f"Error getting search tool from DB: {str(e)}" + ) raise Exception(f"Error getting search tool from DB: {str(e)}") - diff --git a/litellm/proxy/spend_tracking/cloudzero_endpoints.py b/litellm/proxy/spend_tracking/cloudzero_endpoints.py index 172169f2c7..c7bff7ec64 100644 --- a/litellm/proxy/spend_tracking/cloudzero_endpoints.py +++ b/litellm/proxy/spend_tracking/cloudzero_endpoints.py @@ -303,6 +303,7 @@ def is_cloudzero_setup_in_config() -> bool: bool: True if CloudZero is configured, False otherwise """ import litellm + return "cloudzero" in litellm.callbacks @@ -312,7 +313,7 @@ async def is_cloudzero_setup() -> bool: CloudZero is considered setup if: - CloudZero is configured in config.yaml callbacks, OR - - CloudZero environment variables are set, OR + - CloudZero environment variables are set, OR - CloudZero settings exist in the database Returns: @@ -322,11 +323,11 @@ async def is_cloudzero_setup() -> bool: # Check config.yaml/environment variables first if is_cloudzero_setup_in_config(): return True - + # Check database as fallback if await is_cloudzero_setup_in_db(): return True - + return False except Exception as e: @@ -425,9 +426,7 @@ async def cloudzero_dry_run_export( # Initialize logger with credentials directly logger = CloudZeroLogger() - dry_run_result = await logger.dry_run_export_usage_data( - limit=request.limit - ) + dry_run_result = await logger.dry_run_export_usage_data(limit=request.limit) verbose_proxy_logger.info("CloudZero dry run export completed successfully") @@ -470,7 +469,6 @@ async def cloudzero_export( Only admin users can perform CloudZero exports. """ - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise HTTPException( status_code=403, @@ -500,10 +498,10 @@ async def cloudzero_export( verbose_proxy_logger.info("CloudZero export completed successfully") return CloudZeroExportResponse( - message="CloudZero export completed successfully", + message="CloudZero export completed successfully", status="success", dry_run_data=None, - summary=None + summary=None, ) except Exception as e: diff --git a/litellm/proxy/spend_tracking/cold_storage_handler.py b/litellm/proxy/spend_tracking/cold_storage_handler.py index 262d14fad7..adbbc14123 100644 --- a/litellm/proxy/spend_tracking/cold_storage_handler.py +++ b/litellm/proxy/spend_tracking/cold_storage_handler.py @@ -16,46 +16,52 @@ class ColdStorageHandler: It allows fetching a dict of the proxy server request from s3 or GCS bucket. """ - + async def get_proxy_server_request_from_cold_storage_with_object_key( self, object_key: str, ) -> Optional[dict]: """ Get the proxy server request from cold storage using the object key directly. - + Args: object_key: The S3/GCS object key to retrieve - + Returns: Optional[dict]: The proxy server request dict or None if not found """ - + # select the custom logger to use for cold storage - custom_logger_name: Optional[_custom_logger_compatible_callbacks_literal] = self._select_custom_logger_for_cold_storage() + custom_logger_name: Optional[ + _custom_logger_compatible_callbacks_literal + ] = self._select_custom_logger_for_cold_storage() # if no custom logger name is configured, return None if custom_logger_name is None: return None # get the active/initialized custom logger - custom_logger: Optional[CustomLogger] = litellm.logging_callback_manager.get_active_custom_logger_for_callback_name(custom_logger_name) + custom_logger: Optional[ + CustomLogger + ] = litellm.logging_callback_manager.get_active_custom_logger_for_callback_name( + custom_logger_name + ) # if no custom logger is found, return None if custom_logger is None: - return None - + return None + proxy_server_request = await custom_logger.get_proxy_server_request_from_cold_storage_with_object_key( object_key=object_key, ) return proxy_server_request - - def _select_custom_logger_for_cold_storage( self, ) -> Optional[_custom_logger_compatible_callbacks_literal]: - cold_storage_custom_logger: Optional[_custom_logger_compatible_callbacks_literal] = litellm.cold_storage_custom_logger + cold_storage_custom_logger: Optional[ + _custom_logger_compatible_callbacks_literal + ] = litellm.cold_storage_custom_logger return cold_storage_custom_logger diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 5b58fbe70a..b3b4b55af1 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -1461,11 +1461,21 @@ async def _get_spend_report_for_time_range( dependencies=[Depends(user_api_key_auth)], responses={ 200: { - "cost": { - "description": "The calculated cost", - "example": 0.0, - "type": "float", - } + "description": "The calculated cost", + "content": { + "application/json": { + "schema": { + "type": "object", + "properties": { + "cost": { + "type": "number", + "description": "The calculated cost", + "example": 0.0, + } + }, + } + } + }, } }, ) @@ -1672,7 +1682,8 @@ async def ui_view_spend_logs( # noqa: PLR0915 default=None, description="Filter logs by model" ), model_id: Optional[str] = fastapi.Query( - default=None, description="Filter logs by model ID (litellm model deployment id)" + default=None, + description="Filter logs by model ID (litellm model deployment id)", ), key_alias: Optional[str] = fastapi.Query( default=None, description="Filter logs by key alias" @@ -1726,7 +1737,13 @@ async def ui_view_spend_logs( # noqa: PLR0915 ) # Validate sort_by and sort_order - valid_sort_fields = {"spend", "total_tokens", "startTime", "endTime", "request_duration_ms"} + valid_sort_fields = { + "spend", + "total_tokens", + "startTime", + "endTime", + "request_duration_ms", + } if sort_by not in valid_sort_fields: raise ProxyException( message=f"Invalid sort_by: {sort_by}. Must be one of: {', '.join(sorted(valid_sort_fields))}", @@ -1753,7 +1770,11 @@ async def ui_view_spend_logs( # noqa: PLR0915 return datetime.strptime(date_str, fmt).replace(tzinfo=timezone.utc) except ValueError: continue - expected = "'YYYY-MM-DD' or 'YYYY-MM-DD HH:MM:SS'" if is_v2 else "'YYYY-MM-DD HH:MM:SS'" + expected = ( + "'YYYY-MM-DD' or 'YYYY-MM-DD HH:MM:SS'" + if is_v2 + else "'YYYY-MM-DD HH:MM:SS'" + ) raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"Invalid date format: {date_str}. Expected: {expected}", @@ -1796,22 +1817,28 @@ async def ui_view_spend_logs( # noqa: PLR0915 # Build metadata filters metadata_filters = [] if key_alias is not None: - metadata_filters.append({ - "path": ["user_api_key_alias"], - "string_contains": key_alias, - }) + metadata_filters.append( + { + "path": ["user_api_key_alias"], + "string_contains": key_alias, + } + ) if error_code is not None: - metadata_filters.append({ - "path": ["error_information", "error_code"], - "equals": f'"{error_code}"', - }) + metadata_filters.append( + { + "path": ["error_information", "error_code"], + "equals": f'"{error_code}"', + } + ) if error_message is not None: - metadata_filters.append({ - "path": ["error_information", "error_message"], - "string_contains": error_message, - }) + metadata_filters.append( + { + "path": ["error_information", "error_message"], + "string_contains": error_message, + } + ) if metadata_filters: if len(metadata_filters) == 1: @@ -1919,16 +1946,24 @@ async def ui_view_spend_logs( # noqa: PLR0915 sql_params.append(f"%{key_alias}%") p += 1 if error_code is not None: - sql_conditions.append(f"metadata->'error_information'->>'error_code' = ${p}") + sql_conditions.append( + f"metadata->'error_information'->>'error_code' = ${p}" + ) sql_params.append(error_code) p += 1 if error_message is not None: - sql_conditions.append(f"metadata->'error_information'->>'error_message' LIKE ${p}") + sql_conditions.append( + f"metadata->'error_information'->>'error_message' LIKE ${p}" + ) sql_params.append(f"%{error_message}%") p += 1 # Quote column names that need quoting in SQL - _sql_col = f'"{order_column}"' if order_column in ("startTime", "endTime") else order_column + _sql_col = ( + f'"{order_column}"' + if order_column in ("startTime", "endTime") + else order_column + ) _sql_dir = "ASC" if order_direction == "asc" else "DESC" sql_query = f""" @@ -3218,7 +3253,9 @@ async def ui_view_session_spend_logs( ORDER BY "startTime" ASC LIMIT $2 OFFSET $3 """ - result = await prisma_client.db.query_raw(sql_query, session_id, page_size, skip) + result = await prisma_client.db.query_raw( + sql_query, session_id, page_size, skip + ) total_pages = (total_records + page_size - 1) // page_size @@ -3280,9 +3317,17 @@ async def _build_ui_spend_logs_response( if enrich_session_counts: session_ids = list( { - (row.get("session_id") if isinstance(row, dict) else getattr(row, "session_id", None)) + ( + row.get("session_id") + if isinstance(row, dict) + else getattr(row, "session_id", None) + ) for row in data - if (row.get("session_id") if isinstance(row, dict) else getattr(row, "session_id", None)) + if ( + row.get("session_id") + if isinstance(row, dict) + else getattr(row, "session_id", None) + ) } ) if session_ids: @@ -3304,11 +3349,7 @@ async def _build_ui_spend_logs_response( if enrich_session_counts: enriched: List[dict] = [] for row in data: - row_dict = ( - dict(row) - if isinstance(row, dict) - else row.model_dump() - ) + row_dict = dict(row) if isinstance(row, dict) else row.model_dump() sid = row_dict.get("session_id") row_dict["session_total_count"] = count_map.get(sid, 1) if sid else 1 enriched.append(row_dict) @@ -3383,7 +3424,11 @@ def _can_user_view_spend_log(user_api_key_dict: UserAPIKeyAuth) -> bool: """ user_role = user_api_key_dict.user_role user_id = user_api_key_dict.user_id - return user_role in ( - LitellmUserRoles.INTERNAL_USER, - LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, - ) and user_id is not None + return ( + user_role + in ( + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, + ) + and user_id is not None + ) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index b640eaa370..3eacc19a6d 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -15,21 +15,26 @@ from litellm.constants import ( LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, ) -from litellm.constants import \ - MAX_STRING_LENGTH_PROMPT_IN_DB as DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB +from litellm.constants import ( + MAX_STRING_LENGTH_PROMPT_IN_DB as DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB, +) from litellm.constants import REDACTED_BY_LITELM_STRING from litellm.litellm_core_utils.core_helpers import ( - get_litellm_metadata_from_kwargs, reconstruct_model_name) + get_litellm_metadata_from_kwargs, + reconstruct_model_name, +) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload from litellm.proxy.utils import PrismaClient, hash_token -from litellm.types.utils import (CostBreakdown, - StandardLoggingGuardrailInformation, - StandardLoggingMCPToolCall, - StandardLoggingModelInformation, - StandardLoggingPayload, - StandardLoggingVectorStoreRequest, - VectorStoreSearchResponse) +from litellm.types.utils import ( + CostBreakdown, + StandardLoggingGuardrailInformation, + StandardLoggingMCPToolCall, + StandardLoggingModelInformation, + StandardLoggingPayload, + StandardLoggingVectorStoreRequest, + VectorStoreSearchResponse, +) from litellm.utils import get_end_user_id_for_cost_tracking @@ -121,9 +126,9 @@ def _get_spend_logs_metadata( clean_metadata["applied_guardrails"] = applied_guardrails clean_metadata["batch_models"] = batch_models clean_metadata["mcp_tool_call_metadata"] = mcp_tool_call_metadata - clean_metadata["vector_store_request_metadata"] = ( - _get_vector_store_request_for_spend_logs_payload(vector_store_request_metadata) - ) + clean_metadata[ + "vector_store_request_metadata" + ] = _get_vector_store_request_for_spend_logs_payload(vector_store_request_metadata) clean_metadata["guardrail_information"] = guardrail_information clean_metadata["usage_object"] = usage_object clean_metadata["model_map_information"] = model_map_information @@ -501,7 +506,6 @@ def _get_session_id_for_spend_log( """ from litellm._uuid import uuid - if ( standard_logging_payload is not None and standard_logging_payload.get("trace_id") is not None @@ -782,7 +786,9 @@ def _get_proxy_server_request_for_spend_logs_payload( # Apply message redaction if turn_off_message_logging is enabled if kwargs is not None: from litellm.litellm_core_utils.redact_messages import ( - perform_redaction, should_redact_message_logging) + perform_redaction, + should_redact_message_logging, + ) # Build model_call_details dict to check redaction settings model_call_details = { @@ -853,7 +859,9 @@ def _get_response_for_spend_logs_payload( # Apply message redaction if turn_off_message_logging is enabled if kwargs is not None: from litellm.litellm_core_utils.redact_messages import ( - perform_redaction, should_redact_message_logging) + perform_redaction, + should_redact_message_logging, + ) litellm_params = kwargs.get("litellm_params", {}) model_call_details = { diff --git a/litellm/proxy/types_utils/utils.py b/litellm/proxy/types_utils/utils.py index ed50da3aa1..676d7fb51b 100644 --- a/litellm/proxy/types_utils/utils.py +++ b/litellm/proxy/types_utils/utils.py @@ -12,7 +12,7 @@ def get_instance_fn(value: str, config_file_path: Optional[str] = None) -> Any: # Check if value starts with s3:// or gcs:// if value.startswith("s3://") or value.startswith("gcs://"): return _load_instance_from_remote_storage(value, config_file_path) - + # Split the path by dots to separate module from instance parts = value.split(".") @@ -28,9 +28,7 @@ def get_instance_fn(value: str, config_file_path: Optional[str] = None) -> Any: # Check if the file exists before trying to load it if not os.path.exists(module_file_path): - raise ImportError( - f"Could not find module file {module_file_path}" - ) + raise ImportError(f"Could not find module file {module_file_path}") spec = importlib.util.spec_from_file_location(module_name, module_file_path) # type: ignore if spec is None: @@ -63,18 +61,20 @@ def get_instance_fn(value: str, config_file_path: Optional[str] = None) -> Any: raise e -def _load_instance_from_remote_storage(remote_url: str, config_file_path: Optional[str] = None) -> Any: +def _load_instance_from_remote_storage( + remote_url: str, config_file_path: Optional[str] = None +) -> Any: """ Load custom logger instance from S3 or GCS URL. - + Expected format: - s3://bucket-name/path/to/module.instance_name - gcs://bucket-name/path/to/module.instance_name - + Args: remote_url (str): The s3:// or gcs:// URL config_file_path (str): Optional config file path for temp directory context - + Returns: Any: The loaded instance """ @@ -90,19 +90,21 @@ def _load_instance_from_remote_storage(remote_url: str, config_file_path: Option url_without_prefix = remote_url[6:] # Remove 'gcs://' else: raise ValueError(f"Unsupported URL scheme in {remote_url}") - + # Split bucket and path parts = url_without_prefix.split("/", 1) if len(parts) < 2: - raise ValueError(f"Invalid URL format: {remote_url}. Expected: {storage_type}://bucket-name/path/to/module.instance") - + raise ValueError( + f"Invalid URL format: {remote_url}. Expected: {storage_type}://bucket-name/path/to/module.instance" + ) + bucket_name = parts[0] path_and_module = parts[1] - + # Extract module path and instance name # Example: "loggers/custom_callbacks.proxy_handler_instance" # Handle case where user accidentally includes .py extension - if path_and_module.endswith('.py'): + if path_and_module.endswith(".py"): module_name_without_py = path_and_module[:-3] # Remove .py raise ValueError( f"Invalid URL format in {remote_url}. " @@ -110,18 +112,20 @@ def _load_instance_from_remote_storage(remote_url: str, config_file_path: Option f"Expected format: {storage_type}://{bucket_name}/{module_name_without_py}.instance_name " f"(e.g., {storage_type}://{bucket_name}/{module_name_without_py}.proxy_handler_instance)" ) - + # Split by last dot to separate module from instance module_parts = path_and_module.split(".") if len(module_parts) < 2: - raise ValueError(f"Invalid module specification in {remote_url}. Expected: path/to/module.instance_name") - + raise ValueError( + f"Invalid module specification in {remote_url}. Expected: path/to/module.instance_name" + ) + instance_name = module_parts[-1] module_path = ".".join(module_parts[:-1]) - + # Create object key (file path in bucket) object_key = f"{module_path}.py" - + verbose_proxy_logger.debug( f"Loading custom logger from {storage_type}: bucket={bucket_name}, " f"object_key={object_key}, instance={instance_name}" @@ -130,65 +134,80 @@ def _load_instance_from_remote_storage(remote_url: str, config_file_path: Option import tempfile # Create temporary file for the downloaded module using the actual module name - temp_file = tempfile.NamedTemporaryFile(suffix='.py', delete=False) + temp_file = tempfile.NamedTemporaryFile(suffix=".py", delete=False) local_file_path = temp_file.name temp_file.close() # Close the file so we can write to it - + # Download the file if storage_type == "s3": from litellm.proxy.common_utils.load_config_utils import ( download_python_file_from_s3, ) + success = download_python_file_from_s3( bucket_name=bucket_name, object_key=object_key, local_file_path=local_file_path, ) else: # gcs - success = asyncio.run(_download_gcs_file_wrapper(bucket_name, object_key, local_file_path)) - + success = asyncio.run( + _download_gcs_file_wrapper(bucket_name, object_key, local_file_path) + ) + if not success: - raise ImportError(f"Failed to download {object_key} from {storage_type} bucket {bucket_name}") - + raise ImportError( + f"Failed to download {object_key} from {storage_type} bucket {bucket_name}" + ) + # Load the module from the downloaded file using the actual module name spec = importlib.util.spec_from_file_location(module_path, local_file_path) if spec is None or spec.loader is None: raise ImportError(f"Could not create module spec for {local_file_path}") - + module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) - + # Get the instance instance = getattr(module, instance_name) - + # Clean up the temporary file try: os.remove(local_file_path) except Exception as cleanup_error: - verbose_proxy_logger.warning(f"Could not clean up temporary file {local_file_path}: {cleanup_error}") - - verbose_proxy_logger.info(f"Successfully loaded custom logger from {remote_url}") + verbose_proxy_logger.warning( + f"Could not clean up temporary file {local_file_path}: {cleanup_error}" + ) + + verbose_proxy_logger.info( + f"Successfully loaded custom logger from {remote_url}" + ) return instance - + except Exception as e: - raise ImportError(f"Failed to load custom logger from {remote_url}: {str(e)}") from e + raise ImportError( + f"Failed to load custom logger from {remote_url}: {str(e)}" + ) from e -async def _download_gcs_file_wrapper(bucket_name: str, object_key: str, local_file_path: str) -> bool: +async def _download_gcs_file_wrapper( + bucket_name: str, object_key: str, local_file_path: str +) -> bool: """Wrapper for GCS download to handle async properly""" try: from litellm.proxy.common_utils.load_config_utils import ( download_python_file_from_gcs, ) - return await download_python_file_from_gcs(bucket_name, object_key, local_file_path) + + return await download_python_file_from_gcs( + bucket_name, object_key, local_file_path + ) except Exception as e: from litellm._logging import verbose_proxy_logger + verbose_proxy_logger.error(f"Error downloading from GCS: {str(e)}") return False - - def validate_custom_validate_return_type( fn: Optional[Callable[..., Any]], ) -> Optional[Callable[..., Literal[True]]]: diff --git a/litellm/proxy/ui_crud_endpoints/__init__.py b/litellm/proxy/ui_crud_endpoints/__init__.py index 2af6220183..1f4b379626 100644 --- a/litellm/proxy/ui_crud_endpoints/__init__.py +++ b/litellm/proxy/ui_crud_endpoints/__init__.py @@ -1,3 +1,3 @@ from .proxy_setting_endpoints import router as ui_crud_endpoints_router -__all__ = ["ui_crud_endpoints_router"] \ No newline at end of file +__all__ = ["ui_crud_endpoints_router"] diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 7245212dfa..8df215d998 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -91,7 +91,7 @@ class UISettings(BaseModel): require_auth_for_public_ai_hub: bool = Field( default=False, - description="If true, requires authentication for accessing the public AI Hub." + description="If true, requires authentication for accessing the public AI Hub.", ) forward_client_headers_to_llm_api: bool = Field( @@ -423,7 +423,9 @@ async def update_default_team_member_budget( async def _update_litellm_setting( - settings: Union[DefaultInternalUserParams, DefaultTeamSSOParams, MCPSemanticFilterSettings], + settings: Union[ + DefaultInternalUserParams, DefaultTeamSSOParams, MCPSemanticFilterSettings + ], settings_key: str, in_memory_var: Any, success_message: str, @@ -864,9 +866,7 @@ async def update_ui_theme_settings(theme_config: UIThemeConfig): verbose_proxy_logger.debug("Removed LITELLM_FAVICON_URL from config") if "LITELLM_FAVICON_URL" in os.environ: del os.environ["LITELLM_FAVICON_URL"] - verbose_proxy_logger.debug( - "Removed LITELLM_FAVICON_URL from environment" - ) + verbose_proxy_logger.debug("Removed LITELLM_FAVICON_URL from environment") # Handle environment variable encryption if needed stored_config = config.copy() @@ -1062,7 +1062,9 @@ async def get_ui_settings(): # Sync runtime flags into general_settings so the proxy picks them up # at runtime (covers server restart scenarios). - _flags_to_sync = {k: ui_settings[k] for k in _RUNTIME_GENERAL_SETTINGS_FLAGS if k in ui_settings} + _flags_to_sync = { + k: ui_settings[k] for k in _RUNTIME_GENERAL_SETTINGS_FLAGS if k in ui_settings + } if _flags_to_sync: from litellm.proxy.proxy_server import general_settings @@ -1153,7 +1155,9 @@ async def update_ui_settings( # Sync runtime flags to general_settings so the proxy picks them up # at runtime (general_settings is checked in pre-call utils). - _flags_to_sync = {k: ui_settings[k] for k in _RUNTIME_GENERAL_SETTINGS_FLAGS if k in ui_settings} + _flags_to_sync = { + k: ui_settings[k] for k in _RUNTIME_GENERAL_SETTINGS_FLAGS if k in ui_settings + } if _flags_to_sync: from litellm.proxy.proxy_server import general_settings diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index f9fa422680..00ca9043ca 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1796,8 +1796,7 @@ class ProxyLogging: if route is None: return False if not ( - RouteChecks.is_llm_api_route(route) or - RouteChecks.is_info_route(route) + RouteChecks.is_llm_api_route(route) or RouteChecks.is_info_route(route) ): return False @@ -1993,7 +1992,9 @@ class ProxyLogging: merged_headers: Dict[str, str] = {} try: # Build litellm_call_info — normalized routing metadata for callbacks - litellm_call_info = self._build_litellm_call_info(data=data, response=response) + litellm_call_info = self._build_litellm_call_info( + data=data, response=response + ) for callback in litellm.callbacks: _callback: Optional[CustomLogger] = None @@ -2030,9 +2031,7 @@ class ProxyLogging: return merged_headers @staticmethod - def _build_litellm_call_info( - data: dict, response: Any - ) -> Dict[str, Any]: + def _build_litellm_call_info(data: dict, response: Any) -> Dict[str, Any]: """ Build a normalized dict of routing metadata from response._hidden_params and data, abstracting away the metadata vs litellm_metadata split. @@ -2096,8 +2095,10 @@ class ProxyLogging: ## CHECK FOR MODEL-LEVEL GUARDRAILS (cached per-request) if not _guardrail_data_computed: - _cached_guardrail_data = _check_and_merge_model_level_guardrails( - data=data, llm_router=llm_router + _cached_guardrail_data = ( + _check_and_merge_model_level_guardrails( + data=data, llm_router=llm_router + ) ) _guardrail_data_computed = True @@ -3688,13 +3689,15 @@ class PrismaClient: probe_pid, _ = os.waitpid(pid, os.WNOHANG) except ChildProcessError: verbose_proxy_logger.debug( - "PID %s is not a child process; skipping waitpid watch.", pid, + "PID %s is not a child process; skipping waitpid watch.", + pid, ) return False if probe_pid == pid: verbose_proxy_logger.warning( - "prisma-query-engine PID %s already dead at watch start.", pid, + "prisma-query-engine PID %s already dead at watch start.", + pid, ) self._engine_confirmed_dead = True self._reap_all_zombies() @@ -3871,11 +3874,17 @@ class PrismaClient: waitpid thread nor pidfd are available. """ - if self._watching_engine or self._engine_pidfd >= 0 or self._engine_wait_thread is not None: + if ( + self._watching_engine + or self._engine_pidfd >= 0 + or self._engine_wait_thread is not None + ): return pid = self._get_engine_pid() if pid == 0: - verbose_proxy_logger.debug("Could not find prisma-query-engine PID; engine death detection unavailable.") + verbose_proxy_logger.debug( + "Could not find prisma-query-engine PID; engine death detection unavailable." + ) return self._engine_pid = pid self._engine_confirmed_dead = False @@ -3884,15 +3893,18 @@ class PrismaClient: pidfd_ok = False if waitpid_ok else self._try_pidfd_watch(pid) if waitpid_ok: verbose_proxy_logger.info( - "Watching engine PID %s via waitpid thread.", pid, + "Watching engine PID %s via waitpid thread.", + pid, ) elif pidfd_ok: verbose_proxy_logger.info( - "Watching engine PID %s via pidfd.", pid, + "Watching engine PID %s via pidfd.", + pid, ) else: verbose_proxy_logger.info( - "Watching engine PID %s via os.kill polling.", pid, + "Watching engine PID %s via os.kill polling.", + pid, ) self._watching_engine = True asyncio.create_task(self._poll_engine_proc()) @@ -3915,7 +3927,9 @@ class PrismaClient: blip -- disconnect, connect, SELECT 1). """ effective_timeout = ( - timeout_seconds if timeout_seconds is not None else self._db_watchdog_reconnect_timeout_seconds + timeout_seconds + if timeout_seconds is not None + else self._db_watchdog_reconnect_timeout_seconds ) engine_is_dead = self._engine_confirmed_dead or ( @@ -3935,14 +3949,18 @@ class PrismaClient: async def _do_heavy_reconnect() -> None: db_url = os.getenv("DATABASE_URL", "") if not db_url: - verbose_proxy_logger.error("DATABASE_URL not set; cannot recreate Prisma client.") + verbose_proxy_logger.error( + "DATABASE_URL not set; cannot recreate Prisma client." + ) raise RuntimeError("DATABASE_URL not set") await self.db.recreate_prisma_client(db_url) await self._start_engine_watcher() await asyncio.wait_for(_do_heavy_reconnect(), timeout=effective_timeout) else: - verbose_proxy_logger.debug("Performing Prisma DB reconnect (engine alive or unknown).") + verbose_proxy_logger.debug( + "Performing Prisma DB reconnect (engine alive or unknown)." + ) async def _do_direct_reconnect() -> None: try: @@ -4041,7 +4059,9 @@ class PrismaClient: if lock_timeout_seconds is None: async with self._db_reconnect_lock: - return await self._attempt_reconnect_inside_lock(force, reason, timeout_seconds) + return await self._attempt_reconnect_inside_lock( + force, reason, timeout_seconds + ) lock_acquired_by_timeout_task = False @@ -4090,14 +4110,17 @@ class PrismaClient: return False try: - return await self._attempt_reconnect_inside_lock(force, reason, timeout_seconds) + return await self._attempt_reconnect_inside_lock( + force, reason, timeout_seconds + ) finally: self._db_reconnect_lock.release() async def start_db_health_watchdog_task(self) -> None: """Start background tasks that monitor DB health: - A periodic SELECT 1 probe that triggers reconnect on network/connection failure. - - A process-level watcher that detects engine death via waitpid thread, pidfd, or os.kill polling.""" + - A process-level watcher that detects engine death via waitpid thread, pidfd, or os.kill polling. + """ if self._db_health_watchdog_enabled is not True: verbose_proxy_logger.debug( "Prisma DB health watchdog disabled via PRISMA_HEALTH_WATCHDOG_ENABLED" @@ -4557,9 +4580,9 @@ class ProxyUpdateSpend: :MAX_LOGS_PER_INTERVAL ] # Remove the logs we're about to process - prisma_client.spend_log_transactions = prisma_client.spend_log_transactions[ - len(logs_to_process) : - ] + prisma_client.spend_log_transactions = ( + prisma_client.spend_log_transactions[len(logs_to_process) :] + ) popped_batch = True if len(logs_to_process) > 0: verbose_proxy_logger.info( @@ -4713,9 +4736,7 @@ async def update_spend_logs_job( return async with prisma_client._spend_log_transactions_lock: - logs_to_process = prisma_client.spend_log_transactions[ - :MAX_LOGS_PER_INTERVAL - ] + logs_to_process = prisma_client.spend_log_transactions[:MAX_LOGS_PER_INTERVAL] prisma_client.spend_log_transactions = prisma_client.spend_log_transactions[ len(logs_to_process) : ] @@ -4733,6 +4754,7 @@ async def update_spend_logs_job( from litellm.proxy.guardrails.usage_tracking import ( process_spend_logs_guardrail_usage, ) + await process_spend_logs_guardrail_usage( prisma_client=prisma_client, logs_to_process=logs_to_process, @@ -4746,6 +4768,7 @@ async def update_spend_logs_job( # Tool usage tracking (same batch): SpendLogToolIndex for "last N requests for tool X" try: from litellm.proxy.db.spend_log_tool_index import process_spend_logs_tool_usage + await process_spend_logs_tool_usage( prisma_client=prisma_client, logs_to_process=logs_to_process, @@ -5340,7 +5363,9 @@ async def get_available_models_for_user( user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) - await validate_membership(user_api_key_dict=user_api_key_dict, team_table=team_object) + await validate_membership( + user_api_key_dict=user_api_key_dict, team_table=team_object + ) team_models = team_object.models team_models = get_team_models( diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index b43ca29e3a..63ce5c104d 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -24,30 +24,30 @@ def _check_vector_store_access( ) -> bool: """ Check if the user has access to the vector store based on team membership. - + Args: vector_store: The vector store to check access for user_api_key_dict: User API key authentication info - + Returns: True if user has access, False otherwise - + Access rules: - If vector store has no team_id, it's accessible to all (legacy behavior) - If user's team_id matches the vector store's team_id, access is granted - Otherwise, access is denied """ vector_store_team_id = vector_store.get("team_id") - + # If vector store has no team_id, it's accessible to all (legacy behavior) if vector_store_team_id is None: return True - + # Check if user's team matches the vector store's team user_team_id = user_api_key_dict.team_id if user_team_id == vector_store_team_id: return True - + return False @@ -58,30 +58,32 @@ def _update_request_data_with_litellm_managed_vector_store_registry( ) -> Dict: """ Update the request data with the litellm managed vector store registry. - + Args: data: Request data to update vector_store_id: ID of the vector store user_api_key_dict: User API key authentication info for access control - + Raises: HTTPException: If user doesn't have access to the vector store """ if litellm.vector_store_registry is not None: - vector_store_to_run: Optional[LiteLLM_ManagedVectorStore] = ( - litellm.vector_store_registry.get_litellm_managed_vector_store_from_registry( - vector_store_id=vector_store_id - ) + vector_store_to_run: Optional[ + LiteLLM_ManagedVectorStore + ] = litellm.vector_store_registry.get_litellm_managed_vector_store_from_registry( + vector_store_id=vector_store_id ) if vector_store_to_run is not None: # Check access control if user_api_key_dict is provided if user_api_key_dict is not None: - if not _check_vector_store_access(vector_store_to_run, user_api_key_dict): + if not _check_vector_store_access( + vector_store_to_run, user_api_key_dict + ): raise HTTPException( status_code=403, detail="Access denied: You do not have permission to access this vector store", ) - + if "custom_llm_provider" in vector_store_to_run: data["custom_llm_provider"] = vector_store_to_run.get( "custom_llm_provider" @@ -103,7 +105,8 @@ def _update_request_data_with_litellm_managed_vector_store_registry( dependencies=[Depends(user_api_key_auth)], ) @router.post( - "/vector_stores/{vector_store_id:path}/search", dependencies=[Depends(user_api_key_auth)] + "/vector_stores/{vector_store_id:path}/search", + dependencies=[Depends(user_api_key_auth)], ) async def vector_store_search( request: Request, @@ -146,7 +149,7 @@ async def vector_store_search( # 2. Extracting model and provider resource ID # 3. Setting up proper routing # 4. Authentication checks - + processor = ProxyBaseLLMRequestProcessing(data=data) try: return await processor.base_process_llm_request( @@ -188,7 +191,7 @@ async def vector_store_create( API Reference: https://platform.openai.com/docs/api-reference/vector-stores/create - + Supports target_model_names parameter for creating vector stores across multiple models: ```json { @@ -213,10 +216,10 @@ async def vector_store_create( ) data = await _read_request_body(request=request) - + # Check for target_model_names parameter target_model_names = data.pop("target_model_names", None) - + if target_model_names: # Use managed vector stores for multi-model support if isinstance(target_model_names, str): @@ -228,21 +231,23 @@ async def vector_store_create( status_code=400, detail="target_model_names must be a comma-separated string or list of model names", ) - + # Get managed vector stores hook - managed_vector_stores: Any = proxy_logging_obj.get_proxy_hook("managed_vector_stores") + managed_vector_stores: Any = proxy_logging_obj.get_proxy_hook( + "managed_vector_stores" + ) if managed_vector_stores is None: raise HTTPException( status_code=500, detail="Managed vector stores not configured. Please ensure the proxy is initialized with database support.", ) - + if llm_router is None: raise HTTPException( status_code=500, detail="LLM Router not initialized. Ensure models are added to proxy.", ) - + # Create vector store across multiple models response = await managed_vector_stores.acreate_vector_store( create_request=data, @@ -251,9 +256,9 @@ async def vector_store_create( litellm_parent_otel_span=user_api_key_dict.parent_otel_span, user_api_key_dict=user_api_key_dict, ) - + return response - + processor = ProxyBaseLLMRequestProcessing(data=data) try: return await processor.base_process_llm_request( @@ -283,8 +288,12 @@ async def vector_store_create( ) -@router.get("/v1/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)]) -@router.get("/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)]) +@router.get( + "/v1/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)] +) +@router.get( + "/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)] +) async def vector_store_retrieve( request: Request, vector_store_id: str, @@ -416,8 +425,12 @@ async def vector_store_list( ) -@router.post("/v1/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)]) -@router.post("/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)]) +@router.post( + "/v1/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)] +) +@router.post( + "/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)] +) async def vector_store_update( request: Request, vector_store_id: str, @@ -482,8 +495,12 @@ async def vector_store_update( ) -@router.delete("/v1/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)]) -@router.delete("/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)]) +@router.delete( + "/v1/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)] +) +@router.delete( + "/vector_stores/{vector_store_id}", dependencies=[Depends(user_api_key_auth)] +) async def vector_store_delete( request: Request, vector_store_id: str, diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index 068f4217e0..cf57999366 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -43,21 +43,21 @@ def _resolve_embedding_config_from_router( ) -> Optional[Dict[str, Any]]: """ Resolve embedding config from router's config-defined models. - + Config-defined models (from proxy_config.yaml) are stored in the router's model_list, not in the database. This function looks up the model in the router and extracts api_key, api_base, and api_version from the deployment's litellm_params. - + Args: embedding_model: The embedding model string (e.g., "text-embedding-ada-002" or "azure/text-embedding-3-large") llm_router: The LiteLLM router instance - + Returns: Dictionary with api_key, api_base, and api_version if model found, None otherwise """ if not embedding_model or llm_router is None: return None - + # Extract model name candidates - could be "text-embedding-ada-002" or "azure/text-embedding-3-large" # Try exact match first, then try without provider prefix model_name_candidates = [embedding_model] @@ -65,7 +65,7 @@ def _resolve_embedding_config_from_router( # If it has a provider prefix, also try without it _, model_name = embedding_model.split("/", 1) model_name_candidates.append(model_name) - + # Try to find model in router for model_name in model_name_candidates: try: @@ -73,13 +73,13 @@ def _resolve_embedding_config_from_router( deployment = llm_router.get_deployment_by_model_group_name( model_group_name=model_name ) - + if deployment is not None and deployment.litellm_params is not None: litellm_params = deployment.litellm_params - + # Build embedding config from model params embedding_config: Dict[str, Any] = {} - + # Extract api_key api_key = getattr(litellm_params, "api_key", None) if api_key: @@ -87,7 +87,7 @@ def _resolve_embedding_config_from_router( if isinstance(api_key, str) and api_key.startswith("os.environ/"): api_key = get_secret(api_key) embedding_config["api_key"] = api_key - + # Extract api_base api_base = getattr(litellm_params, "api_base", None) if api_base: @@ -95,7 +95,7 @@ def _resolve_embedding_config_from_router( if isinstance(api_base, str) and api_base.startswith("os.environ/"): api_base = get_secret(api_base) embedding_config["api_base"] = api_base - + # Extract api_version api_version = getattr(litellm_params, "api_version", None) if api_version: @@ -104,7 +104,7 @@ def _resolve_embedding_config_from_router( project_id = getattr(litellm_params, "project_id", None) if project_id: embedding_config["project_id"] = project_id - + # Only return config if we have at least api_key or api_base if embedding_config: verbose_proxy_logger.debug( @@ -116,7 +116,7 @@ def _resolve_embedding_config_from_router( f"Error resolving embedding config from router for model {model_name}: {str(e)}" ) continue - + return None @@ -125,21 +125,21 @@ async def _resolve_embedding_config_from_db( ) -> Optional[Dict[str, Any]]: """ Resolve embedding config from database model configuration. - + If litellm_embedding_model is provided but litellm_embedding_config is not, this function looks up the model in the database and extracts api_key, api_base, and api_version from the model's litellm_params to build the embedding config. - + Args: embedding_model: The embedding model string (e.g., "text-embedding-ada-002" or "azure/text-embedding-3-large") prisma_client: The Prisma client instance - + Returns: Dictionary with api_key, api_base, and api_version if model found, None otherwise """ if not embedding_model: return None - + # Extract model name - could be "text-embedding-ada-002" or "azure/text-embedding-3-large" # Try to find model by exact match first, then try without provider prefix model_name_candidates = [embedding_model] @@ -147,20 +147,20 @@ async def _resolve_embedding_config_from_db( # If it has a provider prefix, also try without it _, model_name = embedding_model.split("/", 1) model_name_candidates.append(model_name) - + # Try to find model in database for model_name in model_name_candidates: try: db_model = await prisma_client.db.litellm_proxymodeltable.find_first( where={"model_name": model_name} ) - + if db_model and db_model.litellm_params: # Extract litellm_params (could be dict or JSON string) model_params = db_model.litellm_params if isinstance(model_params, str): model_params = json.loads(model_params) - + # Decrypt values from database (similar to how proxy_server.py does it) # Values stored in DB are encrypted, so we need to decrypt them first decrypted_params = {} @@ -176,10 +176,10 @@ async def _resolve_embedding_config_from_db( decrypted_params[k] = v else: decrypted_params = model_params - + # Build embedding config from model params embedding_config = {} - + # Extract api_key api_key = decrypted_params.get("api_key") if api_key: @@ -187,7 +187,7 @@ async def _resolve_embedding_config_from_db( if isinstance(api_key, str) and api_key.startswith("os.environ/"): api_key = get_secret(api_key) embedding_config["api_key"] = api_key - + # Extract api_base api_base = decrypted_params.get("api_base") if api_base: @@ -195,12 +195,12 @@ async def _resolve_embedding_config_from_db( if isinstance(api_base, str) and api_base.startswith("os.environ/"): api_base = get_secret(api_base) embedding_config["api_base"] = api_base - + # Extract api_version api_version = decrypted_params.get("api_version") if api_version: embedding_config["api_version"] = api_version - + # Only return config if we have at least api_key or api_base if embedding_config: verbose_proxy_logger.debug( @@ -212,7 +212,7 @@ async def _resolve_embedding_config_from_db( f"Error resolving embedding config for model {model_name}: {str(e)}" ) continue - + return None @@ -221,52 +221,50 @@ async def _resolve_embedding_config( ) -> Optional[Dict[str, Any]]: """ Resolve embedding config from either router (config-defined) or database models. - + This function first checks the router for config-defined models, then falls back to the database. This allows users to use models defined in either location. - + Args: embedding_model: The embedding model string (e.g., "text-embedding-ada-002" or "azure/text-embedding-3-large") prisma_client: The Prisma client instance llm_router: The LiteLLM router instance (optional, will be imported if not provided) - + Returns: Dictionary with api_key, api_base, and api_version if model found, None otherwise """ if not embedding_model: return None - + # Import llm_router if not provided if llm_router is None: try: from litellm.proxy.proxy_server import llm_router except ImportError: llm_router = None - + # First try to resolve from router (config-defined models) if llm_router is not None: router_config = _resolve_embedding_config_from_router( - embedding_model=embedding_model, - llm_router=llm_router + embedding_model=embedding_model, llm_router=llm_router ) if router_config: verbose_proxy_logger.debug( f"Resolved embedding config from router for model {embedding_model}" ) return router_config - + # Fall back to database if prisma_client is not None: db_config = await _resolve_embedding_config_from_db( - embedding_model=embedding_model, - prisma_client=prisma_client + embedding_model=embedding_model, prisma_client=prisma_client ) if db_config: verbose_proxy_logger.debug( f"Resolved embedding config from database for model {embedding_model}" ) return db_config - + verbose_proxy_logger.debug( f"Could not resolve embedding config for model {embedding_model} from router or database" ) @@ -282,30 +280,30 @@ def _check_vector_store_access( ) -> bool: """ Check if the user has access to the vector store based on team membership. - + Args: vector_store: The vector store to check access for user_api_key_dict: User API key authentication info - + Returns: True if user has access, False otherwise - + Access rules: - If vector store has no team_id, it's accessible to all (legacy behavior) - If user's team_id matches the vector store's team_id, access is granted - Otherwise, access is denied """ vector_store_team_id = vector_store.get("team_id") - + # If vector store has no team_id, it's accessible to all (legacy behavior) if vector_store_team_id is None: return True - + # Check if user's team matches the vector store's team user_team_id = user_api_key_dict.team_id if user_team_id == vector_store_team_id: return True - + return False @@ -323,23 +321,23 @@ async def create_vector_store_in_db( ) -> LiteLLM_ManagedVectorStore: """ Helper function to create a vector store in the database. - + This function handles: - Checking if vector store already exists - Creating the vector store in the database - Adding it to the vector store registry - + Returns: LiteLLM_ManagedVectorStore: The created vector store object - + Raises: HTTPException: If vector store already exists or database error occurs """ from litellm.types.router import GenericLiteLLMParams - + if prisma_client is None: raise HTTPException(status_code=500, detail="Database not connected") - + # Check if vector store already exists existing_vector_store = ( await prisma_client.db.litellm_managedvectorstorestable.find_unique( @@ -351,13 +349,13 @@ async def create_vector_store_in_db( status_code=400, detail=f"Vector store with ID {vector_store_id} already exists", ) - + # Prepare data for database data_to_create: Dict[str, Any] = { "vector_store_id": vector_store_id, "custom_llm_provider": custom_llm_provider, } - + if vector_store_name is not None: data_to_create["vector_store_name"] = vector_store_name if vector_store_description is not None: @@ -370,51 +368,48 @@ async def create_vector_store_in_db( data_to_create["team_id"] = team_id if user_id is not None: data_to_create["user_id"] = user_id - + # Handle litellm_params - always provide at least an empty dict if litellm_params: # Auto-resolve embedding config if embedding model is provided but config is not embedding_model = litellm_params.get("litellm_embedding_model") if embedding_model and not litellm_params.get("litellm_embedding_config"): resolved_config = await _resolve_embedding_config( - embedding_model=embedding_model, - prisma_client=prisma_client + embedding_model=embedding_model, prisma_client=prisma_client ) if resolved_config: litellm_params["litellm_embedding_config"] = resolved_config verbose_proxy_logger.info( f"Auto-resolved embedding config for model {embedding_model}" ) - - litellm_params_dict = GenericLiteLLMParams( - **litellm_params - ).model_dump(exclude_none=True) + + litellm_params_dict = GenericLiteLLMParams(**litellm_params).model_dump( + exclude_none=True + ) data_to_create["litellm_params"] = safe_dumps(litellm_params_dict) else: # Provide empty dict if no litellm_params provided data_to_create["litellm_params"] = safe_dumps({}) - + # Create in database - _new_vector_store = ( - await prisma_client.db.litellm_managedvectorstorestable.create( - data=data_to_create - ) + _new_vector_store = await prisma_client.db.litellm_managedvectorstorestable.create( + data=data_to_create ) - + new_vector_store: LiteLLM_ManagedVectorStore = LiteLLM_ManagedVectorStore( **_new_vector_store.model_dump() ) - + # Add vector store to registry if litellm.vector_store_registry is not None: litellm.vector_store_registry.add_vector_store_to_registry( vector_store=new_vector_store ) - + verbose_proxy_logger.info( f"Vector store {vector_store_id} created in database successfully" ) - + return new_vector_store @@ -447,19 +442,19 @@ async def new_vector_store( try: vector_store_id = vector_store.get("vector_store_id") custom_llm_provider = vector_store.get("custom_llm_provider") - + if not vector_store_id or not custom_llm_provider: raise HTTPException( status_code=400, - detail="vector_store_id and custom_llm_provider are required" + detail="vector_store_id and custom_llm_provider are required", ) - + # Extract and validate metadata metadata = vector_store.get("vector_store_metadata") validated_metadata: Optional[Dict] = None if metadata is not None and isinstance(metadata, dict): validated_metadata = metadata - + new_vector_store = await create_vector_store_in_db( vector_store_id=vector_store_id, custom_llm_provider=custom_llm_provider, @@ -521,27 +516,27 @@ async def list_vector_stores( vector_stores_from_db = await VectorStoreRegistry._get_vector_stores_from_db( prisma_client=prisma_client ) - + # Build map from database vector stores for vector_store in vector_stores_from_db: vector_store_id = vector_store.get("vector_store_id", None) if vector_store_id: vector_store_map[vector_store_id] = vector_store db_vector_store_ids.add(vector_store_id) - + # Process in-memory vector stores if litellm.vector_store_registry is not None: in_memory_vector_stores = copy.deepcopy( litellm.vector_store_registry.vector_stores ) - + vector_stores_to_delete_from_memory: List[str] = [] - + for vector_store in in_memory_vector_stores: vector_store_id = vector_store.get("vector_store_id", None) if not vector_store_id: continue - + # If vector store is in memory but NOT in database, it was deleted if vector_store_id not in db_vector_store_ids: verbose_proxy_logger.info( @@ -551,7 +546,7 @@ async def list_vector_stores( # If not in our map yet, add it (only in-memory, not in DB) elif vector_store_id not in vector_store_map: vector_store_map[vector_store_id] = vector_store - + # Synchronize in-memory registry with database # 1. Remove deleted vector stores from memory for vs_id in vector_stores_to_delete_from_memory: @@ -561,22 +556,22 @@ async def list_vector_stores( verbose_proxy_logger.debug( f"Removed deleted vector store {vs_id} from in-memory registry" ) - + # 2. Update in-memory registry with database versions (for updates) for vector_store in vector_stores_from_db: vector_store_id = vector_store.get("vector_store_id", None) if vector_store_id: litellm.vector_store_registry.update_vector_store_in_registry( - vector_store_id=vector_store_id, - updated_data=vector_store + vector_store_id=vector_store_id, updated_data=vector_store ) # Filter vector stores based on team access accessible_vector_stores = [ - vs for vs in vector_store_map.values() + vs + for vs in vector_store_map.values() if _check_vector_store_access(vs, user_api_key_dict) ] - + total_count = len(accessible_vector_stores) total_pages = (total_count + page_size - 1) // page_size @@ -622,7 +617,7 @@ async def delete_vector_store( db_vector_store_exists = False memory_vector_store_exists = False vector_store_to_check = None - + existing_vector_store = ( await prisma_client.db.litellm_managedvectorstorestable.find_unique( where={"vector_store_id": data.vector_store_id} @@ -633,7 +628,7 @@ async def delete_vector_store( vector_store_to_check = LiteLLM_ManagedVectorStore( **existing_vector_store.model_dump() ) - + # Check in-memory registry if litellm.vector_store_registry is not None: memory_vector_store = litellm.vector_store_registry.get_litellm_managed_vector_store_from_registry( @@ -643,14 +638,14 @@ async def delete_vector_store( memory_vector_store_exists = True if vector_store_to_check is None: vector_store_to_check = memory_vector_store - + # If not found in either location, raise 404 if not db_vector_store_exists and not memory_vector_store_exists: raise HTTPException( status_code=404, detail=f"Vector store with ID {data.vector_store_id} not found", ) - + # Check access control if vector_store_to_check and not _check_vector_store_access( vector_store_to_check, user_api_key_dict @@ -674,7 +669,7 @@ async def delete_vector_store( return { "status": "success", - "message": f"Vector store {data.vector_store_id} deleted successfully" + "message": f"Vector store {data.vector_store_id} deleted successfully", } except HTTPException: raise @@ -713,7 +708,7 @@ async def get_vector_store_info( status_code=403, detail="Access denied: You do not have permission to access this vector store", ) - + vector_store_metadata = vector_store.get("vector_store_metadata") # Parse metadata if it's a JSON string parsed_metadata: Optional[dict] = None @@ -750,7 +745,7 @@ async def get_vector_store_info( status_code=404, detail=f"Vector store with ID {data.vector_store_id} not found", ) - + # Check access control for DB vector store vector_store_dict = vector_store.model_dump() # type: ignore[attr-defined] vector_store_typed = LiteLLM_ManagedVectorStore(**vector_store_dict) @@ -790,30 +785,31 @@ async def update_vector_store( try: update_data = data.model_dump(exclude_unset=True) vector_store_id = update_data.pop("vector_store_id") - + # Handle metadata serialization if update_data.get("vector_store_metadata") is not None: update_data["vector_store_metadata"] = safe_dumps( update_data["vector_store_metadata"] ) - + # Handle litellm_params if provided if "litellm_params" in update_data: _input_litellm_params: dict = update_data.get("litellm_params", {}) or {} - + # Auto-resolve embedding config if embedding model is provided but config is not embedding_model = _input_litellm_params.get("litellm_embedding_model") - if embedding_model and not _input_litellm_params.get("litellm_embedding_config"): + if embedding_model and not _input_litellm_params.get( + "litellm_embedding_config" + ): resolved_config = await _resolve_embedding_config( - embedding_model=embedding_model, - prisma_client=prisma_client + embedding_model=embedding_model, prisma_client=prisma_client ) if resolved_config: _input_litellm_params["litellm_embedding_config"] = resolved_config verbose_proxy_logger.info( f"Auto-resolved embedding config for model {embedding_model}" ) - + litellm_params_dict = GenericLiteLLMParams( **_input_litellm_params ).model_dump(exclude_none=True) @@ -840,7 +836,7 @@ async def update_vector_store( return { "status": "success", "message": f"Vector store {vector_store_id} updated successfully", - "vector_store": updated_vs + "vector_store": updated_vs, } except Exception as e: verbose_proxy_logger.exception(f"Error updating vector store: {str(e)}") diff --git a/litellm/proxy/vector_store_files_endpoints/endpoints.py b/litellm/proxy/vector_store_files_endpoints/endpoints.py index e3e022c9cf..7cdf865692 100644 --- a/litellm/proxy/vector_store_files_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_files_endpoints/endpoints.py @@ -35,22 +35,22 @@ def _update_request_data_with_managed_file_id( ) -> tuple[Dict, Optional[str]]: """ Update request data with model routing information from managed file ID. - + This function handles two types of file IDs: 1. Simple encoded file IDs (format: litellm:{file_id};model,{model}) 2. Unified managed file IDs (format: litellm_proxy:{mime};unified_id,{uuid};...;llm_output_file_id,{file_id};...) - + For unified managed file IDs, it: - Decodes the unified ID to extract the actual provider file ID (llm_output_file_id) - Extracts the model routing information (target_model_names) - Updates data with credentials for the correct deployment - + Args: data: Request data to update file_id: File ID (can be managed/encoded or regular) request: FastAPI request object llm_router: LiteLLM router for credential lookup (required for managed files) - + Returns: Tuple of (updated request data, original_managed_file_id) - original_managed_file_id is the original file_id if it was managed/encoded, None otherwise @@ -65,19 +65,17 @@ def _update_request_data_with_managed_file_id( # First, check if this is a unified managed file ID (base64 encoded) decoded_id = is_base64_encoded_unified_id(file_id) - + if decoded_id: # This is a unified managed file ID - verbose_logger.debug( - f"Processing unified managed file ID: {file_id}" - ) - + verbose_logger.debug(f"Processing unified managed file ID: {file_id}") + # Parse the unified ID to extract components parsed_id = parse_unified_id(file_id) - + if parsed_id: target_model_names = parsed_id.get("target_model_names", []) - + # Extract the actual provider file ID from llm_output_file_id field # Format: litellm_proxy:...;llm_output_file_id,{actual_file_id};... llm_output_file_id = None @@ -87,16 +85,16 @@ def _update_request_data_with_managed_file_id( llm_output_file_id = match.group(1).strip() except Exception: pass - + verbose_logger.debug( f"Decoded unified file ID - target_model_names: {target_model_names}, llm_output_file_id: {llm_output_file_id}" ) - + # Set the model for routing if target_model_names and len(target_model_names) > 0: routing_model = target_model_names[0] data["model"] = routing_model - + # Get credentials for the model if llm_router: credentials = llm_router.get_deployment_credentials_with_provider( @@ -112,7 +110,7 @@ def _update_request_data_with_managed_file_id( f"Routing vector store file operation to model: {routing_model}, file_id: {file_id} -> {llm_output_file_id}" ) return data, file_id # Return original managed file ID - + # If we extracted the provider file ID but no routing, still use it if llm_output_file_id: data["file_id"] = llm_output_file_id @@ -120,18 +118,23 @@ def _update_request_data_with_managed_file_id( f"Replaced unified file ID with provider file ID: {llm_output_file_id}" ) return data, file_id # Return original managed file ID - + return data, file_id if decoded_id else None - + # Fall back to simple encoded file ID handling (format: litellm:{file_id};model,{model}) - should_route, model_used, original_file_id, credentials = handle_model_based_routing( + ( + should_route, + model_used, + original_file_id, + credentials, + ) = handle_model_based_routing( file_id=file_id, request=request, llm_router=llm_router, data=data, check_file_id_encoding=True, ) - + if should_route: # Use model-based routing with credentials from config prepare_data_with_credentials( @@ -139,33 +142,37 @@ def _update_request_data_with_managed_file_id( credentials=credentials, # type: ignore file_id=original_file_id, # Use decoded file ID if from encoded ID ) - + verbose_logger.debug( f"Routing vector store file operation using model: {model_used}" - + (f", file_id: {file_id} -> {original_file_id}" if original_file_id else "") + + ( + f", file_id: {file_id} -> {original_file_id}" + if original_file_id + else "" + ) ) return data, file_id # Return original file ID for response replacement - + return data, None def _replace_file_id_in_response(response, original_file_id: str): """ Replace the provider file ID in the response with the original managed file ID. - + This ensures that when a user sends a managed file ID, they get back the same managed file ID in the response, not the decoded provider file ID. - + Args: response: The response object from the provider original_file_id: The original managed file ID to restore - + Returns: Modified response with original file ID """ if response is None: return response - + # Handle different response types if isinstance(response, dict): # For dict responses (e.g., VectorStoreFileDeleteResponse) @@ -178,7 +185,7 @@ def _replace_file_id_in_response(response, original_file_id: str): response.id = original_file_id elif hasattr(response, "file_id"): response.file_id = original_file_id - + return response @@ -189,22 +196,22 @@ def _update_request_data_with_litellm_managed_vector_store_registry( ) -> Dict: """ Update request data with model routing information from managed vector store. - + This function handles two types of vector stores: 1. Legacy vector stores from registry (non-managed) 2. Managed vector stores with unified IDs (requires decoding) - + For managed vector stores, this function: - Decodes the unified vector store ID - Extracts the model_id and provider resource ID - Sets data["model"] so the router can use the correct deployment credentials - Replaces the unified ID with the provider-specific ID - + Args: data: Request data to update vector_store_id: Vector store ID (can be unified or legacy) llm_router: LiteLLM router for credential lookup (required for managed vector stores) - + Returns: Updated request data with model routing information """ @@ -216,24 +223,22 @@ def _update_request_data_with_litellm_managed_vector_store_registry( # Check if this is a managed vector store ID (base64 encoded unified ID) decoded_id = is_base64_encoded_unified_id(vector_store_id) - + if decoded_id: # This is a managed vector store - decode and extract routing information - verbose_logger.debug( - f"Processing managed vector store ID: {vector_store_id}" - ) - + verbose_logger.debug(f"Processing managed vector store ID: {vector_store_id}") + parsed_id = parse_unified_id(vector_store_id) - + if parsed_id: model_id = parsed_id.get("model_id") provider_resource_id = parsed_id.get("provider_resource_id") target_model_names = parsed_id.get("target_model_names", []) - + verbose_logger.debug( f"Decoded vector store - model_id: {model_id}, provider_resource_id: {provider_resource_id}, target_model_names: {target_model_names}" ) - + # Set the model for routing - this tells the router which deployment to use # The router will automatically get the credentials from the deployment routing_model = None @@ -241,28 +246,26 @@ def _update_request_data_with_litellm_managed_vector_store_registry( routing_model = model_id elif target_model_names and len(target_model_names) > 0: routing_model = target_model_names[0] - + if routing_model: data["model"] = routing_model verbose_logger.info( f"Routing vector store files operation to model: {routing_model}" ) - + # Replace unified vector store ID with provider resource ID if provider_resource_id: data["vector_store_id"] = provider_resource_id verbose_logger.debug( f"Replaced unified vector store ID with provider resource ID: {provider_resource_id}" ) - + return data - + # Legacy path: Check vector store registry for non-managed vector stores if litellm.vector_store_registry is not None: - vector_store_to_run = ( - litellm.vector_store_registry.get_litellm_managed_vector_store_from_registry( - vector_store_id=vector_store_id - ) + vector_store_to_run = litellm.vector_store_registry.get_litellm_managed_vector_store_from_registry( + vector_store_id=vector_store_id ) if vector_store_to_run is not None: if "custom_llm_provider" in vector_store_to_run: @@ -276,7 +279,7 @@ def _update_request_data_with_litellm_managed_vector_store_registry( if "litellm_params" in vector_store_to_run: litellm_params = vector_store_to_run.get("litellm_params", {}) or {} data.update(litellm_params) - + return data @@ -406,11 +409,11 @@ async def vector_store_file_create( user_api_base=user_api_base, version=version, ) - + # Replace provider file ID with original managed file ID in response if original_managed_file_id: response = _replace_file_id_in_response(response, original_managed_file_id) - + return response except Exception as e: # noqa: BLE001 raise await processor._handle_llm_api_exception( @@ -580,11 +583,11 @@ async def vector_store_file_retrieve( user_api_base=user_api_base, version=version, ) - + # Replace provider file ID with original managed file ID in response if original_managed_file_id: response = _replace_file_id_in_response(response, original_managed_file_id) - + return response except Exception as e: # noqa: BLE001 raise await processor._handle_llm_api_exception( @@ -674,11 +677,11 @@ async def vector_store_file_content( user_api_base=user_api_base, version=version, ) - + # Replace provider file ID with original managed file ID in response if original_managed_file_id: response = _replace_file_id_in_response(response, original_managed_file_id) - + return response except Exception as e: # noqa: BLE001 raise await processor._handle_llm_api_exception( @@ -768,11 +771,11 @@ async def vector_store_file_update( user_api_base=user_api_base, version=version, ) - + # Replace provider file ID with original managed file ID in response if original_managed_file_id: response = _replace_file_id_in_response(response, original_managed_file_id) - + return response except Exception as e: # noqa: BLE001 raise await processor._handle_llm_api_exception( @@ -862,11 +865,11 @@ async def vector_store_file_delete( user_api_base=user_api_base, version=version, ) - + # Replace provider file ID with original managed file ID in response if original_managed_file_id: response = _replace_file_id_in_response(response, original_managed_file_id) - + return response except Exception as e: # noqa: BLE001 raise await processor._handle_llm_api_exception( diff --git a/litellm/proxy/video_endpoints/endpoints.py b/litellm/proxy/video_endpoints/endpoints.py index 5e00eb5845..f7a71c1033 100644 --- a/litellm/proxy/video_endpoints/endpoints.py +++ b/litellm/proxy/video_endpoints/endpoints.py @@ -256,7 +256,9 @@ async def video_status( # Resolve model_name from model_id if available # This allows the router to automatically inject litellm_params from the model config if model_id_from_decoded and llm_router: - resolved_model = llm_router.resolve_model_name_from_model_id(model_id_from_decoded) + resolved_model = llm_router.resolve_model_name_from_model_id( + model_id_from_decoded + ) if resolved_model: data["model"] = resolved_model @@ -341,7 +343,7 @@ async def video_content( decoded = decode_video_id_with_provider(video_id) provider_from_id = decoded.get("custom_llm_provider") model_id_from_decoded = decoded.get("model_id") - + custom_llm_provider = ( get_custom_llm_provider_from_request_headers(request=request) or get_custom_llm_provider_from_request_query(request=request) @@ -354,7 +356,9 @@ async def video_content( # Resolve model_name from model_id if available # This allows the router to automatically inject litellm_params from the model config if model_id_from_decoded and llm_router: - resolved_model = llm_router.resolve_model_name_from_model_id(model_id_from_decoded) + resolved_model = llm_router.resolve_model_name_from_model_id( + model_id_from_decoded + ) if resolved_model: data["model"] = resolved_model # Process request using ProxyBaseLLMRequestProcessing @@ -379,14 +383,14 @@ async def video_content( user_api_base=user_api_base, version=version, ) - + # Return raw video bytes with proper content type return Response( content=video_bytes, media_type="video/mp4", headers={ "Content-Disposition": f"attachment; filename=video_{video_id}.mp4" - } + }, ) except Exception as e: raise await processor._handle_llm_api_exception( @@ -466,7 +470,9 @@ async def video_remix( # Resolve model_name from model_id if available # This allows the router to automatically inject litellm_params from the model config if model_id_from_decoded and llm_router: - resolved_model = llm_router.resolve_model_name_from_model_id(model_id_from_decoded) + resolved_model = llm_router.resolve_model_name_from_model_id( + model_id_from_decoded + ) if resolved_model: data["model"] = resolved_model diff --git a/litellm/rag/__init__.py b/litellm/rag/__init__.py index 54f4d3ccaa..e387cf837e 100644 --- a/litellm/rag/__init__.py +++ b/litellm/rag/__init__.py @@ -19,4 +19,3 @@ async def arag_ingest(*args, **kwargs): def rag_ingest(*args, **kwargs): """Alias for ingest.""" return ingest(*args, **kwargs) - diff --git a/litellm/rag/ingestion/__init__.py b/litellm/rag/ingestion/__init__.py index 264bd6b5e4..3be3fd5c1d 100644 --- a/litellm/rag/ingestion/__init__.py +++ b/litellm/rag/ingestion/__init__.py @@ -17,4 +17,3 @@ __all__ = [ "S3VectorsRAGIngestion", "VertexAIRAGIngestion", ] - diff --git a/litellm/rag/ingestion/base_ingestion.py b/litellm/rag/ingestion/base_ingestion.py index 3daa767188..0d12bdfffc 100644 --- a/litellm/rag/ingestion/base_ingestion.py +++ b/litellm/rag/ingestion/base_ingestion.py @@ -76,7 +76,9 @@ class BaseRAGIngestion(ABC): credential_name = self.vector_store_config.get("litellm_credential_name") if credential_name and litellm.credential_list: - credential_values = CredentialAccessor.get_credential_values(credential_name) + credential_values = CredentialAccessor.get_credential_values( + credential_name + ) # Merge credentials into vector_store_config (don't overwrite existing values) for key, value in credential_values.items(): if key not in self.vector_store_config: @@ -114,7 +116,9 @@ class BaseRAGIngestion(ABC): response.raise_for_status() file_content = response.content filename = file_url.split("/")[-1] or "document" - content_type = response.headers.get("content-type", "application/octet-stream") + content_type = response.headers.get( + "content-type", "application/octet-stream" + ) return filename, file_content, content_type, None if file_id: @@ -352,4 +356,3 @@ class BaseRAGIngestion(ABC): file_id=None, error=str(e), ) - diff --git a/litellm/rag/ingestion/bedrock_ingestion.py b/litellm/rag/ingestion/bedrock_ingestion.py index 3c880b8849..6cf41c82f1 100644 --- a/litellm/rag/ingestion/bedrock_ingestion.py +++ b/litellm/rag/ingestion/bedrock_ingestion.py @@ -40,14 +40,14 @@ def _get_int(value: Any, default: int) -> int: def _normalize_principal_arn(caller_arn: str, account_id: str) -> str: """ Normalize a caller ARN to the format required by OpenSearch data access policies. - + OpenSearch Serverless data access policies require: - IAM users: arn:aws:iam::account-id:user/user-name - IAM roles: arn:aws:iam::account-id:role/role-name - + But get_caller_identity() returns for assumed roles: - arn:aws:sts::account-id:assumed-role/role-name/session-name - + This function converts assumed-role ARNs to the proper IAM role ARN format. """ if ":assumed-role/" in caller_arn: @@ -99,13 +99,22 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): # Optional config self._data_source_id = self.vector_store_config.get("data_source_id") self._s3_bucket = self.vector_store_config.get("s3_bucket") - self._s3_prefix: Optional[str] = str(self.vector_store_config.get("s3_prefix")) if self.vector_store_config.get("s3_prefix") else None - self.embedding_model = self.vector_store_config.get( - "embedding_model" - ) or "amazon.titan-embed-text-v2:0" + self._s3_prefix: Optional[str] = ( + str(self.vector_store_config.get("s3_prefix")) + if self.vector_store_config.get("s3_prefix") + else None + ) + self.embedding_model = ( + self.vector_store_config.get("embedding_model") + or "amazon.titan-embed-text-v2:0" + ) - self.wait_for_ingestion = self.vector_store_config.get("wait_for_ingestion", False) - self.ingestion_timeout: int = _get_int(self.vector_store_config.get("ingestion_timeout"), 300) + self.wait_for_ingestion = self.vector_store_config.get( + "wait_for_ingestion", False + ) + self.ingestion_timeout: int = _get_int( + self.vector_store_config.get("ingestion_timeout"), 300 + ) # Get AWS region using BaseAWSLLM method _aws_region = self.vector_store_config.get("aws_region_name") @@ -219,7 +228,9 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): await self._create_opensearch_index(collection_name) # Step 4: Create IAM role for Bedrock - role_arn = await self._create_bedrock_role(unique_id, account_id, collection_arn) + role_arn = await self._create_bedrock_role( + unique_id, account_id, collection_arn + ) # Step 5: Create Knowledge Base self.knowledge_base_id = await self._create_knowledge_base( @@ -260,49 +271,84 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): oss = self._get_boto3_client("opensearchserverless") collection_name = f"litellm-kb-{unique_id}" - verbose_logger.debug(f"Creating OpenSearch Serverless collection: {collection_name}") + verbose_logger.debug( + f"Creating OpenSearch Serverless collection: {collection_name}" + ) # Create encryption policy oss.create_security_policy( name=f"{collection_name}-enc", type="encryption", - policy=json.dumps({ - "Rules": [{"ResourceType": "collection", "Resource": [f"collection/{collection_name}"]}], - "AWSOwnedKey": True, - }), + policy=json.dumps( + { + "Rules": [ + { + "ResourceType": "collection", + "Resource": [f"collection/{collection_name}"], + } + ], + "AWSOwnedKey": True, + } + ), ) # Create network policy (public access for simplicity) oss.create_security_policy( name=f"{collection_name}-net", type="network", - policy=json.dumps([{ - "Rules": [{"ResourceType": "collection", "Resource": [f"collection/{collection_name}"]}, - {"ResourceType": "dashboard", "Resource": [f"collection/{collection_name}"]}], - "AllowFromPublic": True, - }]), + policy=json.dumps( + [ + { + "Rules": [ + { + "ResourceType": "collection", + "Resource": [f"collection/{collection_name}"], + }, + { + "ResourceType": "dashboard", + "Resource": [f"collection/{collection_name}"], + }, + ], + "AllowFromPublic": True, + } + ] + ), ) # Create data access policy - include both root and actual caller ARN # This ensures the credentials being used have access to the collection # Normalize the caller ARN (convert assumed-role ARN to IAM role ARN if needed) normalized_caller_arn = _normalize_principal_arn(caller_arn, account_id) - verbose_logger.debug(f"Caller ARN: {caller_arn}, Normalized: {normalized_caller_arn}") - + verbose_logger.debug( + f"Caller ARN: {caller_arn}, Normalized: {normalized_caller_arn}" + ) + principals = [f"arn:aws:iam::{account_id}:root", normalized_caller_arn] # Deduplicate in case caller is root principals = list(set(principals)) - + oss.create_access_policy( name=f"{collection_name}-access", type="data", - policy=json.dumps([{ - "Rules": [ - {"ResourceType": "index", "Resource": [f"index/{collection_name}/*"], "Permission": ["aoss:*"]}, - {"ResourceType": "collection", "Resource": [f"collection/{collection_name}"], "Permission": ["aoss:*"]}, - ], - "Principal": principals, - }]), + policy=json.dumps( + [ + { + "Rules": [ + { + "ResourceType": "index", + "Resource": [f"index/{collection_name}/*"], + "Permission": ["aoss:*"], + }, + { + "ResourceType": "collection", + "Resource": [f"collection/{collection_name}"], + "Permission": ["aoss:*"], + }, + ], + "Principal": principals, + } + ] + ), ) # Create collection @@ -341,9 +387,15 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): # Get credentials for signing credentials = self.get_credentials( - aws_access_key_id=_get_str_or_none(self.vector_store_config.get("aws_access_key_id")), - aws_secret_access_key=_get_str_or_none(self.vector_store_config.get("aws_secret_access_key")), - aws_session_token=_get_str_or_none(self.vector_store_config.get("aws_session_token")), + aws_access_key_id=_get_str_or_none( + self.vector_store_config.get("aws_access_key_id") + ), + aws_secret_access_key=_get_str_or_none( + self.vector_store_config.get("aws_secret_access_key") + ), + aws_session_token=_get_str_or_none( + self.vector_store_config.get("aws_session_token") + ), aws_region_name=self.aws_region_name, ) @@ -371,15 +423,17 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): index_name = "bedrock-kb-index" index_body = { - "settings": { - "index": {"knn": True, "knn.algo_param.ef_search": 512} - }, + "settings": {"index": {"knn": True, "knn.algo_param.ef_search": 512}}, "mappings": { "properties": { "bedrock-knowledge-base-default-vector": { "type": "knn_vector", "dimension": 1024, - "method": {"engine": "faiss", "name": "hnsw", "space_type": "l2"}, + "method": { + "engine": "faiss", + "name": "hnsw", + "space_type": "l2", + }, }, "AMAZON_BEDROCK_METADATA": {"type": "text", "index": False}, "AMAZON_BEDROCK_TEXT_CHUNK": {"type": "text"}, @@ -391,7 +445,7 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): max_retries = 8 retry_delay = 20 # seconds last_error = None - + for attempt in range(max_retries): try: client.indices.create(index=index_name, body=index_body) @@ -400,7 +454,10 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): except Exception as e: last_error = e error_str = str(e) - if "authorization_exception" in error_str.lower() or "security_exception" in error_str.lower(): + if ( + "authorization_exception" in error_str.lower() + or "security_exception" in error_str.lower() + ): verbose_logger.warning( f"OpenSearch index creation attempt {attempt + 1}/{max_retries} failed due to authorization. " f"Waiting {retry_delay}s for policy propagation..." @@ -409,7 +466,7 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): else: # Non-auth error, raise immediately raise - + # All retries exhausted raise RuntimeError( f"Failed to create OpenSearch index after {max_retries} attempts. " @@ -427,15 +484,19 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): trust_policy = { "Version": "2012-10-17", - "Statement": [{ - "Effect": "Allow", - "Principal": {"Service": "bedrock.amazonaws.com"}, - "Action": "sts:AssumeRole", - "Condition": { - "StringEquals": {"aws:SourceAccount": account_id}, - "ArnLike": {"aws:SourceArn": f"arn:aws:bedrock:{self.aws_region_name}:{account_id}:knowledge-base/*"}, - }, - }], + "Statement": [ + { + "Effect": "Allow", + "Principal": {"Service": "bedrock.amazonaws.com"}, + "Action": "sts:AssumeRole", + "Condition": { + "StringEquals": {"aws:SourceAccount": account_id}, + "ArnLike": { + "aws:SourceArn": f"arn:aws:bedrock:{self.aws_region_name}:{account_id}:knowledge-base/*" + }, + }, + } + ], } response = iam.create_role( @@ -452,7 +513,9 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): { "Effect": "Allow", "Action": ["bedrock:InvokeModel"], - "Resource": [f"arn:aws:bedrock:{self.aws_region_name}::foundation-model/{self.embedding_model}"], + "Resource": [ + f"arn:aws:bedrock:{self.aws_region_name}::foundation-model/{self.embedding_model}" + ], }, { "Effect": "Allow", @@ -462,7 +525,10 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): { "Effect": "Allow", "Action": ["s3:GetObject", "s3:ListBucket"], - "Resource": [f"arn:aws:s3:::{self.s3_bucket}", f"arn:aws:s3:::{self.s3_bucket}/*"], + "Resource": [ + f"arn:aws:s3:::{self.s3_bucket}", + f"arn:aws:s3:::{self.s3_bucket}/*", + ], }, ], } @@ -554,20 +620,40 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): try: import boto3 except ImportError: - raise ImportError("boto3 is required for Bedrock ingestion. Install with: pip install boto3") + raise ImportError( + "boto3 is required for Bedrock ingestion. Install with: pip install boto3" + ) # Get credentials using BaseAWSLLM's get_credentials method credentials = self.get_credentials( - aws_access_key_id=_get_str_or_none(self.vector_store_config.get("aws_access_key_id")), - aws_secret_access_key=_get_str_or_none(self.vector_store_config.get("aws_secret_access_key")), - aws_session_token=_get_str_or_none(self.vector_store_config.get("aws_session_token")), + aws_access_key_id=_get_str_or_none( + self.vector_store_config.get("aws_access_key_id") + ), + aws_secret_access_key=_get_str_or_none( + self.vector_store_config.get("aws_secret_access_key") + ), + aws_session_token=_get_str_or_none( + self.vector_store_config.get("aws_session_token") + ), aws_region_name=self.aws_region_name, - aws_session_name=_get_str_or_none(self.vector_store_config.get("aws_session_name")), - aws_profile_name=_get_str_or_none(self.vector_store_config.get("aws_profile_name")), - aws_role_name=_get_str_or_none(self.vector_store_config.get("aws_role_name")), - aws_web_identity_token=_get_str_or_none(self.vector_store_config.get("aws_web_identity_token")), - aws_sts_endpoint=_get_str_or_none(self.vector_store_config.get("aws_sts_endpoint")), - aws_external_id=_get_str_or_none(self.vector_store_config.get("aws_external_id")), + aws_session_name=_get_str_or_none( + self.vector_store_config.get("aws_session_name") + ), + aws_profile_name=_get_str_or_none( + self.vector_store_config.get("aws_profile_name") + ), + aws_role_name=_get_str_or_none( + self.vector_store_config.get("aws_role_name") + ), + aws_web_identity_token=_get_str_or_none( + self.vector_store_config.get("aws_web_identity_token") + ), + aws_sts_endpoint=_get_str_or_none( + self.vector_store_config.get("aws_sts_endpoint") + ), + aws_external_id=_get_str_or_none( + self.vector_store_config.get("aws_external_id") + ), ) # Create session with credentials @@ -623,7 +709,9 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): await self._ensure_config_initialized() if not file_content or not filename: - verbose_logger.warning("No file content or filename provided for Bedrock ingestion") + verbose_logger.warning( + "No file content or filename provided for Bedrock ingestion" + ) return _get_str_or_none(self.knowledge_base_id), None # Step 1: Upload file to S3 @@ -655,6 +743,7 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): # Step 3: Wait for ingestion (optional) - use asyncio.sleep to avoid blocking if self.wait_for_ingestion: import time as time_module + start_time = time_module.time() while time_module.time() - start_time < self.ingestion_timeout: job_status = bedrock_agent.get_ingestion_job( @@ -672,7 +761,9 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): ) break elif status == "FAILED": - failure_reasons = job_status["ingestionJob"].get("failureReasons", []) + failure_reasons = job_status["ingestionJob"].get( + "failureReasons", [] + ) verbose_logger.error(f"Ingestion failed: {failure_reasons}") break elif status in ("STARTING", "IN_PROGRESS"): @@ -682,4 +773,3 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): break return str(self.knowledge_base_id) if self.knowledge_base_id else None, s3_key - diff --git a/litellm/rag/ingestion/file_parsers/pdf_parser.py b/litellm/rag/ingestion/file_parsers/pdf_parser.py index 9a533ccf13..cb42dfd5d8 100644 --- a/litellm/rag/ingestion/file_parsers/pdf_parser.py +++ b/litellm/rag/ingestion/file_parsers/pdf_parser.py @@ -25,46 +25,52 @@ def extract_text_from_pdf(file_content: bytes) -> Optional[str]: # Try pypdf first (most common) try: from pypdf import PdfReader as PypdfReader - + pdf_file = BytesIO(file_content) reader = PypdfReader(pdf_file) - + text_parts = [] for page in reader.pages: text = page.extract_text() if text: text_parts.append(text) - + if text_parts: extracted_text = "\n\n".join(text_parts) - verbose_logger.debug(f"Extracted {len(extracted_text)} characters from PDF using pypdf") + verbose_logger.debug( + f"Extracted {len(extracted_text)} characters from PDF using pypdf" + ) return extracted_text - + except ImportError: verbose_logger.debug("pypdf not available, trying PyPDF2") - + # Fallback to PyPDF2 try: from PyPDF2 import PdfReader as PyPDF2Reader - + pdf_file = BytesIO(file_content) reader = PyPDF2Reader(pdf_file) - + text_parts = [] for page in reader.pages: text = page.extract_text() if text: text_parts.append(text) - + if text_parts: extracted_text = "\n\n".join(text_parts) - verbose_logger.debug(f"Extracted {len(extracted_text)} characters from PDF using PyPDF2") + verbose_logger.debug( + f"Extracted {len(extracted_text)} characters from PDF using PyPDF2" + ) return extracted_text - + except ImportError: - verbose_logger.debug("PyPDF2 not available, PDF extraction requires OCR or pypdf/PyPDF2 library") - + verbose_logger.debug( + "PyPDF2 not available, PDF extraction requires OCR or pypdf/PyPDF2 library" + ) + except Exception as e: verbose_logger.debug(f"PDF text extraction failed: {e}") - + return None diff --git a/litellm/rag/ingestion/gemini_ingestion.py b/litellm/rag/ingestion/gemini_ingestion.py index 1e74bcf9c3..96495b9f3f 100644 --- a/litellm/rag/ingestion/gemini_ingestion.py +++ b/litellm/rag/ingestion/gemini_ingestion.py @@ -80,16 +80,24 @@ class GeminiRAGIngestion(BaseRAGIngestion): Tuple of (vector_store_id, file_id) """ vector_store_id = self.vector_store_config.get("vector_store_id") - + vector_store_config = cast(Dict[str, Any], self.vector_store_config) # Get API credentials - api_key = cast(Optional[str], vector_store_config.get("api_key")) or GeminiModelInfo.get_api_key() - api_base = cast(Optional[str], vector_store_config.get("api_base")) or GeminiModelInfo.get_api_base() - + api_key = ( + cast(Optional[str], vector_store_config.get("api_key")) + or GeminiModelInfo.get_api_key() + ) + api_base = ( + cast(Optional[str], vector_store_config.get("api_base")) + or GeminiModelInfo.get_api_base() + ) + if not api_key: - raise ValueError("GEMINI_API_KEY or GOOGLE_API_KEY is required for Gemini File Search") - + raise ValueError( + "GEMINI_API_KEY or GOOGLE_API_KEY is required for Gemini File Search" + ) + if not api_base: raise ValueError("GEMINI_API_BASE is required") @@ -136,11 +144,9 @@ class GeminiRAGIngestion(BaseRAGIngestion): Store name (format: fileSearchStores/xxxxxxx) """ url = f"{base_url}/fileSearchStores?key={api_key}" - - request_body = { - "displayName": display_name - } - + + request_body = {"displayName": display_name} + client = get_async_httpx_client( llm_provider=httpxSpecialProvider.RAG, params={"timeout": 60.0}, @@ -150,15 +156,15 @@ class GeminiRAGIngestion(BaseRAGIngestion): json=request_body, headers={"Content-Type": "application/json"}, ) - + if response.status_code != 200: error_msg = f"Failed to create File Search store: {response.text}" verbose_logger.error(error_msg) raise Exception(error_msg) - + response_data = response.json() store_name = response_data.get("name", "") - + verbose_logger.debug(f"Created File Search store: {store_name}") return store_name @@ -223,11 +229,9 @@ class GeminiRAGIngestion(BaseRAGIngestion): # We need: https://generativelanguage.googleapis.com/upload/v1beta/{store_id}:uploadToFileSearchStore api_base = base_url.replace("/v1beta", "") # Get base without version url = f"{api_base}/upload/v1beta/{vector_store_id}:uploadToFileSearchStore?key={api_key}" - + # Build request body with chunking config and metadata if provided - request_body: Dict[str, Any] = { - "displayName": filename - } + request_body: Dict[str, Any] = {"displayName": filename} # Add chunking configuration if provided chunking_strategy = self.chunking_strategy @@ -236,13 +240,20 @@ class GeminiRAGIngestion(BaseRAGIngestion): if white_space_config: request_body["chunkingConfig"] = { "whiteSpaceConfig": { - "maxTokensPerChunk": white_space_config.get("max_tokens_per_chunk", 800), - "maxOverlapTokens": white_space_config.get("max_overlap_tokens", 400), + "maxTokensPerChunk": white_space_config.get( + "max_tokens_per_chunk", 800 + ), + "maxOverlapTokens": white_space_config.get( + "max_overlap_tokens", 400 + ), } } # Add custom metadata if provided in vector_store_config - custom_metadata = cast(Optional[List[Dict[str, Any]]], self.vector_store_config.get("custom_metadata")) + custom_metadata = cast( + Optional[List[Dict[str, Any]]], + self.vector_store_config.get("custom_metadata"), + ) if custom_metadata: request_body["customMetadata"] = custom_metadata @@ -317,11 +328,12 @@ class GeminiRAGIngestion(BaseRAGIngestion): try: response_data = response.json() # The response should contain the document name or file reference - file_id = response_data.get("name", "") or response_data.get("file", {}).get("name", "") + file_id = response_data.get("name", "") or response_data.get( + "file", {} + ).get("name", "") verbose_logger.debug(f"Upload complete. File ID: {file_id}") return file_id except Exception as e: verbose_logger.warning(f"Could not parse upload response: {e}") # Return a placeholder if we can't get the ID return "uploaded" - diff --git a/litellm/rag/ingestion/openai_ingestion.py b/litellm/rag/ingestion/openai_ingestion.py index 33fe8c06ec..891e3d0e91 100644 --- a/litellm/rag/ingestion/openai_ingestion.py +++ b/litellm/rag/ingestion/openai_ingestion.py @@ -84,7 +84,9 @@ class OpenAIRAGIngestion(BaseRAGIngestion): # Create vector store if not provided if not vector_store_id: - expires_after = {"anchor": "last_active_at", "days": ttl_days} if ttl_days else None + expires_after = ( + {"anchor": "last_active_at", "days": ttl_days} if ttl_days else None + ) create_response = await vector_store_acreate( name=self.ingest_name or "litellm-rag-ingest", custom_llm_provider="openai", @@ -99,7 +101,11 @@ class OpenAIRAGIngestion(BaseRAGIngestion): if file_content and filename and vector_store_id: # Upload file to OpenAI file_response = await litellm.acreate_file( - file=(filename, file_content, content_type or "application/octet-stream"), + file=( + filename, + file_content, + content_type or "application/octet-stream", + ), purpose="assistants", custom_llm_provider="openai", api_key=api_key, @@ -112,10 +118,11 @@ class OpenAIRAGIngestion(BaseRAGIngestion): vector_store_id=vector_store_id, file_id=result_file_id, custom_llm_provider="openai", - chunking_strategy=cast(Optional[Dict[str, Any]], self.chunking_strategy), + chunking_strategy=cast( + Optional[Dict[str, Any]], self.chunking_strategy + ), api_key=api_key, api_base=api_base, ) return vector_store_id, result_file_id - diff --git a/litellm/rag/ingestion/s3_vectors_ingestion.py b/litellm/rag/ingestion/s3_vectors_ingestion.py index e6c166a101..2845a6737b 100644 --- a/litellm/rag/ingestion/s3_vectors_ingestion.py +++ b/litellm/rag/ingestion/s3_vectors_ingestion.py @@ -75,7 +75,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): "non_filterable_metadata_keys", S3_VECTORS_DEFAULT_NON_FILTERABLE_METADATA_KEYS, ) - + # Get dimension from config (will be auto-detected on first use if not provided) self.dimension = self._get_dimension_from_config() @@ -100,26 +100,30 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): async def _get_dimension_from_embedding_request(self) -> int: """ Auto-detect dimension by making a test embedding request. - + Makes a single embedding request with a test string to determine the output dimension of the embedding model. """ if not self.embedding_config or "model" not in self.embedding_config: return S3_VECTORS_DEFAULT_DIMENSION - + try: model_name = self.embedding_config["model"] verbose_logger.debug( f"Auto-detecting dimension by making test embedding request to {model_name}" ) - + # Make a test embedding request test_input = "test" if self.router: - response = await self.router.aembedding(model=model_name, input=[test_input]) + response = await self.router.aembedding( + model=model_name, input=[test_input] + ) else: - response = await litellm.aembedding(model=model_name, input=[test_input]) - + response = await litellm.aembedding( + model=model_name, input=[test_input] + ) + # Get dimension from the response if response.data and len(response.data) > 0: dimension = len(response.data[0]["embedding"]) @@ -132,13 +136,13 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): f"Could not auto-detect dimension from embedding model: {e}. " f"Using default dimension of {S3_VECTORS_DEFAULT_DIMENSION}." ) - + return S3_VECTORS_DEFAULT_DIMENSION - + def _get_dimension_from_config(self) -> Optional[int]: """ Get vector dimension from config if explicitly provided. - + Returns None if dimension should be auto-detected. """ if "dimension" in self.vector_store_config: @@ -197,7 +201,9 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): aws_session_name=self.vector_store_config.get("aws_session_name"), aws_profile_name=self.vector_store_config.get("aws_profile_name"), aws_role_name=self.vector_store_config.get("aws_role_name"), - aws_web_identity_token=self.vector_store_config.get("aws_web_identity_token"), + aws_web_identity_token=self.vector_store_config.get( + "aws_web_identity_token" + ), aws_sts_endpoint=self.vector_store_config.get("aws_sts_endpoint"), aws_external_id=self.vector_store_config.get("aws_external_id"), ) @@ -253,7 +259,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): verbose_logger.debug( f"Ensuring S3 vector bucket exists: {self.vector_bucket_name}" ) - + # Validate bucket name (AWS S3 naming rules) if len(self.vector_bucket_name) < 3: raise ValueError( @@ -273,7 +279,9 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): get_body = safe_dumps({"vectorBucketName": self.vector_bucket_name}) try: - response = await self._sign_and_execute_request("POST", get_url, data=get_body) + response = await self._sign_and_execute_request( + "POST", get_url, data=get_body + ) if response.status_code == 200: verbose_logger.debug(f"Vector bucket {self.vector_bucket_name} exists") return @@ -285,12 +293,14 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): # Create vector bucket using CreateVectorBucket API try: verbose_logger.debug(f"Creating vector bucket: {self.vector_bucket_name}") - create_url = f"https://s3vectors.{self.aws_region_name}.api.aws/CreateVectorBucket" - create_body = safe_dumps({ - "vectorBucketName": self.vector_bucket_name - }) - - response = await self._sign_and_execute_request("POST", create_url, data=create_body) + create_url = ( + f"https://s3vectors.{self.aws_region_name}.api.aws/CreateVectorBucket" + ) + create_body = safe_dumps({"vectorBucketName": self.vector_bucket_name}) + + response = await self._sign_and_execute_request( + "POST", create_url, data=create_body + ) if response.status_code in (200, 201): verbose_logger.info(f"Created vector bucket: {self.vector_bucket_name}") @@ -300,7 +310,9 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): f"Vector bucket {self.vector_bucket_name} already exists" ) else: - verbose_logger.error(f"CreateVectorBucket failed: {response.status_code} - {response.text}") + verbose_logger.error( + f"CreateVectorBucket failed: {response.status_code} - {response.text}" + ) response.raise_for_status() except Exception as e: verbose_logger.exception(f"Error creating vector bucket: {e}") @@ -314,13 +326,14 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): # Try to get index info using GetIndex API get_url = f"https://s3vectors.{self.aws_region_name}.api.aws/GetIndex" - get_body = safe_dumps({ - "vectorBucketName": self.vector_bucket_name, - "indexName": self.index_name - }) + get_body = safe_dumps( + {"vectorBucketName": self.vector_bucket_name, "indexName": self.index_name} + ) try: - response = await self._sign_and_execute_request("POST", get_url, data=get_body) + response = await self._sign_and_execute_request( + "POST", get_url, data=get_body + ) if response.status_code == 200: verbose_logger.debug(f"Vector index {self.index_name} exists") return @@ -359,7 +372,9 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): elif response.status_code == 409: verbose_logger.debug(f"Vector index {self.index_name} already exists") else: - verbose_logger.error(f"CreateIndex failed: {response.status_code} - {response.text}") + verbose_logger.error( + f"CreateIndex failed: {response.status_code} - {response.text}" + ) response.raise_for_status() except Exception as e: verbose_logger.exception(f"Error creating vector index: {e}") @@ -382,7 +397,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): request_body = { "vectorBucketName": self.vector_bucket_name, "indexName": self.index_name, - "vectors": vectors + "vectors": vectors, } try: @@ -430,11 +445,15 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): # Convert to list to ensure type compatibility input_chunks: List[str] = list(chunks) - + if self.router: - response = await self.router.aembedding(model=embedding_model, input=input_chunks) + response = await self.router.aembedding( + model=embedding_model, input=input_chunks + ) else: - response = await litellm.aembedding(model=embedding_model, input=input_chunks) + response = await litellm.aembedding( + model=embedding_model, input=input_chunks + ) return [item["embedding"] for item in response.data] @@ -488,10 +507,10 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): "source_text": chunk, # Non-filterable (for reference) "chunk_index": str(i), # Filterable } - + if filename: metadata["filename"] = filename # Filterable - + vector_obj = { "key": f"{filename}_{i}" if filename else f"chunk_{i}", "data": {"float32": embedding}, @@ -551,7 +570,9 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): if response.status_code == 200: results = response.json() - verbose_logger.debug(f"Query returned {len(results.get('vectors', []))} results") + verbose_logger.debug( + f"Query returned {len(results.get('vectors', []))} results" + ) # Check if query terms appear in results if results.get("vectors"): diff --git a/litellm/rag/ingestion/vertex_ai_ingestion.py b/litellm/rag/ingestion/vertex_ai_ingestion.py index 7394ec7a61..d95d2d56ce 100644 --- a/litellm/rag/ingestion/vertex_ai_ingestion.py +++ b/litellm/rag/ingestion/vertex_ai_ingestion.py @@ -48,7 +48,7 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): # Extract Vertex AI specific configs from vector_store_config litellm_params = dict(self.vector_store_config) - + # Get project, location, and credentials using VertexBase methods self.project_id = self.safe_get_vertex_ai_project(litellm_params) self.location = self.get_vertex_ai_location(litellm_params) or "us-central1" @@ -170,9 +170,7 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): if "vectorDbConfig" not in request_body: request_body["vectorDbConfig"] = {} request_body["vectorDbConfig"]["ragEmbeddingModelConfig"] = { - "vertexPredictionEndpoint": { - "endpoint": embedding_model - } + "vertexPredictionEndpoint": {"endpoint": embedding_model} } verbose_logger.debug(f"Creating RAG corpus: {url}") @@ -197,8 +195,10 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): raise Exception(error_msg) response_data = response.json() - verbose_logger.debug(f"Create corpus response: {json.dumps(response_data, indent=2)}") - + verbose_logger.debug( + f"Create corpus response: {json.dumps(response_data, indent=2)}" + ) + # The response is a long-running operation # Check if it's already done or if we need to poll if response_data.get("done"): @@ -264,21 +264,25 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): raise Exception(error_msg) operation_data = response.json() - + if operation_data.get("done"): # Check for errors if "error" in operation_data: error = operation_data["error"] raise Exception(f"Operation failed: {error}") - + # Extract corpus name from response corpus_name = operation_data.get("response", {}).get("name", "") if corpus_name: return corpus_name else: - raise Exception(f"No corpus name in operation response: {operation_data}") - - verbose_logger.debug(f"Operation not done yet, attempt {attempt + 1}/{max_retries}") + raise Exception( + f"No corpus name in operation response: {operation_data}" + ) + + verbose_logger.debug( + f"Operation not done yet, attempt {attempt + 1}/{max_retries}" + ) await asyncio.sleep(retry_delay) raise Exception(f"Operation timed out after {max_retries} attempts") @@ -311,10 +315,7 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): # Construct upload URL using vertex base URL helper base_url = get_vertex_base_url(self.location) - url = ( - f"{base_url}/upload/v1beta1/" - f"{rag_corpus_id}/ragFiles:upload" - ) + url = f"{base_url}/upload/v1beta1/" f"{rag_corpus_id}/ragFiles:upload" # Build metadata for the file with snake_case keys (as per upload API docs) metadata: Dict[str, Any] = { @@ -333,21 +334,19 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): if chunking_strategy and isinstance(chunking_strategy, dict): chunk_size = chunking_strategy.get("chunk_size") chunk_overlap = chunking_strategy.get("chunk_overlap") - + if chunk_size or chunk_overlap: if "upload_rag_file_config" not in metadata: metadata["upload_rag_file_config"] = {} - + metadata["upload_rag_file_config"]["rag_file_transformation_config"] = { - "rag_file_chunking_config": { - "fixed_length_chunking": {} - } + "rag_file_chunking_config": {"fixed_length_chunking": {}} } - + chunking_config = metadata["upload_rag_file_config"][ "rag_file_transformation_config" ]["rag_file_chunking_config"]["fixed_length_chunking"] - + if chunk_size: chunking_config["chunk_size"] = chunk_size if chunk_overlap: @@ -359,7 +358,11 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): # Prepare multipart form data files = { "metadata": (None, json.dumps(metadata), "application/json"), - "file": (filename, file_content, content_type or "application/octet-stream"), + "file": ( + filename, + file_content, + content_type or "application/octet-stream", + ), } client = get_async_httpx_client( llm_provider=httpxSpecialProvider.RAG, @@ -387,7 +390,7 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): file_id = response_data.get("ragFile", {}).get("name", "") if not file_id: file_id = response_data.get("name", "") - + verbose_logger.debug(f"Upload complete. File ID: {file_id}") return file_id except Exception as e: @@ -418,18 +421,11 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): # Construct import URL using vertex base URL helper base_url = get_vertex_base_url(self.location) - url = ( - f"{base_url}/v1beta1/" - f"{rag_corpus_id}/ragFiles:import" - ) + url = f"{base_url}/v1beta1/" f"{rag_corpus_id}/ragFiles:import" # Build request body with camelCase keys (Vertex AI API format) request_body: Dict[str, Any] = { - "importRagFilesConfig": { - "gcsSource": { - "uris": gcs_uris - } - } + "importRagFilesConfig": {"gcsSource": {"uris": gcs_uris}} } # Add chunking configuration if provided @@ -437,7 +433,7 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): if chunking_strategy and isinstance(chunking_strategy, dict): chunk_size = chunking_strategy.get("chunk_size") chunk_overlap = chunking_strategy.get("chunk_overlap") - + if chunk_size or chunk_overlap: request_body["importRagFilesConfig"]["ragFileChunkingConfig"] = { "chunkSize": chunk_size or 1024, @@ -445,9 +441,13 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): } # Add max embedding requests per minute if specified - max_embedding_qpm = self.vector_store_config.get("max_embedding_requests_per_min") + max_embedding_qpm = self.vector_store_config.get( + "max_embedding_requests_per_min" + ) if max_embedding_qpm: - request_body["importRagFilesConfig"]["maxEmbeddingRequestsPerMin"] = max_embedding_qpm + request_body["importRagFilesConfig"][ + "maxEmbeddingRequestsPerMin" + ] = max_embedding_qpm verbose_logger.debug(f"Importing files from GCS: {url}") verbose_logger.debug(f"Request body: {json.dumps(request_body, indent=2)}") @@ -473,6 +473,6 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): response_data = response.json() operation_name = response_data.get("name", "") - + verbose_logger.debug(f"Import operation started: {operation_name}") return operation_name diff --git a/litellm/rag/main.py b/litellm/rag/main.py index 6091d30025..e3d354b6c3 100644 --- a/litellm/rag/main.py +++ b/litellm/rag/main.py @@ -184,7 +184,9 @@ async def aingest( except Exception as e: raise litellm.exception_type( model=None, - custom_llm_provider=ingest_options.get("vector_store", {}).get("custom_llm_provider"), + custom_llm_provider=ingest_options.get("vector_store", {}).get( + "custom_llm_provider" + ), original_exception=e, completion_kwargs=local_vars, extra_kwargs=kwargs, @@ -430,7 +432,9 @@ def ingest( except Exception as e: raise litellm.exception_type( model=None, - custom_llm_provider=ingest_options.get("vector_store", {}).get("custom_llm_provider"), + custom_llm_provider=ingest_options.get("vector_store", {}).get( + "custom_llm_provider" + ), original_exception=e, completion_kwargs=local_vars, extra_kwargs=kwargs, diff --git a/litellm/rag/rag_query.py b/litellm/rag/rag_query.py index bf346efb3f..ebbf209e81 100644 --- a/litellm/rag/rag_query.py +++ b/litellm/rag/rag_query.py @@ -56,8 +56,8 @@ class RAGQuery: content_text: Optional[str] = content_item.get("text") if content_text: context_content += content_text + "\n\n" - elif "text" in chunk: # Fallback for simple dict with text - context_content += chunk["text"] + "\n\n" + elif "text" in chunk: # Fallback for simple dict with text + context_content += chunk["text"] + "\n\n" elif isinstance(chunk, str): context_content += chunk + "\n\n" @@ -107,7 +107,9 @@ class RAGQuery: return documents @staticmethod - def get_top_chunks_from_rerank(search_response: Any, rerank_response: Any) -> List[Any]: + def get_top_chunks_from_rerank( + search_response: Any, rerank_response: Any + ) -> List[Any]: """Get the original search results corresponding to the top reranked results.""" top_chunks = [] original_results = search_response.get("data", []) diff --git a/litellm/rag/text_splitters/__init__.py b/litellm/rag/text_splitters/__init__.py index 04802b438c..18e84afc17 100644 --- a/litellm/rag/text_splitters/__init__.py +++ b/litellm/rag/text_splitters/__init__.py @@ -7,4 +7,3 @@ from litellm.rag.text_splitters.recursive_character_text_splitter import ( ) __all__ = ["RecursiveCharacterTextSplitter"] - diff --git a/litellm/rag/text_splitters/recursive_character_text_splitter.py b/litellm/rag/text_splitters/recursive_character_text_splitter.py index 2107b1f068..edf6b84312 100644 --- a/litellm/rag/text_splitters/recursive_character_text_splitter.py +++ b/litellm/rag/text_splitters/recursive_character_text_splitter.py @@ -31,13 +31,18 @@ class RecursiveCharacterTextSplitter: """Split text into chunks.""" return self._split_text(text, self.separators) - def _split_text(self, text: str, separators: List[str], depth: int = 0) -> List[str]: + def _split_text( + self, text: str, separators: List[str], depth: int = 0 + ) -> List[str]: """Recursively split text using separators.""" from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH if depth > DEFAULT_MAX_RECURSE_DEPTH: # Max depth reached, return text as-is split into chunk_size pieces - return [text[i:i + self.chunk_size] for i in range(0, len(text), self.chunk_size)] + return [ + text[i : i + self.chunk_size] + for i in range(0, len(text), self.chunk_size) + ] final_chunks: List[str] = [] @@ -104,7 +109,9 @@ class RecursiveCharacterTextSplitter: chunks.append(chunk_text) # Handle overlap - while current_length > self.chunk_overlap and len(current_chunk) > 1: + while ( + current_length > self.chunk_overlap and len(current_chunk) > 1 + ): removed = current_chunk.pop(0) current_length -= len(removed) + len(separator) @@ -132,4 +139,3 @@ class RecursiveCharacterTextSplitter: start = end - self.chunk_overlap if end < len(text) else len(text) return chunks - diff --git a/litellm/rag/utils.py b/litellm/rag/utils.py index e8ab9c7517..49b8037de5 100644 --- a/litellm/rag/utils.py +++ b/litellm/rag/utils.py @@ -62,4 +62,3 @@ def get_rag_transformation_class(custom_llm_provider: str): # OpenAI and Bedrock don't need special transformations return None - diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 81f29ca6e3..7dcdc0d8d9 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -36,7 +36,11 @@ base_llm_http_handler = BaseLLMHTTPHandler() def _build_litellm_metadata(kwargs: dict) -> dict: """Build the litellm_metadata dict for guardrail checking (internal only, not forwarded to provider).""" metadata: dict = {**(kwargs.get("litellm_metadata") or {})} - guardrails = (kwargs.get("metadata") or {}).get("guardrails") or kwargs.get("guardrails") or [] + guardrails = ( + (kwargs.get("metadata") or {}).get("guardrails") + or kwargs.get("guardrails") + or [] + ) if guardrails: metadata["guardrails"] = guardrails return metadata @@ -74,11 +78,7 @@ def _get_realtime_http_provider_config( resolved_api_key = provider_config.get_api_key(api_key=raw_api_key) else: # Fallback for providers without a dedicated HTTP config (treated as OpenAI-compatible). - resolved_api_base = ( - raw_api_base - or litellm.api_base - or "https://api.openai.com" - ) + resolved_api_base = raw_api_base or litellm.api_base or "https://api.openai.com" resolved_api_key = ( raw_api_key or litellm.api_key @@ -111,12 +111,21 @@ async def acreate_realtime_client_secret( litellm_logging_obj: LiteLLMLogging = kwargs.get("litellm_logging_obj") # type: ignore litellm_params = GenericLiteLLMParams(**kwargs) - model_name, custom_llm_provider, dynamic_api_key, dynamic_api_base = get_llm_provider( + ( + model_name, + custom_llm_provider, + dynamic_api_key, + dynamic_api_base, + ) = get_llm_provider( model=model_name, api_base=litellm_params.api_base, api_key=litellm_params.api_key, ) - provider_config, resolved_api_base, resolved_api_key = _get_realtime_http_provider_config( + ( + provider_config, + resolved_api_base, + resolved_api_key, + ) = _get_realtime_http_provider_config( custom_llm_provider=custom_llm_provider, dynamic_api_base=dynamic_api_base, dynamic_api_key=dynamic_api_key, @@ -156,7 +165,12 @@ async def arealtime_calls( litellm_logging_obj: LiteLLMLogging = kwargs.get("litellm_logging_obj") # type: ignore litellm_params = GenericLiteLLMParams(**kwargs) - model_name, custom_llm_provider, dynamic_api_key, dynamic_api_base = get_llm_provider( + ( + model_name, + custom_llm_provider, + dynamic_api_key, + dynamic_api_base, + ) = get_llm_provider( model=model_name, api_base=litellm_params.api_base, api_key=litellm_params.api_key, @@ -271,12 +285,8 @@ async def _arealtime( # noqa: PLR0915 or get_secret_str("AZURE_API_KEY") ) - api_version = ( - api_version - or litellm_params.api_version - or "2024-10-01-preview" - ) - + api_version = api_version or litellm_params.api_version or "2024-10-01-preview" + realtime_protocol = ( kwargs.get("realtime_protocol") or litellm_params.get("realtime_protocol") @@ -365,11 +375,7 @@ async def _arealtime( # noqa: PLR0915 or "https://api.x.ai/v1" ) # set API KEY - api_key = ( - dynamic_api_key - or litellm.api_key - or get_secret_str("XAI_API_KEY") - ) + api_key = dynamic_api_key or litellm.api_key or get_secret_str("XAI_API_KEY") await xai_realtime.async_realtime( model=model, @@ -406,7 +412,10 @@ async def _arealtime( # noqa: PLR0915 vertex_region=vertex_location, model=model ) - access_token, resolved_project = await vertex_llm_base._ensure_access_token_async( + ( + access_token, + resolved_project, + ) = await vertex_llm_base._ensure_access_token_async( credentials=vertex_credentials, project_id=vertex_project, custom_llm_provider="vertex_ai", @@ -471,7 +480,8 @@ async def _realtime_health_check( ) elif custom_llm_provider == "openai": url = openai_realtime._construct_url( - api_base=api_base or "https://api.openai.com/", query_params={"model": model} + api_base=api_base or "https://api.openai.com/", + query_params={"model": model}, ) elif custom_llm_provider == "xai": url = xai_realtime._construct_url( @@ -482,7 +492,10 @@ async def _realtime_health_check( resolved_location = vertex_llm_base.get_vertex_region( vertex_region=vertex_location, model=model ) - access_token, resolved_project = await vertex_llm_base._ensure_access_token_async( + ( + access_token, + resolved_project, + ) = await vertex_llm_base._ensure_access_token_async( credentials=None, project_id=litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT"), custom_llm_provider="vertex_ai", diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index 871b18062f..e9766771dd 100644 --- a/litellm/rerank_api/main.py +++ b/litellm/rerank_api/main.py @@ -30,7 +30,11 @@ async def arerank( model: str, query: str, documents: List[Union[str, Dict[str, Any]]], - custom_llm_provider: Optional[Literal["cohere", "together_ai", "deepinfra", "fireworks_ai", "voyage", "watsonx"]] = None, + custom_llm_provider: Optional[ + Literal[ + "cohere", "together_ai", "deepinfra", "fireworks_ai", "voyage", "watsonx" + ] + ] = None, top_n: Optional[int] = None, rank_fields: Optional[List[str]] = None, return_documents: Optional[bool] = None, @@ -177,7 +181,10 @@ def rerank( # noqa: PLR0915 ) # Implement rerank logic here based on the custom_llm_provider - if _custom_llm_provider == litellm.LlmProviders.COHERE or _custom_llm_provider == litellm.LlmProviders.LITELLM_PROXY: + if ( + _custom_llm_provider == litellm.LlmProviders.COHERE + or _custom_llm_provider == litellm.LlmProviders.LITELLM_PROXY + ): # Implement Cohere rerank logic api_key: Optional[str] = ( dynamic_api_key or optional_params.api_key or litellm.api_key @@ -497,7 +504,9 @@ def rerank( # noqa: PLR0915 ) elif _custom_llm_provider == litellm.LlmProviders.WATSONX: credentials = IBMWatsonXMixin.get_watsonx_credentials( - optional_params=dict(optional_params), api_key=dynamic_api_key, api_base=dynamic_api_base + optional_params=dict(optional_params), + api_key=dynamic_api_key, + api_base=dynamic_api_base, ) api_key = credentials["api_key"] diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index 74ad6675e1..5faa8b587c 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -21,7 +21,6 @@ from litellm.types.utils import ModelResponse class LiteLLMCompletionTransformationHandler: - def response_api_handler( self, model: str, @@ -39,16 +38,14 @@ class LiteLLMCompletionTransformationHandler: Any, Any, Union[ResponsesAPIResponse, BaseResponsesAPIStreamingIterator] ], ]: - litellm_completion_request: dict = ( - LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( - model=model, - input=input, - responses_api_request=responses_api_request, - custom_llm_provider=custom_llm_provider, - stream=stream, - extra_headers=extra_headers, - **kwargs, - ) + litellm_completion_request: dict = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model=model, + input=input, + responses_api_request=responses_api_request, + custom_llm_provider=custom_llm_provider, + stream=stream, + extra_headers=extra_headers, + **kwargs, ) if _is_async: @@ -71,12 +68,10 @@ class LiteLLMCompletionTransformationHandler: ) if isinstance(litellm_completion_response, ModelResponse): - responses_api_response: ResponsesAPIResponse = ( - LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - chat_completion_response=litellm_completion_response, - request_input=input, - responses_api_request=responses_api_request, - ) + responses_api_response: ResponsesAPIResponse = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + chat_completion_response=litellm_completion_response, + request_input=input, + responses_api_request=responses_api_request, ) return responses_api_response @@ -90,7 +85,9 @@ class LiteLLMCompletionTransformationHandler: custom_llm_provider=custom_llm_provider, litellm_metadata=kwargs.get("litellm_metadata", {}), ) - raise ValueError(f"Unexpected response type: {type(litellm_completion_response)}") + raise ValueError( + f"Unexpected response type: {type(litellm_completion_response)}" + ) async def async_response_api_handler( self, @@ -99,7 +96,6 @@ class LiteLLMCompletionTransformationHandler: responses_api_request: ResponsesAPIOptionalRequestParams, **kwargs, ) -> Union[ResponsesAPIResponse, BaseResponsesAPIStreamingIterator]: - previous_response_id: Optional[str] = responses_api_request.get( "previous_response_id" ) @@ -120,12 +116,10 @@ class LiteLLMCompletionTransformationHandler: ) if isinstance(litellm_completion_response, ModelResponse): - responses_api_response: ResponsesAPIResponse = ( - LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - chat_completion_response=litellm_completion_response, - request_input=request_input, - responses_api_request=responses_api_request, - ) + responses_api_response: ResponsesAPIResponse = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + chat_completion_response=litellm_completion_response, + request_input=request_input, + responses_api_request=responses_api_request, ) return responses_api_response @@ -141,4 +135,6 @@ class LiteLLMCompletionTransformationHandler: ), litellm_metadata=kwargs.get("litellm_metadata", {}), ) - raise ValueError(f"Unexpected response type: {type(litellm_completion_response)}") + raise ValueError( + f"Unexpected response type: {type(litellm_completion_response)}" + ) diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 4f2c51edc5..45ab16b0d4 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -27,6 +27,7 @@ else: COLD_STORAGE_HANDLER = ColdStorageHandler() ######################################################## + class ResponsesSessionHandler: @staticmethod async def get_chat_completion_message_history_for_previous_response_id( @@ -78,7 +79,7 @@ class ResponsesSessionHandler: messages=chat_completion_message_history, litellm_session_id=litellm_session_id, ) - + @staticmethod async def extend_chat_completion_message_with_spend_log_payload( spend_log: SpendLogsPayload, @@ -90,7 +91,7 @@ class ResponsesSessionHandler: ChatCompletionResponseMessage, Message, ] - ] + ], ): """ Extend the chat completion message history with the spend log payload @@ -99,8 +100,10 @@ class ResponsesSessionHandler: LiteLLMCompletionResponsesConfig, ) - proxy_server_request_dict = await ResponsesSessionHandler.get_proxy_server_request_from_spend_log( - spend_log=spend_log, + proxy_server_request_dict = ( + await ResponsesSessionHandler.get_proxy_server_request_from_spend_log( + spend_log=spend_log, + ) ) response_input_param: Optional[Union[str, ResponseInputParam]] = None _messages: Optional[Union[str, ResponseInputParam]] = None @@ -114,9 +117,7 @@ class ResponsesSessionHandler: if isinstance(_response_input_param, str): response_input_param = _response_input_param elif isinstance(_response_input_param, dict): - response_input_param = cast( - ResponseInputParam, _response_input_param - ) + response_input_param = cast(ResponseInputParam, _response_input_param) if response_input_param: chat_completion_messages = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( @@ -141,16 +142,18 @@ class ResponsesSessionHandler: # Add Output messages for this Spend Log ############################################################ _response_output = spend_log.get("response", "{}") - if isinstance(_response_output, dict) and _response_output and _response_output != {}: + if ( + isinstance(_response_output, dict) + and _response_output + and _response_output != {} + ): # transform `ChatCompletion Response` to `ResponsesAPIResponse` model_response = ModelResponse(**_response_output) for choice in model_response.choices: if hasattr(choice, "message"): - chat_completion_message_history.append( - getattr(choice, "message") - ) + chat_completion_message_history.append(getattr(choice, "message")) return chat_completion_message_history - + @staticmethod async def get_proxy_server_request_from_spend_log( spend_log: SpendLogsPayload, @@ -166,15 +169,20 @@ class ResponsesSessionHandler: proxy_server_request_dict = proxy_server_request else: proxy_server_request_dict = json.loads(proxy_server_request) - ############################################################ # Check if user has setup cold storage for session handling ############################################################ - if ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_server_request_dict): + if ResponsesSessionHandler._should_check_cold_storage_for_full_payload( + proxy_server_request_dict + ): # Try to get cold storage object key from spend log metadata _proxy_server_request_dict: Optional[dict] = None - cold_storage_object_key = ResponsesSessionHandler._get_cold_storage_object_key_from_spend_log(spend_log) + cold_storage_object_key = ( + ResponsesSessionHandler._get_cold_storage_object_key_from_spend_log( + spend_log + ) + ) if cold_storage_object_key: # Use the object key directly from metadata _proxy_server_request_dict = await ResponsesSessionHandler.get_proxy_server_request_from_cold_storage_with_object_key( @@ -182,17 +190,19 @@ class ResponsesSessionHandler: ) if _proxy_server_request_dict: proxy_server_request_dict = _proxy_server_request_dict - + return proxy_server_request_dict - + @staticmethod - def _get_cold_storage_object_key_from_spend_log(spend_log: SpendLogsPayload) -> Optional[str]: + def _get_cold_storage_object_key_from_spend_log( + spend_log: SpendLogsPayload, + ) -> Optional[str]: """ Extract the cold storage object key from spend log metadata. - + Args: spend_log: The spend log payload containing metadata - + Returns: Optional[str]: The cold storage object key if found, None otherwise """ @@ -205,7 +215,9 @@ class ResponsesSessionHandler: return metadata_str.get("cold_storage_object_key") return None except (json.JSONDecodeError, TypeError, AttributeError): - verbose_proxy_logger.debug("Failed to parse metadata from spend log to extract cold storage object key") + verbose_proxy_logger.debug( + "Failed to parse metadata from spend log to extract cold storage object key" + ) return None @staticmethod @@ -214,14 +226,16 @@ class ResponsesSessionHandler: ) -> Optional[dict]: """ Get the proxy server request from cold storage using the object key directly. - + Args: object_key: The S3/GCS object key to retrieve - + Returns: Optional[dict]: The proxy server request dict or None if not found """ - verbose_proxy_logger.debug("inside get_proxy_server_request_from_cold_storage_with_object_key...") + verbose_proxy_logger.debug( + "inside get_proxy_server_request_from_cold_storage_with_object_key..." + ) proxy_server_request_dict = await COLD_STORAGE_HANDLER.get_proxy_server_request_from_cold_storage_with_object_key( object_key=object_key, @@ -234,11 +248,12 @@ class ResponsesSessionHandler: proxy_server_request_dict: Optional[dict], ) -> bool: """ - Only check cold storage when both are true + Only check cold storage when both are true 1. `LITELLM_TRUNCATED_PAYLOAD_FIELD` is in the proxy server request dict 2. `litellm.cold_storage_custom_logger` is not None """ from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD + configured_cold_storage_custom_logger = litellm.cold_storage_custom_logger if configured_cold_storage_custom_logger is None: return False @@ -250,8 +265,6 @@ class ResponsesSessionHandler: return True return False - - @staticmethod async def get_all_spend_logs_for_previous_response_id( previous_response_id: str, diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index e7866ae0f0..ce037850b8 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -92,9 +92,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._tool_args_by_call_id: dict[str, str] = {} self._tool_call_id_by_index: dict[int, str] = {} self._ambiguous_tool_call_indexes: set[int] = set() - self._next_tool_output_index: int = 1 # output_index=0 reserved for the message item + self._next_tool_output_index: int = ( + 1 # output_index=0 reserved for the message item + ) self._final_tool_events_queued: bool = False - self._sequence_number: int = 0 + self._sequence_number: int = 0 self._cached_reasoning_item_id: Optional[str] = None self._sent_reasoning_summary_text_done_event: bool = False self._sent_reasoning_summary_part_done_event: bool = False @@ -106,7 +108,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._reasoning_item_id: Optional[str] = None self._accumulated_reasoning_content_parts: List[str] = [] - def _get_or_assign_tool_output_index(self, call_id: str) -> int: existing = self._tool_output_index_by_call_id.get(call_id) if existing is not None: @@ -129,7 +130,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): except (TypeError, ValueError): return None - def _is_reasoning_end(self, chunk): delta = chunk.choices[0].delta @@ -153,7 +153,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): We emit: - response.output_item.added (function_call) - response.function_call_arguments.delta (split into smaller chunks to match OpenAI behavior) - + Note: Some providers (like Bedrock) send tool call arguments in one large chunk. We split these into smaller deltas to match OpenAI's token-by-token streaming behavior. """ @@ -162,7 +162,9 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): for tc in tool_calls: tc_index = self._normalize_tool_call_index(tc) - call_id_raw = tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None) + call_id_raw = ( + tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None) + ) call_id = "" if call_id_raw: @@ -184,7 +186,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if not call_id: continue - fn = tc.get("function") if isinstance(tc, dict) else getattr(tc, "function", None) + fn = ( + tc.get("function") + if isinstance(tc, dict) + else getattr(tc, "function", None) + ) fn_name = "" fn_args_delta = "" if isinstance(fn, dict): @@ -213,29 +219,33 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): } ), ) - event.__dict__['sequence_number'] = self._sequence_number + event.__dict__["sequence_number"] = self._sequence_number self._pending_tool_events.append(event) if fn_args_delta: self._tool_args_by_call_id[call_id] += fn_args_delta - + # Split large argument deltas into smaller chunks to match OpenAI's streaming behavior # This is especially important for providers like Bedrock that send complete arguments at once chunk_size = 10 # Match typical OpenAI delta size for i in range(0, len(fn_args_delta), chunk_size): - delta_chunk = fn_args_delta[i:i + chunk_size] + delta_chunk = fn_args_delta[i : i + chunk_size] self._sequence_number += 1 - delta_event: BaseLiteLLMOpenAIResponseObject = FunctionCallArgumentsDeltaEvent( - type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA, - item_id=call_id, - output_index=output_index, - delta=delta_chunk, + delta_event: BaseLiteLLMOpenAIResponseObject = ( + FunctionCallArgumentsDeltaEvent( + type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA, + item_id=call_id, + output_index=output_index, + delta=delta_chunk, + ) ) # Add sequence_number as extra field (BaseLiteLLMOpenAIResponseObject allows extra fields) - delta_event.__dict__['sequence_number'] = self._sequence_number + delta_event.__dict__["sequence_number"] = self._sequence_number self._pending_tool_events.append(delta_event) - def _queue_final_tool_call_done_events(self, litellm_complete_object: ModelResponse) -> None: + def _queue_final_tool_call_done_events( + self, litellm_complete_object: ModelResponse + ) -> None: """ Ensure tool calls that were not streamed as deltas still get emitted before response.completed. """ @@ -253,13 +263,19 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return for tc in tool_calls: - call_id_raw = tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None) + call_id_raw = ( + tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None) + ) if not call_id_raw: continue call_id = str(call_id_raw) output_index = self._get_or_assign_tool_output_index(call_id) - fn = tc.get("function") if isinstance(tc, dict) else getattr(tc, "function", None) + fn = ( + tc.get("function") + if isinstance(tc, dict) + else getattr(tc, "function", None) + ) fn_name = "" fn_args = "" if isinstance(fn, dict): @@ -271,7 +287,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): # Track if this is a new tool call that wasn't streamed is_new_tool_call = call_id not in self._tool_args_by_call_id - + # If we never sent output_item.added for this call_id, emit it now. if is_new_tool_call: self._tool_args_by_call_id[call_id] = "" @@ -290,21 +306,21 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): } ), ) - event.__dict__['sequence_number'] = self._sequence_number + event.__dict__["sequence_number"] = self._sequence_number self._pending_tool_events.append(event) final_args = fn_args or self._tool_args_by_call_id.get(call_id, "") - + # Emit delta events for arguments that weren't streamed yet # This handles cases where Bedrock sends the complete tool call at the end already_streamed = self._tool_args_by_call_id.get(call_id, "") - remaining_args = final_args[len(already_streamed):] if final_args else "" - + remaining_args = final_args[len(already_streamed) :] if final_args else "" + if remaining_args: # Split into smaller chunks to match OpenAI's streaming behavior chunk_size = 10 # Match typical OpenAI delta size for i in range(0, len(remaining_args), chunk_size): - delta_chunk = remaining_args[i:i + chunk_size] + delta_chunk = remaining_args[i : i + chunk_size] self._sequence_number += 1 delta_event = FunctionCallArgumentsDeltaEvent( type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA, @@ -312,9 +328,9 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): output_index=output_index, delta=delta_chunk, ) - delta_event.__dict__['sequence_number'] = self._sequence_number + delta_event.__dict__["sequence_number"] = self._sequence_number self._pending_tool_events.append(delta_event) - + self._sequence_number += 1 done_event = FunctionCallArgumentsDoneEvent( type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE, @@ -322,7 +338,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): output_index=output_index, arguments=final_args, ) - done_event.__dict__['sequence_number'] = self._sequence_number + done_event.__dict__["sequence_number"] = self._sequence_number self._pending_tool_events.append(done_event) self._sequence_number += 1 @@ -347,7 +363,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): # Use cached response ID if available, otherwise generate a new one if self._cached_response_id is None: self._cached_response_id = f"resp_{str(uuid.uuid4())}" - + response_created_event_data = { "id": self._cached_response_id, "object": "response", @@ -372,9 +388,12 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): response_created_event_data["text"] = self.responses_api_request["text"] if "tool_choice" in self.responses_api_request: # Transform tool_choice from dict format (e.g., {"type": "auto"}) to string format - response_created_event_data["tool_choice"] = LiteLLMCompletionResponsesConfig._transform_tool_choice( - self.responses_api_request["tool_choice"] - ) or "auto" + response_created_event_data["tool_choice"] = ( + LiteLLMCompletionResponsesConfig._transform_tool_choice( + self.responses_api_request["tool_choice"] + ) + or "auto" + ) else: response_created_event_data["tool_choice"] = "auto" if "tools" in self.responses_api_request: @@ -408,7 +427,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): type=ResponsesAPIStreamEvents.RESPONSE_CREATED, response=ResponsesAPIResponse(**response_created_event_data), ) - event.__dict__['sequence_number'] = self._sequence_number + event.__dict__["sequence_number"] = self._sequence_number return event def create_response_in_progress_event(self) -> ResponseInProgressEvent: @@ -419,13 +438,13 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): type=ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, response=ResponsesAPIResponse(**response_in_progress_event_data), ) - event.__dict__['sequence_number'] = self._sequence_number + event.__dict__["sequence_number"] = self._sequence_number return event def create_output_item_added_event(self) -> OutputItemAddedEvent: if self._cached_item_id is None: self._cached_item_id = f"msg_{str(uuid.uuid4())}" - + self._sequence_number += 1 event = OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, @@ -440,13 +459,13 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): } ), ) - event.__dict__['sequence_number'] = self._sequence_number + event.__dict__["sequence_number"] = self._sequence_number return event def create_content_part_added_event(self) -> ContentPartAddedEvent: if self._cached_item_id is None: self._cached_item_id = f"msg_{str(uuid.uuid4())}" - + self._sequence_number += 1 event = ContentPartAddedEvent( type=ResponsesAPIStreamEvents.CONTENT_PART_ADDED, @@ -457,7 +476,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): **{"type": "output_text", "text": "", "annotations": []} ), ) - event.__dict__['sequence_number'] = self._sequence_number + event.__dict__["sequence_number"] = self._sequence_number return event def create_litellm_model_response( @@ -483,7 +502,9 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): hidden_params = getattr(chunk, "_hidden_params", None) if hidden_params is not None: chunk_dict["_hidden_params"] = ( - dict(hidden_params) if isinstance(hidden_params, dict) else hidden_params + dict(hidden_params) + if isinstance(hidden_params, dict) + else hidden_params ) return chunk_dict @@ -495,7 +516,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ) -> ReasoningSummaryTextDoneEvent: """ Create response.reasoning_summary_text.done event. - + Example: { "type": "response.reasoning_summary_text.done", @@ -516,14 +537,14 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ) def create_reasoning_summary_part_done_event( - self, + self, reasoning_item_id: str, reasoning_content: str, sequence_number: int, ) -> ReasoningSummaryPartDoneEvent: """ Create response.reasoning_summary_part.done event. - + Example: { "type": "response.reasoning_summary_part.done", @@ -556,7 +577,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ) -> OutputTextDoneEvent: if self._cached_item_id is None: self._cached_item_id = f"msg_{str(uuid.uuid4())}" - + return OutputTextDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE, item_id=self._cached_item_id, @@ -607,7 +628,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ) -> OutputItemDoneEvent: if self._cached_item_id is None: self._cached_item_id = f"msg_{str(uuid.uuid4())}" - + text = self.litellm_model_response.choices[0].message.content or "" # type: ignore annotations = getattr(self.litellm_model_response.choices[0].message, "annotations", None) # type: ignore @@ -643,7 +664,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ) -> OutputItemDoneEvent: """ Create response.output_item.done event for reasoning items. - + Example: { "type": "response.output_item.done", @@ -776,7 +797,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): } ), ) - event.__dict__['sequence_number'] = self._sequence_number + event.__dict__["sequence_number"] = self._sequence_number self._pending_response_events.append(event) return @@ -801,7 +822,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): } ), ) - event.__dict__['sequence_number'] = self._sequence_number + event.__dict__["sequence_number"] = self._sequence_number self._pending_response_events.append(event) return @@ -840,41 +861,61 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): # Incrementally accumulate reasoning content instead of # calling stream_chunk_builder on every chunk (O(n²)) delta = chunk.choices[0].delta if chunk.choices else None - if delta and hasattr(delta, "reasoning_content") and delta.reasoning_content: - self._accumulated_reasoning_content_parts.append(delta.reasoning_content) + if ( + delta + and hasattr(delta, "reasoning_content") + and delta.reasoning_content + ): + self._accumulated_reasoning_content_parts.append( + delta.reasoning_content + ) if self._is_reasoning_end(chunk): - reasoning_content = "".join(self._accumulated_reasoning_content_parts) - + reasoning_content = "".join( + self._accumulated_reasoning_content_parts + ) + # Ensure we have a valid reasoning_item_id - reasoning_item_id = self._reasoning_item_id or self._cached_reasoning_item_id or f"rs_{uuid.uuid4()}" - + reasoning_item_id = ( + self._reasoning_item_id + or self._cached_reasoning_item_id + or f"rs_{uuid.uuid4()}" + ) + # Create text.done event first with its own sequence number self._sequence_number += 1 - text_done_event = self.create_reasoning_summary_text_done_event( - reasoning_item_id=reasoning_item_id, - reasoning_content=reasoning_content, - sequence_number=self._sequence_number + text_done_event = ( + self.create_reasoning_summary_text_done_event( + reasoning_item_id=reasoning_item_id, + reasoning_content=reasoning_content, + sequence_number=self._sequence_number, + ) ) - + # Create part.done event second with its own sequence number self._sequence_number += 1 - part_done_event = self.create_reasoning_summary_part_done_event( - reasoning_item_id=reasoning_item_id, - reasoning_content=reasoning_content, - sequence_number=self._sequence_number + part_done_event = ( + self.create_reasoning_summary_part_done_event( + reasoning_item_id=reasoning_item_id, + reasoning_content=reasoning_content, + sequence_number=self._sequence_number, + ) ) - + self._sequence_number += 1 - reasoning_output_item_done_event = self.create_reasoning_output_item_done_event( - reasoning_item_id=reasoning_item_id, - reasoning_content=reasoning_content, - sequence_number=self._sequence_number + reasoning_output_item_done_event = ( + self.create_reasoning_output_item_done_event( + reasoning_item_id=reasoning_item_id, + reasoning_content=reasoning_content, + sequence_number=self._sequence_number, + ) + ) + self._pending_response_events.extend( + [ + text_done_event, + part_done_event, + reasoning_output_item_done_event, + ] ) - self._pending_response_events.extend([ - text_done_event, - part_done_event, - reasoning_output_item_done_event, - ]) self._reasoning_done_emitted = True self._reasoning_active = False @@ -885,7 +926,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ) if response_api_chunk: self._pending_response_events.append(response_api_chunk) - + if self._pending_response_events: return self._pending_response_events.pop(0) @@ -959,7 +1000,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if self._cached_item_id is None and chunk.id: self._cached_item_id = chunk.id item_id = self._cached_item_id or chunk.id - + # Check if this chunk has annotations first (before processing text/reasoning) # This ensures we detect and queue annotation events from the annotation chunk if chunk.choices and hasattr(chunk.choices[0].delta, "annotations"): @@ -967,14 +1008,17 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if annotations and self.sent_annotation_events is False: self.sent_annotation_events = True # Store annotation events to emit them one by one - if not hasattr(self, '_pending_annotation_events'): - + if not hasattr(self, "_pending_annotation_events"): response_annotations = LiteLLMCompletionResponsesConfig._transform_chat_completion_annotations_to_response_output_annotations( annotations=annotations - ) + ) self._pending_annotation_events = [] for idx, annotation in enumerate(response_annotations): - annotation_dict = annotation.model_dump() if hasattr(annotation, 'model_dump') else dict(annotation) + annotation_dict = ( + annotation.model_dump() + if hasattr(annotation, "model_dump") + else dict(annotation) + ) event = OutputTextAnnotationAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED, item_id=item_id, @@ -983,7 +1027,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): annotation_index=idx, annotation=annotation_dict, ) - self._pending_annotation_events.append(event) + self._pending_annotation_events.append(event) # Priority 1: Handle reasoning content (highest priority) if ( chunk.choices @@ -998,7 +1042,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): output_index=0, delta=reasoning_content, ) - + # Priority 2: Handle text deltas delta_content = self._get_delta_string_from_streaming_choices(chunk.choices) if delta_content: @@ -1010,7 +1054,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): content_index=0, delta=delta_content, ) - text_delta_event.__dict__['sequence_number'] = self._sequence_number + text_delta_event.__dict__["sequence_number"] = self._sequence_number return text_delta_event # Priority 3: Handle tool call deltas (if any) -> queue events and emit them @@ -1024,10 +1068,13 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): # Return one pending tool event at a time if self._pending_tool_events: return self._pending_tool_events.pop(0) - + # Priority 4: If we have pending annotation events, emit the next one # This happens when the current chunk has no text/reasoning content - if hasattr(self, '_pending_annotation_events') and self._pending_annotation_events: + if ( + hasattr(self, "_pending_annotation_events") + and self._pending_annotation_events + ): event = self._pending_annotation_events.pop(0) return event @@ -1054,7 +1101,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): def _emit_response_completed_event( self, litellm_model_response: ModelResponse ) -> Optional[ResponseCompletedEvent]: - if litellm_model_response: # Add cost to usage object if include_cost_in_streaming_usage is True if ( diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 19845d7c49..0310d75895 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -292,22 +292,21 @@ class LiteLLMCompletionResponsesConfig: ) _messages = litellm_completion_request.get("messages") or [] session_messages = chat_completion_session.get("messages") or [] - + # If session messages are empty (e.g., no database in test environment), # we still need to process the new input messages # Store original _messages before combining for safety check original_new_messages = _messages.copy() if _messages else [] - + combined_messages = session_messages + _messages - + # Fix: Ensure tool_results have corresponding tool_calls in previous assistant message # Pass tools parameter to help reconstruct tool_calls if not in cache tools = litellm_completion_request.get("tools") or [] combined_messages = LiteLLMCompletionResponsesConfig._ensure_tool_results_have_corresponding_tool_calls( - messages=combined_messages, - tools=tools + messages=combined_messages, tools=tools ) - + # Safety check: Ensure we don't end up with empty messages # This can happen when using previous_response_id without a database (e.g., in tests) # and session messages are empty but new input messages exist @@ -326,6 +325,7 @@ class LiteLLMCompletionResponsesConfig: # Both are empty - this likely means function_call_output had empty/invalid call_id # Provide a helpful error message import litellm + raise litellm.BadRequestError( message=( f"Unable to create messages for completion request. " @@ -336,9 +336,11 @@ class LiteLLMCompletionResponsesConfig: f"Original request: previous_response_id={previous_response_id}" ), model=litellm_completion_request.get("model", ""), - llm_provider=litellm_completion_request.get("custom_llm_provider", ""), + llm_provider=litellm_completion_request.get( + "custom_llm_provider", "" + ), ) - + litellm_completion_request["messages"] = combined_messages litellm_completion_request["litellm_trace_id"] = chat_completion_session.get( "litellm_session_id" @@ -384,10 +386,45 @@ class LiteLLMCompletionResponsesConfig: if call_id_raw: existing_tool_call_ids.add(str(call_id_raw)) + ######################################################### + # Merge consecutive function_call items into a single assistant + # message. Anthropic requires that all tool_use blocks appear in + # ONE assistant message immediately followed by the tool_result + # blocks. Without this merging, each function_call creates its own + # assistant message, producing back-to-back assistant messages that + # Anthropic rejects with "tool_use ids were found without + # tool_result blocks immediately after". + ######################################################### + if messages: + last_msg = messages[-1] + last_role = ( + last_msg.get("role") + if isinstance(last_msg, dict) + else getattr(last_msg, "role", None) + ) + if last_role == "assistant": + for new_msg in chat_completion_messages: + new_role = ( + new_msg.get("role") + if isinstance(new_msg, dict) + else getattr(new_msg, "role", None) + ) + if new_role == "assistant": + new_tcs = ( + new_msg.get("tool_calls") + if isinstance(new_msg, dict) + else getattr(new_msg, "tool_calls", None) + ) or [] + for tc in new_tcs: + LiteLLMCompletionResponsesConfig._add_tool_call_to_assistant( + last_msg, tc + ) + continue + ######################################################### # If Input Item is a Tool Call Output, add it to the tool_call_output_messages list - # preserving the ordering of tool call outputs. Some models require the tool - # result to immediately follow the assistant tool call. + # preserving the ordering of tool call outputs. Some models require the tool + # result to immediately follow the assistant tool call. ######################################################### if LiteLLMCompletionResponsesConfig._is_input_item_tool_call_output( input_item=_input @@ -541,7 +578,11 @@ class LiteLLMCompletionResponsesConfig: if isinstance(assistant_message, dict) else getattr(assistant_message, "tool_calls", None) ) - if tool_calls_raw and isinstance(tool_calls_raw, list) and len(tool_calls_raw) > 0: + if ( + tool_calls_raw + and isinstance(tool_calls_raw, list) + and len(tool_calls_raw) > 0 + ): first_tool_call = tool_calls_raw[0] if isinstance(first_tool_call, dict): tool_call_id_raw = first_tool_call.get("id", "") @@ -633,8 +674,10 @@ class LiteLLMCompletionResponsesConfig: function_name_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( function_raw, "name" ) - function_arguments_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( - function_raw, "arguments" + function_arguments_raw = ( + LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( + function_raw, "arguments" + ) ) function: Dict[str, Any] = { "name": function_name_raw or "", @@ -675,11 +718,15 @@ class LiteLLMCompletionResponsesConfig: if isinstance(tool_use_definition, dict): normalized_definition: Dict[str, Any] = dict(tool_use_definition) else: - tool_use_id_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( - tool_use_definition, "id" + tool_use_id_raw = ( + LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( + tool_use_definition, "id" + ) ) - tool_use_type_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( - tool_use_definition, "type" + tool_use_type_raw = ( + LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( + tool_use_definition, "type" + ) ) function_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( tool_use_definition, "function" @@ -701,11 +748,15 @@ class LiteLLMCompletionResponsesConfig: function_raw = normalized_definition.get("function") if function_raw is not None and not isinstance(function_raw, dict): - function_name_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( - function_raw, "name" + function_name_raw = ( + LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( + function_raw, "name" + ) ) - function_arguments_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( - function_raw, "arguments" + function_arguments_raw = ( + LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value( + function_raw, "arguments" + ) ) if function_name_raw is not None or function_arguments_raw is not None: normalized_definition["function"] = { @@ -714,9 +765,7 @@ class LiteLLMCompletionResponsesConfig: } normalized_definition["id"] = normalized_definition.get("id") or tool_call_id - normalized_definition["type"] = ( - normalized_definition.get("type") or "function" - ) + normalized_definition["type"] = normalized_definition.get("type") or "function" return normalized_definition @staticmethod @@ -760,14 +809,14 @@ class LiteLLMCompletionResponsesConfig: ]: """ Ensure that tool_result messages have corresponding tool_calls in the previous assistant message. - + This is critical for Anthropic API which requires that each tool_result block has a corresponding tool_use block in the previous assistant message. - + Args: messages: List of messages that may include tool_result messages tools: Optional list of tools that can be used to reconstruct tool_calls if not in cache - + Returns: List of messages with tool_calls added to assistant messages when needed """ @@ -776,6 +825,7 @@ class LiteLLMCompletionResponsesConfig: # Create a deep copy to avoid modifying the original (use list() so we can mutate and return List) import copy + fixed_messages: List[ Union[ AllMessageValues, @@ -786,29 +836,35 @@ class LiteLLMCompletionResponsesConfig: ] ] = list(copy.deepcopy(messages)) messages_to_remove = [] - + # Count non-tool messages to avoid removing all messages # This prevents empty messages list when using previous_response_id without a database non_tool_messages_count = sum( 1 for msg in fixed_messages if msg.get("role") != "tool" ) - + for i, message in enumerate(fixed_messages): # Only process tool messages - check role first to narrow the type if message.get("role") != "tool": continue - + # At this point, we know it's a tool message, so it should have tool_call_id # Use get() with default to safely access tool_call_id - tool_call_id_raw = message.get("tool_call_id") if isinstance(message, dict) else getattr(message, "tool_call_id", None) + tool_call_id_raw = ( + message.get("tool_call_id") + if isinstance(message, dict) + else getattr(message, "tool_call_id", None) + ) tool_call_id: str = ( str(tool_call_id_raw) if tool_call_id_raw is not None else "" ) - - prev_assistant_idx = LiteLLMCompletionResponsesConfig._find_previous_assistant_idx( - fixed_messages, i + + prev_assistant_idx = ( + LiteLLMCompletionResponsesConfig._find_previous_assistant_idx( + fixed_messages, i + ) ) - + # Try to recover empty tool_call_id from previous assistant message if not tool_call_id and prev_assistant_idx is not None: prev_assistant = fixed_messages[prev_assistant_idx] @@ -823,7 +879,7 @@ class LiteLLMCompletionResponsesConfig: message_dict["tool_call_id"] = tool_call_id elif hasattr(message, "tool_call_id"): setattr(message, "tool_call_id", tool_call_id) - + # Only remove messages with empty tool_call_id if we have other non-tool messages # This prevents ending up with an empty messages list when using previous_response_id # without a database (e.g., in tests where session messages are empty) @@ -835,7 +891,7 @@ class LiteLLMCompletionResponsesConfig: # If no non-tool messages, keep the tool message even with empty call_id # The API will return a proper error message about the missing tool_use block continue - + # Check if the previous assistant message has the corresponding tool_call # This needs to run for ALL tool messages with a valid tool_call_id, # not just those that had an empty tool_call_id initially @@ -844,17 +900,15 @@ class LiteLLMCompletionResponsesConfig: tool_calls = LiteLLMCompletionResponsesConfig._get_tool_calls_list( prev_assistant ) - + if not LiteLLMCompletionResponsesConfig._check_tool_call_exists( tool_calls, tool_call_id ): _tool_use_definition = TOOL_CALLS_CACHE.get_cache(key=tool_call_id) - + if not _tool_use_definition and tools: - _tool_use_definition = ( - LiteLLMCompletionResponsesConfig._reconstruct_tool_call_from_tools( - tool_call_id, tools - ) + _tool_use_definition = LiteLLMCompletionResponsesConfig._reconstruct_tool_call_from_tools( + tool_call_id, tools ) normalized_tool_use_definition = ( @@ -874,11 +928,11 @@ class LiteLLMCompletionResponsesConfig: LiteLLMCompletionResponsesConfig._add_tool_call_to_assistant( prev_assistant, tool_call_chunk ) - + # Remove messages with empty tool_call_id that couldn't be fixed for idx in reversed(messages_to_remove): fixed_messages.pop(idx) - + return fixed_messages @staticmethod @@ -1014,7 +1068,10 @@ class LiteLLMCompletionResponsesConfig: ) elif isinstance(image_url_val, str) and image_url_val: normalized_blocks.append( - {"type": "image_url", "image_url": {"url": image_url_val}} + { + "type": "image_url", + "image_url": {"url": image_url_val}, + } ) # Prefer structured blocks if we have images; otherwise return a string. @@ -1078,7 +1135,9 @@ class LiteLLMCompletionResponsesConfig: function: dict = _tool_use_definition.get("function") or {} tool_call_chunk = ChatCompletionToolCallChunk( id=_tool_use_definition.get("id") or "", - type=cast(Literal["function"], _tool_use_definition.get("type") or "function"), + type=cast( + Literal["function"], _tool_use_definition.get("type") or "function" + ), function=ChatCompletionToolCallFunctionChunk( name=function.get("name") or "", arguments=str(function.get("arguments") or ""), @@ -1314,7 +1373,7 @@ class LiteLLMCompletionResponsesConfig: "description": typed_tool.get("description") or "", "parameters": parameters, "strict": typed_tool.get("strict", False) or False, - } + }, } if tool.get("cache_control"): chat_completion_tool["cache_control"] = tool.get("cache_control") # type: ignore @@ -1328,7 +1387,9 @@ class LiteLLMCompletionResponsesConfig: cast(ChatCompletionToolParam, chat_completion_tool) ) else: - chat_completion_tools.append(cast(Union[ChatCompletionToolParam, OpenAIMcpServerTool], tool)) + chat_completion_tools.append( + cast(Union[ChatCompletionToolParam, OpenAIMcpServerTool], tool) + ) return chat_completion_tools, web_search_options @staticmethod @@ -1523,6 +1584,39 @@ class LiteLLMCompletionResponsesConfig: return tool_call_dict + @staticmethod + def convert_apply_patch_tool_call_to_chat_completion_tool_call( + tool_call_item: Any, + index: int = 0, + ) -> Dict[str, Any]: + """ + Convert ResponseApplyPatchToolCall to ChatCompletionToolCallChunk format. + + The operation (create_file / update_file / delete_file) is serialised + as JSON so it appears in function.arguments, just like any other + tool call. + + Args: + tool_call_item: ResponseApplyPatchToolCall object with call_id and operation + index: The index of this tool call + + Returns: + Dictionary in ChatCompletionToolCallChunk format + """ + import json + + operation_dict = tool_call_item.operation.model_dump() + tool_call_dict: Dict[str, Any] = { + "id": tool_call_item.call_id, + "function": { + "name": "apply_patch", + "arguments": json.dumps(operation_dict), + }, + "type": "function", + "index": index, + } + return tool_call_dict + @staticmethod def transform_chat_completion_response_to_responses_api_response( request_input: Union[str, ResponseInputParam], @@ -1579,10 +1673,14 @@ class LiteLLMCompletionResponsesConfig: ), user=getattr(chat_completion_response, "user", None), ) - responses_api_response._hidden_params = getattr(chat_completion_response, "_hidden_params", {}) + responses_api_response._hidden_params = getattr( + chat_completion_response, "_hidden_params", {} + ) # Surface provider-specific fields (generic passthrough from any provider) - provider_fields = responses_api_response._hidden_params.get("provider_specific_fields") + provider_fields = responses_api_response._hidden_params.get( + "provider_specific_fields" + ) if provider_fields: setattr(responses_api_response, "provider_specific_fields", provider_fields) @@ -1954,9 +2052,9 @@ class LiteLLMCompletionResponsesConfig: hasattr(completion_details, "reasoning_tokens") and completion_details.reasoning_tokens is not None ): - output_details_dict["reasoning_tokens"] = ( - completion_details.reasoning_tokens - ) + output_details_dict[ + "reasoning_tokens" + ] = completion_details.reasoning_tokens else: output_details_dict["reasoning_tokens"] = 0 diff --git a/litellm/responses/main.py b/litellm/responses/main.py index e6160d95d4..6f7e38dc8b 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -185,10 +185,13 @@ async def aresponses_api_with_mcp( mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None secret_fields = kwargs.get("secret_fields") if secret_fields and isinstance(secret_fields, dict): - mcp_auth_header, mcp_server_auth_headers, _, _ = ( - ResponsesAPIRequestUtils.extract_mcp_headers_from_request( - secret_fields=secret_fields, tools=tools - ) + ( + mcp_auth_header, + mcp_server_auth_headers, + _, + _, + ) = ResponsesAPIRequestUtils.extract_mcp_headers_from_request( + secret_fields=secret_fields, tools=tools ) # Get original MCP tools (for events) and OpenAI tools (for LLM) by reusing existing methods @@ -1714,12 +1717,15 @@ async def _aresponses_websocket( litellm_params = GenericLiteLLMParams(**kwargs) litellm_params_dict = get_litellm_params(**kwargs) - model, _custom_llm_provider, dynamic_api_key, dynamic_api_base = ( - litellm.get_llm_provider( - model=model, - api_base=api_base, - api_key=api_key, - ) + ( + model, + _custom_llm_provider, + dynamic_api_key, + dynamic_api_base, + ) = litellm.get_llm_provider( + model=model, + api_base=api_base, + api_key=api_key, ) litellm_logging_obj.update_environment_variables( @@ -1740,10 +1746,7 @@ async def _aresponses_websocket( ) resolved_api_base = ( - dynamic_api_base - or litellm_params.api_base - or litellm.api_base - or None + dynamic_api_base or litellm_params.api_base or litellm.api_base or None ) resolved_api_key = ( dynamic_api_key @@ -1754,7 +1757,11 @@ async def _aresponses_websocket( ) # Extract params that we're passing explicitly to avoid duplicates in **kwargs - remaining_kwargs = {k: v for k, v in kwargs.items() if k not in {"user_api_key_dict", "litellm_metadata"}} + remaining_kwargs = { + k: v + for k, v in kwargs.items() + if k not in {"user_api_key_dict", "litellm_metadata"} + } await base_llm_http_handler.async_responses_websocket( model=model, diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index bacc627cc8..24b5db2857 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -24,13 +24,13 @@ def _add_mcp_metadata_to_response( ) -> None: """ Add MCP metadata to response's provider_specific_fields. - + This function adds MCP-related information to the response so that clients can access which tools were available, which were called, and what results were returned. - + For ModelResponse: adds to choices[].message.provider_specific_fields - For CustomStreamWrapper: stores in _hidden_params and automatically adds to + For CustomStreamWrapper: stores in _hidden_params and automatically adds to final chunk's delta.provider_specific_fields via CustomStreamWrapper._add_mcp_metadata_to_final_chunk() """ if isinstance(response, CustomStreamWrapper): @@ -39,7 +39,7 @@ def _add_mcp_metadata_to_response( # add it to the final chunk's delta.provider_specific_fields if not hasattr(response, "_hidden_params"): response._hidden_params = {} - + mcp_metadata = {} if openai_tools: mcp_metadata["mcp_list_tools"] = openai_tools @@ -47,26 +47,24 @@ def _add_mcp_metadata_to_response( mcp_metadata["mcp_tool_calls"] = tool_calls if tool_results: mcp_metadata["mcp_call_results"] = tool_results - + if mcp_metadata: response._hidden_params["mcp_metadata"] = mcp_metadata return - + if not isinstance(response, ModelResponse): return - + if not hasattr(response, "choices") or not response.choices: return - + # Add MCP metadata to all choices' messages for choice in response.choices: message = getattr(choice, "message", None) if message is not None: # Get existing provider_specific_fields or create new dict - provider_fields = ( - getattr(message, "provider_specific_fields", None) or {} - ) - + provider_fields = getattr(message, "provider_specific_fields", None) or {} + # Add MCP metadata if openai_tools: provider_fields["mcp_list_tools"] = openai_tools @@ -74,7 +72,7 @@ def _add_mcp_metadata_to_response( provider_fields["mcp_tool_calls"] = tool_calls if tool_results: provider_fields["mcp_call_results"] = tool_results - + # Set the provider_specific_fields setattr(message, "provider_specific_fields", provider_fields) @@ -207,10 +205,22 @@ async def acompletion_with_mcp( # noqa: PLR0915 class MCPStreamingIterator: """Custom iterator that collects chunks, detects tool calls, and adds MCP metadata to final chunk.""" - - def __init__(self, stream_wrapper, messages, tool_server_map, user_api_key_auth, - mcp_auth_header, mcp_server_auth_headers, oauth2_headers, raw_headers, - litellm_call_id, litellm_trace_id, openai_tools, base_call_args): + + def __init__( + self, + stream_wrapper, + messages, + tool_server_map, + user_api_key_auth, + mcp_auth_header, + mcp_server_auth_headers, + oauth2_headers, + raw_headers, + litellm_call_id, + litellm_trace_id, + openai_tools, + base_call_args, + ): self.stream_wrapper = stream_wrapper self.messages = messages self.tool_server_map = tool_server_map @@ -236,93 +246,116 @@ async def acompletion_with_mcp( # noqa: PLR0915 async def __aiter__(self): return self - def _add_mcp_list_tools_to_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: + def _add_mcp_list_tools_to_chunk( + self, chunk: ModelResponseStream + ) -> ModelResponseStream: """Add mcp_list_tools to the first chunk.""" from litellm.types.utils import ( StreamingChoices, add_provider_specific_fields, ) - + if not self.openai_tools: return chunk - + if hasattr(chunk, "choices") and chunk.choices: for choice in chunk.choices: - if isinstance(choice, StreamingChoices) and hasattr(choice, "delta") and choice.delta: + if ( + isinstance(choice, StreamingChoices) + and hasattr(choice, "delta") + and choice.delta + ): # Get existing provider_specific_fields or create new dict - existing_fields = getattr(choice.delta, "provider_specific_fields", None) or {} - provider_fields = dict(existing_fields) # Create a copy to avoid mutating the original - + existing_fields = ( + getattr(choice.delta, "provider_specific_fields", None) + or {} + ) + provider_fields = dict( + existing_fields + ) # Create a copy to avoid mutating the original + # Add only mcp_list_tools to first chunk provider_fields["mcp_list_tools"] = self.openai_tools - + # Use add_provider_specific_fields to ensure proper setting # This function handles Pydantic model attribute setting correctly add_provider_specific_fields(choice.delta, provider_fields) - + return chunk - def _add_mcp_tool_metadata_to_final_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: + def _add_mcp_tool_metadata_to_final_chunk( + self, chunk: ModelResponseStream + ) -> ModelResponseStream: """Add mcp_tool_calls and mcp_call_results to the final chunk.""" from litellm.types.utils import ( StreamingChoices, add_provider_specific_fields, ) - + if hasattr(chunk, "choices") and chunk.choices: for choice in chunk.choices: - if isinstance(choice, StreamingChoices) and hasattr(choice, "delta") and choice.delta: + if ( + isinstance(choice, StreamingChoices) + and hasattr(choice, "delta") + and choice.delta + ): # Get existing provider_specific_fields or create new dict # Access the attribute directly to handle Pydantic model attributes correctly existing_fields = {} if hasattr(choice.delta, "provider_specific_fields"): - attr_value = getattr(choice.delta, "provider_specific_fields", None) + attr_value = getattr( + choice.delta, "provider_specific_fields", None + ) if attr_value is not None: # Create a copy to avoid mutating the original - existing_fields = dict(attr_value) if isinstance(attr_value, dict) else {} - + existing_fields = ( + dict(attr_value) + if isinstance(attr_value, dict) + else {} + ) + provider_fields = existing_fields - + # Add tool_calls and tool_results if available if self.tool_calls: provider_fields["mcp_tool_calls"] = self.tool_calls if self.tool_results: provider_fields["mcp_call_results"] = self.tool_results - + # Use add_provider_specific_fields to ensure proper setting # This function handles Pydantic model attribute setting correctly add_provider_specific_fields(choice.delta, provider_fields) - + return chunk async def __anext__(self): # Phase 1: Collect and yield initial stream chunks if not self.stream_exhausted: # Get the iterator from the stream wrapper - if not hasattr(self, '_stream_iterator'): + if not hasattr(self, "_stream_iterator"): self._stream_iterator = self.stream_wrapper.__aiter__() # Add mcp_list_tools to the first chunk (available from the start) _add_mcp_metadata_to_response( response=self.stream_wrapper, openai_tools=self.openai_tools, ) - + try: chunk = await self._stream_iterator.__anext__() self.collected_chunks.append(chunk) - + # Add mcp_list_tools to the first chunk if len(self.collected_chunks) == 1: chunk = self._add_mcp_list_tools_to_chunk(chunk) - + # Check if this is the final chunk (has finish_reason) is_final = ( - hasattr(chunk, "choices") - and chunk.choices + hasattr(chunk, "choices") + and chunk.choices and hasattr(chunk.choices[0], "finish_reason") and chunk.choices[0].finish_reason is not None ) - + if is_final: # This is the final chunk, mark stream as exhausted self.stream_exhausted = True @@ -333,7 +366,7 @@ async def acompletion_with_mcp( # noqa: PLR0915 # If we have tool results, prepare follow-up call immediately if self.tool_results and self.complete_response: await self._prepare_follow_up_call() - + return chunk except StopAsyncIteration: self.stream_exhausted = True @@ -342,50 +375,61 @@ async def acompletion_with_mcp( # noqa: PLR0915 # If we have chunks, yield the final one with metadata if self.collected_chunks: final_chunk = self.collected_chunks[-1] - final_chunk = self._add_mcp_tool_metadata_to_final_chunk(final_chunk) + final_chunk = self._add_mcp_tool_metadata_to_final_chunk( + final_chunk + ) # If we have tool results, prepare follow-up call if self.tool_results and self.complete_response: await self._prepare_follow_up_call() return final_chunk - + # Phase 2: Yield follow-up stream chunks if available if self.follow_up_stream and not self.follow_up_exhausted: if not self.follow_up_iterator: self.follow_up_iterator = self.follow_up_stream.__aiter__() from litellm._logging import verbose_logger + verbose_logger.debug("Follow-up stream iterator created") - + try: chunk = await self.follow_up_iterator.__anext__() from litellm._logging import verbose_logger + verbose_logger.debug(f"Follow-up chunk yielded: {chunk}") return chunk except StopAsyncIteration: self.follow_up_exhausted = True from litellm._logging import verbose_logger + verbose_logger.debug("Follow-up stream exhausted") # After follow-up stream is exhausted, check if we need to raise StopAsyncIteration raise StopAsyncIteration - + # If we're here and follow_up_stream is None but we expected it, log a warning - if self.stream_exhausted and self.tool_results and self.complete_response and self.follow_up_stream is None: + if ( + self.stream_exhausted + and self.tool_results + and self.complete_response + and self.follow_up_stream is None + ): from litellm._logging import verbose_logger + verbose_logger.warning( "Follow-up stream was not created despite having tool results" ) - + raise StopAsyncIteration async def _process_tool_calls(self): """Process tool calls after streaming completes.""" if self.tool_execution_done: return - + self.tool_execution_done = True - + if not self.collected_chunks: return - + # Build complete response from chunks complete_response = stream_chunk_builder( chunks=self.collected_chunks, @@ -401,31 +445,35 @@ async def acompletion_with_mcp( # noqa: PLR0915 if self.tool_calls: # Execute tool calls - self.tool_results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( - tool_server_map=self.tool_server_map, - tool_calls=self.tool_calls, - user_api_key_auth=self.user_api_key_auth, - mcp_auth_header=self.mcp_auth_header, - mcp_server_auth_headers=self.mcp_server_auth_headers, - oauth2_headers=self.oauth2_headers, - raw_headers=self.raw_headers, - litellm_call_id=self.litellm_call_id, - litellm_trace_id=self.litellm_trace_id, + self.tool_results = ( + await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( + tool_server_map=self.tool_server_map, + tool_calls=self.tool_calls, + user_api_key_auth=self.user_api_key_auth, + mcp_auth_header=self.mcp_auth_header, + mcp_server_auth_headers=self.mcp_server_auth_headers, + oauth2_headers=self.oauth2_headers, + raw_headers=self.raw_headers, + litellm_call_id=self.litellm_call_id, + litellm_trace_id=self.litellm_trace_id, + ) ) async def _prepare_follow_up_call(self): """Prepare and initiate follow-up call with tool results.""" if self.follow_up_stream is not None: return # Already prepared - + if not self.tool_results or not self.complete_response: return - + # Create follow-up messages with tool results - follow_up_messages = LiteLLM_Proxy_MCP_Handler._create_follow_up_messages_for_chat( - original_messages=self.messages, - response=self.complete_response, - tool_results=self.tool_results, + follow_up_messages = ( + LiteLLM_Proxy_MCP_Handler._create_follow_up_messages_for_chat( + original_messages=self.messages, + response=self.complete_response, + tool_results=self.tool_results, + ) ) # Make follow-up call with streaming @@ -438,16 +486,19 @@ async def acompletion_with_mcp( # noqa: PLR0915 # Import litellm here to ensure we get the patched version # This ensures the patch works correctly in tests import litellm + follow_up_response = await litellm.acompletion(**follow_up_call_args) - + # Ensure follow-up response is a CustomStreamWrapper if isinstance(follow_up_response, CustomStreamWrapper): self.follow_up_stream = follow_up_response from litellm._logging import verbose_logger + verbose_logger.debug("Follow-up stream created successfully") else: # Unexpected response type - log and set to None from litellm._logging import verbose_logger + verbose_logger.warning( f"Follow-up response is not a CustomStreamWrapper: {type(follow_up_response)}" ) @@ -478,10 +529,14 @@ async def acompletion_with_mcp( # noqa: PLR0915 completion_stream=None, model=getattr(original_wrapper, "model", "unknown"), logging_obj=getattr(original_wrapper, "logging_obj", None), - custom_llm_provider=getattr(original_wrapper, "custom_llm_provider", None), + custom_llm_provider=getattr( + original_wrapper, "custom_llm_provider", None + ), stream_options=getattr(original_wrapper, "stream_options", None), make_call=getattr(original_wrapper, "make_call", None), - _response_headers=getattr(original_wrapper, "_response_headers", None), + _response_headers=getattr( + original_wrapper, "_response_headers", None + ), ) self._original_wrapper = original_wrapper self._custom_iterator = custom_iterator @@ -499,12 +554,15 @@ async def acompletion_with_mcp( # noqa: PLR0915 # For synchronous iteration, create a sync wrapper if self._sync_iterator is None: import asyncio + try: self._sync_loop = asyncio.get_event_loop() except RuntimeError: self._sync_loop = asyncio.new_event_loop() asyncio.set_event_loop(self._sync_loop) - self._sync_iterator = _SyncIteratorWrapper(self._custom_iterator, self._sync_loop) + self._sync_iterator = _SyncIteratorWrapper( + self._custom_iterator, self._sync_loop + ) return self._sync_iterator def __next__(self): @@ -531,7 +589,7 @@ async def acompletion_with_mcp( # noqa: PLR0915 if self._iterator is None: # __aiter__ might be async, so we need to await it aiter_result = self._async_iterator.__aiter__() - if hasattr(aiter_result, '__await__'): + if hasattr(aiter_result, "__await__"): # It's a coroutine, await it self._iterator = self._loop.run_until_complete(aiter_result) else: diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index 0b0d9744df..7aed48c2f9 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -305,10 +305,10 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): # Mark as async iterator self.is_async = True - + # Track if we've emitted initial OpenAI lifecycle events self.initial_events_emitted = False - + # Cache the response ID to ensure consistency across all events self._cached_response_id: Optional[str] = None @@ -489,7 +489,9 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): response_obj = getattr(chunk, "response", None) if response_obj and hasattr(response_obj, "id"): self._cached_response_id = response_obj.id - verbose_logger.debug(f"Cached response ID: {self._cached_response_id}") + verbose_logger.debug( + f"Cached response ID: {self._cached_response_id}" + ) # After emitting response.output_item.added, transition to MCP discovery if not self.initial_events_emitted and hasattr(chunk, "type"): @@ -542,15 +544,17 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): """ if not self.base_iterator or not hasattr(self.base_iterator, "__anext__"): raise StopAsyncIteration - + chunk = await cast(Any, self.base_iterator).__anext__() # type: ignore[attr-defined] # Ensure response ID consistency - update chunk if needed - if self._cached_response_id and hasattr(chunk, 'response'): - response_obj = getattr(chunk, 'response', None) - if response_obj and hasattr(response_obj, 'id'): + if self._cached_response_id and hasattr(chunk, "response"): + response_obj = getattr(chunk, "response", None) + if response_obj and hasattr(response_obj, "id"): if response_obj.id != self._cached_response_id: - verbose_logger.debug(f"Updating response ID from {response_obj.id} to {self._cached_response_id}") + verbose_logger.debug( + f"Updating response ID from {response_obj.id} to {self._cached_response_id}" + ) response_obj.id = self._cached_response_id # If auto-execution is enabled, check for completed responses diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 5c5a955ae0..073ee92606 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -133,19 +133,16 @@ class BaseResponsesAPIStreamingIterator: # if "response" in parsed_chunk, then encode litellm specific information like custom_llm_provider response_object = getattr(openai_responses_api_chunk, "response", None) if response_object: - response = ( - ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( - responses_api_response=response_object, - litellm_metadata=self.litellm_metadata, - custom_llm_provider=self.custom_llm_provider, - ) + response = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( + responses_api_response=response_object, + litellm_metadata=self.litellm_metadata, + custom_llm_provider=self.custom_llm_provider, ) setattr(openai_responses_api_chunk, "response", response) # Wrap encrypted_content in streaming events (output_item.added, output_item.done) - if ( - self.litellm_metadata - and self.litellm_metadata.get("encrypted_content_affinity_enabled") + if self.litellm_metadata and self.litellm_metadata.get( + "encrypted_content_affinity_enabled" ): event_type = getattr(openai_responses_api_chunk, "type", None) if event_type in ( @@ -157,7 +154,9 @@ class BaseResponsesAPIStreamingIterator: encrypted_content = getattr(item, "encrypted_content", None) if encrypted_content and isinstance(encrypted_content, str): model_id = ( - self.litellm_metadata.get("model_info", {}).get("id") + self.litellm_metadata.get("model_info", {}).get( + "id" + ) if self.litellm_metadata else None ) @@ -188,10 +187,10 @@ class BaseResponsesAPIStreamingIterator: ) if usage_obj is not None: try: - cost: Optional[float] = ( - self.logging_obj._response_cost_calculator( - result=response_obj - ) + cost: Optional[ + float + ] = self.logging_obj._response_cost_calculator( + result=response_obj ) if cost is not None: setattr(usage_obj, "cost", cost) @@ -231,7 +230,9 @@ class BaseResponsesAPIStreamingIterator: typed_call_type = None if typed_call_type is None: try: - typed_call_type = CallTypes(getattr(self.logging_obj, "call_type", None)) + typed_call_type = CallTypes( + getattr(self.logging_obj, "call_type", None) + ) except Exception: typed_call_type = None @@ -333,7 +334,7 @@ class BaseResponsesAPIStreamingIterator: if self._failure_handled: return self._failure_handled = True - + traceback_exception = traceback.format_exc() try: run_async_function( @@ -445,7 +446,9 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): # Use model_dump + model_validate instead of deepcopy to avoid pickle errors with # Pydantic ValidatorIterator when response contains tool_choice with allowed_tools (fixes #17192) logging_response = self.completed_response - if self.completed_response is not None and hasattr(self.completed_response, 'model_dump'): + if self.completed_response is not None and hasattr( + self.completed_response, "model_dump" + ): try: logging_response = type(self.completed_response).model_validate( self.completed_response.model_dump() @@ -550,7 +553,9 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): # Use model_dump + model_validate instead of deepcopy to avoid pickle errors with # Pydantic ValidatorIterator when response contains tool_choice with allowed_tools (fixes #17192) logging_response = self.completed_response - if self.completed_response is not None and hasattr(self.completed_response, 'model_dump'): + if self.completed_response is not None and hasattr( + self.completed_response, "model_dump" + ): try: logging_response = type(self.completed_response).model_validate( self.completed_response.model_dump() @@ -633,9 +638,7 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): # Add cost to usage object if include_cost_in_streaming_usage is True if litellm.include_cost_in_streaming_usage and logging_obj is not None: - usage_obj: Optional[ResponseAPIUsage] = getattr( - transformed, "usage", None - ) + usage_obj: Optional[ResponseAPIUsage] = getattr(transformed, "usage", None) if usage_obj is not None: try: cost: Optional[float] = logging_obj._response_cost_calculator( @@ -801,9 +804,7 @@ class ResponsesWebSocketStreaming: if self.input_messages: self.logging_obj.model_call_details["messages"] = self.input_messages if self.messages: - asyncio.create_task( - self.logging_obj.async_success_handler(self.messages) - ) + asyncio.create_task(self.logging_obj.async_success_handler(self.messages)) _ws_executor.submit(self.logging_obj.success_handler, self.messages) async def backend_to_client(self) -> None: @@ -826,13 +827,9 @@ class ResponsesWebSocketStreaming: await self.websocket.send_text(response_str) except websockets.exceptions.ConnectionClosed as e: # type: ignore - verbose_logger.debug( - "Responses WS backend connection closed: %s", e - ) + verbose_logger.debug("Responses WS backend connection closed: %s", e) except Exception as e: - verbose_logger.exception( - "Error in responses WS backend_to_client: %s", e - ) + verbose_logger.exception("Error in responses WS backend_to_client: %s", e) finally: await self._log_messages() @@ -874,16 +871,17 @@ class ResponsesWebSocketStreaming: # --------------------------------------------------------------------------- _RESPONSE_CREATE_PARAMS: frozenset = ( - ResponsesAPIRequestParams.__required_keys__ | ResponsesAPIRequestParams.__optional_keys__ + ResponsesAPIRequestParams.__required_keys__ + | ResponsesAPIRequestParams.__optional_keys__ ) _MANAGED_WS_SKIP_KWARGS: frozenset = frozenset( { - "litellm_logging_obj", - "litellm_call_id", - "aresponses", - "_aresponses_websocket", - "user_api_key_dict", + "litellm_logging_obj", + "litellm_call_id", + "aresponses", + "_aresponses_websocket", + "user_api_key_dict", } ) @@ -953,13 +951,17 @@ class ManagedResponsesWebSocketHandler: return json.dumps(chunk, default=str) return json.dumps(str(chunk)) except Exception as exc: - verbose_logger.debug("ManagedResponsesWS: failed to serialize chunk: %s", exc) + verbose_logger.debug( + "ManagedResponsesWS: failed to serialize chunk: %s", exc + ) return None async def _send_error(self, message: str, error_type: str = "server_error") -> None: try: await self.websocket.send_text( - json.dumps({"type": "error", "error": {"type": error_type, "message": message}}) + json.dumps( + {"type": "error", "error": {"type": error_type, "message": message}} + ) ) except Exception: pass @@ -993,14 +995,18 @@ class ManagedResponsesWebSocketHandler: Returns *None* if the event doesn't contain a usable ID. """ resp_obj = completed_event.get("response", {}) - encoded_id: Optional[str] = resp_obj.get("id") if isinstance(resp_obj, dict) else None + encoded_id: Optional[str] = ( + resp_obj.get("id") if isinstance(resp_obj, dict) else None + ) if not encoded_id: return None decoded = ResponsesAPIRequestUtils._decode_responses_api_response_id(encoded_id) return decoded.get("response_id", encoded_id) @staticmethod - def _extract_output_messages(completed_event: Dict[str, Any]) -> List[Dict[str, Any]]: + def _extract_output_messages( + completed_event: Dict[str, Any] + ) -> List[Dict[str, Any]]: """ Convert the output items in a ``response.completed`` event into Responses API message dicts suitable for the next turn's ``input``. @@ -1023,7 +1029,13 @@ class ManagedResponsesWebSocketHandler: ] text = "".join(text_parts) if text: - messages.append({"type": "message", "role": role, "content": [{"type": "output_text", "text": text}]}) + messages.append( + { + "type": "message", + "role": role, + "content": [{"type": "output_text", "text": text}], + } + ) elif item_type == "function_call": messages.append(item) return messages @@ -1035,7 +1047,13 @@ class ManagedResponsesWebSocketHandler: of Responses API message dicts. """ if isinstance(input_val, str): - return [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": input_val}]}] + return [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": input_val}], + } + ] if isinstance(input_val, list): return [item for item in input_val if isinstance(item, dict)] return [] @@ -1049,7 +1067,9 @@ class ManagedResponsesWebSocketHandler: try: msg_obj = json.loads(raw_message) except json.JSONDecodeError: - await self._send_error("Invalid JSON in response.create event", "invalid_request_error") + await self._send_error( + "Invalid JSON in response.create event", "invalid_request_error" + ) return None if msg_obj.get("type") != "response.create": # Silently ignore non-response.create messages (e.g. warmup pings) @@ -1233,7 +1253,9 @@ class ManagedResponsesWebSocketHandler: event_model: Optional[str] = call_kwargs.pop("model", None) model = event_model or self.model - previous_response_id: Optional[str] = call_kwargs.pop("previous_response_id", None) + previous_response_id: Optional[str] = call_kwargs.pop( + "previous_response_id", None + ) current_messages = self._input_to_messages(call_kwargs.get("input")) # Fetch history once; reused in both _apply_history and _save_turn_history @@ -1243,7 +1265,9 @@ class ManagedResponsesWebSocketHandler: else [] ) - self._apply_history(call_kwargs, previous_response_id, current_messages, prior_history) + self._apply_history( + call_kwargs, previous_response_id, current_messages, prior_history + ) self._inject_credentials(call_kwargs, event_model) self._update_proxy_request(call_kwargs, model) call_kwargs.update(self.extra_kwargs) diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 89e8971170..1109786422 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -202,7 +202,7 @@ class ResponsesAPIRequestUtils: response_id = responses_api_response.get("id") else: response_id = getattr(responses_api_response, "id", None) - + # If no response_id, return the response as-is (likely an error response) if response_id is None: return responses_api_response @@ -248,7 +248,7 @@ class ResponsesAPIRequestUtils: if not encoded_id.startswith("encitem_"): return None try: - cleaned = encoded_id[len("encitem_"):] + cleaned = encoded_id[len("encitem_") :] # Restore any padding that may have been stripped in transit missing = len(cleaned) % 4 if missing: @@ -346,10 +346,10 @@ class ResponsesAPIRequestUtils: if encrypted_content and isinstance(encrypted_content, str): # Always wrap encrypted_content with model_id for redundancy - item["encrypted_content"] = ( - ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( - encrypted_content, model_id - ) + item[ + "encrypted_content" + ] = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( + encrypted_content, model_id ) # Also encode the ID if present if item_id and isinstance(item_id, str): @@ -363,10 +363,8 @@ class ResponsesAPIRequestUtils: if encrypted_content and isinstance(encrypted_content, str): # Always wrap encrypted_content with model_id for redundancy try: - item.encrypted_content = ( - ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( - encrypted_content, model_id - ) + item.encrypted_content = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( + encrypted_content, model_id ) except AttributeError: pass @@ -399,16 +397,19 @@ class ResponsesAPIRequestUtils: if isinstance(item, dict): item_id = item.get("id") if item_id and isinstance(item_id, str): - decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id(item_id) + decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id( + item_id + ) if decoded: item["id"] = decoded["item_id"] encrypted_content = item.get("encrypted_content") if encrypted_content and isinstance(encrypted_content, str): - _, unwrapped = ( - ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id( - encrypted_content - ) + ( + _, + unwrapped, + ) = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id( + encrypted_content ) if unwrapped != encrypted_content: item["encrypted_content"] = unwrapped @@ -579,17 +580,23 @@ class ResponsesAPIRequestUtils: raw_headers_from_request: Optional[Dict[str, str]] = None if secret_fields and isinstance(secret_fields, dict): raw_headers_from_request = secret_fields.get("raw_headers") - + # Extract MCP-specific headers using MCPRequestHandler methods mcp_auth_header: Optional[str] = None mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None oauth2_headers: Optional[Dict[str, str]] = None - + if raw_headers_from_request: headers_obj = Headers(raw_headers_from_request) - mcp_auth_header = MCPRequestHandler._get_mcp_auth_header_from_headers(headers_obj) - mcp_server_auth_headers = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers_obj) - oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers_obj) + mcp_auth_header = MCPRequestHandler._get_mcp_auth_header_from_headers( + headers_obj + ) + mcp_server_auth_headers = ( + MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers_obj) + ) + oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers( + headers_obj + ) if tools: for tool in tools: @@ -599,21 +606,35 @@ class ResponsesAPIRequestUtils: # Merge tool headers into mcp_server_auth_headers # Extract server-specific headers from tool.headers headers_obj_from_tool = Headers(tool_headers) - tool_mcp_server_auth_headers = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers_obj_from_tool) + tool_mcp_server_auth_headers = ( + MCPRequestHandler._get_mcp_server_auth_headers_from_headers( + headers_obj_from_tool + ) + ) if tool_mcp_server_auth_headers: if mcp_server_auth_headers is None: mcp_server_auth_headers = {} # Merge the headers from tool into existing headers - for server_alias, headers_dict in tool_mcp_server_auth_headers.items(): + for ( + server_alias, + headers_dict, + ) in tool_mcp_server_auth_headers.items(): if server_alias not in mcp_server_auth_headers: mcp_server_auth_headers[server_alias] = {} - mcp_server_auth_headers[server_alias].update(headers_dict) + mcp_server_auth_headers[server_alias].update( + headers_dict + ) # Also merge raw headers (non-prefixed headers from tool.headers) if raw_headers_from_request is None: raw_headers_from_request = {} raw_headers_from_request.update(tool_headers) - - return mcp_auth_header, mcp_server_auth_headers, oauth2_headers, raw_headers_from_request + + return ( + mcp_auth_header, + mcp_server_auth_headers, + oauth2_headers, + raw_headers_from_request, + ) class ResponseAPILoggingUtils: @@ -664,20 +685,32 @@ class ResponseAPILoggingUtils: ) else: prompt_tokens_details = PromptTokensDetailsWrapper( - cached_tokens=getattr(response_api_usage.input_tokens_details, "cached_tokens", None), - audio_tokens=getattr(response_api_usage.input_tokens_details, "audio_tokens", None), - text_tokens=getattr(response_api_usage.input_tokens_details, "text_tokens", None), - image_tokens=getattr(response_api_usage.input_tokens_details, "image_tokens", None), + cached_tokens=getattr( + response_api_usage.input_tokens_details, "cached_tokens", None + ), + audio_tokens=getattr( + response_api_usage.input_tokens_details, "audio_tokens", None + ), + text_tokens=getattr( + response_api_usage.input_tokens_details, "text_tokens", None + ), + image_tokens=getattr( + response_api_usage.input_tokens_details, "image_tokens", None + ), ) completion_tokens_details: Optional[CompletionTokensDetailsWrapper] = None - output_tokens_details = getattr(response_api_usage, "output_tokens_details", None) + output_tokens_details = getattr( + response_api_usage, "output_tokens_details", None + ) if output_tokens_details: completion_tokens_details = CompletionTokensDetailsWrapper( - reasoning_tokens=getattr(output_tokens_details, "reasoning_tokens", None), + reasoning_tokens=getattr( + output_tokens_details, "reasoning_tokens", None + ), image_tokens=getattr(output_tokens_details, "image_tokens", None), text_tokens=getattr(output_tokens_details, "text_tokens", None), ) - + chat_usage = Usage( prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, diff --git a/litellm/router.py b/litellm/router.py index 47de15655a..585cec682d 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -5561,6 +5561,10 @@ class Router: return response except Exception as e: + # Always track the latest error so we raise the most + # recent exception instead of the first one. + original_exception = e + ## LOGGING kwargs = self.log_retry(kwargs=kwargs, e=e) remaining_retries = num_retries - current_attempt - 1 @@ -5575,6 +5579,24 @@ class Router: ) else: _healthy_deployments = [] + + # Check if this error is non-retryable (e.g., 400 context + # window exceeded). If so, raise immediately instead of + # continuing the retry loop. Respect retry policy + # precedence - only check when no retry policy applies. + if not _retry_policy_applies: + try: + self.should_retry_this_error( + error=e, + healthy_deployments=_healthy_deployments, + all_deployments=_all_deployments, + context_window_fallbacks=context_window_fallbacks, + regular_fallbacks=fallbacks, + content_policy_fallbacks=content_policy_fallbacks, + ) + except Exception: + raise e + _timeout = self._time_to_sleep_before_retry( e=e, remaining_retries=remaining_retries, diff --git a/litellm/router_strategy/auto_router/auto_router.py b/litellm/router_strategy/auto_router/auto_router.py index c508d5b46c..6a78611519 100644 --- a/litellm/router_strategy/auto_router/auto_router.py +++ b/litellm/router_strategy/auto_router/auto_router.py @@ -19,6 +19,7 @@ else: class AutoRouter(CustomLogger): DEFAULT_AUTO_SYNC_VALUE = "local" + def __init__( self, model_name: str, @@ -27,7 +28,7 @@ class AutoRouter(CustomLogger): litellm_router_instance: "Router", auto_router_config_path: Optional[str] = None, auto_router_config: Optional[str] = None, - ): + ): """ Auto-Router class that uses a semantic router to route requests to the appropriate model. @@ -49,22 +50,22 @@ class AutoRouter(CustomLogger): self.default_model = default_model self.embedding_model: str = embedding_model self.litellm_router_instance: "Router" = litellm_router_instance - + def _load_semantic_routing_routes(self) -> List[Route]: from semantic_router.routers import SemanticRouter + if self.auto_router_config_path: return SemanticRouter.from_json(self.auto_router_config_path).routes elif self.auto_router_config: return self._load_auto_router_routes_from_config_json() else: raise ValueError("No router config provided") - def _load_auto_router_routes_from_config_json(self) -> List[Route]: import json from semantic_router.routers.base import Route - + if self.auto_router_config is None: raise ValueError("No auto router config provided") auto_router_routes: List[Route] = [] @@ -75,12 +76,11 @@ class AutoRouter(CustomLogger): name=route.get("name"), description=route.get("description"), utterances=route.get("utterances", []), - score_threshold=route.get("score_threshold") + score_threshold=route.get("score_threshold"), ) ) return auto_router_routes - async def async_pre_routing_hook( self, model: str, @@ -101,34 +101,36 @@ class AutoRouter(CustomLogger): LiteLLMRouterEncoder, ) from litellm.types.router import PreRoutingHookResponse + if messages is None: # do nothing, return same inputs return None - + if self.routelayer is None: ####################### # Create the route layer ####################### self.routelayer = SemanticRouter( - routes=self.loaded_routes, - encoder=LiteLLMRouterEncoder( - litellm_router_instance=self.litellm_router_instance, - model_name=self.embedding_model, - ), - auto_sync=self.auto_sync_value, + routes=self.loaded_routes, + encoder=LiteLLMRouterEncoder( + litellm_router_instance=self.litellm_router_instance, + model_name=self.embedding_model, + ), + auto_sync=self.auto_sync_value, ) - + user_message: Dict[str, str] = messages[-1] message_content: str = user_message.get("content", "") - route_choice: Optional[Union[RouteChoice, List[RouteChoice]]] = self.routelayer(text=message_content) + route_choice: Optional[Union[RouteChoice, List[RouteChoice]]] = self.routelayer( + text=message_content + ) verbose_router_logger.debug(f"route_choice: {route_choice}") if isinstance(route_choice, RouteChoice): model = route_choice.name or self.default_model elif isinstance(route_choice, list): model = route_choice[0].name or self.default_model - + return PreRoutingHookResponse( model=model, messages=messages, ) - diff --git a/litellm/router_strategy/auto_router/litellm_encoder.py b/litellm/router_strategy/auto_router/litellm_encoder.py index e0fd7c3625..1fe22eafdf 100644 --- a/litellm/router_strategy/auto_router/litellm_encoder.py +++ b/litellm/router_strategy/auto_router/litellm_encoder.py @@ -28,13 +28,13 @@ def litellm_to_list(embeds: litellm.EmbeddingResponse) -> list[list[float]]: class CustomDenseEncoder(DenseEncoder): - model_config = ConfigDict(extra='allow') + model_config = ConfigDict(extra="allow") def __init__(self, litellm_router_instance: Optional["Router"] = None, **kwargs): # Extract litellm_router_instance from kwargs if passed there - if 'litellm_router_instance' in kwargs: - litellm_router_instance = kwargs.pop('litellm_router_instance') - + if "litellm_router_instance" in kwargs: + litellm_router_instance = kwargs.pop("litellm_router_instance") + super().__init__(**kwargs) self.litellm_router_instance = litellm_router_instance @@ -91,9 +91,7 @@ class LiteLLMRouterEncoder(CustomDenseEncoder, AsymmetricDenseMixin): raise ValueError("litellm_router_instance is not set") try: embeds = self.litellm_router_instance.embedding( - input=docs, - model=self.model_name, - **kwargs + input=docs, model=self.model_name, **kwargs ) return litellm_to_list(embeds) except Exception as e: @@ -106,9 +104,7 @@ class LiteLLMRouterEncoder(CustomDenseEncoder, AsymmetricDenseMixin): raise ValueError("litellm_router_instance is not set") try: embeds = self.litellm_router_instance.embedding( - input=docs, - model=self.model_name, - **kwargs + input=docs, model=self.model_name, **kwargs ) return litellm_to_list(embeds) except Exception as e: @@ -121,9 +117,7 @@ class LiteLLMRouterEncoder(CustomDenseEncoder, AsymmetricDenseMixin): raise ValueError("litellm_router_instance is not set") try: embeds = await self.litellm_router_instance.aembedding( - input=docs, - model=self.model_name, - **kwargs + input=docs, model=self.model_name, **kwargs ) return litellm_to_list(embeds) except Exception as e: @@ -136,9 +130,7 @@ class LiteLLMRouterEncoder(CustomDenseEncoder, AsymmetricDenseMixin): raise ValueError("litellm_router_instance is not set") try: embeds = await self.litellm_router_instance.aembedding( - input=docs, - model=self.model_name, - **kwargs + input=docs, model=self.model_name, **kwargs ) return litellm_to_list(embeds) except Exception as e: diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 6ad2160666..29bed360fa 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -33,9 +33,9 @@ else: class DimensionScore: """Represents a score for a single dimension with optional signal.""" - + __slots__ = ("name", "score", "signal") - + def __init__(self, name: str, score: float, signal: Optional[str] = None): self.name = name self.score = score @@ -45,7 +45,7 @@ class DimensionScore: class ComplexityRouter(CustomLogger): """ Rule-based complexity router that classifies requests and routes to appropriate models. - + Handles requests in <1ms with zero external API calls by using weighted scoring across multiple dimensions: - Token count (short=simple, long=complex) @@ -56,7 +56,7 @@ class ComplexityRouter(CustomLogger): - Multi-step patterns ("first...then", numbered steps) - Question complexity (multiple questions) """ - + def __init__( self, model_name: str, @@ -66,7 +66,7 @@ class ComplexityRouter(CustomLogger): ): """ Initialize ComplexityRouter. - + Args: model_name: The name of the model/deployment using this router. litellm_router_instance: The LiteLLM Router instance. @@ -75,23 +75,27 @@ class ComplexityRouter(CustomLogger): """ self.model_name = model_name self.litellm_router_instance = litellm_router_instance - + # Parse config - always create a new instance to avoid singleton mutation if complexity_router_config: self.config = ComplexityRouterConfig(**complexity_router_config) else: self.config = ComplexityRouterConfig() - + # Override default_model if provided if default_model: self.config.default_model = default_model - + # Build effective keyword lists (use config overrides or defaults) self.code_keywords = self.config.code_keywords or DEFAULT_CODE_KEYWORDS - self.reasoning_keywords = self.config.reasoning_keywords or DEFAULT_REASONING_KEYWORDS - self.technical_keywords = self.config.technical_keywords or DEFAULT_TECHNICAL_KEYWORDS + self.reasoning_keywords = ( + self.config.reasoning_keywords or DEFAULT_REASONING_KEYWORDS + ) + self.technical_keywords = ( + self.config.technical_keywords or DEFAULT_TECHNICAL_KEYWORDS + ) self.simple_keywords = self.config.simple_keywords or DEFAULT_SIMPLE_KEYWORDS - + # Pre-compile regex patterns for efficiency # Use non-greedy .*? to prevent ReDoS on pathological inputs self._multi_step_patterns = [ @@ -100,38 +104,34 @@ class ComplexityRouter(CustomLogger): re.compile(r"\d+\.\s"), re.compile(r"[a-z]\)\s", re.IGNORECASE), ] - + verbose_router_logger.debug( f"ComplexityRouter initialized for {model_name} with tiers: {self.config.tiers}" ) - + def _estimate_tokens(self, text: str) -> int: """ Estimate token count from text. Uses a simple heuristic: ~4 characters per token on average. """ return len(text) // 4 - + def _score_token_count(self, estimated_tokens: int) -> DimensionScore: """Score based on token count.""" thresholds = self.config.token_thresholds simple_threshold = thresholds.get("simple", 15) complex_threshold = thresholds.get("complex", 400) - + if estimated_tokens < simple_threshold: return DimensionScore( - "tokenCount", - -1.0, - f"short ({estimated_tokens} tokens)" + "tokenCount", -1.0, f"short ({estimated_tokens} tokens)" ) if estimated_tokens > complex_threshold: return DimensionScore( - "tokenCount", - 1.0, - f"long ({estimated_tokens} tokens)" + "tokenCount", 1.0, f"long ({estimated_tokens} tokens)" ) return DimensionScore("tokenCount", 0, None) - + def _keyword_matches(self, text: str, keyword: str) -> bool: """ Check if a keyword matches in text using word boundary matching. @@ -145,12 +145,12 @@ class ComplexityRouter(CustomLogger): # For single-word keywords, use word boundary matching to avoid false positives # e.g., "api" should not match "capital", "error" should not match "terrorism" if " " not in kw_lower: - pattern = r'\b' + re.escape(kw_lower) + r'\b' + pattern = r"\b" + re.escape(kw_lower) + r"\b" return bool(re.search(pattern, text)) # For multi-word phrases, substring matching is fine return kw_lower in text - + def _score_keyword_match( self, text: str, @@ -172,49 +172,45 @@ class ComplexityRouter(CustomLogger): match_count = len(matches) if match_count >= high_threshold: - return DimensionScore( - name, - score_high, - f"{signal_label} ({', '.join(matches[:3])})" - ), match_count + return ( + DimensionScore( + name, score_high, f"{signal_label} ({', '.join(matches[:3])})" + ), + match_count, + ) if match_count >= low_threshold: - return DimensionScore( - name, - score_low, - f"{signal_label} ({', '.join(matches[:3])})" - ), match_count + return ( + DimensionScore( + name, score_low, f"{signal_label} ({', '.join(matches[:3])})" + ), + match_count, + ) return DimensionScore(name, score_none, None), match_count - + def _score_multi_step(self, text: str) -> DimensionScore: """Score based on multi-step patterns.""" hits = sum(1 for p in self._multi_step_patterns if p.search(text)) if hits > 0: return DimensionScore("multiStepPatterns", 0.5, "multi-step") return DimensionScore("multiStepPatterns", 0, None) - + def _score_question_complexity(self, text: str) -> DimensionScore: """Score based on number of question marks.""" count = text.count("?") if count > 3: - return DimensionScore( - "questionComplexity", - 0.5, - f"{count} questions" - ) + return DimensionScore("questionComplexity", 0.5, f"{count} questions") return DimensionScore("questionComplexity", 0, None) - + def classify( - self, - prompt: str, - system_prompt: Optional[str] = None + self, prompt: str, system_prompt: Optional[str] = None ) -> Tuple[ComplexityTier, float, List[str]]: """ Classify a prompt by complexity. - + Args: prompt: The user's prompt/message. system_prompt: Optional system prompt for context. - + Returns: Tuple of (tier, score, signals) where: - tier: The ComplexityTier (SIMPLE, MEDIUM, COMPLEX, REASONING) @@ -228,26 +224,42 @@ class ComplexityRouter(CustomLogger): # user_text only to prevent system prompts from forcing REASONING tier. full_text = f"{system_prompt or ''} {prompt}".lower() user_text = prompt.lower() - + # Estimate tokens estimated_tokens = self._estimate_tokens(prompt) - + # Score all dimensions, capturing match counts where needed code_score, _ = self._score_keyword_match( - full_text, self.code_keywords, "codePresence", "code", - (1, 2), (0, 0.5, 1.0), + full_text, + self.code_keywords, + "codePresence", + "code", + (1, 2), + (0, 0.5, 1.0), ) reasoning_score, reasoning_match_count = self._score_keyword_match( - user_text, self.reasoning_keywords, "reasoningMarkers", "reasoning", - (1, 2), (0, 0.7, 1.0), + user_text, + self.reasoning_keywords, + "reasoningMarkers", + "reasoning", + (1, 2), + (0, 0.7, 1.0), ) technical_score, _ = self._score_keyword_match( - full_text, self.technical_keywords, "technicalTerms", "technical", - (2, 4), (0, 0.5, 1.0), + full_text, + self.technical_keywords, + "technicalTerms", + "technical", + (2, 4), + (0, 0.5, 1.0), ) simple_score, _ = self._score_keyword_match( - full_text, self.simple_keywords, "simpleIndicators", "simple", - (1, 2), (0, -1.0, -1.0), + full_text, + self.simple_keywords, + "simpleIndicators", + "simple", + (1, 2), + (0, -1.0, -1.0), ) dimensions: List[DimensionScore] = [ @@ -265,22 +277,19 @@ class ComplexityRouter(CustomLogger): # Compute weighted score weights = self.config.dimension_weights - weighted_score = sum( - d.score * weights.get(d.name, 0) - for d in dimensions - ) + weighted_score = sum(d.score * weights.get(d.name, 0) for d in dimensions) # Check for reasoning override (2+ reasoning markers) # Reuse match count from _score_keyword_match to avoid scanning twice if reasoning_match_count >= 2: return ComplexityTier.REASONING, weighted_score, signals - + # Map score to tier boundaries = self.config.tier_boundaries simple_medium = boundaries.get("simple_medium", 0.15) medium_complex = boundaries.get("medium_complex", 0.35) complex_reasoning = boundaries.get("complex_reasoning", 0.60) - + if weighted_score < simple_medium: tier = ComplexityTier.SIMPLE elif weighted_score < medium_complex: @@ -289,39 +298,39 @@ class ComplexityRouter(CustomLogger): tier = ComplexityTier.COMPLEX else: tier = ComplexityTier.REASONING - + return tier, weighted_score, signals - + def get_model_for_tier(self, tier: ComplexityTier) -> str: """ Get the model name for a given complexity tier. - + Args: tier: The complexity tier. - + Returns: The model name configured for that tier. """ tier_key = tier.value if isinstance(tier, ComplexityTier) else tier - + # Check config tiers mapping model = self.config.tiers.get(tier_key) if model: return model - + # Fallback to default model if configured if self.config.default_model: return self.config.default_model - + # Last resort: return MEDIUM tier model or error medium_model = self.config.tiers.get(ComplexityTier.MEDIUM.value) if medium_model: return medium_model - + raise ValueError( f"No model configured for tier {tier_key} and no default_model set" ) - + async def async_pre_routing_hook( self, model: str, @@ -332,31 +341,31 @@ class ComplexityRouter(CustomLogger): ) -> Optional["PreRoutingHookResponse"]: """ Pre-routing hook called before the routing decision. - + Classifies the request by complexity and returns the appropriate model. - + Args: model: The original model name requested. request_kwargs: The request kwargs. messages: The messages in the request. input: Optional input for embeddings. specific_deployment: Whether a specific deployment was requested. - + Returns: PreRoutingHookResponse with the routed model, or None if no routing needed. """ from litellm.types.router import PreRoutingHookResponse - + if messages is None or len(messages) == 0: verbose_router_logger.debug( "ComplexityRouter: No messages provided, skipping routing" ) return None - + # Extract the last user message and the last system prompt user_message: Optional[str] = None system_prompt: Optional[str] = None - + for msg in reversed(messages): role = msg.get("role", "") content = msg.get("content") or "" @@ -373,27 +382,28 @@ class ComplexityRouter(CustomLogger): user_message = content elif role == "system" and system_prompt is None: system_prompt = content - + if user_message is None: verbose_router_logger.debug( "ComplexityRouter: No user message found, routing to default model" ) return PreRoutingHookResponse( - model=self.config.default_model or self.get_model_for_tier(ComplexityTier.MEDIUM), + model=self.config.default_model + or self.get_model_for_tier(ComplexityTier.MEDIUM), messages=messages, ) - + # Classify the request tier, score, signals = self.classify(user_message, system_prompt) - + # Get the model for this tier routed_model = self.get_model_for_tier(tier) - + verbose_router_logger.info( f"ComplexityRouter: tier={tier.value}, score={score:.3f}, " f"signals={signals}, routed_model={routed_model}" ) - + return PreRoutingHookResponse( model=routed_model, messages=messages, diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 755f834ac8..a8a21e3f30 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -13,6 +13,7 @@ from pydantic import BaseModel, ConfigDict, Field class ComplexityTier(str, Enum): """Complexity tiers for routing decisions.""" + SIMPLE = "SIMPLE" MEDIUM = "MEDIUM" COMPLEX = "COMPLEX" @@ -24,47 +25,134 @@ class ComplexityTier(str, Enum): # The matching logic uses word boundary detection for single-word keywords. DEFAULT_CODE_KEYWORDS: List[str] = [ - "function", "class", "def", "const", "let", "var", - "import", "export", "return", "async", "await", - "try", "catch", "exception", "error", "debug", - "api", "endpoint", "request", "response", - "database", "sql", "query", "schema", - "algorithm", "implement", "refactor", "optimize", - "python", "javascript", "typescript", "java", "rust", "golang", - "react", "vue", "angular", "node", "docker", "kubernetes", - "git", "commit", "merge", "branch", "pull request", + "function", + "class", + "def", + "const", + "let", + "var", + "import", + "export", + "return", + "async", + "await", + "try", + "catch", + "exception", + "error", + "debug", + "api", + "endpoint", + "request", + "response", + "database", + "sql", + "query", + "schema", + "algorithm", + "implement", + "refactor", + "optimize", + "python", + "javascript", + "typescript", + "java", + "rust", + "golang", + "react", + "vue", + "angular", + "node", + "docker", + "kubernetes", + "git", + "commit", + "merge", + "branch", + "pull request", ] DEFAULT_REASONING_KEYWORDS: List[str] = [ - "step by step", "think through", "let's think", - "reason through", "analyze this", "break down", - "explain your reasoning", "show your work", - "chain of thought", "think carefully", - "consider all", "evaluate", "pros and cons", - "compare and contrast", "weigh the options", - "logical", "deduce", "infer", "conclude", + "step by step", + "think through", + "let's think", + "reason through", + "analyze this", + "break down", + "explain your reasoning", + "show your work", + "chain of thought", + "think carefully", + "consider all", + "evaluate", + "pros and cons", + "compare and contrast", + "weigh the options", + "logical", + "deduce", + "infer", + "conclude", ] DEFAULT_TECHNICAL_KEYWORDS: List[str] = [ - "architecture", "distributed", "scalable", "microservice", - "machine learning", "neural network", "deep learning", - "encryption", "authentication", "authorization", - "performance", "latency", "throughput", "benchmark", - "concurrency", "parallel", "threading", - "memory", "cpu", "gpu", "optimization", - "protocol", "tcp", "http", "grpc", "websocket", - "container", "orchestration", + "architecture", + "distributed", + "scalable", + "microservice", + "machine learning", + "neural network", + "deep learning", + "encryption", + "authentication", + "authorization", + "performance", + "latency", + "throughput", + "benchmark", + "concurrency", + "parallel", + "threading", + "memory", + "cpu", + "gpu", + "optimization", + "protocol", + "tcp", + "http", + "grpc", + "websocket", + "container", + "orchestration", # Note: "async", "kubernetes", "docker" are in DEFAULT_CODE_KEYWORDS ] DEFAULT_SIMPLE_KEYWORDS: List[str] = [ - "what is", "what's", "define", "definition of", - "who is", "who was", "when did", "when was", - "where is", "where was", "how many", "how much", - "yes or no", "true or false", - "simple", "brief", "short", "quick", - "hello", "hi", "hey", "thanks", "thank you", - "goodbye", "bye", "okay", + "what is", + "what's", + "define", + "definition of", + "who is", + "who was", + "when did", + "when was", + "where is", + "where was", + "how many", + "how much", + "yes or no", + "true or false", + "simple", + "brief", + "short", + "quick", + "hello", + "hi", + "hey", + "thanks", + "thank you", + "goodbye", + "bye", + "okay", # Note: "ok" removed due to false positives (matches "token", "book", etc.) ] @@ -72,10 +160,10 @@ DEFAULT_SIMPLE_KEYWORDS: List[str] = [ # ─── Default Dimension Weights ─── DEFAULT_DIMENSION_WEIGHTS: Dict[str, float] = { - "tokenCount": 0.10, # Reduced - length is less important than content - "codePresence": 0.30, # High - code requests need capable models + "tokenCount": 0.10, # Reduced - length is less important than content + "codePresence": 0.30, # High - code requests need capable models "reasoningMarkers": 0.25, # High - explicit reasoning requests - "technicalTerms": 0.25, # High - technical content matters + "technicalTerms": 0.25, # High - technical content matters "simpleIndicators": 0.05, # Low - don't over-penalize simple patterns "multiStepPatterns": 0.03, "questionComplexity": 0.02, @@ -85,8 +173,8 @@ DEFAULT_DIMENSION_WEIGHTS: Dict[str, float] = { # ─── Default Tier Boundaries ─── DEFAULT_TIER_BOUNDARIES: Dict[str, float] = { - "simple_medium": 0.15, # Lower threshold to catch more MEDIUM cases - "medium_complex": 0.35, # Lower threshold to catch technical COMPLEX cases + "simple_medium": 0.15, # Lower threshold to catch more MEDIUM cases + "medium_complex": 0.35, # Lower threshold to catch technical COMPLEX cases "complex_reasoning": 0.60, # Reasoning tier reserved for explicit reasoning markers } @@ -94,7 +182,7 @@ DEFAULT_TIER_BOUNDARIES: Dict[str, float] = { # ─── Default Token Thresholds ─── DEFAULT_TOKEN_THRESHOLDS: Dict[str, int] = { - "simple": 15, # Only very short prompts (<15 tokens) are penalized + "simple": 15, # Only very short prompts (<15 tokens) are penalized "complex": 400, # Long prompts (>400 tokens) get complexity boost } @@ -111,31 +199,31 @@ DEFAULT_TIER_MODELS: Dict[str, str] = { class ComplexityRouterConfig(BaseModel): """Configuration for the ComplexityRouter.""" - + # Tier to model mapping tiers: Dict[str, str] = Field( default_factory=lambda: DEFAULT_TIER_MODELS.copy(), description="Mapping of complexity tiers to model names", ) - + # Tier boundaries (normalized scores) tier_boundaries: Dict[str, float] = Field( default_factory=lambda: DEFAULT_TIER_BOUNDARIES.copy(), description="Score boundaries between tiers", ) - + # Token count thresholds token_thresholds: Dict[str, int] = Field( default_factory=lambda: DEFAULT_TOKEN_THRESHOLDS.copy(), description="Token count thresholds for simple/complex classification", ) - + # Dimension weights dimension_weights: Dict[str, float] = Field( default_factory=lambda: DEFAULT_DIMENSION_WEIGHTS.copy(), description="Weights for each scoring dimension", ) - + # Keyword lists (overridable) code_keywords: Optional[List[str]] = Field( default=None, @@ -153,13 +241,13 @@ class ComplexityRouterConfig(BaseModel): default=None, description="Keywords indicating simple/basic queries", ) - + # Default model if scoring fails default_model: Optional[str] = Field( default=None, description="Default model to use if tier cannot be determined", ) - + model_config = ConfigDict(extra="allow") # Allow additional fields diff --git a/litellm/router_strategy/complexity_router/evals/eval_complexity_router.py b/litellm/router_strategy/complexity_router/evals/eval_complexity_router.py index e671b97c57..a361d95a0a 100644 --- a/litellm/router_strategy/complexity_router/evals/eval_complexity_router.py +++ b/litellm/router_strategy/complexity_router/evals/eval_complexity_router.py @@ -14,7 +14,9 @@ from dataclasses import dataclass from typing import List, Optional, Tuple from unittest.mock import MagicMock -sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../.."))) +sys.path.insert( + 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../..")) +) from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter from litellm.router_strategy.complexity_router.config import ComplexityTier @@ -23,6 +25,7 @@ from litellm.router_strategy.complexity_router.config import ComplexityTier @dataclass class EvalCase: """A single evaluation case.""" + prompt: str expected_tier: ComplexityTier description: str @@ -85,7 +88,6 @@ EVAL_CASES: List[EvalCase] = [ expected_tier=ComplexityTier.SIMPLE, description="Simple time zone question", ), - # === MEDIUM tier cases === EvalCase( prompt="Explain how REST APIs work and when to use them", @@ -117,86 +119,91 @@ EVAL_CASES: List[EvalCase] = [ description="Debugging help", acceptable_tiers=[ComplexityTier.MEDIUM, ComplexityTier.COMPLEX], ), - # === COMPLEX tier cases === EvalCase( prompt="Design a distributed microservice architecture for a high-throughput " - "real-time data processing pipeline with Kubernetes orchestration, " - "implementing proper authentication and encryption protocols", + "real-time data processing pipeline with Kubernetes orchestration, " + "implementing proper authentication and encryption protocols", expected_tier=ComplexityTier.COMPLEX, description="Complex architecture design", acceptable_tiers=[ComplexityTier.COMPLEX, ComplexityTier.REASONING], ), EvalCase( prompt="Write a Python function that implements a binary search tree with " - "insert, delete, and search operations. Include proper error handling " - "and optimize for memory efficiency.", + "insert, delete, and search operations. Include proper error handling " + "and optimize for memory efficiency.", expected_tier=ComplexityTier.COMPLEX, description="Complex coding task", acceptable_tiers=[ComplexityTier.MEDIUM, ComplexityTier.COMPLEX], ), EvalCase( prompt="Explain the differences between TCP and UDP protocols, including " - "use cases for each, performance implications, and how they handle " - "packet loss in distributed systems", + "use cases for each, performance implications, and how they handle " + "packet loss in distributed systems", expected_tier=ComplexityTier.COMPLEX, description="Deep technical explanation", acceptable_tiers=[ComplexityTier.MEDIUM, ComplexityTier.COMPLEX], ), EvalCase( prompt="Create a comprehensive database schema for an e-commerce platform " - "that handles users, products, orders, payments, shipping, reviews, " - "and inventory management with proper indexing strategies", + "that handles users, products, orders, payments, shipping, reviews, " + "and inventory management with proper indexing strategies", expected_tier=ComplexityTier.COMPLEX, description="Complex database design", - acceptable_tiers=[ComplexityTier.MEDIUM, ComplexityTier.COMPLEX, ComplexityTier.REASONING], + acceptable_tiers=[ + ComplexityTier.MEDIUM, + ComplexityTier.COMPLEX, + ComplexityTier.REASONING, + ], ), EvalCase( prompt="Implement a rate limiter using the token bucket algorithm in Python " - "that supports multiple rate limit tiers and can be used across " - "distributed systems with Redis as the backend", + "that supports multiple rate limit tiers and can be used across " + "distributed systems with Redis as the backend", expected_tier=ComplexityTier.COMPLEX, description="Complex distributed systems coding", - acceptable_tiers=[ComplexityTier.MEDIUM, ComplexityTier.COMPLEX, ComplexityTier.REASONING], + acceptable_tiers=[ + ComplexityTier.MEDIUM, + ComplexityTier.COMPLEX, + ComplexityTier.REASONING, + ], ), - # === REASONING tier cases === EvalCase( prompt="Think step by step about how to solve this: A farmer has 17 sheep. " - "All but 9 die. How many are left? Explain your reasoning.", + "All but 9 die. How many are left? Explain your reasoning.", expected_tier=ComplexityTier.REASONING, description="Explicit reasoning request", ), EvalCase( prompt="Let's think through this carefully. Analyze the pros and cons of " - "microservices vs monolithic architecture for a startup with 5 engineers. " - "Consider scalability, development speed, and operational complexity.", + "microservices vs monolithic architecture for a startup with 5 engineers. " + "Consider scalability, development speed, and operational complexity.", expected_tier=ComplexityTier.REASONING, description="Multiple reasoning markers + analysis", ), EvalCase( prompt="Reason through this problem: If I have a function that's O(n^2) and " - "I need to process 1 million items, what are my options to optimize it? " - "Walk me through each approach step by step.", + "I need to process 1 million items, what are my options to optimize it? " + "Walk me through each approach step by step.", expected_tier=ComplexityTier.REASONING, description="Algorithm reasoning", ), EvalCase( prompt="I need you to think carefully and analyze this code for potential " - "security vulnerabilities. Consider injection attacks, authentication " - "bypasses, and data exposure risks. Show your reasoning process.", + "security vulnerabilities. Consider injection attacks, authentication " + "bypasses, and data exposure risks. Show your reasoning process.", expected_tier=ComplexityTier.REASONING, description="Security analysis with reasoning", acceptable_tiers=[ComplexityTier.COMPLEX, ComplexityTier.REASONING], ), EvalCase( prompt="Step by step, explain your reasoning as you evaluate whether we should " - "use PostgreSQL or MongoDB for our new project. Consider our requirements: " - "complex queries, high write volume, and eventual consistency is acceptable.", + "use PostgreSQL or MongoDB for our new project. Consider our requirements: " + "complex queries, high write volume, and eventual consistency is acceptable.", expected_tier=ComplexityTier.REASONING, description="Database decision with explicit reasoning", ), - # === Edge cases / regression tests === EvalCase( prompt="What is the capital of France?", @@ -227,7 +234,7 @@ EVAL_CASES: List[EvalCase] = [ def run_eval() -> Tuple[int, int, List[dict]]: """ Run the evaluation suite. - + Returns: Tuple of (passed, total, failures) """ @@ -237,82 +244,95 @@ def run_eval() -> Tuple[int, int, List[dict]]: model_name="eval-router", litellm_router_instance=mock_router, ) - + passed = 0 total = len(EVAL_CASES) failures = [] - + print("=" * 70) # noqa: T201 print("COMPLEXITY ROUTER EVALUATION") # noqa: T201 print("=" * 70) # noqa: T201 print() # noqa: T201 - + for i, case in enumerate(EVAL_CASES, 1): tier, score, signals = router.classify(case.prompt, case.system_prompt) - + # Check if pass is_exact_match = tier == case.expected_tier is_acceptable = ( - case.acceptable_tiers is not None and - tier in case.acceptable_tiers + case.acceptable_tiers is not None and tier in case.acceptable_tiers ) is_pass = is_exact_match or is_acceptable - + if is_pass: passed += 1 status = "✓ PASS" else: status = "✗ FAIL" - failures.append({ - "case": i, - "description": case.description, - "prompt": case.prompt[:80] + "..." if len(case.prompt) > 80 else case.prompt, - "expected": case.expected_tier.value, - "actual": tier.value, - "score": round(score, 3), - "signals": signals, - "acceptable": [t.value for t in case.acceptable_tiers] if case.acceptable_tiers else None, - }) - + failures.append( + { + "case": i, + "description": case.description, + "prompt": case.prompt[:80] + "..." + if len(case.prompt) > 80 + else case.prompt, + "expected": case.expected_tier.value, + "actual": tier.value, + "score": round(score, 3), + "signals": signals, + "acceptable": [t.value for t in case.acceptable_tiers] + if case.acceptable_tiers + else None, + } + ) + # Print result print(f"[{i:2d}] {status} | {case.description}") # noqa: T201 - print(f" Expected: {case.expected_tier.value:10s} | Got: {tier.value:10s} | Score: {score:+.3f}") # noqa: T201 + print( + f" Expected: {case.expected_tier.value:10s} | Got: {tier.value:10s} | Score: {score:+.3f}" + ) # noqa: T201 if signals: print(f" Signals: {', '.join(signals)}") # noqa: T201 if not is_pass: print(f" Prompt: {case.prompt[:60]}...") # noqa: T201 print() # noqa: T201 - + # Summary print("=" * 70) # noqa: T201 print(f"RESULTS: {passed}/{total} passed ({100*passed/total:.1f}%)") # noqa: T201 print("=" * 70) # noqa: T201 - + if failures: print("\nFAILURES:") # noqa: T201 print("-" * 70) # noqa: T201 for f in failures: print(f"Case {f['case']}: {f['description']}") # noqa: T201 - print(f" Expected: {f['expected']}, Got: {f['actual']} (score: {f['score']})") # noqa: T201 + print( + f" Expected: {f['expected']}, Got: {f['actual']} (score: {f['score']})" + ) # noqa: T201 print(f" Signals: {f['signals']}") # noqa: T201 - if f['acceptable']: + if f["acceptable"]: print(f" Acceptable: {f['acceptable']}") # noqa: T201 print() # noqa: T201 - + return passed, total, failures def main(): """Main entry point.""" passed, total, failures = run_eval() - + # Exit with error code if too many failures pass_rate = passed / total if pass_rate < 0.80: - print(f"\n❌ EVAL FAILED: Pass rate {pass_rate:.1%} is below 80% threshold") # noqa: T201 + print( + f"\n❌ EVAL FAILED: Pass rate {pass_rate:.1%} is below 80% threshold" + ) # noqa: T201 sys.exit(1) elif pass_rate < 0.90: - print(f"\n⚠️ EVAL WARNING: Pass rate {pass_rate:.1%} is below 90%") # noqa: T201 + print( + f"\n⚠️ EVAL WARNING: Pass rate {pass_rate:.1%} is below 90%" + ) # noqa: T201 sys.exit(0) else: print(f"\n✅ EVAL PASSED: Pass rate {pass_rate:.1%}") # noqa: T201 diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index ae0f8433d8..e161438837 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -21,7 +21,6 @@ class LeastBusyLoggingHandler(CustomLogger): def __init__(self, router_cache: DualCache): self.router_cache = router_cache - def log_pre_api_call(self, model, messages, kwargs): """ Log when a model is being used. diff --git a/litellm/router_strategy/lowest_cost.py b/litellm/router_strategy/lowest_cost.py index b0612069df..54498363f5 100644 --- a/litellm/router_strategy/lowest_cost.py +++ b/litellm/router_strategy/lowest_cost.py @@ -15,9 +15,7 @@ class LowestCostLoggingHandler(CustomLogger): logged_success: int = 0 logged_failure: int = 0 - def __init__( - self, router_cache: DualCache, routing_args: dict = {} - ): + def __init__(self, router_cache: DualCache, routing_args: dict = {}): self.router_cache = router_cache def log_success_event(self, kwargs, response_obj, start_time, end_time): diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py index 0449a843bd..20db28fa10 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -31,9 +31,7 @@ class LowestLatencyLoggingHandler(CustomLogger): logged_success: int = 0 logged_failure: int = 0 - def __init__( - self, router_cache: DualCache, routing_args: dict = {} - ): + def __init__(self, router_cache: DualCache, routing_args: dict = {}): self.router_cache = router_cache self.routing_args = RoutingArgs(**routing_args) @@ -96,14 +94,16 @@ class LowestLatencyLoggingHandler(CustomLogger): if _usage is not None: completion_tokens = _usage.completion_tokens total_tokens = _usage.total_tokens - + # Handle both timedelta and float response times if isinstance(response_ms, timedelta): response_seconds = response_ms.total_seconds() else: response_seconds = response_ms - - final_value = safe_divide_seconds(response_seconds, completion_tokens) + + final_value = safe_divide_seconds( + response_seconds, completion_tokens + ) if final_value is not None: final_value = float(final_value) else: @@ -111,7 +111,9 @@ class LowestLatencyLoggingHandler(CustomLogger): if time_to_first_token_response_time is not None: if isinstance(time_to_first_token_response_time, timedelta): - ttft_seconds = time_to_first_token_response_time.total_seconds() + ttft_seconds = ( + time_to_first_token_response_time.total_seconds() + ) else: ttft_seconds = time_to_first_token_response_time time_to_first_token = safe_divide_seconds( @@ -204,7 +206,9 @@ class LowestLatencyLoggingHandler(CustomLogger): "model_group", None ) - id = (kwargs["litellm_params"].get("model_info") or {}).get("id", None) + id = (kwargs["litellm_params"].get("model_info") or {}).get( + "id", None + ) if model_group is None or id is None: return elif isinstance(id, int): @@ -317,14 +321,16 @@ class LowestLatencyLoggingHandler(CustomLogger): if _usage is not None: completion_tokens = _usage.completion_tokens total_tokens = _usage.total_tokens - + # Handle both timedelta and float response times if isinstance(response_ms, timedelta): response_seconds = response_ms.total_seconds() else: response_seconds = response_ms - - final_value = safe_divide_seconds(response_seconds, completion_tokens) + + final_value = safe_divide_seconds( + response_seconds, completion_tokens + ) if final_value is not None: final_value = float(final_value) else: @@ -332,7 +338,9 @@ class LowestLatencyLoggingHandler(CustomLogger): if time_to_first_token_response_time is not None: if isinstance(time_to_first_token_response_time, timedelta): - ttft_seconds = time_to_first_token_response_time.total_seconds() + ttft_seconds = ( + time_to_first_token_response_time.total_seconds() + ) else: ttft_seconds = time_to_first_token_response_time time_to_first_token = safe_divide_seconds( @@ -490,20 +498,22 @@ class LowestLatencyLoggingHandler(CustomLogger): # get average latency or average ttft (depending on streaming/non-streaming) total: float = 0.0 - if ( + use_ttft = ( request_kwargs is not None and request_kwargs.get("stream", None) is not None and request_kwargs["stream"] is True and len(item_ttft_latency) > 0 - ): + ) + if use_ttft: for _call_latency in item_ttft_latency: if isinstance(_call_latency, float): total += _call_latency + item_latency = total / len(item_ttft_latency) else: for _call_latency in item_latency: if isinstance(_call_latency, float): total += _call_latency - item_latency = total / len(item_latency) + item_latency = total / len(item_latency) # -------------- # # Debugging Logic diff --git a/litellm/router_strategy/lowest_tpm_rpm.py b/litellm/router_strategy/lowest_tpm_rpm.py index 93d3c8e041..488f845094 100644 --- a/litellm/router_strategy/lowest_tpm_rpm.py +++ b/litellm/router_strategy/lowest_tpm_rpm.py @@ -22,9 +22,7 @@ class LowestTPMLoggingHandler(CustomLogger): logged_failure: int = 0 default_cache_time_seconds: int = 1 * 60 * 60 # 1 hour - def __init__( - self, router_cache: DualCache, routing_args: dict = {} - ): + def __init__(self, router_cache: DualCache, routing_args: dict = {}): self.router_cache = router_cache self.routing_args = RoutingArgs(**routing_args) diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index 70d4c6751d..23e8896cd5 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -47,9 +47,7 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): logged_failure: int = 0 default_cache_time_seconds: int = 1 * 60 * 60 # 1 hour - def __init__( - self, router_cache: DualCache, routing_args: dict = {} - ): + def __init__(self, router_cache: DualCache, routing_args: dict = {}): self.router_cache = router_cache self.routing_args = RoutingArgs(**routing_args) BaseRoutingStrategy.__init__( diff --git a/litellm/router_strategy/simple_shuffle.py b/litellm/router_strategy/simple_shuffle.py index ca82ddc6aa..9827522747 100644 --- a/litellm/router_strategy/simple_shuffle.py +++ b/litellm/router_strategy/simple_shuffle.py @@ -43,7 +43,9 @@ def simple_shuffle( for weight_by in ["weight", "rpm", "tpm"]: weight = healthy_deployments[0].get("litellm_params").get(weight_by, None) if weight is not None: - weights = [m["litellm_params"].get(weight_by, 0) for m in healthy_deployments] + weights = [ + m["litellm_params"].get(weight_by, 0) for m in healthy_deployments + ] verbose_router_logger.debug(f"\nweight {weights}") total_weight = sum(weights) weights = [weight / total_weight for weight in weights] @@ -57,7 +59,6 @@ def simple_shuffle( ) return deployment or deployment[0] - ############## No RPM/TPM passed, we do a random pick ################# item = random.choice(healthy_deployments) return item or item[0] diff --git a/litellm/router_strategy/tag_based_routing.py b/litellm/router_strategy/tag_based_routing.py index e960e00a68..e7156bf128 100644 --- a/litellm/router_strategy/tag_based_routing.py +++ b/litellm/router_strategy/tag_based_routing.py @@ -75,7 +75,9 @@ async def get_deployments_for_tag( ) return healthy_deployments - verbose_logger.debug("request metadata: %s", request_kwargs.get(metadata_variable_name)) + verbose_logger.debug( + "request metadata: %s", request_kwargs.get(metadata_variable_name) + ) if metadata_variable_name in request_kwargs: metadata = request_kwargs[metadata_variable_name] request_tags = metadata.get("tags") @@ -112,7 +114,11 @@ async def get_deployments_for_tag( f"{RouterErrors.no_deployments_with_tag_routing.value}. Passed model={model} and tags={request_tags}" ) - return new_healthy_deployments if len(new_healthy_deployments) > 0 else default_deployments + return ( + new_healthy_deployments + if len(new_healthy_deployments) > 0 + else default_deployments + ) # for Untagged requests use default deployments if set _default_deployments_with_tags = [] diff --git a/litellm/router_utils/batch_utils.py b/litellm/router_utils/batch_utils.py index aef8f5cc01..5e58479825 100644 --- a/litellm/router_utils/batch_utils.py +++ b/litellm/router_utils/batch_utils.py @@ -97,7 +97,6 @@ def replace_model_in_jsonl(file_content: FileTypes, new_model_name: str) -> File elif isinstance(file_content_bytes, str): file_content_str = file_content_bytes else: - return file_content # Parse JSONL properly, handling potential multiline JSON objects diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 3b0273f4c5..7530247ce7 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -32,9 +32,9 @@ def add_model_file_id_mappings( model_file_id_mapping = {} if isinstance(healthy_deployments, list): for deployment, response in zip(healthy_deployments, responses): - model_file_id_mapping[deployment.get("model_info", {}).get("id")] = ( - response.id - ) + model_file_id_mapping[ + deployment.get("model_info", {}).get("id") + ] = response.id elif isinstance(healthy_deployments, dict): for model_id, file_id in healthy_deployments.items(): model_file_id_mapping[model_id] = file_id @@ -75,6 +75,7 @@ def filter_team_based_models( if deployment.get("model_info", {}).get("id") not in ids_to_remove ] + def _deployment_supports_web_search(deployment: Dict) -> bool: """ Check if a deployment supports web search. @@ -112,7 +113,7 @@ def filter_web_search_deployments( is_web_search_request = False tools = request_kwargs.get("tools") or [] for tool in tools: - # These are the two websearch tools for OpenAI / Azure. + # These are the two websearch tools for OpenAI / Azure. if tool.get("type") == "web_search" or tool.get("type") == "web_search_preview": is_web_search_request = True break @@ -121,7 +122,9 @@ def filter_web_search_deployments( return healthy_deployments # Filter out deployments that don't support web search - final_deployments = [d for d in healthy_deployments if _deployment_supports_web_search(d)] + final_deployments = [ + d for d in healthy_deployments if _deployment_supports_web_search(d) + ] if len(healthy_deployments) > 0 and len(final_deployments) == 0: verbose_logger.warning("No deployments support web search for request") return final_deployments diff --git a/litellm/router_utils/cooldown_cache.py b/litellm/router_utils/cooldown_cache.py index edbcacca27..b210ea4459 100644 --- a/litellm/router_utils/cooldown_cache.py +++ b/litellm/router_utils/cooldown_cache.py @@ -129,7 +129,7 @@ class CooldownCache: if results is None or all(v is None for v in results): return active_cooldowns - + # Process the results for model_id, result in zip(model_ids, results): if result and isinstance(result, dict): @@ -142,7 +142,9 @@ class CooldownCache: self, model_ids: List[str], parent_otel_span: Optional[Span] ) -> List[Tuple[str, CooldownCacheValue]]: # Generate the keys for the deployments - keys = [CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids] + keys = [ + CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids + ] # Retrieve the values for the keys using mget results = ( self.cache.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) diff --git a/litellm/router_utils/cooldown_callbacks.py b/litellm/router_utils/cooldown_callbacks.py index 343328dacf..32777a1dd4 100644 --- a/litellm/router_utils/cooldown_callbacks.py +++ b/litellm/router_utils/cooldown_callbacks.py @@ -59,9 +59,9 @@ async def router_cooldown_event_callback( pass # get the prometheus logger from in memory loggers - prometheusLogger: Optional[PrometheusLogger] = ( - _get_prometheus_logger_from_callbacks() - ) + prometheusLogger: Optional[ + PrometheusLogger + ] = _get_prometheus_logger_from_callbacks() if prometheusLogger is not None: prometheusLogger.set_deployment_complete_outage( diff --git a/litellm/router_utils/prompt_caching_cache.py b/litellm/router_utils/prompt_caching_cache.py index 69698f282b..eab342e540 100644 --- a/litellm/router_utils/prompt_caching_cache.py +++ b/litellm/router_utils/prompt_caching_cache.py @@ -53,30 +53,32 @@ class PromptCachingCache: return str(obj) @staticmethod - def extract_cacheable_prefix(messages: List[AllMessageValues]) -> List[AllMessageValues]: + def extract_cacheable_prefix( + messages: List[AllMessageValues], + ) -> List[AllMessageValues]: """ Extract the cacheable prefix from messages. - + The cacheable prefix is everything UP TO AND INCLUDING the LAST content block (across all messages) that has cache_control. This includes ALL blocks before the last cacheable block (even if they don't have cache_control). - + Args: messages: List of messages to extract cacheable prefix from - + Returns: List of messages containing only the cacheable prefix """ if not messages: return messages - + # Find the last content block (across all messages) that has cache_control last_cacheable_message_idx = None last_cacheable_content_idx = None - + for msg_idx, message in enumerate(messages): content = message.get("content") - + # Check for cache_control at message level (when content is a string) # This handles the case where cache_control is a sibling of string content: # {"role": "user", "content": "...", "cache_control": {"type": "ephemeral"}} @@ -90,11 +92,11 @@ class PromptCachingCache: # Set to None to indicate the entire message content is cacheable # (not a specific content block index within a list) last_cacheable_content_idx = None - + # Also check for cache_control within content blocks (when content is a list) if not isinstance(content, list): continue - + for content_idx, content_block in enumerate(content): if isinstance(content_block, dict): cache_control = content_block.get("cache_control") @@ -105,14 +107,14 @@ class PromptCachingCache: ): last_cacheable_message_idx = msg_idx last_cacheable_content_idx = content_idx - + # If no cacheable block found, return empty list (no cacheable prefix) if last_cacheable_message_idx is None: return [] - + # Build the cacheable prefix: all messages up to and including the last cacheable message cacheable_prefix = [] - + for msg_idx, message in enumerate(messages): if msg_idx < last_cacheable_message_idx: # Include entire message (comes before last cacheable block) @@ -124,7 +126,10 @@ class PromptCachingCache: # Create a copy of the message with only cacheable content blocks message_copy = cast( AllMessageValues, - {**message, "content": content[: last_cacheable_content_idx + 1]}, + { + **message, + "content": content[: last_cacheable_content_idx + 1], + }, ) cacheable_prefix.append(message_copy) else: @@ -133,7 +138,7 @@ class PromptCachingCache: else: # Message comes after last cacheable block, don't include break - + return cacheable_prefix @staticmethod @@ -143,7 +148,7 @@ class PromptCachingCache: ) -> Optional[str]: if messages is None and tools is None: return None - + # Extract cacheable prefix from messages (only include up to last cache_control block) cacheable_messages = None if messages is not None: @@ -151,11 +156,13 @@ class PromptCachingCache: # If no cacheable prefix found, return None (can't cache) if not cacheable_messages: return None - + # Use serialize_object for consistent and stable serialization data_to_hash = {} if cacheable_messages is not None: - serialized_messages = PromptCachingCache.serialize_object(cacheable_messages) + serialized_messages = PromptCachingCache.serialize_object( + cacheable_messages + ) data_to_hash["messages"] = serialized_messages if tools is not None: serialized_tools = PromptCachingCache.serialize_object(tools) @@ -219,7 +226,7 @@ class PromptCachingCache: ) -> Optional[PromptCachingCacheValue]: """ Get model ID from cache using the cacheable prefix. - + The cache key is based on the cacheable prefix (everything up to and including the last cache_control block), so requests with the same cacheable prefix but different user messages will have the same cache key. diff --git a/litellm/router_utils/search_api_router.py b/litellm/router_utils/search_api_router.py index 9a1907fa55..491a25e58e 100644 --- a/litellm/router_utils/search_api_router.py +++ b/litellm/router_utils/search_api_router.py @@ -16,7 +16,7 @@ from litellm._logging import verbose_router_logger class SearchAPIRouter: """ Static utility class for routing search API calls through the LiteLLM router. - + Provides methods for search tool selection, load balancing, and fallback handling. """ @@ -24,18 +24,20 @@ class SearchAPIRouter: async def update_router_search_tools(router_instance: Any, search_tools: list): """ Update the router with search tools from the database. - + This method is called by a cron job to sync search tools from DB to router. - + Args: router_instance: The Router instance to update search_tools: List of search tool configurations from the database """ try: from litellm.types.router import SearchToolTypedDict - - verbose_router_logger.debug(f"Adding {len(search_tools)} search tools to router") - + + verbose_router_logger.debug( + f"Adding {len(search_tools)} search tools to router" + ) + # Convert search tools to the format expected by the router router_search_tools: list = [] for tool in search_tools: @@ -47,14 +49,14 @@ class SearchAPIRouter: "search_tool_info": tool.get("search_tool_info"), } router_search_tools.append(router_search_tool) - + # Update the router's search_tools list router_instance.search_tools = router_search_tools - + verbose_router_logger.info( f"Successfully updated router with {len(router_search_tools)} search tool(s)" ) - + except Exception as e: verbose_router_logger.exception( f"Error updating router with search tools: {str(e)}" @@ -68,25 +70,28 @@ class SearchAPIRouter: ) -> list: """ Get all search tools matching the given name. - + Args: router_instance: The Router instance search_tool_name: Name of the search tool to find - + Returns: List of matching search tool configurations - + Raises: ValueError: If no matching search tools are found """ matching_tools = [ - tool for tool in router_instance.search_tools + tool + for tool in router_instance.search_tools if tool.get("search_tool_name") == search_tool_name ] - + if not matching_tools: - raise ValueError(f"Search tool '{search_tool_name}' not found in router.search_tools") - + raise ValueError( + f"Search tool '{search_tool_name}' not found in router.search_tools" + ) + return matching_tools @staticmethod @@ -98,47 +103,55 @@ class SearchAPIRouter: """ Helper function to make a search API call through the router with load balancing and fallbacks. Reuses the router's retry/fallback infrastructure. - + Args: router_instance: The Router instance original_function: The original litellm.asearch function **kwargs: Search parameters including search_tool_name, query, etc. - + Returns: SearchResponse from the search API """ try: search_tool_name = kwargs.get("search_tool_name", kwargs.get("model")) - + if not search_tool_name: - raise ValueError("search_tool_name or model parameter is required for search") - + raise ValueError( + "search_tool_name or model parameter is required for search" + ) + # Set up kwargs for the fallback system - kwargs["model"] = search_tool_name # Use model field for compatibility with fallback system + kwargs[ + "model" + ] = search_tool_name # Use model field for compatibility with fallback system kwargs["original_generic_function"] = original_function # Bind router_instance to the helper method using partial kwargs["original_function"] = partial( SearchAPIRouter.async_search_with_fallbacks_helper, router_instance=router_instance, ) - + # Update kwargs before fallbacks (for logging, metadata, etc) router_instance._update_kwargs_before_fallbacks( - model=search_tool_name, kwargs=kwargs, metadata_variable_name="litellm_metadata" + model=search_tool_name, + kwargs=kwargs, + metadata_variable_name="litellm_metadata", ) - - available_search_tool_names = [tool.get("search_tool_name") for tool in router_instance.search_tools] + + available_search_tool_names = [ + tool.get("search_tool_name") for tool in router_instance.search_tools + ] verbose_router_logger.debug( f"Inside SearchAPIRouter.async_search_with_fallbacks() - search_tool_name: {search_tool_name}, Available Search Tools: {available_search_tool_names}, kwargs: {kwargs}" ) - + # Use the existing retry/fallback infrastructure response = await router_instance.async_function_with_fallbacks(**kwargs) return response - + except Exception as e: from litellm.router_utils.handle_error import send_llm_exception_alert - + asyncio.create_task( send_llm_exception_alert( litellm_router_instance=router_instance, @@ -148,7 +161,7 @@ class SearchAPIRouter: ) ) raise e - + @staticmethod async def async_search_with_fallbacks_helper( router_instance: Any, @@ -159,42 +172,44 @@ class SearchAPIRouter: """ Helper function for search API calls - selects a search tool and calls the original function. Called by async_function_with_fallbacks for each retry attempt. - + Args: router_instance: The Router instance model: The search tool name (passed as model for compatibility) original_generic_function: The original litellm.asearch function **kwargs: Search parameters - + Returns: SearchResponse from the selected search provider """ search_tool_name = model # model field contains the search_tool_name - + try: # Find matching search tools matching_tools = SearchAPIRouter.get_matching_search_tools( router_instance=router_instance, search_tool_name=search_tool_name, ) - + # Simple random selection for load balancing across multiple providers with same name # For search tools, we use simple random choice since they don't have TPM/RPM constraints selected_tool = random.choice(matching_tools) - + # Extract search provider and other params from litellm_params litellm_params = selected_tool.get("litellm_params", {}) search_provider = litellm_params.get("search_provider") api_key = litellm_params.get("api_key") api_base = litellm_params.get("api_base") - + if not search_provider: - raise ValueError(f"search_provider not found in litellm_params for search tool '{search_tool_name}'") - + raise ValueError( + f"search_provider not found in litellm_params for search tool '{search_tool_name}'" + ) + verbose_router_logger.debug( f"Selected search tool with provider: {search_provider}" ) - + # Call the original search function with the provider config response = await original_generic_function( search_provider=search_provider, @@ -202,12 +217,11 @@ class SearchAPIRouter: api_base=api_base, **kwargs, ) - + return response - + except Exception as e: verbose_router_logger.error( f"Error in SearchAPIRouter.async_search_with_fallbacks_helper for {search_tool_name}: {str(e)}" ) raise e - diff --git a/litellm/scheduler.py b/litellm/scheduler.py index 0221e24984..5309971eed 100644 --- a/litellm/scheduler.py +++ b/litellm/scheduler.py @@ -101,7 +101,9 @@ class Scheduler: filtered_queue = [item for item in queue if item[1] != request_id] heapq.heapify(filtered_queue) # restore heap invariant after filtering await self.save_queue(queue=filtered_queue, model_name=model_name) - print_verbose(f"Removed request_id: {request_id} from queue for model: {model_name}") + print_verbose( + f"Removed request_id: {request_id} from queue for model: {model_name}" + ) async def peek(self, id: str, model_name: str, health_deployments: list) -> bool: """Return if the id is at the top of the queue. Don't pop the value from heap.""" diff --git a/litellm/search/__init__.py b/litellm/search/__init__.py index a91dff7060..a3ebb3d870 100644 --- a/litellm/search/__init__.py +++ b/litellm/search/__init__.py @@ -5,4 +5,3 @@ from litellm.search.cost_calculator import search_provider_cost_per_query from litellm.search.main import asearch, search __all__ = ["search", "asearch", "search_provider_cost_per_query"] - diff --git a/litellm/search/cost_calculator.py b/litellm/search/cost_calculator.py index 1dc155d748..9821c12ae4 100644 --- a/litellm/search/cost_calculator.py +++ b/litellm/search/cost_calculator.py @@ -14,27 +14,29 @@ def search_provider_cost_per_query( ) -> Tuple[float, float]: """ Calculate cost for search-only providers. - + Returns (input_cost, output_cost) where input_cost = queries * cost_per_query Supports tiered pricing based on max_results parameter. - + Args: model: Model name (e.g., "exa_ai/search", "tavily/search") custom_llm_provider: Provider name (e.g., "exa_ai", "tavily") number_of_queries: Number of search queries performed (default: 1) optional_params: Optional parameters including max_results for tiered pricing - + Returns: Tuple of (input_cost, output_cost) where output_cost is always 0.0 """ model_info = get_model_info(model=model, custom_llm_provider=custom_llm_provider) - + # Check for tiered pricing (e.g., Exa AI based on max_results) tiered_pricing = model_info.get("tiered_pricing") if tiered_pricing and isinstance(tiered_pricing, list): - max_results = (optional_params or {}).get("max_results", 10) # default 10 results + max_results = (optional_params or {}).get( + "max_results", 10 + ) # default 10 results cost_per_query = 0.0 - + for tier in tiered_pricing: range_min, range_max = tier["max_results_range"] if range_min <= max_results <= range_max: @@ -46,7 +48,6 @@ def search_provider_cost_per_query( else: # Simple flat rate cost_per_query = float(model_info.get("input_cost_per_query") or 0.0) - + total_cost = number_of_queries * cost_per_query return (total_cost, 0.0) # (input_cost, output_cost) - diff --git a/litellm/search/main.py b/litellm/search/main.py index c87694e70f..fdfe9a1fee 100644 --- a/litellm/search/main.py +++ b/litellm/search/main.py @@ -30,18 +30,18 @@ def _build_search_optional_params( ) -> Dict[str, Any]: """ Helper function to build optional_params dict from Perplexity Search API parameters. - + Args: max_results: Maximum number of results (1-20) search_domain_filter: List of domains to filter (max 20) max_tokens_per_page: Max tokens per page country: Country code filter - + Returns: Dict with non-None optional parameters """ optional_params: Dict[str, Any] = {} - + if max_results is not None: optional_params["max_results"] = max_results if search_domain_filter is not None: @@ -50,7 +50,7 @@ def _build_search_optional_params( optional_params["max_tokens_per_page"] = max_tokens_per_page if country is not None: optional_params["country"] = country - + return optional_params @@ -70,7 +70,7 @@ async def asearch( ) -> SearchResponse: """ Async Search function. - + Args: query: Search query (string or list of strings) search_provider: Provider name (e.g., "perplexity") @@ -83,20 +83,20 @@ async def asearch( timeout: Optional timeout extra_headers: Optional extra headers **kwargs: Additional parameters - + Returns: SearchResponse with results list following Perplexity format - + Example: ```python import litellm - + # Basic search response = await litellm.asearch( query="latest AI developments 2024", search_provider="perplexity" ) - + # Search with options response = await litellm.asearch( query="AI developments", @@ -106,7 +106,7 @@ async def asearch( max_tokens_per_page=1024, country="US" ) - + # Access results for result in response.results: print(f"{result.title}: {result.url}") @@ -175,7 +175,7 @@ def search( ) -> Union[SearchResponse, Coroutine[Any, Any, SearchResponse]]: """ Synchronous Search function. - + Args: query: Search query (string or list of strings) search_provider: Provider name (e.g., "perplexity") @@ -188,20 +188,20 @@ def search( timeout: Optional timeout extra_headers: Optional extra headers **kwargs: Additional parameters - + Returns: SearchResponse with results list following Perplexity format - + Example: ```python import litellm - + # Basic search response = litellm.search( query="latest AI developments 2024", search_provider="perplexity" ) - + # Search with options response = litellm.search( query="AI developments", @@ -211,13 +211,13 @@ def search( max_tokens_per_page=1024, country="US" ) - + # Multi-query search response = litellm.search( query=["AI developments", "machine learning trends"], search_provider="perplexity" ) - + # Access results for result in response.results: print(f"{result.title}: {result.url}") @@ -231,29 +231,27 @@ def search( litellm_logging_obj: LiteLLMLoggingObj = kwargs.pop("litellm_logging_obj") # type: ignore litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) _is_async = kwargs.pop("asearch", False) is True - + # Validate query parameter if not isinstance(query, (str, list)): - raise ValueError(f"query must be a string or list of strings, got {type(query)}") - + raise ValueError( + f"query must be a string or list of strings, got {type(query)}" + ) + if isinstance(query, list) and not all(isinstance(q, str) for q in query): raise ValueError("All items in query list must be strings") # Get provider config - search_provider_config: Optional[BaseSearchConfig] = ( - ProviderConfigManager.get_provider_search_config( - provider=SearchProviders(search_provider), - ) + search_provider_config: Optional[ + BaseSearchConfig + ] = ProviderConfigManager.get_provider_search_config( + provider=SearchProviders(search_provider), ) if search_provider_config is None: - raise ValueError( - f"Search is not supported for provider: {search_provider}" - ) + raise ValueError(f"Search is not supported for provider: {search_provider}") - verbose_logger.debug( - f"Search call - provider: {search_provider}" - ) + verbose_logger.debug(f"Search call - provider: {search_provider}") # Build optional_params from explicit parameters optional_params = _build_search_optional_params( @@ -262,15 +260,15 @@ def search( max_tokens_per_page=max_tokens_per_page, country=country, ) - + # Filter out internal LiteLLM parameters from kwargs filtered_kwargs = filter_out_litellm_params(kwargs=kwargs) - + # Add remaining kwargs to optional_params (for provider-specific params) for key, value in filtered_kwargs.items(): if key not in optional_params: optional_params[key] = value - + verbose_logger.debug(f"Search optional_params: {optional_params}") # Validate environment and get headers @@ -322,4 +320,3 @@ def search( completion_kwargs=local_vars, extra_kwargs=kwargs, ) - diff --git a/litellm/secret_managers/base_secret_manager.py b/litellm/secret_managers/base_secret_manager.py index af77c5f45b..32a244c5ea 100644 --- a/litellm/secret_managers/base_secret_manager.py +++ b/litellm/secret_managers/base_secret_manager.py @@ -59,7 +59,7 @@ class BaseSecretManager(ABC): description: Optional[str] = None, optional_params: Optional[dict] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, - tags: Optional[Union[dict, list]] = None + tags: Optional[Union[dict, list]] = None, ) -> Dict[str, Any]: """ Asynchronously write a secret to the secret manager. diff --git a/litellm/secret_managers/custom_secret_manager_loader.py b/litellm/secret_managers/custom_secret_manager_loader.py index bbd0e78686..8405740062 100644 --- a/litellm/secret_managers/custom_secret_manager_loader.py +++ b/litellm/secret_managers/custom_secret_manager_loader.py @@ -17,38 +17,38 @@ from litellm.types.secret_managers.main import KeyManagementSystem def load_custom_secret_manager(config_file_path: Optional[str] = None) -> None: """ Load and initialize a custom secret manager from a python file. - + Similar to how custom guardrails are loaded - loads the class from the custom_secret_manager field in key_management_settings. - + Args: config_file_path: Path to the config.yaml file - + Raises: ValueError: If required configuration is missing ImportError: If the custom secret manager module cannot be loaded """ - + if not config_file_path: raise ValueError( "CustomSecretManagerException - config_file_path is required to load custom secret manager" ) - + # Get the custom_secret_manager class path from settings if litellm._key_management_settings is None: raise ValueError( "CustomSecretManagerException - key_management_settings is required with custom_secret_manager field" ) - + custom_secret_manager_path = getattr( litellm._key_management_settings, "custom_secret_manager", None ) - + if not custom_secret_manager_path: raise ValueError( "CustomSecretManagerException - custom_secret_manager field is required in key_management_settings" ) - + # Split into file_name and class_name (e.g., "my_secret_manager.InMemorySecretManager") _file_name, _class_name = custom_secret_manager_path.split(".") verbose_proxy_logger.debug( @@ -57,38 +57,37 @@ def load_custom_secret_manager(config_file_path: Optional[str] = None) -> None: _file_name, _class_name, ) - + # Load the module from the same directory as config.yaml directory = os.path.dirname(config_file_path) module_file_path = os.path.join(directory, _file_name) + ".py" - + spec = importlib.util.spec_from_file_location(_class_name, module_file_path) # type: ignore if not spec: raise ImportError( f"Could not find a module specification for {module_file_path}" ) - + module = importlib.util.module_from_spec(spec) # type: ignore spec.loader.exec_module(module) # type: ignore _secret_manager_class = getattr(module, _class_name) - + # Validate that it's a CustomSecretManager subclass if not issubclass(_secret_manager_class, CustomSecretManager): raise TypeError( f"CustomSecretManagerException - {_class_name} must be a subclass of CustomSecretManager" ) - + # Instantiate the custom secret manager _secret_manager_instance = _secret_manager_class() - + # Set it as the secret manager client litellm.secret_manager_client = _secret_manager_instance - + # Set the key management system to CUSTOM so get_secret knows to use it litellm._key_management_system = KeyManagementSystem.CUSTOM - + verbose_proxy_logger.info( "Successfully initialized custom secret manager: %s", custom_secret_manager_path, ) - diff --git a/litellm/secret_managers/cyberark_secret_manager.py b/litellm/secret_managers/cyberark_secret_manager.py index 2745df0077..11b853412f 100644 --- a/litellm/secret_managers/cyberark_secret_manager.py +++ b/litellm/secret_managers/cyberark_secret_manager.py @@ -39,9 +39,7 @@ class CyberArkSecretManager(BaseSecretManager): self.ssl_verify: bool = ssl_verify_env if ssl_verify_env is not None else True # Validate environment - if not self.conjur_api_key and not ( - self.tls_cert_path and self.tls_key_path - ): + if not self.conjur_api_key and not (self.tls_cert_path and self.tls_key_path): raise ValueError( "Missing CyberArk credentials. Please set CYBERARK_API_KEY or both CYBERARK_CLIENT_CERT and CYBERARK_CLIENT_KEY in your environment." ) @@ -318,7 +316,6 @@ class CyberArkSecretManager(BaseSecretManager): verbose_logger.exception(f"Error writing secret to CyberArk Conjur: {e}") return {"status": "error", "message": str(e)} - async def async_delete_secret( self, secret_name: str, @@ -351,4 +348,3 @@ class CyberArkSecretManager(BaseSecretManager): "status": "not_supported", "message": "CyberArk Conjur does not support direct secret deletion. Use policy updates to remove variables.", } - diff --git a/litellm/secret_managers/hashicorp_secret_manager.py b/litellm/secret_managers/hashicorp_secret_manager.py index ccee5018ee..8bb3f801a1 100644 --- a/litellm/secret_managers/hashicorp_secret_manager.py +++ b/litellm/secret_managers/hashicorp_secret_manager.py @@ -470,7 +470,9 @@ class HashicorpSecretManager(BaseSecretManager): try: # First verify the old secret exists using _build_secret_target - current_target = self._build_secret_target(current_secret_name, optional_params) + current_target = self._build_secret_target( + current_secret_name, optional_params + ) try: response = await async_client.get( url=current_target["url"], @@ -480,8 +482,13 @@ class HashicorpSecretManager(BaseSecretManager): # Secret exists, we can proceed except httpx.HTTPStatusError as e: if e.response.status_code == 404: - verbose_logger.exception(f"Current secret {current_secret_name} not found") - return {"status": "error", "message": f"Current secret {current_secret_name} not found"} + verbose_logger.exception( + f"Current secret {current_secret_name} not found" + ) + return { + "status": "error", + "message": f"Current secret {current_secret_name} not found", + } verbose_logger.exception( f"Error checking current secret existence: {e.response.text if hasattr(e, 'response') else str(e)}" ) @@ -490,8 +497,13 @@ class HashicorpSecretManager(BaseSecretManager): "message": f"HTTP error occurred while checking current secret: {e.response.text if hasattr(e, 'response') else str(e)}", } except Exception as e: - verbose_logger.exception(f"Error checking current secret existence: {e}") - return {"status": "error", "message": f"Error checking current secret: {e}"} + verbose_logger.exception( + f"Error checking current secret existence: {e}" + ) + return { + "status": "error", + "message": f"Error checking current secret: {e}", + } # Create new secret with new name and value # Use _build_secret_target to handle optional_params @@ -504,7 +516,10 @@ class HashicorpSecretManager(BaseSecretManager): ) # Check if async_write_secret returned an error - if isinstance(create_response, dict) and create_response.get("status") == "error": + if ( + isinstance(create_response, dict) + and create_response.get("status") == "error" + ): return create_response # Verify new secret was created successfully using _build_secret_target @@ -518,7 +533,9 @@ class HashicorpSecretManager(BaseSecretManager): json_resp = response.json() # Use data_key from target to get the correct value data_key = new_target["data_key"] - new_secret_value_from_vault = json_resp.get("data", {}).get("data", {}).get(data_key, None) + new_secret_value_from_vault = ( + json_resp.get("data", {}).get("data", {}).get(data_key, None) + ) if new_secret_value_from_vault != new_secret_value: verbose_logger.exception( f"New secret value mismatch. Expected: {new_secret_value}, Got: {new_secret_value_from_vault}" @@ -529,8 +546,13 @@ class HashicorpSecretManager(BaseSecretManager): } except httpx.HTTPStatusError as e: if e.response.status_code == 404: - verbose_logger.exception(f"Failed to verify new secret {new_secret_name}") - return {"status": "error", "message": f"Failed to verify new secret {new_secret_name}"} + verbose_logger.exception( + f"Failed to verify new secret {new_secret_name}" + ) + return { + "status": "error", + "message": f"Failed to verify new secret {new_secret_name}", + } verbose_logger.exception( f"Error verifying new secret: {e.response.text if hasattr(e, 'response') else str(e)}" ) @@ -540,7 +562,10 @@ class HashicorpSecretManager(BaseSecretManager): } except Exception as e: verbose_logger.exception(f"Error verifying new secret: {e}") - return {"status": "error", "message": f"Error verifying new secret: {e}"} + return { + "status": "error", + "message": f"Error verifying new secret: {e}", + } # If everything is successful, delete the old secret # Only delete if the names are different (same name means we're just updating the value) @@ -552,7 +577,10 @@ class HashicorpSecretManager(BaseSecretManager): timeout=timeout, ) # Check if async_delete_secret returned an error - if isinstance(delete_response, dict) and delete_response.get("status") == "error": + if ( + isinstance(delete_response, dict) + and delete_response.get("status") == "error" + ): # Log the error but don't fail the rotation since new secret was created successfully verbose_logger.warning( f"Failed to delete old secret {current_secret_name} after rotation: {delete_response.get('message')}" diff --git a/litellm/secret_managers/main.py b/litellm/secret_managers/main.py index 38405f058c..2aca1cd9dd 100644 --- a/litellm/secret_managers/main.py +++ b/litellm/secret_managers/main.py @@ -21,10 +21,10 @@ def _get_oidc_http_handler(timeout: Optional[httpx.Timeout] = None) -> HTTPHandl """ Factory function to create HTTPHandler for OIDC requests. This function can be mocked in tests. - + Args: timeout: Optional timeout for HTTP requests. Defaults to 600.0 seconds with 5.0 connect timeout. - + Returns: HTTPHandler instance configured for OIDC requests. """ @@ -148,7 +148,10 @@ def get_secret( # noqa: PLR0915 # https://docs.github.com/en/actions/deployment/security-hardening-your-deployments/configuring-openid-connect-in-cloud-providers#using-custom-actions actions_id_token_request_url = os.getenv("ACTIONS_ID_TOKEN_REQUEST_URL") actions_id_token_request_token = os.getenv("ACTIONS_ID_TOKEN_REQUEST_TOKEN") - if actions_id_token_request_url is None or actions_id_token_request_token is None: + if ( + actions_id_token_request_url is None + or actions_id_token_request_token is None + ): raise ValueError( "ACTIONS_ID_TOKEN_REQUEST_URL or ACTIONS_ID_TOKEN_REQUEST_TOKEN not found in environment" ) @@ -215,7 +218,10 @@ def get_secret( # noqa: PLR0915 raise ValueError("Unsupported OIDC provider") try: - if _should_read_secret_from_secret_manager() and litellm.secret_manager_client is not None: + if ( + _should_read_secret_from_secret_manager() + and litellm.secret_manager_client is not None + ): try: client = litellm.secret_manager_client key_manager = "local" @@ -253,7 +259,9 @@ def get_secret( # noqa: PLR0915 else: secret = os.environ.get(secret_name) secret_value_as_bool = str_to_bool(secret) if secret is not None else None - if secret_value_as_bool is not None and isinstance(secret_value_as_bool, bool): + if secret_value_as_bool is not None and isinstance( + secret_value_as_bool, bool + ): return secret_value_as_bool else: return secret diff --git a/litellm/secret_managers/secret_manager_handler.py b/litellm/secret_managers/secret_manager_handler.py index b93503a864..eb90dda0e9 100644 --- a/litellm/secret_managers/secret_manager_handler.py +++ b/litellm/secret_managers/secret_manager_handler.py @@ -15,13 +15,14 @@ from litellm.types.secret_managers.main import KeyManagementSystem def _is_base64(s): """Check if a string is valid base64.""" import binascii + try: return base64.b64encode(base64.b64decode(s)).decode() == s except binascii.Error: return False -def get_secret_from_manager( # noqa: PLR0915 +def get_secret_from_manager( # noqa: PLR0915 client: Any, key_manager: str, secret_name: str, @@ -29,36 +30,38 @@ def get_secret_from_manager( # noqa: PLR0915 ) -> Optional[str]: """ Get a secret from the configured secret manager. - + Args: client: The secret manager client instance key_manager: The type of key manager (e.g., "azure_key_vault", "google_kms", etc.) secret_name: The name/path of the secret to retrieve key_management_settings: Optional settings for the key management system - + Returns: The secret value as a string, or None if not found - + Raises: ValueError: If the secret cannot be retrieved or required parameters are missing Exception: For other errors during secret retrieval """ secret = None - + if ( key_manager == KeyManagementSystem.AZURE_KEY_VAULT.value or type(client).__module__ + "." + type(client).__name__ == "azure.keyvault.secrets._client.SecretClient" ): # support Azure Secret Client - from azure.keyvault.secrets import SecretClient secret = client.get_secret(secret_name).value - + elif ( key_manager == KeyManagementSystem.GOOGLE_KMS.value or client.__class__.__name__ == "KeyManagementServiceClient" ): encrypted_secret: Any = os.getenv(secret_name) if encrypted_secret is None: - raise ValueError("Google KMS requires the encrypted secret to be in the environment!") + raise ValueError( + "Google KMS requires the encrypted secret to be in the environment!" + ) b64_flag = _is_base64(encrypted_secret) if b64_flag is True: # if passed in as encoded b64 string encrypted_secret = base64.b64decode(encrypted_secret) @@ -73,15 +76,19 @@ def get_secret_from_manager( # noqa: PLR0915 "ciphertext": ciphertext, } ) - secret = response.plaintext.decode("utf-8") # assumes the original value was encoded with utf-8 - + secret = response.plaintext.decode( + "utf-8" + ) # assumes the original value was encoded with utf-8 + elif key_manager == KeyManagementSystem.AWS_KMS.value: """ Only check the tokens which start with 'aws_kms/'. This prevents latency impact caused by checking all keys. """ encrypted_value = os.getenv(secret_name, None) if encrypted_value is None: - raise Exception("AWS KMS - Encrypted Value of Key={} is None".format(secret_name)) + raise Exception( + "AWS KMS - Encrypted Value of Key={} is None".format(secret_name) + ) # Decode the base64 encoded ciphertext ciphertext_blob = base64.b64decode(encrypted_value) @@ -95,7 +102,7 @@ def get_secret_from_manager( # noqa: PLR0915 secret = plaintext.decode("utf-8") if isinstance(secret, str): secret = secret.strip() - + elif key_manager == KeyManagementSystem.AWS_SECRET_MANAGER.value: from litellm.secret_managers.aws_secret_manager_v2 import ( AWSSecretsManagerV2, @@ -105,62 +112,71 @@ def get_secret_from_manager( # noqa: PLR0915 primary_secret_name = None if key_management_settings is not None: primary_secret_name = key_management_settings.primary_secret_name - + secret = client.sync_read_secret( secret_name=secret_name, primary_secret_name=primary_secret_name, ) print_verbose(f"get_secret_value_response: {secret}") - + elif key_manager == KeyManagementSystem.GOOGLE_SECRET_MANAGER.value: try: secret = client.get_secret_from_google_secret_manager(secret_name) print_verbose(f"secret from google secret manager: {secret}") if secret is None: - raise ValueError(f"No secret found in Google Secret Manager for {secret_name}") + raise ValueError( + f"No secret found in Google Secret Manager for {secret_name}" + ) except Exception as e: print_verbose(f"An error occurred - {str(e)}") raise e - + elif key_manager == KeyManagementSystem.HASHICORP_VAULT.value: try: secret = client.sync_read_secret(secret_name=secret_name) if secret is None: - raise ValueError(f"No secret found in Hashicorp Secret Manager for {secret_name}") + raise ValueError( + f"No secret found in Hashicorp Secret Manager for {secret_name}" + ) except Exception as e: print_verbose(f"An error occurred - {str(e)}") raise e - + elif key_manager == KeyManagementSystem.CYBERARK.value: try: secret = client.sync_read_secret(secret_name=secret_name) if secret is None: - raise ValueError(f"No secret found in CyberArk Secret Manager for {secret_name}") + raise ValueError( + f"No secret found in CyberArk Secret Manager for {secret_name}" + ) except Exception as e: print_verbose(f"An error occurred - {str(e)}") raise e - + elif key_manager == KeyManagementSystem.CUSTOM.value: # Check if client is a CustomSecretManager instance from litellm.integrations.custom_secret_manager import CustomSecretManager - + if isinstance(client, CustomSecretManager): secret = client.sync_read_secret( secret_name=secret_name, - optional_params=key_management_settings.model_dump() if key_management_settings else None, + optional_params=key_management_settings.model_dump() + if key_management_settings + else None, ) if secret is None: - raise ValueError(f"No secret found in Custom Secret Manager for {secret_name}") + raise ValueError( + f"No secret found in Custom Secret Manager for {secret_name}" + ) else: raise ValueError( f"Custom secret manager client must be an instance of CustomSecretManager, got {type(client).__name__}" ) - + elif key_manager == "local": secret = os.getenv(secret_name) - + else: # assume the default is infisicial client secret = client.get_secret(secret_name).secret_value - - return secret + return secret diff --git a/litellm/skills/__init__.py b/litellm/skills/__init__.py index 5a5f332068..96147d5a10 100644 --- a/litellm/skills/__init__.py +++ b/litellm/skills/__init__.py @@ -21,4 +21,3 @@ __all__ = [ "delete_skill", "adelete_skill", ] - diff --git a/litellm/skills/main.py b/litellm/skills/main.py index f6abd9043d..f3b670d5d6 100644 --- a/litellm/skills/main.py +++ b/litellm/skills/main.py @@ -41,6 +41,7 @@ def _get_litellm_skills_handler(): from litellm.llms.litellm_proxy.skills.transformation import ( LiteLLMSkillsTransformationHandler, ) + _litellm_skills_handler = LiteLLMSkillsTransformationHandler() return _litellm_skills_handler @@ -58,7 +59,7 @@ async def acreate_skill( ) -> Skill: """ Async: Create a new skill - + Args: files: Files to upload for the skill. All files must be in the same top-level directory and must include a SKILL.md file at the root. display_title: Optional display title for the skill @@ -68,7 +69,7 @@ async def acreate_skill( timeout: Request timeout custom_llm_provider: Provider name (e.g., 'anthropic') **kwargs: Additional parameters - + Returns: Skill object """ @@ -121,7 +122,7 @@ def create_skill( ) -> Union[Skill, Coroutine[Any, Any, Skill]]: """ Create a new skill - + Args: files: Files to upload for the skill. All files must be in the same top-level directory and must include a SKILL.md file at the root. display_title: Optional display title for the skill @@ -131,7 +132,7 @@ def create_skill( timeout: Request timeout custom_llm_provider: Provider name (e.g., 'anthropic') **kwargs: Additional parameters - + Returns: Skill object """ @@ -172,16 +173,14 @@ def create_skill( ) # Get provider config for external providers (Anthropic, etc.) - skills_api_provider_config: Optional[BaseSkillsAPIConfig] = ( - ProviderConfigManager.get_provider_skills_api_config( - provider=litellm.LlmProviders(custom_llm_provider), - ) + skills_api_provider_config: Optional[ + BaseSkillsAPIConfig + ] = ProviderConfigManager.get_provider_skills_api_config( + provider=litellm.LlmProviders(custom_llm_provider), ) if skills_api_provider_config is None: - raise ValueError( - f"CREATE skill is not supported for {custom_llm_provider}" - ) + raise ValueError(f"CREATE skill is not supported for {custom_llm_provider}") # Validate environment and get headers headers = extra_headers or {} @@ -253,7 +252,7 @@ async def alist_skills( ) -> ListSkillsResponse: """ Async: List all skills - + Args: limit: Number of results to return per page (max 100, default 20) page: Pagination token for fetching a specific page of results @@ -263,7 +262,7 @@ async def alist_skills( timeout: Request timeout custom_llm_provider: Provider name (e.g., 'anthropic') **kwargs: Additional parameters - + Returns: ListSkillsResponse object """ @@ -316,7 +315,7 @@ def list_skills( ) -> Union[ListSkillsResponse, Coroutine[Any, Any, ListSkillsResponse]]: """ List all skills - + Args: limit: Number of results to return per page (max 100, default 20) page: Pagination token for fetching a specific page of results @@ -326,7 +325,7 @@ def list_skills( timeout: Request timeout custom_llm_provider: Provider name (e.g., 'anthropic') **kwargs: Additional parameters - + Returns: ListSkillsResponse object """ @@ -354,10 +353,10 @@ def list_skills( ) # Get provider config for external providers (Anthropic, etc.) - skills_api_provider_config: Optional[BaseSkillsAPIConfig] = ( - ProviderConfigManager.get_provider_skills_api_config( - provider=litellm.LlmProviders(custom_llm_provider), - ) + skills_api_provider_config: Optional[ + BaseSkillsAPIConfig + ] = ProviderConfigManager.get_provider_skills_api_config( + provider=litellm.LlmProviders(custom_llm_provider), ) if skills_api_provider_config is None: @@ -436,7 +435,7 @@ async def aget_skill( ) -> Skill: """ Async: Get a skill by ID - + Args: skill_id: The ID of the skill to fetch extra_headers: Additional headers for the request @@ -444,7 +443,7 @@ async def aget_skill( timeout: Request timeout custom_llm_provider: Provider name (e.g., 'anthropic') **kwargs: Additional parameters - + Returns: Skill object """ @@ -493,7 +492,7 @@ def get_skill( ) -> Union[Skill, Coroutine[Any, Any, Skill]]: """ Get a skill by ID - + Args: skill_id: The ID of the skill to fetch extra_headers: Additional headers for the request @@ -501,7 +500,7 @@ def get_skill( timeout: Request timeout custom_llm_provider: Provider name (e.g., 'anthropic') **kwargs: Additional parameters - + Returns: Skill object """ @@ -528,10 +527,10 @@ def get_skill( ) # Get provider config for external providers (Anthropic, etc.) - skills_api_provider_config: Optional[BaseSkillsAPIConfig] = ( - ProviderConfigManager.get_provider_skills_api_config( - provider=litellm.LlmProviders(custom_llm_provider), - ) + skills_api_provider_config: Optional[ + BaseSkillsAPIConfig + ] = ProviderConfigManager.get_provider_skills_api_config( + provider=litellm.LlmProviders(custom_llm_provider), ) if skills_api_provider_config is None: @@ -602,7 +601,7 @@ async def adelete_skill( ) -> DeleteSkillResponse: """ Async: Delete a skill by ID - + Args: skill_id: The ID of the skill to delete extra_headers: Additional headers for the request @@ -610,7 +609,7 @@ async def adelete_skill( timeout: Request timeout custom_llm_provider: Provider name (e.g., 'anthropic') **kwargs: Additional parameters - + Returns: DeleteSkillResponse object """ @@ -659,7 +658,7 @@ def delete_skill( ) -> Union[DeleteSkillResponse, Coroutine[Any, Any, DeleteSkillResponse]]: """ Delete a skill by ID - + Args: skill_id: The ID of the skill to delete extra_headers: Additional headers for the request @@ -667,7 +666,7 @@ def delete_skill( timeout: Request timeout custom_llm_provider: Provider name (e.g., 'anthropic') **kwargs: Additional parameters - + Returns: DeleteSkillResponse object """ @@ -694,16 +693,14 @@ def delete_skill( ) # Get provider config for external providers (Anthropic, etc.) - skills_api_provider_config: Optional[BaseSkillsAPIConfig] = ( - ProviderConfigManager.get_provider_skills_api_config( - provider=litellm.LlmProviders(custom_llm_provider), - ) + skills_api_provider_config: Optional[ + BaseSkillsAPIConfig + ] = ProviderConfigManager.get_provider_skills_api_config( + provider=litellm.LlmProviders(custom_llm_provider), ) if skills_api_provider_config is None: - raise ValueError( - f"DELETE skill is not supported for {custom_llm_provider}" - ) + raise ValueError(f"DELETE skill is not supported for {custom_llm_provider}") # Validate environment and get headers headers = extra_headers or {} @@ -757,4 +754,3 @@ def delete_skill( completion_kwargs=local_vars, extra_kwargs=kwargs, ) - diff --git a/litellm/types/caching.py b/litellm/types/caching.py index 7126ba3e9b..c8194ce2e7 100644 --- a/litellm/types/caching.py +++ b/litellm/types/caching.py @@ -113,6 +113,7 @@ class HealthCheckCacheParams(BaseModel): class CachedEmbedding(TypedDict): """Type definition for cached embedding objects""" + embedding: Optional[List[float]] index: Optional[int] object: Optional[str] diff --git a/litellm/types/containers/main.py b/litellm/types/containers/main.py index 66aa7dc5fa..df8c05a74c 100644 --- a/litellm/types/containers/main.py +++ b/litellm/types/containers/main.py @@ -6,12 +6,14 @@ from typing_extensions import TypedDict class ExpiresAfter(BaseModel): """Container expiration settings.""" + anchor: Literal["last_active_at"] minutes: int class ContainerObject(BaseModel): """Represents a container object.""" + id: str object: Literal["container"] created_at: int @@ -43,6 +45,7 @@ class ContainerObject(BaseModel): class DeleteContainerResult(BaseModel): """Result of a delete container request.""" + id: str object: Literal["container.deleted"] deleted: bool @@ -65,6 +68,7 @@ class DeleteContainerResult(BaseModel): class ContainerListResponse(BaseModel): """Response object for list containers request.""" + object: Literal["list"] data: List[ContainerObject] first_id: Optional[str] = None @@ -90,9 +94,10 @@ class ContainerListResponse(BaseModel): class ContainerCreateOptionalRequestParams(TypedDict, total=False): """ TypedDict for Optional parameters supported by OpenAI's container creation API. - + Params here: https://platform.openai.com/docs/api-reference/containers/create """ + expires_after: Optional[Dict[str, Any]] # ExpiresAfter object file_ids: Optional[List[str]] extra_headers: Optional[Dict[str, str]] @@ -102,18 +107,20 @@ class ContainerCreateOptionalRequestParams(TypedDict, total=False): class ContainerCreateRequestParams(ContainerCreateOptionalRequestParams, total=False): """ TypedDict for request parameters supported by OpenAI's container creation API. - + Params here: https://platform.openai.com/docs/api-reference/containers/create """ + name: str class ContainerListOptionalRequestParams(TypedDict, total=False): """ TypedDict for Optional parameters supported by OpenAI's container list API. - + Params here: https://platform.openai.com/docs/api-reference/containers/list """ + after: Optional[str] limit: Optional[int] order: Optional[str] @@ -123,8 +130,11 @@ class ContainerListOptionalRequestParams(TypedDict, total=False): class ContainerFileObject(BaseModel): """Represents a container file object.""" + id: str - object: Literal["container.file", "container_file"] # OpenAI returns "container.file" + object: Literal[ + "container.file", "container_file" + ] # OpenAI returns "container.file" container_id: str bytes: Optional[int] = None # Can be null for some files created_at: int @@ -150,6 +160,7 @@ class ContainerFileObject(BaseModel): class ContainerFileListResponse(BaseModel): """Response object for list container files request.""" + object: Literal["list"] data: List[ContainerFileObject] first_id: Optional[str] = None @@ -174,6 +185,7 @@ class ContainerFileListResponse(BaseModel): class DeleteContainerFileResponse(BaseModel): """Response object for delete container file request.""" + id: str object: Literal["container_file.deleted"] deleted: bool @@ -192,4 +204,3 @@ class DeleteContainerFileResponse(BaseModel): return self.model_dump(**kwargs) except Exception: return self.dict() - diff --git a/litellm/types/files.py b/litellm/types/files.py index 8b87b33cd1..bf56894329 100644 --- a/litellm/types/files.py +++ b/litellm/types/files.py @@ -293,10 +293,10 @@ Two-Step File Upload Types class TwoStepFileUploadRequest(TypedDict): """ Request structure for two-step file upload process. - + Step 1: Initial request to get upload URL Step 2: Upload file content to the upload URL - + Used by providers like Manus and Google Cloud Storage. """ @@ -309,7 +309,7 @@ class TwoStepFileUploadRequest(TypedDict): class TwoStepFileUploadConfig(TypedDict, total=False): """ Configuration for two-step file upload process. - + Properties: initial_request: Request to create file record and get upload URL upload_request: Request to upload actual file content diff --git a/litellm/types/google_genai/__init__.py b/litellm/types/google_genai/__init__.py index f510f3cdbe..9f74df91e2 100644 --- a/litellm/types/google_genai/__init__.py +++ b/litellm/types/google_genai/__init__.py @@ -7,7 +7,7 @@ from .main import ( __all__ = [ "ContentListUnion", - "ContentListUnionDict", + "ContentListUnionDict", "GenerateContentConfigOrDict", "GenerateContentResponse", -] \ No newline at end of file +] diff --git a/litellm/types/google_genai/main.py b/litellm/types/google_genai/main.py index 0a26f266a6..89781860d2 100644 --- a/litellm/types/google_genai/main.py +++ b/litellm/types/google_genai/main.py @@ -21,11 +21,12 @@ if TYPE_CHECKING: class GenerateContentRequestDict(GenerateContentRequestParametersDict): # type: ignore[misc] generationConfig: Optional[Any] - tools: Optional[ToolConfigDict] # type: ignore[assignment] + tools: Optional[ToolConfigDict] # type: ignore[assignment] - class GenerateContentResponse(GoogleGenAIGenerateContentResponse, BaseLiteLLMOpenAIResponseObject): # type: ignore[misc] + class GenerateContentResponse(GoogleGenAIGenerateContentResponse, BaseLiteLLMOpenAIResponseObject): # type: ignore[misc] _hidden_params: dict = {} pass + else: # Fallback types when google.genai is not available ContentListUnion = Any @@ -48,11 +49,11 @@ else: class GenerateContentRequestDict(GenerateContentRequestParametersDict): # type: ignore[misc] def __init__(self, **kwargs): # type: ignore # Extract specific fields - self.generationConfig = kwargs.get('generationConfig') - self.tools = kwargs.get('tools') + self.generationConfig = kwargs.get("generationConfig") + self.tools = kwargs.get("tools") super().__init__(**kwargs) - class GenerateContentResponse(BaseLiteLLMOpenAIResponseObject): # type: ignore[misc] + class GenerateContentResponse(BaseLiteLLMOpenAIResponseObject): # type: ignore[misc] def __init__(self, **kwargs): # type: ignore super().__init__(**kwargs) - self._hidden_params = kwargs.get('_hidden_params', {}) \ No newline at end of file + self._hidden_params = kwargs.get("_hidden_params", {}) diff --git a/litellm/types/integrations/azure_sentinel.py b/litellm/types/integrations/azure_sentinel.py index f821dc9733..8460c7c0be 100644 --- a/litellm/types/integrations/azure_sentinel.py +++ b/litellm/types/integrations/azure_sentinel.py @@ -9,4 +9,3 @@ class AzureSentinelInitParams(StandardCustomLoggerInitParams): """ pass - diff --git a/litellm/types/integrations/cloudzero.py b/litellm/types/integrations/cloudzero.py index e79500e08d..36db7df135 100644 --- a/litellm/types/integrations/cloudzero.py +++ b/litellm/types/integrations/cloudzero.py @@ -3,12 +3,12 @@ from typing import Any, Dict class CBFRecord(Dict[str, Any]): """CloudZero Billing Format (CBF) record structure. - - This class represents a CBF record that is created from LiteLLM usage data - for CloudZero integration. Since CBF field names contain forward slashes - (e.g., 'time/usage_start', 'cost/cost'), we use a Dict base class rather + + This class represents a CBF record that is created from LiteLLM usage data + for CloudZero integration. Since CBF field names contain forward slashes + (e.g., 'time/usage_start', 'cost/cost'), we use a Dict base class rather than TypedDict to accommodate the special characters in field names. - + Expected CBF fields (per LIT-1907): - time/usage_start: ISO-formatted UTC datetime (Optional[str]) - cost/cost: Billed cost (float) @@ -28,8 +28,9 @@ class CBFRecord(Dict[str, Any]): - resource/tag:user_alias: User alias if available (Optional[str]) - resource/tag:{key}: Various resource tags for dimensions and metrics (Optional[str]) """ + pass # Type alias for better readability in function signatures -CBFRecordDict = Dict[str, Any] \ No newline at end of file +CBFRecordDict = Dict[str, Any] diff --git a/litellm/types/integrations/custom_logger.py b/litellm/types/integrations/custom_logger.py index 96952404b7..0698940922 100644 --- a/litellm/types/integrations/custom_logger.py +++ b/litellm/types/integrations/custom_logger.py @@ -7,4 +7,5 @@ class StandardCustomLoggerInitParams(BaseModel): """ Params for initializing a CustomLogger. """ - turn_off_message_logging: Optional[bool] = False \ No newline at end of file + + turn_off_message_logging: Optional[bool] = False diff --git a/litellm/types/integrations/langfuse_otel.py b/litellm/types/integrations/langfuse_otel.py index 53d84c4005..17b5a78edf 100644 --- a/litellm/types/integrations/langfuse_otel.py +++ b/litellm/types/integrations/langfuse_otel.py @@ -11,7 +11,8 @@ else: class LangfuseOtelConfig(BaseModel): otlp_auth_headers: Optional[str] = None - protocol: Protocol = "otlp_http" + protocol: Protocol = "otlp_http" + class LangfuseSpanAttributes(str, Enum): LANGFUSE_ENVIRONMENT = "langfuse.environment" @@ -41,4 +42,4 @@ class LangfuseSpanAttributes(str, Enum): UPDATE_TRACE_KEYS = "langfuse.trace.update_keys" # ---- Misc / flags ---- - DEBUG_LANGFUSE = "langfuse.debug" \ No newline at end of file + DEBUG_LANGFUSE = "langfuse.debug" diff --git a/litellm/types/integrations/weave_otel.py b/litellm/types/integrations/weave_otel.py index 5b40ff8534..d3cf489435 100644 --- a/litellm/types/integrations/weave_otel.py +++ b/litellm/types/integrations/weave_otel.py @@ -24,8 +24,7 @@ class WeaveSpanAttributes(str, Enum): """ DISPLAY_NAME = "wandb.display_name" - + # Thread organization, similar to OpenInference session_id. THREAD_ID = "wandb.thread_id" IS_TURN = "wandb.is_turn" - diff --git a/litellm/types/interactions/generated.py b/litellm/types/interactions/generated.py index 30e4ff4722..ed626b0b7c 100644 --- a/litellm/types/interactions/generated.py +++ b/litellm/types/interactions/generated.py @@ -13,14 +13,14 @@ from pydantic import AwareDatetime, Base64Str, BaseModel, Field, RootModel class Annotation(BaseModel): start_index: Optional[int] = Field( None, - description='Start of segment of the response that is attributed to this source.\n\nIndex indicates the start of the segment, measured in bytes.', + description="Start of segment of the response that is attributed to this source.\n\nIndex indicates the start of the segment, measured in bytes.", ) end_index: Optional[int] = Field( - None, description='End of the attributed segment, exclusive.' + None, description="End of the attributed segment, exclusive." ) source: Optional[str] = Field( None, - description='Source attributed for a portion of the text. Could be a URL, title, or\nother identifier.', + description="Source attributed for a portion of the text. Could be a URL, title, or\nother identifier.", ) @@ -28,105 +28,105 @@ class DocumentContent(BaseModel): data: Optional[Base64Str] = None uri: Optional[str] = None mime_type: Optional[str] = None - type: Literal['document'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["document"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) class FunctionCallContent(BaseModel): - name: str = Field(..., description='The name of the tool to call.') + name: str = Field(..., description="The name of the tool to call.") arguments: Dict[str, Any] = Field( - ..., description='The arguments to pass to the function.' + ..., description="The arguments to pass to the function." ) - type: Literal['function_call'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["function_call"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) - id: str = Field(..., description='A unique ID for this specific tool call.') + id: str = Field(..., description="A unique ID for this specific tool call.") class Language(Enum): - python = 'python' + python = "python" class CodeExecutionCallArguments(BaseModel): language: Optional[Language] = Field( - None, description='Programming language of the `code`.' + None, description="Programming language of the `code`." ) - code: Optional[str] = Field(None, description='The code to be executed.') + code: Optional[str] = Field(None, description="The code to be executed.") class UrlContextCallArguments(BaseModel): - urls: Optional[List[str]] = Field(None, description='The URLs to fetch.') + urls: Optional[List[str]] = Field(None, description="The URLs to fetch.") class McpServerToolCallContent(BaseModel): - name: str = Field(..., description='The name of the tool which was called.') - server_name: str = Field(..., description='The name of the used MCP server.') + name: str = Field(..., description="The name of the tool which was called.") + server_name: str = Field(..., description="The name of the used MCP server.") arguments: Dict[str, Any] = Field( - ..., description='The JSON object of arguments for the function.' + ..., description="The JSON object of arguments for the function." ) - type: Literal['mcp_server_tool_call'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["mcp_server_tool_call"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) - id: str = Field(..., description='A unique ID for this specific tool call.') + id: str = Field(..., description="A unique ID for this specific tool call.") class GoogleSearchCallArguments(BaseModel): queries: Optional[List[str]] = Field( - None, description='Web search queries for the following-up web search.' + None, description="Web search queries for the following-up web search." ) class CodeExecutionResultContent(BaseModel): - result: Optional[str] = Field(None, description='The output of the code execution.') + result: Optional[str] = Field(None, description="The output of the code execution.") is_error: Optional[bool] = Field( - None, description='Whether the code execution resulted in an error.' + None, description="Whether the code execution resulted in an error." ) signature: Optional[str] = Field( - None, description='A signature hash for backend validation.' + None, description="A signature hash for backend validation." ) - type: Literal['code_execution_result'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["code_execution_result"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) call_id: Optional[str] = Field( - None, description='ID to match the ID from the code execution call block.' + None, description="ID to match the ID from the code execution call block." ) class Status(Enum): - success = 'success' - error = 'error' - paywall = 'paywall' - unsafe = 'unsafe' + success = "success" + error = "error" + paywall = "paywall" + unsafe = "unsafe" class UrlContextResult(BaseModel): - url: Optional[str] = Field(None, description='The URL that was fetched.') + url: Optional[str] = Field(None, description="The URL that was fetched.") status: Optional[Status] = Field( - None, description='The status of the URL retrieval.' + None, description="The status of the URL retrieval." ) class GoogleSearchResult(BaseModel): - url: Optional[str] = Field(None, description='URI reference of the search result.') - title: Optional[str] = Field(None, description='Title of the search result.') + url: Optional[str] = Field(None, description="URI reference of the search result.") + title: Optional[str] = Field(None, description="Title of the search result.") rendered_content: Optional[str] = Field( None, - description='Web content snippet that can be embedded in a web page or an app webview.', + description="Web content snippet that can be embedded in a web page or an app webview.", ) class FileSearchResult(BaseModel): - title: Optional[str] = Field(None, description='The title of the search result.') - text: Optional[str] = Field(None, description='The text of the search result.') + title: Optional[str] = Field(None, description="The title of the search result.") + text: Optional[str] = Field(None, description="The text of the search result.") file_search_store: Optional[str] = Field( - None, description='The name of the file search store.' + None, description="The name of the file search store." ) class SpeechConfig(BaseModel): - voice: Optional[str] = Field(None, description='The voice of the speaker.') - language: Optional[str] = Field(None, description='The language of the speech.') + voice: Optional[str] = Field(None, description="The voice of the speaker.") + language: Optional[str] = Field(None, description="The language of the speech.") speaker: Optional[str] = Field( None, description="The speaker's name, it should match the speaker name given in the prompt.", @@ -134,94 +134,94 @@ class SpeechConfig(BaseModel): class DynamicAgentConfig(BaseModel): - type: Literal['dynamic'] = Field( - 'dynamic', - description='Used as the OpenAPI type discriminator for the content oneof.', + type: Literal["dynamic"] = Field( + "dynamic", + description="Used as the OpenAPI type discriminator for the content oneof.", ) class Function(BaseModel): - name: Optional[str] = Field(None, description='The name of the function.') + name: Optional[str] = Field(None, description="The name of the function.") description: Optional[str] = Field( - None, description='A description of the function.' + None, description="A description of the function." ) parameters: Optional[Any] = Field( None, description="The JSON Schema for the function's parameters." ) - type: Literal['function'] + type: Literal["function"] class CodeExecution(BaseModel): - type: Literal['code_execution'] + type: Literal["code_execution"] class UrlContext(BaseModel): - type: Literal['url_context'] + type: Literal["url_context"] class Environment(Enum): - browser = 'browser' + browser = "browser" class ComputerUse(BaseModel): - type: Literal['computer_use'] + type: Literal["computer_use"] environment: Optional[Environment] = Field( - None, description='The environment being operated.' + None, description="The environment being operated." ) excludedPredefinedFunctions: Optional[List[str]] = Field( None, - description='The list of predefined functions that are excluded from the model call.', + description="The list of predefined functions that are excluded from the model call.", ) class GoogleSearch(BaseModel): - type: Literal['google_search'] + type: Literal["google_search"] class FileSearch(BaseModel): file_search_store_names: Optional[List[str]] = Field( - None, description='The file search store names to search.' + None, description="The file search store names to search." ) top_k: Optional[int] = Field( - None, description='The number of semantic retrieval chunks to retrieve.' + None, description="The number of semantic retrieval chunks to retrieve." ) metadata_filter: Optional[str] = Field( None, - description='Metadata filter to apply to the semantic retrieval documents and chunks.', + description="Metadata filter to apply to the semantic retrieval documents and chunks.", ) - type: Literal['file_search'] + type: Literal["file_search"] class EventType(Enum): - interaction_start = 'interaction.start' - interaction_complete = 'interaction.complete' + interaction_start = "interaction.start" + interaction_complete = "interaction.complete" class Status1(Enum): - in_progress = 'in_progress' - requires_action = 'requires_action' - completed = 'completed' - failed = 'failed' - cancelled = 'cancelled' + in_progress = "in_progress" + requires_action = "requires_action" + completed = "completed" + failed = "failed" + cancelled = "cancelled" class InteractionStatusUpdate(BaseModel): interaction_id: Optional[str] = None status: Optional[Status1] = None - event_type: Literal['interaction.status_update'] = 'interaction.status_update' + event_type: Literal["interaction.status_update"] = "interaction.status_update" event_id: Optional[str] = Field( None, - description='The event_id token to be used to resume the interaction stream, from\nthis event.', + description="The event_id token to be used to resume the interaction stream, from\nthis event.", ) class TextDelta(BaseModel): text: Optional[str] = None - type: Literal['text'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["text"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) annotations: Optional[List[Annotation]] = Field( - None, description='Citation information for model-generated content.' + None, description="Citation information for model-generated content." ) @@ -229,59 +229,59 @@ class DocumentDelta(BaseModel): data: Optional[Base64Str] = None uri: Optional[str] = None mime_type: Optional[str] = None - type: Literal['document'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["document"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) class ThoughtSignatureDelta(BaseModel): signature: Optional[Base64Str] = Field( None, - description='Signature to match the backend source to be part of the generation.', + description="Signature to match the backend source to be part of the generation.", ) - type: Literal['thought_signature'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["thought_signature"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) class FunctionCallDelta(BaseModel): name: Optional[str] = None arguments: Optional[Dict[str, Any]] = None - type: Literal['function_call'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["function_call"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) id: Optional[str] = Field( - None, description='A unique ID for this specific tool call.' + None, description="A unique ID for this specific tool call." ) class CodeExecutionCallDelta(BaseModel): arguments: Optional[CodeExecutionCallArguments] = None - type: Literal['code_execution_call'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["code_execution_call"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) id: Optional[str] = Field( - None, description='A unique ID for this specific tool call.' + None, description="A unique ID for this specific tool call." ) class UrlContextCallDelta(BaseModel): arguments: Optional[UrlContextCallArguments] = None - type: Literal['url_context_call'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["url_context_call"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) id: Optional[str] = Field( - None, description='A unique ID for this specific tool call.' + None, description="A unique ID for this specific tool call." ) class GoogleSearchCallDelta(BaseModel): arguments: Optional[GoogleSearchCallArguments] = None - type: Literal['google_search_call'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["google_search_call"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) id: Optional[str] = Field( - None, description='A unique ID for this specific tool call.' + None, description="A unique ID for this specific tool call." ) @@ -289,11 +289,11 @@ class McpServerToolCallDelta(BaseModel): name: Optional[str] = None server_name: Optional[str] = None arguments: Optional[Dict[str, Any]] = None - type: Literal['mcp_server_tool_call'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["mcp_server_tool_call"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) id: Optional[str] = Field( - None, description='A unique ID for this specific tool call.' + None, description="A unique ID for this specific tool call." ) @@ -301,11 +301,11 @@ class CodeExecutionResultDelta(BaseModel): result: Optional[str] = None is_error: Optional[bool] = None signature: Optional[str] = None - type: Literal['code_execution_result'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["code_execution_result"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) call_id: Optional[str] = Field( - None, description='ID to match the ID from the function call block.' + None, description="ID to match the ID from the function call block." ) @@ -313,11 +313,11 @@ class UrlContextResultDelta(BaseModel): signature: Optional[str] = None result: Optional[List[UrlContextResult]] = None is_error: Optional[bool] = None - type: Literal['url_context_result'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["url_context_result"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) call_id: Optional[str] = Field( - None, description='ID to match the ID from the function call block.' + None, description="ID to match the ID from the function call block." ) @@ -325,113 +325,113 @@ class GoogleSearchResultDelta(BaseModel): signature: Optional[str] = None result: Optional[List[GoogleSearchResult]] = None is_error: Optional[bool] = None - type: Literal['google_search_result'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["google_search_result"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) call_id: Optional[str] = Field( - None, description='ID to match the ID from the function call block.' + None, description="ID to match the ID from the function call block." ) class FileSearchResultDelta(BaseModel): result: Optional[List[FileSearchResult]] = None - type: Literal['file_search_result'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["file_search_result"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) class ContentStop(BaseModel): index: Optional[int] = None - event_type: Literal['content.stop'] = 'content.stop' + event_type: Literal["content.stop"] = "content.stop" event_id: Optional[str] = Field( None, - description='The event_id token to be used to resume the interaction stream, from\nthis event.', + description="The event_id token to be used to resume the interaction stream, from\nthis event.", ) class Error(BaseModel): code: Optional[str] = Field( - None, description='A URI that identifies the error type.' + None, description="A URI that identifies the error type." ) - message: Optional[str] = Field(None, description='A human-readable error message.') + message: Optional[str] = Field(None, description="A human-readable error message.") class MediaResolution(Enum): - low = 'low' - medium = 'medium' - high = 'high' + low = "low" + medium = "medium" + high = "high" class ToolChoiceType(Enum): - auto = 'auto' - any = 'any' - none = 'none' - validated = 'validated' + auto = "auto" + any = "any" + none = "none" + validated = "validated" class ThinkingLevel(Enum): - low = 'low' - high = 'high' + low = "low" + high = "high" class ThinkingSummaries(Enum): - auto = 'auto' - none = 'none' + auto = "auto" + none = "none" class ResponseModality(Enum): - text = 'text' - image = 'image' - audio = 'audio' + text = "text" + image = "image" + audio = "audio" class Status3(Enum): - UNSPECIFIED = 'UNSPECIFIED' - IN_PROGRESS = 'IN_PROGRESS' - REQUIRES_ACTION = 'REQUIRES_ACTION' - COMPLETED = 'COMPLETED' - FAILED = 'FAILED' - CANCELLED = 'CANCELLED' - INCOMPLETE = 'INCOMPLETE' + UNSPECIFIED = "UNSPECIFIED" + IN_PROGRESS = "IN_PROGRESS" + REQUIRES_ACTION = "REQUIRES_ACTION" + COMPLETED = "COMPLETED" + FAILED = "FAILED" + CANCELLED = "CANCELLED" + INCOMPLETE = "INCOMPLETE" class ModelOption(RootModel[str]): root: str = Field( ..., - description='The model that will complete your prompt.\\n\\nSee [models](https://ai.google.dev/gemini-api/docs/models) for additional details.', - title='Model', + description="The model that will complete your prompt.\\n\\nSee [models](https://ai.google.dev/gemini-api/docs/models) for additional details.", + title="Model", ) class AgentOption(RootModel[str]): - root: str = Field(..., description='The agent to interact with.', title='Agent') + root: str = Field(..., description="The agent to interact with.", title="Agent") class ImageMimeTypeOption(RootModel[str]): root: str = Field( - ..., description='The mime type of the image.', title='ImageMimeType' + ..., description="The mime type of the image.", title="ImageMimeType" ) class AudioMimeTypeOption(RootModel[str]): root: str = Field( - ..., description='The mime type of the audio.', title='AudioMimeType' + ..., description="The mime type of the audio.", title="AudioMimeType" ) class VideoMimeTypeOption(RootModel[str]): root: str = Field( - ..., description='The mime type of the video.', title='VideoMimeType' + ..., description="The mime type of the video.", title="VideoMimeType" ) class TextContent(BaseModel): - text: Optional[str] = Field(None, description='The text content.') - type: Literal['text'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + text: Optional[str] = Field(None, description="The text content.") + type: Literal["text"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) annotations: Optional[List[Annotation]] = Field( - None, description='Citation information for model-generated content.' + None, description="Citation information for model-generated content." ) @@ -439,11 +439,11 @@ class ImageContent(BaseModel): data: Optional[Base64Str] = None uri: Optional[str] = None mime_type: Optional[ImageMimeTypeOption] = None - type: Literal['image'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["image"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) resolution: Optional[MediaResolution] = Field( - None, description='The resolution of the media.' + None, description="The resolution of the media." ) @@ -451,8 +451,8 @@ class AudioContent(BaseModel): data: Optional[Base64Str] = None uri: Optional[str] = None mime_type: Optional[AudioMimeTypeOption] = None - type: Literal['audio'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["audio"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) @@ -460,55 +460,55 @@ class VideoContent(BaseModel): data: Optional[Base64Str] = None uri: Optional[str] = None mime_type: Optional[VideoMimeTypeOption] = None - type: Literal['video'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["video"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) resolution: Optional[MediaResolution] = Field( - None, description='The resolution of the media.' + None, description="The resolution of the media." ) class ThoughtSummary1(RootModel[Union[TextContent, ImageContent]]): - root: Union[TextContent, ImageContent] = Field(..., discriminator='type') + root: Union[TextContent, ImageContent] = Field(..., discriminator="type") class ThoughtSummary(RootModel[List[ThoughtSummary1]]): - root: List[ThoughtSummary1] = Field(..., description='A summary of the thought.') + root: List[ThoughtSummary1] = Field(..., description="A summary of the thought.") class CodeExecutionCallContent(BaseModel): arguments: Optional[CodeExecutionCallArguments] = Field( - None, description='The arguments to pass to the code execution.' + None, description="The arguments to pass to the code execution." ) - type: Literal['code_execution_call'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["code_execution_call"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) id: Optional[str] = Field( - None, description='A unique ID for this specific tool call.' + None, description="A unique ID for this specific tool call." ) class UrlContextCallContent(BaseModel): arguments: Optional[UrlContextCallArguments] = Field( - None, description='The arguments to pass to the URL context.' + None, description="The arguments to pass to the URL context." ) - type: Literal['url_context_call'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["url_context_call"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) id: Optional[str] = Field( - None, description='A unique ID for this specific tool call.' + None, description="A unique ID for this specific tool call." ) class GoogleSearchCallContent(BaseModel): arguments: Optional[GoogleSearchCallArguments] = Field( - None, description='The arguments to pass to Google Search.' + None, description="The arguments to pass to Google Search." ) - type: Literal['google_search_call'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["google_search_call"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) id: Optional[str] = Field( - None, description='A unique ID for this specific tool call.' + None, description="A unique ID for this specific tool call." ) @@ -518,127 +518,127 @@ class Result(BaseModel): class FunctionResultContent(BaseModel): name: Optional[str] = Field( - None, description='The name of the tool that was called.' + None, description="The name of the tool that was called." ) is_error: Optional[bool] = Field( - None, description='Whether the tool call resulted in an error.' + None, description="Whether the tool call resulted in an error." ) - type: Literal['function_result'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["function_result"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) result: Union[Result, Dict[str, Any], str] = Field( - ..., description='The result of the tool call.' + ..., description="The result of the tool call." ) call_id: str = Field( - ..., description='ID to match the ID from the function call block.' + ..., description="ID to match the ID from the function call block." ) class UrlContextResultContent(BaseModel): signature: Optional[str] = Field( - None, description='The signature of the URL context result.' + None, description="The signature of the URL context result." ) result: Optional[List[UrlContextResult]] = Field( - None, description='The results of the URL context.' + None, description="The results of the URL context." ) is_error: Optional[bool] = Field( - None, description='Whether the URL context resulted in an error.' + None, description="Whether the URL context resulted in an error." ) - type: Literal['url_context_result'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["url_context_result"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) call_id: Optional[str] = Field( - None, description='ID to match the ID from the url context call block.' + None, description="ID to match the ID from the url context call block." ) class GoogleSearchResultContent(BaseModel): signature: Optional[str] = Field( - None, description='The signature of the Google Search result.' + None, description="The signature of the Google Search result." ) result: Optional[List[GoogleSearchResult]] = Field( - None, description='The results of the Google Search.' + None, description="The results of the Google Search." ) is_error: Optional[bool] = Field( - None, description='Whether the Google Search resulted in an error.' + None, description="Whether the Google Search resulted in an error." ) - type: Literal['google_search_result'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["google_search_result"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) call_id: Optional[str] = Field( - None, description='ID to match the ID from the google search call block.' + None, description="ID to match the ID from the google search call block." ) class McpServerToolResultContent(BaseModel): name: Optional[str] = Field( None, - description='Name of the tool which is called for this specific tool call.', + description="Name of the tool which is called for this specific tool call.", ) server_name: Optional[str] = Field( - None, description='The name of the used MCP server.' + None, description="The name of the used MCP server." ) - type: Literal['mcp_server_tool_result'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["mcp_server_tool_result"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) result: Union[Result, Dict[str, Any], str] = Field( - ..., description='The result of the tool call.' + ..., description="The result of the tool call." ) call_id: str = Field( - ..., description='ID to match the ID from the MCP server tool call block.' + ..., description="ID to match the ID from the MCP server tool call block." ) class FileSearchResultContent(BaseModel): result: Optional[List[FileSearchResult]] = Field( - None, description='The results of the File Search.' + None, description="The results of the File Search." ) - type: Literal['file_search_result'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["file_search_result"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) class AllowedTools(BaseModel): mode: Optional[ToolChoiceType] = Field( - None, description='The mode of the tool choice.' + None, description="The mode of the tool choice." ) tools: Optional[List[str]] = Field( - None, description='The names of the allowed tools.' + None, description="The names of the allowed tools." ) class DeepResearchAgentConfig(BaseModel): - type: Literal['deep-research'] = Field( - 'deep-research', - description='Used as the OpenAPI type discriminator for the content oneof.', + type: Literal["deep-research"] = Field( + "deep-research", + description="Used as the OpenAPI type discriminator for the content oneof.", ) thinking_summaries: Optional[ThinkingSummaries] = Field( - None, description='Whether to include thought summaries in the response.' + None, description="Whether to include thought summaries in the response." ) class McpServer(BaseModel): - type: Literal['mcp_server'] - name: Optional[str] = Field(None, description='The name of the MCPServer.') + type: Literal["mcp_server"] + name: Optional[str] = Field(None, description="The name of the MCPServer.") url: Optional[str] = Field( None, description='The full URL for the MCPServer endpoint.\nExample: "https://api.example.com/mcp"', ) headers: Optional[Dict[str, str]] = Field( None, - description='Optional: Fields for authentication headers, timeouts, etc., if needed.', + description="Optional: Fields for authentication headers, timeouts, etc., if needed.", ) allowed_tools: Optional[List[AllowedTools]] = Field( - None, description='The allowed tools.' + None, description="The allowed tools." ) class ModalityTokens(BaseModel): modality: Optional[ResponseModality] = Field( - None, description='The modality associated with the token count.' + None, description="The modality associated with the token count." ) tokens: Optional[int] = Field( - None, description='Number of tokens for the modality.' + None, description="Number of tokens for the modality." ) @@ -646,11 +646,11 @@ class ImageDelta(BaseModel): data: Optional[Base64Str] = None uri: Optional[str] = None mime_type: Optional[ImageMimeTypeOption] = None - type: Literal['image'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["image"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) resolution: Optional[MediaResolution] = Field( - None, description='The resolution of the media.' + None, description="The resolution of the media." ) @@ -658,8 +658,8 @@ class AudioDelta(BaseModel): data: Optional[Base64Str] = None uri: Optional[str] = None mime_type: Optional[AudioMimeTypeOption] = None - type: Literal['audio'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["audio"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) @@ -667,57 +667,57 @@ class VideoDelta(BaseModel): data: Optional[Base64Str] = None uri: Optional[str] = None mime_type: Optional[VideoMimeTypeOption] = None - type: Literal['video'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["video"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) resolution: Optional[MediaResolution] = Field( - None, description='The resolution of the media.' + None, description="The resolution of the media." ) class ThoughtSummaryDelta(BaseModel): - type: Literal['thought_summary'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["thought_summary"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) content: Optional[Union[TextContent, ImageContent]] = Field( - None, discriminator='type' + None, discriminator="type" ) class FunctionResultDelta(BaseModel): name: Optional[str] = None is_error: Optional[bool] = None - type: Literal['function_result'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["function_result"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) result: Optional[Union[Result, str]] = Field( - None, description='Tool call result delta.' + None, description="Tool call result delta." ) call_id: Optional[str] = Field( - None, description='ID to match the ID from the function call block.' + None, description="ID to match the ID from the function call block." ) class McpServerToolResultDelta(BaseModel): name: Optional[str] = None server_name: Optional[str] = None - type: Literal['mcp_server_tool_result'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["mcp_server_tool_result"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) result: Optional[Union[Result, str]] = Field( - None, description='Tool call result delta.' + None, description="Tool call result delta." ) call_id: Optional[str] = Field( - None, description='ID to match the ID from the function call block.' + None, description="ID to match the ID from the function call block." ) class ErrorEvent(BaseModel): - event_type: Literal['error'] = 'error' + event_type: Literal["error"] = "error" error: Optional[Error] = None event_id: Optional[str] = Field( None, - description='The event_id token to be used to resume the interaction stream, from\nthis event.', + description="The event_id token to be used to resume the interaction stream, from\nthis event.", ) @@ -746,69 +746,69 @@ class Tool( ComputerUse, McpServer, FileSearch, - ] = Field(..., discriminator='type') + ] = Field(..., discriminator="type") class ThoughtContent(BaseModel): signature: Optional[Base64Str] = Field( None, - description='Signature to match the backend source to be part of the generation.', + description="Signature to match the backend source to be part of the generation.", ) - type: Literal['thought'] = Field( - ..., description='Used as the OpenAPI type discriminator for the content oneof.' + type: Literal["thought"] = Field( + ..., description="Used as the OpenAPI type discriminator for the content oneof." ) summary: Optional[ThoughtSummary] = Field( - None, description='A summary of the thought.' + None, description="A summary of the thought." ) class ToolChoice(RootModel[Union[ToolChoiceType, ToolChoiceConfig]]): root: Union[ToolChoiceType, ToolChoiceConfig] = Field( - ..., description='The configuration for tool choice.' + ..., description="The configuration for tool choice." ) class Usage(BaseModel): total_input_tokens: Optional[int] = Field( - None, description='Number of tokens in the prompt (context).' + None, description="Number of tokens in the prompt (context)." ) input_tokens_by_modality: Optional[List[ModalityTokens]] = Field( - None, description='A breakdown of input token usage by modality.' + None, description="A breakdown of input token usage by modality." ) total_cached_tokens: Optional[int] = Field( None, - description='Number of tokens in the cached part of the prompt (the cached content).', + description="Number of tokens in the cached part of the prompt (the cached content).", ) cached_tokens_by_modality: Optional[List[ModalityTokens]] = Field( - None, description='A breakdown of cached token usage by modality.' + None, description="A breakdown of cached token usage by modality." ) total_output_tokens: Optional[int] = Field( - None, description='Total number of tokens across all the generated responses.' + None, description="Total number of tokens across all the generated responses." ) output_tokens_by_modality: Optional[List[ModalityTokens]] = Field( - None, description='A breakdown of output token usage by modality.' + None, description="A breakdown of output token usage by modality." ) total_tool_use_tokens: Optional[int] = Field( - None, description='Number of tokens present in tool-use prompt(s).' + None, description="Number of tokens present in tool-use prompt(s)." ) tool_use_tokens_by_modality: Optional[List[ModalityTokens]] = Field( - None, description='A breakdown of tool-use token usage by modality.' + None, description="A breakdown of tool-use token usage by modality." ) total_reasoning_tokens: Optional[int] = Field( - None, description='Number of tokens of thoughts for thinking models.' + None, description="Number of tokens of thoughts for thinking models." ) total_tokens: Optional[int] = Field( None, - description='Total token count for the interaction request (prompt + responses + other\ninternal tokens).', + description="Total token count for the interaction request (prompt + responses + other\ninternal tokens).", ) class ContentDelta(BaseModel): index: Optional[int] = None - event_type: Literal['content.delta'] = 'content.delta' + event_type: Literal["content.delta"] = "content.delta" event_id: Optional[str] = Field( None, - description='The event_id token to be used to resume the interaction stream, from\nthis event.', + description="The event_id token to be used to resume the interaction stream, from\nthis event.", ) delta: Optional[ Union[ @@ -831,7 +831,7 @@ class ContentDelta(BaseModel): McpServerToolResultDelta, FileSearchResultDelta, ] - ] = Field(None, discriminator='type') + ] = Field(None, discriminator="type") class Content( @@ -875,102 +875,102 @@ class Content( McpServerToolCallContent, McpServerToolResultContent, FileSearchResultContent, - ] = Field(..., description='The content of the response.', discriminator='type') + ] = Field(..., description="The content of the response.", discriminator="type") class Turn(BaseModel): role: Optional[str] = Field( None, - description='The originator of this turn. Must be user for input or model for\nmodel output.', + description="The originator of this turn. Must be user for input or model for\nmodel output.", ) content: Optional[Union[str, List[Content]]] = Field( - None, description='The content of the turn.' + None, description="The content of the turn." ) class GenerationConfig(BaseModel): temperature: Optional[float] = Field( - None, description='Controls the randomness of the output.' + None, description="Controls the randomness of the output." ) top_p: Optional[float] = Field( None, - description='The maximum cumulative probability of tokens to consider when sampling.', + description="The maximum cumulative probability of tokens to consider when sampling.", ) seed: Optional[int] = Field( - None, description='Seed used in decoding for reproducibility.' + None, description="Seed used in decoding for reproducibility." ) stop_sequences: Optional[List[str]] = Field( None, - description='A list of character sequences that will stop output interaction.', + description="A list of character sequences that will stop output interaction.", ) tool_choice: Optional[ToolChoice] = Field( - None, description='The tool choice for the interaction.' + None, description="The tool choice for the interaction." ) thinking_level: Optional[ThinkingLevel] = Field( - None, description='The level of thought tokens that the model should generate.' + None, description="The level of thought tokens that the model should generate." ) thinking_summaries: Optional[ThinkingSummaries] = Field( - None, description='Whether to include thought summaries in the response.' + None, description="Whether to include thought summaries in the response." ) max_output_tokens: Optional[int] = Field( - None, description='The maximum number of tokens to include in the response.' + None, description="The maximum number of tokens to include in the response." ) speech_config: Optional[List[SpeechConfig]] = Field( - None, description='Configuration for speech interaction.' + None, description="Configuration for speech interaction." ) class ContentStart(BaseModel): index: Optional[int] = None content: Optional[Content] = None - event_type: Literal['content.start'] = 'content.start' + event_type: Literal["content.start"] = "content.start" event_id: Optional[str] = Field( None, - description='The event_id token to be used to resume the interaction stream, from\nthis event.', + description="The event_id token to be used to resume the interaction stream, from\nthis event.", ) class Interaction(BaseModel): model: Optional[ModelOption] = Field( - None, description='The name of the `Model` used for generating the interaction.' + None, description="The name of the `Model` used for generating the interaction." ) agent: Optional[AgentOption] = Field( - None, description='The name of the `Agent` used for generating the interaction.' + None, description="The name of the `Agent` used for generating the interaction." ) id: str = Field( ..., - description='Output only. A unique identifier for the interaction completion.', + description="Output only. A unique identifier for the interaction completion.", ) status: Status1 = Field( - ..., description='Output only. The status of the interaction.' + ..., description="Output only. The status of the interaction." ) created: Optional[AwareDatetime] = Field( None, - description='Output only. The time at which the response was created in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).', + description="Output only. The time at which the response was created in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).", ) updated: Optional[AwareDatetime] = Field( None, - description='Output only. The time at which the response was last updated in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).', + description="Output only. The time at which the response was last updated in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).", ) role: Optional[str] = Field( - None, description='Output only. The role of the interaction.' + None, description="Output only. The role of the interaction." ) outputs: Optional[List[Content]] = Field( - None, description='Output only. Responses from the model.' + None, description="Output only. Responses from the model." ) system_instruction: Optional[str] = Field( - None, description='System instruction for the interaction.' + None, description="System instruction for the interaction." ) tools: Optional[List[Tool]] = Field( None, - description='A list of tool declarations the model may call during interaction.', + description="A list of tool declarations the model may call during interaction.", ) background: Optional[bool] = Field( - None, description='Whether to run the model interaction in the background.' + None, description="Whether to run the model interaction in the background." ) - object: Literal['interaction'] = Field( - 'interaction', - description='Output only. The object type of the interaction. Always set to `interaction`.', + object: Literal["interaction"] = Field( + "interaction", + description="Output only. The object type of the interaction. Always set to `interaction`.", ) usage: Optional[Usage] = Field( None, @@ -978,72 +978,72 @@ class Interaction(BaseModel): ) response_modalities: Optional[List[ResponseModality]] = Field( None, - description='The requested modalities of the response (TEXT, IMAGE, AUDIO).', + description="The requested modalities of the response (TEXT, IMAGE, AUDIO).", ) response_format: Optional[Any] = Field( None, - description='Enforces that the generated response is a JSON object that complies with\nthe JSON schema specified in this field.', + description="Enforces that the generated response is a JSON object that complies with\nthe JSON schema specified in this field.", ) response_mime_type: Optional[str] = Field( None, - description='The mime type of the response. This is required if response_format is set.', + description="The mime type of the response. This is required if response_format is set.", ) previous_interaction_id: Optional[str] = Field( - None, description='The ID of the previous interaction, if any.' + None, description="The ID of the previous interaction, if any." ) input: Optional[Union[str, List[Content], List[Turn], Content]] = Field( - None, description='The inputs for the interaction.' + None, description="The inputs for the interaction." ) generation_config: Optional[GenerationConfig] = Field( None, - description='Input only. Configuration parameters for the model interaction.', + description="Input only. Configuration parameters for the model interaction.", ) agent_config: Optional[Union[DynamicAgentConfig, DeepResearchAgentConfig]] = Field( - None, description='Configuration for the agent.', discriminator='type' + None, description="Configuration for the agent.", discriminator="type" ) class CreateModelInteractionParams(BaseModel): model: ModelOption = Field( - ..., description='The name of the `Model` used for generating the interaction.' + ..., description="The name of the `Model` used for generating the interaction." ) stream: Optional[bool] = Field( - None, description='Input only. Whether the interaction will be streamed.' + None, description="Input only. Whether the interaction will be streamed." ) store: Optional[bool] = Field( None, - description='Input only. Whether to store the response and request for later retrieval.', + description="Input only. Whether to store the response and request for later retrieval.", ) id: Optional[str] = Field( None, - description='Output only. A unique identifier for the interaction completion.', + description="Output only. A unique identifier for the interaction completion.", ) status: Optional[Status3] = Field( - None, description='Output only. The status of the interaction.' + None, description="Output only. The status of the interaction." ) created: Optional[AwareDatetime] = Field( None, - description='Output only. The time at which the response was created in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).', + description="Output only. The time at which the response was created in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).", ) updated: Optional[AwareDatetime] = Field( None, - description='Output only. The time at which the response was last updated in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).', + description="Output only. The time at which the response was last updated in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).", ) role: Optional[str] = Field( - None, description='Output only. The role of the interaction.' + None, description="Output only. The role of the interaction." ) outputs: Optional[List[Content]] = Field( - None, description='Output only. Responses from the model.' + None, description="Output only. Responses from the model." ) system_instruction: Optional[str] = Field( - None, description='System instruction for the interaction.' + None, description="System instruction for the interaction." ) tools: Optional[List[Tool]] = Field( None, - description='A list of tool declarations the model may call during interaction.', + description="A list of tool declarations the model may call during interaction.", ) background: Optional[bool] = Field( - None, description='Whether to run the model interaction in the background.' + None, description="Whether to run the model interaction in the background." ) usage: Optional[Usage] = Field( None, @@ -1051,69 +1051,69 @@ class CreateModelInteractionParams(BaseModel): ) response_modalities: Optional[List[ResponseModality]] = Field( None, - description='The requested modalities of the response (TEXT, IMAGE, AUDIO).', + description="The requested modalities of the response (TEXT, IMAGE, AUDIO).", ) response_format: Optional[Any] = Field( None, - description='Enforces that the generated response is a JSON object that complies with\nthe JSON schema specified in this field.', + description="Enforces that the generated response is a JSON object that complies with\nthe JSON schema specified in this field.", ) response_mime_type: Optional[str] = Field( None, - description='The mime type of the response. This is required if response_format is set.', + description="The mime type of the response. This is required if response_format is set.", ) previous_interaction_id: Optional[str] = Field( - None, description='The ID of the previous interaction, if any.' + None, description="The ID of the previous interaction, if any." ) input: Union[str, List[Content], List[Turn], Content] = Field( - ..., description='The inputs for the interaction.' + ..., description="The inputs for the interaction." ) generation_config: Optional[GenerationConfig] = Field( None, - description='Input only. Configuration parameters for the model interaction.', + description="Input only. Configuration parameters for the model interaction.", ) class CreateAgentInteractionParams(BaseModel): agent: AgentOption = Field( - ..., description='The name of the `Agent` used for generating the interaction.' + ..., description="The name of the `Agent` used for generating the interaction." ) stream: Optional[bool] = Field( - None, description='Input only. Whether the interaction will be streamed.' + None, description="Input only. Whether the interaction will be streamed." ) store: Optional[bool] = Field( None, - description='Input only. Whether to store the response and request for later retrieval.', + description="Input only. Whether to store the response and request for later retrieval.", ) id: Optional[str] = Field( None, - description='Output only. A unique identifier for the interaction completion.', + description="Output only. A unique identifier for the interaction completion.", ) status: Optional[Status3] = Field( - None, description='Output only. The status of the interaction.' + None, description="Output only. The status of the interaction." ) created: Optional[AwareDatetime] = Field( None, - description='Output only. The time at which the response was created in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).', + description="Output only. The time at which the response was created in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).", ) updated: Optional[AwareDatetime] = Field( None, - description='Output only. The time at which the response was last updated in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).', + description="Output only. The time at which the response was last updated in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).", ) role: Optional[str] = Field( - None, description='Output only. The role of the interaction.' + None, description="Output only. The role of the interaction." ) outputs: Optional[List[Content]] = Field( - None, description='Output only. Responses from the model.' + None, description="Output only. Responses from the model." ) system_instruction: Optional[str] = Field( - None, description='System instruction for the interaction.' + None, description="System instruction for the interaction." ) tools: Optional[List[Tool]] = Field( None, - description='A list of tool declarations the model may call during interaction.', + description="A list of tool declarations the model may call during interaction.", ) background: Optional[bool] = Field( - None, description='Whether to run the model interaction in the background.' + None, description="Whether to run the model interaction in the background." ) usage: Optional[Usage] = Field( None, @@ -1121,33 +1121,33 @@ class CreateAgentInteractionParams(BaseModel): ) response_modalities: Optional[List[ResponseModality]] = Field( None, - description='The requested modalities of the response (TEXT, IMAGE, AUDIO).', + description="The requested modalities of the response (TEXT, IMAGE, AUDIO).", ) response_format: Optional[Any] = Field( None, - description='Enforces that the generated response is a JSON object that complies with\nthe JSON schema specified in this field.', + description="Enforces that the generated response is a JSON object that complies with\nthe JSON schema specified in this field.", ) response_mime_type: Optional[str] = Field( None, - description='The mime type of the response. This is required if response_format is set.', + description="The mime type of the response. This is required if response_format is set.", ) previous_interaction_id: Optional[str] = Field( - None, description='The ID of the previous interaction, if any.' + None, description="The ID of the previous interaction, if any." ) input: Union[str, List[Content], List[Turn], Content] = Field( - ..., description='The inputs for the interaction.' + ..., description="The inputs for the interaction." ) agent_config: Optional[Union[DynamicAgentConfig, DeepResearchAgentConfig]] = Field( - None, description='Configuration for the agent.', discriminator='type' + None, description="Configuration for the agent.", discriminator="type" ) class InteractionEvent(BaseModel): - event_type: Literal['interaction.start', 'interaction.complete'] + event_type: Literal["interaction.start", "interaction.complete"] interaction: Optional[Interaction] = None event_id: Optional[str] = Field( None, - description='The event_id token to be used to resume the interaction stream, from\nthis event.', + description="The event_id token to be used to resume the interaction stream, from\nthis event.", ) @@ -1170,13 +1170,13 @@ class InteractionSseEvent( ContentDelta, ContentStop, ErrorEvent, - ] = Field(..., discriminator='event_type') + ] = Field(..., discriminator="event_type") # ============================================================ # LiteLLM-specific types (added manually after generation) # ============================================================ -# +# # When regenerating this file, copy these types to the end. # See README.md for regeneration instructions. @@ -1191,9 +1191,10 @@ InteractionInput = Union[str, Content, List[Content], List[Turn]] class InteractionsAPIResponse(BaseLiteLLMOpenAIResponseObject): """ Response from the Interactions API. - + Wraps the API response with LiteLLM-specific hidden params. """ + id: Optional[str] = None object: Optional[str] = "interaction" model: Optional[str] = None @@ -1204,19 +1205,20 @@ class InteractionsAPIResponse(BaseLiteLLMOpenAIResponseObject): role: Optional[str] = None outputs: Optional[List[Dict[str, Any]]] = None usage: Optional[Dict[str, Any]] = None - + _hidden_params: dict = PrivateAttr(default_factory=dict) class InteractionsAPIStreamingResponse(BaseLiteLLMOpenAIResponseObject): """ Streaming response chunk from the Interactions API. - + Event types per OpenAPI spec: - interaction.start, interaction.status_update, interaction.complete - content.start, content.delta, content.stop - error """ + event_type: Optional[str] = None id: Optional[str] = None object: Optional[str] = "interaction" @@ -1229,23 +1231,25 @@ class InteractionsAPIStreamingResponse(BaseLiteLLMOpenAIResponseObject): outputs: Optional[List[Dict[str, Any]]] = None usage: Optional[Dict[str, Any]] = None delta: Optional[Dict[str, Any]] = None - + _hidden_params: dict = PrivateAttr(default_factory=dict) class DeleteInteractionResult(BaseLiteLLMOpenAIResponseObject): """Result of deleting an interaction.""" + success: bool = True id: Optional[str] = None - + _hidden_params: dict = PrivateAttr(default_factory=dict) class CancelInteractionResult(BaseLiteLLMOpenAIResponseObject): """Result of cancelling an interaction.""" + id: Optional[str] = None status: Optional[str] = None - + _hidden_params: dict = PrivateAttr(default_factory=dict) diff --git a/litellm/types/llms/aiml.py b/litellm/types/llms/aiml.py index 3d42518b8b..d5781add18 100644 --- a/litellm/types/llms/aiml.py +++ b/litellm/types/llms/aiml.py @@ -5,6 +5,7 @@ from typing_extensions import TypedDict class AimlImageSize(TypedDict, total=False): """Custom image size specification for AI/ML API""" + width: int # Must be multiple of 32, min 256, max 1440 height: int # Must be multiple of 32, min 256, max 1440 @@ -12,12 +13,15 @@ class AimlImageSize(TypedDict, total=False): class AimlImageGenerationRequestParams(TypedDict, total=False): """ TypedDict for AI/ML flux image generation request parameters. - + Based on AI/ML API docs: https://api.aimlapi.com/v1/images/generations """ + model: str # Required: flux-pro/v1.1 prompt: str # Required: Text prompt (max 4000 chars) - image_size: Union[AimlImageSize, str] # Custom size or predefined: square_hd, square, portrait_4_3, portrait_16_9, landscape_4_3, landscape_16_9 + image_size: Union[ + AimlImageSize, str + ] # Custom size or predefined: square_hd, square, portrait_4_3, portrait_16_9, landscape_4_3, landscape_16_9 safety_tolerance: Optional[str] # 1-6, default 2 (1=strict, 6=permissive) output_format: Optional[str] # jpeg or png, default jpeg num_images: Optional[int] # 1-4, default 1 diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index 5b8044911e..478fcbdbd1 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -38,6 +38,7 @@ class AnthropicOutputSchema(TypedDict, total=False): class AnthropicOutputConfig(TypedDict, total=False): """Configuration for controlling Claude's output behavior.""" + effort: Literal["high", "medium", "low"] @@ -109,12 +110,14 @@ class AnthropicMemoryTool(TypedDict, total=False): class AnthropicToolSearchToolRegex(TypedDict, total=False): """Tool search tool using regex patterns for tool discovery.""" + type: Required[Literal["tool_search_tool_regex_20251119"]] name: Required[str] class AnthropicToolSearchToolBM25(TypedDict, total=False): """Tool search tool using BM25 algorithm for tool discovery.""" + type: Required[Literal["tool_search_tool_bm25_20251119"]] name: Required[str] cache_control: Optional[Union[dict, ChatCompletionCachedContent]] @@ -125,17 +128,20 @@ class AnthropicToolSearchToolBM25(TypedDict, total=False): class ToolReference(TypedDict, total=False): """Reference to a tool that should be expanded from deferred tools.""" + type: Required[Literal["tool_reference"]] tool_name: Required[str] class DirectToolCaller(TypedDict, total=False): """Indicates a tool was called directly by Claude.""" + type: Required[Literal["direct"]] class CodeExecutionToolCaller(TypedDict, total=False): """Indicates a tool was called programmatically from code execution.""" + type: Required[Literal["code_execution_20250825"]] tool_id: Required[str] # ID of the code execution tool that made the call @@ -145,6 +151,7 @@ ToolCaller = Union[DirectToolCaller, CodeExecutionToolCaller] class AnthropicContainer(TypedDict, total=False): """Container metadata for code execution.""" + id: Required[str] expires_at: Optional[str] # ISO 8601 timestamp @@ -359,10 +366,14 @@ class AnthropicMessagesRequestOptionalParams(TypedDict, total=False): top_p: Optional[float] mcp_servers: Optional[List[AnthropicMcpServerTool]] context_management: Optional[Dict[str, Any]] - container: Optional[Dict[str, Any]] # Container config with skills for code execution + container: Optional[ + Dict[str, Any] + ] # Container config with skills for code execution output_format: Optional[AnthropicOutputSchema] # Structured outputs support speed: Optional[str] # Fast mode support for Opus models - output_config: Optional[AnthropicOutputConfig] # Configuration for Claude's output behavior + output_config: Optional[ + AnthropicOutputConfig + ] # Configuration for Claude's output behavior cache_control: Optional[Dict[str, Any]] # Automatic prompt caching @@ -543,7 +554,7 @@ class AnthropicResponseContentBlockToolUse(BaseModel): input: dict provider_specific_fields: Optional[Dict[str, Any]] = None - model_config = ConfigDict(extra="allow") # Allow provider_specific_fields + model_config = ConfigDict(extra="allow") # Allow provider_specific_fields class AnthropicResponseContentBlockThinking(BaseModel): @@ -634,6 +645,7 @@ class ANTHROPIC_BETA_HEADER_VALUES(str, Enum): """ Known beta header values for Anthropic. """ + WEB_FETCH_2025_09_10 = "web-fetch-2025-09-10" WEB_SEARCH_2025_03_05 = "web-search-2025-03-05" CONTEXT_MANAGEMENT_2025_06_27 = "context-management-2025-06-27" @@ -653,4 +665,4 @@ ANTHROPIC_EFFORT_BETA_HEADER = "effort-2025-11-24" ANTHROPIC_OAUTH_TOKEN_PREFIX = "sk-ant-oat" ANTHROPIC_OAUTH_BETA_HEADER = "oauth-2025-04-20" -ANTHROPIC_PROMPT_CACHING_SCOPE_BETA_HEADER = "prompt-caching-scope-2026-01-05" \ No newline at end of file +ANTHROPIC_PROMPT_CACHING_SCOPE_BETA_HEADER = "prompt-caching-scope-2026-01-05" diff --git a/litellm/types/llms/anthropic_messages/anthropic_request.py b/litellm/types/llms/anthropic_messages/anthropic_request.py index 00b2590ba7..4f31e9a509 100644 --- a/litellm/types/llms/anthropic_messages/anthropic_request.py +++ b/litellm/types/llms/anthropic_messages/anthropic_request.py @@ -9,5 +9,5 @@ class AnthropicMetadata(BaseModel): https://docs.anthropic.com/en/api/messages#body-metadata-user-id """ - user_id: Optional[str] = None + user_id: Optional[str] = None diff --git a/litellm/types/llms/anthropic_skills.py b/litellm/types/llms/anthropic_skills.py index c7ccf2faab..2225788849 100644 --- a/litellm/types/llms/anthropic_skills.py +++ b/litellm/types/llms/anthropic_skills.py @@ -156,4 +156,3 @@ class DeleteSkillVersionResponse(BaseModel): deleted: bool """Whether the version was successfully deleted""" - diff --git a/litellm/types/llms/anthropic_tool_search.py b/litellm/types/llms/anthropic_tool_search.py index d8656ce8bb..7cdaec2e7c 100644 --- a/litellm/types/llms/anthropic_tool_search.py +++ b/litellm/types/llms/anthropic_tool_search.py @@ -30,7 +30,5 @@ def get_tool_search_beta_header(custom_llm_provider: str) -> str: Get the tool search beta header for a given provider. """ return TOOL_SEARCH_BETA_HEADER_BY_PROVIDER.get( - custom_llm_provider, - TOOL_SEARCH_BETA_HEADER_ANTHROPIC + custom_llm_provider, TOOL_SEARCH_BETA_HEADER_ANTHROPIC ) - diff --git a/litellm/types/llms/base.py b/litellm/types/llms/base.py index ec0d3ed95d..625d044172 100644 --- a/litellm/types/llms/base.py +++ b/litellm/types/llms/base.py @@ -74,4 +74,4 @@ class HiddenParams(OpenAIObject): # Override model_dump to include private attributes data = super().model_dump(**kwargs) data["_response_ms"] = self._response_ms - return data \ No newline at end of file + return data diff --git a/litellm/types/llms/bedrock_agentcore.py b/litellm/types/llms/bedrock_agentcore.py index 49c3bfb2d5..cd6b75f2ac 100644 --- a/litellm/types/llms/bedrock_agentcore.py +++ b/litellm/types/llms/bedrock_agentcore.py @@ -132,4 +132,3 @@ class AgentCoreParsedResponse(TypedDict): content: str usage: Optional[AgentCoreUsage] final_message: Optional[AgentCoreMessage] - diff --git a/litellm/types/llms/gemini.py b/litellm/types/llms/gemini.py index e29a2cc19a..9e3fea1bbb 100644 --- a/litellm/types/llms/gemini.py +++ b/litellm/types/llms/gemini.py @@ -158,82 +158,91 @@ from pydantic import BaseModel class GeminiImageGenerationInstance(TypedDict): """Instance data for Gemini image generation request""" + prompt: str class GeminiImageGenerationParameters(BaseModel): """Parameters for Gemini image generation request""" + sampleCount: Optional[int] = None """Number of images to generate (maps to OpenAI 'n' parameter)""" - + aspectRatio: Optional[str] = None """Aspect ratio for generated images (e.g., '1:1', '16:9', '9:16', '4:3', '3:4')""" - + personGeneration: Optional[str] = None """Controls person generation in images""" - + # Additional parameters that might be passed through background: Optional[str] = None """Background specification""" - + input_fidelity: Optional[str] = None """Input fidelity specification""" - + moderation: Optional[str] = None """Moderation settings""" - + output_compression: Optional[str] = None """Output compression settings""" - + output_format: Optional[str] = None """Output format specification""" - + quality: Optional[str] = None """Quality settings""" - + response_format: Optional[str] = None """Response format specification""" - + style: Optional[str] = None """Style specification""" - + user: Optional[str] = None """User specification""" class GeminiImageGenerationRequest(BaseModel): """Complete request body for Gemini image generation""" + instances: List[GeminiImageGenerationInstance] parameters: GeminiImageGenerationParameters class GeminiGeneratedImage(TypedDict): """Individual generated image data from Gemini response""" + bytesBase64Encoded: str """Base64 encoded image data""" class GeminiImageGenerationPrediction(TypedDict): """Prediction object containing generated images""" + generatedImages: List[GeminiGeneratedImage] class GeminiImageGenerationResponse(TypedDict): """Complete response body from Gemini image generation API""" + predictions: List[GeminiImageGenerationPrediction] + # Video Generation Types class GeminiVideoGenerationInstance(TypedDict): """Instance data for Gemini video generation request""" + prompt: str class GeminiVideoGenerationParameters(BaseModel): """ Parameters for Gemini video generation request. - + See: Veo 3/3.1 parameter guide. """ + aspectRatio: Optional[str] = None """Aspect ratio for generated video (e.g., '16:9', '9:16').""" @@ -286,6 +295,7 @@ class GeminiVideoGenerationParameters(BaseModel): class GeminiVideoGenerationRequest(BaseModel): """Complete request body for Gemini video generation""" + instances: List[GeminiVideoGenerationInstance] parameters: Optional[GeminiVideoGenerationParameters] = None @@ -293,30 +303,35 @@ class GeminiVideoGenerationRequest(BaseModel): # Video Generation Operation Response Types class GeminiVideoUri(BaseModel): """Video URI in the generated sample""" + uri: str """File URI of the generated video (e.g., 'files/abc123...')""" class GeminiGeneratedVideoSample(BaseModel): """Individual generated video sample""" + video: GeminiVideoUri """Video object containing the URI""" class GeminiGenerateVideoResponse(BaseModel): """Generate video response containing the samples""" + generatedSamples: List[GeminiGeneratedVideoSample] """List of generated video samples""" class GeminiOperationResponse(BaseModel): """Response object in the operation when done""" + generateVideoResponse: GeminiGenerateVideoResponse """Video generation response""" class GeminiOperationMetadata(BaseModel): """Metadata for the operation""" + createTime: Optional[str] = None """Creation timestamp""" model: Optional[str] = None @@ -326,20 +341,21 @@ class GeminiOperationMetadata(BaseModel): class GeminiLongRunningOperationResponse(BaseModel): """ Complete response for a long-running operation. - + Used when polling operation status and extracting results. """ + name: str """Operation name (e.g., 'operations/generate_1234567890')""" - + done: bool = False """Whether the operation is complete""" - + metadata: Optional[GeminiOperationMetadata] = None """Operation metadata""" - + response: Optional[GeminiOperationResponse] = None """Response object when operation is complete""" - + error: Optional[Dict[str, Any]] = None """Error details if operation failed""" diff --git a/litellm/types/llms/langgraph.py b/litellm/types/llms/langgraph.py index cdf5d67b51..9286ca463e 100644 --- a/litellm/types/llms/langgraph.py +++ b/litellm/types/llms/langgraph.py @@ -65,4 +65,3 @@ class LangGraphParsedResponse(TypedDict): content: str role: str usage: Optional[Dict[str, int]] - diff --git a/litellm/types/llms/oci.py b/litellm/types/llms/oci.py index cb1dd39143..e041810158 100644 --- a/litellm/types/llms/oci.py +++ b/litellm/types/llms/oci.py @@ -112,6 +112,7 @@ class OCIServingMode(BaseModel): endpointId: Optional[str] = None modelId: Optional[str] = None + class OCICompletionPayload(BaseModel): """Pydantic model for the complete OCI chat request body.""" @@ -194,6 +195,7 @@ class OCIStreamChunk(BaseModel): # --- Cohere-Specific Models --- + class CohereStreamChunk(BaseModel): """Model for a single SSE event chunk from OCI Cohere API.""" @@ -204,6 +206,7 @@ class CohereStreamChunk(BaseModel): pad: Optional[str] = None index: Optional[int] = None + class CohereMessage(BaseModel): """Base model for Cohere messages.""" @@ -305,7 +308,13 @@ class CohereChatRequest(BaseModel): seed: Optional[int] = None tools: Optional[List[CohereTool]] = None toolChoice: Optional[Union[str, Dict[str, Any]]] = None - responseFormat: Optional[Union[CohereResponseTextFormat, CohereResponseJSONSchemaFormat, CohereResponseFormat]] = None + responseFormat: Optional[ + Union[ + CohereResponseTextFormat, + CohereResponseJSONSchemaFormat, + CohereResponseFormat, + ] + ] = None preambleOverride: Optional[str] = None documents: Optional[List[Dict[str, Any]]] = None searchQueriesOnly: Optional[bool] = None @@ -355,7 +364,9 @@ class CohereChatResponse(BaseModel): # Required fields text: str apiFormat: Literal["COHERE"] = "COHERE" - finishReason: Literal["COMPLETE", "ERROR_TOXIC", "ERROR_LIMIT", "ERROR", "USER_CANCEL", "MAX_TOKENS"] + finishReason: Literal[ + "COMPLETE", "ERROR_TOXIC", "ERROR_LIMIT", "ERROR", "USER_CANCEL", "MAX_TOKENS" + ] # Optional fields chatHistory: Optional[List[CohereMessage]] = None diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index be2792e859..0ca48611e1 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -55,7 +55,7 @@ from openai.types.responses.response import ( # Handle OpenAI SDK version compatibility for Text type try: - from openai.types.responses.response_create_params import ( Text as ResponseText ) # type: ignore[attr-defined] # fmt: skip # isort: skip + from openai.types.responses.response_create_params import Text as ResponseText # type: ignore[attr-defined] # fmt: skip # isort: skip except (ImportError, AttributeError): # Fall back to the concrete config type available in all SDK versions from openai.types.responses.response_text_config_param import ( @@ -287,6 +287,7 @@ OpenAIFilesPurpose = Literal[ "fine-tune-results", "vision", "user_data", + "messages", ] @@ -352,7 +353,7 @@ class OpenAIFileObject(BaseModel): return self.dict() -CREATE_FILE_REQUESTS_PURPOSE = Literal["assistants", "batch", "fine-tune"] +CREATE_FILE_REQUESTS_PURPOSE = Literal["assistants", "batch", "fine-tune", "messages"] # File expiration policy @@ -373,11 +374,11 @@ class FileExpiresAfter(TypedDict): class CreateFileRequest(TypedDict, total=False): """ CreateFileRequest - Used by Assistants API, Batches API, and Fine-Tunes API + Used by Assistants API, Batches API, Fine-Tunes API, and Anthropic Files API Required Params: file: FileTypes - purpose: Literal['assistants', 'batch', 'fine-tune'] + purpose: Literal['assistants', 'batch', 'fine-tune', 'messages'] Optional Params: expires_after: Optional[FileExpiresAfter] - The expiration policy for a file @@ -663,7 +664,9 @@ class ChatCompletionFileObjectFile(TypedDict, total=False): filename: str format: str detail: str # For video/image resolution control (low, medium, high, ultra_high) - video_metadata: Dict[str, Any] # For video-specific metadata (fps, start_offset, end_offset) + video_metadata: Dict[ + str, Any + ] # For video-specific metadata (fps, start_offset, end_offset) class ChatCompletionFileObject(TypedDict): @@ -966,16 +969,14 @@ class OpenAIChatCompletionChunk(ChatCompletionChunk): class Hyperparameters(BaseModel): batch_size: Optional[Union[str, int]] = None # "Number of examples in each batch." - learning_rate_multiplier: Optional[Union[str, float]] = ( - None # Scaling factor for the learning rate - ) - n_epochs: Optional[Union[str, int]] = ( - None # "The number of epochs to train the model for" - ) - - model_config = { - "extra": "allow" - } + learning_rate_multiplier: Optional[ + Union[str, float] + ] = None # Scaling factor for the learning rate + n_epochs: Optional[ + Union[str, int] + ] = None # "The number of epochs to train the model for" + + model_config = {"extra": "allow"} class FineTuningJobCreate(BaseModel): @@ -1002,18 +1003,18 @@ class FineTuningJobCreate(BaseModel): model: str # "The name of the model to fine-tune." training_file: str # "The ID of an uploaded file that contains training data." - hyperparameters: Optional[Hyperparameters] = ( - None # "The hyperparameters used for the fine-tuning job." - ) - suffix: Optional[str] = ( - None # "A string of up to 18 characters that will be added to your fine-tuned model name." - ) - validation_file: Optional[str] = ( - None # "The ID of an uploaded file that contains validation data." - ) - integrations: Optional[List[str]] = ( - None # "A list of integrations to enable for your fine-tuning job." - ) + hyperparameters: Optional[ + Hyperparameters + ] = None # "The hyperparameters used for the fine-tuning job." + suffix: Optional[ + str + ] = None # "A string of up to 18 characters that will be added to your fine-tuned model name." + validation_file: Optional[ + str + ] = None # "The ID of an uploaded file that contains validation data." + integrations: Optional[ + List[str] + ] = None # "A list of integrations to enable for your fine-tuning job." seed: Optional[int] = None # "The seed controls the reproducibility of the job." @@ -1056,8 +1057,7 @@ OpenAIImageGenerationOptionalParams = Literal[ OpenAIImageEditOptionalParams = Literal[ "background", "n", - "mask" - "output_compression", + "mask" "output_compression", "output_format", "quality", "partial_images", @@ -1067,6 +1067,7 @@ OpenAIImageEditOptionalParams = Literal[ "user", ] + class ComputerToolParam(TypedDict, total=False): display_height: Required[float] """The height of the computer display.""" @@ -1302,8 +1303,7 @@ class ResponsesAPIResponse(BaseLiteLLMOpenAIResponseObject): { k: v for k, v in item.items() - if v is not None - or k not in ("status", "content", "encrypted_content") + if v is not None or k not in ("status", "content", "encrypted_content") } if isinstance(item, dict) and item.get("type") == "reasoning" else item @@ -2118,7 +2118,15 @@ class OpenAIBatchResult(TypedDict, total=False): OpenAIChatCompletionFinishReason = Literal[ - "stop", "content_filter", "function_call", "tool_calls", "length" + "stop", + "content_filter", + "function_call", + "tool_calls", + "length", + "guardrail_intervened", + "eos", + "finish_reason_unspecified", + "malformed_function_call", # last 2 are vertex ai specific, guardrail_intervened is bedrock specific ] diff --git a/litellm/types/llms/openai_evals.py b/litellm/types/llms/openai_evals.py index 4f475665eb..431dd34647 100644 --- a/litellm/types/llms/openai_evals.py +++ b/litellm/types/llms/openai_evals.py @@ -42,7 +42,9 @@ class DataSourceConfigStoredCompletions(TypedDict, total=False): """Optional metadata for filtering stored completions""" -DataSourceConfig = Union[DataSourceConfigCustom, DataSourceConfigLogs, DataSourceConfigStoredCompletions] +DataSourceConfig = Union[ + DataSourceConfigCustom, DataSourceConfigLogs, DataSourceConfigStoredCompletions +] class LLMAsJudgeGraderConfig(TypedDict, total=False): @@ -78,7 +80,9 @@ class CustomGraderConfig(TypedDict, total=False): """ID of the custom grading function""" -GraderConfig = Union[LLMAsJudgeGraderConfig, GroundTruthGraderConfig, CustomGraderConfig] +GraderConfig = Union[ + LLMAsJudgeGraderConfig, GroundTruthGraderConfig, CustomGraderConfig +] class CreateEvalRequest(TypedDict, total=False): diff --git a/litellm/types/llms/recraft.py b/litellm/types/llms/recraft.py index c5345b855f..35e4101ee0 100644 --- a/litellm/types/llms/recraft.py +++ b/litellm/types/llms/recraft.py @@ -20,9 +20,10 @@ class RecraftImageGenerationRequestParams(TypedDict, total=False): class RecraftImageEditRequestParams(TypedDict, total=False): """ TypedDict for Recraft image edit request parameters. - + Based on Recraft API docs: https://www.recraft.ai/docs#image-to-image """ + prompt: str # required - A text description of areas to change. Max 1000 bytes strength: float # required - Defines difference with original image, [0, 1] model: Optional[str] # The model to use, default is recraftv3 diff --git a/litellm/types/llms/stability.py b/litellm/types/llms/stability.py index 7dd92e380c..c439a3b59e 100644 --- a/litellm/types/llms/stability.py +++ b/litellm/types/llms/stability.py @@ -18,9 +18,12 @@ class StabilityImageGenerationRequest(TypedDict, total=False): - /v2beta/stable-image/generate/ultra - /v2beta/stable-image/generate/core """ + prompt: str # Required - text prompt for image generation negative_prompt: Optional[str] # What to avoid in the image - aspect_ratio: Optional[str] # e.g., "1:1", "16:9", "9:16", "4:3", "3:4", "21:9", "9:21" + aspect_ratio: Optional[ + str + ] # e.g., "1:1", "16:9", "9:16", "4:3", "3:4", "21:9", "9:21" seed: Optional[int] # Random seed for reproducibility (0 to 4294967294) output_format: Optional[Literal["jpeg", "png", "webp"]] # Output format model: Optional[str] # Model variant (e.g., "sd3.5-large", "sd3.5-medium") @@ -29,18 +32,22 @@ class StabilityImageGenerationRequest(TypedDict, total=False): strength: Optional[float] # How much to transform the image (0-1) style_preset: Optional[str] # Style preset name + class StabilityImageEditRequest(StabilityImageGenerationRequest): """ Request parameters for Stability AI image edit endpoint. Endpoint: /v2beta/stable-image/edit/inpaint """ + mask: Optional[str] # Base64-encoded mask (white = edit, black = keep) + class StabilityImageGenerationResponse(TypedDict, total=False): """ Response from Stability AI image generation endpoints. """ + image: str # Base64-encoded image finish_reason: str # "SUCCESS", "CONTENT_FILTERED", etc. seed: int # The seed used for generation @@ -55,6 +62,7 @@ class StabilityUpscaleRequest(TypedDict, total=False): - /v2beta/stable-image/upscale/conservative - /v2beta/stable-image/upscale/creative """ + image: str # Required - Base64-encoded image to upscale prompt: Optional[str] # Text prompt (required for creative upscale) negative_prompt: Optional[str] # What to avoid @@ -69,6 +77,7 @@ class StabilityInpaintRequest(TypedDict, total=False): Endpoint: /v2beta/stable-image/edit/inpaint """ + image: str # Required - Base64-encoded image to edit prompt: str # Required - Description of desired changes mask: Optional[str] # Base64-encoded mask (white = edit, black = keep) @@ -84,6 +93,7 @@ class StabilityOutpaintRequest(TypedDict, total=False): Endpoint: /v2beta/stable-image/edit/outpaint """ + image: str # Required - Base64-encoded image to expand prompt: Optional[str] # Description of content to generate negative_prompt: Optional[str] # What to avoid @@ -102,6 +112,7 @@ class StabilityEraseRequest(TypedDict, total=False): Endpoint: /v2beta/stable-image/edit/erase """ + image: str # Required - Base64-encoded image mask: Optional[str] # Base64-encoded mask (white = erase) seed: Optional[int] # Random seed @@ -115,6 +126,7 @@ class StabilitySearchReplaceRequest(TypedDict, total=False): Endpoint: /v2beta/stable-image/edit/search-and-replace """ + image: str # Required - Base64-encoded image prompt: str # Required - Description of object to add search_prompt: str # Required - Description of object to find and replace @@ -130,8 +142,11 @@ class StabilityRemoveBackgroundRequest(TypedDict, total=False): Endpoint: /v2beta/stable-image/edit/remove-background """ + image: str # Required - Base64-encoded image - output_format: Optional[Literal["png", "webp"]] # Output format (no jpeg - needs transparency) + output_format: Optional[ + Literal["png", "webp"] + ] # Output format (no jpeg - needs transparency) class StabilityControlRequest(TypedDict, total=False): @@ -143,6 +158,7 @@ class StabilityControlRequest(TypedDict, total=False): - /v2beta/stable-image/control/structure - /v2beta/stable-image/control/style """ + image: str # Required - Base64-encoded control image (sketch/structure/style reference) prompt: str # Required - Description of desired output negative_prompt: Optional[str] # What to avoid @@ -155,6 +171,7 @@ class StabilityEditResponse(TypedDict, total=False): """ Response from Stability AI edit/upscale/control endpoints. """ + image: str # Base64-encoded result image finish_reason: str # "SUCCESS", "CONTENT_FILTERED", etc. seed: int # The seed used diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index d94ba5f805..201854369f 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -174,7 +174,9 @@ class GeminiThinkingConfig(TypedDict, total=False): GeminiResponseModalities = Literal["TEXT", "IMAGE", "AUDIO", "VIDEO"] -GeminiImageAspectRatio = Literal["1:1", "2:3", "3:2", "3:4", "4:3", "9:16", "16:9", "21:9"] +GeminiImageAspectRatio = Literal[ + "1:1", "2:3", "3:2", "3:4", "4:3", "9:16", "16:9", "21:9" +] GeminiImageSize = Literal["1K", "2K", "4K"] @@ -220,6 +222,7 @@ class GenerationConfig(TypedDict, total=False): class VertexToolName(str, Enum): """Enum for Vertex AI tool field names.""" + GOOGLE_SEARCH = "googleSearch" GOOGLE_SEARCH_RETRIEVAL = "googleSearchRetrieval" ENTERPRISE_WEB_SEARCH = "enterpriseWebSearch" @@ -265,7 +268,9 @@ class UsageMetadata(TypedDict, total=False): cacheTokensDetails: List[PromptTokensDetails] thoughtsTokenCount: int responseTokensDetails: List[PromptTokensDetails] - candidatesTokensDetails: List[PromptTokensDetails] # Alternative key name used in some responses + candidatesTokensDetails: List[ + PromptTokensDetails + ] # Alternative key name used in some responses class TokenCountDetailsResponse(TypedDict): diff --git a/litellm/types/llms/vertex_ai_text_to_speech.py b/litellm/types/llms/vertex_ai_text_to_speech.py index e65b75356b..8ac3e35216 100644 --- a/litellm/types/llms/vertex_ai_text_to_speech.py +++ b/litellm/types/llms/vertex_ai_text_to_speech.py @@ -12,9 +12,10 @@ from typing_extensions import TypedDict class VertexTextToSpeechInput(TypedDict, total=False): """ Input for Vertex AI Text-to-Speech synthesis. - + Exactly one of text or ssml must be provided. """ + text: Optional[str] ssml: Optional[str] @@ -22,11 +23,12 @@ class VertexTextToSpeechInput(TypedDict, total=False): class VertexTextToSpeechVoice(TypedDict, total=False): """ Voice configuration for Vertex AI Text-to-Speech. - + Attributes: languageCode: The language code (e.g., "en-US", "de-DE") name: The voice name (e.g., "en-US-Studio-O", "en-US-Wavenet-D") """ + languageCode: str name: str @@ -34,11 +36,12 @@ class VertexTextToSpeechVoice(TypedDict, total=False): class VertexTextToSpeechAudioConfig(TypedDict, total=False): """ Audio configuration for Vertex AI Text-to-Speech. - + Attributes: audioEncoding: The audio encoding format (e.g., "LINEAR16", "MP3", "OGG_OPUS") speakingRate: The speaking rate (0.25 to 4.0, default "1") """ + audioEncoding: str speakingRate: str @@ -46,9 +49,10 @@ class VertexTextToSpeechAudioConfig(TypedDict, total=False): class VertexTextToSpeechRequest(TypedDict, total=False): """ Request body for Vertex AI Text-to-Speech API. - + Reference: https://cloud.google.com/text-to-speech/docs/reference/rest/v1/text/synthesize """ + input: VertexTextToSpeechInput voice: VertexTextToSpeechVoice audioConfig: Optional[VertexTextToSpeechAudioConfig] diff --git a/litellm/types/llms/xai.py b/litellm/types/llms/xai.py index 9de8dfc7ad..8500e218d8 100644 --- a/litellm/types/llms/xai.py +++ b/litellm/types/llms/xai.py @@ -3,21 +3,26 @@ from typing import List, Literal, Optional, TypedDict class XAIWebSearchFilters(TypedDict, total=False): """Filters for XAI web search tool""" + allowed_domains: Optional[List[str]] # Max 5 domains excluded_domains: Optional[List[str]] # Max 5 domains - + + class XAIWebSearchTool(TypedDict, total=False): """XAI web search tool configuration""" + type: Literal["web_search"] filters: Optional[XAIWebSearchFilters] enable_image_understanding: Optional[bool] + class XAIXSearchTool(TypedDict, total=False): """XAI X (Twitter) search tool configuration""" + type: Literal["x_search"] allowed_x_handles: Optional[List[str]] # Max 10 handles excluded_x_handles: Optional[List[str]] # Max 10 handles from_date: Optional[str] # ISO8601 format: YYYY-MM-DD to_date: Optional[str] # ISO8601 format: YYYY-MM-DD enable_image_understanding: Optional[bool] - enable_video_understanding: Optional[bool] \ No newline at end of file + enable_video_understanding: Optional[bool] diff --git a/litellm/types/management_endpoints/__init__.py b/litellm/types/management_endpoints/__init__.py index edcc2f5133..5c5bcb2e75 100644 --- a/litellm/types/management_endpoints/__init__.py +++ b/litellm/types/management_endpoints/__init__.py @@ -21,4 +21,3 @@ __all__ = [ "REDIS_TYPE_DESCRIPTIONS", "CacheSettingsField", ] - diff --git a/litellm/types/management_endpoints/cache_settings_endpoints.py b/litellm/types/management_endpoints/cache_settings_endpoints.py index d2a8a39e4f..fd68f43e7b 100644 --- a/litellm/types/management_endpoints/cache_settings_endpoints.py +++ b/litellm/types/management_endpoints/cache_settings_endpoints.py @@ -13,10 +13,14 @@ class CacheSettingsField(BaseModel): field_value: Any field_description: str field_default: Any = None - options: Optional[List[str]] = None # For fields with predefined options/enum values + options: Optional[ + List[str] + ] = None # For fields with predefined options/enum values ui_field_name: str # User-friendly display name link: Optional[str] = None # Documentation link for the field - redis_type: Optional[str] = None # Which Redis type this field applies to (node, cluster, sentinel) + redis_type: Optional[ + str + ] = None # Which Redis type this field applies to (node, cluster, sentinel) # Redis type descriptions @@ -207,4 +211,3 @@ CACHE_SETTINGS_FIELDS: List[CacheSettingsField] = [ redis_type=None, ), ] - diff --git a/litellm/types/management_endpoints/router_settings_endpoints.py b/litellm/types/management_endpoints/router_settings_endpoints.py index 5024fe39b3..4f09b7da85 100644 --- a/litellm/types/management_endpoints/router_settings_endpoints.py +++ b/litellm/types/management_endpoints/router_settings_endpoints.py @@ -8,6 +8,7 @@ from pydantic import BaseModel, Field, field_validator # Fallback Management Types + class FallbackCreateRequest(BaseModel): """Request model for creating/updating fallbacks""" @@ -74,7 +75,9 @@ class RouterSettingsField(BaseModel): field_value: Any field_description: str field_default: Any = None - options: Optional[List[str]] = None # For fields with predefined options/enum values + options: Optional[ + List[str] + ] = None # For fields with predefined options/enum values ui_field_name: str # User-friendly display name link: Optional[str] = None # Documentation link for the field @@ -245,7 +248,7 @@ ROUTER_SETTINGS_FIELDS: List[RouterSettingsField] = [ field_default=False, ui_field_name="Enable Tag Filtering", link="https://docs.litellm.ai/docs/proxy/tag_routing", - ), + ), RouterSettingsField( field_name="tag_filtering_match_any", field_type="Boolean", @@ -263,4 +266,3 @@ ROUTER_SETTINGS_FIELDS: List[RouterSettingsField] = [ ui_field_name="Disable Cooldowns", ), ] - diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 511cfc958a..ed391f0af6 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -95,17 +95,17 @@ class MCPServer(BaseModel): This includes: - OAuth2 servers without client credentials - Servers with auth_type=none but extra_headers configured for auth passthrough - + Health checks should be skipped for these servers since they cannot authenticate without user-provided credentials. """ # OAuth2 without client credentials if self.needs_user_oauth_token: return True - + # PAT passthrough: auth_type is none but extra_headers includes auth headers if self.auth_type == MCPAuth.none and self.extra_headers: auth_header_names = {"authorization", "x-api-key", "api-key", "apikey"} return any(h.lower() in auth_header_names for h in self.extra_headers) - + return False diff --git a/litellm/types/mcp_server/tool_registry.py b/litellm/types/mcp_server/tool_registry.py index 36c0919282..8e3f1d9657 100644 --- a/litellm/types/mcp_server/tool_registry.py +++ b/litellm/types/mcp_server/tool_registry.py @@ -4,7 +4,6 @@ from pydantic import BaseModel, ConfigDict class MCPTool(BaseModel): - model_config: ClassVar[ConfigDict] = ConfigDict(arbitrary_types_allowed=True) name: str description: str diff --git a/litellm/types/proxy/cloudzero_endpoints.py b/litellm/types/proxy/cloudzero_endpoints.py index fc48717e80..3ac5795b35 100644 --- a/litellm/types/proxy/cloudzero_endpoints.py +++ b/litellm/types/proxy/cloudzero_endpoints.py @@ -10,50 +10,75 @@ from pydantic import BaseModel, Field class CloudZeroInitRequest(BaseModel): """Request model for initializing CloudZero settings""" - + api_key: str = Field(..., description="CloudZero API key for authentication") - connection_id: str = Field(..., description="CloudZero connection ID for data submission") - timezone: str = Field(default="UTC", description="Timezone for date handling (default: UTC)") + connection_id: str = Field( + ..., description="CloudZero connection ID for data submission" + ) + timezone: str = Field( + default="UTC", description="Timezone for date handling (default: UTC)" + ) class CloudZeroInitResponse(BaseModel): """Response model for CloudZero initialization""" - + message: str status: str class CloudZeroExportRequest(BaseModel): """Request model for CloudZero export operations""" - - limit: Optional[int] = Field(None, description="Optional limit on number of records to export") - operation: str = Field(default="replace_hourly", description="CloudZero operation type (replace_hourly or sum)") - start_time_utc: Optional[datetime] = Field(None, description="Start time for data export in UTC") - end_time_utc: Optional[datetime] = Field(None, description="End time for data export in UTC") + + limit: Optional[int] = Field( + None, description="Optional limit on number of records to export" + ) + operation: str = Field( + default="replace_hourly", + description="CloudZero operation type (replace_hourly or sum)", + ) + start_time_utc: Optional[datetime] = Field( + None, description="Start time for data export in UTC" + ) + end_time_utc: Optional[datetime] = Field( + None, description="End time for data export in UTC" + ) class CloudZeroExportResponse(BaseModel): """Response model for CloudZero export operations""" - + message: str status: str records_exported: Optional[int] = None - dry_run_data: Optional[Dict[str, Any]] = Field(None, description="Dry run data including usage data and CBF transformed data") - summary: Optional[Dict[str, Any]] = Field(None, description="Summary statistics for dry run") + dry_run_data: Optional[Dict[str, Any]] = Field( + None, description="Dry run data including usage data and CBF transformed data" + ) + summary: Optional[Dict[str, Any]] = Field( + None, description="Summary statistics for dry run" + ) class CloudZeroSettingsView(BaseModel): """Response model for viewing CloudZero settings with masked API key""" - - api_key_masked: Optional[str] = Field(None, description="Masked API key showing only first 4 and last 4 characters") - connection_id: Optional[str] = Field(None, description="CloudZero connection ID for data submission") + + api_key_masked: Optional[str] = Field( + None, description="Masked API key showing only first 4 and last 4 characters" + ) + connection_id: Optional[str] = Field( + None, description="CloudZero connection ID for data submission" + ) timezone: Optional[str] = Field(None, description="Timezone for date handling") status: Optional[str] = Field(None, description="Configuration status") class CloudZeroSettingsUpdate(BaseModel): """Request model for updating CloudZero settings""" - - api_key: Optional[str] = Field(None, description="New CloudZero API key for authentication") - connection_id: Optional[str] = Field(None, description="New CloudZero connection ID for data submission") - timezone: Optional[str] = Field(None, description="New timezone for date handling") \ No newline at end of file + + api_key: Optional[str] = Field( + None, description="New CloudZero API key for authentication" + ) + connection_id: Optional[str] = Field( + None, description="New CloudZero connection ID for data submission" + ) + timezone: Optional[str] = Field(None, description="New timezone for date handling") diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_text_moderation.py b/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_text_moderation.py index 308b2d68ac..d73b73502b 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_text_moderation.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_text_moderation.py @@ -76,7 +76,6 @@ class AzureContentSafetyTextModerationConfigModel( AzureContentSafetyConfigModel, GuardrailConfigModel[AzureTextModerationOptionalParams], ): - optional_params: AzureTextModerationOptionalParams = Field( description="Optional parameters for the Azure Content Safety Text Moderation guardrail", ) diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/dynamoai.py b/litellm/types/proxy/guardrails/guardrail_hooks/dynamoai.py index 04123e3964..8d08931364 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/dynamoai.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/dynamoai.py @@ -10,12 +10,15 @@ from .base import GuardrailConfigModel class DynamoAIMessage(TypedDict): """Message structure for DynamoAI API""" + role: str content: str + class DynamoRequestMetadata(TypedDict): endUserId: Optional[str] + class DynamoTextType(str, enum.Enum): MODEL_INPUT = "MODEL_INPUT" MODEL_RESPONSE = "MODEL_RESPONSE" @@ -37,6 +40,7 @@ class PolicyApplicableTo(str, enum.Enum): class DynamoAIRequest(TypedDict, total=False): """Request structure for DynamoAI /moderation/analyze endpoint""" + messages: List[Dict[str, Any]] textType: Optional[DynamoTextType] policyIds: List[str] @@ -47,6 +51,7 @@ class DynamoAIRequest(TypedDict, total=False): class PolicyInfo(TypedDict, total=False): """Policy information from DynamoAI response""" + id: str name: str description: str @@ -61,12 +66,14 @@ class PolicyInfo(TypedDict, total=False): class PolicyOutputs(TypedDict, total=False): """Outputs from the policy""" + action: Literal["BLOCK", "WARN", "REDACT", "SANITIZE", "NONE"] message: Optional[str] class AppliedPolicyDto(TypedDict, total=False): """Applied policy details from DynamoAI response""" + policy: PolicyInfo outputs: Optional[Dict[str, Any]] action: Optional[str] @@ -74,6 +81,7 @@ class AppliedPolicyDto(TypedDict, total=False): class DynamoAIResponse(TypedDict, total=False): """Response structure from DynamoAI /moderation/analyze endpoint""" + text: str textType: DynamoTextType finalAction: Literal["BLOCK", "WARN", "REDACT", "SANITIZE", "NONE"] @@ -83,14 +91,14 @@ class DynamoAIResponse(TypedDict, total=False): class DynamoAIProcessedResult(TypedDict): """Processed result from DynamoAI guardrail check""" + violations_detected: List[str] violation_details: Dict[str, Any] - class DynamoAIGuardrailConfigModel(GuardrailConfigModel): """Configuration model for DynamoAI Guardrails""" - + api_key: Optional[str] = Field( default=None, description="API key for DynamoAI Guardrails. If not provided, the `DYNAMOAI_API_KEY` environment variable is checked.", @@ -111,8 +119,7 @@ class DynamoAIGuardrailConfigModel(GuardrailConfigModel): default=None, description="Name of the guardrail for identification in logs and traces.", ) - + @staticmethod def ui_friendly_name() -> str: return "DynamoAI Guardrails" - diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/enkryptai.py b/litellm/types/proxy/guardrails/guardrail_hooks/enkryptai.py index 4ab49d5276..cebac4d826 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/enkryptai.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/enkryptai.py @@ -23,7 +23,7 @@ class EnkryptAIPIIDetail(TypedDict, total=False): class EnkryptAIToxicityDetail(TypedDict, total=False): """Details for toxicity detection. - + Contains scores for different types of toxicity: - toxic - severe_toxic @@ -55,7 +55,7 @@ class EnkryptAIBiasDetail(TypedDict, total=False): class EnkryptAIResponseSummary(TypedDict, total=False): """Summary of detected violations in EnkryptAI response. - + Each key represents a type of violation: - toxicity: List (non-empty if detected) - policy_violation: 0 or 1 @@ -98,12 +98,23 @@ class EnkryptAIProcessedResult(TypedDict): """Processed result from EnkryptAI guardrail response.""" attacks_detected: List[str] - attack_details: Dict[str, Union[EnkryptAIPolicyViolationDetail, EnkryptAIPIIDetail, EnkryptAIToxicityDetail, EnkryptAIKeywordDetail, EnkryptAIBiasDetail, Dict[str, Any]]] + attack_details: Dict[ + str, + Union[ + EnkryptAIPolicyViolationDetail, + EnkryptAIPIIDetail, + EnkryptAIToxicityDetail, + EnkryptAIKeywordDetail, + EnkryptAIBiasDetail, + Dict[str, Any], + ], + ] # Pydantic Config Model class EnkryptAIGuardrailConfigs(BaseModel): """Configuration parameters for the EnkryptAI guardrail""" + api_key: Optional[str] = Field( default=None, description="The EnkryptAI API key. Reads from ENKRYPTAI_API_KEY env var if None.", @@ -129,8 +140,8 @@ class EnkryptAIGuardrailConfigs(BaseModel): description="Whether to block requests when violations are detected. Defaults to True.", ) + class EnkryptAIGuardrailConfigModel(GuardrailConfigModel, EnkryptAIGuardrailConfigs): @staticmethod def ui_friendly_name() -> str: return "EnkryptAI" - diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py index 94f219a5fc..c87086bdce 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py @@ -60,9 +60,9 @@ class GenericGuardrailAPIRequest(BaseModel): input_type: Literal["request", "response"] litellm_call_id: Optional[str] = None # the call id of the individual LLM call - litellm_trace_id: Optional[str] = ( - None # the trace id of the LLM call - useful if there are multiple LLM calls for the same conversation - ) + litellm_trace_id: Optional[ + str + ] = None # the trace id of the LLM call - useful if there are multiple LLM calls for the same conversation structured_messages: Optional[List[AllMessageValues]] = None images: Optional[List[str]] = None tools: Optional[List[ChatCompletionToolParam]] = None diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/litellm_content_filter.py b/litellm/types/proxy/guardrails/guardrail_hooks/litellm_content_filter.py index b27789fd7f..992becdb13 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/litellm_content_filter.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/litellm_content_filter.py @@ -4,8 +4,7 @@ from typing import Any, Dict, List, Literal, Optional, TypedDict, Union from pydantic import Field from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject -from litellm.types.proxy.guardrails.guardrail_hooks.base import \ - GuardrailConfigModel +from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel # --- Competitor intent blocker (generic, industry-agnostic) --- diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/model_armor.py b/litellm/types/proxy/guardrails/guardrail_hooks/model_armor.py index 2d8fa0606b..5fa701574c 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/model_armor.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/model_armor.py @@ -32,4 +32,4 @@ class ModelArmorGuardrailConfigModel(GuardrailConfigModel): @staticmethod def ui_friendly_name() -> str: """Return the UI-friendly name for Model Armor guardrail""" - return "Google Cloud Model Armor" \ No newline at end of file + return "Google Cloud Model Armor" diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/openai/openai_moderation.py b/litellm/types/proxy/guardrails/guardrail_hooks/openai/openai_moderation.py index 355430ef2f..ee67626967 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/openai/openai_moderation.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/openai/openai_moderation.py @@ -7,26 +7,28 @@ from ..base import GuardrailConfigModel class BaseOpenAIModerationGuardrailConfigModel(GuardrailConfigModel): """Base configuration model for the OpenAI Moderation guardrail""" - model: Optional[Literal["omni-moderation-latest", "text-moderation-latest"]] = Field( + + model: Optional[ + Literal["omni-moderation-latest", "text-moderation-latest"] + ] = Field( default="omni-moderation-latest", description="The OpenAI moderation model to use. 'omni-moderation-latest' supports more categorization options and multi-modal inputs. Defaults to 'omni-moderation-latest'.", ) + class OpenAIModerationGuardrailConfigModel(BaseOpenAIModerationGuardrailConfigModel): """Configuration model for the OpenAI Moderation guardrail""" - + api_key: Optional[str] = Field( default=None, description="OpenAI API key. Can also be set via OPENAI_API_KEY environment variable.", ) - + api_base: Optional[str] = Field( default="https://api.openai.com/v1", description="OpenAI API base URL. Defaults to 'https://api.openai.com/v1'.", ) - - @staticmethod def ui_friendly_name() -> str: - return "OpenAI Moderation" \ No newline at end of file + return "OpenAI Moderation" diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/zscaler_ai_guard.py b/litellm/types/proxy/guardrails/guardrail_hooks/zscaler_ai_guard.py index f522f5b470..a0cd280202 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/zscaler_ai_guard.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/zscaler_ai_guard.py @@ -108,7 +108,9 @@ class ZscalerAIGuardConfigModel(GuardrailConfigModel): # Check for configuration issues assert api_base is not None # always set via env default above is_resolve_policy = api_base.endswith("/resolve-and-execute-policy") - is_execute_policy = api_base.endswith("/execute-policy") and not is_resolve_policy + is_execute_policy = ( + api_base.endswith("/execute-policy") and not is_resolve_policy + ) # Scenario A: execute-policy without policy_id if is_execute_policy and (policy_id is None or policy_id < 1): diff --git a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py index 6023094a92..4770877dab 100644 --- a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py @@ -25,13 +25,13 @@ class UserListResponse(BaseModel): class BulkUpdateUserRequest(BaseModel): """Request for bulk user updates""" - users: Optional[List[UpdateUserRequest]] = ( - None # List of specific user update requests - ) + users: Optional[ + List[UpdateUserRequest] + ] = None # List of specific user update requests all_users: Optional[bool] = False # Flag to update all users - user_updates: Optional[UpdateUserRequestNoUserIDorEmail] = ( - None # Updates to apply to all users when all_users=True - ) + user_updates: Optional[ + UpdateUserRequestNoUserIDorEmail + ] = None # Updates to apply to all users when all_users=True @field_validator("users", "all_users", "user_updates") @classmethod diff --git a/litellm/types/proxy/management_endpoints/model_management_endpoints.py b/litellm/types/proxy/management_endpoints/model_management_endpoints.py index 6f07e5c6de..be4d730e93 100644 --- a/litellm/types/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/model_management_endpoints.py @@ -21,8 +21,12 @@ class UpdateUsefulLinksRequest(BaseModel): class NewModelGroupRequest(BaseModel): access_group: str # The access group name (e.g., "production-models") - model_names: Optional[List[str]] = None # Existing model groups to include - tags ALL deployments for each name - model_ids: Optional[List[str]] = None # Specific deployment IDs to tag (more precise than model_names) + model_names: Optional[ + List[str] + ] = None # Existing model groups to include - tags ALL deployments for each name + model_ids: Optional[ + List[str] + ] = None # Specific deployment IDs to tag (more precise than model_names) class NewModelGroupResponse(BaseModel): @@ -33,8 +37,12 @@ class NewModelGroupResponse(BaseModel): class UpdateModelGroupRequest(BaseModel): - model_names: Optional[List[str]] = None # Updated list of model groups to include - tags ALL deployments for each name - model_ids: Optional[List[str]] = None # Specific deployment IDs to tag (more precise than model_names) + model_names: Optional[ + List[str] + ] = None # Updated list of model groups to include - tags ALL deployments for each name + model_ids: Optional[ + List[str] + ] = None # Specific deployment IDs to tag (more precise than model_names) class DeleteModelGroupResponse(BaseModel): @@ -50,4 +58,4 @@ class AccessGroupInfo(BaseModel): class ListAccessGroupsResponse(BaseModel): - access_groups: List[AccessGroupInfo] \ No newline at end of file + access_groups: List[AccessGroupInfo] diff --git a/litellm/types/proxy/management_endpoints/scim_v2.py b/litellm/types/proxy/management_endpoints/scim_v2.py index c4d95d99ed..c5fdc66154 100644 --- a/litellm/types/proxy/management_endpoints/scim_v2.py +++ b/litellm/types/proxy/management_endpoints/scim_v2.py @@ -101,9 +101,7 @@ class SCIMFeature(BaseModel): class SCIMServiceProviderConfig(BaseModel): - schemas: List[str] = [ - "urn:ietf:params:scim:schemas:core:2.0:ServiceProviderConfig" - ] + schemas: List[str] = ["urn:ietf:params:scim:schemas:core:2.0:ServiceProviderConfig"] patch: SCIMFeature = SCIMFeature(supported=True) bulk: SCIMFeature = SCIMFeature(supported=False) filter: SCIMFeature = SCIMFeature(supported=False) @@ -130,9 +128,7 @@ class SCIMSchemaExtension(BaseModel): class SCIMResourceType(BaseModel): model_config = ConfigDict(populate_by_name=True) - schemas: List[str] = [ - "urn:ietf:params:scim:schemas:core:2.0:ResourceType" - ] + schemas: List[str] = ["urn:ietf:params:scim:schemas:core:2.0:ResourceType"] id: str name: str description: Optional[str] = None diff --git a/litellm/types/proxy/policy_engine/__init__.py b/litellm/types/proxy/policy_engine/__init__.py index 4df9f21e80..84d354b82a 100644 --- a/litellm/types/proxy/policy_engine/__init__.py +++ b/litellm/types/proxy/policy_engine/__init__.py @@ -11,28 +11,52 @@ Configuration: """ from litellm.types.proxy.policy_engine.pipeline_types import ( - GuardrailPipeline, PipelineExecutionResult, PipelineStep, - PipelineStepResult) -from litellm.types.proxy.policy_engine.policy_types import (Policy, - PolicyAttachment, - PolicyCondition, - PolicyConfig, - PolicyGuardrails, - PolicyScope) + GuardrailPipeline, + PipelineExecutionResult, + PipelineStep, + PipelineStepResult, +) +from litellm.types.proxy.policy_engine.policy_types import ( + Policy, + PolicyAttachment, + PolicyCondition, + PolicyConfig, + PolicyGuardrails, + PolicyScope, +) from litellm.types.proxy.policy_engine.resolver_types import ( - AttachmentImpactResponse, PipelineTestRequest, - PolicyAttachmentCreateRequest, PolicyAttachmentDBResponse, - PolicyAttachmentListResponse, PolicyConditionRequest, PolicyCreateRequest, - PolicyDBResponse, PolicyGuardrailsResponse, PolicyInfoResponse, - PolicyListDBResponse, PolicyListResponse, PolicyMatchContext, - PolicyMatchDetail, PolicyResolveRequest, PolicyResolveResponse, - PolicyScopeResponse, PolicySummaryItem, PolicyTestResponse, - PolicyUpdateRequest, PolicyVersionCompareResponse, - PolicyVersionCreateRequest, PolicyVersionListResponse, - PolicyVersionStatusUpdateRequest, ResolvedPolicy) + AttachmentImpactResponse, + PipelineTestRequest, + PolicyAttachmentCreateRequest, + PolicyAttachmentDBResponse, + PolicyAttachmentListResponse, + PolicyConditionRequest, + PolicyCreateRequest, + PolicyDBResponse, + PolicyGuardrailsResponse, + PolicyInfoResponse, + PolicyListDBResponse, + PolicyListResponse, + PolicyMatchContext, + PolicyMatchDetail, + PolicyResolveRequest, + PolicyResolveResponse, + PolicyScopeResponse, + PolicySummaryItem, + PolicyTestResponse, + PolicyUpdateRequest, + PolicyVersionCompareResponse, + PolicyVersionCreateRequest, + PolicyVersionListResponse, + PolicyVersionStatusUpdateRequest, + ResolvedPolicy, +) from litellm.types.proxy.policy_engine.validation_types import ( - PolicyValidateRequest, PolicyValidationError, PolicyValidationErrorType, - PolicyValidationResponse) + PolicyValidateRequest, + PolicyValidationError, + PolicyValidationErrorType, + PolicyValidationResponse, +) __all__ = [ # Pipeline types diff --git a/litellm/types/proxy/policy_engine/resolver_types.py b/litellm/types/proxy/policy_engine/resolver_types.py index 2df450dc2b..cb6590688d 100644 --- a/litellm/types/proxy/policy_engine/resolver_types.py +++ b/litellm/types/proxy/policy_engine/resolver_types.py @@ -207,7 +207,8 @@ class PolicyDBResponse(BaseModel): default=None, description="Policy ID this version was cloned from." ) is_latest: bool = Field( - default=True, description="True if this is the latest version by version_number." + default=True, + description="True if this is the latest version by version_number.", ) published_at: Optional[datetime] = Field( default=None, description="When this version was published." @@ -235,7 +236,9 @@ class PolicyDBResponse(BaseModel): updated_at: Optional[datetime] = Field( default=None, description="When the policy was last updated." ) - created_by: Optional[str] = Field(default=None, description="Who created the policy.") + created_by: Optional[str] = Field( + default=None, description="Who created the policy." + ) updated_by: Optional[str] = Field( default=None, description="Who last updated the policy." ) @@ -382,12 +385,8 @@ class PolicyResolveRequest(BaseModel): key_alias: Optional[str] = Field( default=None, description="Key alias to resolve for." ) - model: Optional[str] = Field( - default=None, description="Model name to resolve for." - ) - tags: Optional[List[str]] = Field( - default=None, description="Tags to resolve for." - ) + model: Optional[str] = Field(default=None, description="Model name to resolve for.") + tags: Optional[List[str]] = Field(default=None, description="Tags to resolve for.") class PolicyMatchDetail(BaseModel): @@ -425,10 +424,12 @@ class AttachmentImpactResponse(BaseModel): """Response for estimating the impact of a policy attachment.""" affected_keys_count: int = Field( - default=0, description="Number of keys that would be affected (named + unnamed)." + default=0, + description="Number of keys that would be affected (named + unnamed).", ) affected_teams_count: int = Field( - default=0, description="Number of teams that would be affected (named + unnamed)." + default=0, + description="Number of teams that would be affected (named + unnamed).", ) unnamed_keys_count: int = Field( default=0, description="Number of affected keys without an alias." diff --git a/litellm/types/proxy/prompt_endpoints.py b/litellm/types/proxy/prompt_endpoints.py index 620a565b0a..609a6e55c9 100644 --- a/litellm/types/proxy/prompt_endpoints.py +++ b/litellm/types/proxy/prompt_endpoints.py @@ -7,4 +7,3 @@ class TestPromptRequest(BaseModel): dotprompt_content: str prompt_variables: Optional[Dict[str, Any]] = None conversation_history: Optional[List[Dict[str, str]]] = None - diff --git a/litellm/types/proxy/ui_sso.py b/litellm/types/proxy/ui_sso.py index 04523e88b1..bef200952b 100644 --- a/litellm/types/proxy/ui_sso.py +++ b/litellm/types/proxy/ui_sso.py @@ -23,6 +23,7 @@ class ParsedOpenIDResult(TypedDict, total=False): """ Parsed OpenID result """ + user_email: Optional[str] user_id: Optional[str] - user_role: Optional[str] \ No newline at end of file + user_role: Optional[str] diff --git a/litellm/types/rag.py b/litellm/types/rag.py index cae0770868..29e35d5fe8 100644 --- a/litellm/types/rag.py +++ b/litellm/types/rag.py @@ -53,7 +53,9 @@ class OpenAIVectorStoreOptions(TypedDict, total=False): ttl_days: Optional[int] # Time-to-live in days for indexed content # Credentials (loaded from litellm.credential_list if litellm_credential_name is provided) - litellm_credential_name: Optional[str] # Credential name to load from litellm.credential_list + litellm_credential_name: Optional[ + str + ] # Credential name to load from litellm.credential_list api_key: Optional[str] # Direct API key (alternative to litellm_credential_name) api_base: Optional[str] # Direct API base (alternative to litellm_credential_name) @@ -81,13 +83,21 @@ class BedrockVectorStoreOptions(TypedDict, total=False): # Bedrock-specific options s3_bucket: Optional[str] # S3 bucket (auto-created if not provided) s3_prefix: Optional[str] # S3 key prefix (default: "data/") - embedding_model: Optional[str] # Embedding model (default: amazon.titan-embed-text-v2:0) + embedding_model: Optional[ + str + ] # Embedding model (default: amazon.titan-embed-text-v2:0) data_source_id: Optional[str] # For existing KB: override auto-detected DS - wait_for_ingestion: Optional[bool] # Wait for completion (default: False - returns immediately) - ingestion_timeout: Optional[int] # Timeout in seconds if wait_for_ingestion=True (default: 300) + wait_for_ingestion: Optional[ + bool + ] # Wait for completion (default: False - returns immediately) + ingestion_timeout: Optional[ + int + ] # Timeout in seconds if wait_for_ingestion=True (default: 300) # Credentials (loaded from litellm.credential_list if litellm_credential_name is provided) - litellm_credential_name: Optional[str] # Credential name to load from litellm.credential_list + litellm_credential_name: Optional[ + str + ] # Credential name to load from litellm.credential_list # AWS auth (uses BaseAWSLLM) aws_access_key_id: Optional[str] @@ -119,10 +129,14 @@ class VertexAIVectorStoreOptions(TypedDict, total=False): vector_store_id: str # RAG corpus ID (required for Vertex AI) # GCP config - vertex_project: Optional[str] # GCP project ID (uses env VERTEXAI_PROJECT if not set) + vertex_project: Optional[ + str + ] # GCP project ID (uses env VERTEXAI_PROJECT if not set) vertex_location: Optional[str] # GCP region (default: us-central1) vertex_credentials: Optional[str] # Path to credentials JSON (uses ADC if not set) - gcs_bucket: Optional[str] # GCS bucket for file uploads (uses env GCS_BUCKET_NAME if not set) + gcs_bucket: Optional[ + str + ] # GCS bucket for file uploads (uses env GCS_BUCKET_NAME if not set) # Import settings wait_for_import: Optional[bool] # Wait for import to complete (default: True) @@ -153,12 +167,18 @@ class S3VectorsVectorStoreOptions(TypedDict, total=False): index_name: Optional[str] # Vector index name (auto-creates if not provided) # Index configuration (for auto-creation) - dimension: Optional[int] # Vector dimension (auto-detected from embedding model, or default: 1024) + dimension: Optional[ + int + ] # Vector dimension (auto-detected from embedding model, or default: 1024) distance_metric: Optional[Literal["cosine", "euclidean"]] # Default: cosine - non_filterable_metadata_keys: Optional[List[str]] # Keys excluded from filtering (e.g., ["source_text"]) + non_filterable_metadata_keys: Optional[ + List[str] + ] # Keys excluded from filtering (e.g., ["source_text"]) # Credentials (loaded from litellm.credential_list if litellm_credential_name is provided) - litellm_credential_name: Optional[str] # Credential name to load from litellm.credential_list + litellm_credential_name: Optional[ + str + ] # Credential name to load from litellm.credential_list # AWS auth (uses BaseAWSLLM) aws_access_key_id: Optional[str] @@ -175,7 +195,10 @@ class S3VectorsVectorStoreOptions(TypedDict, total=False): # Union type for vector store options RAGIngestVectorStoreOptions = Union[ - OpenAIVectorStoreOptions, BedrockVectorStoreOptions, VertexAIVectorStoreOptions, S3VectorsVectorStoreOptions + OpenAIVectorStoreOptions, + BedrockVectorStoreOptions, + VertexAIVectorStoreOptions, + S3VectorsVectorStoreOptions, ] @@ -209,10 +232,13 @@ class RAGIngestOptions(TypedDict, total=False): name: Optional[str] # Optional pipeline name for logging ocr: Optional[RAGIngestOCROptions] # Optional OCR step - chunking_strategy: Optional[RAGChunkingStrategy] # RecursiveCharacterTextSplitter args + chunking_strategy: Optional[ + RAGChunkingStrategy + ] # RecursiveCharacterTextSplitter args embedding: Optional[RAGIngestEmbeddingOptions] # Embedding model config vector_store: RAGIngestVectorStoreOptions # OpenAI or Bedrock config + class RAGIngestResponse(TypedDict, total=False): """Response from RAG ingest API.""" @@ -223,7 +249,6 @@ class RAGIngestResponse(TypedDict, total=False): error: Optional[str] # Error message if status is "failed" - class RAGIngestRequest(BaseModel): """Request body for RAG ingest API (for validation).""" @@ -268,4 +293,3 @@ class RAGQueryResponse(ModelResponse): """Response from RAG query API.""" pass - diff --git a/litellm/types/responses/main.py b/litellm/types/responses/main.py index 449a5ac49c..7a666d5e65 100644 --- a/litellm/types/responses/main.py +++ b/litellm/types/responses/main.py @@ -6,7 +6,8 @@ from typing_extensions import Any, List, Optional, TypedDict from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject -Phase = Optional[Literal["commentary", "final_answer"]] +Phase = Optional[Literal["commentary", "final_answer"]] + class GenericResponseOutputItemContentAnnotation(BaseLiteLLMOpenAIResponseObject): """Annotation for content in a message""" diff --git a/litellm/types/search.py b/litellm/types/search.py index b0ce0636ae..bbac1237a1 100644 --- a/litellm/types/search.py +++ b/litellm/types/search.py @@ -15,11 +15,11 @@ SearchProvider = SearchProviders __all__ = ["SearchProvider", "SearchProviders"] - class SearchToolLiteLLMParams(TypedDict, total=False): """ LiteLLM params for search tools configuration. """ + search_provider: Required[str] api_key: Optional[str] api_base: Optional[str] @@ -30,7 +30,7 @@ class SearchToolLiteLLMParams(TypedDict, total=False): class SearchTool(TypedDict, total=False): """ Search tool configuration. - + Example: { "search_tool_id": "123e4567-e89b-12d3-a456-426614174000", @@ -44,6 +44,7 @@ class SearchTool(TypedDict, total=False): } } """ + search_tool_id: Optional[str] search_tool_name: Required[str] litellm_params: Required[SearchToolLiteLLMParams] @@ -54,23 +55,26 @@ class SearchTool(TypedDict, total=False): class SearchToolInfoResponse(TypedDict, total=False): """Response model for search tool information.""" + search_tool_id: Optional[str] search_tool_name: str litellm_params: dict search_tool_info: Optional[dict] created_at: Optional[str] updated_at: Optional[str] - is_from_config: Optional[bool] # True if this tool is defined in config file, False if from DB + is_from_config: Optional[ + bool + ] # True if this tool is defined in config file, False if from DB class ListSearchToolsResponse(TypedDict): """Response model for listing search tools.""" + search_tools: List[SearchToolInfoResponse] class AvailableSearchProvider(TypedDict): """Information about an available search provider.""" + provider_name: str ui_friendly_name: str - - diff --git a/litellm/types/secret_managers/main.py b/litellm/types/secret_managers/main.py index e4c7d76573..b0a294188c 100644 --- a/litellm/types/secret_managers/main.py +++ b/litellm/types/secret_managers/main.py @@ -71,4 +71,4 @@ class KeyManagementSettings(LiteLLMPydanticObjectBase): """Web identity token for OIDC/IRSA authentication""" aws_sts_endpoint: Optional[str] = None - """Custom STS endpoint URL (useful for VPC endpoints or testing)""" \ No newline at end of file + """Custom STS endpoint URL (useful for VPC endpoints or testing)""" diff --git a/litellm/types/tag_management.py b/litellm/types/tag_management.py index a9b58ace02..3bf70c73fc 100644 --- a/litellm/types/tag_management.py +++ b/litellm/types/tag_management.py @@ -1,4 +1,3 @@ -from datetime import datetime from typing import Dict, List, Optional from pydantic import BaseModel @@ -47,23 +46,3 @@ class TagDeleteRequest(BaseModel): class TagInfoRequest(BaseModel): names: List[str] - - -class LiteLLM_DailyTagSpendTable(BaseModel): - id: str - tag: str - date: str - api_key: str - model: str - model_group: Optional[str] - custom_llm_provider: Optional[str] - prompt_tokens: int - completion_tokens: int - cache_read_input_tokens: int - cache_creation_input_tokens: int - spend: float - api_requests: int - successful_requests: int - failed_requests: int - created_at: datetime - updated_at: datetime diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 74afb5fe2e..de8e707423 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1339,7 +1339,9 @@ class Choices(SafeAttributeModel, OpenAIObject): mapped = map_finish_reason(finish_reason) params["finish_reason"] = mapped if finish_reason != mapped: - provider_specific_fields = dict(provider_specific_fields) if provider_specific_fields else {} + provider_specific_fields = ( + dict(provider_specific_fields) if provider_specific_fields else {} + ) provider_specific_fields["native_finish_reason"] = finish_reason else: params["finish_reason"] = "stop" @@ -1706,7 +1708,6 @@ class StreamingChatCompletionChunk(OpenAIChatCompletionChunk): super().__init__(**kwargs) - class ModelResponseBase(OpenAIObject): id: str """A unique identifier for the completion.""" diff --git a/litellm/types/vector_stores.py b/litellm/types/vector_stores.py index 58a6a3272c..ce247fc900 100644 --- a/litellm/types/vector_stores.py +++ b/litellm/types/vector_stores.py @@ -38,7 +38,7 @@ class LiteLLM_ManagedVectorStore(TypedDict, total=False): # litellm_params litellm_params: Optional[Dict[str, Any]] - + # access control fields team_id: Optional[str] user_id: Optional[str] @@ -243,6 +243,8 @@ class VectorStoreIndexEndpoints(TypedDict): write: List[ Tuple[Literal["GET", "POST", "PUT", "DELETE", "PATCH"], str] ] # endpoints for writing a vector store index + + VECTOR_STORE_OPENAI_PARAMS = Literal[ "filters", "max_num_results", @@ -251,14 +253,14 @@ VECTOR_STORE_OPENAI_PARAMS = Literal[ ] - @dataclass class VectorStoreToolParams: """Parameters extracted from a file_search tool definition""" + filters: Optional[Dict] = None max_num_results: Optional[int] = None ranking_options: Optional[Dict] = None - + def to_dict(self) -> Dict: """Convert to dict, excluding None values""" return { diff --git a/litellm/types/videos/main.py b/litellm/types/videos/main.py index 8e595db39f..b6357f3273 100644 --- a/litellm/types/videos/main.py +++ b/litellm/types/videos/main.py @@ -8,6 +8,7 @@ from litellm.types.utils import FileTypes class VideoObject(BaseModel): """Represents a generated video object.""" + id: str object: Literal["video"] status: str @@ -43,10 +44,9 @@ class VideoObject(BaseModel): return self.dict() - - class VideoResponse(BaseModel): """Response object for video generation requests.""" + data: List[VideoObject] hidden_params: Dict[str, Any] = {} @@ -72,9 +72,14 @@ class VideoCreateOptionalRequestParams(TypedDict, total=False): Params here: https://platform.openai.com/docs/api-reference/videos/create """ + input_reference: Optional[FileTypes] # File reference for input image - image: Optional[Any] # Image for image-to-video; dict with gcsUri/bytesBase64Encoded, or file-like object - parameters: Optional[Dict[str, Any]] # Provider-specific parameters block passed directly to the API + image: Optional[ + Any + ] # Image for image-to-video; dict with gcsUri/bytesBase64Encoded, or file-like object + parameters: Optional[ + Dict[str, Any] + ] # Provider-specific parameters block passed directly to the API model: Optional[str] seconds: Optional[str] size: Optional[str] @@ -89,11 +94,13 @@ class VideoCreateRequestParams(VideoCreateOptionalRequestParams, total=False): Params here: https://platform.openai.com/docs/api-reference/videos/create """ + prompt: str + class DecodedVideoId(TypedDict, total=False): """Structure representing a decoded video ID""" custom_llm_provider: Optional[str] model_id: Optional[str] - video_id: str \ No newline at end of file + video_id: str diff --git a/litellm/types/videos/utils.py b/litellm/types/videos/utils.py index 7f2148bb96..4916394e7e 100644 --- a/litellm/types/videos/utils.py +++ b/litellm/types/videos/utils.py @@ -15,14 +15,12 @@ VIDEO_ID_PREFIX = "video_" def encode_video_id_with_provider( - video_id: str, - provider: str, - model_id: Optional[str] = None + video_id: str, provider: str, model_id: Optional[str] = None ) -> str: """Encode provider and model_id into video_id using base64.""" if not provider or not video_id: return video_id - + # Try to decode the ID first to check if it's already encoded # This handles the case where Azure/OpenAI return IDs that start with "video_" # but are not yet encoded with provider information @@ -30,14 +28,16 @@ def encode_video_id_with_provider( if decoded.get("custom_llm_provider") is not None: # ID is already encoded, return as-is return video_id - + # ID is not encoded (even if it starts with video_), so encode it - assembled_id = str( - SpecialEnums.LITELLM_MANAGED_VIDEO_COMPLETE_STR.value - ).format(provider, model_id or "", video_id) - - base64_encoded_id: str = base64.b64encode(assembled_id.encode("utf-8")).decode("utf-8") - + assembled_id = str(SpecialEnums.LITELLM_MANAGED_VIDEO_COMPLETE_STR.value).format( + provider, model_id or "", video_id + ) + + base64_encoded_id: str = base64.b64encode(assembled_id.encode("utf-8")).decode( + "utf-8" + ) + return f"{VIDEO_ID_PREFIX}{base64_encoded_id}" @@ -49,14 +49,14 @@ def decode_video_id_with_provider(encoded_video_id: str) -> DecodedVideoId: model_id=None, video_id=encoded_video_id, ) - + if not encoded_video_id.startswith(VIDEO_ID_PREFIX): return DecodedVideoId( custom_llm_provider=None, model_id=None, video_id=encoded_video_id, ) - + try: cleaned_id = encoded_video_id.replace(VIDEO_ID_PREFIX, "") decoded_id = base64.b64decode(cleaned_id.encode("utf-8")).decode("utf-8") diff --git a/litellm/utils.py b/litellm/utils.py index 17dd6f91e9..860e33f047 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1143,7 +1143,9 @@ def function_setup( # noqa: PLR0915 litellm_params: Dict[str, Any] = {"api_base": ""} if "metadata" in kwargs: litellm_params["metadata"] = kwargs["metadata"] - if "litellm_metadata" in kwargs and isinstance(kwargs["litellm_metadata"], dict): + if "litellm_metadata" in kwargs and isinstance( + kwargs["litellm_metadata"], dict + ): litellm_params["litellm_metadata"] = kwargs["litellm_metadata"].copy() # For endpoints like /v1/messages that use "litellm_metadata" instead # of "metadata" (to avoid conflicting with provider API metadata fields), @@ -8348,7 +8350,9 @@ class ProviderConfigManager: from litellm.llms.openai_like.json_loader import JSONProviderRegistry # Resolve provider string for JSON lookup - provider_str = provider.value if isinstance(provider, LlmProviders) else str(provider) + provider_str = ( + provider.value if isinstance(provider, LlmProviders) else str(provider) + ) # Try to convert to enum for Python class lookup first. # Python classes take priority over JSON (they have custom overrides). @@ -8369,7 +8373,9 @@ class ProviderConfigManager: return result # Fall back to JSON providers (generic OpenAI-compatible) - if JSONProviderRegistry.exists(provider_str) and JSONProviderRegistry.supports_responses_api(provider_str): + if JSONProviderRegistry.exists( + provider_str + ) and JSONProviderRegistry.supports_responses_api(provider_str): provider_config = JSONProviderRegistry.get(provider_str) if provider_config is not None: return create_responses_config_class(provider_config)() @@ -8585,6 +8591,12 @@ class ProviderConfigManager: from litellm.llms.manus.files.transformation import ManusFilesConfig return ManusFilesConfig() + elif LlmProviders.ANTHROPIC == provider: + from litellm.llms.anthropic.files.transformation import ( + AnthropicFilesConfig, + ) + + return AnthropicFilesConfig() return None @staticmethod diff --git a/litellm/vector_store_files/main.py b/litellm/vector_store_files/main.py index 7d8fbbcd5f..de191bd212 100644 --- a/litellm/vector_store_files/main.py +++ b/litellm/vector_store_files/main.py @@ -130,9 +130,7 @@ def create( _prepare_registry_credentials(vector_store_id=vector_store_id, kwargs=kwargs) - litellm_params = GenericLiteLLMParams( - vector_store_id=vector_store_id, **kwargs - ) + litellm_params = GenericLiteLLMParams(vector_store_id=vector_store_id, **kwargs) provider_config = ProviderConfigManager.get_provider_vector_store_files_config( provider=LlmProviders(custom_llm_provider) @@ -253,7 +251,9 @@ def list( timeout: Optional[Union[float, httpx.Timeout]] = None, custom_llm_provider: Optional[str] = None, **kwargs, -) -> Union[VectorStoreFileListResponse, Coroutine[Any, Any, VectorStoreFileListResponse]]: +) -> Union[ + VectorStoreFileListResponse, Coroutine[Any, Any, VectorStoreFileListResponse] +]: local_vars = locals() try: litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore @@ -264,9 +264,7 @@ def list( _prepare_registry_credentials(vector_store_id=vector_store_id, kwargs=kwargs) - litellm_params = GenericLiteLLMParams( - vector_store_id=vector_store_id, **kwargs - ) + litellm_params = GenericLiteLLMParams(vector_store_id=vector_store_id, **kwargs) provider_config = ProviderConfigManager.get_provider_vector_store_files_config( provider=LlmProviders(custom_llm_provider) @@ -379,9 +377,7 @@ def retrieve( _prepare_registry_credentials(vector_store_id=vector_store_id, kwargs=kwargs) - litellm_params = GenericLiteLLMParams( - vector_store_id=vector_store_id, **kwargs - ) + litellm_params = GenericLiteLLMParams(vector_store_id=vector_store_id, **kwargs) provider_config = ProviderConfigManager.get_provider_vector_store_files_config( provider=LlmProviders(custom_llm_provider) @@ -492,9 +488,7 @@ def retrieve_content( _prepare_registry_credentials(vector_store_id=vector_store_id, kwargs=kwargs) - litellm_params = GenericLiteLLMParams( - vector_store_id=vector_store_id, **kwargs - ) + litellm_params = GenericLiteLLMParams(vector_store_id=vector_store_id, **kwargs) provider_config = ProviderConfigManager.get_provider_vector_store_files_config( provider=LlmProviders(custom_llm_provider) @@ -609,9 +603,7 @@ def update( _prepare_registry_credentials(vector_store_id=vector_store_id, kwargs=kwargs) - litellm_params = GenericLiteLLMParams( - vector_store_id=vector_store_id, **kwargs - ) + litellm_params = GenericLiteLLMParams(vector_store_id=vector_store_id, **kwargs) provider_config = ProviderConfigManager.get_provider_vector_store_files_config( provider=LlmProviders(custom_llm_provider) @@ -731,9 +723,7 @@ def delete( _prepare_registry_credentials(vector_store_id=vector_store_id, kwargs=kwargs) - litellm_params = GenericLiteLLMParams( - vector_store_id=vector_store_id, **kwargs - ) + litellm_params = GenericLiteLLMParams(vector_store_id=vector_store_id, **kwargs) provider_config = ProviderConfigManager.get_provider_vector_store_files_config( provider=LlmProviders(custom_llm_provider) diff --git a/litellm/vector_store_files/utils.py b/litellm/vector_store_files/utils.py index cfc932f0cb..ffe73516bd 100644 --- a/litellm/vector_store_files/utils.py +++ b/litellm/vector_store_files/utils.py @@ -29,9 +29,7 @@ class VectorStoreFileRequestUtils: return cast(VectorStoreFileCreateRequest, filtered) @staticmethod - def get_list_query_params( - params: Dict[str, Any] - ) -> VectorStoreFileListQueryParams: + def get_list_query_params(params: Dict[str, Any]) -> VectorStoreFileListQueryParams: filtered = VectorStoreFileRequestUtils._filter_params( params=params, model=VectorStoreFileListQueryParams ) diff --git a/litellm/vector_stores/__init__.py b/litellm/vector_stores/__init__.py index 6bcc654032..011c620f13 100644 --- a/litellm/vector_stores/__init__.py +++ b/litellm/vector_stores/__init__.py @@ -1,4 +1,4 @@ from .main import acreate, asearch, create, search from .vector_store_registry import VectorStoreRegistry -__all__ = ["search", "asearch", "create", "acreate", "VectorStoreRegistry"] +__all__ = ["search", "asearch", "create", "acreate", "VectorStoreRegistry"] diff --git a/litellm/vector_stores/vector_store_registry.py b/litellm/vector_stores/vector_store_registry.py index cf0bf89d70..2596f968a0 100644 --- a/litellm/vector_stores/vector_store_registry.py +++ b/litellm/vector_stores/vector_store_registry.py @@ -19,13 +19,14 @@ if TYPE_CHECKING: else: PrismaClient = Any + class VectorStoreIndexRegistry: def __init__( self, vector_store_indexes: List[LiteLLM_ManagedVectorStoreIndex] = [] ): - self.vector_store_indexes: List[LiteLLM_ManagedVectorStoreIndex] = ( - vector_store_indexes - ) + self.vector_store_indexes: List[ + LiteLLM_ManagedVectorStoreIndex + ] = vector_store_indexes def get_vector_store_indexes(self) -> List[LiteLLM_ManagedVectorStoreIndex]: """ @@ -114,15 +115,15 @@ class VectorStoreRegistry: def _extract_tool_params(self, tool: Dict) -> VectorStoreToolParams: """ Extract supported parameters from a tool definition. - + Dynamically extracts all parameters defined in VECTOR_STORE_OPENAI_PARAMS. """ # Get the list of supported param names from the Literal type supported_params = get_args(VECTOR_STORE_OPENAI_PARAMS) - + # Extract only the params that exist in the tool kwargs = {param: tool.get(param) for param in supported_params if param in tool} - + return VectorStoreToolParams(**kwargs) def get_vector_store_ids_to_run( @@ -148,51 +149,53 @@ class VectorStoreRegistry: return list(dict.fromkeys(vector_store_ids)) def get_and_pop_recognised_vector_store_tools( - self, tools: Optional[List[Dict]] = None, vector_store_ids: Optional[List[str]] = None + self, + tools: Optional[List[Dict]] = None, + vector_store_ids: Optional[List[str]] = None, ) -> Dict[str, VectorStoreToolParams]: """ Returns and pops recognized vector store tools from the tools list. - + Args: tools: The tools to extract and remove vector store IDs from vector_store_ids: Mutable list to append found vector_store_ids to - + Returns: Dict mapping vector_store_id to its extracted tool parameters """ params_by_id: Dict[str, VectorStoreToolParams] = {} - + if not tools: return params_by_id - + if vector_store_ids is None: vector_store_ids = [] - + tools_to_remove: List[int] = [] - + for i, tool in enumerate(tools): tool_vector_store_ids = tool.get("vector_store_ids", []) if not tool_vector_store_ids: continue - + # Check if all vector_store_ids are recognized in the registry recognised = all( any(vs.get("vector_store_id") == vs_id for vs in self.vector_stores) for vs_id in tool_vector_store_ids ) - + if recognised: tools_to_remove.append(i) vector_store_ids.extend(tool_vector_store_ids) - + # Extract and store params for each vector store tool_params = self._extract_tool_params(tool) for vs_id in tool_vector_store_ids: params_by_id[vs_id] = tool_params - + # Remove recognized tools from the original list remove_items_at_indices(items=tools, indices=tools_to_remove) - + return params_by_id def get_vector_store_to_run( @@ -241,10 +244,12 @@ class VectorStoreRegistry: This ensures synchronization across multiple instances. """ # First check in-memory registry - vector_store = self.get_litellm_managed_vector_store_from_registry(vector_store_id) + vector_store = self.get_litellm_managed_vector_store_from_registry( + vector_store_id + ) if vector_store is not None: return vector_store - + # Fall back to database if not found in memory if prisma_client is not None: try: @@ -260,7 +265,7 @@ class VectorStoreRegistry: verbose_logger.debug( f"Error fetching vector store from database: {str(e)}" ) - + return None def get_litellm_managed_vector_store_from_registry_by_name( @@ -279,87 +284,91 @@ class VectorStoreRegistry: ) -> List[LiteLLM_ManagedVectorStore]: """ Pops the vector stores to run with their tool parameters merged. - + Primary function to use for vector store pre call hook. - + Args: non_default_params: Parameters dict to pop vector_store_ids from tools: Optional list of tools to extract vector store params from - + Returns: List of vector stores with tool parameters merged into litellm_params """ # Pop vector_store_ids from params - vector_store_ids: List[str] = non_default_params.pop("vector_store_ids", None) or [] - + vector_store_ids: List[str] = ( + non_default_params.pop("vector_store_ids", None) or [] + ) + # Extract params from tools and collect IDs params_by_id = self.get_and_pop_recognised_vector_store_tools( - tools=tools, - vector_store_ids=vector_store_ids + tools=tools, vector_store_ids=vector_store_ids ) - + vector_stores_to_run: List[LiteLLM_ManagedVectorStore] = [] - + for vector_store_id in vector_store_ids: for vector_store in self.vector_stores: if vector_store.get("vector_store_id") == vector_store_id: # Create a copy to avoid modifying the registry vector_store_copy = vector_store.copy() - + # Merge tool params if they exist if vector_store_id in params_by_id: - existing_params = vector_store_copy.get("litellm_params", {}) or {} + existing_params = ( + vector_store_copy.get("litellm_params", {}) or {} + ) tool_params_dict = params_by_id[vector_store_id].to_dict() # Tool params take precedence over existing params tool_params_dict.update(existing_params) vector_store_copy["litellm_params"] = tool_params_dict - + vector_stores_to_run.append(vector_store_copy) break - + return vector_stores_to_run async def pop_vector_stores_to_run_with_db_fallback( - self, - non_default_params: Dict, + self, + non_default_params: Dict, tools: Optional[List[Dict]] = None, - prisma_client: Optional[PrismaClient] = None + prisma_client: Optional[PrismaClient] = None, ) -> List[LiteLLM_ManagedVectorStore]: """ Pops the vector stores to run with their tool parameters merged. Falls back to database if vector stores are not found in memory. This ensures synchronization across multiple instances. - + Primary function to use for vector store pre call hook. - + Args: non_default_params: Parameters dict to pop vector_store_ids from tools: Optional list of tools to extract vector store params from prisma_client: Optional database client for fallback lookup - + Returns: List of vector stores with tool parameters merged into litellm_params """ # Pop vector_store_ids from params - vector_store_ids: List[str] = non_default_params.pop("vector_store_ids", None) or [] - + vector_store_ids: List[str] = ( + non_default_params.pop("vector_store_ids", None) or [] + ) + # Extract params from tools and collect IDs params_by_id = self.get_and_pop_recognised_vector_store_tools( - tools=tools, - vector_store_ids=vector_store_ids + tools=tools, vector_store_ids=vector_store_ids ) - + vector_stores_to_run: List[LiteLLM_ManagedVectorStore] = [] - + for vector_store_id in vector_store_ids: vector_store = None - + # First check in-memory registry for vs in self.vector_stores: if vs.get("vector_store_id") == vector_store_id: vector_store = vs break - + # Verify vector store still exists in database (if we have DB access) # This ensures deleted vector stores are removed from cache if vector_store is not None and prisma_client is not None: @@ -373,29 +382,32 @@ class VectorStoreRegistry: verbose_logger.debug( f"Vector store {vector_store_id} found in memory but deleted from database, removing from cache" ) - self.delete_vector_store_from_registry(vector_store_id=vector_store_id) + self.delete_vector_store_from_registry( + vector_store_id=vector_store_id + ) vector_store = None except Exception as e: verbose_logger.debug( f"Error verifying vector store {vector_store_id} in database: {str(e)}" ) - + # Fall back to database if not found in memory (or was deleted) if vector_store is None and prisma_client is not None: try: - vector_store = await self.get_litellm_managed_vector_store_from_registry_or_db( - vector_store_id=vector_store_id, - prisma_client=prisma_client + vector_store = ( + await self.get_litellm_managed_vector_store_from_registry_or_db( + vector_store_id=vector_store_id, prisma_client=prisma_client + ) ) except Exception as e: verbose_logger.debug( f"Error fetching vector store {vector_store_id} from database: {str(e)}" ) - + if vector_store is not None: # Create a copy to avoid modifying the registry vector_store_copy = vector_store.copy() - + # Merge tool params if they exist if vector_store_id in params_by_id: existing_params = vector_store_copy.get("litellm_params", {}) or {} @@ -403,9 +415,9 @@ class VectorStoreRegistry: # Tool params take precedence over existing params tool_params_dict.update(existing_params) vector_store_copy["litellm_params"] = tool_params_dict - + vector_stores_to_run.append(vector_store_copy) - + return vector_stores_to_run def _get_vector_store_ids_from_tool_calls( diff --git a/litellm/videos/main.py b/litellm/videos/main.py index 2225b9eec7..f6c9bb0057 100644 --- a/litellm/videos/main.py +++ b/litellm/videos/main.py @@ -25,6 +25,7 @@ from litellm.videos.utils import VideoGenerationRequestUtils #################### Initialize provider clients #################### llm_http_handler: BaseLLMHTTPHandler = BaseLLMHTTPHandler() + ##### Video Generation ####################### @client async def avideo_generation( @@ -71,7 +72,8 @@ async def avideo_generation( # get custom llm provider so we can use this for mapping exceptions if custom_llm_provider is None: _, custom_llm_provider, _, _ = litellm.get_llm_provider( - model=model or DEFAULT_VIDEO_ENDPOINT_MODEL, api_base=local_vars.get("api_base", None) + model=model or DEFAULT_VIDEO_ENDPOINT_MODEL, + api_base=local_vars.get("api_base", None), ) func = partial( @@ -170,10 +172,7 @@ def video_generation( # noqa: PLR0915 extra_query: Optional[Dict[str, Any]] = None, extra_body: Optional[Dict[str, Any]] = None, **kwargs, -) -> Union[ - VideoObject, - Coroutine[Any, Any, VideoObject], -]: +) -> Union[VideoObject, Coroutine[Any, Any, VideoObject],]: """ Maps the https://api.openai.com/v1/videos endpoint. @@ -202,20 +201,24 @@ def video_generation( # noqa: PLR0915 ) # get provider config - video_generation_provider_config: Optional[BaseVideoConfig] = ( - ProviderConfigManager.get_provider_video_config( - model=model, - provider=litellm.LlmProviders(custom_llm_provider), - ) + video_generation_provider_config: Optional[ + BaseVideoConfig + ] = ProviderConfigManager.get_provider_video_config( + model=model, + provider=litellm.LlmProviders(custom_llm_provider), ) if video_generation_provider_config is None: - raise ValueError(f"video generation is not supported for {custom_llm_provider}") + raise ValueError( + f"video generation is not supported for {custom_llm_provider}" + ) local_vars.update(kwargs) # Get VideoGenerationOptionalRequestParams with only valid parameters video_generation_optional_params: VideoCreateOptionalRequestParams = ( - VideoGenerationRequestUtils.get_requested_video_generation_optional_param(local_vars) + VideoGenerationRequestUtils.get_requested_video_generation_optional_param( + local_vars + ) ) # Get optional parameters for the video generation API @@ -280,10 +283,7 @@ def video_content( extra_query: Optional[Dict[str, Any]] = None, extra_body: Optional[Dict[str, Any]] = None, **kwargs, -) -> Union[ - bytes, - Coroutine[Any, Any, bytes], -]: +) -> Union[bytes, Coroutine[Any, Any, bytes],]: """ Download video content from OpenAI's video API. @@ -328,15 +328,17 @@ def video_content( litellm_params = GenericLiteLLMParams(**kwargs) # get provider config - video_provider_config: Optional[BaseVideoConfig] = ( - ProviderConfigManager.get_provider_video_config( - model=None, - provider=litellm.LlmProviders(custom_llm_provider), - ) + video_provider_config: Optional[ + BaseVideoConfig + ] = ProviderConfigManager.get_provider_video_config( + model=None, + provider=litellm.LlmProviders(custom_llm_provider), ) if video_provider_config is None: - raise ValueError(f"video support download is not supported for {custom_llm_provider}") + raise ValueError( + f"video support download is not supported for {custom_llm_provider}" + ) local_vars.update(kwargs) # For video content download, we don't need complex optional parameter handling @@ -451,6 +453,7 @@ async def avideo_content( extra_kwargs=kwargs, ) + ##### Video Remix ####################### @client async def avideo_remix( @@ -567,10 +570,7 @@ def video_remix( # noqa: PLR0915 extra_query: Optional[Dict[str, Any]] = None, extra_body: Optional[Dict[str, Any]] = None, **kwargs, -) -> Union[ - VideoObject, - Coroutine[Any, Any, VideoObject], -]: +) -> Union[VideoObject, Coroutine[Any, Any, VideoObject],]: """ Maps the https://api.openai.com/v1/videos/{video_id}/remix endpoint. @@ -600,11 +600,11 @@ def video_remix( # noqa: PLR0915 litellm_params = GenericLiteLLMParams(**kwargs) # get provider config - video_remix_provider_config: Optional[BaseVideoConfig] = ( - ProviderConfigManager.get_provider_video_config( - model=None, - provider=litellm.LlmProviders(custom_llm_provider), - ) + video_remix_provider_config: Optional[ + BaseVideoConfig + ] = ProviderConfigManager.get_provider_video_config( + model=None, + provider=litellm.LlmProviders(custom_llm_provider), ) if video_remix_provider_config is None: @@ -788,10 +788,7 @@ def video_list( # noqa: PLR0915 extra_query: Optional[Dict[str, Any]] = None, extra_body: Optional[Dict[str, Any]] = None, **kwargs, -) -> Union[ - List[VideoObject], - Coroutine[Any, Any, List[VideoObject]], -]: +) -> Union[List[VideoObject], Coroutine[Any, Any, List[VideoObject]],]: """ Maps the https://api.openai.com/v1/videos endpoint. @@ -818,11 +815,11 @@ def video_list( # noqa: PLR0915 litellm_params = GenericLiteLLMParams(**kwargs) # get provider config - video_list_provider_config: Optional[BaseVideoConfig] = ( - ProviderConfigManager.get_provider_video_config( - model=None, - provider=litellm.LlmProviders(custom_llm_provider), - ) + video_list_provider_config: Optional[ + BaseVideoConfig + ] = ProviderConfigManager.get_provider_video_config( + model=None, + provider=litellm.LlmProviders(custom_llm_provider), ) if video_list_provider_config is None: @@ -911,7 +908,6 @@ async def avideo_status( loop = asyncio.get_event_loop() kwargs["async_call"] = True - func = partial( video_status, video_id=video_id, @@ -989,10 +985,7 @@ def video_status( # noqa: PLR0915 extra_query: Optional[Dict[str, Any]] = None, extra_body: Optional[Dict[str, Any]] = None, **kwargs, -) -> Union[ - VideoObject, - Coroutine[Any, Any, VideoObject], -]: +) -> Union[VideoObject, Coroutine[Any, Any, VideoObject],]: """ Retrieve video status from OpenAI's video API. @@ -1044,11 +1037,11 @@ def video_status( # noqa: PLR0915 litellm_params = GenericLiteLLMParams(**kwargs) # get provider config - video_status_provider_config: Optional[BaseVideoConfig] = ( - ProviderConfigManager.get_provider_video_config( - model=None, - provider=litellm.LlmProviders(custom_llm_provider), - ) + video_status_provider_config: Optional[ + BaseVideoConfig + ] = ProviderConfigManager.get_provider_video_config( + model=None, + provider=litellm.LlmProviders(custom_llm_provider), ) if video_status_provider_config is None: diff --git a/litellm/videos/utils.py b/litellm/videos/utils.py index e04ab9fe18..06dfa5d139 100644 --- a/litellm/videos/utils.py +++ b/litellm/videos/utils.py @@ -69,7 +69,8 @@ class VideoGenerationRequestUtils: base_params_raw = { key: value for key, value in params.items() - if key not in {"kwargs", "extra_body", "prompt", "model"} and value is not None + if key not in {"kwargs", "extra_body", "prompt", "model"} + and value is not None } base_params = filter_out_litellm_params(kwargs=base_params_raw) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d0c250fb0d..9b1d81fee4 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -2565,32 +2565,6 @@ "supports_parallel_function_calling": true, "supports_tool_choice": true }, - "azure/gpt-35-turbo-0301": { - "deprecation_date": "2025-02-13", - "input_cost_per_token": 2e-07, - "litellm_provider": "azure", - "max_input_tokens": 4097, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 2e-06, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "azure/gpt-35-turbo-0613": { - "deprecation_date": "2025-02-13", - "input_cost_per_token": 1.5e-06, - "litellm_provider": "azure", - "max_input_tokens": 4097, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 2e-06, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, "azure/gpt-35-turbo-1106": { "deprecation_date": "2025-03-31", "input_cost_per_token": 1e-06, @@ -8185,72 +8159,6 @@ "supports_reasoning": true, "supports_tool_choice": true }, - "chat-bison": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-chat-models", - "max_input_tokens": 8192, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_tool_choice": true - }, - "chat-bison-32k": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-chat-models", - "max_input_tokens": 32000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_tool_choice": true - }, - "chat-bison-32k@002": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-chat-models", - "max_input_tokens": 32000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_tool_choice": true - }, - "chat-bison@001": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-chat-models", - "max_input_tokens": 8192, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_tool_choice": true - }, - "chat-bison@002": { - "deprecation_date": "2025-04-09", - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-chat-models", - "max_input_tokens": 8192, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_tool_choice": true - }, "chatdolphin": { "input_cost_per_token": 5e-07, "litellm_provider": "nlp_cloud", @@ -8288,60 +8196,6 @@ "/v1/audio/transcriptions" ] }, - "claude-3-5-haiku-20241022": { - "cache_creation_input_token_cost": 1e-06, - "cache_creation_input_token_cost_above_1hr": 6e-06, - "cache_read_input_token_cost": 8e-08, - "deprecation_date": "2025-10-01", - "input_cost_per_token": 8e-07, - "litellm_provider": "anthropic", - "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 4e-06, - "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 - }, - "supports_assistant_prefill": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tool_use_system_prompt_tokens": 264 - }, - "claude-3-5-haiku-latest": { - "cache_creation_input_token_cost": 1.25e-06, - "cache_creation_input_token_cost_above_1hr": 6e-06, - "cache_read_input_token_cost": 1e-07, - "deprecation_date": "2025-10-01", - "input_cost_per_token": 1e-06, - "litellm_provider": "anthropic", - "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 5e-06, - "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 - }, - "supports_assistant_prefill": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tool_use_system_prompt_tokens": 264 - }, "claude-haiku-4-5-20251001": { "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, @@ -8384,83 +8238,6 @@ "supports_tool_choice": true, "supports_vision": true }, - "claude-3-5-sonnet-20240620": { - "cache_creation_input_token_cost": 3.75e-06, - "cache_creation_input_token_cost_above_1hr": 6e-06, - "cache_read_input_token_cost": 3e-07, - "deprecation_date": "2025-06-01", - "input_cost_per_token": 3e-06, - "litellm_provider": "anthropic", - "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.5e-05, - "supports_assistant_prefill": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 - }, - "claude-3-5-sonnet-20241022": { - "cache_creation_input_token_cost": 3.75e-06, - "cache_creation_input_token_cost_above_1hr": 6e-06, - "cache_read_input_token_cost": 3e-07, - "deprecation_date": "2025-10-01", - "input_cost_per_token": 3e-06, - "litellm_provider": "anthropic", - "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.5e-05, - "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 - }, - "supports_assistant_prefill": true, - "supports_computer_use": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tool_use_system_prompt_tokens": 159 - }, - "claude-3-5-sonnet-latest": { - "cache_creation_input_token_cost": 3.75e-06, - "cache_creation_input_token_cost_above_1hr": 6e-06, - "cache_read_input_token_cost": 3e-07, - "deprecation_date": "2025-06-01", - "input_cost_per_token": 3e-06, - "litellm_provider": "anthropic", - "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.5e-05, - "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 - }, - "supports_assistant_prefill": true, - "supports_computer_use": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tool_use_system_prompt_tokens": 159 - }, "claude-3-7-sonnet-20250219": { "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, @@ -8490,34 +8267,6 @@ "supports_web_search": true, "tool_use_system_prompt_tokens": 159 }, - "claude-3-7-sonnet-latest": { - "cache_creation_input_token_cost": 3.75e-06, - "cache_creation_input_token_cost_above_1hr": 6e-06, - "cache_read_input_token_cost": 3e-07, - "deprecation_date": "2025-06-01", - "input_cost_per_token": 3e-06, - "litellm_provider": "anthropic", - "max_input_tokens": 200000, - "max_output_tokens": 64000, - "max_tokens": 64000, - "mode": "chat", - "output_cost_per_token": 1.5e-05, - "search_context_cost_per_query": { - "search_context_size_high": 0.01, - "search_context_size_low": 0.01, - "search_context_size_medium": 0.01 - }, - "supports_assistant_prefill": true, - "supports_computer_use": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 159 - }, "claude-3-haiku-20240307": { "cache_creation_input_token_cost": 3e-07, "cache_creation_input_token_cost_above_1hr": 6e-06, @@ -8557,26 +8306,6 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 395 }, - "claude-3-opus-latest": { - "cache_creation_input_token_cost": 1.875e-05, - "cache_creation_input_token_cost_above_1hr": 6e-06, - "cache_read_input_token_cost": 1.5e-06, - "deprecation_date": "2025-03-01", - "input_cost_per_token": 1.5e-05, - "litellm_provider": "anthropic", - "max_input_tokens": 200000, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 7.5e-05, - "supports_assistant_prefill": true, - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_tool_choice": true, - "supports_vision": true, - "tool_use_system_prompt_tokens": 395 - }, "claude-4-opus-20250514": { "cache_creation_input_token_cost": 1.875e-05, "cache_read_input_token_cost": 1.5e-06, @@ -9025,185 +8754,6 @@ "mode": "chat", "output_cost_per_token": 1.923e-06 }, - "code-bison": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-text-models", - "max_input_tokens": 6144, - "max_output_tokens": 1024, - "max_tokens": 1024, - "mode": "chat", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_tool_choice": true - }, - "code-bison-32k@002": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-text-models", - "max_input_tokens": 6144, - "max_output_tokens": 1024, - "max_tokens": 1024, - "mode": "completion", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "code-bison32k": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-text-models", - "max_input_tokens": 6144, - "max_output_tokens": 1024, - "max_tokens": 1024, - "mode": "completion", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "code-bison@001": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-text-models", - "max_input_tokens": 6144, - "max_output_tokens": 1024, - "max_tokens": 1024, - "mode": "completion", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "code-bison@002": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-text-models", - "max_input_tokens": 6144, - "max_output_tokens": 1024, - "max_tokens": 1024, - "mode": "completion", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "code-gecko": { - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-text-models", - "max_input_tokens": 2048, - "max_output_tokens": 64, - "max_tokens": 64, - "mode": "completion", - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "code-gecko-latest": { - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-text-models", - "max_input_tokens": 2048, - "max_output_tokens": 64, - "max_tokens": 64, - "mode": "completion", - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "code-gecko@001": { - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-text-models", - "max_input_tokens": 2048, - "max_output_tokens": 64, - "max_tokens": 64, - "mode": "completion", - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "code-gecko@002": { - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-text-models", - "max_input_tokens": 2048, - "max_output_tokens": 64, - "max_tokens": 64, - "mode": "completion", - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "codechat-bison": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-chat-models", - "max_input_tokens": 6144, - "max_output_tokens": 1024, - "max_tokens": 1024, - "mode": "chat", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_tool_choice": true - }, - "codechat-bison-32k": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-chat-models", - "max_input_tokens": 32000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_tool_choice": true - }, - "codechat-bison-32k@002": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-chat-models", - "max_input_tokens": 32000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_tool_choice": true - }, - "codechat-bison@001": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-chat-models", - "max_input_tokens": 6144, - "max_output_tokens": 1024, - "max_tokens": 1024, - "mode": "chat", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_tool_choice": true - }, - "codechat-bison@002": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-chat-models", - "max_input_tokens": 6144, - "max_output_tokens": 1024, - "max_tokens": 1024, - "mode": "chat", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_tool_choice": true - }, - "codechat-bison@latest": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-code-chat-models", - "max_input_tokens": 6144, - "max_output_tokens": 1024, - "max_tokens": 1024, - "mode": "chat", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_tool_choice": true - }, "codestral/codestral-2405": { "input_cost_per_token": 0.0, "litellm_provider": "codestral", @@ -13718,475 +13268,6 @@ "supports_response_schema": true, "supports_tool_choice": true }, - "gemini-1.0-pro": { - "input_cost_per_character": 1.25e-07, - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "input_cost_per_video_per_second": 0.002, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 32760, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 3.75e-07, - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-1.0-pro-001": { - "deprecation_date": "2025-04-09", - "input_cost_per_character": 1.25e-07, - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "input_cost_per_video_per_second": 0.002, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 32760, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 3.75e-07, - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-1.0-pro-002": { - "deprecation_date": "2025-04-09", - "input_cost_per_character": 1.25e-07, - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "input_cost_per_video_per_second": 0.002, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 32760, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 3.75e-07, - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-1.0-pro-vision": { - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "litellm_provider": "vertex_ai-vision-models", - "max_images_per_prompt": 16, - "max_input_tokens": 16384, - "max_output_tokens": 2048, - "max_tokens": 2048, - "max_video_length": 2, - "max_videos_per_prompt": 1, - "mode": "chat", - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.0-pro-vision-001": { - "deprecation_date": "2025-04-09", - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "litellm_provider": "vertex_ai-vision-models", - "max_images_per_prompt": 16, - "max_input_tokens": 16384, - "max_output_tokens": 2048, - "max_tokens": 2048, - "max_video_length": 2, - "max_videos_per_prompt": 1, - "mode": "chat", - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.0-ultra": { - "input_cost_per_character": 1.25e-07, - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "input_cost_per_video_per_second": 0.002, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 8192, - "max_output_tokens": 2048, - "max_tokens": 2048, - "mode": "chat", - "output_cost_per_character": 3.75e-07, - "output_cost_per_token": 1.5e-06, - "source": "As of Jun, 2024. There is no available doc on vertex ai pricing gemini-1.0-ultra-001. Using gemini-1.0-pro pricing. Got max_tokens info here: https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-1.0-ultra-001": { - "input_cost_per_character": 1.25e-07, - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "input_cost_per_video_per_second": 0.002, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 8192, - "max_output_tokens": 2048, - "max_tokens": 2048, - "mode": "chat", - "output_cost_per_character": 3.75e-07, - "output_cost_per_token": 1.5e-06, - "source": "As of Jun, 2024. There is no available doc on vertex ai pricing gemini-1.0-ultra-001. Using gemini-1.0-pro pricing. Got max_tokens info here: https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-1.5-flash": { - "deprecation_date": "2025-09-29", - "input_cost_per_audio_per_second": 2e-06, - "input_cost_per_audio_per_second_above_128k_tokens": 4e-06, - "input_cost_per_character": 1.875e-08, - "input_cost_per_character_above_128k_tokens": 2.5e-07, - "input_cost_per_image": 2e-05, - "input_cost_per_image_above_128k_tokens": 4e-05, - "input_cost_per_token": 7.5e-08, - "input_cost_per_token_above_128k_tokens": 1e-06, - "input_cost_per_video_per_second": 2e-05, - "input_cost_per_video_per_second_above_128k_tokens": 4e-05, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_character": 7.5e-08, - "output_cost_per_character_above_128k_tokens": 1.5e-07, - "output_cost_per_token": 3e-07, - "output_cost_per_token_above_128k_tokens": 6e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.5-flash-001": { - "deprecation_date": "2025-05-24", - "input_cost_per_audio_per_second": 2e-06, - "input_cost_per_audio_per_second_above_128k_tokens": 4e-06, - "input_cost_per_character": 1.875e-08, - "input_cost_per_character_above_128k_tokens": 2.5e-07, - "input_cost_per_image": 2e-05, - "input_cost_per_image_above_128k_tokens": 4e-05, - "input_cost_per_token": 7.5e-08, - "input_cost_per_token_above_128k_tokens": 1e-06, - "input_cost_per_video_per_second": 2e-05, - "input_cost_per_video_per_second_above_128k_tokens": 4e-05, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_character": 7.5e-08, - "output_cost_per_character_above_128k_tokens": 1.5e-07, - "output_cost_per_token": 3e-07, - "output_cost_per_token_above_128k_tokens": 6e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.5-flash-002": { - "deprecation_date": "2025-09-24", - "input_cost_per_audio_per_second": 2e-06, - "input_cost_per_audio_per_second_above_128k_tokens": 4e-06, - "input_cost_per_character": 1.875e-08, - "input_cost_per_character_above_128k_tokens": 2.5e-07, - "input_cost_per_image": 2e-05, - "input_cost_per_image_above_128k_tokens": 4e-05, - "input_cost_per_token": 7.5e-08, - "input_cost_per_token_above_128k_tokens": 1e-06, - "input_cost_per_video_per_second": 2e-05, - "input_cost_per_video_per_second_above_128k_tokens": 4e-05, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_character": 7.5e-08, - "output_cost_per_character_above_128k_tokens": 1.5e-07, - "output_cost_per_token": 3e-07, - "output_cost_per_token_above_128k_tokens": 6e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-1.5-flash", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.5-flash-exp-0827": { - "deprecation_date": "2025-09-29", - "input_cost_per_audio_per_second": 2e-06, - "input_cost_per_audio_per_second_above_128k_tokens": 4e-06, - "input_cost_per_character": 1.875e-08, - "input_cost_per_character_above_128k_tokens": 2.5e-07, - "input_cost_per_image": 2e-05, - "input_cost_per_image_above_128k_tokens": 4e-05, - "input_cost_per_token": 4.688e-09, - "input_cost_per_token_above_128k_tokens": 1e-06, - "input_cost_per_video_per_second": 2e-05, - "input_cost_per_video_per_second_above_128k_tokens": 4e-05, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_character": 1.875e-08, - "output_cost_per_character_above_128k_tokens": 3.75e-08, - "output_cost_per_token": 4.6875e-09, - "output_cost_per_token_above_128k_tokens": 9.375e-09, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.5-flash-preview-0514": { - "deprecation_date": "2025-09-29", - "input_cost_per_audio_per_second": 2e-06, - "input_cost_per_audio_per_second_above_128k_tokens": 4e-06, - "input_cost_per_character": 1.875e-08, - "input_cost_per_character_above_128k_tokens": 2.5e-07, - "input_cost_per_image": 2e-05, - "input_cost_per_image_above_128k_tokens": 4e-05, - "input_cost_per_token": 7.5e-08, - "input_cost_per_token_above_128k_tokens": 1e-06, - "input_cost_per_video_per_second": 2e-05, - "input_cost_per_video_per_second_above_128k_tokens": 4e-05, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_character": 1.875e-08, - "output_cost_per_character_above_128k_tokens": 3.75e-08, - "output_cost_per_token": 4.6875e-09, - "output_cost_per_token_above_128k_tokens": 9.375e-09, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.5-pro": { - "deprecation_date": "2025-09-29", - "input_cost_per_audio_per_second": 3.125e-05, - "input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05, - "input_cost_per_character": 3.125e-07, - "input_cost_per_character_above_128k_tokens": 6.25e-07, - "input_cost_per_image": 0.00032875, - "input_cost_per_image_above_128k_tokens": 0.0006575, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_128k_tokens": 2.5e-06, - "input_cost_per_video_per_second": 0.00032875, - "input_cost_per_video_per_second_above_128k_tokens": 0.0006575, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 2097152, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 1.25e-06, - "output_cost_per_character_above_128k_tokens": 2.5e-06, - "output_cost_per_token": 5e-06, - "output_cost_per_token_above_128k_tokens": 1e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.5-pro-001": { - "deprecation_date": "2025-05-24", - "input_cost_per_audio_per_second": 3.125e-05, - "input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05, - "input_cost_per_character": 3.125e-07, - "input_cost_per_character_above_128k_tokens": 6.25e-07, - "input_cost_per_image": 0.00032875, - "input_cost_per_image_above_128k_tokens": 0.0006575, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_128k_tokens": 2.5e-06, - "input_cost_per_video_per_second": 0.00032875, - "input_cost_per_video_per_second_above_128k_tokens": 0.0006575, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 1.25e-06, - "output_cost_per_character_above_128k_tokens": 2.5e-06, - "output_cost_per_token": 5e-06, - "output_cost_per_token_above_128k_tokens": 1e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.5-pro-002": { - "deprecation_date": "2025-09-24", - "input_cost_per_audio_per_second": 3.125e-05, - "input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05, - "input_cost_per_character": 3.125e-07, - "input_cost_per_character_above_128k_tokens": 6.25e-07, - "input_cost_per_image": 0.00032875, - "input_cost_per_image_above_128k_tokens": 0.0006575, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_128k_tokens": 2.5e-06, - "input_cost_per_video_per_second": 0.00032875, - "input_cost_per_video_per_second_above_128k_tokens": 0.0006575, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 2097152, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 1.25e-06, - "output_cost_per_character_above_128k_tokens": 2.5e-06, - "output_cost_per_token": 5e-06, - "output_cost_per_token_above_128k_tokens": 1e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-1.5-pro", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gemini-1.5-pro-preview-0215": { - "deprecation_date": "2025-09-29", - "input_cost_per_audio_per_second": 3.125e-05, - "input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05, - "input_cost_per_character": 3.125e-07, - "input_cost_per_character_above_128k_tokens": 6.25e-07, - "input_cost_per_image": 0.00032875, - "input_cost_per_image_above_128k_tokens": 0.0006575, - "input_cost_per_token": 7.8125e-08, - "input_cost_per_token_above_128k_tokens": 1.5625e-07, - "input_cost_per_video_per_second": 0.00032875, - "input_cost_per_video_per_second_above_128k_tokens": 0.0006575, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 1.25e-06, - "output_cost_per_character_above_128k_tokens": 2.5e-06, - "output_cost_per_token": 3.125e-07, - "output_cost_per_token_above_128k_tokens": 6.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, - "gemini-1.5-pro-preview-0409": { - "deprecation_date": "2025-09-29", - "input_cost_per_audio_per_second": 3.125e-05, - "input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05, - "input_cost_per_character": 3.125e-07, - "input_cost_per_character_above_128k_tokens": 6.25e-07, - "input_cost_per_image": 0.00032875, - "input_cost_per_image_above_128k_tokens": 0.0006575, - "input_cost_per_token": 7.8125e-08, - "input_cost_per_token_above_128k_tokens": 1.5625e-07, - "input_cost_per_video_per_second": 0.00032875, - "input_cost_per_video_per_second_above_128k_tokens": 0.0006575, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 1.25e-06, - "output_cost_per_character_above_128k_tokens": 2.5e-06, - "output_cost_per_token": 3.125e-07, - "output_cost_per_token_above_128k_tokens": 6.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_tool_choice": true - }, - "gemini-1.5-pro-preview-0514": { - "deprecation_date": "2025-09-29", - "input_cost_per_audio_per_second": 3.125e-05, - "input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05, - "input_cost_per_character": 3.125e-07, - "input_cost_per_character_above_128k_tokens": 6.25e-07, - "input_cost_per_image": 0.00032875, - "input_cost_per_image_above_128k_tokens": 0.0006575, - "input_cost_per_token": 7.8125e-08, - "input_cost_per_token_above_128k_tokens": 1.5625e-07, - "input_cost_per_video_per_second": 0.00032875, - "input_cost_per_video_per_second_above_128k_tokens": 0.0006575, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 1.25e-06, - "output_cost_per_character_above_128k_tokens": 2.5e-06, - "output_cost_per_token": 3.125e-07, - "output_cost_per_token_above_128k_tokens": 6.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, "gemini-2.0-flash": { "cache_read_input_token_cost": 2.5e-08, "deprecation_date": "2026-06-01", @@ -14265,54 +13346,6 @@ "supports_vision": true, "supports_web_search": true }, - "gemini-2.0-flash-exp": { - "cache_read_input_token_cost": 3.75e-08, - "input_cost_per_audio_per_second": 0, - "input_cost_per_audio_per_second_above_128k_tokens": 0, - "input_cost_per_character": 0, - "input_cost_per_character_above_128k_tokens": 0, - "input_cost_per_image": 0, - "input_cost_per_image_above_128k_tokens": 0, - "input_cost_per_token": 1.5e-07, - "input_cost_per_token_above_128k_tokens": 0, - "input_cost_per_video_per_second": 0, - "input_cost_per_video_per_second_above_128k_tokens": 0, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_character": 0, - "output_cost_per_character_above_128k_tokens": 0, - "output_cost_per_token": 6e-07, - "output_cost_per_token_above_128k_tokens": 0, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_audio_output": true, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true - }, "gemini-2.0-flash-lite": { "cache_read_input_token_cost": 1.875e-08, "deprecation_date": "2026-06-01", @@ -14385,235 +13418,6 @@ "supports_vision": true, "supports_web_search": true }, - "gemini-2.0-flash-live-preview-04-09": { - "cache_read_input_token_cost": 7.5e-08, - "input_cost_per_audio_token": 3e-06, - "input_cost_per_image": 3e-06, - "input_cost_per_token": 5e-07, - "input_cost_per_video_per_second": 3e-06, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_audio_token": 1.2e-05, - "output_cost_per_token": 2e-06, - "rpm": 10, - "source": "https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/gemini#gemini-2-0-flash-live-preview-04-09", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "audio" - ], - "supports_audio_output": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 250000 - }, - "gemini-2.0-flash-preview-image-generation": { - "deprecation_date": "2025-11-14", - "cache_read_input_token_cost": 2.5e-08, - "input_cost_per_audio_token": 7e-07, - "input_cost_per_token": 1e-07, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 4e-07, - "source": "https://ai.google.dev/pricing#2_0flash", - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_audio_input": true, - "supports_audio_output": true, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true - }, - "gemini-2.0-flash-thinking-exp": { - "deprecation_date": "2025-12-02", - "cache_read_input_token_cost": 0.0, - "input_cost_per_audio_per_second": 0, - "input_cost_per_audio_per_second_above_128k_tokens": 0, - "input_cost_per_character": 0, - "input_cost_per_character_above_128k_tokens": 0, - "input_cost_per_image": 0, - "input_cost_per_image_above_128k_tokens": 0, - "input_cost_per_token": 0, - "input_cost_per_token_above_128k_tokens": 0, - "input_cost_per_video_per_second": 0, - "input_cost_per_video_per_second_above_128k_tokens": 0, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_character": 0, - "output_cost_per_character_above_128k_tokens": 0, - "output_cost_per_token": 0, - "output_cost_per_token_above_128k_tokens": 0, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-2.0-flash", - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_audio_output": true, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true - }, - "gemini-2.0-flash-thinking-exp-01-21": { - "deprecation_date": "2025-12-02", - "cache_read_input_token_cost": 0.0, - "input_cost_per_audio_per_second": 0, - "input_cost_per_audio_per_second_above_128k_tokens": 0, - "input_cost_per_character": 0, - "input_cost_per_character_above_128k_tokens": 0, - "input_cost_per_image": 0, - "input_cost_per_image_above_128k_tokens": 0, - "input_cost_per_token": 0, - "input_cost_per_token_above_128k_tokens": 0, - "input_cost_per_video_per_second": 0, - "input_cost_per_video_per_second_above_128k_tokens": 0, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65536, - "max_pdf_size_mb": 30, - "max_tokens": 65536, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_character": 0, - "output_cost_per_character_above_128k_tokens": 0, - "output_cost_per_token": 0, - "output_cost_per_token_above_128k_tokens": 0, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-2.0-flash", - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_audio_output": false, - "supports_function_calling": false, - "supports_parallel_function_calling": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": false, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true - }, - "gemini-2.0-pro-exp-02-05": { - "cache_read_input_token_cost": 3.125e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 2097152, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_input": true, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_video_input": true, - "supports_vision": true, - "supports_web_search": true - }, "gemini-2.5-flash": { "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, @@ -14708,57 +13512,6 @@ "supports_web_search": false, "tpm": 8000000 }, - "gemini-2.5-flash-image-preview": { - "deprecation_date": "2026-01-15", - "cache_read_input_token_cost": 7.5e-08, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_image_token": 3e-07, - "input_cost_per_token": 3e-07, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "image_generation", - "output_cost_per_image": 0.039, - "output_cost_per_image_token": 3e-05, - "output_cost_per_reasoning_token": 3e-05, - "output_cost_per_token": 3e-05, - "rpm": 100000, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 8000000 - }, "gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -15181,96 +13934,6 @@ "supports_vision": true, "supports_web_search": true }, - "gemini-2.5-flash-preview-04-17": { - "cache_read_input_token_cost": 3.75e-08, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 1.5e-07, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_reasoning_token": 3.5e-06, - "output_cost_per_token": 6e-07, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true - }, - "gemini-2.5-flash-preview-05-20": { - "deprecation_date": "2025-11-18", - "cache_read_input_token_cost": 7.5e-08, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_reasoning_token": 2.5e-06, - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true - }, "gemini-2.5-pro": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, @@ -15703,193 +14366,6 @@ "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, "supports_service_tier": true }, - "gemini-2.5-pro-exp-03-25": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_input": true, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_video_input": true, - "supports_vision": true, - "supports_web_search": true - }, - "gemini-2.5-pro-preview-03-25": { - "deprecation_date": "2025-12-02", - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_audio_token": 1.25e-06, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true - }, - "gemini-2.5-pro-preview-05-06": { - "deprecation_date": "2025-12-02", - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_audio_token": 1.25e-06, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supported_regions": [ - "global" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true - }, - "gemini-2.5-pro-preview-06-05": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_audio_token": 1.25e-06, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true - }, "gemini-2.5-pro-preview-tts": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, @@ -16072,63 +14548,23 @@ "max_input_tokens": 1000000, "max_output_tokens": 8192, "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 0, + "mode": "embedding", "output_cost_per_token": 0, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/gemini-experimental", - "supports_function_calling": false, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-pro": { - "input_cost_per_character": 1.25e-07, - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "input_cost_per_video_per_second": 0.002, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 32760, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 3.75e-07, - "output_cost_per_token": 1.5e-06, + "output_vector_size": 3072, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true + "uses_embed_content": true }, - "gemini-pro-experimental": { - "input_cost_per_character": 0, - "input_cost_per_token": 0, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 1000000, - "max_output_tokens": 8192, + "vertex_ai/gemini-embedding-2-preview": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "vertex_ai", + "max_input_tokens": 8192, "max_tokens": 8192, - "mode": "chat", - "output_cost_per_character": 0, + "mode": "embedding", "output_cost_per_token": 0, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/gemini-experimental", - "supports_function_calling": false, - "supports_parallel_function_calling": true, - "supports_tool_choice": true - }, - "gemini-pro-vision": { - "input_cost_per_image": 0.0025, - "input_cost_per_token": 5e-07, - "litellm_provider": "vertex_ai-vision-models", - "max_images_per_prompt": 16, - "max_input_tokens": 16384, - "max_output_tokens": 2048, - "max_tokens": 2048, - "max_video_length": 2, - "max_videos_per_prompt": 1, - "mode": "chat", - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_tool_choice": true, - "supports_vision": true + "output_vector_size": 3072, + "source": "https://ai.google.dev/gemini-api/docs/embeddings#multimodal", + "supports_multimodal": true, + "uses_embed_content": true }, "gemini/gemini-embedding-001": { "input_cost_per_token": 1.5e-07, @@ -16163,339 +14599,15 @@ "input_cost_per_token": 7.5e-08, "input_cost_per_token_above_128k_tokens": 1.5e-07, "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, + "max_input_tokens": 8192, "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 3e-07, - "output_cost_per_token_above_128k_tokens": 6e-07, - "rpm": 2000, - "source": "https://ai.google.dev/pricing", - "supports_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 4000000 - }, - "gemini/gemini-1.5-flash-001": { - "cache_creation_input_token_cost": 1e-06, - "cache_read_input_token_cost": 1.875e-08, - "deprecation_date": "2025-05-24", - "input_cost_per_token": 7.5e-08, - "input_cost_per_token_above_128k_tokens": 1.5e-07, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 3e-07, - "output_cost_per_token_above_128k_tokens": 6e-07, - "rpm": 2000, - "source": "https://ai.google.dev/pricing", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 4000000 - }, - "gemini/gemini-1.5-flash-002": { - "cache_creation_input_token_cost": 1e-06, - "cache_read_input_token_cost": 1.875e-08, - "deprecation_date": "2025-09-24", - "input_cost_per_token": 7.5e-08, - "input_cost_per_token_above_128k_tokens": 1.5e-07, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 3e-07, - "output_cost_per_token_above_128k_tokens": 6e-07, - "rpm": 2000, - "source": "https://ai.google.dev/pricing", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 4000000 - }, - "gemini/gemini-1.5-flash-8b": { - "deprecation_date": "2025-09-29", - "input_cost_per_token": 0, - "input_cost_per_token_above_128k_tokens": 0, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", + "mode": "embedding", "output_cost_per_token": 0, - "output_cost_per_token_above_128k_tokens": 0, - "rpm": 4000, - "source": "https://ai.google.dev/pricing", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 4000000 - }, - "gemini/gemini-1.5-flash-8b-exp-0827": { - "deprecation_date": "2025-09-29", - "input_cost_per_token": 0, - "input_cost_per_token_above_128k_tokens": 0, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 0, - "output_cost_per_token_above_128k_tokens": 0, - "rpm": 4000, - "source": "https://ai.google.dev/pricing", - "supports_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 4000000 - }, - "gemini/gemini-1.5-flash-8b-exp-0924": { - "deprecation_date": "2025-09-29", - "input_cost_per_token": 0, - "input_cost_per_token_above_128k_tokens": 0, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 0, - "output_cost_per_token_above_128k_tokens": 0, - "rpm": 4000, - "source": "https://ai.google.dev/pricing", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 4000000 - }, - "gemini/gemini-1.5-flash-exp-0827": { - "deprecation_date": "2025-09-29", - "input_cost_per_token": 0, - "input_cost_per_token_above_128k_tokens": 0, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 0, - "output_cost_per_token_above_128k_tokens": 0, - "rpm": 2000, - "source": "https://ai.google.dev/pricing", - "supports_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 4000000 - }, - "gemini/gemini-1.5-flash-latest": { - "deprecation_date": "2025-09-29", - "input_cost_per_token": 7.5e-08, - "input_cost_per_token_above_128k_tokens": 1.5e-07, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 3e-07, - "output_cost_per_token_above_128k_tokens": 6e-07, - "rpm": 2000, - "source": "https://ai.google.dev/pricing", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 4000000 - }, - "gemini/gemini-1.5-pro": { - "deprecation_date": "2025-09-29", - "input_cost_per_token": 3.5e-06, - "input_cost_per_token_above_128k_tokens": 7e-06, - "litellm_provider": "gemini", - "max_input_tokens": 2097152, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.05e-05, - "output_cost_per_token_above_128k_tokens": 2.1e-05, - "rpm": 1000, - "source": "https://ai.google.dev/pricing", - "supports_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 4000000 - }, - "gemini/gemini-1.5-pro-001": { - "deprecation_date": "2025-05-24", - "input_cost_per_token": 3.5e-06, - "input_cost_per_token_above_128k_tokens": 7e-06, - "litellm_provider": "gemini", - "max_input_tokens": 2097152, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.05e-05, - "output_cost_per_token_above_128k_tokens": 2.1e-05, - "rpm": 1000, - "source": "https://ai.google.dev/pricing", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 4000000 - }, - "gemini/gemini-1.5-pro-002": { - "deprecation_date": "2025-09-24", - "input_cost_per_token": 3.5e-06, - "input_cost_per_token_above_128k_tokens": 7e-06, - "litellm_provider": "gemini", - "max_input_tokens": 2097152, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.05e-05, - "output_cost_per_token_above_128k_tokens": 2.1e-05, - "rpm": 1000, - "source": "https://ai.google.dev/pricing", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 4000000 - }, - "gemini/gemini-1.5-pro-exp-0801": { - "deprecation_date": "2025-09-29", - "input_cost_per_token": 3.5e-06, - "input_cost_per_token_above_128k_tokens": 7e-06, - "litellm_provider": "gemini", - "max_input_tokens": 2097152, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.05e-05, - "output_cost_per_token_above_128k_tokens": 2.1e-05, - "rpm": 1000, - "source": "https://ai.google.dev/pricing", - "supports_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 4000000 - }, - "gemini/gemini-1.5-pro-exp-0827": { - "deprecation_date": "2025-09-29", - "input_cost_per_token": 0, - "input_cost_per_token_above_128k_tokens": 0, - "litellm_provider": "gemini", - "max_input_tokens": 2097152, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 0, - "output_cost_per_token_above_128k_tokens": 0, - "rpm": 1000, - "source": "https://ai.google.dev/pricing", - "supports_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 4000000 - }, - "gemini/gemini-1.5-pro-latest": { - "deprecation_date": "2025-09-29", - "input_cost_per_token": 3.5e-06, - "input_cost_per_token_above_128k_tokens": 7e-06, - "litellm_provider": "gemini", - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.05e-06, - "output_cost_per_token_above_128k_tokens": 2.1e-05, - "rpm": 1000, - "source": "https://ai.google.dev/pricing", - "supports_function_calling": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 4000000 + "output_vector_size": 3072, + "rpm": 10000, + "source": "https://ai.google.dev/gemini-api/docs/embeddings#multimodal", + "supports_multimodal": true, + "tpm": 10000000 }, "gemini/gemini-2.0-flash": { "cache_read_input_token_cost": 2.5e-08, @@ -16577,55 +14689,6 @@ "supports_web_search": true, "tpm": 10000000 }, - "gemini/gemini-2.0-flash-exp": { - "cache_read_input_token_cost": 0.0, - "input_cost_per_audio_per_second": 0, - "input_cost_per_audio_per_second_above_128k_tokens": 0, - "input_cost_per_character": 0, - "input_cost_per_character_above_128k_tokens": 0, - "input_cost_per_image": 0, - "input_cost_per_image_above_128k_tokens": 0, - "input_cost_per_token": 0, - "input_cost_per_token_above_128k_tokens": 0, - "input_cost_per_video_per_second": 0, - "input_cost_per_video_per_second_above_128k_tokens": 0, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_character": 0, - "output_cost_per_character_above_128k_tokens": 0, - "output_cost_per_token": 0, - "output_cost_per_token_above_128k_tokens": 0, - "rpm": 10, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-2.0-flash", - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_audio_output": true, - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 4000000 - }, "gemini/gemini-2.0-flash-lite": { "cache_read_input_token_cost": 1.875e-08, "deprecation_date": "2026-06-01", @@ -16663,275 +14726,6 @@ "supports_web_search": true, "tpm": 4000000 }, - "gemini/gemini-2.0-flash-lite-preview-02-05": { - "deprecation_date": "2025-12-09", - "cache_read_input_token_cost": 1.875e-08, - "input_cost_per_audio_token": 7.5e-08, - "input_cost_per_token": 7.5e-08, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 3e-07, - "rpm": 60000, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-2.0-flash-lite", - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 10000000 - }, - "gemini/gemini-2.0-flash-live-001": { - "deprecation_date": "2025-12-09", - "cache_read_input_token_cost": 7.5e-08, - "input_cost_per_audio_token": 2.1e-06, - "input_cost_per_image": 2.1e-06, - "input_cost_per_token": 3.5e-07, - "input_cost_per_video_per_second": 2.1e-06, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_audio_token": 8.5e-06, - "output_cost_per_token": 1.5e-06, - "rpm": 10, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2-0-flash-live-001", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "audio" - ], - "supports_audio_output": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 250000 - }, - "gemini/gemini-2.0-flash-preview-image-generation": { - "deprecation_date": "2025-11-14", - "cache_read_input_token_cost": 2.5e-08, - "input_cost_per_audio_token": 7e-07, - "input_cost_per_token": 1e-07, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 4e-07, - "rpm": 10000, - "source": "https://ai.google.dev/pricing#2_0flash", - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_audio_input": true, - "supports_audio_output": true, - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 10000000 - }, - "gemini/gemini-2.0-flash-thinking-exp": { - "deprecation_date": "2025-12-02", - "cache_read_input_token_cost": 0.0, - "input_cost_per_audio_per_second": 0, - "input_cost_per_audio_per_second_above_128k_tokens": 0, - "input_cost_per_character": 0, - "input_cost_per_character_above_128k_tokens": 0, - "input_cost_per_image": 0, - "input_cost_per_image_above_128k_tokens": 0, - "input_cost_per_token": 0, - "input_cost_per_token_above_128k_tokens": 0, - "input_cost_per_video_per_second": 0, - "input_cost_per_video_per_second_above_128k_tokens": 0, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65536, - "max_pdf_size_mb": 30, - "max_tokens": 65536, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_character": 0, - "output_cost_per_character_above_128k_tokens": 0, - "output_cost_per_token": 0, - "output_cost_per_token_above_128k_tokens": 0, - "rpm": 10, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-2.0-flash", - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_audio_output": true, - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 4000000 - }, - "gemini/gemini-2.0-flash-thinking-exp-01-21": { - "deprecation_date": "2025-12-02", - "cache_read_input_token_cost": 0.0, - "input_cost_per_audio_per_second": 0, - "input_cost_per_audio_per_second_above_128k_tokens": 0, - "input_cost_per_character": 0, - "input_cost_per_character_above_128k_tokens": 0, - "input_cost_per_image": 0, - "input_cost_per_image_above_128k_tokens": 0, - "input_cost_per_token": 0, - "input_cost_per_token_above_128k_tokens": 0, - "input_cost_per_video_per_second": 0, - "input_cost_per_video_per_second_above_128k_tokens": 0, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65536, - "max_pdf_size_mb": 30, - "max_tokens": 65536, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_character": 0, - "output_cost_per_character_above_128k_tokens": 0, - "output_cost_per_token": 0, - "output_cost_per_token_above_128k_tokens": 0, - "rpm": 10, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-2.0-flash", - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_audio_output": true, - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 4000000 - }, - "gemini/gemini-2.0-pro-exp-02-05": { - "cache_read_input_token_cost": 0.0, - "input_cost_per_audio_per_second": 0, - "input_cost_per_audio_per_second_above_128k_tokens": 0, - "input_cost_per_character": 0, - "input_cost_per_character_above_128k_tokens": 0, - "input_cost_per_image": 0, - "input_cost_per_image_above_128k_tokens": 0, - "input_cost_per_token": 0, - "input_cost_per_token_above_128k_tokens": 0, - "input_cost_per_video_per_second": 0, - "input_cost_per_video_per_second_above_128k_tokens": 0, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 2097152, - "max_output_tokens": 8192, - "max_pdf_size_mb": 30, - "max_tokens": 8192, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_character": 0, - "output_cost_per_character_above_128k_tokens": 0, - "output_cost_per_token": 0, - "output_cost_per_token_above_128k_tokens": 0, - "rpm": 2, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supports_audio_input": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_video_input": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 1000000 - }, "gemini/gemini-2.5-flash": { "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, @@ -17029,56 +14823,6 @@ "supports_web_search": true, "tpm": 8000000 }, - "gemini/gemini-2.5-flash-image-preview": { - "deprecation_date": "2026-01-15", - "cache_read_input_token_cost": 7.5e-08, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "image_generation", - "output_cost_per_image": 0.039, - "output_cost_per_image_token": 3e-05, - "output_cost_per_reasoning_token": 3e-05, - "output_cost_per_token": 3e-05, - "rpm": 100000, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 8000000 - }, "gemini/gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -17470,96 +15214,6 @@ "supports_web_search": true, "tpm": 250000 }, - "gemini/gemini-2.5-flash-preview-04-17": { - "cache_read_input_token_cost": 3.75e-08, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 1.5e-07, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_reasoning_token": 3.5e-06, - "output_cost_per_token": 6e-07, - "rpm": 10, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 250000 - }, - "gemini/gemini-2.5-flash-preview-05-20": { - "deprecation_date": "2025-11-18", - "cache_read_input_token_cost": 7.5e-08, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_reasoning_token": 2.5e-06, - "output_cost_per_token": 2.5e-06, - "rpm": 10, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 250000 - }, "gemini/gemini-2.5-flash-preview-tts": { "input_cost_per_token": 3e-07, "litellm_provider": "gemini", @@ -17984,177 +15638,6 @@ "cache_read_input_token_cost_priority": 9e-08, "supports_service_tier": true }, - "gemini/gemini-2.5-pro-exp-03-25": { - "cache_read_input_token_cost": 0.0, - "input_cost_per_token": 0.0, - "input_cost_per_token_above_200k_tokens": 0.0, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 0.0, - "output_cost_per_token_above_200k_tokens": 0.0, - "rpm": 5, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_input": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_video_input": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 250000 - }, - "gemini/gemini-2.5-pro-preview-03-25": { - "deprecation_date": "2025-12-02", - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_audio_token": 7e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "rpm": 10000, - "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-pro-preview", - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 10000000 - }, - "gemini/gemini-2.5-pro-preview-05-06": { - "deprecation_date": "2025-12-02", - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_audio_token": 7e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "rpm": 10000, - "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-pro-preview", - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 10000000 - }, - "gemini/gemini-2.5-pro-preview-06-05": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_audio_token": 7e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "rpm": 10000, - "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-pro-preview", - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 10000000 - }, "gemini/gemini-2.5-pro-preview-tts": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, @@ -18278,41 +15761,6 @@ "tpm": 250000, "rpm": 10 }, - "gemini/gemini-pro": { - "input_cost_per_token": 3.5e-07, - "input_cost_per_token_above_128k_tokens": 7e-07, - "litellm_provider": "gemini", - "max_input_tokens": 32760, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.05e-06, - "output_cost_per_token_above_128k_tokens": 2.1e-06, - "rpd": 30000, - "rpm": 360, - "source": "https://ai.google.dev/gemini-api/docs/models/gemini", - "supports_function_calling": true, - "supports_tool_choice": true, - "tpm": 120000 - }, - "gemini/gemini-pro-vision": { - "input_cost_per_token": 3.5e-07, - "input_cost_per_token_above_128k_tokens": 7e-07, - "litellm_provider": "gemini", - "max_input_tokens": 30720, - "max_output_tokens": 2048, - "max_tokens": 2048, - "mode": "chat", - "output_cost_per_token": 1.05e-06, - "output_cost_per_token_above_128k_tokens": 2.1e-06, - "rpd": 30000, - "rpm": 360, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", - "supports_function_calling": true, - "supports_tool_choice": true, - "supports_vision": true, - "tpm": 120000 - }, "gemini/gemma-3-27b-it": { "input_cost_per_audio_per_second": 0, "input_cost_per_audio_per_second_above_128k_tokens": 0, @@ -18420,36 +15868,6 @@ "video" ] }, - "gemini/veo-3.0-fast-generate-preview": { - "deprecation_date": "2025-11-12", - "litellm_provider": "gemini", - "max_input_tokens": 1024, - "max_tokens": 1024, - "mode": "video_generation", - "output_cost_per_second": 0.4, - "source": "https://ai.google.dev/gemini-api/docs/video", - "supported_modalities": [ - "text" - ], - "supported_output_modalities": [ - "video" - ] - }, - "gemini/veo-3.0-generate-preview": { - "deprecation_date": "2025-11-12", - "litellm_provider": "gemini", - "max_input_tokens": 1024, - "max_tokens": 1024, - "mode": "video_generation", - "output_cost_per_second": 0.75, - "source": "https://ai.google.dev/gemini-api/docs/video", - "supported_modalities": [ - "text" - ], - "supported_output_modalities": [ - "video" - ] - }, "gemini/veo-3.1-fast-generate-preview": { "litellm_provider": "gemini", "max_input_tokens": 1024, @@ -19373,31 +16791,6 @@ "supports_system_messages": true, "supports_tool_choice": true }, - "gpt-3.5-turbo-0301": { - "input_cost_per_token": 1.5e-06, - "litellm_provider": "openai", - "max_input_tokens": 4097, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 2e-06, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, - "gpt-3.5-turbo-0613": { - "input_cost_per_token": 1.5e-06, - "litellm_provider": "openai", - "max_input_tokens": 4097, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 2e-06, - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, "gpt-3.5-turbo-1106": { "deprecation_date": "2026-09-28", "input_cost_per_token": 1e-06, @@ -19425,18 +16818,6 @@ "supports_system_messages": true, "supports_tool_choice": true }, - "gpt-3.5-turbo-16k-0613": { - "input_cost_per_token": 3e-06, - "litellm_provider": "openai", - "max_input_tokens": 16385, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 4e-06, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, "gpt-3.5-turbo-instruct": { "input_cost_per_token": 1.5e-06, "litellm_provider": "text-completion-openai", @@ -19483,18 +16864,6 @@ "supports_system_messages": true, "supports_tool_choice": true }, - "gpt-4-0314": { - "input_cost_per_token": 3e-05, - "litellm_provider": "openai", - "max_input_tokens": 8192, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 6e-05, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, "gpt-4-0613": { "deprecation_date": "2025-06-06", "input_cost_per_token": 3e-05, @@ -19524,57 +16893,6 @@ "supports_system_messages": true, "supports_tool_choice": true }, - "gpt-4-1106-vision-preview": { - "deprecation_date": "2024-12-06", - "input_cost_per_token": 1e-05, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 3e-05, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gpt-4-32k": { - "input_cost_per_token": 6e-05, - "litellm_provider": "openai", - "max_input_tokens": 32768, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 0.00012, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, - "gpt-4-32k-0314": { - "input_cost_per_token": 6e-05, - "litellm_provider": "openai", - "max_input_tokens": 32768, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 0.00012, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, - "gpt-4-32k-0613": { - "input_cost_per_token": 6e-05, - "litellm_provider": "openai", - "max_input_tokens": 32768, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 0.00012, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, "gpt-4-turbo": { "input_cost_per_token": 1e-05, "litellm_provider": "openai", @@ -19622,21 +16940,6 @@ "supports_system_messages": true, "supports_tool_choice": true }, - "gpt-4-vision-preview": { - "deprecation_date": "2024-12-06", - "input_cost_per_token": 1e-05, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 3e-05, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, "gpt-4.1": { "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_priority": 8.75e-07, @@ -19854,47 +17157,6 @@ "supports_service_tier": true, "supports_vision": true }, - "gpt-4.5-preview": { - "cache_read_input_token_cost": 3.75e-05, - "input_cost_per_token": 7.5e-05, - "input_cost_per_token_batches": 3.75e-05, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 16384, - "max_tokens": 16384, - "mode": "chat", - "output_cost_per_token": 0.00015, - "output_cost_per_token_batches": 7.5e-05, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "gpt-4.5-preview-2025-02-27": { - "cache_read_input_token_cost": 3.75e-05, - "deprecation_date": "2025-07-14", - "input_cost_per_token": 7.5e-05, - "input_cost_per_token_batches": 3.75e-05, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 16384, - "max_tokens": 16384, - "mode": "chat", - "output_cost_per_token": 0.00015, - "output_cost_per_token_batches": 7.5e-05, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true - }, "gpt-4o": { "cache_read_input_token_cost": 1.25e-06, "cache_read_input_token_cost_priority": 2.125e-06, @@ -19998,23 +17260,6 @@ "supports_system_messages": true, "supports_tool_choice": true }, - "gpt-4o-audio-preview-2024-10-01": { - "input_cost_per_audio_token": 4e-05, - "input_cost_per_token": 2.5e-06, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 16384, - "max_tokens": 16384, - "mode": "chat", - "output_cost_per_audio_token": 8e-05, - "output_cost_per_token": 1e-05, - "supports_audio_input": true, - "supports_audio_output": true, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, "gpt-4o-audio-preview-2024-12-17": { "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 2.5e-06, @@ -20478,25 +17723,6 @@ "supports_system_messages": true, "supports_tool_choice": true }, - "gpt-4o-realtime-preview-2024-10-01": { - "cache_creation_input_audio_token_cost": 2e-05, - "cache_read_input_token_cost": 2.5e-06, - "input_cost_per_audio_token": 0.0001, - "input_cost_per_token": 5e-06, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_audio_token": 0.0002, - "output_cost_per_token": 2e-05, - "supports_audio_input": true, - "supports_audio_output": true, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, "gpt-4o-realtime-preview-2024-12-17": { "cache_read_input_token_cost": 2.5e-06, "input_cost_per_audio_token": 4e-05, @@ -25700,62 +22926,6 @@ "supports_tool_choice": true, "supports_vision": true }, - "o1-mini": { - "cache_read_input_token_cost": 5.5e-07, - "input_cost_per_token": 1.1e-06, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 65536, - "max_tokens": 65536, - "mode": "chat", - "output_cost_per_token": 4.4e-06, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_vision": true - }, - "o1-mini-2024-09-12": { - "deprecation_date": "2025-10-27", - "cache_read_input_token_cost": 1.5e-06, - "input_cost_per_token": 3e-06, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 65536, - "max_tokens": 65536, - "mode": "chat", - "output_cost_per_token": 1.2e-05, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_vision": true - }, - "o1-preview": { - "cache_read_input_token_cost": 7.5e-06, - "input_cost_per_token": 1.5e-05, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 32768, - "max_tokens": 32768, - "mode": "chat", - "output_cost_per_token": 6e-05, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_vision": true - }, - "o1-preview-2024-09-12": { - "cache_read_input_token_cost": 7.5e-06, - "input_cost_per_token": 1.5e-05, - "litellm_provider": "openai", - "max_input_tokens": 128000, - "max_output_tokens": 32768, - "max_tokens": 32768, - "mode": "chat", - "output_cost_per_token": 6e-05, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_vision": true - }, "o1-pro": { "input_cost_per_token": 0.00015, "input_cost_per_token_batches": 7.5e-05, @@ -26622,15 +23792,6 @@ "mode": "moderation", "output_cost_per_token": 0.0 }, - "omni-moderation-latest-intents": { - "input_cost_per_token": 0.0, - "litellm_provider": "openai", - "max_input_tokens": 32768, - "max_output_tokens": 0, - "max_tokens": 0, - "mode": "moderation", - "output_cost_per_token": 0.0 - }, "openai.gpt-oss-120b-1:0": { "input_cost_per_token": 1.5e-07, "litellm_provider": "bedrock_converse", @@ -28380,56 +25541,6 @@ "mode": "chat", "output_cost_per_token": 2e-07 }, - "perplexity/llama-3.1-sonar-huge-128k-online": { - "deprecation_date": "2025-02-22", - "input_cost_per_token": 5e-06, - "litellm_provider": "perplexity", - "max_input_tokens": 127072, - "max_output_tokens": 127072, - "max_tokens": 127072, - "mode": "chat", - "output_cost_per_token": 5e-06 - }, - "perplexity/llama-3.1-sonar-large-128k-chat": { - "deprecation_date": "2025-02-22", - "input_cost_per_token": 1e-06, - "litellm_provider": "perplexity", - "max_input_tokens": 131072, - "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", - "output_cost_per_token": 1e-06 - }, - "perplexity/llama-3.1-sonar-large-128k-online": { - "deprecation_date": "2025-02-22", - "input_cost_per_token": 1e-06, - "litellm_provider": "perplexity", - "max_input_tokens": 127072, - "max_output_tokens": 127072, - "max_tokens": 127072, - "mode": "chat", - "output_cost_per_token": 1e-06 - }, - "perplexity/llama-3.1-sonar-small-128k-chat": { - "deprecation_date": "2025-02-22", - "input_cost_per_token": 2e-07, - "litellm_provider": "perplexity", - "max_input_tokens": 131072, - "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", - "output_cost_per_token": 2e-07 - }, - "perplexity/llama-3.1-sonar-small-128k-online": { - "deprecation_date": "2025-02-22", - "input_cost_per_token": 2e-07, - "litellm_provider": "perplexity", - "max_input_tokens": 127072, - "max_output_tokens": 127072, - "max_tokens": 127072, - "mode": "chat", - "output_cost_per_token": 2e-07 - }, "perplexity/mistral-7b-instruct": { "input_cost_per_token": 7e-08, "litellm_provider": "perplexity", @@ -30212,60 +27323,6 @@ "litellm_provider": "tavily", "mode": "search" }, - "text-bison": { - "input_cost_per_character": 2.5e-07, - "litellm_provider": "vertex_ai-text-models", - "max_input_tokens": 8192, - "max_output_tokens": 2048, - "max_tokens": 2048, - "mode": "completion", - "output_cost_per_character": 5e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "text-bison32k": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-text-models", - "max_input_tokens": 8192, - "max_output_tokens": 1024, - "max_tokens": 1024, - "mode": "completion", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "text-bison32k@002": { - "input_cost_per_character": 2.5e-07, - "input_cost_per_token": 1.25e-07, - "litellm_provider": "vertex_ai-text-models", - "max_input_tokens": 8192, - "max_output_tokens": 1024, - "max_tokens": 1024, - "mode": "completion", - "output_cost_per_character": 5e-07, - "output_cost_per_token": 1.25e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "text-bison@001": { - "input_cost_per_character": 2.5e-07, - "litellm_provider": "vertex_ai-text-models", - "max_input_tokens": 8192, - "max_output_tokens": 1024, - "max_tokens": 1024, - "mode": "completion", - "output_cost_per_character": 5e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "text-bison@002": { - "input_cost_per_character": 2.5e-07, - "litellm_provider": "vertex_ai-text-models", - "max_input_tokens": 8192, - "max_output_tokens": 1024, - "max_tokens": 1024, - "mode": "completion", - "output_cost_per_character": 5e-07, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, "text-completion-codestral/codestral-2405": { "input_cost_per_token": 0.0, "litellm_provider": "text-completion-codestral", @@ -30410,16 +27467,6 @@ "output_vector_size": 768, "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models" }, - "text-multilingual-embedding-preview-0409": { - "input_cost_per_token": 6.25e-09, - "litellm_provider": "vertex_ai-embedding-models", - "max_input_tokens": 3072, - "max_tokens": 3072, - "mode": "embedding", - "output_cost_per_token": 0, - "output_vector_size": 768, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, "text-unicorn": { "input_cost_per_token": 1e-05, "litellm_provider": "vertex_ai-text-models", @@ -30440,61 +27487,6 @@ "output_cost_per_token": 2.8e-05, "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" }, - "textembedding-gecko": { - "input_cost_per_character": 2.5e-08, - "input_cost_per_token": 1e-07, - "litellm_provider": "vertex_ai-embedding-models", - "max_input_tokens": 3072, - "max_tokens": 3072, - "mode": "embedding", - "output_cost_per_token": 0, - "output_vector_size": 768, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "textembedding-gecko-multilingual": { - "input_cost_per_character": 2.5e-08, - "input_cost_per_token": 1e-07, - "litellm_provider": "vertex_ai-embedding-models", - "max_input_tokens": 3072, - "max_tokens": 3072, - "mode": "embedding", - "output_cost_per_token": 0, - "output_vector_size": 768, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "textembedding-gecko-multilingual@001": { - "input_cost_per_character": 2.5e-08, - "input_cost_per_token": 1e-07, - "litellm_provider": "vertex_ai-embedding-models", - "max_input_tokens": 3072, - "max_tokens": 3072, - "mode": "embedding", - "output_cost_per_token": 0, - "output_vector_size": 768, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "textembedding-gecko@001": { - "input_cost_per_character": 2.5e-08, - "input_cost_per_token": 1e-07, - "litellm_provider": "vertex_ai-embedding-models", - "max_input_tokens": 3072, - "max_tokens": 3072, - "mode": "embedding", - "output_cost_per_token": 0, - "output_vector_size": 768, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, - "textembedding-gecko@003": { - "input_cost_per_character": 2.5e-08, - "input_cost_per_token": 1e-07, - "litellm_provider": "vertex_ai-embedding-models", - "max_input_tokens": 3072, - "max_tokens": 3072, - "mode": "embedding", - "output_cost_per_token": 0, - "output_vector_size": 768, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" - }, "together-ai-21.1b-41b": { "input_cost_per_token": 8e-07, "litellm_provider": "together_ai", @@ -32896,36 +29888,6 @@ "supports_tool_choice": true, "supports_vision": true }, - "vertex_ai/claude-3-5-sonnet-v2": { - "input_cost_per_token": 3e-06, - "litellm_provider": "vertex_ai-anthropic_models", - "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.5e-05, - "supports_assistant_prefill": true, - "supports_computer_use": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_tool_choice": true, - "supports_vision": true - }, - "vertex_ai/claude-3-5-sonnet-v2@20241022": { - "input_cost_per_token": 3e-06, - "litellm_provider": "vertex_ai-anthropic_models", - "max_input_tokens": 200000, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.5e-05, - "supports_assistant_prefill": true, - "supports_computer_use": true, - "supports_function_calling": true, - "supports_pdf_input": true, - "supports_tool_choice": true, - "supports_vision": true - }, "vertex_ai/claude-3-5-sonnet@20240620": { "input_cost_per_token": 3e-06, "litellm_provider": "vertex_ai-anthropic_models", @@ -32943,7 +29905,7 @@ "vertex_ai/claude-3-7-sonnet@20250219": { "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, - "deprecation_date": "2025-06-01", + "deprecation_date": "2026-05-11", "input_cost_per_token": 3e-06, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, @@ -33959,6 +30921,9 @@ "mode": "chat", "output_cost_per_token": 2.2e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#partner-models", + "supported_regions": [ + "global" + ], "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true @@ -33973,6 +30938,7 @@ "mode": "chat", "output_cost_per_token": 3.2e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#glm-models", + "supported_regions": ["global"], "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -34224,36 +31190,6 @@ "video" ] }, - "vertex_ai/veo-3.0-fast-generate-preview": { - "deprecation_date": "2025-11-12", - "litellm_provider": "vertex_ai-video-models", - "max_input_tokens": 1024, - "max_tokens": 1024, - "mode": "video_generation", - "output_cost_per_second": 0.15, - "source": "https://ai.google.dev/gemini-api/docs/video", - "supported_modalities": [ - "text" - ], - "supported_output_modalities": [ - "video" - ] - }, - "vertex_ai/veo-3.0-generate-preview": { - "deprecation_date": "2025-11-12", - "litellm_provider": "vertex_ai-video-models", - "max_input_tokens": 1024, - "max_tokens": 1024, - "mode": "video_generation", - "output_cost_per_second": 0.4, - "source": "https://ai.google.dev/gemini-api/docs/video", - "supported_modalities": [ - "text" - ], - "supported_output_modalities": [ - "video" - ] - }, "vertex_ai/veo-3.0-fast-generate-001": { "litellm_provider": "vertex_ai-video-models", "max_input_tokens": 1024, diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 64942636a9..2f3302bb57 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -458,6 +458,24 @@ "interactions": true } }, + "charity_engine": { + "display_name": "Charity Engine (`charity_engine`)", + "url": "https://docs.litellm.ai/docs/providers/charity_engine", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false, + "interactions": false + } + }, "chutes": { "display_name": "Chutes (`chutes`)", "endpoints": { diff --git a/ruff.toml b/ruff.toml index 76acb5dc93..55d008a7dd 100644 --- a/ruff.toml +++ b/ruff.toml @@ -17,3 +17,4 @@ exclude = ["litellm/types/*", "litellm/__init__.py", "litellm/proxy/example_conf "litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py" = ["PLR0915"] "litellm/proxy/guardrails/guardrail_hooks/guardrail_benchmarks/test_eval.py" = ["PLR0915"] "litellm/responses/streaming_iterator.py" = ["PLR0915"] +"litellm/files/main.py" = ["PLR0915"] diff --git a/tests/llm_translation/test_skills_api.py b/tests/llm_translation/test_skills_api.py index 7565ba7440..76eb274293 100644 --- a/tests/llm_translation/test_skills_api.py +++ b/tests/llm_translation/test_skills_api.py @@ -44,19 +44,25 @@ def create_skill_zip(skill_name: str, unique_suffix: Optional[str] = None): skill_dir = test_dir / skill_name # Create a zip file containing the skill directory + # When unique_suffix is set, folder name must match skill name in SKILL.md (Anthropic requirement) + zip_folder_name = f"{skill_name}-{unique_suffix}" if unique_suffix else skill_name zip_path = test_dir / f"{skill_name}.zip" with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf: - zf.write(skill_dir, arcname=skill_name) - if unique_suffix is not None: - # Rewrite SKILL.md with a unique name to avoid API conflicts + # Rewrite SKILL.md with a unique name and use matching folder name skill_md = (skill_dir / "SKILL.md").read_text() skill_md = skill_md.replace( f"name: {skill_name}", - f"name: {skill_name}-{unique_suffix}", + f"name: {zip_folder_name}", ) - zf.writestr(f"{skill_name}/SKILL.md", skill_md) + zf.writestr(f"{zip_folder_name}/SKILL.md", skill_md) + # Add any other files in the skill dir (e.g. subdirs) under the new folder name + for f in skill_dir.rglob("*"): + if f.is_file() and f.name != "SKILL.md": + rel = f.relative_to(skill_dir) + zf.write(f, arcname=f"{zip_folder_name}/{rel}") else: + zf.write(skill_dir, arcname=skill_name) zf.write(skill_dir / "SKILL.md", arcname=f"{skill_name}/SKILL.md") try: @@ -268,17 +274,11 @@ class BaseSkillsAPITest(ABC): print(f"Deleted skill response: {response}") -class TestAnthropicSkillsAPI(BaseSkillsAPITest): - """ - Test Anthropic Skills API implementation. - """ - - def get_custom_llm_provider(self) -> str: - return "anthropic" - - def get_api_key(self) -> Optional[str]: - return os.environ.get("ANTHROPIC_API_KEY") - - def get_api_base(self) -> Optional[str]: - return os.environ.get("ANTHROPIC_API_BASE") +# Live integration tests for the Anthropic Skills API are not run in CI because +# the Skills API requires beta access (anthropic-beta: skills-2025-10-02) that +# is not available on the standard API key used in CI. +# +# Transformation logic (URL construction, headers, request/response parsing) is +# covered by unit tests in: +# tests/test_litellm/test_anthropic_skills_transformation.py diff --git a/tests/local_testing/test_custom_callback_input.py b/tests/local_testing/test_custom_callback_input.py index a28151d47a..c498de15d7 100644 --- a/tests/local_testing/test_custom_callback_input.py +++ b/tests/local_testing/test_custom_callback_input.py @@ -1310,9 +1310,11 @@ def test_logging_async_cache_hit_sync_call(turn_off_message_logging): "redacted-by-litellm" == standard_logging_object["messages"][0]["content"] ) - assert {"text": "redacted-by-litellm"} == standard_logging_object[ - "response" - ] + # response is a full ModelResponse dict (choices format) since d84e5e381acf + assert ( + standard_logging_object["response"]["choices"][0]["message"]["content"] + == "redacted-by-litellm" + ) def test_logging_standard_payload_failure_call(): diff --git a/tests/logging_callback_tests/test_logging_redaction_e2e_test.py b/tests/logging_callback_tests/test_logging_redaction_e2e_test.py index 0536ec7205..0391a5a895 100644 --- a/tests/logging_callback_tests/test_logging_redaction_e2e_test.py +++ b/tests/logging_callback_tests/test_logging_redaction_e2e_test.py @@ -45,7 +45,8 @@ async def test_global_redaction_on(): await asyncio.sleep(1) standard_logging_payload = test_custom_logger.logged_standard_logging_payload assert standard_logging_payload is not None - assert standard_logging_payload["response"] == {"text": "redacted-by-litellm"} + response = standard_logging_payload["response"] + assert response["choices"][0]["message"]["content"] == "redacted-by-litellm" assert standard_logging_payload["messages"][0]["content"] == "redacted-by-litellm" print( "logged standard logging payload", @@ -75,7 +76,8 @@ async def test_global_redaction_with_dynamic_params(turn_off_message_logging): ) if turn_off_message_logging is True: - assert standard_logging_payload["response"] == {"text": "redacted-by-litellm"} + response = standard_logging_payload["response"] + assert response["choices"][0]["message"]["content"] == "redacted-by-litellm" assert ( standard_logging_payload["messages"][0]["content"] == "redacted-by-litellm" ) @@ -108,7 +110,8 @@ async def test_global_redaction_off_with_dynamic_params(turn_off_message_logging json.dumps(standard_logging_payload, indent=2), ) if turn_off_message_logging is True: - assert standard_logging_payload["response"] == {"text": "redacted-by-litellm"} + response = standard_logging_payload["response"] + assert response["choices"][0]["message"]["content"] == "redacted-by-litellm" assert ( standard_logging_payload["messages"][0]["content"] == "redacted-by-litellm" ) @@ -390,7 +393,8 @@ async def test_redaction_with_streaming_response(): assert standard_logging_payload is not None # Verify that redaction worked without pickle errors - assert standard_logging_payload["response"] == {"text": "redacted-by-litellm"} + response = standard_logging_payload["response"] + assert response["choices"][0]["message"]["content"] == "redacted-by-litellm" assert standard_logging_payload["messages"][0]["content"] == "redacted-by-litellm" print( "logged standard logging payload for streaming with coroutine handling", @@ -477,5 +481,6 @@ async def test_redaction_with_metadata_completion_api(): # Verify the helper function works correctly - with get_metadata_variable_name_from_kwargs, # the system checks the appropriate field for headers - assert standard_logging_payload["response"] == {"text": "redacted-by-litellm"} + response = standard_logging_payload["response"] + assert response["choices"][0]["message"]["content"] == "redacted-by-litellm" assert standard_logging_payload["messages"][0]["content"] == "redacted-by-litellm" diff --git a/tests/test_litellm/caching/test_dual_cache.py b/tests/test_litellm/caching/test_dual_cache.py index 9974c23e4b..606f25ddf4 100644 --- a/tests/test_litellm/caching/test_dual_cache.py +++ b/tests/test_litellm/caching/test_dual_cache.py @@ -1,9 +1,11 @@ import asyncio +import time from unittest.mock import AsyncMock, MagicMock, patch import pytest from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache @@ -56,3 +58,104 @@ async def test_dual_cache_async_batch_get_cache_rolls_back_redis_reservation_on_ assert mock_async_batch_get_cache.call_count == 2 assert "shared_a" not in dual_cache.last_redis_batch_access_time assert "shared_b" not in dual_cache.last_redis_batch_access_time + + +@pytest.mark.asyncio +async def test_dual_cache_async_set_cache_injects_default_in_memory_ttl(): + """ + Test that async_set_cache injects default_in_memory_ttl into kwargs + when no explicit ttl is provided, matching the sync set_cache behavior. + + Regression test for: async_set_cache was missing the TTL injection that + sync set_cache has, causing InMemoryCache to use its own default_ttl (600s) + instead of DualCache's default_in_memory_ttl. + """ + in_memory_cache = InMemoryCache(default_ttl=600) + dual_cache = DualCache( + in_memory_cache=in_memory_cache, + default_in_memory_ttl=60, + ) + + before = time.time() + await dual_cache.async_set_cache(key="test_key", value="test_value") + after = time.time() + + # The TTL stored should reflect default_in_memory_ttl (60s), not + # InMemoryCache's default_ttl (600s) + expiry = in_memory_cache.ttl_dict["test_key"] + assert expiry >= before + 60 + assert expiry <= after + 60 + + +@pytest.mark.asyncio +async def test_dual_cache_async_set_cache_respects_explicit_ttl(): + """ + Test that async_set_cache does NOT override an explicitly provided ttl. + """ + in_memory_cache = InMemoryCache(default_ttl=600) + dual_cache = DualCache( + in_memory_cache=in_memory_cache, + default_in_memory_ttl=60, + ) + + before = time.time() + await dual_cache.async_set_cache(key="test_key", value="test_value", ttl=30) + after = time.time() + + # The explicit ttl=30 should be used, not default_in_memory_ttl (60) + expiry = in_memory_cache.ttl_dict["test_key"] + assert expiry >= before + 30 + assert expiry <= after + 30 + + +@pytest.mark.asyncio +async def test_dual_cache_async_set_cache_pipeline_injects_default_in_memory_ttl(): + """ + Test that async_set_cache_pipeline injects default_in_memory_ttl into kwargs + when no explicit ttl is provided. + """ + in_memory_cache = InMemoryCache(default_ttl=600) + dual_cache = DualCache( + in_memory_cache=in_memory_cache, + default_in_memory_ttl=60, + ) + + cache_list = [("key_a", "value_a"), ("key_b", "value_b")] + + before = time.time() + await dual_cache.async_set_cache_pipeline(cache_list=cache_list) + after = time.time() + + for key in ["key_a", "key_b"]: + expiry = in_memory_cache.ttl_dict[key] + assert expiry >= before + 60 + assert expiry <= after + 60 + + +@pytest.mark.asyncio +async def test_dual_cache_sync_and_async_set_cache_use_same_ttl(): + """ + Test that sync set_cache and async async_set_cache produce the same TTL + when no explicit ttl is provided, ensuring parity between the two paths. + """ + in_memory_sync = InMemoryCache(default_ttl=600) + dual_cache_sync = DualCache( + in_memory_cache=in_memory_sync, + default_in_memory_ttl=60, + ) + + in_memory_async = InMemoryCache(default_ttl=600) + dual_cache_async = DualCache( + in_memory_cache=in_memory_async, + default_in_memory_ttl=60, + ) + + dual_cache_sync.set_cache(key="test_key", value="test_value") + await dual_cache_async.async_set_cache(key="test_key", value="test_value") + + sync_expiry = in_memory_sync.ttl_dict["test_key"] + async_expiry = in_memory_async.ttl_dict["test_key"] + + # Both should use default_in_memory_ttl=60, so their expiry times + # should be within a small tolerance of each other + assert abs(sync_expiry - async_expiry) < 1.0 diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index ef3d7534d9..8c72b7725a 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -738,7 +738,58 @@ def test_response_completed_with_message_only_emits_stop_finish_reason(): ) -def test_function_call_done_does_not_emit_finish_reason(): + +def test_response_completed_preserves_usage_with_cached_tokens(): + """ + Test that response.completed correctly translates Responses API usage + (input_tokens_details) to chat completion usage (prompt_tokens_details). + + This is a regression test for an issue where streaming with models that + use the Responses API bridge (e.g. gpt-5.2-codex) would drop + prompt_tokens_details, causing cached_tokens to always be None. + """ + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + OpenAiResponsesToChatCompletionStreamIterator, + ) + + iterator = OpenAiResponsesToChatCompletionStreamIterator(streaming_response=None, sync_stream=True) + + chunk = { + "type": "response.completed", + "response": { + "id": "resp_789", + "status": "completed", + "output": [ + { + "type": "message", + "id": "msg_abc", + "role": "assistant", + "content": [{"type": "output_text", "text": "Six"}], + "status": "completed", + } + ], + "usage": { + "input_tokens": 1226, + "output_tokens": 5, + "total_tokens": 1231, + "input_tokens_details": {"cached_tokens": 1024}, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + }, + } + + result = iterator.chunk_parser(chunk) + + assert result.usage is not None, "usage should be set on response.completed chunk" + assert result.usage.prompt_tokens == 1226, "prompt_tokens should map from input_tokens" + assert result.usage.completion_tokens == 5, "completion_tokens should map from output_tokens" + assert result.usage.prompt_tokens_details is not None, "prompt_tokens_details should be set" + assert result.usage.prompt_tokens_details.cached_tokens == 1024, ( + "cached_tokens should be preserved from input_tokens_details" + ) + + +def test_function_call_done_emits_is_finished(): """ Test that OUTPUT_ITEM_DONE for a function_call does NOT emit finish_reason. The response.completed event handles the terminal finish_reason correctly. @@ -1327,6 +1378,138 @@ def test_transform_response_preserves_annotations(): print("✓ Annotations from Responses API are correctly preserved in Chat Completions format") +def test_apply_patch_tool_call_converted_to_chat_completion_tool_call(): + """ + Test that ResponseApplyPatchToolCall items from the Responses API are + correctly converted to ChatCompletions-style tool calls by the bridge. + + This is a regression test for a bug where litellm.completion() with a + responses/ model prefix crashed when the model returned an + apply_patch_call, because _convert_response_output_to_choices did not + handle ResponseApplyPatchToolCall items. The model DID use the tool, + but the bridge silently dropped it (or raised an error), while the + native litellm.responses() path worked correctly. + """ + import json + from unittest.mock import Mock + + from openai.types.responses.response_apply_patch_tool_call import ( + OperationCreateFile, + ) + from openai.types.responses.response_output_item import ( + ResponseApplyPatchToolCall, + ) + + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + from litellm.types.llms.openai import ( + InputTokensDetails, + OutputTokensDetails, + ResponseAPIUsage, + ResponsesAPIResponse, + ) + from litellm.types.utils import ModelResponse, Usage + + handler = LiteLLMResponsesTransformationHandler() + + # Build an apply_patch_call item like the model would return + operation = OperationCreateFile( + diff="--- /dev/null\n+++ b/hello.py\n@@ -0,0 +1 @@\n+print('hello world')\n", + path="hello.py", + type="create_file", + ) + apply_patch_item = ResponseApplyPatchToolCall( + id="apc_001", + call_id="call_patch_hello", + operation=operation, + status="completed", + type="apply_patch_call", + ) + + # Minimal usage + usage = ResponseAPIUsage( + input_tokens=30, + input_tokens_details=InputTokensDetails(cached_tokens=0), + output_tokens=40, + output_tokens_details=OutputTokensDetails(reasoning_tokens=0), + total_tokens=70, + ) + + raw_response = ResponsesAPIResponse( + id="resp_apply_patch_test", + created_at=1234567890, + error=None, + incomplete_details=None, + instructions=None, + metadata={}, + model="gpt-5.2-codex", + object="response", + output=[apply_patch_item], + parallel_tool_calls=True, + temperature=1.0, + tool_choice="auto", + tools=[], + top_p=1.0, + max_output_tokens=None, + previous_response_id=None, + reasoning=None, + status="completed", + text=None, + truncation="disabled", + usage=usage, + user=None, + store=True, + background=False, + ) + + model_response = ModelResponse( + id="chatcmpl-apply-patch", + created=1234567890, + model=None, + object="chat.completion", + choices=[], + usage=Usage(completion_tokens=0, prompt_tokens=0, total_tokens=0), + ) + + logging_obj = Mock() + + result = handler.transform_response( + model="gpt-5.2-codex", + raw_response=raw_response, + model_response=model_response, + logging_obj=logging_obj, + request_data={"model": "gpt-5.2-codex"}, + messages=[ + {"role": "system", "content": "You are a coding assistant."}, + {"role": "user", "content": "Create hello.py"}, + ], + optional_params={}, + litellm_params={}, + encoding=Mock(), + ) + + # Should have exactly one choice with finish_reason="tool_calls" + assert len(result.choices) == 1, f"Expected 1 choice, got {len(result.choices)}" + + choice = result.choices[0] + assert choice.finish_reason == "tool_calls" + + # The choice should contain one tool call for apply_patch + tool_calls = choice.message.tool_calls + assert tool_calls is not None, "tool_calls should not be None" + assert len(tool_calls) == 1, f"Expected 1 tool_call, got {len(tool_calls)}" + + tc = tool_calls[0] + assert tc["id"] == "call_patch_hello" + assert tc["type"] == "function" + assert tc["function"]["name"] == "apply_patch" + + # The operation should be serialised as JSON in arguments + args = json.loads(tc["function"]["arguments"]) + assert args["type"] == "create_file" + assert args["path"] == "hello.py" + assert "print('hello world')" in args["diff"] def test_multi_tool_call_stream_no_premature_finish(): """ Regression test for multi-tool-call streaming bug. @@ -1778,3 +1961,35 @@ def test_parallel_tool_calls_comprehensive_streaming_integration(): ) print("✓ Parallel tool calls with split argument deltas stream correctly end-to-end") + + +def test_map_optional_params_preserves_reasoning_summary(): + """Test that reasoning_effort dict with summary field is preserved. + + Regression test for: User reported that summary field was being dropped + when routing to Responses API. The dict format should be fully preserved. + """ + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams + + handler = LiteLLMResponsesTransformationHandler() + + optional_params = { + "stream": False, + "tools": [{"type": "function", "function": {"name": "test_tool"}}], + "tool_choice": "auto", + "reasoning_effort": {"effort": "high", "summary": "detailed"}, + } + + responses_api_request = ResponsesAPIOptionalRequestParams() + handler._map_optional_params_to_responses_api_request( + optional_params, responses_api_request + ) + + # Verify reasoning_effort dict with summary was fully preserved + assert "reasoning" in responses_api_request + assert responses_api_request["reasoning"] == {"effort": "high", "summary": "detailed"} + assert responses_api_request["reasoning"]["effort"] == "high" + assert responses_api_request["reasoning"]["summary"] == "detailed" diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 00dc0c4c4a..e907e92e66 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -357,7 +357,7 @@ def test_generic_cost_per_token_anthropic_prompt_caching(): def test_generic_cost_per_token_anthropic_prompt_caching_with_cache_creation(): - model = "claude-3-5-haiku-20241022" + model = "claude-haiku-4-5-20251001" usage = Usage( completion_tokens=90, prompt_tokens=28436, @@ -382,7 +382,7 @@ def test_generic_cost_per_token_anthropic_prompt_caching_with_cache_creation(): ) print(f"prompt_cost: {prompt_cost}") - assert round(prompt_cost, 3) == 0.023 + assert round(prompt_cost, 3) == 0.029 def test_string_cost_values(): diff --git a/tests/test_litellm/llms/anthropic/files/__init__.py b/tests/test_litellm/llms/anthropic/files/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/test_litellm/llms/anthropic/files/test_anthropic_files_transformation.py b/tests/test_litellm/llms/anthropic/files/test_anthropic_files_transformation.py new file mode 100644 index 0000000000..e9509be9e1 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/files/test_anthropic_files_transformation.py @@ -0,0 +1,379 @@ +""" +Test Anthropic Files API transformation functionality. + +Tests the AnthropicFilesConfig class which transforms between +OpenAI-compatible file operations and Anthropic's Files API format. +""" + +import io +import time + +import httpx +import pytest +from unittest.mock import Mock, patch + +from litellm.llms.anthropic.files.transformation import ( + AnthropicFilesConfig, + ANTHROPIC_FILES_API_BASE, + ANTHROPIC_FILES_BETA_HEADER, +) +from litellm.types.llms.openai import OpenAIFileObject +from litellm.types.utils import LlmProviders + + +class TestAnthropicFilesConfig: + """Test AnthropicFilesConfig transformation methods.""" + + def setup_method(self): + self.config = AnthropicFilesConfig() + + def test_custom_llm_provider(self): + assert self.config.custom_llm_provider == LlmProviders.ANTHROPIC + + def test_get_complete_url_default(self): + url = self.config.get_complete_url( + api_base=None, + api_key="test-key", + model="", + optional_params={}, + litellm_params={}, + ) + assert url == f"{ANTHROPIC_FILES_API_BASE}/v1/files" + + def test_get_complete_url_custom_base(self): + url = self.config.get_complete_url( + api_base="https://custom.api.com", + api_key="test-key", + model="", + optional_params={}, + litellm_params={}, + ) + assert url == "https://custom.api.com/v1/files" + + def test_get_complete_url_strips_trailing_slash(self): + url = self.config.get_complete_url( + api_base="https://custom.api.com/", + api_key="test-key", + model="", + optional_params={}, + litellm_params={}, + ) + assert url == "https://custom.api.com/v1/files" + + def test_validate_environment_sets_headers(self): + headers = {} + result = self.config.validate_environment( + headers=headers, + model="", + messages=[], + optional_params={}, + litellm_params={}, + api_key="sk-ant-test-key", + ) + assert result["x-api-key"] == "sk-ant-test-key" + assert result["anthropic-version"] == "2023-06-01" + assert result["anthropic-beta"] == ANTHROPIC_FILES_BETA_HEADER + + @patch.dict("os.environ", {}, clear=True) + @patch( + "litellm.llms.anthropic.common_utils.AnthropicModelInfo.get_api_key", + return_value=None, + ) + def test_validate_environment_missing_api_key(self, mock_get_key): + with pytest.raises(ValueError, match="Anthropic API key is required"): + self.config.validate_environment( + headers={}, + model="", + messages=[], + optional_params={}, + litellm_params={}, + api_key=None, + ) + + def test_get_supported_openai_params(self): + params = self.config.get_supported_openai_params(model="") + assert "purpose" in params + + def test_transform_create_file_request(self): + file_content = b"test file content" + file_tuple = ("test.txt", file_content, "text/plain") + + result = self.config.transform_create_file_request( + model="", + create_file_data={ + "file": file_tuple, + "purpose": "messages", + }, + optional_params={}, + litellm_params={}, + ) + + assert "file" in result + assert "purpose" in result + # file should be a tuple (filename, content, content_type) + assert result["file"][0] == "test.txt" + assert result["file"][1] == file_content + assert result["file"][2] == "text/plain" + # purpose should be (None, value) for multipart form field + assert result["purpose"] == (None, "messages") + + def test_transform_create_file_request_missing_file(self): + with pytest.raises(ValueError, match="File data is required"): + self.config.transform_create_file_request( + model="", + create_file_data={"purpose": "messages"}, + optional_params={}, + litellm_params={}, + ) + + def test_transform_create_file_request_default_purpose(self): + file_tuple = ("test.txt", b"content", "text/plain") + result = self.config.transform_create_file_request( + model="", + create_file_data={"file": file_tuple}, + optional_params={}, + litellm_params={}, + ) + assert result["purpose"] == (None, "messages") + + def test_transform_create_file_response(self): + mock_response = Mock(spec=httpx.Response) + mock_response.json.return_value = { + "id": "file-abc123", + "type": "file", + "filename": "document.pdf", + "mime_type": "application/pdf", + "size_bytes": 12345, + "created_at": "2025-01-15T10:30:00Z", + } + + result = self.config.transform_create_file_response( + model=None, + raw_response=mock_response, + logging_obj=Mock(), + litellm_params={}, + ) + + assert isinstance(result, OpenAIFileObject) + assert result.id == "file-abc123" + assert result.filename == "document.pdf" + assert result.bytes == 12345 + assert result.object == "file" + assert result.purpose == "messages" + assert result.status == "uploaded" + + def test_transform_retrieve_file_request(self): + url, params = self.config.transform_retrieve_file_request( + file_id="file-abc123", + optional_params={}, + litellm_params={}, + ) + assert url == f"{ANTHROPIC_FILES_API_BASE}/v1/files/file-abc123" + assert params == {} + + def test_transform_retrieve_file_request_custom_base(self): + url, params = self.config.transform_retrieve_file_request( + file_id="file-abc123", + optional_params={}, + litellm_params={"api_base": "https://custom.api.com"}, + ) + assert url == "https://custom.api.com/v1/files/file-abc123" + assert params == {} + + def test_transform_retrieve_file_response(self): + mock_response = Mock(spec=httpx.Response) + mock_response.json.return_value = { + "id": "file-abc123", + "type": "file", + "filename": "document.pdf", + "mime_type": "application/pdf", + "size_bytes": 5000, + "created_at": "2025-06-01T12:00:00Z", + } + + result = self.config.transform_retrieve_file_response( + raw_response=mock_response, + logging_obj=Mock(), + litellm_params={}, + ) + + assert isinstance(result, OpenAIFileObject) + assert result.id == "file-abc123" + assert result.bytes == 5000 + + def test_transform_delete_file_request(self): + url, params = self.config.transform_delete_file_request( + file_id="file-abc123", + optional_params={}, + litellm_params={}, + ) + assert url == f"{ANTHROPIC_FILES_API_BASE}/v1/files/file-abc123" + assert params == {} + + def test_transform_delete_file_response(self): + mock_response = Mock(spec=httpx.Response) + mock_response.json.return_value = { + "id": "file-abc123", + "type": "file_deleted", + } + + result = self.config.transform_delete_file_response( + raw_response=mock_response, + logging_obj=Mock(), + litellm_params={}, + ) + + assert result.id == "file-abc123" + assert result.deleted is True + assert result.object == "file" + + def test_transform_list_files_request(self): + url, params = self.config.transform_list_files_request( + purpose=None, + optional_params={}, + litellm_params={}, + ) + assert url == f"{ANTHROPIC_FILES_API_BASE}/v1/files" + assert params == {} + + def test_transform_list_files_request_with_purpose(self): + url, params = self.config.transform_list_files_request( + purpose="messages", + optional_params={}, + litellm_params={}, + ) + assert url == f"{ANTHROPIC_FILES_API_BASE}/v1/files" + assert params == {"purpose": "messages"} + + def test_transform_list_files_response(self): + mock_response = Mock(spec=httpx.Response) + mock_response.json.return_value = { + "data": [ + { + "id": "file-1", + "filename": "a.txt", + "size_bytes": 100, + "created_at": "2025-01-01T00:00:00Z", + }, + { + "id": "file-2", + "filename": "b.txt", + "size_bytes": 200, + "created_at": "2025-01-02T00:00:00Z", + }, + ], + "has_more": False, + } + + result = self.config.transform_list_files_response( + raw_response=mock_response, + logging_obj=Mock(), + litellm_params={}, + ) + + assert len(result) == 2 + assert result[0].id == "file-1" + assert result[0].filename == "a.txt" + assert result[1].id == "file-2" + + def test_transform_list_files_response_empty(self): + mock_response = Mock(spec=httpx.Response) + mock_response.json.return_value = {"data": [], "has_more": False} + + result = self.config.transform_list_files_response( + raw_response=mock_response, + logging_obj=Mock(), + litellm_params={}, + ) + assert result == [] + + def test_transform_file_content_request(self): + url, params = self.config.transform_file_content_request( + file_content_request={"file_id": "file-abc123"}, + optional_params={}, + litellm_params={}, + ) + assert url == f"{ANTHROPIC_FILES_API_BASE}/v1/files/file-abc123/content" + assert params == {} + + def test_transform_file_content_response(self): + mock_response = Mock(spec=httpx.Response) + result = self.config.transform_file_content_response( + raw_response=mock_response, + logging_obj=Mock(), + litellm_params={}, + ) + assert result.response == mock_response + + def test_parse_anthropic_file_with_size_bytes(self): + """Test that size_bytes is correctly mapped to bytes field.""" + result = AnthropicFilesConfig._parse_anthropic_file( + { + "id": "file-test", + "filename": "test.pdf", + "size_bytes": 9999, + "created_at": "2025-03-01T00:00:00Z", + } + ) + assert result.bytes == 9999 + + def test_parse_anthropic_file_fallback_bytes_field(self): + """Test fallback to 'bytes' field when 'size_bytes' is missing.""" + result = AnthropicFilesConfig._parse_anthropic_file( + { + "id": "file-test", + "filename": "test.pdf", + "bytes": 7777, + "created_at": "2025-03-01T00:00:00Z", + } + ) + assert result.bytes == 7777 + + def test_parse_anthropic_file_invalid_timestamp(self): + """Test that invalid timestamps fall back to current time.""" + result = AnthropicFilesConfig._parse_anthropic_file( + { + "id": "file-test", + "filename": "test.pdf", + "size_bytes": 100, + "created_at": "not-a-date", + } + ) + # Should not raise, should use current time + assert isinstance(result.created_at, int) + assert result.created_at > 0 + + def test_parse_anthropic_file_missing_timestamp(self): + """Test that missing timestamps fall back to current time.""" + result = AnthropicFilesConfig._parse_anthropic_file( + { + "id": "file-test", + "filename": "test.pdf", + "size_bytes": 100, + } + ) + assert isinstance(result.created_at, int) + assert result.created_at > 0 + + def test_get_error_class(self): + error = self.config.get_error_class( + error_message="Not found", + status_code=404, + headers={}, + ) + assert error.status_code == 404 + assert error.message == "Not found" + + +class TestProviderConfigRegistration: + """Test that AnthropicFilesConfig is properly registered.""" + + def test_provider_config_returns_anthropic_files_config(self): + from litellm.utils import ProviderConfigManager + + config = ProviderConfigManager.get_provider_files_config( + model="", + provider=LlmProviders.ANTHROPIC, + ) + assert config is not None + assert isinstance(config, AnthropicFilesConfig) diff --git a/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py b/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py index 25f3d1364f..635359563b 100644 --- a/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py +++ b/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py @@ -192,6 +192,23 @@ def test_azure_gpt5_1_series_temperature_handling(config: AzureOpenAIGPT5Config) assert params["temperature"] == 0.6 +def test_azure_gpt5_4_drops_reasoning_effort_when_tools_present(config: AzureOpenAIGPT5Config): + """Azure Chat Completions: gpt-5.4+ drops reasoning_effort when tools are present. + + OpenAI routes tools+reasoning to Responses API; Azure does not, so we drop reasoning_effort. + """ + tools = [{"type": "function", "function": {"name": "test", "description": "test"}}] + params = config.map_openai_params( + non_default_params={"reasoning_effort": "high", "tools": tools}, + optional_params={}, + model="gpt5_series/gpt-5.4", + drop_params=False, + api_version="2024-05-01-preview", + ) + assert "reasoning_effort" not in params + assert params["tools"] == tools + + def test_azure_gpt5_reasoning_effort_none_error(config: AzureOpenAIGPT5Config): """Test that Azure GPT-5 (non-5.1) raises error for reasoning_effort='none' when drop_params=False.""" with pytest.raises(litellm.utils.UnsupportedParamsError): diff --git a/tests/test_litellm/llms/base_llm/test_base_model_iterator.py b/tests/test_litellm/llms/base_llm/test_base_model_iterator.py new file mode 100644 index 0000000000..96cd299b2b --- /dev/null +++ b/tests/test_litellm/llms/base_llm/test_base_model_iterator.py @@ -0,0 +1,117 @@ +""" +Tests for BaseModelResponseIterator - specifically testing that empty SSE lines are filtered +""" + +import pytest +from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator +from litellm.types.utils import GenericStreamingChunk + + +class TestBaseModelResponseIterator: + """Test cases for BaseModelResponseIterator empty line filtering""" + + def test_filter_empty_sse_lines_sync(self): + """ + Test that empty SSE lines (common between SSE events) are filtered out + and don't produce empty chunks. + + This fixes the bug where providers using BaseLLMHTTPHandler (like xAI) + would return extra empty chunks when streaming with include_usage=True. + + Related: GitHub Issue #17136 + """ + # Simulate SSE stream with empty lines between events (normal SSE format) + sse_lines = [ + 'data: {"id":"1","choices":[{"delta":{"content":"Hello"}}]}', + '', # Empty line (SSE separator) + 'data: {"id":"1","choices":[{"delta":{"content":" World"}}]}', + '', # Empty line (SSE separator) + 'data: {"id":"1","choices":[],"usage":{"prompt_tokens":10,"completion_tokens":5}}', + '', # Empty line (SSE separator) + 'data: [DONE]', + '', # Empty line after DONE + ] + + iterator = BaseModelResponseIterator( + streaming_response=iter(sse_lines), + sync_stream=True + ) + + chunks = list(iterator) + + # Should have 4 chunks: 2 content + 1 usage + 1 DONE + # Empty lines should be filtered out + assert len(chunks) == 4, f"Expected 4 chunks, got {len(chunks)}" + + # Verify no empty/None chunks were included + # The base iterator returns ModelResponseStream objects + for i, chunk in enumerate(chunks): + assert chunk is not None, f"Chunk {i} should not be None" + + def test_filter_whitespace_only_lines_sync(self): + """Test that lines with only whitespace are also filtered""" + sse_lines = [ + 'data: {"id":"1","choices":[{"delta":{"content":"Hi"}}]}', + ' ', # Whitespace only + '\t', # Tab only + 'data: [DONE]', + ] + + iterator = BaseModelResponseIterator( + streaming_response=iter(sse_lines), + sync_stream=True + ) + + chunks = list(iterator) + + # Should have 2 chunks: 1 content + 1 DONE + assert len(chunks) == 2, f"Expected 2 chunks, got {len(chunks)}" + + def test_valid_chunks_not_filtered_sync(self): + """Test that valid data chunks are not filtered""" + sse_lines = [ + 'data: {"id":"1","choices":[{"delta":{"content":"A"}}]}', + 'data: {"id":"1","choices":[{"delta":{"content":"B"}}]}', + 'data: {"id":"1","choices":[{"delta":{"content":"C"}}]}', + 'data: [DONE]', + ] + + iterator = BaseModelResponseIterator( + streaming_response=iter(sse_lines), + sync_stream=True + ) + + chunks = list(iterator) + + # All 4 chunks should be present + assert len(chunks) == 4, f"Expected 4 chunks, got {len(chunks)}" + + +@pytest.mark.asyncio +async def test_filter_empty_sse_lines_async(): + """ + Test async version: empty SSE lines should be filtered out + """ + async def async_sse_generator(): + lines = [ + 'data: {"id":"1","choices":[{"delta":{"content":"Hello"}}]}', + '', # Empty line + 'data: {"id":"1","choices":[{"delta":{"content":" World"}}]}', + '', # Empty line + 'data: [DONE]', + '', # Empty line + ] + for line in lines: + yield line + + iterator = BaseModelResponseIterator( + streaming_response=async_sse_generator(), + sync_stream=False + ) + + chunks = [] + async for chunk in iterator: + chunks.append(chunk) + + # Should have 3 chunks: 2 content + 1 DONE + assert len(chunks) == 3, f"Expected 3 chunks, got {len(chunks)}" diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 317faa5457..a305009659 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -43,6 +43,29 @@ def test_transform_usage(): ) assert openai_usage._cache_creation_input_tokens == usage["cacheWriteInputTokens"] assert openai_usage._cache_read_input_tokens == usage["cacheReadInputTokens"] + # completion_tokens_details should always be populated + assert openai_usage.completion_tokens_details is not None + assert openai_usage.completion_tokens_details.reasoning_tokens == 0 + assert openai_usage.completion_tokens_details.text_tokens == usage["outputTokens"] + + +def test_transform_usage_with_reasoning_content(): + """Test that completion_tokens_details correctly tracks reasoning vs text tokens.""" + usage = ConverseTokenUsageBlock( + **{ + "inputTokens": 10, + "outputTokens": 100, + "totalTokens": 110, + } + ) + config = AmazonConverseConfig() + reasoning_text = "Let me think about this step by step." + openai_usage = config._transform_usage(usage, reasoning_content=reasoning_text) + assert openai_usage.completion_tokens_details is not None + assert openai_usage.completion_tokens_details.reasoning_tokens > 0 + assert openai_usage.completion_tokens_details.text_tokens == ( + usage["outputTokens"] - openai_usage.completion_tokens_details.reasoning_tokens + ) def test_transform_system_message(): diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 8006ffdff1..5d5aaa64c8 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -110,6 +110,60 @@ def test_get_supported_openai_params_reasoning_effort(): assert "reasoning_effort" not in unsupported_params +@pytest.mark.parametrize( + "api_base, expected_url_prefix", + [ + ( + "https://api.fireworks.ai/inference/v1", + "https://api.fireworks.ai/inference/v1/accounts/", + ), + ( + "https://api.fireworks.ai/inference/v1/", + "https://api.fireworks.ai/inference/v1/accounts/", + ), + ( + "https://custom-host.example.com/v1", + "https://custom-host.example.com/v1/accounts/", + ), + ( + "https://custom-host.example.com/api", + "https://custom-host.example.com/api/v1/accounts/", + ), + ], + ids=["default", "trailing-slash", "custom-with-v1", "custom-without-v1"], +) +def test_get_models_url_no_double_v1(api_base, expected_url_prefix): + """Ensure get_models never produces a /v1/v1/ URL segment (fixes #23106).""" + config = FireworksAIConfig() + account_id = "fireworks" + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "models": [{"name": "accounts/fireworks/models/llama-v3-70b"}] + } + + with ( + patch("litellm.module_level_client.get", return_value=mock_response) as mock_get, + patch( + "litellm.llms.fireworks_ai.chat.transformation.get_secret_str", + side_effect=lambda key: { + "FIREWORKS_API_KEY": "test-key", + "FIREWORKS_API_BASE": api_base, + "FIREWORKS_ACCOUNT_ID": account_id, + }.get(key), + ), + ): + result = config.get_models(api_key="test-key", api_base=api_base) + + called_url = mock_get.call_args.kwargs.get("url") or mock_get.call_args[1].get("url", "") + assert "/v1/v1/" not in called_url, f"Double /v1/ detected in URL: {called_url}" + assert called_url.startswith(expected_url_prefix), ( + f"URL {called_url} does not start with {expected_url_prefix}" + ) + assert result == ["fireworks_ai/accounts/fireworks/models/llama-v3-70b"] + + def test_transform_messages_helper_removes_provider_specific_fields(): """ Test that _transform_messages_helper removes provider_specific_fields from messages. diff --git a/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py b/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py index 39ff0a4f4d..0f743b1a93 100644 --- a/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py +++ b/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py @@ -13,6 +13,7 @@ from litellm.llms.openai.chat.gpt_transformation import ( OpenAIChatCompletionStreamingHandler, OpenAIGPTConfig, ) +from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config class TestOpenAIGPTConfig: @@ -324,3 +325,192 @@ class TestPromptCacheParams: ) assert optional_params.get("prompt_cache_key") == "my-cache-key" assert optional_params.get("prompt_cache_retention") == "24h" + + +class TestGPT5ReasoningEffortPreservation: + """Tests for GPT-5 reasoning_effort dict preservation for Responses API.""" + + def setup_method(self): + self.config = OpenAIGPT5Config() + + def test_reasoning_effort_string_preserved(self): + """Test that reasoning_effort as string is preserved.""" + non_default_params = {"reasoning_effort": "high"} + optional_params = {} + + self.config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model="gpt-5.4", + drop_params=False, + ) + + # String format should be preserved + assert non_default_params.get("reasoning_effort") == "high" + + def test_reasoning_effort_dict_with_only_effort_normalized(self): + """Test that reasoning_effort dict with only 'effort' key is normalized to string.""" + non_default_params = {"reasoning_effort": {"effort": "high"}} + optional_params = {} + + self.config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model="gpt-5.4", + drop_params=False, + ) + + # Dict with only 'effort' should be normalized to string + assert non_default_params.get("reasoning_effort") == "high" + + def test_reasoning_effort_dict_with_summary_preserved(self): + """Test that reasoning_effort dict with 'summary' field is preserved for Responses API. + + Regression test for: User reported that summary field was being dropped when + routing to Responses API. The dict format with additional fields should be + preserved so it can be properly handled by the Responses API transformation. + """ + non_default_params = {"reasoning_effort": {"effort": "high", "summary": "detailed"}} + optional_params = {} + + self.config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model="gpt-5.4", + drop_params=False, + ) + + # Dict with additional fields should be preserved + assert non_default_params.get("reasoning_effort") == {"effort": "high", "summary": "detailed"} + assert isinstance(non_default_params.get("reasoning_effort"), dict) + assert non_default_params["reasoning_effort"]["effort"] == "high" + assert non_default_params["reasoning_effort"]["summary"] == "detailed" + + def test_reasoning_effort_dict_with_generate_summary_preserved(self): + """Test that reasoning_effort dict with 'generate_summary' field is preserved.""" + non_default_params = {"reasoning_effort": {"effort": "medium", "generate_summary": "auto"}} + optional_params = {} + + self.config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model="gpt-5.4", + drop_params=False, + ) + + # Dict with additional fields should be preserved + assert non_default_params.get("reasoning_effort") == {"effort": "medium", "generate_summary": "auto"} + assert isinstance(non_default_params.get("reasoning_effort"), dict) + + def test_reasoning_effort_dict_with_all_fields_preserved(self): + """Test that reasoning_effort dict with all fields is preserved.""" + non_default_params = { + "reasoning_effort": { + "effort": "high", + "summary": "detailed", + "generate_summary": "concise" + } + } + optional_params = {} + + self.config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model="gpt-5.4", + drop_params=False, + ) + + # Dict with all fields should be preserved + reasoning = non_default_params.get("reasoning_effort") + assert isinstance(reasoning, dict) + assert reasoning["effort"] == "high" + assert reasoning["summary"] == "detailed" + assert reasoning["generate_summary"] == "concise" + + def test_reasoning_effort_dict_xhigh_triggers_validation(self): + """xhigh-dict: effective effort is extracted for model-support validation. + + When reasoning_effort={"effort": "xhigh", "summary": "detailed"} is passed to a model + that doesn't support xhigh (e.g. gpt-5.1), the xhigh guard must fire. + """ + import litellm + + non_default_params = {"reasoning_effort": {"effort": "xhigh", "summary": "detailed"}} + optional_params = {} + + with pytest.raises(litellm.utils.UnsupportedParamsError): + self.config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model="gpt-5.1", + drop_params=False, + ) + + def test_reasoning_effort_dict_xhigh_dropped_when_requested(self): + """xhigh-dict with drop_params=True: reasoning_effort is dropped.""" + non_default_params = {"reasoning_effort": {"effort": "xhigh", "summary": "detailed"}} + optional_params = {} + + self.config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model="gpt-5.1", + drop_params=True, + ) + + assert "reasoning_effort" not in non_default_params + + def test_reasoning_effort_dict_none_dropped_for_gpt5_4_with_tools(self): + """none-dict with tools on gpt-5.4: reasoning_effort is dropped.""" + tools = [{"type": "function", "function": {"name": "test", "description": "test"}}] + non_default_params = {"reasoning_effort": {"effort": "none", "summary": "detailed"}, "tools": tools} + optional_params = {} + + self.config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model="gpt-5.4", + drop_params=False, + ) + + assert "reasoning_effort" not in non_default_params + assert non_default_params.get("tools") == tools + + def test_reasoning_effort_dict_none_treated_as_none_for_sampling(self): + """none-dict: {"effort": "none", "summary": "detailed"} allows logprobs/top_p. + + Sampling-param guard should NOT fire; logprobs should be kept. + """ + non_default_params = { + "reasoning_effort": {"effort": "none", "summary": "detailed"}, + "logprobs": True, + } + optional_params = {} + + self.config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model="gpt-5.1", + drop_params=False, + ) + + assert non_default_params.get("reasoning_effort") == {"effort": "none", "summary": "detailed"} + assert non_default_params.get("logprobs") is True + + def test_reasoning_effort_dict_none_allows_temperature(self): + """none-dict: {"effort": "none", "summary": "detailed"} allows non-default temperature.""" + non_default_params = { + "reasoning_effort": {"effort": "none", "summary": "detailed"}, + "temperature": 0.5, + } + optional_params = {} + + self.config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model="gpt-5.1", + drop_params=False, + ) + + assert optional_params.get("temperature") == 0.5 + assert non_default_params.get("reasoning_effort") == {"effort": "none", "summary": "detailed"} diff --git a/tests/test_litellm/llms/openai/test_gpt5_transformation.py b/tests/test_litellm/llms/openai/test_gpt5_transformation.py index b136f8774b..7c731e4e00 100644 --- a/tests/test_litellm/llms/openai/test_gpt5_transformation.py +++ b/tests/test_litellm/llms/openai/test_gpt5_transformation.py @@ -324,10 +324,11 @@ def test_gpt5_4_pro_allows_reasoning_effort_xhigh(config: OpenAIConfig): assert params["reasoning_effort"] == "xhigh" -def test_gpt5_normalizes_reasoning_effort_dict_to_string(config: OpenAIConfig): - """Chat completion API expects reasoning_effort as a string, not a dict. +def test_gpt5_preserves_reasoning_effort_dict_with_summary(config: OpenAIConfig): + """Dict with summary/generate_summary is preserved for Responses API. Config/deployments may pass Responses API format: {'effort': 'high', 'summary': 'detailed'}. + We preserve the full dict so it reaches the Responses API transformation. """ params = config.map_openai_params( non_default_params={"reasoning_effort": {"effort": "high", "summary": "detailed"}}, @@ -335,21 +336,85 @@ def test_gpt5_normalizes_reasoning_effort_dict_to_string(config: OpenAIConfig): model="gpt-5.4", drop_params=False, ) - assert params["reasoning_effort"] == "high" + assert params["reasoning_effort"] == {"effort": "high", "summary": "detailed"} -def test_gpt5_normalizes_reasoning_effort_dict_from_optional_params(config: OpenAIConfig): - """reasoning_effort dict in optional_params (e.g. from model config) is normalized.""" +def test_gpt5_xhigh_dict_triggers_validation(config: OpenAIConfig): + """Dict with effort='xhigh' triggers xhigh model-support validation. + + Regression: when reasoning_effort is a dict, effective_effort must be used for + the xhigh guard so validation is not silently skipped. + """ + with pytest.raises(litellm.utils.UnsupportedParamsError): + config.map_openai_params( + non_default_params={"reasoning_effort": {"effort": "xhigh", "summary": "detailed"}}, + optional_params={}, + model="gpt-5.1", + drop_params=False, + ) + + +def test_gpt5_xhigh_dict_accepted_for_supported_model(config: OpenAIConfig): + """Dict with effort='xhigh' passes through for gpt-5.4+.""" + params = config.map_openai_params( + non_default_params={"reasoning_effort": {"effort": "xhigh", "summary": "detailed"}}, + optional_params={}, + model="gpt-5.4", + drop_params=False, + ) + assert params["reasoning_effort"] == {"effort": "xhigh", "summary": "detailed"} + + +def test_gpt5_none_dict_with_tools_no_tool_drop(config: OpenAIConfig): + """Dict with effort='none' and tools: reasoning_effort dropped for gpt-5.4. + + gpt-5.4 drops all reasoning_effort when tools are present, + since that combination is only supported in the Responses API. + """ + tools = [{"type": "function", "function": {"name": "test", "description": "test"}}] + params = config.map_openai_params( + non_default_params={"reasoning_effort": {"effort": "none", "summary": "detailed"}, "tools": tools}, + optional_params={}, + model="gpt-5.4", + drop_params=False, + ) + assert "reasoning_effort" not in params + assert params["tools"] == tools + + +def test_gpt5_none_dict_with_sampling_params_allowed(config: OpenAIConfig): + """Dict with effort='none' allows logprobs/top_p/top_logprobs. + + Regression: effective_effort='none' must be used for sampling guard so + {"effort": "none", "summary": "detailed"} does not incorrectly trigger sampling errors. + """ + params = config.map_openai_params( + non_default_params={ + "reasoning_effort": {"effort": "none", "summary": "detailed"}, + "logprobs": True, + "top_p": 0.9, + }, + optional_params={}, + model="gpt-5.1", + drop_params=False, + ) + assert params["reasoning_effort"] == {"effort": "none", "summary": "detailed"} + assert params["logprobs"] is True + assert params["top_p"] == 0.9 + + +def test_gpt5_preserves_reasoning_effort_dict_with_summary_from_optional_params(config: OpenAIConfig): + """reasoning_effort dict with summary in optional_params is preserved.""" params = config.map_openai_params( non_default_params={}, optional_params={"reasoning_effort": {"effort": "medium", "summary": "detailed"}}, model="gpt-5.4", drop_params=False, ) - assert params["reasoning_effort"] == "medium" + assert params["reasoning_effort"] == {"effort": "medium", "summary": "detailed"} -def test_gpt5_4_drops_reasoning_effort_when_tools_present(config: OpenAIConfig): +def test_gpt5_4_drops_reasoning_effort_when_user_sends_reasoning_and_tools(config: OpenAIConfig): """gpt-5.4: function calls not supported with reasoning_effort != 'none'. Drop reasoning_effort.""" tools = [{"type": "function", "function": {"name": "test", "description": "test"}}] params = config.map_openai_params( @@ -373,8 +438,8 @@ def test_gpt5_4_keeps_reasoning_effort_when_no_tools(config: OpenAIConfig): assert params["reasoning_effort"] == "high" -def test_gpt5_4_keeps_reasoning_effort_none_with_tools(config: OpenAIConfig): - """reasoning_effort='none' is kept when tools are present.""" +def test_gpt5_4_drops_reasoning_effort_none_with_tools(config: OpenAIConfig): + """reasoning_effort='none' is also dropped when tools are present for gpt-5.4.""" tools = [{"type": "function", "function": {"name": "test", "description": "test"}}] params = config.map_openai_params( non_default_params={"reasoning_effort": "none", "tools": tools}, @@ -382,7 +447,7 @@ def test_gpt5_4_keeps_reasoning_effort_none_with_tools(config: OpenAIConfig): model="gpt-5.4", drop_params=False, ) - assert params["reasoning_effort"] == "none" + assert "reasoning_effort" not in params assert params["tools"] == tools diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_embedding_role_assumption.py b/tests/test_litellm/llms/sagemaker/test_sagemaker_embedding_role_assumption.py new file mode 100644 index 0000000000..82c84af5e2 --- /dev/null +++ b/tests/test_litellm/llms/sagemaker/test_sagemaker_embedding_role_assumption.py @@ -0,0 +1,243 @@ +""" +Test cases for SageMaker embedding role assumption support + +This module tests that the SageMaker embedding handler properly supports +AWS IAM role assumption via aws_role_name and aws_session_name parameters, +matching the behavior of the completion handler. +""" + +import json +import os +import sys +from datetime import timezone +from unittest.mock import MagicMock, call, patch + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from botocore.credentials import Credentials + +from litellm.llms.sagemaker.completion.handler import SagemakerLLM +from litellm.types.utils import EmbeddingResponse + + +class TestSagemakerEmbeddingRoleAssumption: + """Test that SageMaker embedding supports role assumption like completion does""" + + def setup_method(self): + self.sagemaker_llm = SagemakerLLM() + + def test_embedding_uses_load_credentials(self): + """ + Test that embedding() calls _load_credentials() to support role assumption. + This ensures aws_role_name and aws_session_name parameters are properly handled. + """ + # Mock credentials that would be returned after role assumption + mock_credentials = Credentials( + access_key="assumed-access-key", + secret_key="assumed-secret-key", + token="assumed-session-token", + ) + + # Mock the SageMaker client response + mock_sagemaker_client = MagicMock() + mock_sagemaker_client.invoke_endpoint.return_value = { + "Body": MagicMock( + read=MagicMock(return_value=json.dumps({"embedding": [[0.1, 0.2, 0.3]]}).encode()) + ) + } + + # Mock boto3.Session to return our mock client + mock_session = MagicMock() + mock_session.client.return_value = mock_sagemaker_client + + with patch.object( + self.sagemaker_llm, "_load_credentials", return_value=(mock_credentials, "us-east-1") + ) as mock_load_creds, patch("boto3.Session", return_value=mock_session): + + # Create mock logging object + mock_logging = MagicMock() + + optional_params = { + "aws_role_name": "arn:aws:iam::123456789012:role/TestRole", + "aws_session_name": "test-session", + } + + self.sagemaker_llm.embedding( + model="test-endpoint", + input=["hello world"], + model_response=EmbeddingResponse(), + print_verbose=print, + encoding=None, + logging_obj=mock_logging, + optional_params=optional_params, + ) + + # Verify _load_credentials was called with the optional_params + mock_load_creds.assert_called_once() + + # Verify boto3.Session was created with the assumed credentials + mock_session_calls = mock_session.client.call_args_list + assert len(mock_session_calls) == 1 + assert mock_session_calls[0] == call(service_name="sagemaker-runtime") + + def test_embedding_role_assumption_with_sts(self): + """ + Test the full role assumption flow for embeddings, similar to completion. + Verifies that STS assume_role is called when aws_role_name is provided. + """ + # Mock the STS client for role assumption + mock_sts_client = MagicMock() + + # Mock the STS response with proper expiration handling + mock_expiry = MagicMock() + mock_expiry.tzinfo = timezone.utc + time_diff = MagicMock() + time_diff.total_seconds.return_value = 3600 + mock_expiry.__sub__ = MagicMock(return_value=time_diff) + + mock_sts_response = { + "Credentials": { + "AccessKeyId": "assumed-access-key", + "SecretAccessKey": "assumed-secret-key", + "SessionToken": "assumed-session-token", + "Expiration": mock_expiry, + } + } + mock_sts_client.assume_role.return_value = mock_sts_response + + # Mock the SageMaker client response + mock_sagemaker_client = MagicMock() + mock_sagemaker_client.invoke_endpoint.return_value = { + "Body": MagicMock( + read=MagicMock(return_value=json.dumps({"embedding": [[0.1, 0.2, 0.3]]}).encode()) + ) + } + + # Mock boto3.Session for SageMaker client creation + mock_session = MagicMock() + mock_session.client.return_value = mock_sagemaker_client + + def mock_boto3_client(service_name, **kwargs): + if service_name == "sts": + return mock_sts_client + return mock_sagemaker_client + + with patch("boto3.client", side_effect=mock_boto3_client), \ + patch("boto3.Session", return_value=mock_session): + + mock_logging = MagicMock() + + optional_params = { + "aws_role_name": "arn:aws:iam::123456789012:role/CrossAccountRole", + "aws_session_name": "litellm-embedding-session", + "aws_region_name": "us-east-1", + } + + self.sagemaker_llm.embedding( + model="test-endpoint", + input=["hello world"], + model_response=EmbeddingResponse(), + print_verbose=print, + encoding=None, + logging_obj=mock_logging, + optional_params=optional_params, + ) + + # Verify STS assume_role was called with correct parameters + mock_sts_client.assume_role.assert_called_once() + call_args = mock_sts_client.assume_role.call_args + assert call_args[1]["RoleArn"] == "arn:aws:iam::123456789012:role/CrossAccountRole" + assert call_args[1]["RoleSessionName"] == "litellm-embedding-session" + + def test_embedding_without_role_assumption(self): + """ + Test that embedding works without role assumption when aws_role_name is not provided. + Should use default credentials from environment/instance profile. + """ + # Mock the SageMaker client response + mock_sagemaker_client = MagicMock() + mock_sagemaker_client.invoke_endpoint.return_value = { + "Body": MagicMock( + read=MagicMock(return_value=json.dumps({"embedding": [[0.1, 0.2, 0.3]]}).encode()) + ) + } + + mock_session = MagicMock() + mock_session.client.return_value = mock_sagemaker_client + + # Mock credentials returned from environment + mock_credentials = Credentials( + access_key="env-access-key", + secret_key="env-secret-key", + token=None, + ) + + with patch.object( + self.sagemaker_llm, "_load_credentials", return_value=(mock_credentials, "us-west-2") + ), patch("boto3.Session", return_value=mock_session): + + mock_logging = MagicMock() + + # No aws_role_name provided + optional_params = { + "aws_region_name": "us-west-2", + } + + result = self.sagemaker_llm.embedding( + model="test-endpoint", + input=["hello world"], + model_response=EmbeddingResponse(), + print_verbose=print, + encoding=None, + logging_obj=mock_logging, + optional_params=optional_params, + ) + + # Should still work and return embeddings + assert result is not None + + def test_embedding_session_created_with_assumed_credentials(self): + """ + Test that boto3.Session is created with the credentials from role assumption. + This verifies the credentials flow from _load_credentials to the SageMaker client. + """ + mock_credentials = Credentials( + access_key="assumed-key", + secret_key="assumed-secret", + token="assumed-token", + ) + + mock_sagemaker_client = MagicMock() + mock_sagemaker_client.invoke_endpoint.return_value = { + "Body": MagicMock( + read=MagicMock(return_value=json.dumps({"embedding": [[0.1, 0.2, 0.3]]}).encode()) + ) + } + + with patch.object( + self.sagemaker_llm, "_load_credentials", return_value=(mock_credentials, "us-east-1") + ), patch("boto3.Session") as mock_session_class: + + mock_session = MagicMock() + mock_session.client.return_value = mock_sagemaker_client + mock_session_class.return_value = mock_session + + mock_logging = MagicMock() + + self.sagemaker_llm.embedding( + model="test-endpoint", + input=["hello world"], + model_response=EmbeddingResponse(), + print_verbose=print, + encoding=None, + logging_obj=mock_logging, + optional_params={}, + ) + + # Verify Session was created with the assumed credentials + mock_session_class.assert_called_once_with( + aws_access_key_id="assumed-key", + aws_secret_access_key="assumed-secret", + aws_session_token="assumed-token", + region_name="us-east-1", + ) diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py index 444125dffa..ce3d2daa74 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py @@ -128,6 +128,75 @@ def test_vertex_ai_includes_labels(): +def test_extra_body_cache_not_forwarded_to_vertex_ai(): + """ + 'cache' inside extra_body is a LiteLLM-internal proxy caching control. + It must NOT be forwarded to the Vertex AI request body. + + Regression test for: "Invalid JSON payload received. Unknown name \"cache\": Cannot find field." + Vertex AI enforces a strict JSON schema and rejects any unknown field. + """ + messages = [{"role": "user", "content": "test"}] + optional_params = { + "extra_body": { + "cache": {"use-cache": True, "ttl": 86400}, # LiteLLM-internal + "some_vertex_param": "value", # legitimate provider extra + }, + } + litellm_params = {} + + result = _transform_request_body( + messages=messages, + model="gemini-2.5-pro", + optional_params=optional_params, + custom_llm_provider="vertex_ai", + litellm_params=litellm_params, + cached_content=None, + ) + + # 'cache' must be stripped — Vertex AI has no such field + assert "cache" not in result, ( + "extra_body.cache must not be forwarded to Vertex AI. " + "Vertex AI rejects it with 400: Unknown name \"cache\": Cannot find field." + ) + + # Other legitimate extra_body keys should still pass through + assert "some_vertex_param" in result + assert result["some_vertex_param"] == "value" + + # Core request fields must be present + assert "contents" in result + + +def test_extra_body_tags_not_forwarded_to_vertex_ai(): + """ + 'tags' inside extra_body is a LiteLLM-internal param for logging/tracking. + It must NOT be forwarded to the Vertex AI request body. + Documented in litellm_proxy.md: "Send tags by including them in the extra_body parameter" + """ + messages = [{"role": "user", "content": "test"}] + optional_params = { + "extra_body": { + "tags": ["user:alice", "env:prod"], + "custom_param": "allowed", + }, + } + litellm_params = {} + + result = _transform_request_body( + messages=messages, + model="gemini-2.5-pro", + optional_params=optional_params, + custom_llm_provider="vertex_ai", + litellm_params=litellm_params, + cached_content=None, + ) + + assert "tags" not in result + assert "custom_param" in result + assert result["custom_param"] == "allowed" + + def test_metadata_to_labels_vertex_only(): """Test that metadata->labels conversion only happens for Vertex AI""" messages = [{"role": "user", "content": "test"}] diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py index 94323e0690..d483a81a34 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -558,69 +558,49 @@ def test_get_vertex_url_global_region(stream, expected_endpoint_suffix): assert url == expected_url -@pytest.mark.parametrize( - "supported_regions, expected_result", - [ - (None, False), # get_supported_regions returns None - ([], False), # empty list, no global region - (["us-central1"], False), # only regional, no global - (["global"], True), # only global region - (["global", "us-central1"], True), # global and other regions - ( - ["us-central1", "global", "europe-west1"], - True, - ), # global among multiple regions - ], -) -def test_is_global_only_vertex_model(supported_regions, expected_result): - """Test is_global_only_vertex_model with various supported regions scenarios""" - from litellm.llms.vertex_ai.common_utils import is_global_only_vertex_model - - with patch("litellm.utils.get_supported_regions") as mock_get_supported_regions: - mock_get_supported_regions.return_value = supported_regions - - result = is_global_only_vertex_model("test-model") - - assert result == expected_result - mock_get_supported_regions.assert_called_once_with( - model="test-model", custom_llm_provider="vertex_ai" - ) - @pytest.mark.parametrize( - "model_is_global_only, vertex_region, expected_region", + "model_cost_entry, vertex_region, expected_region", [ - (True, None, "global"), # Global-only model with no region specified - (True, "us-central1", "global"), # Global-only model overrides specified region - (True, "europe-west1", "global"), # Global-only model overrides any region - (False, None, "us-central1"), # Non-global model defaults to us-central1 - ( - False, - "europe-west1", - "europe-west1", - ), # Non-global model uses specified region - (False, "us-east1", "us-east1"), # Non-global model uses specified region + # Model with supported_regions=["global"], no user region -> use "global" + ({"supported_regions": ["global"]}, None, "global"), + # Model with supported_regions=["global"], user passes unsupported region -> override to "global" + ({"supported_regions": ["global"]}, "us-central1", "global"), + # Model with supported_regions=["global"], user passes unsupported region -> override to "global" + ({"supported_regions": ["global"]}, "europe-west1", "global"), + # Model with supported_regions=["us-west2"], no user region -> use "us-west2" + ({"supported_regions": ["us-west2"]}, None, "us-west2"), + # Model with supported_regions=["us-west2", "us-central1"], user passes supported region -> respect it + ({"supported_regions": ["us-west2", "us-central1"]}, "us-central1", "us-central1"), + # Model with supported_regions=["us-west2", "us-central1"], user passes unsupported region -> override + ({"supported_regions": ["us-west2", "us-central1"]}, "europe-west1", "us-west2"), + # No model_cost entry, no user region -> default us-central1 + ({}, None, "us-central1"), + # No model_cost entry, user specifies region -> use specified region + ({}, "europe-west1", "europe-west1"), + # No model_cost entry, user specifies region -> use specified region + ({}, "us-east1", "us-east1"), ], ) def test_get_vertex_region_global_only_model( - model_is_global_only, vertex_region, expected_region + model_cost_entry, vertex_region, expected_region ): - """Test get_vertex_region ensures global-only models default to 'global' region""" + """Test get_vertex_region resolves region from model_cost supported_regions""" + import litellm from litellm.llms.vertex_ai.vertex_llm_base import VertexBase vertex_base = VertexBase() - with patch( - "litellm.llms.vertex_ai.vertex_llm_base.is_global_only_vertex_model" - ) as mock_is_global_only: - mock_is_global_only.return_value = model_is_global_only - + with patch.dict( + litellm.model_cost, + {"vertex_ai/test-model": model_cost_entry}, + clear=False, + ): result = vertex_base.get_vertex_region( vertex_region=vertex_region, model="test-model" ) assert result == expected_region - mock_is_global_only.assert_called_once_with("test-model") def test_vertex_filter_format_uri(): diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py index dafd58b06d..53e42a519b 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py @@ -2,8 +2,8 @@ Tests for Vertex AI Qwen MaaS models that require the global endpoint. These tests verify that: -1. Qwen models are correctly identified as global-only models -2. The correct global URL is constructed (https://aiplatform.googleapis.com) +1. The correct global URL is constructed (https://aiplatform.googleapis.com) +2. The get_vertex_region method resolves regions from model_cost supported_regions 3. The completion() and responses() API work with Qwen models """ @@ -19,7 +19,6 @@ sys.path.insert( ) # Adds the parent directory to the system path import litellm -from litellm.llms.vertex_ai.common_utils import is_global_only_vertex_model from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.types.llms.vertex_ai import VertexPartnerProvider @@ -48,66 +47,36 @@ def clean_vertex_env(): os.environ[var] = value -class TestQwenGlobalOnlyDetection: - """Test that Qwen models are correctly identified as global-only.""" - - @pytest.mark.parametrize( - "model", - [ - "vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas", - "vertex_ai/qwen/qwen3-next-80b-a3b-thinking-maas", - "vertex_ai/qwen/qwen3-235b-a22b-instruct-2507-maas", - "vertex_ai/qwen/qwen3-coder-480b-a35b-instruct-maas", - ], - ) - def test_qwen_models_are_global_only(self, model): - """Test that Qwen MaaS models are identified as global-only.""" - # This test requires the model_cost to have supported_regions: ["global"] - # If the model is not in model_cost, it should return False (fallback behavior) - result = is_global_only_vertex_model(model) - # Note: This will return True only if the model is in model_cost with supported_regions: ["global"] - # If running without the updated model_cost, this may return False - assert isinstance(result, bool) - - def test_non_global_model_returns_false(self): - """Test that non-global models return False.""" - result = is_global_only_vertex_model("vertex_ai/gemini-1.5-pro") - assert result is False - - def test_unknown_model_returns_false(self): - """Test that unknown models return False (fallback behavior).""" - result = is_global_only_vertex_model("vertex_ai/unknown-model-xyz") - assert result is False - - class TestVertexBaseGetVertexRegion: - """Test the get_vertex_region method.""" + """Test the get_vertex_region method using model_cost lookup.""" - def test_global_only_model_returns_global(self): - """Test that global-only models return 'global' regardless of input.""" + def test_global_model_no_user_region_returns_global(self): + """Test that global-only models return 'global' when user doesn't specify region.""" vertex_base = VertexBase() - with patch( - "litellm.llms.vertex_ai.vertex_llm_base.is_global_only_vertex_model", - return_value=True, - ): - result = vertex_base.get_vertex_region( - vertex_region="us-central1", - model="vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas", - ) - assert result == "global" - - def test_global_only_model_with_none_returns_global(self): - """Test that global-only models return 'global' even with None input.""" - vertex_base = VertexBase() - - with patch( - "litellm.llms.vertex_ai.vertex_llm_base.is_global_only_vertex_model", - return_value=True, + with patch.dict( + litellm.model_cost, + {"vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas": {"supported_regions": ["global"]}}, + clear=False, ): result = vertex_base.get_vertex_region( vertex_region=None, - model="vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas", + model="qwen/qwen3-next-80b-a3b-instruct-maas", + ) + assert result == "global" + + def test_global_model_with_unsupported_user_region_overrides(self): + """Test that unsupported user region is overridden for global-only models.""" + vertex_base = VertexBase() + + with patch.dict( + litellm.model_cost, + {"vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas": {"supported_regions": ["global"]}}, + clear=False, + ): + result = vertex_base.get_vertex_region( + vertex_region="us-central1", + model="qwen/qwen3-next-80b-a3b-instruct-maas", ) assert result == "global" @@ -115,13 +84,10 @@ class TestVertexBaseGetVertexRegion: """Test that non-global models use the provided region.""" vertex_base = VertexBase() - with patch( - "litellm.llms.vertex_ai.vertex_llm_base.is_global_only_vertex_model", - return_value=False, - ): + with patch.dict(litellm.model_cost, {}, clear=False): result = vertex_base.get_vertex_region( vertex_region="europe-west1", - model="vertex_ai/gemini-1.5-pro", + model="gemini-1.5-pro", ) assert result == "europe-west1" @@ -129,13 +95,10 @@ class TestVertexBaseGetVertexRegion: """Test that non-global models with None region fallback to us-central1.""" vertex_base = VertexBase() - with patch( - "litellm.llms.vertex_ai.vertex_llm_base.is_global_only_vertex_model", - return_value=False, - ): + with patch.dict(litellm.model_cost, {}, clear=False): result = vertex_base.get_vertex_region( vertex_region=None, - model="vertex_ai/gemini-1.5-pro", + model="unknown-model-xyz", ) assert result == "us-central1" @@ -217,15 +180,15 @@ async def test_vertex_ai_qwen_global_endpoint_url(): client, "post", side_effect=mock_post_func ) as mock_post, patch.object( VertexLLM, "_ensure_access_token", return_value=("fake-token", "test-project") - ), patch( - "litellm.llms.vertex_ai.vertex_llm_base.is_global_only_vertex_model", - return_value=True, + ), patch.dict( + litellm.model_cost, + {"vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas": {"supported_regions": ["global"]}}, + clear=False, ): response = await litellm.acompletion( model="vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas", messages=[{"role": "user", "content": "Hello"}], vertex_ai_project="test-project", - vertex_ai_location="us-central1", client=client, ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 314eed9598..edc69ad6a4 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -2095,6 +2095,153 @@ async def test_get_tools_from_mcp_servers_logs_list_tools_to_spendlogs_when_enab assert spend_meta["per_server_tool_counts"]["server_a"] == 1 +def test_tool_name_matches_case_insensitive(): + """Test that _tool_name_matches performs case-insensitive comparison. + + This is critical for OpenAPI-based MCP servers where: + 1. operationIds are often in camelCase (e.g., 'addPet', 'updatePet') + 2. Tool names are lowercased during registration (e.g., 'addpet', 'updatepet') + 3. allowed_tools configuration may use the original camelCase names + + Without case-insensitive matching, all tools would be filtered out. + """ + try: + from litellm.proxy._experimental.mcp_server.server import _tool_name_matches + except ImportError: + pytest.skip("MCP server not available") + + # Test case 1: Unprefixed tool name with camelCase in filter list + assert _tool_name_matches("addpet", ["addPet", "updatePet"]) is True + assert _tool_name_matches("updatepet", ["addPet", "updatePet"]) is True + assert _tool_name_matches("deletepet", ["addPet", "updatePet"]) is False + + # Test case 2: Prefixed tool name with camelCase in filter list + assert _tool_name_matches("per_store-addpet", ["addPet", "updatePet"]) is True + assert _tool_name_matches("per_store-updatepet", ["addPet", "updatePet"]) is True + assert _tool_name_matches("per_store-deletepet", ["addPet", "updatePet"]) is False + + # Test case 3: Mixed case variations + assert _tool_name_matches("findPetsByStatus", ["findpetsbystatus"]) is True + assert _tool_name_matches("findpetsbystatus", ["findPetsByStatus"]) is True + assert _tool_name_matches("FINDPETSBYSTATUS", ["findPetsByStatus"]) is True + + # Test case 4: Full prefixed name in filter list (case-insensitive) + assert _tool_name_matches("server-addPet", ["server-addpet"]) is True + assert _tool_name_matches("server-addpet", ["server-addPet"]) is True + + # Test case 5: Ensure non-matching names still don't match + assert _tool_name_matches("addpet", ["deletePet", "updatePet"]) is False + assert _tool_name_matches("server-addpet", ["deletePet", "updatePet"]) is False + + +def test_filter_tools_by_allowed_tools_case_insensitive(): + """Test that filter_tools_by_allowed_tools handles case-insensitive matching. + + Ensures that OpenAPI tools with lowercase names can be filtered using + camelCase allowed_tools configuration from the OpenAPI spec. + """ + try: + from litellm.proxy._experimental.mcp_server.server import ( + filter_tools_by_allowed_tools, + ) + from litellm.types.mcp_server.tool_registry import MCPTool + except ImportError: + pytest.skip("MCP server not available") + + # Mock handler function + def mock_handler(**kwargs): + return kwargs + + # Create mock tools with lowercase names (as registered from OpenAPI) + tools = [ + MCPTool( + name="per_store-addpet", + description="Add a pet", + input_schema={"type": "object"}, + handler=mock_handler, + ), + MCPTool( + name="per_store-updatepet", + description="Update a pet", + input_schema={"type": "object"}, + handler=mock_handler, + ), + MCPTool( + name="per_store-deletepet", + description="Delete a pet", + input_schema={"type": "object"}, + handler=mock_handler, + ), + MCPTool( + name="per_store-findpetsbystatus", + description="Find pets by status", + input_schema={"type": "object"}, + handler=mock_handler, + ), + ] + + # Create mock server with camelCase allowed_tools (as from OpenAPI spec) + server = MCPServer( + server_id="test-server", + name="per_store", + transport=MCPTransport.http, + allowed_tools=["addPet", "updatePet", "findPetsByStatus"], + ) + + # Filter tools + filtered_tools = filter_tools_by_allowed_tools(tools, server) + + # Should return 3 tools (case-insensitive match) + assert len(filtered_tools) == 3 + assert any(t.name == "per_store-addpet" for t in filtered_tools) + assert any(t.name == "per_store-updatepet" for t in filtered_tools) + assert any(t.name == "per_store-findpetsbystatus" for t in filtered_tools) + assert not any(t.name == "per_store-deletepet" for t in filtered_tools) + + +def test_filter_tools_by_allowed_tools_no_filter(): + """Test that filter_tools_by_allowed_tools returns all tools when no filter is set.""" + try: + from litellm.proxy._experimental.mcp_server.server import ( + filter_tools_by_allowed_tools, + ) + from litellm.types.mcp_server.tool_registry import MCPTool + except ImportError: + pytest.skip("MCP server not available") + + # Mock handler function + def mock_handler(**kwargs): + return kwargs + + tools = [ + MCPTool( + name="fusion_litellm_mcp-model_list", + description="List models", + input_schema={"type": "object"}, + handler=mock_handler, + ), + MCPTool( + name="fusion_litellm_mcp-chat_completion", + description="Chat completion", + input_schema={"type": "object"}, + handler=mock_handler, + ), + ] + + # Server with no allowed_tools filter + server = MCPServer( + server_id="test-server", + name="fusion_litellm_mcp", + transport=MCPTransport.http, + allowed_tools=None, + ) + + filtered_tools = filter_tools_by_allowed_tools(tools, server) + + # Should return all tools when no filter is configured + assert len(filtered_tools) == 2 + + @pytest.mark.asyncio async def test_get_tools_from_mcp_servers_injects_stored_oauth2_token(): """ diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index 193b014f03..c43621d7f7 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -21,6 +21,140 @@ def test_get_team_models_for_all_models_and_team_only_models(): assert set(result) == set(combined_models) +def test_get_team_models_all_proxy_models_includes_access_groups(): + """ + When a team has 'all-proxy-models' and include_model_access_groups=True, + the result should include model access group names (e.g. 'claude-model-group') + in addition to individual model names. + """ + from litellm.proxy.auth.model_checks import get_team_models + + team_models = ["all-proxy-models"] + proxy_model_list = ["model1", "model2"] + model_access_groups = { + "group-a": ["model1"], + "group-b": ["model2"], + } + + result = get_team_models( + team_models, proxy_model_list, model_access_groups, include_model_access_groups=True + ) + assert "group-a" in result + assert "group-b" in result + assert "model1" in result + assert "model2" in result + assert len(result) == len(set(result)), "result should have no duplicates" + + +def test_get_team_models_all_proxy_models_without_include_flag(): + """ + When include_model_access_groups=False, access group names should NOT + appear in the result even with 'all-proxy-models'. + """ + from litellm.proxy.auth.model_checks import get_team_models + + team_models = ["all-proxy-models"] + proxy_model_list = ["model1", "model2"] + model_access_groups = { + "group-a": ["model1"], + "group-b": ["model2"], + } + + result = get_team_models( + team_models, proxy_model_list, model_access_groups, include_model_access_groups=False + ) + assert "group-a" not in result + assert "group-b" not in result + assert "model1" in result + assert "model2" in result + + +def test_get_key_models_all_proxy_models_includes_access_groups(): + """ + When a key has 'all-proxy-models' and include_model_access_groups=True, + the result should include model access group names. + """ + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.model_checks import get_key_models + + user_api_key_dict = UserAPIKeyAuth( + models=["all-proxy-models"], + api_key="test-key", + ) + proxy_model_list = ["model1", "model2"] + model_access_groups = { + "group-a": ["model1"], + } + + result = get_key_models( + user_api_key_dict=user_api_key_dict, + proxy_model_list=proxy_model_list, + model_access_groups=model_access_groups, + include_model_access_groups=True, + ) + assert "group-a" in result + assert "model1" in result + assert "model2" in result + assert len(result) == len(set(result)), "result should have no duplicates" + + +def test_get_key_models_passes_include_model_access_groups(): + """ + When a key explicitly has an access group name in its models list and + include_model_access_groups=True, the group name should be retained + (not stripped by _get_models_from_access_groups). + """ + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.model_checks import get_key_models + + user_api_key_dict = UserAPIKeyAuth( + models=["group-a"], + api_key="test-key", + ) + proxy_model_list = ["model1", "model2"] + model_access_groups = { + "group-a": ["model1", "model2"], + } + + result = get_key_models( + user_api_key_dict=user_api_key_dict, + proxy_model_list=proxy_model_list, + model_access_groups=model_access_groups, + include_model_access_groups=True, + ) + assert "group-a" in result + assert "model1" in result + assert "model2" in result + + +def test_get_key_models_does_not_mutate_input(): + """ + get_key_models must not mutate user_api_key_dict.models in-place. + _get_models_from_access_groups uses .pop()/.extend() which would corrupt + cached UserAPIKeyAuth objects if all_models were an alias instead of a copy. + """ + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.model_checks import get_key_models + + original_models = ["group-a", "extra-model"] + user_api_key_dict = UserAPIKeyAuth( + models=list(original_models), # give it a list + api_key="test-key", + ) + model_access_groups = { + "group-a": ["model1", "model2"], + } + + _ = get_key_models( + user_api_key_dict=user_api_key_dict, + proxy_model_list=["model1", "model2"], + model_access_groups=model_access_groups, + include_model_access_groups=False, + ) + # The original models list on the auth object must be unchanged + assert user_api_key_dict.models == original_models + + @pytest.mark.parametrize( "key_models,team_models,proxy_model_list,model_list,expected", [ diff --git a/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py index 6da3d1f918..f330b40282 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py @@ -244,6 +244,120 @@ async def test_delete_tag(): app.dependency_overrides.clear() +@pytest.mark.asyncio +async def test_list_tags_with_dynamic_tags(): + """ + Test that list_tags uses group_by to get distinct dynamic tags efficiently + and merges them with stored tags, excluding duplicates. + """ + from datetime import datetime + from unittest.mock import AsyncMock, Mock + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + mock_user_auth = UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth + + try: + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + mock_db = Mock() + mock_prisma.db = mock_db + + # Setup stored tags + stored_tag = Mock() + stored_tag.tag_name = "stored-tag" + stored_tag.description = "A stored tag" + stored_tag.models = ["model-1"] + stored_tag.model_info = {} + stored_tag.spend = 0.0 + stored_tag.budget_id = None + stored_tag.created_at = datetime(2025, 1, 1) + stored_tag.updated_at = datetime(2025, 1, 1) + stored_tag.created_by = "user-123" + stored_tag.litellm_budget_table = None + mock_db.litellm_tagtable.find_many = AsyncMock(return_value=[stored_tag]) + + # Setup dynamic tags via group_by — includes one that overlaps with stored + mock_db.litellm_dailytagspend.group_by = AsyncMock(return_value=[ + {"tag": "dynamic-tag-1", "_min": {"created_at": datetime(2025, 2, 1)}, "_max": {"updated_at": datetime(2025, 3, 1)}}, + {"tag": "dynamic-tag-2", "_min": {"created_at": datetime(2025, 2, 2)}, "_max": {"updated_at": datetime(2025, 3, 2)}}, + {"tag": "stored-tag", "_min": {"created_at": datetime(2025, 1, 1)}, "_max": {"updated_at": datetime(2025, 1, 1)}}, # duplicate, should be excluded + ]) + + headers = {"Authorization": "Bearer sk-1234"} + response = client.get("/tag/list", headers=headers) + + assert response.status_code == 200 + result = response.json() + + # Should have 1 stored + 2 dynamic (the duplicate excluded) + assert len(result) == 3 + + tag_names = [t["name"] for t in result] + assert "stored-tag" in tag_names + assert "dynamic-tag-1" in tag_names + assert "dynamic-tag-2" in tag_names + + # Verify dynamic tags include created_at/updated_at + dynamic_tags = {t["name"]: t for t in result if t["name"].startswith("dynamic-")} + assert dynamic_tags["dynamic-tag-1"]["created_at"] is not None + assert dynamic_tags["dynamic-tag-1"]["updated_at"] is not None + + finally: + app.dependency_overrides.clear() + + +@pytest.mark.asyncio +async def test_list_tags_no_dynamic_tags(): + """ + Test list_tags when there are no dynamic tags in the spend table. + """ + from datetime import datetime + from unittest.mock import AsyncMock, Mock + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + mock_user_auth = UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth + + try: + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + mock_db = Mock() + mock_prisma.db = mock_db + + stored_tag = Mock() + stored_tag.tag_name = "stored-tag" + stored_tag.description = "A stored tag" + stored_tag.models = [] + stored_tag.model_info = None + stored_tag.spend = 0.0 + stored_tag.budget_id = None + stored_tag.created_at = datetime(2025, 1, 1) + stored_tag.updated_at = datetime(2025, 1, 1) + stored_tag.created_by = "user-123" + stored_tag.litellm_budget_table = None + mock_db.litellm_tagtable.find_many = AsyncMock(return_value=[stored_tag]) + + mock_db.litellm_dailytagspend.group_by = AsyncMock(return_value=[]) + + headers = {"Authorization": "Bearer sk-1234"} + response = client.get("/tag/list", headers=headers) + + assert response.status_code == 200 + result = response.json() + assert len(result) == 1 + assert result[0]["name"] == "stored-tag" + + finally: + app.dependency_overrides.clear() + + @pytest.mark.asyncio async def test_get_deployments_by_model_id(): """ diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 4d949cfbe6..1aee1d4965 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -6091,3 +6091,69 @@ async def test_list_available_teams_returns_empty_list_when_none_configured(): assert result == [] litellm.default_internal_user_params = original + + +@pytest.mark.asyncio +async def test_list_team_v1_batches_key_queries(): + """ + Test that list_team fetches all keys in a single batched query + instead of issuing one query per team (N+1). + """ + from unittest.mock import AsyncMock, MagicMock, Mock, patch + + from fastapi import Request + + from litellm.proxy._types import ( + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LitellmUserRoles, + TeamListResponseObject, + UserAPIKeyAuth, + ) + from litellm.proxy.management_endpoints.team_endpoints import list_team + + mock_request = Mock(spec=Request) + + mock_user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="admin_user", + ) + + # Two teams + team1 = LiteLLM_TeamTable(team_id="team-1", team_alias="Team One") + team2 = LiteLLM_TeamTable(team_id="team-2", team_alias="Team Two") + + # Mock keys belonging to different teams + key1 = MagicMock() + key1.team_id = "team-1" + key2 = MagicMock() + key2.team_id = "team-1" + key3 = MagicMock() + key3.team_id = "team-2" + + with patch( + "litellm.proxy.proxy_server.prisma_client" + ) as mock_prisma_client, patch( + "litellm.proxy.management_endpoints.team_endpoints._authorize_and_filter_teams", + new_callable=AsyncMock, + return_value=[team1, team2], + ), patch( + "litellm.proxy.management_endpoints.team_endpoints.get_all_team_memberships", + new_callable=AsyncMock, + return_value=[], + ): + mock_find_many = AsyncMock(return_value=[key1, key2, key3]) + mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many + + result = await list_team( + http_request=mock_request, + user_api_key_dict=mock_user_api_key_dict, + ) + + # Verify keys are correctly distributed + assert len(result) == 2 + # Results are sorted by team_alias + assert result[0].team_id == "team-1" + assert result[0].keys == [key1, key2] + assert result[1].team_id == "team-2" + assert result[1].keys == [key3] diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 50459bf18a..ea68e8566a 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -2410,12 +2410,13 @@ def test_mapped_pass_through_routes_with_server_root_path(): ) + @pytest.mark.asyncio async def test_multipart_passthrough_preserves_boundary(): """ Test that multipart/form-data requests through passthrough preserve the boundary and can be correctly parsed by the upstream server. - + Regression test for multipart boundary stripping issue. """ from io import BytesIO @@ -2426,41 +2427,41 @@ async def test_multipart_passthrough_preserves_boundary(): mock_response.headers = httpx.Headers({"content-type": "application/json"}) mock_response.aread = AsyncMock(return_value=b'{"filename": "test.txt", "size": 17}') mock_response.text = '{"filename": "test.txt", "size": 17}' - + async def mock_httpx_request(method, url, **kwargs): # Verify that files parameter is passed (not json) assert "files" in kwargs, "Files should be passed for multipart requests" assert "file" in kwargs["files"], "File field should be in files dict" - + # Verify content-type is NOT in headers (httpx will set it with correct boundary) headers = kwargs.get("headers", {}) assert "content-type" not in headers, "content-type should be removed for multipart" - + filename, content, content_type = kwargs["files"]["file"] assert filename == "test.txt" assert content == b"test file content" assert content_type == "text/plain" - + return mock_response - + async_client = MagicMock() async_client.request = AsyncMock(side_effect=mock_httpx_request) - + # Create mock request request = MagicMock(spec=Request) request.method = "POST" request.headers = Headers({"content-type": "multipart/form-data; boundary=test123"}) - + # Mock form data file_content = b"test file content" file = BytesIO(file_content) headers = Headers({"content-type": "text/plain"}) upload_file = UploadFile(file=file, filename="test.txt", headers=headers) upload_file.read = AsyncMock(return_value=file_content) - + form_data = {"file": upload_file} request.form = AsyncMock(return_value=form_data) - + # Test the multipart handler directly response = await HttpPassThroughEndpointHelpers.make_multipart_http_request( request=request, @@ -2469,7 +2470,7 @@ async def test_multipart_passthrough_preserves_boundary(): headers={}, requested_query_params=None, ) - + # Verify the response assert response.status_code == 200 async_client.request.assert_called_once() diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 7174253538..30b952cd42 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -1072,9 +1072,10 @@ def test_spend_logs_redacts_request_and_response_when_turn_off_message_logging_e response_result = _get_response_for_spend_logs_payload(payload=payload, kwargs=kwargs) # When redaction is enabled and response is a dict (not ModelResponse), - # perform_redaction returns {"text": "redacted-by-litellm"} + # perform_redaction redacts content in-place within the choices structure parsed_response = json.loads(response_result) - assert parsed_response == {"text": "redacted-by-litellm"} + assert parsed_response["choices"][0]["message"]["content"] == "redacted-by-litellm" + assert parsed_response["choices"][0]["message"]["role"] == "assistant" @patch("litellm.secret_managers.main.get_secret_bool") diff --git a/tests/test_litellm/proxy/test_openapi_schema_validation.py b/tests/test_litellm/proxy/test_openapi_schema_validation.py new file mode 100644 index 0000000000..aafe08f303 --- /dev/null +++ b/tests/test_litellm/proxy/test_openapi_schema_validation.py @@ -0,0 +1,142 @@ +""" +Test that the OpenAPI schema generated by FastAPI is valid for specific endpoints. + +Validates fixes for: +- /spend/calculate response schema (must use proper OpenAPI 3.x content wrapper) +- /credentials/by_model/{model_id} path parameter (must not leak credential_name) + +Related issue: https://github.com/BerriAI/litellm/issues/21305 +""" + +import pytest + + +class TestSpendCalculateOpenAPISchema: + """Test /spend/calculate response schema is valid OpenAPI 3.x.""" + + def test_response_schema_has_description(self): + """The 200 response must have a 'description' field per OpenAPI 3.x spec.""" + from litellm.proxy.spend_tracking.spend_management_endpoints import router + + for route in router.routes: + if hasattr(route, "path") and route.path == "/spend/calculate": + responses = route.responses or {} + response_200 = responses.get(200, {}) + assert "description" in response_200, ( + "/spend/calculate 200 response must have a 'description' field" + ) + break + else: + pytest.fail("/spend/calculate route not found in router") + + def test_response_schema_has_content_wrapper(self): + """The 200 response must use 'content' wrapper, not bare properties.""" + from litellm.proxy.spend_tracking.spend_management_endpoints import router + + for route in router.routes: + if hasattr(route, "path") and route.path == "/spend/calculate": + responses = route.responses or {} + response_200 = responses.get(200, {}) + # Must NOT have 'cost' as a top-level key (invalid OpenAPI) + assert "cost" not in response_200, ( + "/spend/calculate 200 response must not have 'cost' as a " + "top-level property - use 'content' wrapper instead" + ) + # Must have 'content' wrapper + assert "content" in response_200, ( + "/spend/calculate 200 response must have a 'content' field" + ) + content = response_200["content"] + assert "application/json" in content + assert "schema" in content["application/json"] + break + else: + pytest.fail("/spend/calculate route not found in router") + + +class TestCredentialEndpointsOpenAPISchema: + """Test /credentials endpoints have correct path parameters.""" + + def test_by_name_and_by_model_are_separate_handlers(self): + """ + /credentials/by_name/{credential_name} and /credentials/by_model/{model_id} + must be separate handler functions so each only declares its own path params. + """ + from litellm.proxy.credential_endpoints.endpoints import router + + by_name_routes = [] + by_model_routes = [] + for route in router.routes: + if not hasattr(route, "path"): + continue + if "by_name" in route.path: + by_name_routes.append(route) + elif "by_model" in route.path: + by_model_routes.append(route) + + assert len(by_name_routes) == 1, "Expected exactly one by_name route" + assert len(by_model_routes) == 1, "Expected exactly one by_model route" + + # They must be different endpoint functions + by_name_endpoint = by_name_routes[0].endpoint + by_model_endpoint = by_model_routes[0].endpoint + assert by_name_endpoint is not by_model_endpoint, ( + "by_name and by_model must be separate handler functions " + "to avoid path parameter conflicts in OpenAPI spec" + ) + + def test_by_model_route_does_not_require_credential_name(self): + """ + The /credentials/by_model/{model_id} route must NOT have + credential_name as a parameter. + """ + import inspect + from litellm.proxy.credential_endpoints.endpoints import ( + get_credential_by_model, + ) + + sig = inspect.signature(get_credential_by_model) + param_names = list(sig.parameters.keys()) + assert "credential_name" not in param_names, ( + "get_credential_by_model must not have a credential_name parameter" + ) + + def test_by_name_route_does_not_require_model_id(self): + """ + The /credentials/by_name/{credential_name} route must NOT have + model_id as a parameter. + """ + import inspect + from litellm.proxy.credential_endpoints.endpoints import ( + get_credential_by_name, + ) + + sig = inspect.signature(get_credential_by_name) + param_names = list(sig.parameters.keys()) + assert "model_id" not in param_names, ( + "get_credential_by_name must not have a model_id parameter" + ) + + def test_by_model_has_model_id_path_param(self): + """The by_model handler must accept model_id as a path parameter.""" + import inspect + from litellm.proxy.credential_endpoints.endpoints import ( + get_credential_by_model, + ) + + sig = inspect.signature(get_credential_by_model) + assert "model_id" in sig.parameters, ( + "get_credential_by_model must have a model_id parameter" + ) + + def test_by_name_has_credential_name_path_param(self): + """The by_name handler must accept credential_name as a path parameter.""" + import inspect + from litellm.proxy.credential_endpoints.endpoints import ( + get_credential_by_name, + ) + + sig = inspect.signature(get_credential_by_name) + assert "credential_name" in sig.parameters, ( + "get_credential_by_name must have a credential_name parameter" + ) diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 6d6162437c..a931a9bc93 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -1774,3 +1774,128 @@ class TestStreamingIDConsistency: # Verify it matches the cached ID assert iterator._cached_item_id is not None assert iterator._cached_item_id == text_done_id + + def test_parallel_tool_calls_merged_into_single_assistant_message(self): + """ + Regression test: multi-turn parallel tool calls via the Responses API must + produce a single assistant message with all tool_calls, not one assistant + message per function_call item. + + When the model responds with two parallel tool calls (e.g. get_weather for + SF and NYC), the next Responses API request includes two consecutive + function_call items followed by two function_call_output items. + + Without the fix each function_call becomes its own assistant message, + producing back-to-back assistant messages that Anthropic/Vertex AI rejects: + "tool_use ids were found without tool_result blocks immediately after". + """ + input_items = [ + {"type": "message", "role": "user", "content": "Weather in SF and NYC?"}, + # Two parallel tool calls from the previous assistant response + { + "type": "function_call", + "call_id": "toolu_01", + "name": "get_weather", + "arguments": '{"city": "SF"}', + }, + { + "type": "function_call", + "call_id": "toolu_02", + "name": "get_weather", + "arguments": '{"city": "NYC"}', + }, + # Tool results + {"type": "function_call_output", "call_id": "toolu_01", "output": "72°F"}, + {"type": "function_call_output", "call_id": "toolu_02", "output": "55°F"}, + ] + + messages = LiteLLMCompletionResponsesConfig._transform_response_input_param_to_chat_completion_message( + input=input_items + ) + + roles = [ + m.get("role") if isinstance(m, dict) else getattr(m, "role", None) + for m in messages + ] + + # Must not have two consecutive assistant messages + for i in range(len(roles) - 1): + assert not ( + roles[i] == "assistant" and roles[i + 1] == "assistant" + ), f"Consecutive assistant messages at indices {i} and {i+1}: {roles}" + + # The single assistant message must contain BOTH tool_calls + assistant_messages = [ + m for m in messages + if (m.get("role") if isinstance(m, dict) else getattr(m, "role", None)) + == "assistant" + ] + assert len(assistant_messages) == 1, ( + f"Expected 1 assistant message, got {len(assistant_messages)}" + ) + + assistant_msg = assistant_messages[0] + tool_calls = ( + assistant_msg.get("tool_calls") + if isinstance(assistant_msg, dict) + else getattr(assistant_msg, "tool_calls", None) + ) + assert tool_calls is not None and len(tool_calls) == 2, ( + f"Expected 2 tool_calls in the merged assistant message, got: {tool_calls}" + ) + + call_ids = [ + (tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None)) + for tc in tool_calls + ] + assert "toolu_01" in call_ids, f"toolu_01 missing from tool_calls: {call_ids}" + assert "toolu_02" in call_ids, f"toolu_02 missing from tool_calls: {call_ids}" + + # Both tool messages must be present + tool_messages = [ + m for m in messages + if (m.get("role") if isinstance(m, dict) else getattr(m, "role", None)) + == "tool" + ] + assert len(tool_messages) == 2, ( + f"Expected 2 tool messages, got {len(tool_messages)}" + ) + + def test_single_tool_call_still_works_after_merge_fix(self): + """ + Ensure the parallel-tool-call merging fix does not break the existing + single-tool-call path. + """ + input_items = [ + {"type": "message", "role": "user", "content": "Weather in SF?"}, + { + "type": "function_call", + "call_id": "toolu_01", + "name": "get_weather", + "arguments": '{"city": "SF"}', + }, + {"type": "function_call_output", "call_id": "toolu_01", "output": "72°F"}, + ] + + messages = LiteLLMCompletionResponsesConfig._transform_response_input_param_to_chat_completion_message( + input=input_items + ) + + roles = [ + m.get("role") if isinstance(m, dict) else getattr(m, "role", None) + for m in messages + ] + + assert "user" in roles + assert "assistant" in roles + assert "tool" in roles + + assistant_messages = [m for m in messages if (m.get("role") if isinstance(m, dict) else getattr(m, "role", None)) == "assistant"] + assert len(assistant_messages) == 1 + + tool_calls = ( + assistant_messages[0].get("tool_calls") + if isinstance(assistant_messages[0], dict) + else getattr(assistant_messages[0], "tool_calls", None) + ) + assert tool_calls is not None and len(tool_calls) == 1 diff --git a/tests/test_litellm/test_anthropic_skills_transformation.py b/tests/test_litellm/test_anthropic_skills_transformation.py new file mode 100644 index 0000000000..a70b54984e --- /dev/null +++ b/tests/test_litellm/test_anthropic_skills_transformation.py @@ -0,0 +1,338 @@ +""" +Unit tests for Anthropic Skills API request/response transformation. + +These tests validate URL construction, header generation, request payload +building, and response parsing without requiring a live Anthropic API key +or beta access to the Skills API. +""" +from unittest.mock import MagicMock, patch + +import httpx +import pytest + +from litellm.constants import ANTHROPIC_SKILLS_API_BETA_VERSION +from litellm.llms.anthropic.skills.transformation import AnthropicSkillsConfig +from litellm.types.llms.anthropic_skills import ( + CreateSkillRequest, + DeleteSkillResponse, + ListSkillsParams, + ListSkillsResponse, + Skill, +) +from litellm.types.router import GenericLiteLLMParams + + +FAKE_API_KEY = "sk-ant-test-key-1234" +FAKE_API_BASE = "https://api.anthropic.com" + + +def _make_mock_response( + json_data: dict, status_code: int = 200, method: str = "POST" +) -> httpx.Response: + return httpx.Response( + status_code=status_code, + json=json_data, + request=httpx.Request(method, "https://api.anthropic.com/v1/skills"), + ) + + +def _make_skill_payload(**kwargs) -> dict: + defaults = { + "id": "skill_abc123", + "created_at": "2025-10-15T12:00:00Z", + "updated_at": "2025-10-15T12:00:00Z", + "source": "custom", + "type": "skill", + "display_title": "Test Skill", + "latest_version": "v1", + } + defaults.update(kwargs) + return defaults + + +class TestAnthropicSkillsConfigURLConstruction: + def setup_method(self): + self.config = AnthropicSkillsConfig() + + def test_url_without_skill_id(self): + url = self.config.get_complete_url( + api_base=FAKE_API_BASE, + endpoint="skills", + ) + assert url == f"{FAKE_API_BASE}/v1/skills" + + def test_url_with_skill_id(self): + url = self.config.get_complete_url( + api_base=FAKE_API_BASE, + endpoint="skills", + skill_id="skill_abc123", + ) + assert url == f"{FAKE_API_BASE}/v1/skills/skill_abc123" + + def test_url_falls_back_to_anthropic_default(self): + with patch( + "litellm.llms.anthropic.common_utils.AnthropicModelInfo.get_api_base", + return_value="https://api.anthropic.com", + ): + url = self.config.get_complete_url( + api_base=None, + endpoint="skills", + ) + assert url == "https://api.anthropic.com/v1/skills" + + def test_url_with_custom_api_base(self): + custom_base = "https://my-proxy.example.com" + url = self.config.get_complete_url( + api_base=custom_base, + endpoint="skills", + ) + assert url == f"{custom_base}/v1/skills" + + +class TestAnthropicSkillsConfigHeaderValidation: + def setup_method(self): + self.config = AnthropicSkillsConfig() + + def _make_litellm_params(self, api_key=FAKE_API_KEY): + return GenericLiteLLMParams(api_key=api_key) + + def test_sets_api_key_header(self): + with patch( + "litellm.llms.anthropic.common_utils.AnthropicModelInfo.get_api_key", + return_value=FAKE_API_KEY, + ): + headers = self.config.validate_environment( + headers={}, litellm_params=self._make_litellm_params() + ) + assert headers["x-api-key"] == FAKE_API_KEY + + def test_sets_anthropic_version_header(self): + with patch( + "litellm.llms.anthropic.common_utils.AnthropicModelInfo.get_api_key", + return_value=FAKE_API_KEY, + ): + headers = self.config.validate_environment( + headers={}, litellm_params=self._make_litellm_params() + ) + assert headers["anthropic-version"] == "2023-06-01" + + def test_sets_skills_beta_header(self): + with patch( + "litellm.llms.anthropic.common_utils.AnthropicModelInfo.get_api_key", + return_value=FAKE_API_KEY, + ): + headers = self.config.validate_environment( + headers={}, litellm_params=self._make_litellm_params() + ) + assert headers["anthropic-beta"] == ANTHROPIC_SKILLS_API_BETA_VERSION + + def test_merges_existing_beta_header_string(self): + with patch( + "litellm.llms.anthropic.common_utils.AnthropicModelInfo.get_api_key", + return_value=FAKE_API_KEY, + ): + headers = self.config.validate_environment( + headers={"anthropic-beta": "other-beta-2024-01-01"}, + litellm_params=self._make_litellm_params(), + ) + assert isinstance(headers["anthropic-beta"], list) + assert "other-beta-2024-01-01" in headers["anthropic-beta"] + assert ANTHROPIC_SKILLS_API_BETA_VERSION in headers["anthropic-beta"] + + def test_merges_existing_beta_header_list(self): + with patch( + "litellm.llms.anthropic.common_utils.AnthropicModelInfo.get_api_key", + return_value=FAKE_API_KEY, + ): + headers = self.config.validate_environment( + headers={"anthropic-beta": ["other-beta-2024-01-01"]}, + litellm_params=self._make_litellm_params(), + ) + assert ANTHROPIC_SKILLS_API_BETA_VERSION in headers["anthropic-beta"] + assert "other-beta-2024-01-01" in headers["anthropic-beta"] + + def test_does_not_duplicate_beta_header(self): + with patch( + "litellm.llms.anthropic.common_utils.AnthropicModelInfo.get_api_key", + return_value=FAKE_API_KEY, + ): + headers = self.config.validate_environment( + headers={"anthropic-beta": ANTHROPIC_SKILLS_API_BETA_VERSION}, + litellm_params=self._make_litellm_params(), + ) + beta = headers["anthropic-beta"] + if isinstance(beta, list): + assert beta.count(ANTHROPIC_SKILLS_API_BETA_VERSION) == 1 + else: + assert beta == ANTHROPIC_SKILLS_API_BETA_VERSION + + def test_raises_without_api_key(self): + with patch( + "litellm.llms.anthropic.common_utils.AnthropicModelInfo.get_api_key", + return_value=None, + ): + with pytest.raises(ValueError, match="ANTHROPIC_API_KEY"): + self.config.validate_environment( + headers={}, litellm_params=self._make_litellm_params(api_key=None) + ) + + +class TestAnthropicSkillsConfigCreateRequestTransformation: + def setup_method(self): + self.config = AnthropicSkillsConfig() + self.litellm_params = GenericLiteLLMParams(api_key=FAKE_API_KEY) + + def test_display_title_included(self): + create_request: CreateSkillRequest = {"display_title": "My Skill"} + body = self.config.transform_create_skill_request( + create_request=create_request, + litellm_params=self.litellm_params, + headers={}, + ) + assert body["display_title"] == "My Skill" + + def test_none_values_excluded(self): + create_request: CreateSkillRequest = {"display_title": None, "files": None} + body = self.config.transform_create_skill_request( + create_request=create_request, + litellm_params=self.litellm_params, + headers={}, + ) + assert "display_title" not in body + assert "files" not in body + + def test_empty_request_produces_empty_body(self): + create_request: CreateSkillRequest = {} + body = self.config.transform_create_skill_request( + create_request=create_request, + litellm_params=self.litellm_params, + headers={}, + ) + assert body == {} + + +class TestAnthropicSkillsConfigListRequestTransformation: + def setup_method(self): + self.config = AnthropicSkillsConfig() + self.litellm_params = GenericLiteLLMParams(api_key=FAKE_API_KEY) + + def test_limit_included_in_query_params(self): + list_params: ListSkillsParams = {"limit": 25} + with patch( + "litellm.llms.anthropic.common_utils.AnthropicModelInfo.get_api_base", + return_value=FAKE_API_BASE, + ): + url, query_params = self.config.transform_list_skills_request( + list_params=list_params, + litellm_params=self.litellm_params, + headers={}, + ) + assert query_params["limit"] == 25 + assert url == f"{FAKE_API_BASE}/v1/skills" + + def test_source_filter_included(self): + list_params: ListSkillsParams = {"source": "custom"} + with patch( + "litellm.llms.anthropic.common_utils.AnthropicModelInfo.get_api_base", + return_value=FAKE_API_BASE, + ): + _, query_params = self.config.transform_list_skills_request( + list_params=list_params, + litellm_params=self.litellm_params, + headers={}, + ) + assert query_params["source"] == "custom" + + def test_empty_params_produce_empty_query(self): + list_params: ListSkillsParams = {} + with patch( + "litellm.llms.anthropic.common_utils.AnthropicModelInfo.get_api_base", + return_value=FAKE_API_BASE, + ): + _, query_params = self.config.transform_list_skills_request( + list_params=list_params, + litellm_params=self.litellm_params, + headers={}, + ) + assert query_params == {} + + +class TestAnthropicSkillsConfigResponseTransformation: + def setup_method(self): + self.config = AnthropicSkillsConfig() + self.logging_obj = MagicMock() + + def test_create_skill_response_parses_skill(self): + payload = _make_skill_payload() + raw = _make_mock_response(payload) + skill = self.config.transform_create_skill_response( + raw_response=raw, logging_obj=self.logging_obj + ) + assert isinstance(skill, Skill) + assert skill.id == "skill_abc123" + assert skill.source == "custom" + assert skill.display_title == "Test Skill" + + def test_get_skill_response_parses_skill(self): + payload = _make_skill_payload(id="skill_xyz", display_title="Another") + raw = _make_mock_response(payload, method="GET") + skill = self.config.transform_get_skill_response( + raw_response=raw, logging_obj=self.logging_obj + ) + assert isinstance(skill, Skill) + assert skill.id == "skill_xyz" + assert skill.display_title == "Another" + + def test_list_skills_response_parses_list(self): + payload = { + "data": [_make_skill_payload(), _make_skill_payload(id="skill_def456")], + "has_more": False, + "next_page": None, + } + raw = _make_mock_response(payload, method="GET") + result = self.config.transform_list_skills_response( + raw_response=raw, logging_obj=self.logging_obj + ) + assert isinstance(result, ListSkillsResponse) + assert len(result.data) == 2 + assert result.data[0].id == "skill_abc123" + assert result.data[1].id == "skill_def456" + assert result.has_more is False + + def test_list_skills_response_with_pagination(self): + payload = { + "data": [_make_skill_payload()], + "has_more": True, + "next_page": "page_token_xyz", + } + raw = _make_mock_response(payload, method="GET") + result = self.config.transform_list_skills_response( + raw_response=raw, logging_obj=self.logging_obj + ) + assert result.has_more is True + assert result.next_page == "page_token_xyz" + + def test_delete_skill_response_parses_correctly(self): + payload = {"id": "skill_abc123", "type": "skill_deleted"} + raw = _make_mock_response(payload, method="DELETE") + result = self.config.transform_delete_skill_response( + raw_response=raw, logging_obj=self.logging_obj + ) + assert isinstance(result, DeleteSkillResponse) + assert result.id == "skill_abc123" + assert result.type == "skill_deleted" + + def test_skill_response_optional_fields_default(self): + payload = { + "id": "skill_minimal", + "created_at": "2025-10-15T12:00:00Z", + "updated_at": "2025-10-15T12:00:00Z", + "source": "anthropic", + "type": "skill", + } + raw = _make_mock_response(payload) + skill = self.config.transform_create_skill_response( + raw_response=raw, logging_obj=self.logging_obj + ) + assert skill.display_title is None + assert skill.latest_version is None diff --git a/tests/test_litellm/test_router_retry_non_retryable_errors.py b/tests/test_litellm/test_router_retry_non_retryable_errors.py new file mode 100644 index 0000000000..20a1c979a0 --- /dev/null +++ b/tests/test_litellm/test_router_retry_non_retryable_errors.py @@ -0,0 +1,251 @@ +""" +Test that the Router retry loop correctly handles non-retryable errors. + +Verifies that: +1. Non-retryable errors (e.g., 400 ContextWindowExceeded) inside the retry loop + break out immediately instead of being swallowed. +2. original_exception is updated to the latest error, not stuck on the first. +3. Retryable errors (e.g., 429 RateLimitError) still retry normally. + +Regression tests for https://github.com/BerriAI/litellm/issues/21343 +""" + +from unittest.mock import AsyncMock, patch + +import pytest + +import litellm +from litellm import Router + + +def _make_rate_limit_error(message="Rate limited"): + """Create a RateLimitError for testing.""" + return litellm.RateLimitError( + message=message, + llm_provider="bedrock", + model="anthropic.claude-v2", + ) + + +def _make_context_window_error(message="prompt is too long: 1205821 tokens > 200000"): + """Create a ContextWindowExceededError for testing.""" + return litellm.ContextWindowExceededError( + message=message, + llm_provider="vertex_ai", + model="claude-3-opus", + ) + + +def _make_bad_request_error(message="Invalid request"): + """Create a BadRequestError for testing.""" + return litellm.BadRequestError( + message=message, + llm_provider="openai", + model="gpt-4", + ) + + +def _make_not_found_error(message="Model not found"): + """Create a NotFoundError for testing.""" + return litellm.NotFoundError( + message=message, + llm_provider="openai", + model="gpt-99", + ) + + +def _create_router(num_retries=2): + """Create a Router with two deployments for testing.""" + return Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "openai/gpt-4", + "api_key": "fake-key-1", + }, + }, + { + "model_name": "test-model", + "litellm_params": { + "model": "openai/gpt-4", + "api_key": "fake-key-2", + }, + }, + ], + num_retries=num_retries, + ) + + +def _base_kwargs(): + """Return kwargs required by async_function_with_retries.""" + return { + "model": "test-model", + "messages": [{"role": "user", "content": "test"}], + "original_function": AsyncMock(), + "metadata": {}, + } + + +@pytest.mark.asyncio +async def test_non_retryable_error_in_retry_loop_raises_immediately(): + """ + When a non-retryable error (400 ContextWindowExceeded) occurs inside the + retry loop, the router should raise it immediately instead of swallowing it + and raising the original error. + + Scenario: First call -> 429, Retry -> 400 (non-retryable) + Expected: ContextWindowExceededError is raised, NOT RateLimitError + """ + router = _create_router(num_retries=2) + + rate_limit_error = _make_rate_limit_error() + context_window_error = _make_context_window_error() + + call_count = 0 + + async def mock_make_call(*args, **kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + raise rate_limit_error + else: + raise context_window_error + + with patch.object(router, "make_call", side_effect=mock_make_call), \ + patch.object(router, "_async_get_healthy_deployments", + return_value=(["d1", "d2"], ["d1", "d2"])), \ + patch.object(router, "_time_to_sleep_before_retry", return_value=0), \ + patch.object(router, "log_retry", side_effect=lambda kwargs, e: kwargs): + with pytest.raises(litellm.ContextWindowExceededError): + await router.async_function_with_retries( + num_retries=2, + **_base_kwargs(), + ) + + +@pytest.mark.asyncio +async def test_bad_request_error_in_retry_loop_raises_immediately(): + """ + A generic 400 BadRequestError inside the retry loop should also break out + immediately since 400 is not retryable. + """ + router = _create_router(num_retries=2) + + rate_limit_error = _make_rate_limit_error() + bad_request_error = _make_bad_request_error() + + call_count = 0 + + async def mock_make_call(*args, **kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + raise rate_limit_error + else: + raise bad_request_error + + with patch.object(router, "make_call", side_effect=mock_make_call), \ + patch.object(router, "_async_get_healthy_deployments", + return_value=(["d1", "d2"], ["d1", "d2"])), \ + patch.object(router, "_time_to_sleep_before_retry", return_value=0), \ + patch.object(router, "log_retry", side_effect=lambda kwargs, e: kwargs): + with pytest.raises(litellm.BadRequestError): + await router.async_function_with_retries( + num_retries=2, + **_base_kwargs(), + ) + + +@pytest.mark.asyncio +async def test_original_exception_updated_to_latest_error(): + """ + When all retries are exhausted with retryable errors, the LAST error + should be raised, not the first one. + """ + router = _create_router(num_retries=2) + + call_count = 0 + + async def mock_make_call(*args, **kwargs): + nonlocal call_count + call_count += 1 + raise _make_rate_limit_error(f"Rate limit attempt {call_count}") + + with patch.object(router, "make_call", side_effect=mock_make_call), \ + patch.object(router, "_async_get_healthy_deployments", + return_value=(["d1", "d2"], ["d1", "d2"])), \ + patch.object(router, "_time_to_sleep_before_retry", return_value=0), \ + patch.object(router, "log_retry", side_effect=lambda kwargs, e: kwargs): + with pytest.raises(litellm.RateLimitError) as exc_info: + await router.async_function_with_retries( + num_retries=2, + **_base_kwargs(), + ) + # Should be the LAST error, not the first + assert "Rate limit attempt 3" in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_retryable_errors_still_retry_normally(): + """ + Retryable errors (429 RateLimitError) should still be retried the + configured number of times before raising. + """ + router = _create_router(num_retries=3) + + call_count = 0 + + async def mock_make_call(*args, **kwargs): + nonlocal call_count + call_count += 1 + raise _make_rate_limit_error(f"Rate limit attempt {call_count}") + + with patch.object(router, "make_call", side_effect=mock_make_call), \ + patch.object(router, "_async_get_healthy_deployments", + return_value=(["d1", "d2"], ["d1", "d2"])), \ + patch.object(router, "_time_to_sleep_before_retry", return_value=0), \ + patch.object(router, "log_retry", side_effect=lambda kwargs, e: kwargs): + with pytest.raises(litellm.RateLimitError): + await router.async_function_with_retries( + num_retries=3, + **_base_kwargs(), + ) + + # Initial call + 3 retries = 4 total calls + assert call_count == 4 + + +@pytest.mark.asyncio +async def test_not_found_error_in_retry_loop_raises_immediately(): + """ + A 404 NotFoundError inside the retry loop should break out immediately. + """ + router = _create_router(num_retries=2) + + rate_limit_error = _make_rate_limit_error() + not_found_error = _make_not_found_error() + + call_count = 0 + + async def mock_make_call(*args, **kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + raise rate_limit_error + else: + raise not_found_error + + with patch.object(router, "make_call", side_effect=mock_make_call), \ + patch.object(router, "_async_get_healthy_deployments", + return_value=(["d1", "d2"], ["d1", "d2"])), \ + patch.object(router, "_time_to_sleep_before_retry", return_value=0), \ + patch.object(router, "log_retry", side_effect=lambda kwargs, e: kwargs): + with pytest.raises(litellm.NotFoundError): + await router.async_function_with_retries( + num_retries=2, + **_base_kwargs(), + ) + + # Only 2 calls: initial + first retry that hits non-retryable + assert call_count == 2 diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 3f7fbdc9dc..64488e2fb6 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -444,9 +444,6 @@ def test_anthropic_web_search_in_model_info(): supported_models = [ "anthropic/claude-4-sonnet-20250514", "anthropic/claude-sonnet-4-5-20250929", - "anthropic/claude-3-5-sonnet-20241022", - "anthropic/claude-3-5-haiku-20241022", - "anthropic/claude-3-5-haiku-latest", ] for model in supported_models: from litellm.utils import get_model_info diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index d63154f0ce..e6d144df65 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -275,8 +275,8 @@ it("should display user email correctly", async () => { }); }); -it("should show skeleton loaders when isLoading is true", () => { - // Mock loading state +it("should show loading message only on initial load (isPending)", () => { + // Mock initial loading state mockUseKeys.mockReturnValue({ data: null, isPending: true, @@ -296,7 +296,7 @@ it("should show skeleton loaders when isLoading is true", () => { renderWithProviders(); - // Check that loading message is shown + // Check that loading message is shown on initial load expect(screen.getByText("🚅 Loading keys...")).toBeInTheDocument(); // Check that actual key data is not shown @@ -898,3 +898,79 @@ describe("pagination display – total count and page count", () => { }); }); }); + +describe("refetch button", () => { + it("should show Fetch button in normal state", () => { + renderWithProviders(); + + const fetchButton = screen.getByTitle("Fetch data"); + expect(fetchButton).toBeInTheDocument(); + expect(fetchButton).not.toBeDisabled(); + expect(screen.getByText("Fetch")).toBeInTheDocument(); + }); + + it("should show Fetching state and keep table data visible during refetch", () => { + mockUseKeys.mockReturnValue({ + data: { + keys: [mockKey], + total_count: 1, + current_page: 1, + total_pages: 1, + } as KeysResponse, + isPending: false, + isFetching: true, + refetch: vi.fn(), + } as any); + + renderWithProviders(); + + // Button should show "Fetching" and be disabled + expect(screen.getByText("Fetching")).toBeInTheDocument(); + const fetchButton = screen.getByTitle("Fetch data"); + expect(fetchButton).toBeDisabled(); + + // Table data should still be visible (stale data) + expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); + + // "Loading keys..." should NOT appear during refetch + expect(screen.queryByText("🚅 Loading keys...")).not.toBeInTheDocument(); + }); + + it("should call refetch when Fetch button is clicked", () => { + const mockRefetch = vi.fn(); + mockUseKeys.mockReturnValue({ + data: { + keys: [mockKey], + total_count: 1, + current_page: 1, + total_pages: 1, + } as KeysResponse, + isPending: false, + isFetching: false, + refetch: mockRefetch, + } as any); + + renderWithProviders(); + + const fetchButton = screen.getByTitle("Fetch data"); + fireEvent.click(fetchButton); + + expect(mockRefetch).toHaveBeenCalledTimes(1); + }); + + it("should show Fetch button enabled on error so user can retry", () => { + mockUseKeys.mockReturnValue({ + data: null, + isPending: false, + isFetching: false, + isError: true, + refetch: vi.fn(), + } as any); + + renderWithProviders(); + + const fetchButton = screen.getByTitle("Fetch data"); + expect(fetchButton).not.toBeDisabled(); + expect(screen.getByText("Fetch")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index 6091794170..b30d4b6ce5 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -85,6 +85,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo data: keys, isPending: isLoading, isFetching, + isError, refetch, } = useKeys(tablePagination.pageIndex + 1, tablePagination.pageSize, { sortBy: sortBy || undefined, @@ -102,6 +103,15 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo organizations, }); + // Defer the transition so the button stays in loading state until the table + // has rendered with the new data (mirrors the spend-logs pattern) + const isFetchingDeferred = useDeferredValue(isFetching); + const isButtonLoading = (isFetching || isFetchingDeferred) && !isError; + + const handleRefresh = () => { + refetch(); + }; + const totalCount = filteredTotalCount ?? keys?.total_count ?? 0; // Add a useEffect to call refresh when a key is created @@ -684,16 +694,28 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
- {isLoading || isFetching ? ( - - ) : ( - - Showing {rangeLabel} of {totalCount} results - - )} +
+ {isLoading ? ( + + ) : ( + + Showing {rangeLabel} of {totalCount} results + + )} + + } + onClick={handleRefresh} + disabled={isButtonLoading} + title="Fetch data" + > + {isButtonLoading ? "Fetching" : "Fetch"} + +
- {isLoading || isFetching ? ( + {isLoading ? ( ) : ( @@ -701,24 +723,24 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo )} - {isLoading || isFetching ? ( + {isLoading ? ( ) : ( )} - {isLoading || isFetching ? ( + {isLoading ? ( ) : (