From 5e03c9c637c46457c46a058bd45febc298b10445 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Sat, 13 Jan 2024 15:39:54 -0800 Subject: [PATCH 1/5] (feat) bedrock support provisioned throughput --- litellm/llms/bedrock.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/litellm/llms/bedrock.py b/litellm/llms/bedrock.py index b2d04c0a2d..3b9d71f775 100644 --- a/litellm/llms/bedrock.py +++ b/litellm/llms/bedrock.py @@ -436,6 +436,9 @@ def completion( ) model = model + modelId = ( + optional_params.pop("model_id", None) or model + ) # default to model if not passed provider = model.split(".")[0] prompt = convert_messages_to_prompt( model, messages, provider, custom_prompt_dict @@ -510,7 +513,7 @@ def completion( request_str = f""" response = client.invoke_model( body={data}, - modelId={model}, + modelId={modelId}, accept=accept, contentType=contentType ) @@ -525,7 +528,7 @@ def completion( ) response = client.invoke_model( - body=data, modelId=model, accept=accept, contentType=contentType + body=data, modelId=modelId, accept=accept, contentType=contentType ) response = response.get("body").read() @@ -535,7 +538,7 @@ def completion( request_str = f""" response = client.invoke_model_with_response_stream( body={data}, - modelId={model}, + modelId={modelId}, accept=accept, contentType=contentType ) @@ -550,7 +553,7 @@ def completion( ) response = client.invoke_model_with_response_stream( - body=data, modelId=model, accept=accept, contentType=contentType + body=data, modelId=modelId, accept=accept, contentType=contentType ) response = response.get("body") return response @@ -559,7 +562,7 @@ def completion( request_str = f""" response = client.invoke_model( body={data}, - modelId={model}, + modelId={modelId}, accept=accept, contentType=contentType ) @@ -573,7 +576,7 @@ def completion( }, ) response = client.invoke_model( - body=data, modelId=model, accept=accept, contentType=contentType + body=data, modelId=modelId, accept=accept, contentType=contentType ) except client.exceptions.ValidationException as e: if "The provided model identifier is invalid" in str(e): From 0307ba7def22e4a3f29bce28b1157acef3ec2668 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Sat, 13 Jan 2024 15:40:30 -0800 Subject: [PATCH 2/5] (test) bedrock provisioned througput --- litellm/tests/test_bedrock_completion.py | 36 ++++++++++++++++++++++-- 1 file changed, 34 insertions(+), 2 deletions(-) diff --git a/litellm/tests/test_bedrock_completion.py b/litellm/tests/test_bedrock_completion.py index 3ef27b8c61..6e651804e1 100644 --- a/litellm/tests/test_bedrock_completion.py +++ b/litellm/tests/test_bedrock_completion.py @@ -63,7 +63,7 @@ def test_completion_bedrock_claude_completion_auth(): pytest.fail(f"Error occurred: {e}") -test_completion_bedrock_claude_completion_auth() +# test_completion_bedrock_claude_completion_auth() def test_completion_bedrock_claude_2_1_completion_auth(): @@ -100,7 +100,7 @@ def test_completion_bedrock_claude_2_1_completion_auth(): pytest.fail(f"Error occurred: {e}") -test_completion_bedrock_claude_2_1_completion_auth() +# test_completion_bedrock_claude_2_1_completion_auth() def test_completion_bedrock_claude_external_client_auth(): @@ -118,6 +118,8 @@ def test_completion_bedrock_claude_external_client_auth(): try: import boto3 + litellm.set_verbose = True + bedrock = boto3.client( service_name="bedrock-runtime", region_name=aws_region_name, @@ -146,3 +148,33 @@ def test_completion_bedrock_claude_external_client_auth(): 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 + + +# 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} + + +# 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() From 719b051b3dc3e03cedfa76d0fe4db707605c953e Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Sat, 13 Jan 2024 15:56:24 -0800 Subject: [PATCH 3/5] (test) provisioned throughput model --- litellm/tests/test_bedrock_completion.py | 74 +++++++++++++++--------- 1 file changed, 48 insertions(+), 26 deletions(-) diff --git a/litellm/tests/test_bedrock_completion.py b/litellm/tests/test_bedrock_completion.py index 6e651804e1..d96fdfa5a6 100644 --- a/litellm/tests/test_bedrock_completion.py +++ b/litellm/tests/test_bedrock_completion.py @@ -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() From 069d060ec99aedabf3873ee1be0626699fb05e48 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Sat, 13 Jan 2024 21:07:38 -0800 Subject: [PATCH 4/5] (feat) provisioned throughput - bedrock embedding models --- litellm/llms/bedrock.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/litellm/llms/bedrock.py b/litellm/llms/bedrock.py index 3b9d71f775..85b820c9e8 100644 --- a/litellm/llms/bedrock.py +++ b/litellm/llms/bedrock.py @@ -667,6 +667,9 @@ def _embedding_func_single( inference_params.pop( "user", None ) # make sure user is not passed in for bedrock call + modelId = ( + optional_params.pop("model_id", None) or model + ) # default to model if not passed if provider == "amazon": input = input.replace(os.linesep, " ") data = {"inputText": input, **inference_params} @@ -681,7 +684,7 @@ def _embedding_func_single( request_str = f""" response = client.invoke_model( body={body}, - modelId={model}, + modelId={modelId}, accept="*/*", contentType="application/json", )""" # type: ignore @@ -689,14 +692,14 @@ def _embedding_func_single( input=input, api_key="", # boto3 is used for init. additional_args={ - "complete_input_dict": {"model": model, "texts": input}, + "complete_input_dict": {"model": modelId, "texts": input}, "request_str": request_str, }, ) try: response = client.invoke_model( body=body, - modelId=model, + modelId=modelId, accept="*/*", contentType="application/json", ) From 7af3fe68e4591650858aa9a0a6a98f8f7f7214ff Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Sat, 13 Jan 2024 21:12:59 -0800 Subject: [PATCH 5/5] (docs) provisioned throughput bedrock models --- docs/my-website/docs/providers/bedrock.md | 26 +++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/docs/my-website/docs/providers/bedrock.md b/docs/my-website/docs/providers/bedrock.md index ae180ff359..f83e3ed02c 100644 --- a/docs/my-website/docs/providers/bedrock.md +++ b/docs/my-website/docs/providers/bedrock.md @@ -208,6 +208,32 @@ response = completion( ) ``` +## Provisioned throughput models +To use provisioned throughput Bedrock models pass +- `model=bedrock/`, example `model=bedrock/anthropic.claude-v2`. Set `model` to any of the [Supported AWS models](#supported-aws-bedrock-models) +- `model_id=provisioned-model-arn` + +Completion +```python +import litellm +response = litellm.completion( + model="bedrock/anthropic.claude-instant-v1", + model_id="provisioned-model-arn", + messages=[{"content": "Hello, how are you?", "role": "user"}] +) +``` + +Embedding +```python +import litellm +response = litellm.embedding( + model="bedrock/amazon.titan-embed-text-v1", + model_id="provisioned-model-arn", + input=["hi"], +) +``` + + ## Supported AWS Bedrock Models Here's an example of using a bedrock model with LiteLLM