Merge pull request Codium-ai#483 from tmokmss/add-bedrock-support

Add Amazon Bedrock support

Usage.md

@@ -328,6 +328,23 @@ Your [application default credentials](https://cloud.google.com/docs/authenticat
 
 If you do want to set explicit credentials then you can use the `GOOGLE_APPLICATION_CREDENTIALS` environment variable set to a path to a json credentials file.
 
+#### Amazon Bedrock
+
+To use Amazon Bedrock and its foundational models, add the below configuration:
+
+``` 
+[config] # in configuration.toml
+model = "anthropic.claude-v2"
+fallback_models="anthropic.claude-instant-v1"
+
+[aws] # in .secrets.toml
+bedrock_region = "us-east-1"
+```
+
+Note that you have to add access to foundational models before using them. Please refer to [this document](https://docs.aws.amazon.com/bedrock/latest/userguide/setting-up.html) for more details.
+
+AWS session is automatically authenticated from your environment, but you can also explicitly set `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY` environment variables.
+
 ### Working with large PRs
 
 The default mode of CodiumAI is to have a single call per tool, using GPT-4, which has a token limit of 8000 tokens.

pr_agent/algo/__init__.py

@@ -18,4 +18,7 @@
     'vertex_ai/codechat-bison-32k': 32000,
     'codechat-bison': 6144,
     'codechat-bison-32k': 32000,
+    'anthropic.claude-v2': 100000,
+    'anthropic.claude-instant-v1': 100000,
+    'anthropic.claude-v1': 100000,
 }

pr_agent/algo/ai_handler.py

@@ -1,5 +1,6 @@
 import os
 
+import boto3
 import litellm
 import openai
 from litellm import acompletion
@@ -24,6 +25,7 @@ def __init__(self):
         Raises a ValueError if the OpenAI key is missing.
         """
         self.azure = False
+        self.aws_bedrock_client = None
 
         if get_settings().get("OPENAI.KEY", None):
             openai.api_key = get_settings().openai.key
@@ -60,6 +62,12 @@ def __init__(self):
             litellm.vertex_location = get_settings().get(
                 "VERTEXAI.VERTEX_LOCATION", None
             )
+        if get_settings().get("AWS.BEDROCK_REGION", None):
+            litellm.AmazonAnthropicConfig.max_tokens_to_sample = 2000
+            self.aws_bedrock_client = boto3.client(
+                service_name="bedrock-runtime",
+                region_name=get_settings().aws.bedrock_region,
+            )
 
     @property
     def deployment_id(self):
@@ -100,13 +108,16 @@ async def chat_completion(self, model: str, system: str, user: str, temperature:
             if self.azure:
                 model = 'azure/' + model
             messages = [{"role": "system", "content": system}, {"role": "user", "content": user}]
-            response = await acompletion(
-                model=model,
-                deployment_id=deployment_id,
-                messages=messages,
-                temperature=temperature,
-                force_timeout=get_settings().config.ai_timeout
-            )
+            kwargs = {
+                "model": model,
+                "deployment_id": deployment_id,
+                "messages": messages,
+                "temperature": temperature,
+                "force_timeout": get_settings().config.ai_timeout,
+            }
+            if self.aws_bedrock_client:
+                kwargs["aws_bedrock_client"] = self.aws_bedrock_client
+            response = await acompletion(**kwargs)
         except (APIError, Timeout, TryAgain) as e:
             get_logger().error("Error during OpenAI inference: ", e)
             raise

pr_agent/algo/utils.py

@@ -325,7 +325,15 @@ def try_fix_yaml(response_text: str) -> dict:
             break
         except:
             pass
-    return data
+
+    # thrid fallback - try to remove leading and trailing curly brackets
+    response_text_copy = response_text.strip().rstrip().removeprefix('{').removesuffix('}')
+    try:
+        data = yaml.safe_load(response_text_copy,)
+        get_logger().info(f"Successfully parsed AI prediction after removing curly brackets")
+        return data
+    except:
+        pass
 
 
 def set_custom_labels(variables):

pr_agent/settings/.secrets_template.toml

@@ -40,6 +40,9 @@ api_base = "" # the base url for your local Llama 2, Code Llama, and other model
 vertex_project = "" # the google cloud platform project name for your vertexai deployment
 vertex_location = "" # the google cloud platform location for your vertexai deployment
 
+[aws]
+bedrock_region = "" # the AWS region to call Bedrock APIs
+
 [github]
 # ---- Set the following only for deployment type == "user"
 user_token = ""  # A GitHub personal access token with 'repo' scope.

requirements.txt

@@ -14,7 +14,7 @@ GitPython==3.1.32
 PyYAML==6.0
 starlette-context==0.3.6
 litellm==0.12.5
-boto3==1.28.25
+boto3==1.33.1
 google-cloud-storage==2.10.0
 ujson==5.8.0
 azure-devops==7.1.0b3