scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 79

Merge 507b00ea02 into 3dac3f7a36

Browse source
This commit is contained in:
Sainyam Kapoor 2026-09-02 09:01:33 +09:00 • committed by GitHub
commit e5a1857e6b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 160 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/llms/base_llm/chat/transformation.py
View file

@ -385,6 +385,7 @@ class BaseConfig(ABC):
client: AsyncHTTPHandler | None = None,
json_mode: bool | None = None,
signed_json_body: bytes | None = None,
timeout: float | httpx.Timeout | None = None,
) -> "CustomStreamWrapper":
raise NotImplementedError
@ -400,6 +401,7 @@ class BaseConfig(ABC):
client: HTTPHandler | AsyncHTTPHandler | None = None,
json_mode: bool | None = None,
signed_json_body: bytes | None = None,
timeout: float | httpx.Timeout | None = None,
) -> "CustomStreamWrapper":
raise NotImplementedError

2
litellm/llms/bedrock/chat/agentcore/transformation.py
View file

@ -643,6 +643,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
client: Union[HTTPHandler, "AsyncHTTPHandler"] | None = None,
json_mode: bool | None = None,
signed_json_body: bytes | None = None,
timeout: float | httpx.Timeout | None = None,
) -> "CustomStreamWrapper":
"""
Simplified sync streaming - returns a generator that yields ModelResponse chunks.
@ -856,6 +857,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
client: Optional["AsyncHTTPHandler"] = None,
json_mode: bool | None = None,
signed_json_body: bytes | None = None,
timeout: float | httpx.Timeout | None = None,
) -> "CustomStreamWrapper":
"""
Simplified async streaming - returns an async generator that yields ModelResponse chunks.

4
litellm/llms/bedrock/chat/converse_handler.py
View file

@ -35,6 +35,7 @@ def make_sync_call(
json_mode: bool | None = False,
fake_stream: bool = False,
stream_chunk_size: int | None = None,
timeout: float | httpx.Timeout | None = None,
) -> tuple[Any, httpx.Headers]:
if client is None:
client = _get_httpx_client() # Create a new client if none provided
@ -45,6 +46,7 @@ def make_sync_call(
data=data,
stream=not fake_stream,
logging_obj=logging_obj,
timeout=timeout,
)
if response.status_code != 200:
@ -145,6 +147,7 @@ class BedrockConverseLLM(BaseAWSLLM):
fake_stream=fake_stream,
json_mode=json_mode,
stream_chunk_size=stream_chunk_size,
timeout=timeout,
)
streaming_response: Final = CustomStreamWrapper(
completion_stream=completion_stream,
@ -555,6 +558,7 @@ class BedrockConverseLLM(BaseAWSLLM):
json_mode=json_mode,
fake_stream=fake_stream,
stream_chunk_size=stream_chunk_size,
timeout=timeout,
)
streaming_response: Final = CustomStreamWrapper(
completion_stream=completion_stream,

4
litellm/llms/bedrock/chat/invoke_handler.py
View file

@ -163,6 +163,7 @@ async def make_call(
json_mode: bool | None = False,
bedrock_invoke_provider: litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL | None = None,
stream_chunk_size: int | None = None,
timeout: float | httpx.Timeout | None = None,
) -> tuple[Any, httpx.Headers]:
try:
if client is None:
@ -181,6 +182,7 @@ async def make_call(
data=data,
stream=not fake_stream,
logging_obj=logging_obj,
timeout=timeout,
)
if response.status_code != 200:
@ -248,6 +250,7 @@ def make_sync_call(
json_mode: bool | None = False,
bedrock_invoke_provider: litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL | None = None,
stream_chunk_size: int | None = None,
timeout: float | httpx.Timeout | None = None,
) -> tuple[Any, httpx.Headers]:
try:
if client is None:
@ -265,6 +268,7 @@ def make_sync_call(
data=signed_json_body if signed_json_body is not None else data,
stream=not fake_stream,
logging_obj=logging_obj,
timeout=timeout,
)
if response.status_code != 200:

4
litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py
View file

@ -446,6 +446,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
client: AsyncHTTPHandler | None = None,
json_mode: bool | None = None,
signed_json_body: bytes | None = None,
timeout: float | httpx.Timeout | None = None,
) -> CustomStreamWrapper:
completion_stream, response_headers = await make_call(
client=client,
@ -458,6 +459,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
fake_stream=True if "ai21" in api_base else False,
bedrock_invoke_provider=self.get_bedrock_invoke_provider(model),
json_mode=json_mode,
timeout=timeout,
)
streaming_response: Final = CustomStreamWrapper(
completion_stream=completion_stream,
@ -481,6 +483,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
client: HTTPHandler | AsyncHTTPHandler | None = None,
json_mode: bool | None = None,
signed_json_body: bytes | None = None,
timeout: float | httpx.Timeout | None = None,
) -> CustomStreamWrapper:
sync_client: Final = (
_get_httpx_client(params={}) if client is None or isinstance(client, AsyncHTTPHandler) else client
@ -497,6 +500,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
fake_stream=True if "ai21" in api_base else False,
bedrock_invoke_provider=self.get_bedrock_invoke_provider(model),
json_mode=json_mode,
timeout=timeout,
)
streaming_response: Final = CustomStreamWrapper(
completion_stream=completion_stream,

2
litellm/llms/bytez/chat/transformation.py
View file

@ -259,6 +259,7 @@ class BytezChatConfig(BaseConfig):
client: HTTPHandler | AsyncHTTPHandler | None = None,
json_mode: bool | None = None,
signed_json_body: bytes | None = None,
timeout: float | httpx.Timeout | None = None,
) -> "BytezCustomStreamWrapper":
if client is None or isinstance(client, AsyncHTTPHandler):
client = _get_httpx_client(params={})
@ -301,6 +302,7 @@ class BytezChatConfig(BaseConfig):
client: HTTPHandler | AsyncHTTPHandler | None = None,
json_mode: bool | None = None,
signed_json_body: bytes | None = None,
timeout: float | httpx.Timeout | None = None,
) -> "BytezCustomStreamWrapper":
if client is None or isinstance(client, HTTPHandler):
client = get_async_httpx_client(llm_provider=LlmProviders.BYTEZ, params={})

