mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
fix(bedrock): route an ARN in the invoke model string and map Bedrock's 200 error body to a 400
This commit is contained in:
mateo-berri
parent
e81dd348a4
commit
ab02277d2d
8 changed files with 338 additions and 20 deletions
|
|
@ -669,7 +669,7 @@ class BaseAWSLLM(SignsRequestsWithAWS):
|
|||
# like "bedrock/invoke/arn:..." are fully stripped in one call.
|
||||
from litellm.llms.bedrock.common_utils import strip_bedrock_routing_prefix
|
||||
|
||||
model_id = strip_bedrock_routing_prefix(model_id)
|
||||
model_id = BaseAWSLLM._without_invoke_provider_spec(strip_bedrock_routing_prefix(model_id))
|
||||
# URL-encode ARNs so colons and slashes are safe in the URL path.
|
||||
if model_id.startswith("arn:"):
|
||||
model_id = BaseAWSLLM.encode_model_id(model_id=model_id)
|
||||
|
|
@ -696,6 +696,11 @@ class BaseAWSLLM(SignsRequestsWithAWS):
|
|||
model_id = BaseAWSLLM._get_model_id_from_model_with_spec(model_id, spec="nova")
|
||||
return model_id
|
||||
|
||||
@staticmethod
|
||||
def _without_invoke_provider_spec(model_id: str) -> str:
|
||||
_, _, arn = model_id.partition("/")
|
||||
return arn if arn.startswith("arn:") else model_id
|
||||
|
||||
@staticmethod
|
||||
def _get_model_id_from_model_with_spec(
|
||||
model: str,
|
||||
|
|
|
|||
|
|
@ -48,6 +48,7 @@ from litellm.types.utils import GenericStreamingChunk as GChunk
|
|||
from ..common_utils import (
|
||||
BedrockError,
|
||||
BedrockEventStreamResponseDict,
|
||||
bedrock_error_in_success_body,
|
||||
bedrock_event_stream_header,
|
||||
bedrock_event_stream_response,
|
||||
bedrock_stream_event_error_status,
|
||||
|
|
@ -192,6 +193,10 @@ def response_format_schema_from_request(data: str, json_mode: bool | None) -> Re
|
|||
return JsonModeRequest.model_validate(body).tool_schema()
|
||||
|
||||
|
||||
def _answered_json(response: httpx.Response) -> bool:
|
||||
return "application/json" in (response.headers.get("content-type") or "")
|
||||
|
||||
|
||||
def _stream_decoder(
|
||||
bedrock_invoke_provider: litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL | None,
|
||||
*,
|
||||
|
|
@ -259,6 +264,12 @@ async def make_call(
|
|||
response=response,
|
||||
)
|
||||
|
||||
if _answered_json(response):
|
||||
await response.aread()
|
||||
success_body_error: Final = bedrock_error_in_success_body(response)
|
||||
if success_body_error is not None:
|
||||
raise success_body_error
|
||||
|
||||
if fake_stream:
|
||||
model_response: Final[ModelResponse] = litellm.AmazonConverseConfig()._transform_response(
|
||||
model=model,
|
||||
|
|
@ -355,6 +366,12 @@ def make_sync_call(
|
|||
response=response,
|
||||
)
|
||||
|
||||
if _answered_json(response):
|
||||
response.read()
|
||||
success_body_error: Final = bedrock_error_in_success_body(response)
|
||||
if success_body_error is not None:
|
||||
raise success_body_error
|
||||
|
||||
if fake_stream:
|
||||
model_response: Final[ModelResponse] = litellm.AmazonConverseConfig()._transform_response(
|
||||
model=model,
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ from litellm.llms.bedrock.common_utils import (
|
|||
apply_bedrock_invoke_structured_output,
|
||||
bedrock_chat_uses_native_structured_output,
|
||||
bedrock_cleaned_output_schema,
|
||||
bedrock_error_in_success_body,
|
||||
bedrock_supports_tool_search,
|
||||
deployment_model_info,
|
||||
get_anthropic_beta_from_headers,
|
||||
|
|
@ -416,6 +417,9 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
|
|||
api_key: str | None = None,
|
||||
json_mode: bool | None = None,
|
||||
) -> ModelResponse:
|
||||
success_body_error: Final = bedrock_error_in_success_body(raw_response)
|
||||
if success_body_error is not None:
|
||||
raise success_body_error
|
||||
response: Final = AnthropicConfig.transform_response(
|
||||
self,
|
||||
model=model,
|
||||
|
|
|
|||
|
|
@ -22,7 +22,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
|
|||
from litellm.litellm_core_utils.url_utils import join_api_base
|
||||
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
|
||||
from litellm.llms.bedrock.chat.invoke_handler import make_call, make_sync_call
|
||||
from litellm.llms.bedrock.common_utils import BedrockError
|
||||
from litellm.llms.bedrock.common_utils import BedrockError, bedrock_error_in_success_body
|
||||
from litellm.llms.bedrock.request_metadata import (
|
||||
bedrock_request_metadata_headers,
|
||||
merge_bedrock_invoke_headers,
|
||||
|
|
@ -390,7 +390,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
|
|||
else: # amazon titan
|
||||
outputText = completion_response.get("results")[0].get("outputText")
|
||||
except Exception as e:
|
||||
raise BedrockError(
|
||||
raise bedrock_error_in_success_body(raw_response) or BedrockError(
|
||||
message=f"Error processing={raw_response.text}, Received error={e}",
|
||||
status_code=422,
|
||||
headers=raw_response.headers,
|
||||
|
|
@ -412,7 +412,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
|
|||
else:
|
||||
raise Exception()
|
||||
except Exception as e:
|
||||
raise BedrockError(
|
||||
raise bedrock_error_in_success_body(raw_response) or BedrockError(
|
||||
message=f"Error parsing received text={outputText}.\nError-{e}",
|
||||
status_code=raw_response.status_code,
|
||||
headers=raw_response.headers,
|
||||
|
|
|
|||
|
|
@ -28,7 +28,7 @@ if TYPE_CHECKING:
|
|||
from litellm.types.llms.bedrock import BedrockCreateBatchRequest
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, TypeAdapter, ValidationError
|
||||
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm import verbose_logger
|
||||
|
|
@ -118,6 +118,35 @@ class BedrockError(BaseLLMException):
|
|||
)
|
||||
|
||||
|
||||
class _BedrockServiceErrorOutput(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
error_type: str = Field(alias="__type")
|
||||
|
||||
|
||||
class _BedrockServiceErrorBody(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
Output: _BedrockServiceErrorOutput
|
||||
|
||||
|
||||
def bedrock_error_in_success_body(response: httpx.Response) -> BedrockError | None:
|
||||
if b'"__type"' not in response.content:
|
||||
return None
|
||||
try:
|
||||
error_type: Final = _BedrockServiceErrorBody.model_validate_json(response.content).Output.error_type
|
||||
except ValidationError:
|
||||
return None
|
||||
return BedrockError(
|
||||
status_code=400,
|
||||
message=(
|
||||
f"Bedrock answered HTTP 200 with the error {error_type}, which usually means the model id in the "
|
||||
"model string does not name a Bedrock model or profile this route can call."
|
||||
),
|
||||
headers=response.headers,
|
||||
)
|
||||
|
||||
|
||||
_BEDROCK_AWS_AUTH_PARAMETER_KEYS: Final[tuple[str, ...]] = (*AWS_AUTH_PARAM_KEYS, "aws_region_name")
|
||||
|
||||
|
||||
|
|
@ -1061,6 +1090,21 @@ def bedrock_capability_model(model: str, base_model: object) -> str:
|
|||
return model
|
||||
|
||||
|
||||
def invoke_model_with_base_model_provider(model: str, base_model: object) -> str:
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
|
||||
model_id: Final = strip_bedrock_routing_prefix(model)
|
||||
if (
|
||||
not isinstance(base_model, str)
|
||||
or not model_id.startswith("arn:")
|
||||
or BedrockModelInfo.get_bedrock_route(model) != "invoke"
|
||||
or BaseAWSLLM.get_bedrock_invoke_provider(model) is not None
|
||||
):
|
||||
return model
|
||||
provider: Final = BaseAWSLLM.get_bedrock_invoke_provider(strip_bedrock_routing_prefix(base_model))
|
||||
return model if provider is None else f"invoke/{provider}/{model_id}"
|
||||
|
||||
|
||||
def strip_bedrock_routing_prefix(model: str) -> str:
|
||||
"""Strip LiteLLM routing prefixes from model name."""
|
||||
for prefix in ["bedrock/", "chat_completions/", "converse/", "invoke/", "openai/", "mantle/", "nova-2/", "nova/"]:
|
||||
|
|
@ -1971,21 +2015,23 @@ class BedrockModelInfo(BaseLLMModelInfo):
|
|||
return None
|
||||
|
||||
|
||||
def get_bedrock_chat_config(model: str):
|
||||
def get_bedrock_chat_config(model: str, deployment_base_model: str | None = None):
|
||||
"""
|
||||
Helper function to get the appropriate Bedrock chat config based on model and route.
|
||||
|
||||
Args:
|
||||
model: The model name/identifier
|
||||
deployment_base_model: The deployment's base_model, which names the invoke provider for a bare ARN
|
||||
|
||||
Returns:
|
||||
The appropriate Bedrock config class instance
|
||||
"""
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
|
||||
bedrock_route: Final = BedrockModelInfo.get_bedrock_route(model)
|
||||
bedrock_invoke_provider: Final = BaseAWSLLM.get_bedrock_invoke_provider(model=model)
|
||||
base_model: Final = BedrockModelInfo.get_base_model(model)
|
||||
routed_model: Final = invoke_model_with_base_model_provider(model, deployment_base_model)
|
||||
bedrock_route: Final = BedrockModelInfo.get_bedrock_route(routed_model)
|
||||
bedrock_invoke_provider: Final = BaseAWSLLM.get_bedrock_invoke_provider(model=routed_model)
|
||||
base_model: Final = BedrockModelInfo.get_base_model(routed_model)
|
||||
|
||||
# Handle explicit routes first
|
||||
if bedrock_route == "claude_platform":
|
||||
|
|
|
|||
|
|
@ -129,6 +129,7 @@ from litellm.llms.bedrock.common_utils import (
|
|||
bedrock_route_for_request,
|
||||
bedrock_strict_tools_fallback,
|
||||
deployment_model_info,
|
||||
invoke_model_with_base_model_provider,
|
||||
without_bedrock_native_output_opt_in,
|
||||
without_bedrock_route_prefix,
|
||||
)
|
||||
|
|
@ -4539,10 +4540,11 @@ def _complete_bedrock(ctx: CompletionDispatchContext) -> _CompletionDispatchResu
|
|||
client=client,
|
||||
)
|
||||
else:
|
||||
invoke_model: Final = invoke_model_with_base_model_provider(model, litellm_params.get("base_model"))
|
||||
|
||||
def invoke_attempt(fallback: NativeOutputFallback | None) -> NativeOutputAttemptResult:
|
||||
return base_llm_http_handler.completion(
|
||||
model=model,
|
||||
model=invoke_model,
|
||||
stream=stream,
|
||||
messages=messages,
|
||||
acompletion=acompletion,
|
||||
|
|
@ -4565,10 +4567,10 @@ def _complete_bedrock(ctx: CompletionDispatchContext) -> _CompletionDispatchResu
|
|||
(
|
||||
partial(_native_output_fallback, ctx, invoke_params, dict(litellm_params))
|
||||
if isinstance(provider_config, AmazonAnthropicClaudeConfig)
|
||||
and provider_config.uses_opt_in_native_output(model, invoke_params, litellm_params)
|
||||
and provider_config.uses_opt_in_native_output(invoke_model, invoke_params, litellm_params)
|
||||
else None
|
||||
),
|
||||
model,
|
||||
invoke_model,
|
||||
)
|
||||
|
||||
return response
|
||||
|
|
|
|||
|
|
@ -9466,7 +9466,7 @@ class ProviderConfigManager:
|
|||
# Format: (factory_function, needs_model_parameter: bool)
|
||||
LlmProviders.OPENAI: (lambda: litellm.OpenAIGPTConfig(), False),
|
||||
LlmProviders.ANTHROPIC: (lambda: litellm.AnthropicConfig(), False),
|
||||
# AZURE is handled as a special case in get_provider_chat_config()
|
||||
# AZURE and BEDROCK are handled as special cases in get_provider_chat_config()
|
||||
# so that base_model can be threaded through for model-type detection.
|
||||
LlmProviders.AZURE_AI: (
|
||||
lambda model: ProviderConfigManager._get_azure_ai_config(model),
|
||||
|
|
@ -9476,10 +9476,6 @@ class ProviderConfigManager:
|
|||
lambda model: ProviderConfigManager._get_vertex_ai_config(model),
|
||||
True,
|
||||
),
|
||||
LlmProviders.BEDROCK: (
|
||||
lambda model: ProviderConfigManager._get_bedrock_config(model),
|
||||
True,
|
||||
),
|
||||
LlmProviders.COHERE: (
|
||||
lambda model: ProviderConfigManager._get_cohere_config(model),
|
||||
True,
|
||||
|
|
@ -9686,11 +9682,11 @@ class ProviderConfigManager:
|
|||
return litellm.VertexAILlama3Config()
|
||||
|
||||
@staticmethod
|
||||
def _get_bedrock_config(model: str) -> BaseConfig:
|
||||
def _get_bedrock_config(model: str, base_model: str | None) -> BaseConfig:
|
||||
"""Get Bedrock config based on model."""
|
||||
from litellm.llms.bedrock.common_utils import get_bedrock_chat_config
|
||||
|
||||
return get_bedrock_chat_config(model=model)
|
||||
return get_bedrock_chat_config(model=model, deployment_base_model=base_model)
|
||||
|
||||
@staticmethod
|
||||
def _get_bigmodel_chat_config() -> BaseConfig:
|
||||
|
|
@ -9782,9 +9778,11 @@ class ProviderConfigManager:
|
|||
if not OpenAIGPTConfig.is_openai_catalog_model(model):
|
||||
return OpenAIUnknownModelConfig()
|
||||
|
||||
# Handle Azure before the generic map so base_model can be threaded through
|
||||
# Handle Azure and Bedrock before the generic map so base_model can be threaded through
|
||||
if provider == LlmProviders.AZURE:
|
||||
return ProviderConfigManager._get_azure_config(model=model, base_model=base_model)
|
||||
if provider == LlmProviders.BEDROCK:
|
||||
return ProviderConfigManager._get_bedrock_config(model=model, base_model=base_model)
|
||||
|
||||
# Initialize provider config map lazily (avoids circular imports)
|
||||
if ProviderConfigManager._PROVIDER_CONFIG_MAP is None:
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ import itertools
|
|||
import json
|
||||
import re
|
||||
import struct
|
||||
import urllib.parse
|
||||
from collections.abc import AsyncIterator, Mapping, Sequence
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, Mock, patch
|
||||
|
|
@ -1982,3 +1983,248 @@ def test_invoke_stream_decoder_carries_the_guardrail_trace_of_the_message_stop_c
|
|||
assert (plain.provider_specific_fields or {}).get("amazon-bedrock-trace") is None
|
||||
assert traced.provider_specific_fields is not None
|
||||
assert traced.provider_specific_fields["amazon-bedrock-trace"] == trace
|
||||
|
||||
|
||||
_ARN_ACCOUNT: Final = "arn:aws:bedrock:us-east-1:123456789012"
|
||||
_APPLICATION_PROFILE_ARN: Final = f"{_ARN_ACCOUNT}:application-inference-profile/a1b2c3d4e5f6"
|
||||
_CLAUDE_INVOKE_ANSWER: Final = json.dumps(
|
||||
{
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude",
|
||||
"content": [{"type": "text", "text": "hi"}],
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1},
|
||||
}
|
||||
).encode()
|
||||
_CLAUDE_INVOKE_STREAM: Final = b"".join(
|
||||
_bedrock_event_stream_frame(chunk)
|
||||
for chunk in (
|
||||
{"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 1, "output_tokens": 0}}},
|
||||
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hi"}},
|
||||
{"type": "content_block_stop", "index": 0},
|
||||
{"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 1}},
|
||||
{"type": "message_stop"},
|
||||
)
|
||||
)
|
||||
_UNKNOWN_OPERATION_BODY: Final = (
|
||||
b'{"Output":{"__type":"com.amazon.coral.service#UnknownOperationException"},"Version":"1.0"}'
|
||||
)
|
||||
|
||||
|
||||
def _claude_invoke_upstream(arn: str, stream: bool) -> httpx.MockTransport:
|
||||
operation: Final = "invoke-with-response-stream" if stream else "invoke"
|
||||
expected_path: Final = f"/model/{urllib.parse.quote(arn, safe='')}/{operation}"
|
||||
|
||||
def answer(request: httpx.Request) -> httpx.Response:
|
||||
path: Final = request.url.raw_path.decode()
|
||||
body: Final = json.loads(request.content)
|
||||
if path != expected_path or "anthropic_version" not in body or body.get("stop_sequences") != ["END"]:
|
||||
return httpx.Response(404, json={"message": f"unexpected request {path} {body}"}, request=request)
|
||||
if stream:
|
||||
return httpx.Response(
|
||||
200,
|
||||
content=_CLAUDE_INVOKE_STREAM,
|
||||
headers={"content-type": "application/vnd.amazon.eventstream"},
|
||||
request=request,
|
||||
)
|
||||
return httpx.Response(
|
||||
200, content=_CLAUDE_INVOKE_ANSWER, headers={"content-type": "application/json"}, request=request
|
||||
)
|
||||
|
||||
return httpx.MockTransport(answer)
|
||||
|
||||
|
||||
def _answer_text(response: object) -> str:
|
||||
if isinstance(response, litellm.ModelResponse):
|
||||
return response.choices[0].message.content or ""
|
||||
assert isinstance(response, CustomStreamWrapper)
|
||||
return "".join(chunk.choices[0].delta.content or "" for chunk in response)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream", [False, True], ids=["non-stream", "stream"])
|
||||
@pytest.mark.parametrize(
|
||||
("model", "arn", "deployment"),
|
||||
[
|
||||
pytest.param(
|
||||
f"bedrock/invoke/anthropic/{_APPLICATION_PROFILE_ARN}",
|
||||
_APPLICATION_PROFILE_ARN,
|
||||
{},
|
||||
id="application-profile",
|
||||
),
|
||||
pytest.param(
|
||||
f"bedrock/invoke/anthropic/{_ARN_ACCOUNT}:provisioned-model/a1b2c3d4e5f6",
|
||||
f"{_ARN_ACCOUNT}:provisioned-model/a1b2c3d4e5f6",
|
||||
{},
|
||||
id="provisioned",
|
||||
),
|
||||
pytest.param(
|
||||
f"bedrock/invoke/anthropic/{_ARN_ACCOUNT}:custom-model/anthropic.claude-haiku-5-5/a1b2c3d4e5f6",
|
||||
f"{_ARN_ACCOUNT}:custom-model/anthropic.claude-haiku-5-5/a1b2c3d4e5f6",
|
||||
{},
|
||||
id="custom-model",
|
||||
),
|
||||
pytest.param(
|
||||
f"bedrock/invoke/{_APPLICATION_PROFILE_ARN}",
|
||||
_APPLICATION_PROFILE_ARN,
|
||||
{"base_model": "bedrock/us.anthropic.claude-haiku-5-5"},
|
||||
id="base-model",
|
||||
),
|
||||
pytest.param(
|
||||
f"bedrock/invoke/{_APPLICATION_PROFILE_ARN}",
|
||||
_APPLICATION_PROFILE_ARN,
|
||||
{"model_info": {"base_model": "us.anthropic.claude-haiku-5-5"}},
|
||||
id="model-info-base-model",
|
||||
),
|
||||
pytest.param(
|
||||
f"bedrock/invoke/{_ARN_ACCOUNT}:custom-model/anthropic.claude-haiku-5-5/a1b2c3d4e5f6",
|
||||
f"{_ARN_ACCOUNT}:custom-model/anthropic.claude-haiku-5-5/a1b2c3d4e5f6",
|
||||
{"base_model": "bedrock/meta.llama3-70b-instruct-v1:0"},
|
||||
id="custom-model-arn-over-other-base-model",
|
||||
),
|
||||
pytest.param(
|
||||
f"bedrock/invoke/claude/{_ARN_ACCOUNT}:custom-model/anthropic.claude-haiku-5-5/a1b2c3d4e5f6",
|
||||
f"{_ARN_ACCOUNT}:custom-model/anthropic.claude-haiku-5-5/a1b2c3d4e5f6",
|
||||
{},
|
||||
id="custom-model-under-unlisted-spec",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_an_arn_in_the_invoke_model_string_reaches_bedrock_as_the_encoded_model_id(
|
||||
_aws_test_credentials: None, model: str, arn: str, deployment: Mapping[str, object], stream: bool
|
||||
) -> None:
|
||||
response: Final = litellm.completion(
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
max_tokens=10,
|
||||
stop=["END"],
|
||||
stream=stream,
|
||||
client=HTTPHandler(client=httpx.Client(transport=_claude_invoke_upstream(arn, stream))),
|
||||
**deployment,
|
||||
)
|
||||
|
||||
assert _answer_text(response) == "hi"
|
||||
|
||||
|
||||
def _no_upstream_call() -> httpx.MockTransport:
|
||||
return httpx.MockTransport(
|
||||
lambda request: httpx.Response(500, json={"message": f"unexpected call to {request.url}"}, request=request)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "deployment"),
|
||||
[
|
||||
pytest.param(f"bedrock/invoke/{_APPLICATION_PROFILE_ARN}", {}, id="arn-without-base-model"),
|
||||
pytest.param(
|
||||
"bedrock/invoke/my-imported-model",
|
||||
{"base_model": "bedrock/us.anthropic.claude-haiku-5-5"},
|
||||
id="non-arn-with-base-model",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_an_invoke_model_with_no_provider_to_route_by_keeps_the_unknown_provider_error(
|
||||
_aws_test_credentials: None, model: str, deployment: Mapping[str, object]
|
||||
) -> None:
|
||||
with pytest.raises(litellm.NotFoundError, match="Unknown provider=None"):
|
||||
litellm.completion(
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
max_tokens=10,
|
||||
client=HTTPHandler(client=httpx.Client(transport=_no_upstream_call())),
|
||||
**deployment,
|
||||
)
|
||||
|
||||
|
||||
def _unknown_operation_upstream() -> httpx.MockTransport:
|
||||
return httpx.MockTransport(
|
||||
lambda request: httpx.Response(
|
||||
200, content=_UNKNOWN_OPERATION_BODY, headers={"content-type": "application/json"}, request=request
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _assert_unknown_operation_is_a_bad_request(error: litellm.BadRequestError) -> None:
|
||||
assert error.status_code == 400
|
||||
assert "com.amazon.coral.service#UnknownOperationException" in error.message
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream", [False, True], ids=["non-stream", "stream"])
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
"bedrock/invoke/anthropic.claude-haiku-5-5",
|
||||
"bedrock/invoke/mistral.mistral-large-2407-v1:0",
|
||||
"bedrock/invoke/cohere.command-r-plus-v1:0",
|
||||
],
|
||||
)
|
||||
def test_a_200_unknown_operation_body_is_a_bad_request(_aws_test_credentials: None, model: str, stream: bool) -> None:
|
||||
with pytest.raises(litellm.BadRequestError) as exc_info:
|
||||
_answer_text(
|
||||
litellm.completion(
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
max_tokens=10,
|
||||
stream=stream,
|
||||
client=HTTPHandler(client=httpx.Client(transport=_unknown_operation_upstream())),
|
||||
)
|
||||
)
|
||||
|
||||
_assert_unknown_operation_is_a_bad_request(exc_info.value)
|
||||
|
||||
|
||||
async def _drain_async_unknown_operation_stream() -> None:
|
||||
stream: Final = await litellm.acompletion(
|
||||
model="bedrock/invoke/anthropic.claude-haiku-5-5",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
max_tokens=10,
|
||||
stream=True,
|
||||
client=AsyncHTTPHandler(transport=_unknown_operation_upstream()),
|
||||
)
|
||||
_ = [chunk async for chunk in stream]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_200_unknown_operation_body_is_a_bad_request_on_an_async_stream(_aws_test_credentials: None) -> None:
|
||||
with pytest.raises(litellm.BadRequestError) as exc_info:
|
||||
await _drain_async_unknown_operation_stream()
|
||||
|
||||
_assert_unknown_operation_is_a_bad_request(exc_info.value)
|
||||
|
||||
|
||||
def test_a_tool_call_whose_arguments_carry_a_type_key_is_not_read_as_a_bedrock_error(
|
||||
_aws_test_credentials: None,
|
||||
) -> None:
|
||||
arguments: Final = {"__type": "Document", "title": "notes"}
|
||||
tool_call_answer: Final = json.dumps(
|
||||
{
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude",
|
||||
"content": [{"type": "tool_use", "id": "toolu_1", "name": "save", "input": arguments}],
|
||||
"stop_reason": "tool_use",
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1},
|
||||
}
|
||||
).encode()
|
||||
upstream: Final = httpx.MockTransport(
|
||||
lambda request: httpx.Response(
|
||||
200, content=tool_call_answer, headers={"content-type": "application/json"}, request=request
|
||||
)
|
||||
)
|
||||
|
||||
response: Final = litellm.completion(
|
||||
model="bedrock/invoke/anthropic.claude-haiku-5-5",
|
||||
messages=[{"role": "user", "content": "save my notes"}],
|
||||
max_tokens=10,
|
||||
client=HTTPHandler(client=httpx.Client(transport=upstream)),
|
||||
)
|
||||
|
||||
assert isinstance(response, litellm.ModelResponse)
|
||||
tool_calls: Final = response.choices[0].message.tool_calls
|
||||
assert tool_calls is not None
|
||||
assert json.loads(tool_calls[0].function.arguments) == arguments
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue