codeflash-ai/codeflash-internal
1
0
Fork 0
Code Issues Pull requests 39 Projects Releases Packages Wiki Activity Actions 2.0k

Implement working create-pr endpoint for cfapi along with refactoring to support it and have main call it

Browse source
This commit is contained in:
afik.cohen 2024-01-16 20:17:49 -08:00
parent a62715918a
commit 745d9064e4
10 changed files with 190 additions and 63 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.idea/runConfigurations/bubble_sort.xml
View file

@ -8,7 +8,7 @@
<env name="PYTHONUNBUFFERED" value="1" />
</envs>
<option name="SDK_HOME" value="" />
<option name="SDK_NAME" value="codeflash311" />
<option name="SDK_NAME" value="$USER_HOME$/miniforge3/envs/codeflash311" />
<option name="WORKING_DIRECTORY" value="$PROJECT_DIR$/cli" />
<option name="IS_MODULE_SDK" value="false" />
<option name="ADD_CONTENT_ROOTS" value="true" />

1
.idea/vcs.xml
View file

@ -2,6 +2,5 @@
<project version="4">
<component name="VcsDirectoryMappings">
<mapping directory="" vcs="Git" />
<mapping directory="$PROJECT_DIR$" vcs="Git" />
</component>
</project>

36
cli/codeflash/api/cfapi.py
View file

