From a627dcd64fe1de6e740e385323bf9b451d6d2772 Mon Sep 17 00:00:00 2001
From: Brian Pham <brian_pham@singaporeair.com.sg>
Date: Tue, 12 Dec 2023 23:28:58 +0800
Subject: [PATCH] Update langchain

---
 .../algo/ai_handlers/langchain_ai_handler.py  | 26 ++++++++-----------
 pr_agent/algo/utils.py                        |  4 +--
 2 files changed, 13 insertions(+), 17 deletions(-)

diff --git a/pr_agent/algo/ai_handlers/langchain_ai_handler.py b/pr_agent/algo/ai_handlers/langchain_ai_handler.py
index 87cf1072..bc26e624 100644
--- a/pr_agent/algo/ai_handlers/langchain_ai_handler.py
+++ b/pr_agent/algo/ai_handlers/langchain_ai_handler.py
@@ -1,41 +1,37 @@
 from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler
 from langchain.chat_models import ChatOpenAI
-from langchain.prompts.chat import (
-    ChatPromptTemplate,
-    HumanMessagePromptTemplate,
-    SystemMessagePromptTemplate,
-)
 from langchain.schema import SystemMessage, HumanMessage
-
-
 from pr_agent.config_loader import get_settings
 from pr_agent.log import get_logger
 
-OPENAI_RETRIES = 5
-chat = ChatOpenAI(openai_api_key = get_settings().openai.key)
-
-class LangChainAIHandler(BaseAiHandler):
+class LangChainOpenAIHandler(BaseAiHandler):
     def __init__(self):
         # Initialize OpenAIHandler specific attributes here
         try:
             super().__init__()            
-        
+            self._chat = ChatOpenAI(openai_api_key=get_settings().openai.key)
+
         except AttributeError as e:
             raise ValueError("OpenAI key is required") from e
+    
+    @property
+    def chat(self):
+        return self._chat
+
     @property
     def deployment_id(self):
         """
         Returns the deployment ID for the OpenAI API.
         """
         return get_settings().get("OPENAI.DEPLOYMENT_ID", None)
-        
+    
     async def chat_completion(self, model: str, system: str, user: str, temperature: float = 0.2):
         try:
-
+            get_logger().info("model: ", model)
             messages=[SystemMessage(content=system), HumanMessage(content=user)]
             
             # get a chat completion from the formatted messages
-            resp = chat(messages, model=model, temperature=temperature)
+            resp = self.chat(messages, model=model, temperature=temperature)
             get_logger().info("AI response: ", resp.content)
             finish_reason="completed"
             return resp.content, finish_reason
diff --git a/pr_agent/algo/utils.py b/pr_agent/algo/utils.py
index d0b86b63..8c66d96e 100644
--- a/pr_agent/algo/utils.py
+++ b/pr_agent/algo/utils.py
@@ -11,7 +11,7 @@ import yaml
 from pr_agent.algo.ai_handlers.litellm_ai_handler import LiteLLMAiHandler
 from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler
 from pr_agent.algo.ai_handlers.openai_ai_handler import OpenAIHandler
-from pr_agent.algo.ai_handlers.langchain_ai_handler import LangChainAIHandler
+from pr_agent.algo.ai_handlers.langchain_ai_handler import LangChainOpenAIHandler
 from starlette_context import context
 from pr_agent.config_loader import get_settings, global_settings
 from pr_agent.log import get_logger
@@ -310,4 +310,4 @@ def try_fix_yaml(review_text: str) -> dict:
     return data
 
 def get_ai_handler() -> BaseAiHandler:
-    return LangChainAIHandler()
\ No newline at end of file
+    return LangChainOpenAIHandler()
\ No newline at end of file