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

Merge remote-tracking branch 'upstream/main' into abstract-BaseAiHandler

Browse Source
This commit is contained in:
Brian Pham
2023-12-09 16:47:13 +00:00
parent f2abe5c73e a7a0de764c
commit c0303ff9ec
104 changed files with 3813 additions and 1068 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
pr_agent/algo/__init__.py
View File

@ -8,9 +8,17 @@ MAX_TOKENS = {
'gpt-4': 8000,
'gpt-4-0613': 8000,
'gpt-4-32k': 32000,
'gpt-4-1106-preview': 128000, # 128K, but may be limited by config.max_model_tokens
'claude-instant-1': 100000,
'claude-2': 100000,
'command-nightly': 4096,
'replicate/llama-2-70b-chat:2c1608e18606fad2812020dc541930f2d0495ce32eee50074220b87300bc16e1': 4096,
'meta-llama/Llama-2-7b-chat-hf': 4096
'meta-llama/Llama-2-7b-chat-hf': 4096,
'vertex_ai/codechat-bison': 6144,
'vertex_ai/codechat-bison-32k': 32000,
'codechat-bison': 6144,
'codechat-bison-32k': 32000,
'anthropic.claude-v2': 100000,
'anthropic.claude-instant-v1': 100000,
'anthropic.claude-v1': 100000,
}

113
pr_agent/algo/ai_handler.py
View File

