mirror of
https://github.com/qodo-ai/pr-agent.git
synced 2025-07-13 09:10:38 +08:00
Merge branch 'base-ai-handler' into abstract-BaseAiHandler
This commit is contained in:
28
pr_agent/algo/ai_handlers/base_ai_handler.py
Normal file
28
pr_agent/algo/ai_handlers/base_ai_handler.py
Normal file
@ -0,0 +1,28 @@
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
class BaseAiHandler(ABC):
|
||||
"""
|
||||
This class defines the interface for an AI handler to be used by the PR Agents.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def deployment_id(self):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def chat_completion(self, model: str, system: str, user: str, temperature: float = 0.2):
|
||||
"""
|
||||
This method should be implemented to return a chat completion from the AI model.
|
||||
Args:
|
||||
model (str): the name of the model to use for the chat completion
|
||||
system (str): the system message string to use for the chat completion
|
||||
user (str): the user message string to use for the chat completion
|
||||
temperature (float): the temperature to use for the chat completion
|
||||
"""
|
||||
pass
|
||||
|
46
pr_agent/algo/ai_handlers/langchain_ai_handler.py
Normal file
46
pr_agent/algo/ai_handlers/langchain_ai_handler.py
Normal file
@ -0,0 +1,46 @@
|
||||
from langchain.chat_models import ChatOpenAI
|
||||
from langchain.schema import SystemMessage, HumanMessage
|
||||
|
||||
from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler
|
||||
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
|
||||
|
||||
OPENAI_RETRIES = 5
|
||||
|
||||
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)
|
||||
@retry(exceptions=(APIError, Timeout, TryAgain, AttributeError, RateLimitError),
|
||||
tries=OPENAI_RETRIES, delay=2, backoff=2, jitter=(1, 3))
|
||||
async def chat_completion(self, model: str, system: str, user: str, temperature: float = 0.2):
|
||||
try:
|
||||
messages=[SystemMessage(content=system), HumanMessage(content=user)]
|
||||
|
||||
# get a chat completion from the formatted messages
|
||||
resp = self.chat(messages, model=model, temperature=temperature)
|
||||
finish_reason="completed"
|
||||
return resp.content, finish_reason
|
||||
|
||||
except (Exception) as e:
|
||||
get_logger().error("Unknown error during OpenAI inference: ", e)
|
||||
raise e
|
@ -6,14 +6,14 @@ import openai
|
||||
from litellm import acompletion
|
||||
from openai.error import APIError, RateLimitError, Timeout, TryAgain
|
||||
from retry import retry
|
||||
from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler
|
||||
from pr_agent.config_loader import get_settings
|
||||
from pr_agent.algo.base_ai_handler import BaseAiHandler
|
||||
from pr_agent.log import get_logger
|
||||
|
||||
OPENAI_RETRIES = 5
|
||||
|
||||
|
||||
class AiHandler(BaseAiHandler):
|
||||
class LiteLLMAIHandler(BaseAiHandler):
|
||||
"""
|
||||
This class handles interactions with the OpenAI API for chat completions.
|
||||
It initializes the API key and other settings from a configuration file,
|
||||
@ -135,4 +135,4 @@ class AiHandler(BaseAiHandler):
|
||||
usage = response.get("usage")
|
||||
get_logger().info("AI response", response=resp, messages=messages, finish_reason=finish_reason,
|
||||
model=model, usage=usage)
|
||||
return resp, finish_reason
|
||||
return resp, finish_reason
|
67
pr_agent/algo/ai_handlers/openai_ai_handler.py
Normal file
67
pr_agent/algo/ai_handlers/openai_ai_handler.py
Normal file
@ -0,0 +1,67 @@
|
||||
from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler
|
||||
import openai
|
||||
from openai.error import APIError, RateLimitError, Timeout, TryAgain
|
||||
from retry import retry
|
||||
|
||||
from pr_agent.config_loader import get_settings
|
||||
from pr_agent.log import get_logger
|
||||
|
||||
OPENAI_RETRIES = 5
|
||||
|
||||
|
||||
class OpenAIHandler(BaseAiHandler):
|
||||
def __init__(self):
|
||||
# Initialize OpenAIHandler specific attributes here
|
||||
try:
|
||||
super().__init__()
|
||||
openai.api_key = get_settings().openai.key
|
||||
if get_settings().get("OPENAI.ORG", None):
|
||||
openai.organization = get_settings().openai.org
|
||||
if get_settings().get("OPENAI.API_TYPE", None):
|
||||
if get_settings().openai.api_type == "azure":
|
||||
self.azure = True
|
||||
openai.azure_key = get_settings().openai.key
|
||||
if get_settings().get("OPENAI.API_VERSION", None):
|
||||
openai.api_version = get_settings().openai.api_version
|
||||
if get_settings().get("OPENAI.API_BASE", None):
|
||||
openai.api_base = get_settings().openai.api_base
|
||||
|
||||
except AttributeError as e:
|
||||
raise ValueError("OpenAI key is required") from e
|
||||
@property
|
||||
def deployment_id(self):
|
||||
"""
|
||||
Returns the deployment ID for the OpenAI API.
|
||||
"""
|
||||
return get_settings().get("OPENAI.DEPLOYMENT_ID", None)
|
||||
|
||||
@retry(exceptions=(APIError, Timeout, TryAgain, AttributeError, RateLimitError),
|
||||
tries=OPENAI_RETRIES, delay=2, backoff=2, jitter=(1, 3))
|
||||
async def chat_completion(self, model: str, system: str, user: str, temperature: float = 0.2):
|
||||
try:
|
||||
deployment_id = self.deployment_id
|
||||
get_logger().info("System: ", system)
|
||||
get_logger().info("User: ", user)
|
||||
messages = [{"role": "system", "content": system}, {"role": "user", "content": user}]
|
||||
|
||||
chat_completion = await openai.ChatCompletion.acreate(
|
||||
model=model,
|
||||
deployment_id=deployment_id,
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
)
|
||||
resp = chat_completion["choices"][0]['message']['content']
|
||||
finish_reason = chat_completion["choices"][0]["finish_reason"]
|
||||
usage = chat_completion.get("usage")
|
||||
get_logger().info("AI response", response=resp, messages=messages, finish_reason=finish_reason,
|
||||
model=model, usage=usage)
|
||||
return resp, finish_reason
|
||||
except (APIError, Timeout, TryAgain) as e:
|
||||
get_logger().error("Error during OpenAI inference: ", e)
|
||||
raise
|
||||
except (RateLimitError) as e:
|
||||
get_logger().error("Rate limit error during OpenAI inference: ", e)
|
||||
raise
|
||||
except (Exception) as e:
|
||||
get_logger().error("Unknown error during OpenAI inference: ", e)
|
||||
raise TryAgain from e
|
@ -59,14 +59,14 @@ def convert_to_markdown(output_data: dict, gfm_supported: bool=True) -> str:
|
||||
if key.lower() == 'code feedback':
|
||||
if gfm_supported:
|
||||
markdown_text += f"\n\n- "
|
||||
markdown_text += f"<details><summary> { emoji } Code feedback:</summary>\n\n"
|
||||
markdown_text += f"<details><summary> { emoji } Code feedback:</summary>"
|
||||
else:
|
||||
markdown_text += f"\n\n- **{emoji} Code feedback:**\n\n"
|
||||
else:
|
||||
markdown_text += f"- {emoji} **{key}:**\n\n"
|
||||
for item in value:
|
||||
for i, item in enumerate(value):
|
||||
if isinstance(item, dict) and key.lower() == 'code feedback':
|
||||
markdown_text += parse_code_suggestion(item, gfm_supported)
|
||||
markdown_text += parse_code_suggestion(item, i, gfm_supported)
|
||||
elif item:
|
||||
markdown_text += f" - {item}\n"
|
||||
if key.lower() == 'code feedback':
|
||||
@ -80,7 +80,7 @@ def convert_to_markdown(output_data: dict, gfm_supported: bool=True) -> str:
|
||||
return markdown_text
|
||||
|
||||
|
||||
def parse_code_suggestion(code_suggestions: dict, gfm_supported: bool=True) -> str:
|
||||
def parse_code_suggestion(code_suggestions: dict, i: int = 0, gfm_supported: bool = True) -> str:
|
||||
"""
|
||||
Convert a dictionary of data into markdown format.
|
||||
|
||||
@ -91,24 +91,52 @@ def parse_code_suggestion(code_suggestions: dict, gfm_supported: bool=True) -> s
|
||||
str: A string containing the markdown formatted text generated from the input dictionary.
|
||||
"""
|
||||
markdown_text = ""
|
||||
for sub_key, sub_value in code_suggestions.items():
|
||||
if isinstance(sub_value, dict): # "code example"
|
||||
markdown_text += f" - **{sub_key}:**\n"
|
||||
for code_key, code_value in sub_value.items(): # 'before' and 'after' code
|
||||
code_str = f"```\n{code_value}\n```"
|
||||
code_str_indented = textwrap.indent(code_str, ' ')
|
||||
markdown_text += f" - **{code_key}:**\n{code_str_indented}\n"
|
||||
else:
|
||||
if "relevant file" in sub_key.lower():
|
||||
markdown_text += f"\n - **{sub_key}:** {sub_value} \n"
|
||||
if gfm_supported and 'relevant line' in code_suggestions:
|
||||
if i == 0:
|
||||
markdown_text += "<hr>"
|
||||
markdown_text += '<table>'
|
||||
for sub_key, sub_value in code_suggestions.items():
|
||||
try:
|
||||
if sub_key.lower() == 'relevant file':
|
||||
relevant_file = sub_value.strip('`').strip('"').strip("'")
|
||||
markdown_text += f"<tr><td>{sub_key}</td><td>{relevant_file}</td></tr>"
|
||||
# continue
|
||||
elif sub_key.lower() == 'suggestion':
|
||||
markdown_text += f"<tr><td>{sub_key} </td><td><strong>{sub_value}</strong></td></tr>"
|
||||
elif sub_key.lower() == 'relevant line':
|
||||
markdown_text += f"<tr><td>relevant line</td>"
|
||||
sub_value_list = sub_value.split('](')
|
||||
relevant_line = sub_value_list[0].lstrip('`').lstrip('[')
|
||||
if len(sub_value_list) > 1:
|
||||
link = sub_value_list[1].rstrip(')').strip('`')
|
||||
markdown_text += f"<td><a href={link}>{relevant_line}</a></td>"
|
||||
else:
|
||||
markdown_text += f"<td>{relevant_line}</td>"
|
||||
markdown_text += "</tr>"
|
||||
except Exception as e:
|
||||
get_logger().exception(f"Failed to parse code suggestion: {e}")
|
||||
pass
|
||||
markdown_text += '</table>'
|
||||
markdown_text += "<hr>"
|
||||
else:
|
||||
for sub_key, sub_value in code_suggestions.items():
|
||||
if isinstance(sub_value, dict): # "code example"
|
||||
markdown_text += f" - **{sub_key}:**\n"
|
||||
for code_key, code_value in sub_value.items(): # 'before' and 'after' code
|
||||
code_str = f"```\n{code_value}\n```"
|
||||
code_str_indented = textwrap.indent(code_str, ' ')
|
||||
markdown_text += f" - **{code_key}:**\n{code_str_indented}\n"
|
||||
else:
|
||||
markdown_text += f" **{sub_key}:** {sub_value} \n"
|
||||
if not gfm_supported:
|
||||
if "relevant line" not in sub_key.lower(): # nicer presentation
|
||||
if "relevant file" in sub_key.lower():
|
||||
markdown_text += f"\n - **{sub_key}:** {sub_value} \n"
|
||||
else:
|
||||
markdown_text += f" **{sub_key}:** {sub_value} \n"
|
||||
if not gfm_supported:
|
||||
if "relevant line" not in sub_key.lower(): # nicer presentation
|
||||
# markdown_text = markdown_text.rstrip('\n') + "\\\n" # works for gitlab
|
||||
markdown_text = markdown_text.rstrip('\n') + " \n" # works for gitlab and bitbucker
|
||||
|
||||
markdown_text += "\n"
|
||||
markdown_text += "\n"
|
||||
return markdown_text
|
||||
|
||||
|
||||
@ -336,7 +364,7 @@ def try_fix_yaml(response_text: str) -> dict:
|
||||
pass
|
||||
|
||||
|
||||
def set_custom_labels(variables):
|
||||
def set_custom_labels(variables, git_provider=None):
|
||||
if not get_settings().config.enable_custom_labels:
|
||||
return
|
||||
|
||||
@ -348,11 +376,8 @@ def set_custom_labels(variables):
|
||||
labels_list = f" - {labels_list}" if labels_list else ""
|
||||
variables["custom_labels"] = labels_list
|
||||
return
|
||||
#final_labels = ""
|
||||
#for k, v in labels.items():
|
||||
# final_labels += f" - {k} ({v['description']})\n"
|
||||
#variables["custom_labels"] = final_labels
|
||||
#variables["custom_labels_examples"] = f" - {list(labels.keys())[0]}"
|
||||
|
||||
# Set custom labels
|
||||
variables["custom_labels_class"] = "class Label(str, Enum):"
|
||||
for k, v in labels.items():
|
||||
description = v['description'].strip('\n').replace('\n', '\\n')
|
||||
@ -422,4 +447,4 @@ def clip_tokens(text: str, max_tokens: int, add_three_dots=True) -> str:
|
||||
return clipped_text
|
||||
except Exception as e:
|
||||
get_logger().warning(f"Failed to clip tokens: {e}")
|
||||
return text
|
||||
return text
|
Reference in New Issue
Block a user