@ -6,8 +6,8 @@ from requests import Response
from codeflash.code_utils.env_utils import get_codeflash_api_key
from codeflash.github.PrComment import PrComment
CFAPI_BASE_URL = "https://app.codeflash.ai"
# CFAPI_BASE_URL = "http://localhost:3001"
# CFAPI_BASE_URL = "https://app.codeflash.ai"
CFAPI_BASE_URL = "http://localhost:3001"
CFAPI_HEADERS = {"Authorization": f"Bearer {get_codeflash_api_key()}"}
@ -38,6 +38,18 @@ def suggest_changes(
pr_comment: PrComment,
generated_tests: str,
) -> Response:
"""
Suggest changes to a pull request.
Will make a review suggestion when possible;
or create a new dependent pull request with the suggested changes.
:param owner: The owner of the repository.
:param repo: The name of the repository.
:param pr_number: The number of the pull request.
:param file_changes: A dictionary of file changes.
:param pr_comment: The pull request comment object, containing the optimization explanation, best runtime, etc.
:param generated_tests: The generated tests.
:return: The response object.
"""
payload = {
"owner": owner,
"repo": repo,
@ -50,19 +62,31 @@ def suggest_changes(
return response
def create_new_pr(
def create_pr(
owner: str,
repo: str,
baseBranch: str,
file_changes: dict[str, dict[str, str]],
pr_comment: PrComment,
generated_tests: str,
) -> bool:
) -> Response:
"""
Create a pull request, targeting the specified branch. (usually 'main')
:param owner: The owner of the repository.
:param repo: The name of the repository.
:param targetBranch: The branch to target.
:param file_changes: A dictionary of file changes.
:param pr_comment: The pull request comment object, containing the optimization explanation, best runtime, etc.
:param generated_tests: The generated tests.
:return: The response object.
"""
payload = {
"owner": owner,
"repo": repo,
"baseBranch": baseBranch,
"diffContents": file_changes,
"prCommentFields": pr_comment.to_json(),
"generatedTests": generated_tests,
}
response = make_cfapi_request(endpoint="/open-new-pr", method="POST", payload=payload)
return response.ok
response = make_cfapi_request(endpoint="/create-pr", method="POST", payload=payload)
return response

6
cli/codeflash/github/PrComment.py
View file

@ -10,9 +10,9 @@ from codeflash.verification.test_results import TestResults
class PrComment:
optimization_explanation: str
best_runtime: int
function_name: str
original_runtime: int
file_path: str
function_name: str
relative_file_path: str
speedup: float
winning_test_results: TestResults
@ -22,7 +22,7 @@ class PrComment:
"best_runtime": f"{(self.best_runtime / 1000):.2f}",
"original_runtime": f"{(self.original_runtime / 1000):.2f}",
"function_name": self.function_name,
"file_path": self.file_path,
"file_path": self.relative_file_path,
"speedup_x": f"{self.speedup:.2f}",
"speedup_pct": f"{self.speedup * 100:.2f}",
"report_table": {

32
cli/codeflash/main.py
View file

@ -609,14 +609,23 @@ class Optimizer:
),
generated_tests=generated_original_test_source,
)
if response.ok:
logging.info("OK")
else:
logging.error(
f"Optimization was successful, but I failed to suggest changes to PR #{pr}."
f" Response from server was: {response.text}"
)
elif self.args.all:
logging.info("Creating a new PR with the optimized code...")
owner, repo = get_repo_owner_and_name()
response = cfapi.create_new_pr(
relative_path = os.path.relpath(path, self.args.root)
response = cfapi.create_pr(
owner=owner,
repo=repo,
baseBranch="main",
file_changes={
path: FileDiffContent(
relative_path: FileDiffContent(
oldContent=original_code, newContent=new_code
).model_dump(mode="json")
},
@ -625,20 +634,19 @@ class Optimizer:
best_runtime=best_runtime,
original_runtime=original_runtime,
function_name=function_name,
file_path=path,
relative_file_path=relative_path,
speedup=speedup,
winning_test_results=winning_test_results,
),
generated_tests=generated_test_source,
generated_tests=generated_original_test_source,
)
if response.ok:
logging.info("OK")
else:
logging.error(
f"Optimization was successful, but I failed to suggest changes to PR #{pr}."
f" Response from server was: {response.text}"
)
if response.ok:
logging.info("OK")
else:
logging.error(
f"Optimization was successful, but I failed to create a PR with the optimized code."
f" Response from server was: {response.text}"
)
else:
# Delete it here to not cause a lot of clutter if we are optimizing with --all option
if os.path.exists(generated_tests_path):

4
cli/codeflash/version.py
View file

@ -1,3 +1,3 @@
# These version placeholders will be replaced by poetry-dynamic-versioning during `poetry build`.
__version__ = "0.1.0"
__version_tuple__ = (0, 1, 0)
__version__ = "0.1.1"
__version_tuple__ = (0, 1, 1)

111
js/cf-api/github/create-pr-from-diffcontents.ts
View file

@ -58,17 +58,15 @@ ${generatedTests}
return PR_TEST_REPORT_TEMPLATE.replace(/\{report_table}/g, reportTableMd)
}
export async function createPullRequestFromDiffContents(
export async function createNewBranchFromDiffContents(
installationOctokit: Octokit,
owner: string,
repo: string,
pullNumber: number,
newBranchName: string,
diffContentsMap: Map<string, FileDiffContent>,
prCommentFields: PrCommentFields,
generatedTests: string,
): Promise<Octokit.Response<Octokit.PullsCreateResponse>> {
commitMessage: string,
): Promise<boolean> {
// Create a new branch
const newBranchName = `codeflash-optimizations-for-pr${pullNumber}-${Date.now()}`
await installationOctokit.rest.git.createRef({
owner,
repo,
@ -129,60 +127,123 @@ export async function createPullRequestFromDiffContents(
const newCommitData = await installationOctokit.rest.git.createCommit({
owner,
repo,
message: `Code Optimizations for PR #${pullNumber}`,
message: commitMessage,
tree: newTreeData.data.sha,
parents: [currentCommitSha],
})
// Update the new branch reference to point to the new commit
await installationOctokit.rest.git.updateRef({
const result = await installationOctokit.rest.git.updateRef({
owner,
repo,
ref: `heads/${newBranchName}`,
sha: newCommitData.data.sha,
})
return result.status === 200
}
async function createNewPullRequest(
installationOctokit: Octokit,
owner: string,
repo: string,
title: string,
body: string,
newBranchName: string,
baseBranch: string,
): Promise<Octokit.Response<Octokit.PullsCreateResponse>> {
return await installationOctokit.rest.pulls.create({
owner,
repo,
title,
head: newBranchName,
base: baseBranch,
body,
})
}
export async function createStandalonePullRequest(
installationOctokit: Octokit,
owner: string,
repo: string,
newBranchName: string,
baseBranch: string,
prCommentFields: PrCommentFields,
generatedTests: string,
): Promise<Octokit.Response<Octokit.PullsCreateResponse>> {
// Open a new standalone CodeFlash PR (that doesn't reference an original PR, likely just to main)
const prCommentHeader = buildResultHeader(prCommentFields)
// Open a new PR from the new branch onto the original PR's head branch
const prCommentBody = buildResultDetails(prCommentFields)
const prCommentTestReport = buildResultTestReport(prCommentFields, generatedTests)
const title = `CodeFlash Optimizations for ${prCommentFields.function_name}() in ${prCommentFields.file_path} ⚡️`
const body = `${prCommentHeader}\n` + `${prCommentBody}\n` + `${prCommentTestReport}`
return await createNewPullRequest(
installationOctokit,
owner,
repo,
title,
body,
newBranchName,
baseBranch,
)
}
export async function createDependentPullRequest(
installationOctokit: Octokit,
owner: string,
repo: string,
origPrNumber: number,
newBranchName: string,
prCommentFields: PrCommentFields,
generatedTests: string,
): Promise<Octokit.Response<Octokit.PullsCreateResponse>> {
// Get the head branch of the original pull request
const originalPrData = await installationOctokit.rest.pulls.get({
owner,
repo,
pull_number: pullNumber,
pull_number: origPrNumber,
})
const targetBranch = originalPrData.data.head.ref
const baseBranch = originalPrData.data.head.ref
const prCommentHeader = buildResultHeader(prCommentFields)
// Open a new PR from the new branch onto the original PR's head branch
const prCommentBody = buildResultDetails(prCommentFields)
const prCommentTestReport = buildResultTestReport(prCommentFields, generatedTests)
const newPrData = await installationOctokit.rest.pulls.create({
const title = `CodeFlash Optimizations for PR #${origPrNumber} ⚡️`
const body =
`## ⚡️ This pull request contains optimizations for PR #${origPrNumber}
If you approve this dependent PR, these changes will be merged into the original PR branch \`${baseBranch}\`.
>This PR will be automatically closed if the original PR is merged.\n` +
`----\n` +
`${prCommentHeader}\n` +
`${prCommentBody}\n` +
`${prCommentTestReport}`
const newPrData = await createNewPullRequest(
installationOctokit,
owner,
repo,
title: `CodeFlash Optimizations for PR #${pullNumber} ⚡️`,
head: newBranchName,
base: targetBranch,
body:
`## ⚡️ This pull request contains optimizations for PR #${pullNumber}
If you approve this dependent PR, these changes will be merged into the original PR branch \`${targetBranch}\`.
>This PR will be automatically closed if the original PR is merged.\n` +
`----\n` +
`${prCommentHeader}\n` +
`${prCommentBody}\n` +
`${prCommentTestReport}`,
})
title,
body,
newBranchName,
baseBranch,
)
// Make a comment on the original PR with a link to the new PR
const result = await installationOctokit.rest.issues.createComment({
owner,
repo,
issue_number: pullNumber,
issue_number: origPrNumber,
body: `## ⚡️ CodeFlash found optimizations for this PR
${prCommentHeader}
#### I created a new dependent PR with the suggested changes. Please review:
- ### #${newPrData.data.number}
If you approve, it will be merged into this PR (branch \`${targetBranch}\`).`,
If you approve, it will be merged into this PR (branch \`${baseBranch}\`).`,
})
return newPrData

2
js/cf-api/github/optimization_results_header.md
View file

@ -1,4 +1,4 @@
### 📄 `{file_path}` : `{function_name}()`
### 📄 `{function_name}()` in `{file_path}`
📈 Performance went up by **`{speedup_pct}%`** (**`{speedup_x}x`**)

56
js/cf-api/index.ts
View file

@ -8,7 +8,9 @@ import {
buildResultDetails,
buildResultHeader,
buildResultTestReport,
createPullRequestFromDiffContents,
createDependentPullRequest,
createNewBranchFromDiffContents,
createStandalonePullRequest,
} from "./github/create-pr-from-diffcontents"
import { ghAppMiddleware, ghAppPathPrefix, githubApp } from "./github/github-app"
import { userNickname } from "./auth0-mgmt"
@ -168,12 +170,26 @@ appExpress.post("/cfapi/suggest-pr-changes", async (req, res) => {
)
console.log(`Making a new dependent PR...`)
const newPrData = await createPullRequestFromDiffContents(
const newBranchName = `codeflash-optimizations-for-pr${pullNumber}-${Date.now()}`
const branchCreated = await createNewBranchFromDiffContents(
installationOctokit,
owner,
repo,
newBranchName,
diffContentsMap,
`Code Optimizations for PR #${pullNumber}`,
)
if (!branchCreated) {
throw new Error(`Failed to create branch ${newBranchName}`)
}
const newPrData = await createDependentPullRequest(
installationOctokit,
owner,
repo,
pullNumber,
diffContentsMap,
newBranchName,
prCommentFields,
generatedTests,
)
@ -232,11 +248,15 @@ appExpress.post("/cfapi/suggest-pr-changes", async (req, res) => {
}
})
appExpress.post("/cfapi/open-new-pr", async (req, res) => {
appExpress.post("/cfapi/create-pr", async (req, res) => {
try {
const { owner, repo, diffContents, prCommentFields, generatedTests } = req.body
const { owner, repo, baseBranch, diffContents, prCommentFields, generatedTests } = req.body
const userId = req.userId
if (!repo || !owner || !baseBranch || !isDiffContentsWellFormed(diffContents)) {
return res.status(400).send("Missing or malformed fields")
}
const nickname: string | null = await userNickname(userId)
if (nickname == null) {
return res.status(401).send("Unauthorized") // Error getting user nickname
@ -270,24 +290,36 @@ appExpress.post("/cfapi/open-new-pr", async (req, res) => {
const diffContentsMap: Map<string, FileDiffContent> = fileDiffsToMap(diffContents)
const pullNumber: number = 0
const newPrData = await createPullRequestFromDiffContents(
const newBranchName = `codeflash-optimizations-for-function-${
prCommentFields.function_name
}-${Date.now()}`
const branchCreated = await createNewBranchFromDiffContents(
installationOctokit,
owner,
repo,
pullNumber,
newBranchName,
diffContentsMap,
`Code Optimizations for ${prCommentFields.function_name}`,
)
if (!branchCreated) {
throw new Error(`Failed to create branch ${newBranchName}`)
}
const newPrData = await createStandalonePullRequest(
installationOctokit,
owner,
repo,
newBranchName,
baseBranch,
prCommentFields,
generatedTests,
)
// Respond with the new PR details
console.log(
`Created new dependent PR #${newPrData.data.number} from branch ${newPrData.data.head.ref}`,
)
console.log(`Created new PR #${newPrData.data.number} with branch ${newPrData.data.head.ref}`)
res.json(newPrData.data.number)
} catch (error) {
console.log(`Error in /cfapi/open-new-pr: ${error}`)
console.log(`Error in /cfapi/create-pr: ${error}`)
console.log(`Error message: ${error.message}`)
console.log(`Error stack: ${error.stack}`)
res.status(500).send(`Error creating pull request: ${error.message}`)

3
js/cf-webapp/src/app/app/install/install.md
View file

@ -18,6 +18,9 @@ repository you want to install CodeFlash on.
pip install codeflash
```
> ℹ️ Don't pin or use a specific version of codeflash - we're constantly improving it, and frequent API changes are
> expected!
### 2. Set up your API key
Grab one from https://app.codeflash.ai/app/apikeys and save it to your environment variables: