From 6201ab2c219ed40bc0d1b245af92a69d33624f8d Mon Sep 17 00:00:00 2001 From: "spdustin@gmail.com" Date: Fri, 5 Jan 2024 23:32:32 +0000 Subject: [PATCH] Update factory (and tests) for Claude 2.1 via Bedrock --- litellm/llms/prompt_templates/factory.py | 2 +- litellm/tests/test_bedrock_completion.py | 37 ++++++++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/litellm/llms/prompt_templates/factory.py b/litellm/llms/prompt_templates/factory.py index e3e0f325d08..61aa71c1bf6 100644 --- a/litellm/llms/prompt_templates/factory.py +++ b/litellm/llms/prompt_templates/factory.py @@ -510,7 +510,7 @@ def prompt_factory( if custom_llm_provider == "ollama": return ollama_pt(model=model, messages=messages) elif custom_llm_provider == "anthropic": - if "claude-2.1" in model: + if any(_ in model for _ in ["claude-2.1","claude-v2:1"]): return claude_2_1_pt(messages=messages) else: return anthropic_pt(messages=messages) diff --git a/litellm/tests/test_bedrock_completion.py b/litellm/tests/test_bedrock_completion.py index 004ae86f12b..3ef27b8c611 100644 --- a/litellm/tests/test_bedrock_completion.py +++ b/litellm/tests/test_bedrock_completion.py @@ -66,6 +66,43 @@ def test_completion_bedrock_claude_completion_auth(): test_completion_bedrock_claude_completion_auth() +def test_completion_bedrock_claude_2_1_completion_auth(): + print("calling bedrock claude 2.1 completion params auth") + import os + + aws_access_key_id = os.environ["AWS_ACCESS_KEY_ID"] + aws_secret_access_key = os.environ["AWS_SECRET_ACCESS_KEY"] + aws_region_name = os.environ["AWS_REGION_NAME"] + + os.environ.pop("AWS_ACCESS_KEY_ID", None) + os.environ.pop("AWS_SECRET_ACCESS_KEY", None) + os.environ.pop("AWS_REGION_NAME", None) + + try: + response = completion( + model="bedrock/anthropic.claude-v2:1", + messages=messages, + max_tokens=10, + temperature=0.1, + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_region_name=aws_region_name, + ) + # Add any assertions here to check the response + print(response) + + os.environ["AWS_ACCESS_KEY_ID"] = aws_access_key_id + os.environ["AWS_SECRET_ACCESS_KEY"] = aws_secret_access_key + os.environ["AWS_REGION_NAME"] = aws_region_name + except RateLimitError: + pass + except Exception as e: + pytest.fail(f"Error occurred: {e}") + + +test_completion_bedrock_claude_2_1_completion_auth() + + def test_completion_bedrock_claude_external_client_auth(): print("\ncalling bedrock claude external client auth") import os