Merge pull request #1828 from Akileox/refactor-langchain-handler

Refactor: Enhance AI Handler Robustness, Interface Compliance, and Asynchronous Operations (Resolves #1784)

pr_agent/algo/ai_handlers/langchain_ai_handler.py

@@ -1,6 +1,9 @@
+_LANGCHAIN_INSTALLED = False
+
 try:
     from langchain_core.messages import HumanMessage, SystemMessage
     from langchain_openai import AzureChatOpenAI, ChatOpenAI
+    _LANGCHAIN_INSTALLED = True
 except:  # we don't enforce langchain as a dependency, so if it's not installed, just move on
     pass
 
@@ -8,6 +11,7 @@ import functools
 
 import openai
 from tenacity import retry, retry_if_exception_type, retry_if_not_exception_type, stop_after_attempt
+from langchain_core.runnables import Runnable
 
 from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler
 from pr_agent.config_loader import get_settings
@@ -18,17 +22,14 @@ OPENAI_RETRIES = 5
 
 class LangChainOpenAIHandler(BaseAiHandler):
     def __init__(self):
-        # Initialize OpenAIHandler specific attributes here
+        if not _LANGCHAIN_INSTALLED:
+            error_msg = "LangChain is not installed. Please install it with `pip install langchain`."
+            get_logger().error(error_msg)
+            raise ImportError(error_msg)
+        
         super().__init__()
         self.azure = get_settings().get("OPENAI.API_TYPE", "").lower() == "azure"
 
-        # Create a default unused chat object to trigger early validation
-        self._create_chat(self.deployment_id)
-
-    def chat(self, messages: list, model: str, temperature: float):
-        chat = self._create_chat(self.deployment_id)
-        return chat.invoke(input=messages, model=model, temperature=temperature)
-
     @property
     def deployment_id(self):
         """
@@ -36,16 +37,66 @@ class LangChainOpenAIHandler(BaseAiHandler):
         """
         return get_settings().get("OPENAI.DEPLOYMENT_ID", None)
 
+    async def _create_chat_async(self, deployment_id=None):
+        try:
+            if self.azure:
+                # Using Azure OpenAI service
+                return AzureChatOpenAI(
+                    openai_api_key=get_settings().openai.key,
+                    openai_api_version=get_settings().openai.api_version,
+                    azure_deployment=deployment_id,
+                    azure_endpoint=get_settings().openai.api_base,
+                )
+            else:
+                # Using standard OpenAI or other LLM services
+                openai_api_base = get_settings().get("OPENAI.API_BASE", None)
+                if openai_api_base is None or len(openai_api_base) == 0:
+                    return ChatOpenAI(openai_api_key=get_settings().openai.key)
+                else:
+                    return ChatOpenAI(
+                        openai_api_key=get_settings().openai.key, 
+                        openai_api_base=openai_api_base
+                    )
+        except AttributeError as e:
+            # Handle configuration errors
+            error_msg = f"OpenAI {e.name} is required" if getattr(e, "name") else str(e)
+            get_logger().error(error_msg)
+            raise ValueError(error_msg) from e
+
     @retry(
         retry=retry_if_exception_type(openai.APIError) & retry_if_not_exception_type(openai.RateLimitError),
         stop=stop_after_attempt(OPENAI_RETRIES),
     )
-    async def chat_completion(self, model: str, system: str, user: str, temperature: float = 0.2):
+    async def chat_completion(self, model: str, system: str, user: str, temperature: float = 0.2, img_path: str = None):
+        if img_path:
+            get_logger().warning(f"Image path is not supported for LangChainOpenAIHandler. Ignoring image path: {img_path}")
         try:
             messages = [SystemMessage(content=system), HumanMessage(content=user)]
+            llm = await self._create_chat_async(deployment_id=self.deployment_id)
+            
+            if not isinstance(llm, Runnable):
+                error_message = (
+                    f"The Langchain LLM object ({type(llm)}) does not implement the Runnable interface. "
+                    f"Please update your Langchain library to the latest version or "
+                    f"check your LLM configuration to support async calls. "
+                    f"PR-Agent is designed to utilize Langchain's async capabilities."
+                )
+                get_logger().error(error_message)
+                raise NotImplementedError(error_message)
+
+            # Handle parameters based on LLM type
+            if isinstance(llm, (ChatOpenAI, AzureChatOpenAI)):
+                # OpenAI models support all parameters
+                resp = await llm.ainvoke(
+                    input=messages,
+                    model=model,
+                    temperature=temperature
+                )
+            else:
+                # Other LLMs (like Gemini) only support input parameter
+                get_logger().info(f"Using simplified ainvoke for {type(llm)}")
+                resp = await llm.ainvoke(input=messages)
 
-            # get a chat completion from the formatted messages
-            resp = self.chat(messages, model=model, temperature=temperature)
             finish_reason = "completed"
             return resp.content, finish_reason
 
@@ -58,27 +109,3 @@ class LangChainOpenAIHandler(BaseAiHandler):
         except Exception as e:
             get_logger().warning(f"Unknown error during LLM inference: {e}")
             raise openai.APIError from e
-
-    def _create_chat(self, deployment_id=None):
-        try:
-            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
-                return AzureChatOpenAI(
-                    openai_api_key=get_settings().openai.key,
-                    openai_api_version=get_settings().openai.api_version,
-                    azure_deployment=deployment_id,
-                    azure_endpoint=get_settings().openai.api_base,
-                )
-            else:
-                # for llms that compatible with openai, should use custom api base
-                openai_api_base = get_settings().get("OPENAI.API_BASE", None)
-                if openai_api_base is None or len(openai_api_base) == 0:
-                    return ChatOpenAI(openai_api_key=get_settings().openai.key)
-                else:
-                    return ChatOpenAI(openai_api_key=get_settings().openai.key, openai_api_base=openai_api_base)
-        except AttributeError as e:
-            if getattr(e, "name"):
-                raise ValueError(f"OpenAI {e.name} is required") from e
-            else:
-                raise e

pr_agent/algo/ai_handlers/openai_ai_handler.py

@@ -42,8 +42,10 @@ class OpenAIHandler(BaseAiHandler):
         retry=retry_if_exception_type(openai.APIError) & retry_if_not_exception_type(openai.RateLimitError),
         stop=stop_after_attempt(OPENAI_RETRIES),
     )
-    async def chat_completion(self, model: str, system: str, user: str, temperature: float = 0.2):
+    async def chat_completion(self, model: str, system: str, user: str, temperature: float = 0.2, img_path: str = None):
         try:
+            if img_path:
+                get_logger().warning(f"Image path is not supported for OpenAIHandler. Ignoring image path: {img_path}")
             get_logger().info("System: ", system)
             get_logger().info("User: ", user)
             messages = [{"role": "system", "content": system}, {"role": "user", "content": user}]