mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
fix: make streaming for bedrock smoother
Fix stream_chunk_size Fixes https://github.com/BerriAI/litellm/issues/11747
This commit is contained in:
Krrish Dholakia
parent
d20ea3482f
commit
e96fdfedbf
6 changed files with 139 additions and 106 deletions
|
|
@ -111,15 +111,19 @@ _DEFAULT_TTL_FOR_HTTPX_CLIENTS = 3600 # 1 hour, re-use the same httpx client fo
|
|||
# Aiohttp connection pooling - prevents memory leaks from unbounded connection growth
|
||||
# Set to 0 for unlimited (not recommended for production)
|
||||
AIOHTTP_CONNECTOR_LIMIT = int(os.getenv("AIOHTTP_CONNECTOR_LIMIT", 300))
|
||||
AIOHTTP_CONNECTOR_LIMIT_PER_HOST = int(os.getenv("AIOHTTP_CONNECTOR_LIMIT_PER_HOST", 50))
|
||||
AIOHTTP_CONNECTOR_LIMIT_PER_HOST = int(
|
||||
os.getenv("AIOHTTP_CONNECTOR_LIMIT_PER_HOST", 50)
|
||||
)
|
||||
AIOHTTP_KEEPALIVE_TIMEOUT = int(os.getenv("AIOHTTP_KEEPALIVE_TIMEOUT", 120))
|
||||
AIOHTTP_TTL_DNS_CACHE = int(os.getenv("AIOHTTP_TTL_DNS_CACHE", 300))
|
||||
# enable_cleanup_closed is only needed for Python versions with the SSL leak bug
|
||||
# Fixed in Python 3.12.7+ and 3.13.1+ (see https://github.com/python/cpython/pull/118960)
|
||||
# Reference: https://github.com/aio-libs/aiohttp/blob/master/aiohttp/connector.py#L74-L78
|
||||
AIOHTTP_NEEDS_CLEANUP_CLOSED = (
|
||||
(3, 13, 0) <= sys.version_info < (3, 13, 1) or sys.version_info < (3, 12, 7)
|
||||
)
|
||||
AIOHTTP_NEEDS_CLEANUP_CLOSED = (3, 13, 0) <= sys.version_info < (
|
||||
3,
|
||||
13,
|
||||
1,
|
||||
) or sys.version_info < (3, 12, 7)
|
||||
|
||||
# WebSocket constants
|
||||
# Default to None (unlimited) to match OpenAI's official agents SDK behavior
|
||||
|
|
@ -157,7 +161,9 @@ REDIS_UPDATE_BUFFER_KEY = "litellm_spend_update_buffer"
|
|||
REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_spend_update_buffer"
|
||||
REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_team_spend_update_buffer"
|
||||
REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_org_spend_update_buffer"
|
||||
REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_end_user_spend_update_buffer"
|
||||
REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY = (
|
||||
"litellm_daily_end_user_spend_update_buffer"
|
||||
)
|
||||
REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_agent_spend_update_buffer"
|
||||
REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_tag_spend_update_buffer"
|
||||
MAX_REDIS_BUFFER_DEQUEUE_COUNT = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100))
|
||||
|
|
@ -283,7 +289,9 @@ MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB = int(
|
|||
DEFAULT_MAX_TOKENS_FOR_TRITON = int(os.getenv("DEFAULT_MAX_TOKENS_FOR_TRITON", 2000))
|
||||
#### Networking settings ####
|
||||
request_timeout: float = float(os.getenv("REQUEST_TIMEOUT", 6000)) # time in seconds
|
||||
DEFAULT_A2A_AGENT_TIMEOUT: float = float(os.getenv("DEFAULT_A2A_AGENT_TIMEOUT", 6000)) # 10 minutes
|
||||
DEFAULT_A2A_AGENT_TIMEOUT: float = float(
|
||||
os.getenv("DEFAULT_A2A_AGENT_TIMEOUT", 6000)
|
||||
) # 10 minutes
|
||||
STREAM_SSE_DONE_STRING: str = "[DONE]"
|
||||
STREAM_SSE_DATA_PREFIX: str = "data: "
|
||||
### SPEND TRACKING ###
|
||||
|
|
@ -319,8 +327,12 @@ DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE = os.getenv(
|
|||
"DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE", "streaming.chunk.yield"
|
||||
)
|
||||
|
||||
EMAIL_BUDGET_ALERT_TTL = int(os.getenv("EMAIL_BUDGET_ALERT_TTL", 24 * 60 * 60)) # 24 hours in seconds
|
||||
EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE = float(os.getenv("EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE", 0.8)) # 80% of max budget
|
||||
EMAIL_BUDGET_ALERT_TTL = int(
|
||||
os.getenv("EMAIL_BUDGET_ALERT_TTL", 24 * 60 * 60)
|
||||
) # 24 hours in seconds
|
||||
EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE = float(
|
||||
os.getenv("EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE", 0.8)
|
||||
) # 80% of max budget
|
||||
############### LLM Provider Constants ###############
|
||||
### ANTHROPIC CONSTANTS ###
|
||||
ANTHROPIC_TOKEN_COUNTING_BETA_VERSION = os.getenv(
|
||||
|
|
@ -1067,7 +1079,15 @@ known_tokenizer_config = {
|
|||
}
|
||||
|
||||
|
||||
OPENAI_FINISH_REASONS = ["stop", "length", "function_call", "content_filter", "null", "finish_reason_unspecified", "malformed_function_call"]
|
||||
OPENAI_FINISH_REASONS = [
|
||||
"stop",
|
||||
"length",
|
||||
"function_call",
|
||||
"content_filter",
|
||||
"null",
|
||||
"finish_reason_unspecified",
|
||||
"malformed_function_call",
|
||||
]
|
||||
HUMANLOOP_PROMPT_CACHE_TTL_SECONDS = int(
|
||||
os.getenv("HUMANLOOP_PROMPT_CACHE_TTL_SECONDS", 60)
|
||||
) # 1 minute
|
||||
|
|
@ -1333,12 +1353,13 @@ MICROSOFT_USER_EMAIL_ATTRIBUTE = str(
|
|||
MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE = str(
|
||||
os.getenv("MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE", "displayName")
|
||||
)
|
||||
MICROSOFT_USER_ID_ATTRIBUTE = str(
|
||||
os.getenv("MICROSOFT_USER_ID_ATTRIBUTE", "id")
|
||||
)
|
||||
MICROSOFT_USER_ID_ATTRIBUTE = str(os.getenv("MICROSOFT_USER_ID_ATTRIBUTE", "id"))
|
||||
MICROSOFT_USER_FIRST_NAME_ATTRIBUTE = str(
|
||||
os.getenv("MICROSOFT_USER_FIRST_NAME_ATTRIBUTE", "givenName")
|
||||
)
|
||||
MICROSOFT_USER_LAST_NAME_ATTRIBUTE = str(
|
||||
os.getenv("MICROSOFT_USER_LAST_NAME_ATTRIBUTE", "surname")
|
||||
)
|
||||
LITELLM_CONSTANT_STREAM_CHUNK_SIZE = int(
|
||||
os.getenv("LITELLM_CONSTANT_STREAM_CHUNK_SIZE", 10)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,11 +1,11 @@
|
|||
# What is this?
|
||||
## Helper utilities
|
||||
from typing import TYPE_CHECKING, Any, Iterable, List, Literal, Optional, Union
|
||||
from typing import TYPE_CHECKING, Any, Iterable, List, Literal, Optional, Union, cast
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.llms.openai import AllMessageValues, OpenAIChatCompletionFinishReason
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from opentelemetry.trace import Span as _Span
|
||||
|
|
@ -60,6 +60,8 @@ def safe_divide(
|
|||
|
||||
def map_finish_reason(
|
||||
finish_reason: str,
|
||||
) -> (
|
||||
OpenAIChatCompletionFinishReason
|
||||
): # openai supports 5 stop sequences - 'stop', 'length', 'function_call', 'content_filter', 'null'
|
||||
# anthropic mapping
|
||||
if finish_reason == "stop_sequence":
|
||||
|
|
@ -96,7 +98,7 @@ def map_finish_reason(
|
|||
return "tool_calls"
|
||||
elif finish_reason == "content_filtered":
|
||||
return "content_filter"
|
||||
return finish_reason
|
||||
return cast(OpenAIChatCompletionFinishReason, finish_reason)
|
||||
|
||||
|
||||
def remove_index_from_tool_calls(
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ from typing import Any, Optional, Union
|
|||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm.constants import LITELLM_CONSTANT_STREAM_CHUNK_SIZE
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
|
|
@ -29,7 +30,7 @@ def make_sync_call(
|
|||
logging_obj: LiteLLMLoggingObject,
|
||||
json_mode: Optional[bool] = False,
|
||||
fake_stream: bool = False,
|
||||
stream_chunk_size: int = 1024,
|
||||
stream_chunk_size: int = LITELLM_CONSTANT_STREAM_CHUNK_SIZE,
|
||||
):
|
||||
if client is None:
|
||||
client = _get_httpx_client() # Create a new client if none provided
|
||||
|
|
@ -67,7 +68,9 @@ def make_sync_call(
|
|||
)
|
||||
else:
|
||||
decoder = AWSEventStreamDecoder(model=model)
|
||||
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(
|
||||
|
|
@ -103,7 +106,7 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
fake_stream: bool = False,
|
||||
json_mode: Optional[bool] = False,
|
||||
api_key: Optional[str] = None,
|
||||
stream_chunk_size: int = 1024,
|
||||
stream_chunk_size: int = LITELLM_CONSTANT_STREAM_CHUNK_SIZE,
|
||||
) -> CustomStreamWrapper:
|
||||
request_data = await litellm.AmazonConverseConfig()._async_transform_request(
|
||||
model=model,
|
||||
|
|
@ -121,7 +124,7 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
endpoint_url=api_base,
|
||||
data=data,
|
||||
headers=headers,
|
||||
api_key=api_key
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
## LOGGING
|
||||
|
|
@ -181,7 +184,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",
|
||||
|
|
@ -189,7 +192,7 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
endpoint_url=api_base,
|
||||
data=data,
|
||||
headers=headers,
|
||||
api_key=api_key
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
## LOGGING
|
||||
|
|
@ -263,7 +266,9 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
):
|
||||
## SETUP ##
|
||||
stream = optional_params.pop("stream", None)
|
||||
stream_chunk_size = optional_params.pop("stream_chunk_size", 1024)
|
||||
stream_chunk_size = optional_params.pop(
|
||||
"stream_chunk_size", LITELLM_CONSTANT_STREAM_CHUNK_SIZE
|
||||
)
|
||||
unencoded_model_id = optional_params.pop("model_id", None)
|
||||
fake_stream = optional_params.pop("fake_stream", False)
|
||||
json_mode = optional_params.get("json_mode", False)
|
||||
|
|
@ -279,7 +284,6 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
custom_llm_provider="bedrock",
|
||||
)
|
||||
|
||||
|
||||
### SET REGION NAME ###
|
||||
aws_region_name = self._get_aws_region_name(
|
||||
optional_params=optional_params,
|
||||
|
|
@ -303,9 +307,9 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
aws_external_id = optional_params.pop("aws_external_id", None)
|
||||
optional_params.pop("aws_region_name", None)
|
||||
|
||||
litellm_params[
|
||||
"aws_region_name"
|
||||
] = aws_region_name # [DO NOT DELETE] important for async calls
|
||||
litellm_params["aws_region_name"] = (
|
||||
aws_region_name # [DO NOT DELETE] important for async calls
|
||||
)
|
||||
|
||||
credentials: Credentials = self.get_credentials(
|
||||
aws_access_key_id=aws_access_key_id,
|
||||
|
|
@ -379,7 +383,7 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
timeout=timeout,
|
||||
client=client,
|
||||
credentials=credentials,
|
||||
api_key=api_key
|
||||
api_key=api_key,
|
||||
) # type: ignore
|
||||
|
||||
## TRANSFORMATION ##
|
||||
|
|
@ -392,7 +396,7 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
headers=extra_headers,
|
||||
)
|
||||
data = json.dumps(_data)
|
||||
|
||||
|
||||
prepped = self.get_request_headers(
|
||||
credentials=credentials,
|
||||
aws_region_name=aws_region_name,
|
||||
|
|
@ -400,7 +404,7 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
endpoint_url=proxy_endpoint_url,
|
||||
data=data,
|
||||
headers=headers,
|
||||
api_key=api_key
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
## LOGGING
|
||||
|
|
|
|||
|
|
@ -6,15 +6,7 @@ import copy
|
|||
import time
|
||||
import types
|
||||
from functools import partial
|
||||
from typing import (
|
||||
AsyncIterator,
|
||||
Callable,
|
||||
Iterator,
|
||||
Optional,
|
||||
Tuple,
|
||||
cast,
|
||||
get_args,
|
||||
)
|
||||
from typing import AsyncIterator, Callable, Iterator, Optional, Tuple, cast, get_args
|
||||
|
||||
import httpx # type: ignore
|
||||
|
||||
|
|
@ -22,6 +14,7 @@ import litellm
|
|||
from litellm import verbose_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.caching.caching import InMemoryCache
|
||||
from litellm.constants import LITELLM_CONSTANT_STREAM_CHUNK_SIZE
|
||||
from litellm.litellm_core_utils.core_helpers import map_finish_reason
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.litellm_core_utils.logging_utils import track_llm_api_timing
|
||||
|
|
@ -51,11 +44,7 @@ from litellm.types.llms.openai import (
|
|||
ChatCompletionToolCallFunctionChunk,
|
||||
ChatCompletionUsageBlock,
|
||||
)
|
||||
from litellm.types.utils import (
|
||||
ChatCompletionMessageToolCall,
|
||||
Choices,
|
||||
Delta,
|
||||
)
|
||||
from litellm.types.utils import ChatCompletionMessageToolCall, Choices, Delta
|
||||
from litellm.types.utils import GenericStreamingChunk as GChunk
|
||||
from litellm.types.utils import (
|
||||
ModelResponse,
|
||||
|
|
@ -192,17 +181,19 @@ async def make_call(
|
|||
fake_stream: bool = False,
|
||||
json_mode: Optional[bool] = False,
|
||||
bedrock_invoke_provider: Optional[litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL] = None,
|
||||
stream_chunk_size: int = 1024,
|
||||
stream_chunk_size: int = LITELLM_CONSTANT_STREAM_CHUNK_SIZE,
|
||||
):
|
||||
try:
|
||||
if client is None:
|
||||
client = get_async_httpx_client(
|
||||
llm_provider=litellm.LlmProviders.BEDROCK,
|
||||
params={"ssl_verify": logging_obj.litellm_params.get("ssl_verify")}
|
||||
if logging_obj
|
||||
and logging_obj.litellm_params
|
||||
and logging_obj.litellm_params.get("ssl_verify")
|
||||
else None,
|
||||
params=(
|
||||
{"ssl_verify": logging_obj.litellm_params.get("ssl_verify")}
|
||||
if logging_obj
|
||||
and logging_obj.litellm_params
|
||||
and logging_obj.litellm_params.get("ssl_verify")
|
||||
else None
|
||||
),
|
||||
) # Create a new client if none provided
|
||||
|
||||
response = await client.post(
|
||||
|
|
@ -287,16 +278,18 @@ def make_sync_call(
|
|||
fake_stream: bool = False,
|
||||
json_mode: Optional[bool] = False,
|
||||
bedrock_invoke_provider: Optional[litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL] = None,
|
||||
stream_chunk_size: int = 1024,
|
||||
stream_chunk_size: int = LITELLM_CONSTANT_STREAM_CHUNK_SIZE,
|
||||
):
|
||||
try:
|
||||
if client is None:
|
||||
client = _get_httpx_client(
|
||||
params={"ssl_verify": logging_obj.litellm_params.get("ssl_verify")}
|
||||
if logging_obj
|
||||
and logging_obj.litellm_params
|
||||
and logging_obj.litellm_params.get("ssl_verify")
|
||||
else None
|
||||
params=(
|
||||
{"ssl_verify": logging_obj.litellm_params.get("ssl_verify")}
|
||||
if logging_obj
|
||||
and logging_obj.litellm_params
|
||||
and logging_obj.litellm_params.get("ssl_verify")
|
||||
else None
|
||||
)
|
||||
)
|
||||
|
||||
response = client.post(
|
||||
|
|
@ -406,9 +399,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
|
||||
]
|
||||
|
||||
|
|
@ -546,9 +539,9 @@ class BedrockLLM(BaseAWSLLM):
|
|||
content=None,
|
||||
)
|
||||
model_response.choices[0].message = _message # type: ignore
|
||||
model_response._hidden_params[
|
||||
"original_response"
|
||||
] = outputText # allow user to access raw anthropic tool calling response
|
||||
model_response._hidden_params["original_response"] = (
|
||||
outputText # allow user to access raw anthropic tool calling response
|
||||
)
|
||||
if (
|
||||
_is_function_call is True
|
||||
and stream is not None
|
||||
|
|
@ -559,8 +552,10 @@ class BedrockLLM(BaseAWSLLM):
|
|||
)
|
||||
# return an iterator
|
||||
streaming_model_response = ModelResponse(stream=True)
|
||||
streaming_model_response.choices[0].finish_reason = getattr(
|
||||
model_response.choices[0], "finish_reason", "stop"
|
||||
cast(
|
||||
Choices, streaming_model_response.choices[0]
|
||||
).finish_reason = map_finish_reason(
|
||||
getattr(model_response.choices[0], "finish_reason", "stop")
|
||||
)
|
||||
# streaming_model_response.choices = [litellm.utils.StreamingChoices()]
|
||||
streaming_choice = litellm.utils.StreamingChoices()
|
||||
|
|
@ -605,8 +600,8 @@ class BedrockLLM(BaseAWSLLM):
|
|||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
model_response.choices[0].finish_reason = map_finish_reason(
|
||||
completion_response.get("stop_reason", "")
|
||||
cast(Choices, model_response.choices[0]).finish_reason = (
|
||||
map_finish_reason(completion_response.get("stop_reason", ""))
|
||||
)
|
||||
_usage = litellm.Usage(
|
||||
prompt_tokens=completion_response["usage"]["input_tokens"],
|
||||
|
|
@ -641,8 +636,8 @@ class BedrockLLM(BaseAWSLLM):
|
|||
|
||||
# Set finish reason
|
||||
if "finish_reason" in choice:
|
||||
model_response.choices[0].finish_reason = map_finish_reason(
|
||||
choice["finish_reason"]
|
||||
cast(Choices, model_response.choices[0]).finish_reason = (
|
||||
map_finish_reason(choice["finish_reason"])
|
||||
)
|
||||
|
||||
# Set usage if available
|
||||
|
|
@ -781,7 +776,9 @@ class BedrockLLM(BaseAWSLLM):
|
|||
|
||||
## SETUP ##
|
||||
stream = optional_params.pop("stream", None)
|
||||
stream_chunk_size = optional_params.pop("stream_chunk_size", 1024)
|
||||
stream_chunk_size = optional_params.pop(
|
||||
"stream_chunk_size", LITELLM_CONSTANT_STREAM_CHUNK_SIZE
|
||||
)
|
||||
|
||||
provider = self.get_bedrock_invoke_provider(model)
|
||||
modelId = self.get_bedrock_model_id(
|
||||
|
|
@ -881,9 +878,9 @@ class BedrockLLM(BaseAWSLLM):
|
|||
): # completion(top_k=3) > anthropic_config(top_k=3) <- allows for dynamic variables to be passed in
|
||||
inference_params[k] = v
|
||||
if stream is True:
|
||||
inference_params[
|
||||
"stream"
|
||||
] = True # cohere requires stream = True in inference params
|
||||
inference_params["stream"] = (
|
||||
True # cohere requires stream = True in inference params
|
||||
)
|
||||
data = json.dumps({"prompt": prompt, **inference_params})
|
||||
elif provider == "anthropic":
|
||||
if self.is_claude_messages_api_model(model):
|
||||
|
|
@ -1222,7 +1219,7 @@ class BedrockLLM(BaseAWSLLM):
|
|||
logger_fn=None,
|
||||
headers={},
|
||||
client: Optional[AsyncHTTPHandler] = None,
|
||||
stream_chunk_size: int = 1024,
|
||||
stream_chunk_size: int = LITELLM_CONSTANT_STREAM_CHUNK_SIZE,
|
||||
) -> CustomStreamWrapper:
|
||||
# The call is not made here; instead, we prepare the necessary objects for the stream.
|
||||
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ from typing import (
|
|||
|
||||
import httpx
|
||||
|
||||
from litellm.constants import LITELLM_CONSTANT_STREAM_CHUNK_SIZE
|
||||
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.transformation import (
|
||||
AnthropicMessagesConfig,
|
||||
|
|
@ -114,9 +115,7 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
stream=stream,
|
||||
)
|
||||
|
||||
def _remove_ttl_from_cache_control(
|
||||
self, anthropic_messages_request: Dict
|
||||
) -> None:
|
||||
def _remove_ttl_from_cache_control(self, anthropic_messages_request: Dict) -> None:
|
||||
"""
|
||||
Remove `ttl` field from cache_control in messages.
|
||||
Bedrock doesn't support the ttl field in cache_control.
|
||||
|
|
@ -132,7 +131,10 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
for item in content:
|
||||
if isinstance(item, dict) and "cache_control" in item:
|
||||
cache_control = item["cache_control"]
|
||||
if isinstance(cache_control, dict) and "ttl" in cache_control:
|
||||
if (
|
||||
isinstance(cache_control, dict)
|
||||
and "ttl" in cache_control
|
||||
):
|
||||
cache_control.pop("ttl", None)
|
||||
|
||||
def _supports_extended_thinking_on_bedrock(self, model: str) -> bool:
|
||||
|
|
@ -154,10 +156,18 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
|
||||
# Supported models on Bedrock for extended thinking
|
||||
supported_patterns = [
|
||||
"opus-4.5", "opus_4.5", "opus-4-5", "opus_4_5", # Opus 4.5
|
||||
"opus-4.1", "opus_4.1", "opus-4-1", "opus_4_1", # Opus 4.1
|
||||
"opus-4", "opus_4", # Opus 4
|
||||
"sonnet-4", "sonnet_4", # Sonnet 4
|
||||
"opus-4.5",
|
||||
"opus_4.5",
|
||||
"opus-4-5",
|
||||
"opus_4_5", # Opus 4.5
|
||||
"opus-4.1",
|
||||
"opus_4.1",
|
||||
"opus-4-1",
|
||||
"opus_4_1", # Opus 4.1
|
||||
"opus-4",
|
||||
"opus_4", # Opus 4
|
||||
"sonnet-4",
|
||||
"sonnet_4", # Sonnet 4
|
||||
]
|
||||
|
||||
return any(pattern in model_lower for pattern in supported_patterns)
|
||||
|
|
@ -230,7 +240,9 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
input_examples_used: Whether input examples are used
|
||||
beta_set: The set of beta headers to modify in-place
|
||||
"""
|
||||
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():
|
||||
beta_set.add("tool-search-tool-2025-10-19")
|
||||
|
|
@ -242,13 +254,13 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
) -> None:
|
||||
"""
|
||||
Convert Anthropic output_format to inline schema in message content.
|
||||
|
||||
|
||||
Bedrock Invoke doesn't support the output_format parameter, so we embed
|
||||
the schema directly into the user message content as text instructions.
|
||||
|
||||
|
||||
This approach adds the schema to the last user message, instructing the model
|
||||
to respond in the specified JSON format.
|
||||
|
||||
|
||||
Args:
|
||||
output_format: The output_format dict with 'type' and 'schema'
|
||||
anthropic_messages_request: The request dict to modify in-place
|
||||
|
|
@ -256,40 +268,37 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
Ref: https://aws.amazon.com/blogs/machine-learning/structured-data-response-with-amazon-bedrock-prompt-engineering-and-tool-use/
|
||||
"""
|
||||
import json
|
||||
|
||||
|
||||
# Extract schema from output_format
|
||||
schema = output_format.get("schema")
|
||||
if not schema:
|
||||
return
|
||||
|
||||
|
||||
# Get messages from the request
|
||||
messages = anthropic_messages_request.get("messages", [])
|
||||
if not messages:
|
||||
return
|
||||
|
||||
|
||||
# Find the last user message
|
||||
last_user_message_idx = None
|
||||
for idx in range(len(messages) - 1, -1, -1):
|
||||
if messages[idx].get("role") == "user":
|
||||
last_user_message_idx = idx
|
||||
break
|
||||
|
||||
|
||||
if last_user_message_idx is None:
|
||||
return
|
||||
|
||||
|
||||
last_user_message = messages[last_user_message_idx]
|
||||
content = last_user_message.get("content", [])
|
||||
|
||||
|
||||
# Ensure content is a list
|
||||
if isinstance(content, str):
|
||||
content = [{"type": "text", "text": content}]
|
||||
last_user_message["content"] = content
|
||||
|
||||
|
||||
# Add schema as text content to the message
|
||||
schema_text = {
|
||||
"type": "text",
|
||||
"text": json.dumps(schema)
|
||||
}
|
||||
schema_text = {"type": "text", "text": json.dumps(schema)}
|
||||
content.append(schema_text)
|
||||
|
||||
def transform_anthropic_messages_request(
|
||||
|
|
@ -336,14 +345,14 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
output_format=output_format,
|
||||
anthropic_messages_request=anthropic_messages_request,
|
||||
)
|
||||
|
||||
|
||||
# 6. AUTO-INJECT beta headers based on features used
|
||||
anthropic_model_info = AnthropicModelInfo()
|
||||
tools = anthropic_messages_optional_request_params.get("tools")
|
||||
messages_typed = cast(List[AllMessageValues], messages)
|
||||
tool_search_used = anthropic_model_info.is_tool_search_used(tools)
|
||||
programmatic_tool_calling_used = anthropic_model_info.is_programmatic_tool_calling_used(
|
||||
tools
|
||||
programmatic_tool_calling_used = (
|
||||
anthropic_model_info.is_programmatic_tool_calling_used(tools)
|
||||
)
|
||||
input_examples_used = anthropic_model_info.is_input_examples_used(tools)
|
||||
|
||||
|
|
@ -376,8 +385,7 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
|
||||
if beta_set:
|
||||
anthropic_messages_request["anthropic_beta"] = list(beta_set)
|
||||
|
||||
|
||||
|
||||
return anthropic_messages_request
|
||||
|
||||
def get_async_streaming_response_iterator(
|
||||
|
|
@ -395,7 +403,7 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
)
|
||||
# Convert decoded Bedrock events to Server-Sent Events expected by Anthropic clients.
|
||||
return self.bedrock_sse_wrapper(
|
||||
completion_stream=completion_stream,
|
||||
completion_stream=completion_stream,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
request_body=request_body,
|
||||
)
|
||||
|
|
@ -414,14 +422,14 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
|
||||
BaseAnthropicMessagesStreamingIterator,
|
||||
)
|
||||
|
||||
handler = BaseAnthropicMessagesStreamingIterator(
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
request_body=request_body,
|
||||
)
|
||||
|
||||
|
||||
async for chunk in handler.async_sse_wrapper(completion_stream):
|
||||
yield chunk
|
||||
|
||||
|
||||
|
||||
class AmazonAnthropicClaudeMessagesStreamDecoder(AWSEventStreamDecoder):
|
||||
|
|
@ -433,7 +441,7 @@ class AmazonAnthropicClaudeMessagesStreamDecoder(AWSEventStreamDecoder):
|
|||
Iterator to return Bedrock invoke response in anthropic /messages format
|
||||
"""
|
||||
super().__init__(model=model)
|
||||
self.DEFAULT_CHUNK_SIZE = 1024
|
||||
self.DEFAULT_CHUNK_SIZE = LITELLM_CONSTANT_STREAM_CHUNK_SIZE
|
||||
|
||||
def _chunk_parser(
|
||||
self, chunk_data: dict
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union, cast
|
|||
import httpx
|
||||
from httpx._models import Headers
|
||||
|
||||
from litellm.constants import LITELLM_CONSTANT_STREAM_CHUNK_SIZE
|
||||
from litellm.litellm_core_utils.logging_utils import track_llm_api_timing
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
|
|
@ -154,7 +155,7 @@ class SagemakerChatConfig(OpenAIGPTConfig, BaseAWSLLM):
|
|||
|
||||
custom_stream_decoder = AWSEventStreamDecoder(model="", is_messages_api=True)
|
||||
completion_stream = custom_stream_decoder.iter_bytes(
|
||||
response.iter_bytes(chunk_size=1024)
|
||||
response.iter_bytes(chunk_size=LITELLM_CONSTANT_STREAM_CHUNK_SIZE)
|
||||
)
|
||||
|
||||
streaming_response = CustomStreamWrapper(
|
||||
|
|
@ -204,7 +205,7 @@ class SagemakerChatConfig(OpenAIGPTConfig, BaseAWSLLM):
|
|||
|
||||
custom_stream_decoder = AWSEventStreamDecoder(model="", is_messages_api=True)
|
||||
completion_stream = custom_stream_decoder.aiter_bytes(
|
||||
response.aiter_bytes(chunk_size=1024)
|
||||
response.aiter_bytes(chunk_size=LITELLM_CONSTANT_STREAM_CHUNK_SIZE)
|
||||
)
|
||||
|
||||
streaming_response = CustomStreamWrapper(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue