diff --git a/aicodebot/cli.py b/aicodebot/cli.py index 2484017..3d734e4 100644 --- a/aicodebot/cli.py +++ b/aicodebot/cli.py @@ -1,13 +1,8 @@ from aicodebot import version as aicodebot_version -from aicodebot.coder import DEFAULT_MAX_TOKENS, Coder -from aicodebot.commands import alignment, commit, configure, debug, learn, sidekick, sidekick_agent +from aicodebot.commands import alignment, commit, configure, debug, learn, review, sidekick, sidekick_agent from aicodebot.config import read_config -from aicodebot.helpers import logger -from aicodebot.output import OurMarkdown as Markdown, RichLiveCallbackHandler, get_console -from aicodebot.prompts import get_prompt -from langchain.chains import LLMChain -from rich.live import Live -import click, json, langchain, os, sys +from aicodebot.output import get_console +import click, langchain, os, sys # -------------------------- Top level command group ------------------------- # @@ -40,75 +35,12 @@ def cli(ctx, debug_output): cli.add_command(configure) cli.add_command(commit) cli.add_command(debug) +cli.add_command(review) cli.add_command(sidekick) if os.getenv("AICODEBOT_ENABLE_EXPERIMENTAL_FEATURES"): cli.add_command(learn) cli.add_command(sidekick_agent) -@cli.command -@click.option("-c", "--commit", help="The commit hash to review (otherwise look at [un]staged changes).") -@click.option("-v", "--verbose", count=True) -@click.option("--output-format", default="text", type=click.Choice(["text", "json"], case_sensitive=False)) -@click.option("-t", "--response-token-size", type=int, default=DEFAULT_MAX_TOKENS * 2) -@click.argument("files", nargs=-1) -def review(commit, verbose, output_format, response_token_size, files): - """Do a code review, with [un]staged changes, or a specified commit.""" - if not Coder.is_inside_git_repo(): - console.print("🛑 This command must be run from within a git repository.", style=console.error_style) - sys.exit(1) - - # If files are specified, only consider those files - # Otherwise, use git to get the list of files - if not files: - files = Coder.git_staged_files() - if not files: - files = Coder.git_unstaged_files() - - diff_context = Coder.git_diff_context(commit, files) - if not diff_context: - console.print("No changes detected for review. 🤷") - return - languages = ",".join(Coder.identify_languages(files)) - - # Load the prompt - prompt = get_prompt("review", structured_output=output_format == "json") - logger.trace(f"Prompt: {prompt}") - - # Check the size of the diff context and adjust accordingly - request_token_size = Coder.get_token_length(diff_context) + Coder.get_token_length(prompt.template) - model_name = Coder.get_llm_model_name(request_token_size + response_token_size) - if model_name is None: - raise click.ClickException(f"The diff is too large to review ({request_token_size} tokens). 😢") - - llm = Coder.get_llm(model_name, verbose, response_token_size, streaming=True) - chain = LLMChain(llm=llm, prompt=prompt, verbose=verbose) - - if output_format == "json": - with console.status("Examining the diff and generating the review", spinner=console.DEFAULT_SPINNER): - response = chain.run({"diff_context": diff_context, "languages": languages}) - - parsed_response = prompt.output_parser.parse(response) - data = { - "review_status": parsed_response.review_status, - "review_comments": parsed_response.review_comments, - } - if commit: - data["commit"] = commit - json_response = json.dumps(data, indent=4) - print(json_response) # noqa: T201 - - else: - # Stream live - console.print( - "Examining the diff and generating the review for the following files:\n\t" + "\n\t".join(files) - ) - with Live(Markdown(""), auto_refresh=True) as live: - llm.streaming = True - llm.callbacks = [RichLiveCallbackHandler(live, console.bot_style)] - - chain.run({"diff_context": diff_context, "languages": languages}) - - if __name__ == "__main__": # pragma: no cover cli() diff --git a/aicodebot/commands/__init__.py b/aicodebot/commands/__init__.py index 7ae7a16..46e8aab 100644 --- a/aicodebot/commands/__init__.py +++ b/aicodebot/commands/__init__.py @@ -3,6 +3,7 @@ from .configure import configure from .debug import debug from .learn import learn +from .review import review from .sidekick import sidekick, sidekick_agent -__all__ = ["alignment", "commit", "configure", "debug", "learn", "sidekick", "sidekick_agent"] +__all__ = ["alignment", "commit", "configure", "debug", "learn", "review", "sidekick", "sidekick_agent"] diff --git a/aicodebot/commands/review.py b/aicodebot/commands/review.py new file mode 100644 index 0000000..4373846 --- /dev/null +++ b/aicodebot/commands/review.py @@ -0,0 +1,72 @@ +from aicodebot.coder import DEFAULT_MAX_TOKENS, Coder +from aicodebot.helpers import logger +from aicodebot.output import OurMarkdown, RichLiveCallbackHandler, get_console +from aicodebot.prompts import get_prompt +from langchain.chains import LLMChain +from rich.live import Live +import click, json, sys + + +@click.command +@click.option("-c", "--commit", help="The commit hash to review (otherwise look at [un]staged changes).") +@click.option("-v", "--verbose", count=True) +@click.option("--output-format", default="text", type=click.Choice(["text", "json"], case_sensitive=False)) +@click.option("-t", "--response-token-size", type=int, default=DEFAULT_MAX_TOKENS * 2) +@click.argument("files", nargs=-1) +def review(commit, verbose, output_format, response_token_size, files): + """Do a code review, with [un]staged changes, or a specified commit.""" + console = get_console() + if not Coder.is_inside_git_repo(): + console.print("🛑 This command must be run from within a git repository.", style=console.error_style) + sys.exit(1) + + # If files are specified, only consider those files + # Otherwise, use git to get the list of files + if not files: + files = Coder.git_staged_files() + if not files: + files = Coder.git_unstaged_files() + + diff_context = Coder.git_diff_context(commit, files) + if not diff_context: + console.print("No changes detected for review. 🤷") + return + languages = ",".join(Coder.identify_languages(files)) + + # Load the prompt + prompt = get_prompt("review", structured_output=output_format == "json") + logger.trace(f"Prompt: {prompt}") + + # Check the size of the diff context and adjust accordingly + request_token_size = Coder.get_token_length(diff_context) + Coder.get_token_length(prompt.template) + model_name = Coder.get_llm_model_name(request_token_size + response_token_size) + if model_name is None: + raise click.ClickException(f"The diff is too large to review ({request_token_size} tokens). 😢") + + llm = Coder.get_llm(model_name, verbose, response_token_size, streaming=True) + chain = LLMChain(llm=llm, prompt=prompt, verbose=verbose) + + if output_format == "json": + with console.status("Examining the diff and generating the review", spinner=console.DEFAULT_SPINNER): + response = chain.run({"diff_context": diff_context, "languages": languages}) + + parsed_response = prompt.output_parser.parse(response) + data = { + "review_status": parsed_response.review_status, + "review_comments": parsed_response.review_comments, + } + if commit: + data["commit"] = commit + json_response = json.dumps(data, indent=4) + print(json_response) # noqa: T201 + + else: + # Stream live + console.print( + "Examining the diff and generating the review for the following files:\n\t" + "\n\t".join(files) + ) + with Live(OurMarkdown(""), auto_refresh=True) as live: + llm.streaming = True + llm.callbacks = [RichLiveCallbackHandler(live, console.bot_style)] + + chain.run({"diff_context": diff_context, "languages": languages})