scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

(test) provisioned throughput model

Browse source
This commit is contained in:
ishaan-jaff 2024-01-13 15:56:24 -08:00
parent 0307ba7def
commit 719b051b3d
1 changed files with 48 additions and 26 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

74
litellm/tests/test_bedrock_completion.py
View file

@ -147,34 +147,56 @@ def test_completion_bedrock_claude_external_client_auth():
pytest.fail(f"Error occurred: {e}")
test_completion_bedrock_claude_external_client_auth()
# def test_provisioned_throughput():
# try:
# litellm.set_verbose=True
# import botocore
# import botocore.session
# from botocore.stub import Stubber
# test_completion_bedrock_claude_external_client_auth()
# bedrock_client = botocore.session.get_session().create_client('bedrock-runtime', region_name='us-east-1')
# request = {
# "body": {"prompt": "\n\nHuman: Write a short poem about the sky\n\nAssistant: ", "max_tokens_to_sample": 10, "temperature": 0.1},
# "modelId": "anthropic.claude-instant-v1",
# "accept":"accept",
# "contentType": "contentType"
# }
# response ={"completion": " Here is a short poem about the sky:", "stop_reason": "max_tokens", "stop": None}
def test_provisioned_throughput():
try:
litellm.set_verbose = True
import botocore, json, io
import botocore.session
from botocore.stub import Stubber
bedrock_client = botocore.session.get_session().create_client(
"bedrock-runtime", region_name="us-east-1"
)
expected_params = {
"accept": "application/json",
"body": '{"prompt": "\\n\\nHuman: Hello, how are you?\\n\\nAssistant: ", '
'"max_tokens_to_sample": 256}',
"contentType": "application/json",
"modelId": "provisioned-model-arn",
}
response_from_bedrock = {
"body": io.StringIO(
json.dumps(
{
"completion": " Here is a short poem about the sky:",
"stop_reason": "max_tokens",
"stop": None,
}
)
),
"contentType": "contentType",
"ResponseMetadata": {"HTTPStatusCode": 200},
}
with Stubber(bedrock_client) as stubber:
stubber.add_response(
"invoke_model",
service_response=response_from_bedrock,
expected_params=expected_params,
)
response = litellm.completion(
model="bedrock/anthropic.claude-instant-v1",
model_id="provisioned-model-arn",
messages=[{"content": "Hello, how are you?", "role": "user"}],
aws_bedrock_client=bedrock_client,
)
print("response stubbed", response)
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# with Stubber(bedrock_client) as stubber:
# stubber.add_response('invoke_model', service_response=response, expected_params=request)
# response = litellm.completion(
# model="bedrock/anthropic.claude-instant-v1",
# model_id = 'provisioned-model-arn',
# messages=[{ "content": "Hello, how are you?","role": "user"}],
# aws_bedrock_client=bedrock_client,
# )
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# test_provisioned_throughput()