From b0f9b96c7553ff20250e88c2e4565a223b491306 Mon Sep 17 00:00:00 2001
From: evalphobia <evalphobia@gmail.com>
Date: Thu, 13 Jun 2024 14:34:14 +0900
Subject: [PATCH] Support models: Anthropic Calude 3 on Vertex AI

---
 pr_agent/algo/__init__.py | 3 +++
 requirements.txt          | 3 ++-
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/pr_agent/algo/__init__.py b/pr_agent/algo/__init__.py
index ad89b052..1649615f 100644
--- a/pr_agent/algo/__init__.py
+++ b/pr_agent/algo/__init__.py
@@ -23,6 +23,9 @@ MAX_TOKENS = {
     'meta-llama/Llama-2-7b-chat-hf': 4096,
     'vertex_ai/codechat-bison': 6144,
     'vertex_ai/codechat-bison-32k': 32000,
+    'vertex_ai/claude-3-haiku@20240307': 100000,
+    'vertex_ai/claude-3-sonnet@20240229': 100000,
+    'vertex_ai/claude-3-opus@20240229': 100000,
     'codechat-bison': 6144,
     'codechat-bison-32k': 32000,
     'anthropic.claude-instant-v1': 100000,
diff --git a/requirements.txt b/requirements.txt
index 283a319a..d55c1e14 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,4 +1,5 @@
 aiohttp==3.9.1
+anthropic[vertex]==0.21.3
 atlassian-python-api==3.41.4
 azure-devops==7.1.0b3
 azure-identity==1.15.0
@@ -9,7 +10,7 @@ GitPython==3.1.32
 google-cloud-aiplatform==1.35.0
 google-cloud-storage==2.10.0
 Jinja2==3.1.2
-litellm==1.31.10
+litellm==1.34.42
 loguru==0.7.2
 msrest==0.7.1
 openai==1.13.3