Feat add Amazon Bedrock support

config/puppeteer-config.json

@@ -1,6 +1,4 @@
 {
-    "executablePath": "/usr/bin/chromium",
-    "args": [
-        "--no-sandbox"
-    ]
-}
+  "executablePath": "/usr/bin/chromium",
+  "args": ["--no-sandbox"]
+}

metagpt/configs/llm_config.py

@@ -32,6 +32,7 @@ class LLMType(Enum):
     MISTRAL = "mistral"
     YI = "yi"  # lingyiwanwu
     OPENROUTER = "openrouter"
+    BEDROCK = "bedrock"
 
     def __missing__(self, key):
         return self.OPENAI
@@ -74,10 +75,14 @@ class LLMConfig(YamlModel):
     best_of: Optional[int] = None
     n: Optional[int] = None
     stream: bool = False
-    logprobs: Optional[bool] = None  # https://cookbook.openai.com/examples/using_logprobs
+    # https://cookbook.openai.com/examples/using_logprobs
+    logprobs: Optional[bool] = None
     top_logprobs: Optional[int] = None
     timeout: int = 600
 
+    # For Amazon Bedrock
+    region_name: str = None
+
     # For Network
     proxy: Optional[str] = None
 

metagpt/provider/__init__.py

@@ -17,6 +17,7 @@
 from metagpt.provider.qianfan_api import QianFanLLM
 from metagpt.provider.dashscope_api import DashScopeLLM
 from metagpt.provider.anthropic_api import AnthropicLLM
+from metagpt.provider.bedrock_api import BedrockLLM
 
 __all__ = [
     "GeminiLLM",
@@ -30,4 +31,5 @@
     "QianFanLLM",
     "DashScopeLLM",
     "AnthropicLLM",
+    "BedrockLLM",
 ]

metagpt/provider/bedrock/base_provider.py

@@ -0,0 +1,28 @@
+import json
+from abc import ABC, abstractmethod
+
+
+class BaseBedrockProvider(ABC):
+    # to handle different generation kwargs
+    max_tokens_field_name = "max_tokens"
+
+    @abstractmethod
+    def _get_completion_from_dict(self, rsp_dict: dict) -> str:
+        ...
+
+    def get_request_body(self, messages: list[dict], const_kwargs, *args, **kwargs) -> str:
+        body = json.dumps({"prompt": self.messages_to_prompt(messages), **const_kwargs})
+        return body
+
+    def get_choice_text(self, response_body: dict) -> str:
+        completions = self._get_completion_from_dict(response_body)
+        return completions
+
+    def get_choice_text_from_stream(self, event) -> str:
+        rsp_dict = json.loads(event["chunk"]["bytes"])
+        completions = self._get_completion_from_dict(rsp_dict)
+        return completions
+
+    def messages_to_prompt(self, messages: list[dict]) -> str:
+        """[{"role": "user", "content": msg}] to user: <msg> etc."""
+        return "\n".join([f"{i['role']}: {i['content']}" for i in messages])

metagpt/provider/bedrock/bedrock_provider.py

