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

Add support for Azure OpenAI in LangChainOpenAIHandler

Browse Source
This commit is contained in:
zmeir
2024-01-04 16:08:10 +02:00
parent b7af45166a
commit ac7aaa0cd3
1 changed files with 25 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
pr_agent/algo/ai_handlers/langchain_ai_handler.py
View File

@ -1,5 +1,5 @@
try: try:
from langchain.chat_models import ChatOpenAI from langchain.chat_models import ChatOpenAI, AzureChatOpenAI
from langchain.schema import SystemMessage, HumanMessage from langchain.schema import SystemMessage, HumanMessage
except: # we don't enforce langchain as a dependency, so if it's not installed, just move on except: # we don't enforce langchain as a dependency, so if it's not installed, just move on
pass pass
@ -9,23 +9,41 @@ from pr_agent.config_loader import get_settings
from pr_agent.log import get_logger from pr_agent.log import get_logger
from openai.error import APIError, RateLimitError, Timeout, TryAgain from openai.error import APIError, RateLimitError, Timeout, TryAgain
from retry import retry from retry import retry
import functools
OPENAI_RETRIES = 5 OPENAI_RETRIES = 5
class LangChainOpenAIHandler(BaseAiHandler): class LangChainOpenAIHandler(BaseAiHandler):
def __init__(self): def __init__(self):
# Initialize OpenAIHandler specific attributes here # Initialize OpenAIHandler specific attributes here
super().__init__()
self.azure = get_settings().get("OPENAI.API_TYPE", "").lower() == "azure"
try: try:
super().__init__() if self.azure:
self._chat = ChatOpenAI(openai_api_key=get_settings().openai.key) # using a partial function so we can set the deployment_id later to support fallback_deployments
# but still need to access the other settings now so we can raise a proper exception if they're missing
self._chat = functools.partial(
lambda **kwargs: AzureChatOpenAI(**kwargs),
openai_api_key=get_settings().openai.key,
openai_api_base=get_settings().openai.api_base,
openai_api_version=get_settings().openai.api_version,
)
else:
self._chat = ChatOpenAI(openai_api_key=get_settings().openai.key)
except AttributeError as e: except AttributeError as e:
raise ValueError("OpenAI key is required") from e if getattr(e, "name"):
raise ValueError(f"OpenAI {e.name} is required") from e
else:
raise e
@property @property
def chat(self): def chat(self):
return self._chat if self.azure:
# we must set the deployment_id only here (instead of the __init__ method) to support fallback_deployments
return self._chat(deployment_name=self.deployment_id)
else:
return self._chat
@property @property
def deployment_id(self): def deployment_id(self):