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

bug fixes and updates

Browse Source
This commit is contained in:
Krrish Dholakia
2023-08-03 16:05:46 -07:00
parent 102edcdcf1
commit ed8554699b
7 changed files with 42 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
pr_agent/algo/ai_handler.py
View File

@ -3,6 +3,7 @@ import logging
import openai
from openai.error import APIError, RateLimitError, Timeout, TryAgain
from retry import retry
import litellm
from litellm import acompletion
from pr_agent.config_loader import get_settings
@ -22,6 +23,7 @@ class AiHandler:
"""
try:
openai.api_key = get_settings().openai.key
litellm.openai_key = get_settings().openai.key
if get_settings().get("OPENAI.ORG", None):
openai.organization = get_settings().openai.org
self.deployment_id = get_settings().get("OPENAI.DEPLOYMENT_ID", None)
@ -31,6 +33,9 @@ class AiHandler:
openai.api_version = get_settings().openai.api_version
if get_settings().get("OPENAI.API_BASE", None):
openai.api_base = get_settings().openai.api_base
litellm.api_base = get_settings().openai.api_base
if get_settings().get("LITE.KEY", None):
self.llm_api_key = get_settings().lite.key
except AttributeError as e:
raise ValueError("OpenAI key is required") from e
@ -65,6 +70,7 @@ class AiHandler:
{"role": "user", "content": user}
],
temperature=temperature,
api_key=self.llm_api_key
)
except (APIError, Timeout, TryAgain) as e:
logging.error("Error during OpenAI inference: ", e)
@ -75,8 +81,9 @@ class AiHandler:
except (Exception) as e:
logging.error("Unknown error during OpenAI inference: ", e)
raise TryAgain from e
if response is None or len(response.choices) == 0:
if response is None or len(response["choices"]) == 0:
raise TryAgain
resp = response.choices[0]['message']['content']
finish_reason = response.choices[0].finish_reason
resp = response["choices"][0]['message']['content']
finish_reason = response["choices"][0]["finish_reason"]
print(resp, finish_reason)
return resp, finish_reason