This commit is contained in:
dfedoryshchev 2026-08-27 18:46:34 -05:00 • committed by GitHub
commit ddd1109d82
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 23 additions and 2 deletions

View file

@ -99,10 +99,10 @@ class AmazonMistralConfig(AmazonInvokeConfig, BaseConfig):
A string with the response of the LLM
"""
if "choices" in completion_response:
if "choices" in completion_response and len(completion_response["choices"]) > 0:
outputText = completion_response["choices"][0]["message"]["content"]
model_response.choices[0].finish_reason = completion_response["choices"][0]["finish_reason"]
elif "outputs" in completion_response:
elif "outputs" in completion_response and len(completion_response["outputs"]) > 0:
outputText = completion_response["outputs"][0]["text"]
model_response.choices[0].finish_reason = completion_response["outputs"][0]["stop_reason"]
else:

View file

@ -1,6 +1,9 @@
import pytest
from litellm.llms.bedrock.chat.invoke_transformations.amazon_mistral_transformation import (
AmazonMistralConfig,
)
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.types.utils import ModelResponse
@ -30,3 +33,21 @@ def test_mistral_get_outputText():
assert outputText == "Hi!"
assert model_response.choices[0].finish_reason == "finish"
def test_mistral_get_outputText_empty_choices():
model_response = ModelResponse()
with pytest.raises(BedrockError):
AmazonMistralConfig.get_outputText(
completion_response={"choices": []}, model_response=model_response
)
def test_mistral_get_outputText_empty_outputs():
model_response = ModelResponse()
with pytest.raises(BedrockError):
AmazonMistralConfig.get_outputText(
completion_response={"outputs": []}, model_response=model_response
)