diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index cfa0390e836..70a50d7f06e 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -4,7 +4,7 @@ /ui/nginx.conf /ui/litellm-dashboard/src/lib/http/schema.d.ts /ui/litellm-dashboard/tsconfig.tsbuildinfo -/model_prices_and_context_window.json @mateo-berri -/litellm/model_prices_and_context_window_backup.json @mateo-berri +/model_prices_and_context_window.json @mateo-berri @ryan-crabbe-berri @kerry-berri +/litellm/model_prices_and_context_window_backup.json @mateo-berri @ryan-crabbe-berri @kerry-berri /litellm-proxy-extras/litellm_proxy_extras/migrations/ @yuneng-berri @ryan-crabbe-berri /.github/CODEOWNERS @yuneng-berri diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 6036f94ed86..bc26e5d162e 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -70,7 +70,7 @@ env: jobs: rust-lint: runs-on: ubuntu-latest - timeout-minutes: 10 + timeout-minutes: 15 defaults: run: working-directory: litellm-rust @@ -81,17 +81,12 @@ jobs: - run: rustup toolchain install --no-self-update - - run: cargo fmt --check + - run: cargo fmt --all --check - - uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 + - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 with: - path: | - ~/.cargo/registry - ~/.cargo/git - litellm-rust/target - key: ${{ runner.os }}-cargo-${{ github.job }}-${{ hashFiles('rust-toolchain.toml', '.cargo/**', 'litellm-rust/Cargo.lock') }} - restore-keys: | - ${{ runner.os }}-cargo-${{ github.job }}- + workspaces: litellm-rust + cache-on-failure: true - run: cargo clippy --workspace --all-targets --locked -- -D warnings @@ -99,8 +94,37 @@ jobs: rust-test: runs-on: ubuntu-latest - timeout-minutes: 30 + timeout-minutes: 20 + defaults: + run: + working-directory: litellm-rust + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + + - run: rustup toolchain install --no-self-update + + - uses: taiki-e/install-action@d438492cf8a250514fa2d34b30bc3c0dc37c65ff # v2.87.8 + with: + tool: cargo-nextest@0.9.143 + + - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 + with: + workspaces: litellm-rust + cache-on-failure: true + + - run: cargo nextest run --workspace --locked + + - run: cargo test --workspace --doc --locked + + rust-wheel: + runs-on: ubuntu-latest + timeout-minutes: 30 steps: - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 with: @@ -116,18 +140,10 @@ jobs: - run: rustup toolchain install --no-self-update - - uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 + - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 with: - path: | - ~/.cargo/registry - ~/.cargo/git - litellm-rust/target - key: ${{ runner.os }}-cargo-${{ github.job }}-${{ hashFiles('rust-toolchain.toml', '.cargo/**', 'litellm-rust/Cargo.lock') }} - restore-keys: | - ${{ runner.os }}-cargo-${{ github.job }}- - - - run: cargo test --workspace --locked - working-directory: litellm-rust + workspaces: litellm-rust + cache-on-failure: true - run: cargo test -p litellm-core --no-default-features --locked working-directory: litellm-rust diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-bridge/src/marshal.rs index ef3ec8392ca..4dac89b1b23 100644 --- a/litellm-rust/crates/python-bridge/src/marshal.rs +++ b/litellm-rust/crates/python-bridge/src/marshal.rs @@ -290,6 +290,7 @@ mod tests { #[test] fn required_shapes_preserve_nested_values_and_existing_errors() { + Python::initialize(); let nested = json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]); assert_eq!( Value::Array(required_array("messages", nested.clone()).unwrap()), diff --git a/litellm-rust/crates/token-counter/src/tiktoken.rs b/litellm-rust/crates/token-counter/src/tiktoken.rs index c479ae01be9..7a9e71ed587 100644 --- a/litellm-rust/crates/token-counter/src/tiktoken.rs +++ b/litellm-rust/crates/token-counter/src/tiktoken.rs @@ -195,14 +195,21 @@ mod tests { } #[test] - fn long_repeated_runs_stay_cheap() { + fn long_repeated_runs_cost_close_to_linear() { let ranks = ranks(); let mut scratch = MergeScratch::default(); - let piece = vec![b' '; 1 << 20]; - let started = std::time::Instant::now(); - let count = ranks.count_piece(&piece, &mut scratch); - assert!(count > 0); - assert!(started.elapsed().as_secs() < 5, "{:?}", started.elapsed()); + let mut time = |len: usize| { + let piece = vec![b' '; len]; + let started = std::time::Instant::now(); + assert!(ranks.count_piece(&piece, &mut scratch) > 0); + started.elapsed() + }; + let small = (0..3).map(|_| time(1 << 14)).min().unwrap(); + let large = time(1 << 18); + assert!( + large < small * 64, + "{small:?} for 2^14 bytes, {large:?} for 2^18" + ); } #[test] diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index 27c90c9d71e..ba8ce7e5625 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from copy import deepcopy from typing import TYPE_CHECKING, Any, Final, cast from urllib.parse import urlparse @@ -14,6 +15,7 @@ from litellm.types.integrations.rag.bedrock_knowledgebase import ( BedrockKBResponse, BedrockKBRetrievalConfiguration, BedrockKBRetrievalQuery, + BedrockKBUserContext, ) from litellm.types.router import GenericLiteLLMParams from litellm.types.vector_stores import ( @@ -242,10 +244,29 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): retrieval_config.setdefault("vectorSearchConfiguration", {})["filter"] = filters if retrieval_config: request_body["retrievalConfiguration"] = cast(BedrockKBRetrievalConfiguration, retrieval_config) + user_context: Final = self._user_context(extra_body=extra_body, litellm_params=litellm_params) + if user_context is not None: + request_body["userContext"] = user_context litellm_logging_obj.model_call_details["query"] = query return url, request_body + @staticmethod + def _user_context( + extra_body: Mapping[str, object] | None, litellm_params: Mapping[str, object] + ) -> BedrockKBUserContext | None: + sources: Final = tuple(source for source in (extra_body, litellm_params) if isinstance(source, Mapping)) + found: Final = next( + ( + source[key] + for source in sources + for key in ("userContext", "user_context") + if source.get(key) is not None + ), + None, + ) + return None if found is None else cast(BedrockKBUserContext, found) + def sign_request( self, headers: dict, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index cd458666539..dd21bbf0b25 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -11313,13 +11313,16 @@ }, "azure_ai/grok-4.3": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "azure_ai", "max_input_tokens": 200000, "max_output_tokens": 200000, "max_tokens": 200000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, "supports_prompt_caching": true, @@ -11331,13 +11334,16 @@ }, "azure_ai/grok-4.6": { "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_200k_tokens": 1e-06, "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, "litellm_provider": "azure_ai", "max_input_tokens": 200000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 6e-06, + "output_cost_per_token_above_200k_tokens": 1.2e-05, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, "supports_prompt_caching": true, @@ -26162,7 +26168,9 @@ "output_vector_size": 3072, "rpm": 10000, "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supports_audio_input": true, "supports_multimodal": true, + "supports_vision": true, "tpm": 10000000 }, "gemini/gemini-1.5-flash": { @@ -26303,8 +26311,8 @@ "input_cost_per_token": 3e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_reasoning_token": 2.5e-06, "output_cost_per_token": 2.5e-06, @@ -26350,6 +26358,7 @@ "output_cost_per_token_batches": 1.25e-06, "output_cost_per_token_flex": 1.25e-06, "output_cost_per_token_priority": 4.5e-06, + "supports_audio_input": true, "supports_image_size": false }, "gemini/gemini-2.5-flash-image": { @@ -26362,7 +26371,7 @@ "input_cost_per_token_priority": 5.4e-07, "litellm_provider": "gemini", "supports_reasoning": false, - "max_input_tokens": 32768, + "max_input_tokens": 65536, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "image_generation", @@ -26388,22 +26397,23 @@ "image" ], "supports_audio_output": false, - "supports_function_calling": true, + "supports_function_calling": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, "supports_vision": true, - "supports_web_search": true, + "supports_web_search": false, "tpm": 8000000, "search_context_cost_per_query": { "search_context_size_low": 0.035, "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, + "supports_audio_input": false, "supports_image_size": false }, "gemini/gemini-3-pro-image": { @@ -26441,7 +26451,7 @@ ], "supports_function_calling": false, "supports_prompt_caching": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_system_messages": true, "supports_vision": true, "supports_web_search": true, @@ -26544,7 +26554,7 @@ "input_cost_per_token": 5e-07, "input_cost_per_token_batches": 2.5e-07, "litellm_provider": "gemini", - "max_input_tokens": 65536, + "max_input_tokens": 131072, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "image_generation", @@ -26571,7 +26581,7 @@ "supports_function_calling": false, "supports_prompt_caching": true, "supports_reasoning": false, - "supports_response_schema": true, + "supports_response_schema": false, "supports_system_messages": true, "supports_vision": true, "supports_web_search": true, @@ -26653,12 +26663,13 @@ "text", "image" ], - "supports_function_calling": true, + "supports_function_calling": false, "supports_prompt_caching": false, "supports_reasoning": false, "supports_response_schema": false, "supports_system_messages": true, "supports_vision": true, + "supports_web_search": false, "tpm": 4000000 }, "gemini/deep-research-pro-preview-12-2025": { @@ -26711,8 +26722,8 @@ "input_cost_per_token": 1e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_reasoning_token": 4e-07, "output_cost_per_token": 4e-07, @@ -26758,6 +26769,7 @@ "output_cost_per_token_batches": 2e-07, "output_cost_per_token_flex": 2e-07, "output_cost_per_token_priority": 7.2e-07, + "supports_audio_input": true, "supports_image_size": false }, "gemini/gemini-2.5-flash-lite-preview-09-2025": { @@ -27005,6 +27017,9 @@ "input_cost_per_token": 5e-07, "input_cost_per_token_batches": 2.5e-07, "litellm_provider": "gemini", + "max_input_tokens": 8192, + "max_output_tokens": 16384, + "max_tokens": 16384, "mode": "audio_speech", "output_cost_per_audio_token": 1e-05, "output_cost_per_token": 1e-05, @@ -27013,7 +27028,11 @@ "/v1/audio/speech" ], "tpm": 4000000, - "rpm": 10 + "rpm": 10, + "supports_audio_input": false, + "supports_function_calling": false, + "supports_response_schema": false, + "supports_web_search": false }, "gemini/gemini-2.5-pro": { "cache_read_input_token_cost": 1.25e-07, @@ -27027,8 +27046,8 @@ "input_cost_per_token_above_200k_tokens_priority": 4.5e-06, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_above_200k_tokens": 1.5e-05, @@ -27338,8 +27357,8 @@ "input_cost_per_token": 5e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_reasoning_token": 3e-06, "output_cost_per_token": 3e-06, @@ -27388,7 +27407,8 @@ "input_cost_per_token_batches": 2.5e-07, "input_cost_per_token_flex": 2.5e-07, "output_cost_per_token_batches": 1.5e-06, - "output_cost_per_token_flex": 1.5e-06 + "output_cost_per_token_flex": 1.5e-06, + "supports_audio_input": true }, "gemini/gemini-3.5-flash": { "prompt_cache_min_tokens": 4096, @@ -27397,8 +27417,8 @@ "input_cost_per_token": 1.5e-06, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_reasoning_token": 9e-06, "output_cost_per_token": 9e-06, @@ -28112,9 +28132,9 @@ "input_cost_per_token": 1e-06, "input_cost_per_token_batches": 5e-07, "litellm_provider": "gemini", - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_input_tokens": 8192, + "max_output_tokens": 16384, + "max_tokens": 16384, "mode": "chat", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2e-05, @@ -28127,19 +28147,20 @@ "audio" ], "supports_audio_output": false, - "supports_function_calling": true, + "supports_function_calling": false, "supports_prompt_caching": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, + "supports_vision": false, + "supports_web_search": false, "tpm": 10000000, "search_context_cost_per_query": { "search_context_size_low": 0.035, "search_context_size_medium": 0.035, "search_context_size_high": 0.035 - } + }, + "supports_audio_input": false }, "gemini/gemini-exp-1114": { "input_cost_per_token": 0, @@ -55917,7 +55938,7 @@ "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, "litellm_provider": "gemini", - "max_input_tokens": 1048576, + "max_input_tokens": 131072, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "realtime", @@ -55937,7 +55958,11 @@ ], "supports_audio_input": true, "supports_audio_output": true, - "gemini_native_audio": true + "gemini_native_audio": true, + "supports_function_calling": true, + "supports_response_schema": false, + "supports_vision": true, + "supports_web_search": true }, "gemini-3.1-flash-live-preview": { "input_cost_per_audio_token": 3e-06, @@ -55971,7 +55996,8 @@ "supports_vision": true, "supports_web_search": true, "gemini_audio_only_live": true, - "input_cost_per_second": 8.33333333333e-05 + "input_cost_per_second": 8.33333333333e-05, + "supports_response_schema": false }, "gemini/gemini-2.5-flash-native-audio-latest": { "input_cost_per_audio_token": 3e-06, @@ -56033,7 +56059,7 @@ "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, "litellm_provider": "gemini", - "max_input_tokens": 1048576, + "max_input_tokens": 131072, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "realtime", @@ -56055,7 +56081,11 @@ "supports_audio_output": true, "tpm": 250000, "rpm": 10, - "gemini_native_audio": true + "gemini_native_audio": true, + "supports_function_calling": true, + "supports_response_schema": false, + "supports_vision": true, + "supports_web_search": true }, "gemini/gemini-3.1-flash-live-preview": { "input_cost_per_audio_token": 3e-06, @@ -56091,7 +56121,8 @@ "tpm": 250000, "rpm": 10, "gemini_audio_only_live": true, - "input_cost_per_second": 8.33333333333e-05 + "input_cost_per_second": 8.33333333333e-05, + "supports_response_schema": false }, "gemini/gemini-3.1-flash-tts-preview": { "input_cost_per_token": 1e-06, @@ -56108,19 +56139,29 @@ "/v1/audio/speech" ], "tpm": 4000000, - "rpm": 10 + "rpm": 10, + "supports_function_calling": false, + "supports_response_schema": false, + "supports_web_search": false }, "gemini-2.5-flash-preview-tts": { "input_cost_per_token": 5e-07, "input_cost_per_token_batches": 2.5e-07, "litellm_provider": "gemini", + "max_input_tokens": 8192, + "max_output_tokens": 16384, + "max_tokens": 16384, "mode": "audio_speech", "output_cost_per_audio_token": 1e-05, "output_cost_per_token": 1e-05, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/audio/speech" - ] + ], + "supports_audio_input": false, + "supports_function_calling": false, + "supports_response_schema": false, + "supports_web_search": false }, "gemini-flash-latest": { "cache_read_input_token_cost": 3e-08, @@ -58694,6 +58735,9 @@ ], "supports_audio_input": true, "supports_audio_output": true, + "supports_function_calling": false, + "supports_response_schema": false, + "supports_web_search": false, "tpm": 250000 }, "gemini/gemini-3.5-transcribe": { @@ -58715,7 +58759,8 @@ ], "supports_audio_input": true, "tpm": 800000, - "rpm": 2000 + "rpm": 2000, + "supports_function_calling": false }, "gemini/gemini-3.5-transcribe-live": { "input_cost_per_audio_token": 3.5e-06, @@ -58735,7 +58780,8 @@ ], "supports_audio_input": true, "tpm": 250000, - "rpm": 10 + "rpm": 10, + "supports_function_calling": false }, "vertex_ai/gemini-3.5-transcribe-preview": { "input_cost_per_audio_token": 2e-06, @@ -61374,7 +61420,7 @@ "input_cost_per_audio_token": 1.5e-06, "input_cost_per_token": 1.5e-06, "litellm_provider": "gemini", - "max_input_tokens": 131072, + "max_input_tokens": 1048576, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index c28b5558c75..1b9f39449cf 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -176,6 +176,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return {**item_kwargs, "name": tool_name, **namespace_kwargs} def _is_reasoning_end(self, chunk): + if not chunk.choices: + return False delta: Final = chunk.choices[0].delta # if this indicates reasoning content, don't consider reasoning ended @@ -897,6 +899,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): # Change: Never return a value, just enqueue output item events if self.sent_output_item_added_event: return + if not chunk.choices: + return delta: Final = chunk.choices[0].delta self._sequence_number += 1 @@ -1224,6 +1228,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): It's unclear how users expect litellm to translate multiple-choices-per-chunk to the responses API output. """ + if not choices: + return "" choice: Final = choices[0] chat_completion_delta: Final[ChatCompletionDelta] = choice.delta return chat_completion_delta.content or "" diff --git a/litellm/types/integrations/rag/bedrock_knowledgebase.py b/litellm/types/integrations/rag/bedrock_knowledgebase.py index e3aba85ed9b..7156d8101e1 100644 --- a/litellm/types/integrations/rag/bedrock_knowledgebase.py +++ b/litellm/types/integrations/rag/bedrock_knowledgebase.py @@ -1,6 +1,6 @@ from typing import Any, Literal -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict class BedrockKBLocation(TypedDict, total=False): @@ -127,6 +127,10 @@ class BedrockKBGuardrailConfiguration(TypedDict, total=False): guardrailVersion: str | None +class BedrockKBUserContext(TypedDict): + userId: ReadOnly[str] + + class BedrockKBRequest(TypedDict, total=False): """Complete request structure for Bedrock Knowledge Base retrieval.""" @@ -134,6 +138,7 @@ class BedrockKBRequest(TypedDict, total=False): nextToken: str | None retrievalConfiguration: BedrockKBRetrievalConfiguration | None retrievalQuery: BedrockKBRetrievalQuery + userContext: ReadOnly[BedrockKBUserContext | None] ######################################################################### diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index cd458666539..dd21bbf0b25 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -11313,13 +11313,16 @@ }, "azure_ai/grok-4.3": { "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "azure_ai", "max_input_tokens": 200000, "max_output_tokens": 200000, "max_tokens": 200000, "mode": "chat", "output_cost_per_token": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, "supports_prompt_caching": true, @@ -11331,13 +11334,16 @@ }, "azure_ai/grok-4.6": { "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_200k_tokens": 1e-06, "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, "litellm_provider": "azure_ai", "max_input_tokens": 200000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 6e-06, + "output_cost_per_token_above_200k_tokens": 1.2e-05, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, "supports_prompt_caching": true, @@ -26162,7 +26168,9 @@ "output_vector_size": 3072, "rpm": 10000, "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supports_audio_input": true, "supports_multimodal": true, + "supports_vision": true, "tpm": 10000000 }, "gemini/gemini-1.5-flash": { @@ -26303,8 +26311,8 @@ "input_cost_per_token": 3e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_reasoning_token": 2.5e-06, "output_cost_per_token": 2.5e-06, @@ -26350,6 +26358,7 @@ "output_cost_per_token_batches": 1.25e-06, "output_cost_per_token_flex": 1.25e-06, "output_cost_per_token_priority": 4.5e-06, + "supports_audio_input": true, "supports_image_size": false }, "gemini/gemini-2.5-flash-image": { @@ -26362,7 +26371,7 @@ "input_cost_per_token_priority": 5.4e-07, "litellm_provider": "gemini", "supports_reasoning": false, - "max_input_tokens": 32768, + "max_input_tokens": 65536, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "image_generation", @@ -26388,22 +26397,23 @@ "image" ], "supports_audio_output": false, - "supports_function_calling": true, + "supports_function_calling": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_system_messages": true, "supports_tool_choice": true, "supports_url_context": true, "supports_vision": true, - "supports_web_search": true, + "supports_web_search": false, "tpm": 8000000, "search_context_cost_per_query": { "search_context_size_low": 0.035, "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, + "supports_audio_input": false, "supports_image_size": false }, "gemini/gemini-3-pro-image": { @@ -26441,7 +26451,7 @@ ], "supports_function_calling": false, "supports_prompt_caching": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_system_messages": true, "supports_vision": true, "supports_web_search": true, @@ -26544,7 +26554,7 @@ "input_cost_per_token": 5e-07, "input_cost_per_token_batches": 2.5e-07, "litellm_provider": "gemini", - "max_input_tokens": 65536, + "max_input_tokens": 131072, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "image_generation", @@ -26571,7 +26581,7 @@ "supports_function_calling": false, "supports_prompt_caching": true, "supports_reasoning": false, - "supports_response_schema": true, + "supports_response_schema": false, "supports_system_messages": true, "supports_vision": true, "supports_web_search": true, @@ -26653,12 +26663,13 @@ "text", "image" ], - "supports_function_calling": true, + "supports_function_calling": false, "supports_prompt_caching": false, "supports_reasoning": false, "supports_response_schema": false, "supports_system_messages": true, "supports_vision": true, + "supports_web_search": false, "tpm": 4000000 }, "gemini/deep-research-pro-preview-12-2025": { @@ -26711,8 +26722,8 @@ "input_cost_per_token": 1e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_reasoning_token": 4e-07, "output_cost_per_token": 4e-07, @@ -26758,6 +26769,7 @@ "output_cost_per_token_batches": 2e-07, "output_cost_per_token_flex": 2e-07, "output_cost_per_token_priority": 7.2e-07, + "supports_audio_input": true, "supports_image_size": false }, "gemini/gemini-2.5-flash-lite-preview-09-2025": { @@ -27005,6 +27017,9 @@ "input_cost_per_token": 5e-07, "input_cost_per_token_batches": 2.5e-07, "litellm_provider": "gemini", + "max_input_tokens": 8192, + "max_output_tokens": 16384, + "max_tokens": 16384, "mode": "audio_speech", "output_cost_per_audio_token": 1e-05, "output_cost_per_token": 1e-05, @@ -27013,7 +27028,11 @@ "/v1/audio/speech" ], "tpm": 4000000, - "rpm": 10 + "rpm": 10, + "supports_audio_input": false, + "supports_function_calling": false, + "supports_response_schema": false, + "supports_web_search": false }, "gemini/gemini-2.5-pro": { "cache_read_input_token_cost": 1.25e-07, @@ -27027,8 +27046,8 @@ "input_cost_per_token_above_200k_tokens_priority": 4.5e-06, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_above_200k_tokens": 1.5e-05, @@ -27338,8 +27357,8 @@ "input_cost_per_token": 5e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_reasoning_token": 3e-06, "output_cost_per_token": 3e-06, @@ -27388,7 +27407,8 @@ "input_cost_per_token_batches": 2.5e-07, "input_cost_per_token_flex": 2.5e-07, "output_cost_per_token_batches": 1.5e-06, - "output_cost_per_token_flex": 1.5e-06 + "output_cost_per_token_flex": 1.5e-06, + "supports_audio_input": true }, "gemini/gemini-3.5-flash": { "prompt_cache_min_tokens": 4096, @@ -27397,8 +27417,8 @@ "input_cost_per_token": 1.5e-06, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_reasoning_token": 9e-06, "output_cost_per_token": 9e-06, @@ -28112,9 +28132,9 @@ "input_cost_per_token": 1e-06, "input_cost_per_token_batches": 5e-07, "litellm_provider": "gemini", - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_input_tokens": 8192, + "max_output_tokens": 16384, + "max_tokens": 16384, "mode": "chat", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2e-05, @@ -28127,19 +28147,20 @@ "audio" ], "supports_audio_output": false, - "supports_function_calling": true, + "supports_function_calling": false, "supports_prompt_caching": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true, - "supports_web_search": true, + "supports_vision": false, + "supports_web_search": false, "tpm": 10000000, "search_context_cost_per_query": { "search_context_size_low": 0.035, "search_context_size_medium": 0.035, "search_context_size_high": 0.035 - } + }, + "supports_audio_input": false }, "gemini/gemini-exp-1114": { "input_cost_per_token": 0, @@ -55917,7 +55938,7 @@ "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, "litellm_provider": "gemini", - "max_input_tokens": 1048576, + "max_input_tokens": 131072, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "realtime", @@ -55937,7 +55958,11 @@ ], "supports_audio_input": true, "supports_audio_output": true, - "gemini_native_audio": true + "gemini_native_audio": true, + "supports_function_calling": true, + "supports_response_schema": false, + "supports_vision": true, + "supports_web_search": true }, "gemini-3.1-flash-live-preview": { "input_cost_per_audio_token": 3e-06, @@ -55971,7 +55996,8 @@ "supports_vision": true, "supports_web_search": true, "gemini_audio_only_live": true, - "input_cost_per_second": 8.33333333333e-05 + "input_cost_per_second": 8.33333333333e-05, + "supports_response_schema": false }, "gemini/gemini-2.5-flash-native-audio-latest": { "input_cost_per_audio_token": 3e-06, @@ -56033,7 +56059,7 @@ "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, "litellm_provider": "gemini", - "max_input_tokens": 1048576, + "max_input_tokens": 131072, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "realtime", @@ -56055,7 +56081,11 @@ "supports_audio_output": true, "tpm": 250000, "rpm": 10, - "gemini_native_audio": true + "gemini_native_audio": true, + "supports_function_calling": true, + "supports_response_schema": false, + "supports_vision": true, + "supports_web_search": true }, "gemini/gemini-3.1-flash-live-preview": { "input_cost_per_audio_token": 3e-06, @@ -56091,7 +56121,8 @@ "tpm": 250000, "rpm": 10, "gemini_audio_only_live": true, - "input_cost_per_second": 8.33333333333e-05 + "input_cost_per_second": 8.33333333333e-05, + "supports_response_schema": false }, "gemini/gemini-3.1-flash-tts-preview": { "input_cost_per_token": 1e-06, @@ -56108,19 +56139,29 @@ "/v1/audio/speech" ], "tpm": 4000000, - "rpm": 10 + "rpm": 10, + "supports_function_calling": false, + "supports_response_schema": false, + "supports_web_search": false }, "gemini-2.5-flash-preview-tts": { "input_cost_per_token": 5e-07, "input_cost_per_token_batches": 2.5e-07, "litellm_provider": "gemini", + "max_input_tokens": 8192, + "max_output_tokens": 16384, + "max_tokens": 16384, "mode": "audio_speech", "output_cost_per_audio_token": 1e-05, "output_cost_per_token": 1e-05, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/audio/speech" - ] + ], + "supports_audio_input": false, + "supports_function_calling": false, + "supports_response_schema": false, + "supports_web_search": false }, "gemini-flash-latest": { "cache_read_input_token_cost": 3e-08, @@ -58694,6 +58735,9 @@ ], "supports_audio_input": true, "supports_audio_output": true, + "supports_function_calling": false, + "supports_response_schema": false, + "supports_web_search": false, "tpm": 250000 }, "gemini/gemini-3.5-transcribe": { @@ -58715,7 +58759,8 @@ ], "supports_audio_input": true, "tpm": 800000, - "rpm": 2000 + "rpm": 2000, + "supports_function_calling": false }, "gemini/gemini-3.5-transcribe-live": { "input_cost_per_audio_token": 3.5e-06, @@ -58735,7 +58780,8 @@ ], "supports_audio_input": true, "tpm": 250000, - "rpm": 10 + "rpm": 10, + "supports_function_calling": false }, "vertex_ai/gemini-3.5-transcribe-preview": { "input_cost_per_audio_token": 2e-06, @@ -61374,7 +61420,7 @@ "input_cost_per_audio_token": 1.5e-06, "input_cost_per_token": 1.5e-06, "litellm_provider": "gemini", - "max_input_tokens": 131072, + "max_input_tokens": 1048576, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", diff --git a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py b/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py index 7b04efa17dc..ab5a2531461 100644 --- a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py +++ b/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py @@ -1,3 +1,4 @@ +from typing import Final from unittest.mock import MagicMock from litellm.llms.bedrock.vector_stores.transformation import BedrockVectorStoreConfig @@ -82,6 +83,7 @@ def test_transform_search_request_uses_only_retrieval_config_from_extra_body(): == "HYBRID" ) assert "unrelatedField" not in body + assert "userContext" not in body def test_transform_search_request_does_not_mutate_extra_body_and_overrides_number_of_results(): @@ -152,3 +154,44 @@ def test_transform_search_request_overrides_filter_without_mutating_extra_body() ]["value"] == "a" ) + + +def _search_body(extra_body: dict[str, object] | None, litellm_params: dict[str, object]) -> dict[str, object]: + config: Final = BedrockVectorStoreConfig() + mock_log: Final = MagicMock() + mock_log.model_call_details = {} + _, body = config.transform_search_vector_store_request( + vector_store_id="kb123", + query="hello", + vector_store_search_optional_params={"max_num_results": 3}, + api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases", + litellm_logging_obj=mock_log, + litellm_params=litellm_params, + extra_body=extra_body, + ) + return body + + +def test_transform_search_request_forwards_user_context_from_extra_body(): + body = _search_body(extra_body={"userContext": {"userId": "alice@example.com"}}, litellm_params={}) + + assert body["userContext"] == {"userId": "alice@example.com"} + assert body["retrievalConfiguration"] == {"vectorSearchConfiguration": {"numberOfResults": 3}} + + +def test_transform_search_request_forwards_top_level_user_context_from_litellm_params(): + body = _search_body( + extra_body=None, + litellm_params={"vector_store_id": "kb123", "user_context": {"userId": "bob@example.com"}}, + ) + + assert body["userContext"] == {"userId": "bob@example.com"} + + +def test_transform_search_request_prefers_extra_body_user_context_over_top_level(): + body = _search_body( + extra_body={"userContext": {"userId": "alice@example.com"}}, + litellm_params={"userContext": {"userId": "bob@example.com"}}, + ) + + assert body["userContext"] == {"userId": "alice@example.com"} diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 5d97b0531d6..343fc873fa4 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -752,6 +752,49 @@ def test_completed_event_restores_usage_hidden_by_stream_options_none(): assert completed.response.usage.output_tokens == 5 +def _empty_choices_chunk(usage: Usage | None = None) -> ModelResponseStream: + return ModelResponseStream(id=CHAT_COMPLETION_ID, model="claude-haiku-4-5", choices=[], usage=usage) + + +@pytest.mark.asyncio +async def test_leading_empty_choices_chunk_does_not_kill_the_stream(): + """ + Azure leads some streams with a `prompt_filter_results` chunk whose `choices` is empty. + The bridge used to index `choices[0]` on it and die before the first token. + """ + iterator = _build_iterator([_empty_choices_chunk(), _chunk("Hello"), _chunk("!", finish_reason="stop")]) + + events = [event async for event in iterator] + + event_types = [getattr(event, "type", None) for event in events] + assert event_types.count(ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED) == 1 + assert "".join(event.delta for event in events if event.type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA) == "Hello!" + assert event_types[-1] == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + + +@pytest.mark.asyncio +async def test_trailing_empty_choices_usage_chunk_reaches_response_completed(): + """ + With `stream_options.include_usage` (which the bridge always sets) the last upstream chunk + carries only usage and an empty `choices`. It must not crash the stream, and its usage must + still land on `response.completed`. + """ + usage: Final = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) + iterator = _build_iterator([_chunk("Hello"), _chunk("", finish_reason="stop"), _empty_choices_chunk(usage)]) + + events = [event async for event in iterator] + + completed = next( + event for event in events if getattr(event, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + ) + assert completed.response.usage.input_tokens == 10 + assert completed.response.usage.output_tokens == 5 + + +def test_is_reasoning_end_ignores_empty_choices_chunk(): + assert _build_iterator([])._is_reasoning_end(_empty_choices_chunk()) is False + + def test_object_tool_call_arguments_stream_as_valid_json(): """A provider that sends decoded object arguments must still stream valid JSON. diff --git a/tests/test_litellm/test_gemini_3_1_flash_lite_image_pricing.py b/tests/test_litellm/test_gemini_3_1_flash_lite_image_pricing.py index 10d1d6fecd1..250b587aaf1 100644 --- a/tests/test_litellm/test_gemini_3_1_flash_lite_image_pricing.py +++ b/tests/test_litellm/test_gemini_3_1_flash_lite_image_pricing.py @@ -3,14 +3,7 @@ from pathlib import Path import pytest -import litellm from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider -from litellm.types.utils import ( - ImageObject, - ImageResponse, - ImageUsage, - ImageUsageInputTokensDetails, -) REPO_ROOT = Path(__file__).parents[2] MAIN_PATH = REPO_ROOT / "model_prices_and_context_window.json" @@ -21,94 +14,12 @@ GEMINI = "gemini/gemini-3.1-flash-lite-image" VERTEX = "vertex_ai/gemini-3.1-flash-lite-image" ALL_KEYS = (UNPREFIXED, GEMINI, VERTEX) -INPUT_COST = 2.5e-07 -INPUT_COST_BATCHES = 1.25e-07 -OUTPUT_TEXT_COST = 1.5e-06 -OUTPUT_TEXT_COST_BATCHES = 7.5e-07 -OUTPUT_IMAGE_TOKEN_COST = 3e-05 -OUTPUT_COST_PER_1K_IMAGE = 0.0336 -INPUT_COST_PER_IMAGE = 0.00028 -CACHE_READ_COST = 2.5e-08 -MAX_INPUT_TOKENS = 65536 -MAX_OUTPUT_TOKENS = 4096 -TOKENS_PER_1K_IMAGE = 1120 - -SHARED_FIELDS = { - "mode": "image_generation", - "input_cost_per_token": INPUT_COST, - "input_cost_per_token_batches": INPUT_COST_BATCHES, - "input_cost_per_image": INPUT_COST_PER_IMAGE, - "output_cost_per_token": OUTPUT_TEXT_COST, - "output_cost_per_token_batches": OUTPUT_TEXT_COST_BATCHES, - "output_cost_per_image": OUTPUT_COST_PER_1K_IMAGE, - "output_cost_per_image_token": OUTPUT_IMAGE_TOKEN_COST, - "max_input_tokens": MAX_INPUT_TOKENS, - "max_output_tokens": MAX_OUTPUT_TOKENS, - "max_tokens": MAX_OUTPUT_TOKENS, - "supported_endpoints": ["/v1/chat/completions", "/v1/completions", "/v1/batch"], - "supported_output_modalities": ["text", "image"], - "supports_reasoning": False, - "supports_response_schema": False, - "supports_system_messages": True, - "supports_vision": True, -} - -VERTEX_ROUTE_FIELDS = { - "litellm_provider": "vertex_ai-language-models", - "cache_read_input_token_cost": CACHE_READ_COST, - "supported_modalities": ["text", "image", "video"], - "supports_function_calling": False, - "supports_pdf_input": True, - "supports_prompt_caching": True, - "supports_video_input": True, -} - -PER_ROUTE_FIELDS = { - UNPREFIXED: VERTEX_ROUTE_FIELDS, - VERTEX: VERTEX_ROUTE_FIELDS, - GEMINI: { - "litellm_provider": "gemini", - "supported_modalities": ["text", "image"], - "supports_function_calling": True, - "supports_prompt_caching": False, - "rpm": 1000, - "tpm": 4000000, - }, -} - -GROUNDING_FIELDS = ( - "supports_web_search", - "search_context_cost_per_query", - "web_search_billing_unit", -) - def _load(path: Path) -> dict: with open(path, encoding="utf-8") as f: return json.load(f) -@pytest.fixture -def local_model_cost_map(monkeypatch): - original_model_cost = litellm.model_cost - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - litellm.model_cost = litellm.get_model_cost_map(url="") - litellm.get_model_info.cache_clear() - try: - yield - finally: - litellm.model_cost = original_model_cost - litellm.get_model_info.cache_clear() - - -@pytest.mark.parametrize("model", ALL_KEYS) -@pytest.mark.parametrize("path", (MAIN_PATH, BACKUP_PATH), ids=("main", "backup")) -def test_per_route_capabilities_match_model_cards(model: str, path: Path): - info = _load(path)[model] - for field, value in PER_ROUTE_FIELDS[model].items(): - assert info[field] == value, f"{model} {field} in {path.name}: {info.get(field)} != {value}" - - @pytest.mark.parametrize("model", ALL_KEYS) def test_backup_matches_main(model: str): assert _load(BACKUP_PATH).get(model) == _load(MAIN_PATH).get(model) @@ -124,18 +35,3 @@ def test_vertex_prefix_routes_to_vertex(): routed_model, provider, _, _ = get_llm_provider(model=VERTEX) assert routed_model == UNPREFIXED assert provider == "vertex_ai" - - -def _one_k_image_response() -> ImageResponse: - return ImageResponse( - data=[ImageObject(b64_json="img1")], - usage=ImageUsage( - input_tokens=50 + TOKENS_PER_1K_IMAGE, - input_tokens_details=ImageUsageInputTokensDetails( - text_tokens=50, - image_tokens=TOKENS_PER_1K_IMAGE, - ), - output_tokens=TOKENS_PER_1K_IMAGE, - total_tokens=50 + TOKENS_PER_1K_IMAGE + TOKENS_PER_1K_IMAGE, - ), - ) diff --git a/tests/test_litellm/vector_stores/test_main.py b/tests/test_litellm/vector_stores/test_main.py index e3575c33b17..1c968126c42 100644 --- a/tests/test_litellm/vector_stores/test_main.py +++ b/tests/test_litellm/vector_stores/test_main.py @@ -7,6 +7,7 @@ executor, and it must never leak into litellm_params/kwargs where logging would model_dump() it (the #19550 serialization trap). """ +import json from unittest.mock import MagicMock, patch import pytest @@ -15,6 +16,7 @@ import litellm.vector_stores.main as vector_stores_main from litellm.llms.base_llm.vector_store.transformation import ( RouterVectorStoreEmbeddingExecutor, ) +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.vector_stores.main import search MOCK_SEARCH_RESPONSE = { @@ -89,3 +91,26 @@ def test_search_router_not_in_litellm_params(): litellm_params = mock_handler.call_args.kwargs["litellm_params"] assert "router" not in litellm_params.model_dump(exclude_none=True) assert getattr(litellm_params, "router", None) is None + + +def test_search_forwards_top_level_user_context_to_bedrock_retrieve(): + """Regression (LIT-4415): a top-level userContext, the shape the OpenAI SDK's extra_body + produces on the proxy path, reaches the Bedrock Retrieve request body.""" + client = MagicMock(spec=HTTPHandler) + client.post.return_value = MagicMock(status_code=200, json=MagicMock(return_value={"retrievalResults": []})) + + search( + vector_store_id="kb123", + query="q", + custom_llm_provider="bedrock", + aws_region_name="us-west-2", + aws_access_key_id="test-key-id", + aws_secret_access_key="test-secret-key", + userContext={"userId": "alice@example.com"}, + client=client, + litellm_logging_obj=MagicMock(), + ) + + posted = json.loads(client.post.call_args.kwargs["data"]) + assert posted["userContext"] == {"userId": "alice@example.com"} + assert posted["retrievalQuery"] == {"text": "q"}