Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add implementation details flag for PR descriptions #10

Merged
merged 9 commits into from
May 20, 2024
78 changes: 59 additions & 19 deletions src/ai_funcs.rs
Original file line number Diff line number Diff line change
@@ -1,19 +1,44 @@
use async_openai::types::{
ChatCompletionRequestAssistantMessageArgs, ChatCompletionRequestSystemMessageArgs,
ChatCompletionRequestUserMessageArgs, CreateChatCompletionRequestArgs,
ChatCompletionRequestUserMessageArgs, ChatCompletionResponseStream,
CreateChatCompletionRequestArgs,
};
use futures::StreamExt;
use std::io::{stdout, Write};

const COMPLETION_TOKENS: u16 = 1024;

const SYSTEM_MESSAGE: &'static str = include_str!("system-message.txt");
const PR_SYSTEM_MESSAGE: &'static str = include_str!("pr-system-message.txt");

const MODEL: &'static str = "gpt-4-1106-preview";

async fn print_stream(
stream: &mut ChatCompletionResponseStream,
) -> Result<(), Box<dyn std::error::Error>> {
let mut lock = stdout().lock();
while let Some(result) = stream.next().await {
match result {
Ok(response) => {
for chat_choice in response.choices.iter() {
if let Some(ref content) = chat_choice.delta.content {
if let Err(e) = write!(lock, "{}", content) {
return Err(e.into());
}
}
}
}
Err(err) => return Err(err.into()),
}
stdout().flush()?;
}
Ok(())
}
pub async fn code_review(output: String) -> Result<(), Box<dyn std::error::Error>> {
let client = async_openai::Client::new();
let request = CreateChatCompletionRequestArgs::default()
.max_tokens(COMPLETION_TOKENS)
.model("gpt-4-1106-preview")
.model(MODEL)
.messages([
ChatCompletionRequestSystemMessageArgs::default()
.content(SYSTEM_MESSAGE)
Expand All @@ -36,22 +61,37 @@ pub async fn code_review(output: String) -> Result<(), Box<dyn std::error::Error

let mut stream = client.chat().create_stream(request).await?;

let mut lock = stdout().lock();
while let Some(result) = stream.next().await {
match result {
Ok(response) => {
for chat_choice in response.choices.iter() {
if let Some(ref content) = chat_choice.delta.content {
if let Err(e) = write!(lock, "{}", content) {
return Err(e.into());
}
}
}
}
Err(err) => return Err(err.into()),
}
stdout().flush()?;
}
print_stream(&mut stream).await
}

Ok(())
pub async fn implementation_details(output: String) -> Result<(), Box<dyn std::error::Error>> {
let client = async_openai::Client::new();
let request = CreateChatCompletionRequestArgs::default()
.max_tokens(COMPLETION_TOKENS)
.model(MODEL)
.messages([
ChatCompletionRequestSystemMessageArgs::default()
.content(PR_SYSTEM_MESSAGE)
.build()?
.into(),
ChatCompletionRequestUserMessageArgs::default()
.content(
"Can please provide the implementation details for the following PR changes?",
)
.build()?
.into(),
ChatCompletionRequestAssistantMessageArgs::default()
.content("Sure thing! What are the changes?")
.build()?
.into(),
ChatCompletionRequestUserMessageArgs::default()
.content(output.as_str())
.build()?
.into(),
])
.build()?;

let mut stream = client.chat().create_stream(request).await?;

print_stream(&mut stream).await
}
48 changes: 31 additions & 17 deletions src/git_funcs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,14 @@ pub struct File {
pub status: String,
}

#[derive(Deserialize, Debug)]
pub struct PRInfo {
pub title: String,
pub body: String,
}

pub struct PR {
pub info: PRInfo,
pub files: Vec<File>,
}

Expand Down Expand Up @@ -67,29 +74,39 @@ pub fn get_git_diff_patch() -> Result<String, git2::Error> {
Ok(diff_str)
}

pub async fn get_pr_info(
pub async fn get_pr(
owner: &str,
repo: &str,
pr_number: u32,
) -> Result<PRInfo, Box<dyn std::error::Error>> {
) -> Result<PR, Box<dyn std::error::Error>> {
let client = reqwest::Client::new();
let pr_url = format!(
"https://api.github.com/repos/{}/{}/pulls/{}",
owner, repo, pr_number
);
let files_url = format!(
"https://api.github.com/repos/{}/{}/pulls/{}/files",
owner, repo, pr_number
);

// Try to get the Bearer token from the environment variable
let token = env::var("GH_PR_TOKEN");
let request = client.get(&pr_url).header("User-Agent", "request");
let pr_request = client.get(&pr_url).header("User-Agent", "request");
let files_request = client.get(&files_url).header("User-Agent", "request");

// If the token exists, add the Authorization header
let request = match token {
Ok(token) => request.header(AUTHORIZATION, format!("Bearer {}", token)),
Err(_) => request,
let pr_request = match &token {
Ok(token) => pr_request.header(AUTHORIZATION, format!("Bearer {}", token)),
Err(_) => pr_request,
};
let files_request = match token {
Ok(token) => files_request.header(AUTHORIZATION, format!("Bearer {}", token)),
Err(_) => files_request,
};

let response = request.send().await?;
if response.status() != StatusCode::OK {
let pr_response = pr_request.send().await?;
let response = files_request.send().await?;
if response.status() != StatusCode::OK || pr_response.status() != StatusCode::OK {
let error_message = if response.status() == StatusCode::NOT_FOUND
|| response.status() == StatusCode::UNAUTHORIZED
{
Expand All @@ -103,14 +120,11 @@ pub async fn get_pr_info(

return Err(error_message.into());
}
let pr_info: PRInfo = pr_response.json().await?;
let files: Vec<File> = response.json().await?;

let files_info: Vec<File> = response.json().await?;

let mut files = Vec::new();

for file_info in files_info {
files.push(file_info);
}

Ok(PRInfo { files })
Ok(PR {
info: pr_info,
files,
})
}
92 changes: 60 additions & 32 deletions src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,15 @@ use colored::Colorize;
use std::process;
use utils::Args;

fn exit_msg(message: &str) -> ! {
eprintln!("{}", message);
process::exit(1);
}

fn exit_err(message: &str, err: Box<dyn std::error::Error>) -> ! {
exit_msg(&format!("{} {}", message, err));
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let args = Args::parse();
Expand All @@ -15,53 +24,72 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
println!("Analysing local changes...\n");
let diff = git_funcs::get_git_diff_patch()?;
let prompt = format!("\n{}", diff);
let review_comments = ai_funcs::code_review(prompt).await;
match review_comments {
Ok(_) => return Ok(()),
Err(e) => {
eprintln!("Failed to analyse code: {}", e);
process::exit(1);

if args.details {
match ai_funcs::implementation_details(prompt.clone()).await {
Ok(_) => {
return Ok(());
}
Err(e) => {
exit_err("Failed to analyse code by implementation details", e);
}
}
} else {
match ai_funcs::code_review(prompt).await {
Ok(_) => {
return Ok(());
}
Err(e) => {
exit_err("Failed to analyse code by code review", e);
}
}
}
}

let owner = args.owner.unwrap_or_else(|| {
eprintln!("Owner argument is missing. Run with --help for usage.");
process::exit(1);
});
let owner = args
.owner
.unwrap_or_else(|| exit_msg("Owner argument is missing. Run with --help for usage."));

let repo = args.repo.unwrap_or_else(|| {
eprintln!("Repo argument is missing. Run with --help for usage.");
process::exit(1);
});
let repo = args
.repo
.unwrap_or_else(|| exit_msg("Repo argument is missing. Run with --help for usage."));

let pr_number = args.pr.unwrap_or_else(|| {
eprintln!("PR number argument is missing. Run with --help for usage.");
process::exit(1);
});
let pr_number = args
.pr
.unwrap_or_else(|| exit_msg("PR number argument is missing. Run with --help for usage."));

println!("Analysing PR changes: {}/{} #{}\n", owner, repo, pr_number);
match git_funcs::get_pr_info(&owner, &repo, pr_number).await {
Ok(pr_info) => {
let mut output = String::new();
for file in pr_info.files {
match git_funcs::get_pr(&owner, &repo, pr_number).await {
Ok(pr) => {
let mut output = String::from(format!("# {}\n\n", pr.info.title));
utils::append_with_newline(
&format!("{}\n # Changed Files:\n", pr.info.body),
&mut output,
);
for file in pr.files {
utils::append_with_newline(
&format!("{} -- {}", &file.filename, &file.status),
&mut output,
);
utils::append_with_newline(&file.patch, &mut output);
}
let review_result = ai_funcs::code_review(output).await;
if let Err(e) = review_result {
eprintln!("Failed to analyze code: {}", e);
process::exit(1);

if args.details {
match ai_funcs::implementation_details(output).await {
Ok(_) => {
return Ok(());
}
Err(e) => exit_err("Failed to analyze code", e),
}
} else {
match ai_funcs::code_review(output).await {
Ok(_) => {
return Ok(());
}
Err(e) => exit_err("Failed to analyze code", e),
}
}
}
Err(e) => {
eprintln!("{}", &format!("Failed to get PR info: {}", e).red());
process::exit(1);
}
Err(e) => exit_msg(&format!("Failed to get PR info: {}", e).red()),
}

Ok(())
}
11 changes: 11 additions & 0 deletions src/pr-system-message.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
Your purpose is to provide the implementation details for a PR description on GitHub. You will be provided with the changes from the PR and you should provide a series of brief notes on what has changed. Structure the response as follows:


## Implementation:

// summarise the changes and reasons for any choices from the point of view of the author of the PR

## Security

// list any security considerations if there are any;
// you can skip the security section if there are no significant security considerations.
12 changes: 11 additions & 1 deletion src/system-message.txt
Original file line number Diff line number Diff line change
@@ -1 +1,11 @@
As a code review bot, focus on identifying areas for code quality and clarity enhancement. Alongside noting areas for improvement, remember to applaud truly commendable implementations. Utilize Markdown for clear and professional feedback, using code blocks for suggested changes. For precision, always reference exact filenames and line numbers. Prioritize providing succinct, insightful and actionable feedback aimed at code performance and maintainability. If a file doesn't warrant any feedback, it's acceptable to bypass it. Aim to keep the review concise, complete with a summarizing conclusion offering a holistic view of the proposed changes.
As a code review bot, focus on identifying areas for code quality and clarity enhancement. Alongside noting areas for improvement, remember to applaud truly commendable implementations. Utilize Markdown for clear and professional feedback, using code blocks for suggested changes. For precision, always reference exact filenames and line numbers. Prioritize providing succinct, insightful and actionable feedback aimed at code performance and maintainability. If a file doesn't warrant any feedback, it's acceptable to bypass it. Aim to keep the review concise, complete with a summarizing conclusion offering a holistic view of the proposed changes. You do not need to highlight issues, if there are no significant issues, but you should provide a short summary of the changes regardless. Structure the response as follows:

# Code Review

## Potential Issues // (if any potential issues)

// list issues by filename and line number, providing code-block suggestions if helpful

## Summary

// provide a summary of the changes with any considerations
3 changes: 3 additions & 0 deletions src/utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,9 @@ pub struct Args {
/// Run locally inside git repository, comparing HEAD against main (branch named `main` must exist)
#[arg(short, long, action = clap::ArgAction::SetTrue)]
pub local: bool,
/// Provide implementation details for the PR description
#[arg(short, long, action = clap::ArgAction::SetTrue)]
pub details: bool,
}

pub fn append_with_newline(new_str: &str, buffer: &mut String) {
Expand Down
Loading