mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
6a0d03914c
* test: drop the cwd-relative sys.path.insert calls from the test suite
TQ003 stands at 1,077 across 1,058 files, and 1,015 of them are the same shape:
sys.path.insert(0, os.path.abspath("../..")) and its deeper siblings. The
argument resolves against the working directory rather than the file, so from
the repo root, where every job runs pytest, it inserts the directory two levels
above the checkout. It has never pointed at litellm. The package is installed
into the environment anyway, which is what actually makes the import work, and
what the rule's message has said all along.
Removing them leaves 1,634 imports of sys and os with no remaining reference,
and those go too, except where another test module imports the name back out of
the file. The rest of TQ003 is 62 call sites that resolve against __file__ or a
variable, which are a different question and are left alone.
Collection is identical either way: 45,871 tests and the same 51 pre-existing
collection errors before and after, and ruff reports no new undefined name.
* test: drop the duplicate imports the sys.path sweep exposed to F811
* test(pre-call-utils): restore the os import the new bedrock tests need
310 lines
12 KiB
Python
310 lines
12 KiB
Python
import json
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
import litellm
|
|
from litellm.llms.bedrock.chat import BedrockConverseLLM
|
|
from litellm.llms.bedrock.chat.converse_handler import make_sync_call
|
|
from litellm.llms.bedrock.common_utils import _get_all_bedrock_regions
|
|
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
|
|
|
|
|
|
|
def test_encode_model_id_with_inference_profile():
|
|
"""
|
|
Test instance profile is properly encoded when used as a model
|
|
"""
|
|
test_model = "arn:aws:bedrock:us-east-1:12345678910:application-inference-profile/ujdtmcirjhevpi"
|
|
expected_model = "arn%3Aaws%3Abedrock%3Aus-east-1%3A12345678910%3Aapplication-inference-profile%2Fujdtmcirjhevpi"
|
|
bedrock_converse_llm = BedrockConverseLLM()
|
|
returned_model = bedrock_converse_llm.encode_model_id(test_model)
|
|
assert expected_model == returned_model
|
|
|
|
|
|
class TestBedrockRegionInModelPath:
|
|
"""
|
|
Tests for region extraction from bedrock/{region}/{model} path format.
|
|
|
|
When a user passes model="bedrock/ap-northeast-1/moonshotai.kimi-k2.5",
|
|
get_llm_provider strips "bedrock/" and passes "ap-northeast-1/moonshotai.kimi-k2.5"
|
|
to the converse handler. The handler must:
|
|
1. Strip the region from modelId (so AWS gets "moonshotai.kimi-k2.5", not "ap-northeast-1%2Fmoonshotai.kimi-k2.5")
|
|
2. Use the extracted region as aws_region_name for the API call
|
|
"""
|
|
|
|
@pytest.mark.parametrize(
|
|
"model,expected_model_id,expected_region",
|
|
[
|
|
# Region embedded in path — both modelId and region must be extracted
|
|
(
|
|
"ap-northeast-1/moonshotai.kimi-k2.5",
|
|
"moonshotai.kimi-k2.5",
|
|
"ap-northeast-1",
|
|
),
|
|
(
|
|
"us-east-1/moonshotai.kimi-k2.5",
|
|
"moonshotai.kimi-k2.5",
|
|
"us-east-1",
|
|
),
|
|
(
|
|
"us-west-2/anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
"anthropic.claude-haiku-4-5-20251001-v1%3A0",
|
|
"us-west-2",
|
|
),
|
|
# No region in path — modelId unchanged, no region injected
|
|
(
|
|
"moonshotai.kimi-k2.5",
|
|
"moonshotai.kimi-k2.5",
|
|
None,
|
|
),
|
|
# Cross-region inference prefix (us., eu., ap.) — not a region path segment
|
|
(
|
|
"us.anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
"us.anthropic.claude-haiku-4-5-20251001-v1%3A0",
|
|
None,
|
|
),
|
|
],
|
|
)
|
|
def test_region_and_model_id_extraction(
|
|
self, model, expected_model_id, expected_region
|
|
):
|
|
"""
|
|
Verify that completion() correctly extracts both modelId and aws_region_name
|
|
from the bedrock/{region}/{model} path format.
|
|
"""
|
|
bedrock_converse_llm = BedrockConverseLLM()
|
|
optional_params: dict = {}
|
|
|
|
# Simulate the modelId + region extraction logic from completion()
|
|
_model_for_id = model
|
|
_stripped = _model_for_id
|
|
for rp in ["bedrock/converse/", "bedrock/", "converse/"]:
|
|
if _stripped.startswith(rp):
|
|
_stripped = _stripped[len(rp) :]
|
|
break
|
|
|
|
_region_from_model = None
|
|
_potential_region = _stripped.split("/", 1)[0]
|
|
if _potential_region in _get_all_bedrock_regions() and "/" in _stripped:
|
|
_region_from_model = _potential_region
|
|
_stripped = _stripped.split("/", 1)[1]
|
|
_model_for_id = _stripped
|
|
|
|
for _nova_prefix in ["nova-2/", "nova/"]:
|
|
if _stripped.startswith(_nova_prefix):
|
|
_model_for_id = _model_for_id.replace(_nova_prefix, "", 1)
|
|
break
|
|
|
|
model_id = bedrock_converse_llm.encode_model_id(model_id=_model_for_id)
|
|
if _region_from_model is not None and "aws_region_name" not in optional_params:
|
|
optional_params["aws_region_name"] = _region_from_model
|
|
|
|
assert (
|
|
model_id == expected_model_id
|
|
), f"modelId mismatch for {model!r}: got {model_id!r}, expected {expected_model_id!r}"
|
|
assert (
|
|
optional_params.get("aws_region_name") == expected_region
|
|
), f"region mismatch for {model!r}: got {optional_params.get('aws_region_name')!r}, expected {expected_region!r}"
|
|
|
|
def test_explicit_aws_region_name_not_overridden(self):
|
|
"""
|
|
If aws_region_name is already set in optional_params, the region in the
|
|
model path must NOT override it.
|
|
"""
|
|
bedrock_converse_llm = BedrockConverseLLM()
|
|
optional_params = {"aws_region_name": "eu-west-1"}
|
|
model = "ap-northeast-1/moonshotai.kimi-k2.5"
|
|
|
|
_model_for_id = model
|
|
_stripped = model
|
|
_region_from_model = None
|
|
_potential_region = _stripped.split("/", 1)[0]
|
|
if _potential_region in _get_all_bedrock_regions() and "/" in _stripped:
|
|
_region_from_model = _potential_region
|
|
_stripped = _stripped.split("/", 1)[1]
|
|
_model_for_id = _stripped
|
|
|
|
model_id = bedrock_converse_llm.encode_model_id(model_id=_model_for_id)
|
|
if _region_from_model is not None and "aws_region_name" not in optional_params:
|
|
optional_params["aws_region_name"] = _region_from_model
|
|
|
|
# modelId is still correctly stripped
|
|
assert model_id == "moonshotai.kimi-k2.5"
|
|
# explicitly set region is preserved
|
|
assert optional_params["aws_region_name"] == "eu-west-1"
|
|
|
|
|
|
def _stream_completion_with_spied_iter_bytes(model: str, **kwargs) -> MagicMock:
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_response.iter_bytes = MagicMock(return_value=iter([]))
|
|
client = HTTPHandler()
|
|
client.post = MagicMock(return_value=mock_response)
|
|
|
|
litellm.completion(
|
|
model=model,
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
stream=True,
|
|
client=client,
|
|
aws_access_key_id="fake",
|
|
aws_secret_access_key="fake",
|
|
aws_region_name="us-east-1",
|
|
**kwargs,
|
|
)
|
|
return mock_response.iter_bytes
|
|
|
|
|
|
def test_make_sync_call_does_not_rechunk_stream_by_default():
|
|
"""Re-chunking the event stream into fixed 1024-byte blocks holds small
|
|
early events in httpx's ByteChunker until 1024 bytes accumulate, delaying
|
|
time-to-first-chunk by the whole generation when Bedrock trickles bytes
|
|
(e.g. buffered tool-use streams)."""
|
|
response = MagicMock()
|
|
response.status_code = 200
|
|
client = MagicMock()
|
|
client.post = MagicMock(return_value=response)
|
|
|
|
make_sync_call(
|
|
client=client,
|
|
api_base="https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-sonnet-4-6/converse-stream",
|
|
headers={},
|
|
data="{}",
|
|
model="anthropic.claude-sonnet-4-6",
|
|
messages=[],
|
|
logging_obj=MagicMock(),
|
|
)
|
|
|
|
response.iter_bytes.assert_called_once_with(chunk_size=None)
|
|
|
|
|
|
def test_make_sync_call_honors_explicit_stream_chunk_size():
|
|
response = MagicMock()
|
|
response.status_code = 200
|
|
client = MagicMock()
|
|
client.post = MagicMock(return_value=response)
|
|
|
|
make_sync_call(
|
|
client=client,
|
|
api_base="https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-sonnet-4-6/converse-stream",
|
|
headers={},
|
|
data="{}",
|
|
model="anthropic.claude-sonnet-4-6",
|
|
messages=[],
|
|
logging_obj=MagicMock(),
|
|
stream_chunk_size=2048,
|
|
)
|
|
|
|
response.iter_bytes.assert_called_once_with(chunk_size=2048)
|
|
|
|
|
|
def _converse_response_body() -> dict:
|
|
return {
|
|
"output": {"message": {"role": "assistant", "content": [{"text": "hi"}]}},
|
|
"stopReason": "end_turn",
|
|
"usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2},
|
|
}
|
|
|
|
|
|
def test_converse_completion_forwards_bedrock_response_headers():
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_response.json = MagicMock(return_value=_converse_response_body())
|
|
mock_response.text = json.dumps(_converse_response_body())
|
|
mock_response.headers = httpx.Headers({"x-amzn-requestid": "req-123"})
|
|
client = HTTPHandler()
|
|
client.post = MagicMock(return_value=mock_response)
|
|
|
|
response = litellm.completion(
|
|
model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
client=client,
|
|
aws_access_key_id="fake",
|
|
aws_secret_access_key="fake",
|
|
aws_region_name="us-east-1",
|
|
)
|
|
|
|
assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-123"
|
|
|
|
|
|
def test_converse_streaming_forwards_bedrock_response_headers():
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_response.iter_bytes = MagicMock(return_value=iter([]))
|
|
mock_response.headers = httpx.Headers({"x-amzn-requestid": "req-456"})
|
|
client = HTTPHandler()
|
|
client.post = MagicMock(return_value=mock_response)
|
|
|
|
response = litellm.completion(
|
|
model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
stream=True,
|
|
client=client,
|
|
aws_access_key_id="fake",
|
|
aws_secret_access_key="fake",
|
|
aws_region_name="us-east-1",
|
|
)
|
|
|
|
assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-456"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_converse_completion_forwards_bedrock_response_headers():
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_response.json = MagicMock(return_value=_converse_response_body())
|
|
mock_response.text = json.dumps(_converse_response_body())
|
|
mock_response.headers = httpx.Headers({"x-amzn-requestid": "req-abc"})
|
|
client = AsyncHTTPHandler()
|
|
client.post = AsyncMock(return_value=mock_response)
|
|
|
|
response = await litellm.acompletion(
|
|
model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
client=client,
|
|
aws_access_key_id="fake",
|
|
aws_secret_access_key="fake",
|
|
aws_region_name="us-east-1",
|
|
)
|
|
|
|
assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-abc"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_converse_streaming_forwards_bedrock_response_headers():
|
|
async def _no_bytes(chunk_size=None):
|
|
return
|
|
yield b""
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_response.aiter_bytes = _no_bytes
|
|
mock_response.headers = httpx.Headers({"x-amzn-requestid": "req-def"})
|
|
client = AsyncHTTPHandler()
|
|
client.post = AsyncMock(return_value=mock_response)
|
|
|
|
response = await litellm.acompletion(
|
|
model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
stream=True,
|
|
client=client,
|
|
aws_access_key_id="fake",
|
|
aws_secret_access_key="fake",
|
|
aws_region_name="us-east-1",
|
|
)
|
|
|
|
assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-def"
|
|
|
|
|
|
def test_completion_plumbs_stream_chunk_size_through_converse():
|
|
iter_bytes_spy = _stream_completion_with_spied_iter_bytes(
|
|
model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0"
|
|
)
|
|
iter_bytes_spy.assert_called_once_with(chunk_size=None)
|
|
|
|
iter_bytes_spy = _stream_completion_with_spied_iter_bytes(
|
|
model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
|
|
stream_chunk_size=2048,
|
|
)
|
|
iter_bytes_spy.assert_called_once_with(chunk_size=2048)
|