@ -1,6 +1,6 @@
import logging
import os
import boto3
import litellm
import openai
from litellm import acompletion
@ -8,6 +8,8 @@ from openai.error import APIError, RateLimitError, Timeout, TryAgain
from retry import retry
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
@ -23,39 +25,50 @@ class AiHandler(BaseAiHandler):
Initializes the OpenAI API key and other settings from a configuration file.
Raises a ValueError if the OpenAI key is missing.
"""
try:
self.azure = False
self.aws_bedrock_client = None
if get_settings().get("OPENAI.KEY", None):
openai.api_key = get_settings().openai.key
litellm.openai_key = get_settings().openai.key
if get_settings().get("litellm.use_client"):
litellm_token = get_settings().get("litellm.LITELLM_TOKEN")
assert litellm_token, "LITELLM_TOKEN is required"
os.environ["LITELLM_TOKEN"] = litellm_token
litellm.use_client = True
self.azure = False
if get_settings().get("OPENAI.ORG", None):
litellm.organization = get_settings().openai.org
if get_settings().get("OPENAI.API_TYPE", None):
if get_settings().openai.api_type == "azure":
self.azure = True
litellm.azure_key = get_settings().openai.key
if get_settings().get("OPENAI.API_VERSION", None):
litellm.api_version = get_settings().openai.api_version
if get_settings().get("OPENAI.API_BASE", None):
litellm.api_base = get_settings().openai.api_base
if get_settings().get("ANTHROPIC.KEY", None):
litellm.anthropic_key = get_settings().anthropic.key
if get_settings().get("COHERE.KEY", None):
litellm.cohere_key = get_settings().cohere.key
if get_settings().get("REPLICATE.KEY", None):
litellm.replicate_key = get_settings().replicate.key
if get_settings().get("REPLICATE.KEY", None):
litellm.replicate_key = get_settings().replicate.key
if get_settings().get("HUGGINGFACE.KEY", None):
litellm.huggingface_key = get_settings().huggingface.key
if get_settings().get("HUGGINGFACE.API_BASE", None):
litellm.api_base = get_settings().huggingface.api_base
except AttributeError as e:
raise ValueError("OpenAI key is required") from e
if get_settings().get("litellm.use_client"):
litellm_token = get_settings().get("litellm.LITELLM_TOKEN")
assert litellm_token, "LITELLM_TOKEN is required"
os.environ["LITELLM_TOKEN"] = litellm_token
litellm.use_client = True
if get_settings().get("OPENAI.ORG", None):
litellm.organization = get_settings().openai.org
if get_settings().get("OPENAI.API_TYPE", None):
if get_settings().openai.api_type == "azure":
self.azure = True
litellm.azure_key = get_settings().openai.key
if get_settings().get("OPENAI.API_VERSION", None):
litellm.api_version = get_settings().openai.api_version
if get_settings().get("OPENAI.API_BASE", None):
litellm.api_base = get_settings().openai.api_base
if get_settings().get("ANTHROPIC.KEY", None):
litellm.anthropic_key = get_settings().anthropic.key
if get_settings().get("COHERE.KEY", None):
litellm.cohere_key = get_settings().cohere.key
if get_settings().get("REPLICATE.KEY", None):
litellm.replicate_key = get_settings().replicate.key
if get_settings().get("REPLICATE.KEY", None):
litellm.replicate_key = get_settings().replicate.key
if get_settings().get("HUGGINGFACE.KEY", None):
litellm.huggingface_key = get_settings().huggingface.key
if get_settings().get("HUGGINGFACE.API_BASE", None):
litellm.api_base = get_settings().huggingface.api_base
if get_settings().get("VERTEXAI.VERTEX_PROJECT", None):
litellm.vertex_project = get_settings().vertexai.vertex_project
litellm.vertex_location = get_settings().get(
"VERTEXAI.VERTEX_LOCATION", None
)
if get_settings().get("AWS.BEDROCK_REGION", None):
litellm.AmazonAnthropicConfig.max_tokens_to_sample = 2000
self.aws_bedrock_client = boto3.client(
service_name="bedrock-runtime",
region_name=get_settings().aws.bedrock_region,
)
@property
def deployment_id(self):
@ -89,33 +102,37 @@ class AiHandler(BaseAiHandler):
try:
deployment_id = self.deployment_id
if get_settings().config.verbosity_level >= 2:
logging.debug(
get_logger().debug(
f"Generating completion with {model}"
f"{(' from deployment ' + deployment_id) if deployment_id else ''}"
)
response = await acompletion(
model=model,
deployment_id=deployment_id,
messages=[
{"role": "system", "content": system},
{"role": "user", "content": user}
],
temperature=temperature,
azure=self.azure,
force_timeout=get_settings().config.ai_timeout
)
if self.azure:
model = 'azure/' + model
messages = [{"role": "system", "content": system}, {"role": "user", "content": user}]
kwargs = {
"model": model,
"deployment_id": deployment_id,
"messages": messages,
"temperature": temperature,
"force_timeout": get_settings().config.ai_timeout,
}
if self.aws_bedrock_client:
kwargs["aws_bedrock_client"] = self.aws_bedrock_client
response = await acompletion(**kwargs)
except (APIError, Timeout, TryAgain) as e:
logging.error("Error during OpenAI inference: ", e)
get_logger().error("Error during OpenAI inference: ", e)
raise
except (RateLimitError) as e:
logging.error("Rate limit error during OpenAI inference: ", e)
get_logger().error("Rate limit error during OpenAI inference: ", e)
raise
except (Exception) as e:
logging.error("Unknown error during OpenAI inference: ", e)
get_logger().error("Unknown error during OpenAI inference: ", e)
raise TryAgain from e
if response is None or len(response["choices"]) == 0:
raise TryAgain
resp = response["choices"][0]['message']['content']
finish_reason = response["choices"][0]["finish_reason"]
print(resp, finish_reason)
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

36
pr_agent/algo/file_filter.py Normal file
View File

@ -0,0 +1,36 @@
import fnmatch
import re
from pr_agent.config_loader import get_settings
def filter_ignored(files):
"""
Filter out files that match the ignore patterns.
"""
try:
# load regex patterns, and translate glob patterns to regex
patterns = get_settings().ignore.regex
if isinstance(patterns, str):
patterns = [patterns]
glob_setting = get_settings().ignore.glob
if isinstance(glob_setting, str): # --ignore.glob=[.*utils.py], --ignore.glob=.*utils.py
glob_setting = glob_setting.strip('[]').split(",")
patterns += [fnmatch.translate(glob) for glob in glob_setting]
# compile all valid patterns
compiled_patterns = []
for r in patterns:
try:
compiled_patterns.append(re.compile(r))
except re.error:
pass
# keep filenames that _don't_ match the ignore regex
for r in compiled_patterns:
files = [f for f in files if (f.filename and not r.match(f.filename))]
except Exception as e:
print(f"Could not filter file list: {e}")
return files

14
pr_agent/algo/git_patch_processing.py
View File

@ -1,8 +1,10 @@
from __future__ import annotations
import logging
import re
from pr_agent.config_loader import get_settings
from pr_agent.git_providers.git_provider import EDIT_TYPE
from pr_agent.log import get_logger
def extend_patch(original_file_str, patch_str, num_lines) -> str:
@ -63,7 +65,7 @@ def extend_patch(original_file_str, patch_str, num_lines) -> str:
extended_patch_lines.append(line)
except Exception as e:
if get_settings().config.verbosity_level >= 2:
logging.error(f"Failed to extend patch: {e}")
get_logger().error(f"Failed to extend patch: {e}")
return patch_str
# finish previous hunk
@ -114,7 +116,7 @@ def omit_deletion_hunks(patch_lines) -> str:
def handle_patch_deletions(patch: str, original_file_content_str: str,
new_file_content_str: str, file_name: str) -> str:
new_file_content_str: str, file_name: str, edit_type: EDIT_TYPE = EDIT_TYPE.UNKNOWN) -> str:
"""
Handle entire file or deletion patches.
@ -131,17 +133,17 @@ def handle_patch_deletions(patch: str, original_file_content_str: str,
str: The modified patch with deletion hunks omitted.
"""
if not new_file_content_str:
if not new_file_content_str and edit_type != EDIT_TYPE.ADDED:
# logic for handling deleted files - don't show patch, just show that the file was deleted
if get_settings().config.verbosity_level > 0:
logging.info(f"Processing file: {file_name}, minimizing deletion file")
get_logger().info(f"Processing file: {file_name}, minimizing deletion file")
patch = None # file was deleted
else:
patch_lines = patch.splitlines()
patch_new = omit_deletion_hunks(patch_lines)
if patch != patch_new:
if get_settings().config.verbosity_level > 0:
logging.info(f"Processing file: {file_name}, hunks were deleted")
get_logger().info(f"Processing file: {file_name}, hunks were deleted")
patch = patch_new
return patch

5
pr_agent/algo/language_handler.py
View File

@ -3,8 +3,7 @@ from typing import Dict
from pr_agent.config_loader import get_settings
language_extension_map_org = get_settings().language_extension_map_org
language_extension_map = {k.lower(): v for k, v in language_extension_map_org.items()}
# Bad Extensions, source: https://github.com/EleutherAI/github-downloader/blob/345e7c4cbb9e0dc8a0615fd995a08bf9d73b3fe6/download_repo_text.py # noqa: E501
bad_extensions = get_settings().bad_extensions.default
@ -29,6 +28,8 @@ def sort_files_by_main_languages(languages: Dict, files: list):
# languages_sorted = sorted(languages, key=lambda x: x[1], reverse=True)
# get all extensions for the languages
main_extensions = []
language_extension_map_org = get_settings().language_extension_map_org
language_extension_map = {k.lower(): v for k, v in language_extension_map_org.items()}
for language in languages_sorted_list:
if language.lower() in language_extension_map:
main_extensions.append(language_extension_map[language.lower()])

120
pr_agent/algo/pr_processing.py
View File

@ -1,27 +1,29 @@
from __future__ import annotations
import difflib
import logging
import re
import traceback
from typing import Any, Callable, List, Tuple
from github import RateLimitExceededException
from pr_agent.algo import MAX_TOKENS
from pr_agent.algo.git_patch_processing import convert_to_hunks_with_lines_numbers, extend_patch, handle_patch_deletions
from pr_agent.algo.language_handler import sort_files_by_main_languages
from pr_agent.algo.token_handler import TokenHandler, get_token_encoder
from pr_agent.algo.file_filter import filter_ignored
from pr_agent.algo.token_handler import TokenHandler
from pr_agent.algo.utils import get_max_tokens
from pr_agent.config_loader import get_settings
from pr_agent.git_providers.git_provider import FilePatchInfo, GitProvider
from pr_agent.git_providers.git_provider import FilePatchInfo, GitProvider, EDIT_TYPE
from pr_agent.log import get_logger
DELETED_FILES_ = "Deleted files:\n"
MORE_MODIFIED_FILES_ = "More modified files:\n"
MORE_MODIFIED_FILES_ = "Additional modified files (insufficient token budget to process):\n"
ADDED_FILES_ = "Additional added files (insufficient token budget to process):\n"
OUTPUT_BUFFER_TOKENS_SOFT_THRESHOLD = 1000
OUTPUT_BUFFER_TOKENS_HARD_THRESHOLD = 600
PATCH_EXTRA_LINES = 3
def get_pr_diff(git_provider: GitProvider, token_handler: TokenHandler, model: str,
add_line_numbers_to_hunks: bool = False, disable_extra_lines: bool = False) -> str:
@ -44,31 +46,37 @@ def get_pr_diff(git_provider: GitProvider, token_handler: TokenHandler, model: s
"""
if disable_extra_lines:
global PATCH_EXTRA_LINES
PATCH_EXTRA_LINES = 0
else:
PATCH_EXTRA_LINES = get_settings().config.patch_extra_lines
try:
diff_files = git_provider.get_diff_files()
except RateLimitExceededException as e:
logging.error(f"Rate limit exceeded for git provider API. original message {e}")
get_logger().error(f"Rate limit exceeded for git provider API. original message {e}")
raise
diff_files = filter_ignored(diff_files)
# get pr languages
pr_languages = sort_files_by_main_languages(git_provider.get_languages(), diff_files)
# generate a standard diff string, with patch extension
patches_extended, total_tokens, patches_extended_tokens = pr_generate_extended_diff(pr_languages, token_handler,
add_line_numbers_to_hunks)
patches_extended, total_tokens, patches_extended_tokens = pr_generate_extended_diff(
pr_languages, token_handler, add_line_numbers_to_hunks, patch_extra_lines=PATCH_EXTRA_LINES)
# if we are under the limit, return the full diff
if total_tokens + OUTPUT_BUFFER_TOKENS_SOFT_THRESHOLD < MAX_TOKENS[model]:
if total_tokens + OUTPUT_BUFFER_TOKENS_SOFT_THRESHOLD < get_max_tokens(model):
return "\n".join(patches_extended)
# if we are over the limit, start pruning
patches_compressed, modified_file_names, deleted_file_names = \
patches_compressed, modified_file_names, deleted_file_names, added_file_names = \
pr_generate_compressed_diff(pr_languages, token_handler, model, add_line_numbers_to_hunks)
final_diff = "\n".join(patches_compressed)
if added_file_names:
added_list_str = ADDED_FILES_ + "\n".join(added_file_names)
final_diff = final_diff + "\n\n" + added_list_str
if modified_file_names:
modified_list_str = MORE_MODIFIED_FILES_ + "\n".join(modified_file_names)
final_diff = final_diff + "\n\n" + modified_list_str
@ -80,7 +88,8 @@ def get_pr_diff(git_provider: GitProvider, token_handler: TokenHandler, model: s
def pr_generate_extended_diff(pr_languages: list,
token_handler: TokenHandler,
add_line_numbers_to_hunks: bool) -> Tuple[list, int, list]:
add_line_numbers_to_hunks: bool,
patch_extra_lines: int = 0) -> Tuple[list, int, list]:
"""
Generate a standard diff string with patch extension, while counting the number of tokens used and applying diff
minimization techniques if needed.
@ -102,7 +111,7 @@ def pr_generate_extended_diff(pr_languages: list,
continue
# extend each patch with extra lines of context
extended_patch = extend_patch(original_file_content_str, patch, num_lines=PATCH_EXTRA_LINES)
extended_patch = extend_patch(original_file_content_str, patch, num_lines=patch_extra_lines)
full_extended_patch = f"\n\n## {file.filename}\n\n{extended_patch}\n"
if add_line_numbers_to_hunks:
@ -118,7 +127,7 @@ def pr_generate_extended_diff(pr_languages: list,
def pr_generate_compressed_diff(top_langs: list, token_handler: TokenHandler, model: str,
convert_hunks_to_line_numbers: bool) -> Tuple[list, list, list]:
convert_hunks_to_line_numbers: bool) -> Tuple[list, list, list, list]:
"""
Generate a compressed diff string for a pull request, using diff minimization techniques to reduce the number of
tokens used.
@ -144,6 +153,7 @@ def pr_generate_compressed_diff(top_langs: list, token_handler: TokenHandler, mo
"""
patches = []
added_files_list = []
modified_files_list = []
deleted_files_list = []
# sort each one of the languages in top_langs by the number of tokens in the diff
@ -161,7 +171,7 @@ def pr_generate_compressed_diff(top_langs: list, token_handler: TokenHandler, mo
# removing delete-only hunks
patch = handle_patch_deletions(patch, original_file_content_str,
new_file_content_str, file.filename)
new_file_content_str, file.filename, file.edit_type)
if patch is None:
if not deleted_files_list:
total_tokens += token_handler.count_tokens(DELETED_FILES_)
@ -175,21 +185,26 @@ def pr_generate_compressed_diff(top_langs: list, token_handler: TokenHandler, mo
new_patch_tokens = token_handler.count_tokens(patch)
# Hard Stop, no more tokens
if total_tokens > MAX_TOKENS[model] - OUTPUT_BUFFER_TOKENS_HARD_THRESHOLD:
logging.warning(f"File was fully skipped, no more tokens: {file.filename}.")
if total_tokens > get_max_tokens(model) - OUTPUT_BUFFER_TOKENS_HARD_THRESHOLD:
get_logger().warning(f"File was fully skipped, no more tokens: {file.filename}.")
continue
# If the patch is too large, just show the file name
if total_tokens + new_patch_tokens > MAX_TOKENS[model] - OUTPUT_BUFFER_TOKENS_SOFT_THRESHOLD:
if total_tokens + new_patch_tokens > get_max_tokens(model) - OUTPUT_BUFFER_TOKENS_SOFT_THRESHOLD:
# Current logic is to skip the patch if it's too large
# TODO: Option for alternative logic to remove hunks from the patch to reduce the number of tokens
# until we meet the requirements
if get_settings().config.verbosity_level >= 2:
logging.warning(f"Patch too large, minimizing it, {file.filename}")
if not modified_files_list:
total_tokens += token_handler.count_tokens(MORE_MODIFIED_FILES_)
modified_files_list.append(file.filename)
total_tokens += token_handler.count_tokens(file.filename) + 1
get_logger().warning(f"Patch too large, minimizing it, {file.filename}")
if file.edit_type == EDIT_TYPE.ADDED:
if not added_files_list:
total_tokens += token_handler.count_tokens(ADDED_FILES_)
added_files_list.append(file.filename)
else:
if not modified_files_list:
total_tokens += token_handler.count_tokens(MORE_MODIFIED_FILES_)
modified_files_list.append(file.filename)
total_tokens += token_handler.count_tokens(file.filename) + 1
continue
if patch:
@ -200,9 +215,9 @@ def pr_generate_compressed_diff(top_langs: list, token_handler: TokenHandler, mo
patches.append(patch_final)
total_tokens += token_handler.count_tokens(patch_final)
if get_settings().config.verbosity_level >= 2:
logging.info(f"Tokens: {total_tokens}, last filename: {file.filename}")
get_logger().info(f"Tokens: {total_tokens}, last filename: {file.filename}")
return patches, modified_files_list, deleted_files_list
return patches, modified_files_list, deleted_files_list, added_files_list
async def retry_with_fallback_models(f: Callable):
@ -214,7 +229,7 @@ async def retry_with_fallback_models(f: Callable):
get_settings().set("openai.deployment_id", deployment_id)
return await f(model)
except Exception as e:
logging.warning(
get_logger().warning(
f"Failed to generate prediction with {model}"
f"{(' from deployment ' + deployment_id) if deployment_id else ''}: "
f"{traceback.format_exc()}"
@ -267,7 +282,7 @@ def find_line_number_of_relevant_line_in_file(diff_files: List[FilePatchInfo],
r"^@@ -(\d+)(?:,(\d+))? \+(\d+)(?:,(\d+))? @@[ ]?(.*)")
for file in diff_files:
if file.filename.strip() == relevant_file:
if file.filename and (file.filename.strip() == relevant_file):
patch = file.patch
patch_lines = patch.splitlines()
@ -311,35 +326,6 @@ def find_line_number_of_relevant_line_in_file(diff_files: List[FilePatchInfo],
return position, absolute_position
def clip_tokens(text: str, max_tokens: int) -> str:
"""
Clip the number of tokens in a string to a maximum number of tokens.
Args:
text (str): The string to clip.
max_tokens (int): The maximum number of tokens allowed in the string.
Returns:
str: The clipped string.
"""
if not text:
return text
try:
encoder = get_token_encoder()
num_input_tokens = len(encoder.encode(text))
if num_input_tokens <= max_tokens:
return text
num_chars = len(text)
chars_per_token = num_chars / num_input_tokens
num_output_chars = int(chars_per_token * max_tokens)
clipped_text = text[:num_output_chars]
return clipped_text
except Exception as e:
logging.warning(f"Failed to clip tokens: {e}")
return text
def get_pr_multi_diffs(git_provider: GitProvider,
token_handler: TokenHandler,
model: str,
@ -347,25 +333,27 @@ def get_pr_multi_diffs(git_provider: GitProvider,
"""
Retrieves the diff files from a Git provider, sorts them by main language, and generates patches for each file.
The patches are split into multiple groups based on the maximum number of tokens allowed for the given model.
Args:
git_provider (GitProvider): An object that provides access to Git provider APIs.
token_handler (TokenHandler): An object that handles tokens in the context of a pull request.
model (str): The name of the model.
max_calls (int, optional): The maximum number of calls to retrieve diff files. Defaults to 5.
Returns:
List[str]: A list of final diff strings, split into multiple groups based on the maximum number of tokens allowed for the given model.
Raises:
RateLimitExceededException: If the rate limit for the Git provider API is exceeded.
"""
try:
diff_files = git_provider.get_diff_files()
except RateLimitExceededException as e:
logging.error(f"Rate limit exceeded for git provider API. original message {e}")
get_logger().error(f"Rate limit exceeded for git provider API. original message {e}")
raise
diff_files = filter_ignored(diff_files)
# Sort files by main language
pr_languages = sort_files_by_main_languages(git_provider.get_languages(), diff_files)
@ -381,7 +369,7 @@ def get_pr_multi_diffs(git_provider: GitProvider,
for file in sorted_files:
if call_number > max_calls:
if get_settings().config.verbosity_level >= 2:
logging.info(f"Reached max calls ({max_calls})")
get_logger().info(f"Reached max calls ({max_calls})")
break
original_file_content_str = file.base_file
@ -391,26 +379,26 @@ def get_pr_multi_diffs(git_provider: GitProvider,
continue
# Remove delete-only hunks
patch = handle_patch_deletions(patch, original_file_content_str, new_file_content_str, file.filename)
patch = handle_patch_deletions(patch, original_file_content_str, new_file_content_str, file.filename, file.edit_type)
if patch is None:
continue
patch = convert_to_hunks_with_lines_numbers(patch, file)
new_patch_tokens = token_handler.count_tokens(patch)
if patch and (total_tokens + new_patch_tokens > MAX_TOKENS[model] - OUTPUT_BUFFER_TOKENS_SOFT_THRESHOLD):
if patch and (total_tokens + new_patch_tokens > get_max_tokens(model) - OUTPUT_BUFFER_TOKENS_SOFT_THRESHOLD):
final_diff = "\n".join(patches)
final_diff_list.append(final_diff)
patches = []
total_tokens = token_handler.prompt_tokens
call_number += 1
if get_settings().config.verbosity_level >= 2:
logging.info(f"Call number: {call_number}")
get_logger().info(f"Call number: {call_number}")
if patch:
patches.append(patch)
total_tokens += new_patch_tokens
if get_settings().config.verbosity_level >= 2:
logging.info(f"Tokens: {total_tokens}, last filename: {file.filename}")
get_logger().info(f"Tokens: {total_tokens}, last filename: {file.filename}")
# Add the last chunk
if patches:

168
pr_agent/algo/utils.py
View File

@ -2,7 +2,6 @@ from __future__ import annotations
import difflib
import json
import logging
import re
import textwrap
from datetime import datetime
@ -10,7 +9,11 @@ from typing import Any, List
import yaml
from starlette_context import context
from pr_agent.algo import MAX_TOKENS
from pr_agent.algo.token_handler import get_token_encoder
from pr_agent.config_loader import get_settings, global_settings
from pr_agent.log import get_logger
def get_setting(key: str) -> Any:
@ -55,14 +58,15 @@ def convert_to_markdown(output_data: dict, gfm_supported: bool=True) -> str:
emoji = emojis.get(key, "")
if key.lower() == 'code feedback':
if gfm_supported:
markdown_text += f"\n\n- **<details><summary> { emoji } Code feedback:**</summary>\n\n"
markdown_text += f"\n\n- "
markdown_text += f"<details><summary> { emoji } Code feedback:</summary>\n\n"
else:
markdown_text += f"\n\n- **{emoji} Code feedback:**\n\n"
else:
markdown_text += f"- {emoji} **{key}:**\n\n"
for item in value:
if isinstance(item, dict) and key.lower() == 'code feedback':
markdown_text += parse_code_suggestion(item)
markdown_text += parse_code_suggestion(item, gfm_supported)
elif item:
markdown_text += f" - {item}\n"
if key.lower() == 'code feedback':
@ -76,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) -> str:
def parse_code_suggestion(code_suggestions: dict, gfm_supported: bool=True) -> str:
"""
Convert a dictionary of data into markdown format.
@ -96,9 +100,13 @@ def parse_code_suggestion(code_suggestions: dict) -> 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"
markdown_text += f"\n - **{sub_key}:** {sub_value} \n"
else:
markdown_text += f" **{sub_key}:** {sub_value}\n"
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"
return markdown_text
@ -156,7 +164,7 @@ def try_fix_json(review, max_iter=10, code_suggestions=False):
iter_count += 1
if not valid_json:
logging.error("Unable to decode JSON response from AI")
get_logger().error("Unable to decode JSON response from AI")
data = {}
return data
@ -227,7 +235,7 @@ def load_large_diff(filename, new_file_content_str: str, original_file_content_s
diff = difflib.unified_diff(original_file_content_str.splitlines(keepends=True),
new_file_content_str.splitlines(keepends=True))
if get_settings().config.verbosity_level >= 2:
logging.warning(f"File was modified, but no patch was found. Manually creating patch: {filename}.")
get_logger().warning(f"File was modified, but no patch was found. Manually creating patch: {filename}.")
patch = ''.join(diff)
except Exception:
pass
@ -259,12 +267,12 @@ def update_settings_from_args(args: List[str]) -> List[str]:
vals = arg.split('=', 1)
if len(vals) != 2:
if len(vals) > 2: # --extended is a valid argument
logging.error(f'Invalid argument format: {arg}')
get_logger().error(f'Invalid argument format: {arg}')
other_args.append(arg)
continue
key, value = _fix_key_value(*vals)
get_settings().set(key, value)
logging.info(f'Updated setting {key} to: "{value}"')
get_logger().info(f'Updated setting {key} to: "{value}"')
else:
other_args.append(arg)
return other_args
@ -276,28 +284,142 @@ def _fix_key_value(key: str, value: str):
try:
value = yaml.safe_load(value)
except Exception as e:
logging.error(f"Failed to parse YAML for config override {key}={value}", exc_info=e)
get_logger().debug(f"Failed to parse YAML for config override {key}={value}", exc_info=e)
return key, value
def load_yaml(review_text: str) -> dict:
review_text = review_text.removeprefix('```yaml').rstrip('`')
def load_yaml(response_text: str) -> dict:
response_text = response_text.removeprefix('```yaml').rstrip('`')
try:
data = yaml.safe_load(review_text)
data = yaml.safe_load(response_text)
except Exception as e:
logging.error(f"Failed to parse AI prediction: {e}")
data = try_fix_yaml(review_text)
get_logger().error(f"Failed to parse AI prediction: {e}")
data = try_fix_yaml(response_text)
return data
def try_fix_yaml(review_text: str) -> dict:
review_text_lines = review_text.split('\n')
def try_fix_yaml(response_text: str) -> dict:
response_text_lines = response_text.split('\n')
keys = ['relevant line:', 'suggestion content:', 'relevant file:']
# first fallback - try to convert 'relevant line: ...' to relevant line: |-\n ...'
response_text_lines_copy = response_text_lines.copy()
for i in range(0, len(response_text_lines_copy)):
for key in keys:
if key in response_text_lines_copy[i] and not '|-' in response_text_lines_copy[i]:
response_text_lines_copy[i] = response_text_lines_copy[i].replace(f'{key}',
f'{key} |-\n ')
try:
data = yaml.safe_load('\n'.join(response_text_lines_copy))
get_logger().info(f"Successfully parsed AI prediction after adding |-\n")
return data
except:
get_logger().info(f"Failed to parse AI prediction after adding |-\n")
# second fallback - try to remove last lines
data = {}
for i in range(1, len(review_text_lines)):
review_text_lines_tmp = '\n'.join(review_text_lines[:-i])
for i in range(1, len(response_text_lines)):
response_text_lines_tmp = '\n'.join(response_text_lines[:-i])
try:
data = yaml.load(review_text_lines_tmp, Loader=yaml.SafeLoader)
logging.info(f"Successfully parsed AI prediction after removing {i} lines")
data = yaml.safe_load(response_text_lines_tmp,)
get_logger().info(f"Successfully parsed AI prediction after removing {i} lines")
break
except:
pass
return data
# thrid fallback - try to remove leading and trailing curly brackets
response_text_copy = response_text.strip().rstrip().removeprefix('{').removesuffix('}')
try:
data = yaml.safe_load(response_text_copy,)
get_logger().info(f"Successfully parsed AI prediction after removing curly brackets")
return data
except:
pass
def set_custom_labels(variables):
if not get_settings().config.enable_custom_labels:
return
labels = get_settings().custom_labels
if not labels:
# set default labels
labels = ['Bug fix', 'Tests', 'Bug fix with tests', 'Enhancement', 'Documentation', 'Other']
labels_list = "\n - ".join(labels) if labels else ""
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]}"
variables["custom_labels_class"] = "class Label(str, Enum):"
for k, v in labels.items():
description = v['description'].strip('\n').replace('\n', '\\n')
variables["custom_labels_class"] += f"\n {k.lower().replace(' ', '_')} = '{k}' # {description}"
def get_user_labels(current_labels: List[str] = None):
"""
Only keep labels that has been added by the user
"""
try:
if current_labels is None:
current_labels = []
user_labels = []
for label in current_labels:
if label.lower() in ['bug fix', 'tests', 'enhancement', 'documentation', 'other']:
continue
if get_settings().config.enable_custom_labels:
if label in get_settings().custom_labels:
continue
user_labels.append(label)
if user_labels:
get_logger().info(f"Keeping user labels: {user_labels}")
except Exception as e:
get_logger().exception(f"Failed to get user labels: {e}")
return current_labels
return user_labels
def get_max_tokens(model):
settings = get_settings()
if model in MAX_TOKENS:
max_tokens_model = MAX_TOKENS[model]
else:
raise Exception(f"MAX_TOKENS must be set for model {model} in ./pr_agent/algo/__init__.py")
if settings.config.max_model_tokens:
max_tokens_model = min(settings.config.max_model_tokens, max_tokens_model)
# get_logger().debug(f"limiting max tokens to {max_tokens_model}")
return max_tokens_model
def clip_tokens(text: str, max_tokens: int, add_three_dots=True) -> str:
"""
Clip the number of tokens in a string to a maximum number of tokens.
Args:
text (str): The string to clip.
max_tokens (int): The maximum number of tokens allowed in the string.
add_three_dots (bool, optional): A boolean indicating whether to add three dots at the end of the clipped
Returns:
str: The clipped string.
"""
if not text:
return text
try:
encoder = get_token_encoder()
num_input_tokens = len(encoder.encode(text))
if num_input_tokens <= max_tokens:
return text
num_chars = len(text)
chars_per_token = num_chars / num_input_tokens
num_output_chars = int(chars_per_token * max_tokens)
clipped_text = text[:num_output_chars]
if add_three_dots:
clipped_text += "...(truncated)"
return clipped_text
except Exception as e:
get_logger().warning(f"Failed to clip tokens: {e}")
return text