From 719b051b3dc3e03cedfa76d0fe4db707605c953e Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Sat, 13 Jan 2024 15:56:24 -0800
Subject: [PATCH] (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 6e651804e1e..d96fdfa5a61 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()