Merge branch 'base-ai-handler' into abstract-BaseAiHandler

This commit is contained in:
Brian Pham
2023-12-14 07:44:13 +08:00
21 changed files with 256 additions and 57 deletions

View 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

View 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

View File

@ -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

View 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

View File

@ -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} &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</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