mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
test: assert the responses bridge forwards aws_region_name
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
03a650c8a9
commit
72abd11b4c
1 changed files with 59 additions and 0 deletions
|
|
@ -0,0 +1,59 @@
|
|||
from datetime import datetime
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.completion_extras.litellm_responses_transformation.handler import (
|
||||
ResponsesToCompletionBridgeHandler,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
MODEL = "openai.gpt-5.5"
|
||||
REGION = "us-east-2"
|
||||
|
||||
|
||||
def _bedrock_mantle_kwargs() -> dict:
|
||||
messages = [{"role": "user", "content": "hi"}]
|
||||
logging_obj = LiteLLMLogging(
|
||||
litellm_call_id="test-call",
|
||||
call_type="acompletion",
|
||||
model=MODEL,
|
||||
messages=messages,
|
||||
function_id="fn-id",
|
||||
stream=False,
|
||||
start_time=datetime.now(),
|
||||
)
|
||||
return {
|
||||
"model": MODEL,
|
||||
"custom_llm_provider": "bedrock_mantle",
|
||||
"messages": messages,
|
||||
"optional_params": {},
|
||||
"litellm_params": {
|
||||
"aws_region_name": REGION,
|
||||
"api_base": "https://bedrock-mantle.us-east-1.api.aws/v1",
|
||||
"custom_llm_provider": "bedrock_mantle",
|
||||
},
|
||||
"headers": {},
|
||||
"model_response": ModelResponse(),
|
||||
"logging_obj": logging_obj,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acompletion_forwards_aws_region_name_to_aresponses():
|
||||
bridge = ResponsesToCompletionBridgeHandler()
|
||||
cached = ModelResponse(id="chatcmpl-cached", model=MODEL)
|
||||
|
||||
async def _fake_aresponses(**kwargs):
|
||||
_fake_aresponses.kwargs = kwargs
|
||||
return cached
|
||||
|
||||
_fake_aresponses.kwargs = {}
|
||||
|
||||
with patch("litellm.aresponses", _fake_aresponses):
|
||||
result = await bridge.acompletion(**_bedrock_mantle_kwargs())
|
||||
|
||||
assert result is cached
|
||||
assert _fake_aresponses.kwargs["aws_region_name"] == REGION
|
||||
assert _fake_aresponses.kwargs["custom_llm_provider"] == "bedrock_mantle"
|
||||
Loading…
Add table
Reference in a new issue