mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
Merge 92827ab9bc into 49affa7c01
This commit is contained in:
commit
ddd1109d82
2 changed files with 23 additions and 2 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue