Merge pull request #7 from JGalego/feat/anthropic-bedrock

feat: AnthropicBedrock support ⛰️
SakanaAI · Aug 16, 2024 · 42be334 · 42be334
2 parents 59e1ff9 + 653f823
commit 42be334
Show file tree

Hide file tree

Showing 7 changed files with 66 additions and 5 deletions.
diff --git a/README.md b/README.md
@@ -53,7 +53,7 @@ conda create -n ai_scientist python=3.11
 conda activate ai_scientist
 
 # LLM APIs
-pip install anthropic aider-chat backoff openai
+pip install anthropic[bedrock] aider-chat backoff openai
 # Viz
 pip install matplotlib pypdf pymupdf4llm
 # Install pdflatex
@@ -67,6 +67,10 @@ We use the following environment variables for the different API providers for d
 
 `OPENAI_API_KEY`, `ANTHROPIC_API_KEY`, `DEEPSEEK_API_KEY`, `OPENROUTER_API_KEY`
 
+For Claude models provided by [Amazon Bedrock](https://aws.amazon.com/bedrock/), please specify a set of valid [AWS Credentials](https://docs.aws.amazon.com/cli/v1/userguide/cli-configure-envvars.html) and the target [AWS Region](https://docs.aws.amazon.com/bedrock/latest/userguide/bedrock-regions.html):
+
+(*required*) `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, (*optional*) `AWS_SESSION_TOKEN`, `AWS_DEFAULT_REGION`
+
 Our code can also optionally use a Semantic Scholar API Key (`S2_API_KEY`) for higher throughput [if you have one](https://www.semanticscholar.org/product/api), though in principle it should work without it.
 
 Be sure to provide the key for the model used for your runs, e.g.

diff --git a/ai_scientist/generate_ideas.py b/ai_scientist/generate_ideas.py
@@ -490,6 +490,14 @@ def check_idea_novelty(
         print(f"Using Anthropic API with model {args.model}.")
         client_model = "claude-3-5-sonnet-20240620"
         client = anthropic.Anthropic()
+    elif args.model.startswith("bedrock") and "claude" in args.model:
+        import anthropic
+
+        # Expects: bedrock/<MODEL_ID>
+        client_model = args.model.split("/")[-1]
+
+        print(f"Using Amazon Bedrock with model {client_model}.")
+        client = anthropic.AnthropicBedrock()
     elif args.model == "gpt-4o-2024-05-13" or args.model == "hybrid":
         import openai
 

diff --git a/ai_scientist/llm.py b/ai_scientist/llm.py
@@ -74,7 +74,7 @@ def get_batch_responses_from_llm(
         new_msg_history = [
             new_msg_history + [{"role": "assistant", "content": c}] for c in content
         ]
-    elif model == "claude-3-5-sonnet-20240620":
+    elif "claude" in model:
         content, new_msg_history = [], []
         for _ in range(n_responses):
             c, hist = get_response_from_llm(
@@ -118,7 +118,7 @@ def get_response_from_llm(
     if msg_history is None:
         msg_history = []
 
-    if model == "claude-3-5-sonnet-20240620":
+    if "claude" in model:
         new_msg_history = msg_history + [
             {
                 "role": "user",
@@ -131,7 +131,7 @@ def get_response_from_llm(
             }
         ]
         response = client.messages.create(
-            model="claude-3-5-sonnet-20240620",
+            model=model,
             max_tokens=3000,
             temperature=temperature,
             system=system_message,

diff --git a/ai_scientist/perform_writeup.py b/ai_scientist/perform_writeup.py
@@ -528,6 +528,11 @@ def perform_writeup(
             "gpt-4o-2024-05-13",
             "deepseek-coder-v2-0724",
             "llama3.1-405b",
+            # Anthropic Claude models via Amazon Bedrock
+            "bedrock/anthropic.claude-3-sonnet-20240229-v1:0",
+            "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
+            "bedrock/anthropic.claude-3-haiku-20240307-v1:0",
+            "bedrock/anthropic.claude-3-opus-20240229-v1:0"
         ],
         help="Model to use for AI Scientist.",
     )
@@ -538,6 +543,14 @@ def perform_writeup(
         print(f"Using Anthropic API with model {args.model}.")
         client_model = "claude-3-5-sonnet-20240620"
         client = anthropic.Anthropic()
+    elif args.model.startswith("bedrock") and "claude" in args.model:
+        import anthropic
+
+        # Expects: bedrock/<MODEL_ID>
+        client_model = args.model.split("/")[-1]
+
+        print(f"Using Amazon Bedrock with model {client_model}.")
+        client = anthropic.AnthropicBedrock()
     elif args.model == "gpt-4o-2024-05-13" or args.model == "hybrid":
         import openai
 

diff --git a/experimental/launch_oe_scientist.py b/experimental/launch_oe_scientist.py
@@ -42,6 +42,11 @@ def parse_arguments():
             "gpt-4o-2024-05-13",
             "deepseek-coder-v2-0724",
             "llama3.1-405b",
+            # Anthropic Claude models via Amazon Bedrock
+            "bedrock/anthropic.claude-3-sonnet-20240229-v1:0",
+            "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
+            "bedrock/anthropic.claude-3-haiku-20240307-v1:0",
+            "bedrock/anthropic.claude-3-opus-20240229-v1:0"
         ],
         help="Model to use for AI Scientist.",
     )
@@ -325,6 +330,14 @@ def do_idea(
         print(f"Using Anthropic API with model {args.model}.")
         client_model = "claude-3-5-sonnet-20240620"
         client = anthropic.Anthropic()
+    elif args.model.startswith("bedrock") and "claude" in args.model:
+        import anthropic
+
+        # Expects: bedrock/<MODEL_ID>
+        client_model = args.model.split("/")[-1]
+
+        print(f"Using Amazon Bedrock with model {client_model}.")
+        client = anthropic.AnthropicBedrock()
     elif args.model == "gpt-4o-2024-05-13" or args.model == "hybrid":
         import openai
 

diff --git a/launch_scientist.py b/launch_scientist.py
@@ -47,7 +47,17 @@ def parse_arguments():
         "--model",
         type=str,
         default="claude-3-5-sonnet-20240620",
-        choices=["claude-3-5-sonnet-20240620", "gpt-4o-2024-05-13", "deepseek-coder-v2-0724", "llama3.1-405b"],
+        choices=[
+            "claude-3-5-sonnet-20240620",
+            "gpt-4o-2024-05-13",
+            "deepseek-coder-v2-0724",
+            "llama3.1-405b",
+            # Anthropic Claude models via Amazon Bedrock
+            "bedrock/anthropic.claude-3-sonnet-20240229-v1:0",
+            "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
+            "bedrock/anthropic.claude-3-haiku-20240307-v1:0",
+            "bedrock/anthropic.claude-3-opus-20240229-v1:0"
+        ],
         help="Model to use for AI Scientist.",
     )
     parser.add_argument(
@@ -266,6 +276,14 @@ def do_idea(
         print(f"Using Anthropic API with model {args.model}.")
         client_model = "claude-3-5-sonnet-20240620"
         client = anthropic.Anthropic()
+    elif args.model.startswith("bedrock") and "claude" in args.model:
+        import anthropic
+
+        # Expects: bedrock/<MODEL_ID>
+        client_model = args.model.split("/")[-1]
+
+        print(f"Using Amazon Bedrock with model {client_model}.")
+        client = anthropic.AnthropicBedrock()
     elif args.model == "gpt-4o-2024-05-13" or args.model == "hybrid":
         import openai
 

diff --git a/review_iclr_bench/iclr_analysis.py b/review_iclr_bench/iclr_analysis.py
@@ -241,6 +241,11 @@ def review_single_paper(
         import anthropic
 
         client = anthropic.Anthropic()
+    elif model.startswith("bedrock") and "claude" in model:
+        import anthropic
+
+        model = model.split("/")[-1]
+        client = anthropic.AnthropicBedrock()
     elif model in [
         "gpt-4o-2024-05-13",
         "gpt-4o-mini-2024-07-18",