Add support for Azure OpenAI in LangChainOpenAIHandler

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.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):