@@ -0,0 +1,121 @@
+import json
+from typing import Literal
+
+from metagpt.provider.bedrock.base_provider import BaseBedrockProvider
+from metagpt.provider.bedrock.utils import (
+    messages_to_prompt_llama2,
+    messages_to_prompt_llama3,
+)
+
+
+class MistralProvider(BaseBedrockProvider):
+    # See https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-mistral.html
+
+    def messages_to_prompt(self, messages: list[dict]):
+        return messages_to_prompt_llama2(messages)
+
+    def _get_completion_from_dict(self, rsp_dict: dict) -> str:
+        return rsp_dict["outputs"][0]["text"]
+
+
+class AnthropicProvider(BaseBedrockProvider):
+    # See https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages.html
+
+    def get_request_body(self, messages: list[dict], generate_kwargs, *args, **kwargs):
+        body = json.dumps({"messages": messages, "anthropic_version": "bedrock-2023-05-31", **generate_kwargs})
+        return body
+
+    def _get_completion_from_dict(self, rsp_dict: dict) -> str:
+        return rsp_dict["content"][0]["text"]
+
+    def get_choice_text_from_stream(self, event) -> str:
+        # https://docs.anthropic.com/claude/reference/messages-streaming
+        rsp_dict = json.loads(event["chunk"]["bytes"])
+        if rsp_dict["type"] == "content_block_delta":
+            completions = rsp_dict["delta"]["text"]
+            return completions
+        else:
+            return ""
+
+
+class CohereProvider(BaseBedrockProvider):
+    # See https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-cohere-command.html
+
+    def _get_completion_from_dict(self, rsp_dict: dict) -> str:
+        return rsp_dict["generations"][0]["text"]
+
+    def get_request_body(self, messages: list[dict], generate_kwargs, *args, **kwargs):
+        body = json.dumps(
+            {"prompt": self.messages_to_prompt(messages), "stream": kwargs.get("stream", False), **generate_kwargs}
+        )
+        return body
+
+    def get_choice_text_from_stream(self, event) -> str:
+        rsp_dict = json.loads(event["chunk"]["bytes"])
+        completions = rsp_dict.get("text", "")
+        return completions
+
+
+class MetaProvider(BaseBedrockProvider):
+    # See https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-meta.html
+
+    max_tokens_field_name = "max_gen_len"
+
+    def __init__(self, llama_version: Literal["llama2", "llama3"]) -> None:
+        self.llama_version = llama_version
+
+    def messages_to_prompt(self, messages: list[dict]):
+        if self.llama_version == "llama2":
+            return messages_to_prompt_llama2(messages)
+        else:
+            return messages_to_prompt_llama3(messages)
+
+    def _get_completion_from_dict(self, rsp_dict: dict) -> str:
+        return rsp_dict["generation"]
+
+
+class Ai21Provider(BaseBedrockProvider):
+    # See https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-jurassic2.html
+
+    max_tokens_field_name = "maxTokens"
+
+    def _get_completion_from_dict(self, rsp_dict: dict) -> str:
+        return rsp_dict["completions"][0]["data"]["text"]
+
+
+class AmazonProvider(BaseBedrockProvider):
+    # See https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-titan-text.html
+
+    max_tokens_field_name = "maxTokenCount"
+
+    def get_request_body(self, messages: list[dict], generate_kwargs, *args, **kwargs):
+        body = json.dumps({"inputText": self.messages_to_prompt(messages), "textGenerationConfig": generate_kwargs})
+        return body
+
+    def _get_completion_from_dict(self, rsp_dict: dict) -> str:
+        return rsp_dict["results"][0]["outputText"]
+
+    def get_choice_text_from_stream(self, event) -> str:
+        rsp_dict = json.loads(event["chunk"]["bytes"])
+        completions = rsp_dict["outputText"]
+        return completions
+
+
+PROVIDERS = {
+    "mistral": MistralProvider,
+    "meta": MetaProvider,
+    "ai21": Ai21Provider,
+    "cohere": CohereProvider,
+    "anthropic": AnthropicProvider,
+    "amazon": AmazonProvider,
+}
+
+
+def get_provider(model_id: str):
+    provider, model_name = model_id.split(".")[0:2]  # meta、mistral……
+    if provider not in PROVIDERS:
+        raise KeyError(f"{provider} is not supported!")
+    if provider == "meta":
+        # distinguish llama2 and llama3
+        return PROVIDERS[provider](model_name[:6])
+    return PROVIDERS[provider]()

metagpt/provider/bedrock/utils.py

