xiaxinhang/pr-agent
1
0
Fork 0
mirror of https://github.com/qodo-ai/pr-agent.git synced 2025-07-04 12:50:38 +08:00
Code Issues Packages Projects Releases Wiki Activity

Update langchain

Browse Source
This commit is contained in:
Brian Pham
2023-12-12 23:28:58 +08:00
parent 0c66554d50
commit a627dcd64f
2 changed files with 13 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -1,41 +1,37 @@
from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler
from langchain.chat_models import ChatOpenAI from langchain.chat_models import ChatOpenAI
from langchain.prompts.chat import (
ChatPromptTemplate,
HumanMessagePromptTemplate,
SystemMessagePromptTemplate,
)
from langchain.schema import SystemMessage, HumanMessage from langchain.schema import SystemMessage, HumanMessage
from pr_agent.config_loader import get_settings from pr_agent.config_loader import get_settings
from pr_agent.log import get_logger from pr_agent.log import get_logger
OPENAI_RETRIES = 5 class LangChainOpenAIHandler(BaseAiHandler):
chat = ChatOpenAI(openai_api_key = get_settings().openai.key)
class LangChainAIHandler(BaseAiHandler):
def __init__(self): def __init__(self):
# Initialize OpenAIHandler specific attributes here # Initialize OpenAIHandler specific attributes here
try: try:
super().__init__() super().__init__()
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 raise ValueError("OpenAI key is required") from e
@property
def chat(self):
return self._chat
@property @property
def deployment_id(self): def deployment_id(self):
""" """
Returns the deployment ID for the OpenAI API. Returns the deployment ID for the OpenAI API.
""" """
return get_settings().get("OPENAI.DEPLOYMENT_ID", None) return get_settings().get("OPENAI.DEPLOYMENT_ID", None)
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):
try: try:
get_logger().info("model: ", model)
messages=[SystemMessage(content=system), HumanMessage(content=user)] messages=[SystemMessage(content=system), HumanMessage(content=user)]
# get a chat completion from the formatted messages # 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) get_logger().info("AI response: ", resp.content)
finish_reason="completed" finish_reason="completed"
return resp.content, finish_reason return resp.content, finish_reason

4
pr_agent/algo/utils.py
View File

@ -11,7 +11,7 @@ import yaml
from pr_agent.algo.ai_handlers.litellm_ai_handler import LiteLLMAiHandler 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.base_ai_handler import BaseAiHandler
from pr_agent.algo.ai_handlers.openai_ai_handler import OpenAIHandler 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 starlette_context import context
from pr_agent.config_loader import get_settings, global_settings from pr_agent.config_loader import get_settings, global_settings
from pr_agent.log import get_logger from pr_agent.log import get_logger
@ -310,4 +310,4 @@ def try_fix_yaml(review_text: str) -> dict:
return data return data
def get_ai_handler() -> BaseAiHandler: def get_ai_handler() -> BaseAiHandler:
return LangChainAIHandler() return LangChainOpenAIHandler()