xiaxinhang/pr-agent
1
0
Fork 0
mirror of https://github.com/qodo-ai/pr-agent.git synced 2025-07-04 12:50:38 +08:00
Code Issues Packages Projects Releases Wiki Activity

Support Google's Vertex AI

Browse Source
This commit is contained in:
Rhys Tyers
2023-11-07 09:13:08 +00:00
parent e878bbbe36
commit 6b4b16dcf9
3 changed files with 42 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
pr_agent/algo/__init__.py
View File

@ -13,5 +13,9 @@ MAX_TOKENS = {
'claude-2': 100000, 'claude-2': 100000,
'command-nightly': 4096, 'command-nightly': 4096,
'replicate/llama-2-70b-chat:2c1608e18606fad2812020dc541930f2d0495ce32eee50074220b87300bc16e1': 4096, 'replicate/llama-2-70b-chat:2c1608e18606fad2812020dc541930f2d0495ce32eee50074220b87300bc16e1': 4096,
'meta-llama/Llama-2-7b-chat-hf': 4096 'meta-llama/Llama-2-7b-chat-hf': 4096,
'vertex_ai/codechat-bison': 6144,
'vertex_ai/codechat-bison-32k': 32000,
'codechat-bison': 6144,
'codechat-bison-32k': 32000,
} }

12
pr_agent/algo/ai_handler.py
View File

@ -23,7 +23,9 @@ class AiHandler:
Initializes the OpenAI API key and other settings from a configuration file. Initializes the OpenAI API key and other settings from a configuration file.
Raises a ValueError if the OpenAI key is missing. Raises a ValueError if the OpenAI key is missing.
""" """
try: self.azure = False
if get_settings().get("OPENAI.KEY", None):
openai.api_key = get_settings().openai.key openai.api_key = get_settings().openai.key
litellm.openai_key = get_settings().openai.key litellm.openai_key = get_settings().openai.key
if get_settings().get("litellm.use_client"): if get_settings().get("litellm.use_client"):
@ -31,7 +33,6 @@ class AiHandler:
assert litellm_token, "LITELLM_TOKEN is required" assert litellm_token, "LITELLM_TOKEN is required"
os.environ["LITELLM_TOKEN"] = litellm_token os.environ["LITELLM_TOKEN"] = litellm_token
litellm.use_client = True litellm.use_client = True
self.azure = False
if get_settings().get("OPENAI.ORG", None): if get_settings().get("OPENAI.ORG", None):
litellm.organization = get_settings().openai.org litellm.organization = get_settings().openai.org
if get_settings().get("OPENAI.API_TYPE", None): if get_settings().get("OPENAI.API_TYPE", None):
@ -54,8 +55,11 @@ class AiHandler:
litellm.huggingface_key = get_settings().huggingface.key litellm.huggingface_key = get_settings().huggingface.key
if get_settings().get("HUGGINGFACE.API_BASE", None): if get_settings().get("HUGGINGFACE.API_BASE", None):
litellm.api_base = get_settings().huggingface.api_base litellm.api_base = get_settings().huggingface.api_base
except AttributeError as e: if get_settings().get("VERTEXAI.VERTEX_PROJECT", None):
raise ValueError("OpenAI key is required") from e litellm.vertex_project = get_settings().vertexai.vertex_project
litellm.vertex_location = get_settings().get(
"VERTEXAI.VERTEX_LOCATION", None
)
@property @property
def deployment_id(self): def deployment_id(self):

3
requirements.txt
View File

@ -13,7 +13,7 @@ atlassian-python-api==3.39.0
GitPython==3.1.32 GitPython==3.1.32
PyYAML==6.0 PyYAML==6.0
starlette-context==0.3.6 starlette-context==0.3.6
litellm~=0.1.574 litellm==0.12.5
boto3==1.28.25 boto3==1.28.25
google-cloud-storage==2.10.0 google-cloud-storage==2.10.0
ujson==5.8.0 ujson==5.8.0
@ -22,3 +22,4 @@ msrest==0.7.1
pinecone-client pinecone-client
pinecone-datasets @ git+https://github.com/mrT23/pinecone-datasets.git@main pinecone-datasets @ git+https://github.com/mrT23/pinecone-datasets.git@main
loguru==0.7.2 loguru==0.7.2
google-cloud-aiplatform==1.35.0