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

bug fixes

Browse Source
This commit is contained in:
Krrish Dholakia
2023-08-05 22:50:41 -07:00
parent ed8554699b
commit 0f975ccf4a
3 changed files with 17 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
pr_agent/algo/ai_handler.py
View File

@ -6,7 +6,7 @@ from retry import retry
import litellm import litellm
from litellm import acompletion from litellm import acompletion
from pr_agent.config_loader import get_settings from pr_agent.config_loader import get_settings
import traceback
OPENAI_RETRIES=5 OPENAI_RETRIES=5
class AiHandler: class AiHandler:
@ -24,18 +24,24 @@ class AiHandler:
try: try:
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
self.azure = False
if get_settings().get("OPENAI.ORG", None): if get_settings().get("OPENAI.ORG", None):
openai.organization = get_settings().openai.org litellm.organization = get_settings().openai.org
self.deployment_id = get_settings().get("OPENAI.DEPLOYMENT_ID", None) self.deployment_id = get_settings().get("OPENAI.DEPLOYMENT_ID", None)
if get_settings().get("OPENAI.API_TYPE", None): if get_settings().get("OPENAI.API_TYPE", None):
openai.api_type = get_settings().openai.api_type if get_settings().openai.api_type == "azure":
self.azure = True
litellm.azure_key = get_settings().openai.key
if get_settings().get("OPENAI.API_VERSION", None): if get_settings().get("OPENAI.API_VERSION", None):
openai.api_version = get_settings().openai.api_version litellm.api_version = get_settings().openai.api_version
if get_settings().get("OPENAI.API_BASE", None): if get_settings().get("OPENAI.API_BASE", None):
openai.api_base = get_settings().openai.api_base
litellm.api_base = get_settings().openai.api_base litellm.api_base = get_settings().openai.api_base
if get_settings().get("LITE.KEY", None): if get_settings().get("ANTHROPIC.KEY", None):
self.llm_api_key = get_settings().lite.key litellm.anthropic_key = get_settings().anthropic.key
if get_settings().get("COHERE.KEY", None):
litellm.cohere_key = get_settings().cohere.key
if get_settings().get("REPLICATE.KEY", None):
litellm.replicate_key = get_settings().replicate.key
except AttributeError as e: except AttributeError as e:
raise ValueError("OpenAI key is required") from e raise ValueError("OpenAI key is required") from e
@ -70,7 +76,7 @@ class AiHandler:
{"role": "user", "content": user} {"role": "user", "content": user}
], ],
temperature=temperature, temperature=temperature,
api_key=self.llm_api_key azure=self.azure
) )
except (APIError, Timeout, TryAgain) as e: except (APIError, Timeout, TryAgain) as e:
logging.error("Error during OpenAI inference: ", e) logging.error("Error during OpenAI inference: ", e)

2
pyproject.toml
View File

@ -42,7 +42,7 @@ dependencies = [
"atlassian-python-api==3.39.0", "atlassian-python-api==3.39.0",
"GitPython~=3.1.32", "GitPython~=3.1.32",
"starlette-context==0.3.6", "starlette-context==0.3.6",
"litellm==0.1.2291" "litellm~=0.1.351"
] ]
[project.urls] [project.urls]

3
requirements.txt
View File

@ -10,4 +10,5 @@ python-gitlab==3.15.0
pytest~=7.4.0 pytest~=7.4.0
aiohttp~=3.8.4 aiohttp~=3.8.4
atlassian-python-api==3.39.0 atlassian-python-api==3.39.0
GitPython~=3.1.32 GitPython~=3.1.32
litellm~=0.1.351