Extract the inline sync streaming post/decode into make_sync_call so it can be
exercised with an injected client, mirroring make_async_call, and add a sync
regression test that each token is forwarded after exactly one pulled frame.
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Mirror the sagemaker_chat fix on the native sagemaker/ streaming path: the
sync and async completion handlers read the invocations-response-stream body
with iter_bytes(chunk_size=1024) / aiter_bytes(chunk_size=1024), so httpx
withholds bytes until 1024 accumulate and tokens arrive in gap-then-burst
waves. Drop the fixed chunk size so each decoded event is forwarded as its
bytes arrive.
Also add a boundary-agnostic decoder test proving frames reassemble correctly
regardless of where transport reads split the stream.
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(sagemaker): use Cohere embed payload for Marketplace endpoints
SageMaker embedding only special-cased Voyage; every other endpoint received
HuggingFace TGI `{"inputs": [...]}`. AWS Marketplace Cohere containers expect
the native Cohere embed payload (`texts`, `input_type`) and reject the HF
shape with `422 EmbedReqV2.inputs is of type string but should be of type
Object`.
Add `SagemakerCohereEmbeddingConfig` that reuses Bedrock/Cohere request and
response transforms, and route SageMaker endpoint names containing `cohere`
or a Cohere embed model fragment (`embed-multilingual`, `embed-english`,
`embed-v3`, `embed-v4`) to it. Supports `input_type`, `dimensions`, and
`encoding_format`. Voyage and HuggingFace SageMaker endpoints are unchanged.
Co-authored-by: Cursor <cursoragent@cursor.com>
* refactor(sagemaker): simplify cohere detection and align with file conventions
- Detect Cohere SageMaker endpoints with a single `"cohere" in model.lower()`
check, mirroring the existing Voyage branch instead of a separate helper
function and marker constant.
- Drop instance caches of sub-configs; instantiate `BedrockCohereEmbeddingConfig`
/ `CohereEmbeddingConfig` per call to match the existing pattern in
`BedrockCohereEmbeddingConfig._transform_request`.
- Match `SagemakerEmbeddingConfig`'s signatures, defaults, and `Any` typing for
`logging_obj`; collapse the input-normalization helper inline.
- Inline `transform_embedding_response` input lookup; no behavior change.
Co-authored-by: Cursor <cursoragent@cursor.com>
* fix(sagemaker): restore provider-supported embedding params after map
Cohere input_type is advertised in get_supported_openai_params but was
filtered out of non_default_params by OPENAI_EMBEDDING_PARAMS before
map_openai_params ran. Merge supported params from passed_params after
map (same path Greptile flagged). Handle input_type explicitly in
SagemakerCohereEmbeddingConfig.map_openai_params and add an integration
test through get_optional_params_embeddings.
Co-authored-by: Cursor <cursoragent@cursor.com>
* fix(embeddings): only restore non-OpenAI supported params after map
The post-map restore loop must skip OPENAI_EMBEDDING_PARAMS so mapped
fields (e.g. dimensions -> output_dimension) are not duplicated under
their OpenAI names. Align SageMaker embedding import order with sibling
files and add a regression test for dimensions mapping.
Co-authored-by: Cursor <cursoragent@cursor.com>
* fix(sagemaker): avoid double post_call on Cohere embedding response
Greptile review on #28613 caught that `CohereEmbeddingConfig._transform_response`
calls `logging_obj.post_call` internally. The SageMaker embedding handler
already calls `post_call` once before invoking the transform, so the Cohere
SageMaker path fired callbacks, cost calculators, and log handlers twice
per request.
Extract the parsing body of `_transform_response` into
`_populate_embedding_response` (pure extract-method, no behavior change
for existing Cohere direct or Bedrock Cohere paths, which keep calling
`_transform_response`). Have `SagemakerCohereEmbeddingConfig` call the
new helper directly so it parses the response without re-logging.
Add a regression test asserting `logging_obj.post_call` is not invoked
by the SageMaker Cohere transform.
Co-authored-by: Cursor <cursoragent@cursor.com>
---------
Co-authored-by: Cursor <cursoragent@cursor.com>
* refactor(bedrock/sagemaker): switch to lazy loading for response stream shapes
- Replace eager loading of BEDROCK_RESPONSE_STREAM_SHAPE and SAGEMAKER_RESPONSE_STREAM_SHAPE with lazy loading via get_bedrock_response_stream_shape() and get_sagemaker_response_stream_shape() respectively.
- This change optimizes performance by avoiding unnecessary imports and logging warnings unless the response stream shapes are actually needed.
- Update relevant classes and tests to utilize the new lazy loading functions, ensuring consistent behavior across the codebase.
* test(bedrock/sagemaker): add fixtures to clear response stream shape cache
- Introduced `_reset_bedrock_response_stream_shape_cache` and `_reset_sagemaker_response_stream_shape_cache` fixtures to prevent lru_cache leakage between tests in their respective modules.
- Updated tests to utilize these fixtures, ensuring that the response stream shape cache is cleared before and after each test run.
- Added `pytest.importorskip("botocore")` to ensure that tests are skipped if the botocore library is not available.
* Refactor Bedrock response stream shape handling
- Introduced a module-level constant `BEDROCK_RESPONSE_STREAM_SHAPE` to cache the response stream shape, eliminating the need for per-instance caching in `BedrockEventStreamDecoderBase`.
- Updated relevant methods to utilize the new constant, improving performance by avoiding redundant loading of the shape.
- Added tests to ensure the shape is loaded correctly at import time and is consistent across different modules.
- Added a new mock server script for testing Bedrock pass-through functionality.
* Refactor response parsing for Bedrock and SageMaker
- Improved code readability by formatting the parsing method calls in `AWSEventStreamDecoder` for both Bedrock and SageMaker response stream shapes.
- Added blank lines for better separation of code blocks in `invoke_handler.py` and `common_utils.py` to enhance maintainability.
* Enhance error handling for Bedrock and SageMaker response stream shape loading
- Wrapped the loading logic in `_load_bedrock_response_stream_shape` and `_load_sagemaker_response_stream_shape` with try-except blocks to gracefully handle exceptions.
- Added logging to warn when the response stream shape cannot be pre-loaded, ensuring the module imports cleanly.
- Updated tests to verify that loading failures return `None` instead of propagating exceptions.
* Implement error handling for missing response stream shapes in Bedrock and SageMaker
- Added checks in `_parse_message_from_event` methods to raise appropriate errors when `BEDROCK_RESPONSE_STREAM_SHAPE` or `SAGEMAKER_RESPONSE_STREAM_SHAPE` is None, ensuring clearer error reporting.
- Updated logging messages to reflect the unavailability of event-stream decoding for both Bedrock and SageMaker.
- Enhanced unit tests to verify that the correct exceptions are raised when the response stream shapes are not loaded.
* feat: add sagemaker_nova provider for Nova models on SageMaker
Add support for custom/fine-tuned Amazon Nova models (Nova Micro, Nova Lite,
Nova 2 Lite) deployed on SageMaker Inference real-time endpoints.
Nova uses OpenAI-compatible request/response format with additional
Nova-specific parameters (top_k, reasoning_effort, allowed_token_ids,
truncate_prompt_tokens) and requires stream:true in the request body.
Nova endpoints also reject 'model' in the request body.
Changes:
- New provider: sagemaker_nova/<endpoint-name>
- SagemakerNovaConfig inherits from SagemakerChatConfig
- Override transform_request to strip 'model' from request body
- Override supports_stream_param_in_request_body (True for Nova)
- Extend get_supported_openai_params with Nova-specific params
- Refactored SagemakerChatConfig to use custom_llm_provider param
instead of hardcoded strings (backwards-compatible)
- Consolidated main.py routing for sagemaker_chat and sagemaker_nova
- 22 unit tests + 9 integration tests (skip-gated)
- Documentation with SDK, streaming, multimodal, and proxy examples
- All tests verified against live SageMaker Nova endpoint
* fix: move integration tests to tests/local_testing/ per test directory policy
* fix: remove unused module-level SagemakerNovaConfig instance
The sagemaker_nova_config singleton was never imported or used — the
ProviderConfigManager creates its own instance via the lambda registered
in utils.py. Removing this leftover boilerplate.
---------
Co-authored-by: Krish Dholakia <krrishdholakia@gmail.com>
The SageMaker embedding handler was not using _load_credentials(),
which meant aws_role_name and aws_session_name parameters were
ignored. This prevented cross-account role assumption for embeddings
while it worked for completions.
Changes:
- Replace direct boto3 client creation with _load_credentials() call
- Create boto3.Session with assumed credentials
- Add comprehensive unit tests for role assumption
This aligns the embedding handler behavior with the completion handler,
which already supports role assumption via the BaseAWSLLM.get_credentials()
method.
Fixes cross-account SageMaker embedding access where users need to
assume a role in another account to invoke endpoints.
HuggingFace Text Embeddings Inference (TEI) returns embeddings as raw
arrays [[0.1, 0.2, ...]] instead of wrapped format {"embedding": [...]}.
This change handles both formats:
- Raw array: [[...]] (TEI, some HF models)
- Wrapped: {"embedding": [[...]]} (standard HF format)
Fixes SagemakerError: "HF response missing 'embedding' field" when using
TEI containers on SageMaker.
Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
* Implement fix for thinking_blocks and converse API calls
This fixes Claude's models via the Converse API, which should also fix
Claude Code.
* Add thinking literal
* Fix mypy issues
* Type fix for redacted thinking
* Add voyage model integration in sagemaker
* Add config file logic
* Use already exiting voyage transformation
* refactor code as per comments
* fix merge error
* refactor code as per comments
* refactor code as per comments
* UI new build
* [Fix] router - regression when adding/removing models (#15451)
* fix(router): update model_name_to_deployment_indices on deployment removal
When a deployment is deleted, the model_name_to_deployment_indices map
was not being updated, causing stale index references. This could lead
to incorrect routing behavior when deployments with the same model_name
were dynamically removed.
Changes:
- Update _update_deployment_indices_after_removal to maintain
model_name_to_deployment_indices mapping
- Remove deleted indices and decrement indices greater than removed index
- Clean up empty entries when no deployments remain for a model name
- Update test to verify proper index shifting and cleanup behavior
* fix(router): remove redundant index building during initialization
Remove duplicate index building operations that were causing unnecessary
work during router initialization:
1. Removed redundant `_build_model_id_to_deployment_index_map` call in
__init__ - `set_model_list` already builds all indices from scratch
2. Removed redundant `_build_model_name_index` call at end of
`set_model_list` - the index is already built incrementally via
`_create_deployment` -> `_add_model_to_list_and_index_map`
Both indices (model_id_to_deployment_index_map and
model_name_to_deployment_indices) are properly maintained as lookup
indexes through existing helper methods. This change eliminates O(N)
duplicate work during initialization without any behavioral changes.
The indices continue to be correctly synchronized with model_list on
all operations (add/remove/upsert).
* fix(prometheus): Fix Prometheus metric collection in a multi-workers environment (#14929)
Co-authored-by: sotazhang <sotazhang@tencent.com>
* Add tiered pricing and cost calculation for xai
* Use generic cost calculator
* Resolve conflicts in generated HTML files
* Remove penalty params as supported params for gemini preview model (#15503)
* fix conversion of thinking block
* add application level encryption in SQS (#15512)
* docs: fix doc
* docs(index.md): bump rc
* [Fix] GEMINI - CLI - add google_routes to llm_api_routes (#15500)
* fix: add google_routes to llm_api_routes
* test: test_virtual_key_llm_api_routes_allows_google_routes
* build: bump version
* bump: version 1.78.0 → 1.78.1
* add application level encryption in SQS
* add application level encryption in SQS
---------
Co-authored-by: Krrish Dholakia <krrishdholakia@gmail.com>
Co-authored-by: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Co-authored-by: deepanshu <deepanshu.lulla@hq.bill.com>
* [Feat] Bedrock Knowledgebase - return search_response when using /chat/completions API with LiteLLM (#15509)
* docs: fix doc
* docs(index.md): bump rc
* [Fix] GEMINI - CLI - add google_routes to llm_api_routes (#15500)
* fix: add google_routes to llm_api_routes
* test: test_virtual_key_llm_api_routes_allows_google_routes
* add AnthropicCitation
* fix async_post_call_success_deployment_hook
* fix add vector_store_custom_logger to global callbacks
* test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call
* async_post_call_success_deployment_hook
* add async_post_call_streaming_deployment_hook
* async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call_streaming(setup_vector_store_registry):
* fix _call_post_streaming_deployment_hook
* fix async_post_call_streaming_deployment_hook
* test update
* docs: Accessing Search Results
* docs KB
* fix chatUI
* fix searchResults
* fix onSearchResults
* fix kb
---------
Co-authored-by: Krrish Dholakia <krrishdholakia@gmail.com>
* [Feat] Add dynamic rate limits on LiteLLM Gateway (#15518)
* docs: fix doc
* docs(index.md): bump rc
* [Fix] GEMINI - CLI - add google_routes to llm_api_routes (#15500)
* fix: add google_routes to llm_api_routes
* test: test_virtual_key_llm_api_routes_allows_google_routes
* build: bump version
* bump: version 1.78.0 → 1.78.1
* fix: KeyRequestBase
* fix rpm_limit_type
* fix dynamic rate limits
* fix use dynamic limits here
* fix _should_enforce_rate_limit
* fix _should_enforce_rate_limit
* fix counter
* test_dynamic_rate_limiting_v3
* use _create_rate_limit_descriptors
---------
Co-authored-by: Krrish Dholakia <krrishdholakia@gmail.com>
* Add google rerank endpoint
* Add docs
* fix mypy error
* fix mypy and lint errors
* Add haiku 4.5 integration
* Add haiku 4.5 integration for other regions as well
* Handle citation field correctly
* Fix filtering headers for signature calcs
* Add haiku 4.5 integration (#15650)
---------
Co-authored-by: Leslie Cheng <leslie.cheng5@gmail.com>
Co-authored-by: Sameer Kankute <sameer@berri.ai>
Co-authored-by: Alexsander Hamir <alexsanderhamirgomesbaptista@gmail.com>
Co-authored-by: Lucas <10226902+LoadingZhang@users.noreply.github.com>
Co-authored-by: sotazhang <sotazhang@tencent.com>
Co-authored-by: Deepanshu Lulla <deepanshu.lulla@gmail.com>
Co-authored-by: Krrish Dholakia <krrishdholakia@gmail.com>
Co-authored-by: deepanshu <deepanshu.lulla@hq.bill.com>
