scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 108

add assertion for test

Browse source
This commit is contained in:
Tim Xia 2024-03-01 23:27:02 -05:00
parent b4dc7f0f17
commit 58ed6e77de
1 changed files with 4 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
litellm/tests/test_bedrock_completion.py
View file

@ -11,7 +11,7 @@ sys.path.insert(
) # Adds the parent directory to the system path
import pytest
import litellm
from litellm import embedding, completion, completion_cost, Timeout
from litellm import embedding, completion, completion_cost, Timeout, ModelResponse
from litellm import RateLimitError
# litellm.num_retries = 3
@ -270,14 +270,15 @@ def test_completion_bedrock_mistral_completion_auth():
# os.environ.pop("AWS_SECRET_ACCESS_KEY", None)
# os.environ.pop("AWS_REGION_NAME", None)
try:
response = completion(
response:ModelResponse = completion(
model="bedrock/mistral.mistral-7b-instruct-v0:2",
messages=messages,
max_tokens=10,
temperature=0.1,
)
# Add any assertions here to check the response
print(response)
assert len(response.choices) > 0
assert len(response.choices[0].message.content) > 0
# os.environ["AWS_ACCESS_KEY_ID"] = aws_access_key_id
# os.environ["AWS_SECRET_ACCESS_KEY"] = aws_secret_access_key