2
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -624,6 +624,7 @@ class BaseLLMHTTPHandler:
messages=messages,
client=client,
json_mode=json_mode,
timeout=timeout,
)
completion_stream, headers = self.make_sync_call(
provider_config=provider_config,
@ -787,6 +788,7 @@ class BaseLLMHTTPHandler:
client=client,
json_mode=json_mode,
signed_json_body=signed_json_body,
timeout=timeout,
)
completion_stream, _response_headers = await self.make_async_call_stream_helper(

2
litellm/llms/langgraph/chat/transformation.py
View file

@ -286,6 +286,7 @@ class LangGraphConfig(BaseConfig):
client: Union[HTTPHandler, "AsyncHTTPHandler"] | None = None,
json_mode: bool | None = None,
signed_json_body: bytes | None = None,
timeout: float | httpx.Timeout | None = None,
) -> CustomStreamWrapper:
"""
Get a CustomStreamWrapper for synchronous streaming.
@ -345,6 +346,7 @@ class LangGraphConfig(BaseConfig):
client: Optional["AsyncHTTPHandler"] = None,
json_mode: bool | None = None,
signed_json_body: bytes | None = None,
timeout: float | httpx.Timeout | None = None,
) -> CustomStreamWrapper:
"""
Get a CustomStreamWrapper for asynchronous streaming.

2
litellm/llms/oci/chat/transformation.py
View file

@ -643,6 +643,7 @@ class OCIChatConfig(BaseConfig):
client: HTTPHandler | AsyncHTTPHandler | None = None,
json_mode: bool | None = None,
signed_json_body: bytes | None = None,
timeout: float | httpx.Timeout | None = None,
) -> "OCIStreamWrapper":
if client is None or isinstance(client, AsyncHTTPHandler):
client = _get_httpx_client(params={})
@ -682,6 +683,7 @@ class OCIChatConfig(BaseConfig):
client: HTTPHandler | AsyncHTTPHandler | None = None,
json_mode: bool | None = None,
signed_json_body: bytes | None = None,
timeout: float | httpx.Timeout | None = None,
) -> "OCIStreamWrapper":
if client is None or isinstance(client, HTTPHandler):
client = get_async_httpx_client(llm_provider=LlmProviders.OCI, params={})

2
litellm/llms/sagemaker/chat/transformation.py
View file

@ -149,6 +149,7 @@ class SagemakerChatConfig(OpenAIGPTConfig, BaseAWSLLM):
client: HTTPHandler | AsyncHTTPHandler | None = None,
json_mode: bool | None = None,
signed_json_body: bytes | None = None,
timeout: float | httpx.Timeout | None = None,
) -> CustomStreamWrapper:
if client is None or isinstance(client, AsyncHTTPHandler):
client = _get_httpx_client(params={})
@ -191,6 +192,7 @@ class SagemakerChatConfig(OpenAIGPTConfig, BaseAWSLLM):
client: HTTPHandler | AsyncHTTPHandler | None = None,
json_mode: bool | None = None,
signed_json_body: bytes | None = None,
timeout: float | httpx.Timeout | None = None,
) -> CustomStreamWrapper:
if client is None or isinstance(client, HTTPHandler):
try:

2
litellm/llms/vertex_ai/agent_engine/transformation.py
View file

@ -366,6 +366,7 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase):
client: Union[HTTPHandler, "AsyncHTTPHandler"] | None = None,
json_mode: bool | None = None,
signed_json_body: bytes | None = None,
timeout: float | httpx.Timeout | None = None,
) -> "CustomStreamWrapper":
"""Get a CustomStreamWrapper for synchronous streaming."""
from litellm.llms.custom_httpx.http_handler import (
@ -424,6 +425,7 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase):
client: Optional["AsyncHTTPHandler"] = None,
json_mode: bool | None = None,
signed_json_body: bytes | None = None,
timeout: float | httpx.Timeout | None = None,
) -> "CustomStreamWrapper":
"""Get a CustomStreamWrapper for asynchronous streaming."""
from litellm.llms.custom_httpx.http_handler import (

66
tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py
View file

@ -6,11 +6,12 @@ extension, and AWS credential resolution is stubbed so nothing reaches STS.
from __future__ import annotations
from unittest.mock import MagicMock, patch
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
import litellm
from botocore.credentials import Credentials
from litellm.llms.bedrock.chat.converse_handler import BedrockConverseLLM
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
@ -487,3 +488,66 @@ def test_post_call_is_not_logged_twice_when_the_sync_rust_call_declines():
assert response.choices[0].message.content == "hi"
assert len(calls["post_call"]) == 1
assert "hi" in calls["post_call"][0]["original_response"]
@pytest.mark.parametrize(
"timeout",
[
pytest.param(2.25, id="numeric"),
pytest.param(httpx.Timeout(2.5), id="httpx-timeout"),
],
)
@pytest.mark.asyncio
async def test_async_converse_streaming_forwards_timeout_to_existing_client(timeout):
async def _no_bytes(chunk_size=None):
return
yield b""
response = MagicMock()
response.status_code = 200
response.aiter_bytes = _no_bytes
response.headers = httpx.Headers()
client = AsyncHTTPHandler()
client.post = AsyncMock(return_value=response)
await litellm.acompletion(
model="bedrock/anthropic.claude-haiku-4-5-20251001-v1:0",
messages=[{"role": "user", "content": "hi"}],
stream=True,
timeout=timeout,
client=client,
aws_access_key_id="fake",
aws_secret_access_key="fake",
aws_region_name="us-east-1",
)
assert client.post.await_args.kwargs["timeout"] == timeout
@pytest.mark.parametrize(
"timeout",
[
pytest.param(3.25, id="numeric"),
pytest.param(httpx.Timeout(3.5), id="httpx-timeout"),
],
)
def test_sync_converse_streaming_forwards_timeout_to_existing_client(timeout):
response = MagicMock()
response.status_code = 200
response.iter_bytes = MagicMock(return_value=iter(()))
response.headers = httpx.Headers()
client = HTTPHandler()
client.post = MagicMock(return_value=response)
litellm.completion(
model="bedrock/anthropic.claude-haiku-4-5-20251001-v1:0",
messages=[{"role": "user", "content": "hi"}],
stream=True,
timeout=timeout,
client=client,
aws_access_key_id="fake",
aws_secret_access_key="fake",
aws_region_name="us-east-1",
)
assert client.post.call_args.kwargs["timeout"] == timeout

67
tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py
View file

@ -496,3 +496,70 @@ async def test_async_invoke_streaming_forwards_bedrock_response_headers():
assert stream._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-987"
@pytest.mark.parametrize(
"timeout",
[
pytest.param(1.25, id="numeric"),
pytest.param(httpx.Timeout(1.5), id="httpx-timeout"),
],
)
@pytest.mark.asyncio
async def test_async_invoke_streaming_forwards_timeout_to_existing_client(timeout):
async def _no_bytes(chunk_size=None):
return
yield b""
response = MagicMock()
response.status_code = 200
response.aiter_bytes = _no_bytes
response.headers = httpx.Headers()
client = AsyncHTTPHandler()
client.post = AsyncMock(return_value=response)
await litellm.acompletion(
model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0",
messages=[{"role": "user", "content": "hi"}],
stream=True,
timeout=timeout,
client=client,
aws_access_key_id="fake",
aws_secret_access_key="fake",
aws_region_name="us-east-1",
)
assert client.post.await_args.kwargs["timeout"] == timeout
@pytest.mark.asyncio
async def test_make_call_sets_timeout_on_httpx_request_transport():
transport = MagicMock(return_value=httpx.Response(200, content=b""))
timeout = httpx.Timeout(connect=1.0, read=2.0, write=3.0, pool=4.0)
client = AsyncHTTPHandler()
await client.client.aclose()
client.client = httpx.AsyncClient(transport=httpx.MockTransport(transport))
try:
await make_call(
client=client,
api_base=(
"https://bedrock-runtime.us-east-1.amazonaws.com/model/"
"anthropic.claude-sonnet-4-6/invoke-with-response-stream"
),
headers={},
data="{}",
model="anthropic.claude-sonnet-4-6",
messages=[],
logging_obj=MagicMock(),
timeout=timeout,
)
finally:
await client.client.aclose()
assert transport.call_args.args[0].extensions["timeout"] == {
"connect": 1.0,
"read": 2.0,
"write": 3.0,
"pool": 4.0,
}