From ac7aaa0cd357deb4d074e9b2f3d4ab1b44acca1f Mon Sep 17 00:00:00 2001 From: zmeir Date: Thu, 4 Jan 2024 16:08:10 +0200 Subject: [PATCH] Add support for Azure OpenAI in LangChainOpenAIHandler --- .../algo/ai_handlers/langchain_ai_handler.py | 32 +++++++++++++++---- 1 file changed, 25 insertions(+), 7 deletions(-) diff --git a/pr_agent/algo/ai_handlers/langchain_ai_handler.py b/pr_agent/algo/ai_handlers/langchain_ai_handler.py index 3e31bcb8..a7c6d345 100644 --- a/pr_agent/algo/ai_handlers/langchain_ai_handler.py +++ b/pr_agent/algo/ai_handlers/langchain_ai_handler.py @@ -1,5 +1,5 @@ try: - from langchain.chat_models import ChatOpenAI + from langchain.chat_models import ChatOpenAI, AzureChatOpenAI from langchain.schema import SystemMessage, HumanMessage except: # we don't enforce langchain as a dependency, so if it's not installed, just move on pass @@ -9,23 +9,41 @@ from pr_agent.config_loader import get_settings from pr_agent.log import get_logger from openai.error import APIError, RateLimitError, Timeout, TryAgain -from retry import retry +from retry import retry +import functools OPENAI_RETRIES = 5 class LangChainOpenAIHandler(BaseAiHandler): def __init__(self): # Initialize OpenAIHandler specific attributes here + super().__init__() + self.azure = get_settings().get("OPENAI.API_TYPE", "").lower() == "azure" try: - super().__init__() - self._chat = ChatOpenAI(openai_api_key=get_settings().openai.key) - + if self.azure: + # 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: - 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 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 def deployment_id(self):