diff --git a/README.md b/README.md index 0f2fbaee951..6478be9f7c3 100644 --- a/README.md +++ b/README.md @@ -258,6 +258,9 @@ The current list of supported models: | | Claude 3 Haiku | --model claude-3-haiku-20240307 | | Meta | Llama 3 70B | --model llama3:70b | | | Llama 3 8B | --model llama3 | +| AWS | Claude 3 Opus | --model bedrock/anthropic.claude-3-opus-20240229-v1:0 | +| | Claude 3 Sonnet | --model bedrock/anthropic.claude-3-sonnet-20240229-v1:0 | +| | Claude 3 Haiku | --model bedrock/anthropic.claude-3-haiku-20240307-v1:0 | > [!NOTE] > Some notes on running ACR with local models such as llama3: diff --git a/app/model/bedrock.py b/app/model/bedrock.py new file mode 100644 index 00000000000..eea904b616b --- /dev/null +++ b/app/model/bedrock.py @@ -0,0 +1,147 @@ +""" +For models other than those from OpenAI, use LiteLLM if possible. +""" + +import os +import sys + +from typing import Literal + +import litellm +from litellm.utils import Choices, Message, ModelResponse +from openai import BadRequestError +from tenacity import retry, stop_after_attempt, wait_random_exponential + +from app.log import log_and_print +from app.model import common +from app.model.common import Model + + +class BedrockModel(Model): + """ + Base class for creating Singleton instances of Amazon Bedrock models. + """ + + _instances = {} + + def __new__(cls): + if cls not in cls._instances: + cls._instances[cls] = super().__new__(cls) + cls._instances[cls]._initialized = False + return cls._instances[cls] + + def __init__( + self, + name: str, + cost_per_input: float, + cost_per_output: float, + parallel_tool_call: bool = False, + ): + if self._initialized: + return + super().__init__(name, cost_per_input, cost_per_output, parallel_tool_call) + self._model_provider = self.name.split(".")[0] + self._initialized = True + + def setup(self) -> None: + """ + Check API key. + """ + self.check_api_key() + + def check_api_key(self) -> str: + # See https://litellm.vercel.app/docs/providers/bedrock + required_env_vars = [ + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", + "AWS_REGION_NAME" + ] + if len(set(os.environ).intersection(required_env_vars)) != len(required_env_vars): + print(f"Missing env vars. Please refer to https://litellm.vercel.app/docs/providers/bedrock") + sys.exit(1) + return os.getenv(required_env_vars[-1]) + + def extract_resp_content(self, chat_message: Message) -> str: + """ + Given a chat completion message, extract the content from it. + """ + content = chat_message.content + if content is None: + return "" + else: + return content + + @retry(wait=wait_random_exponential(min=30, max=600), stop=stop_after_attempt(3)) + def call( + self, + messages: list[dict], + top_p=1, + tools=None, + response_format: Literal["text", "json_object"] = "text", + **kwargs, + ): + try: + if self._model_provider == "bedrock/anthropic": + # antropic models - prefilling response with { increase the success rate + # of producing json output + prefill_content = "{" + if response_format == "json_object": # prefill + messages.append({"role": "assistant", "content": prefill_content}) + + response = litellm.completion( + model=self.name, + messages=messages, + temperature=common.MODEL_TEMP, + max_tokens=1024, + top_p=top_p, + stream=False, + ) + assert isinstance(response, ModelResponse) + resp_usage = response.usage + assert resp_usage is not None + input_tokens = int(resp_usage.prompt_tokens) + output_tokens = int(resp_usage.completion_tokens) + cost = self.calc_cost(input_tokens, output_tokens) + + common.thread_cost.process_cost += cost + common.thread_cost.process_input_tokens += input_tokens + common.thread_cost.process_output_tokens += output_tokens + + first_resp_choice = response.choices[0] + assert isinstance(first_resp_choice, Choices) + resp_msg: Message = first_resp_choice.message + content = self.extract_resp_content(resp_msg) + if response_format == "json_object": + # prepend the prefilled character + if not content.startswith(prefill_content): + content = prefill_content + content + return content, cost, input_tokens, output_tokens + + except BadRequestError as e: + if e.code == "context_length_exceeded": + log_and_print("Context length exceeded") + raise e + + +class AnthropicClaude3Opus(BedrockModel): + def __init__(self): + super().__init__( + "bedrock/anthropic.claude-3-opus-20240229-v1:0", 0.000015, 0.000075, parallel_tool_call=True + ) + self.note = "Most powerful model from Antropic" + + +class AnthropicClaude3Sonnet(BedrockModel): + def __init__(self): + super().__init__( + "bedrock/anthropic.claude-3-sonnet-20240229-v1:0", 0.000003, 0.000015, parallel_tool_call=True + ) + self.note = "Most balanced (intelligence and speed) model from Antropic" + + +class AnthropicClaude3Haiku(BedrockModel): + def __init__(self): + super().__init__( + "bedrock/anthropic.claude-3-haiku-20240307-v1:0", 0.00000025, 0.00000125, parallel_tool_call=True + ) + self.note = "Fastest model from Antropic" diff --git a/app/model/register.py b/app/model/register.py index 772273c128a..1cf22676503 100644 --- a/app/model/register.py +++ b/app/model/register.py @@ -1,4 +1,4 @@ -from app.model import claude, common, gpt, ollama +from app.model import bedrock, claude, common, gpt, ollama def register_all_models() -> None: @@ -17,6 +17,10 @@ def register_all_models() -> None: common.register_model(claude.Claude3Opus()) common.register_model(claude.Claude3Sonnet()) common.register_model(claude.Claude3Haiku()) + + common.register_model(bedrock.AnthropicClaude3Opus()) + common.register_model(bedrock.AnthropicClaude3Sonnet()) + common.register_model(bedrock.AnthropicClaude3Haiku()) common.register_model(ollama.Llama3_8B()) common.register_model(ollama.Llama3_70B()) diff --git a/environment.yml b/environment.yml index 35bb338ed35..09fd26e7516 100644 --- a/environment.yml +++ b/environment.yml @@ -33,6 +33,7 @@ dependencies: - annotated-types==0.6.0 - anyio==4.2.0 - beautifulsoup4==4.12.3 + - boto3==1.34.99 - certifi==2023.11.17 - distro==1.9.0 - docstring-parser==0.15