@@ -0,0 +1,112 @@
+from metagpt.logs import logger
+
+# max_tokens for each model
+NOT_SUUPORT_STREAM_MODELS = {
+    "ai21.j2-grande-instruct": 8000,
+    "ai21.j2-jumbo-instruct": 8000,
+    "ai21.j2-mid": 8000,
+    "ai21.j2-mid-v1": 8000,
+    "ai21.j2-ultra": 8000,
+    "ai21.j2-ultra-v1": 8000,
+}
+
+SUPPORT_STREAM_MODELS = {
+    "amazon.titan-tg1-large": 8000,
+    "amazon.titan-text-express-v1": 8000,
+    "amazon.titan-text-express-v1:0:8k": 8000,
+    "amazon.titan-text-lite-v1:0:4k": 4000,
+    "amazon.titan-text-lite-v1": 4000,
+    "anthropic.claude-instant-v1": 100000,
+    "anthropic.claude-instant-v1:2:100k": 100000,
+    "anthropic.claude-v1": 100000,
+    "anthropic.claude-v2": 100000,
+    "anthropic.claude-v2:1": 200000,
+    "anthropic.claude-v2:0:18k": 18000,
+    "anthropic.claude-v2:1:200k": 200000,
+    "anthropic.claude-3-sonnet-20240229-v1:0": 200000,
+    "anthropic.claude-3-sonnet-20240229-v1:0:28k": 28000,
+    "anthropic.claude-3-sonnet-20240229-v1:0:200k": 200000,
+    "anthropic.claude-3-haiku-20240307-v1:0": 200000,
+    "anthropic.claude-3-haiku-20240307-v1:0:48k": 48000,
+    "anthropic.claude-3-haiku-20240307-v1:0:200k": 200000,
+    # currently (2024-4-29) only available at US West (Oregon) AWS Region.
+    "anthropic.claude-3-opus-20240229-v1:0": 200000,
+    "cohere.command-text-v14": 4000,
+    "cohere.command-text-v14:7:4k": 4000,
+    "cohere.command-light-text-v14": 4000,
+    "cohere.command-light-text-v14:7:4k": 4000,
+    "meta.llama2-13b-chat-v1:0:4k": 4000,
+    "meta.llama2-13b-chat-v1": 2000,
+    "meta.llama2-70b-v1": 4000,
+    "meta.llama2-70b-v1:0:4k": 4000,
+    "meta.llama2-70b-chat-v1": 4000,
+    "meta.llama2-70b-chat-v1:0:4k": 4000,
+    "meta.llama3-8b-instruct-v1:0": 2000,
+    "meta.llama3-70b-instruct-v1:0": 2000,
+    "mistral.mistral-7b-instruct-v0:2": 32000,
+    "mistral.mixtral-8x7b-instruct-v0:1": 32000,
+    "mistral.mistral-large-2402-v1:0": 32000,
+}
+
+# TODO:use a more general function for constructing chat templates.
+
+
+def messages_to_prompt_llama2(messages: list[dict]) -> str:
+    BOS = ("<s>",)
+    B_INST, E_INST = "[INST]", "[/INST]"
+    B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
+
+    prompt = f"{BOS}"
+    for message in messages:
+        role = message.get("role", "")
+        content = message.get("content", "")
+        if role == "system":
+            prompt += f"{B_SYS} {content} {E_SYS}"
+        elif role == "user":
+            prompt += f"{B_INST} {content} {E_INST}"
+        elif role == "assistant":
+            prompt += f"{content}"
+        else:
+            logger.warning(f"Unknown role name {role} when formatting messages")
+            prompt += f"{content}"
+
+    return prompt
+
+
+def messages_to_prompt_llama3(messages: list[dict]) -> str:
+    BOS = "<|begin_of_text|>"
+    GENERAL_TEMPLATE = "<|start_header_id|>{role}<|end_header_id|>\n\n{content}<|eot_id|>"
+
+    prompt = f"{BOS}"
+    for message in messages:
+        role = message.get("role", "")
+        content = message.get("content", "")
+        prompt += GENERAL_TEMPLATE.format(role=role, content=content)
+
+    if role != "assistant":
+        prompt += "<|start_header_id|>assistant<|end_header_id|>"
+
+    return prompt
+
+
+def messages_to_prompt_claude2(messages: list[dict]) -> str:
+    GENERAL_TEMPLATE = "\n\n{role}: {content}"
+    prompt = ""
+    for message in messages:
+        role = message.get("role", "")
+        content = message.get("content", "")
+        prompt += GENERAL_TEMPLATE.format(role=role, content=content)
+
+    if role != "assistant":
+        prompt += "\n\nAssistant:"
+
+    return prompt
+
+
+def get_max_tokens(model_id: str) -> int:
+    try:
+        max_tokens = (NOT_SUUPORT_STREAM_MODELS | SUPPORT_STREAM_MODELS)[model_id]
+    except KeyError:
+        logger.warning(f"Couldn't find model:{model_id} , max tokens has been set to 2048")
+        max_tokens = 2048
+    return max_tokens