diff --git a/docker/Dockerfile b/docker/Dockerfile index 61ab74cf..8d28a9ed 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -9,6 +9,10 @@ FROM base as github_app ADD pr_agent pr_agent CMD ["python", "pr_agent/servers/github_app.py"] +FROM base as bitbucket_app +ADD pr_agent pr_agent +CMD ["python", "pr_agent/servers/bitbucket_app.py"] + FROM base as github_polling ADD pr_agent pr_agent CMD ["python", "pr_agent/servers/github_polling.py"] diff --git a/pr_agent/git_providers/bitbucket_provider.py b/pr_agent/git_providers/bitbucket_provider.py index bee0e351..0cd860fa 100644 --- a/pr_agent/git_providers/bitbucket_provider.py +++ b/pr_agent/git_providers/bitbucket_provider.py @@ -5,16 +5,25 @@ from urllib.parse import urlparse import requests from atlassian.bitbucket import Cloud +from starlette_context import context from ..config_loader import get_settings from .git_provider import FilePatchInfo, GitProvider class BitbucketProvider(GitProvider): - def __init__(self, pr_url: Optional[str] = None, incremental: Optional[bool] = False): + def __init__( + self, pr_url: Optional[str] = None, incremental: Optional[bool] = False + ): s = requests.Session() - s.headers['Authorization'] = f'Bearer {get_settings().get("BITBUCKET.BEARER_TOKEN", None)}' - s.headers['Content-Type'] = 'application/json' + try: + bearer = context.get("bitbucket_bearer_token", None) + s.headers["Authorization"] = f"Bearer {bearer}" + except Exception: + s.headers[ + "Authorization" + ] = f'Bearer {get_settings().get("BITBUCKET.BEARER_TOKEN", None)}' + s.headers["Content-Type"] = "application/json" self.headers = s.headers self.bitbucket_client = Cloud(session=s) self.workspace_slug = None @@ -26,37 +35,44 @@ class BitbucketProvider(GitProvider): self.incremental = incremental if pr_url: self.set_pr(pr_url) - self.bitbucket_comment_api_url = self.pr._BitbucketBase__data['links']['comments']['href'] + self.bitbucket_comment_api_url = self.pr._BitbucketBase__data["links"][ + "comments" + ]["href"] def get_repo_settings(self): try: - contents = self.repo_obj.get_contents(".pr_agent.toml", ref=self.pr.head.sha).decoded_content + contents = self.repo_obj.get_contents( + ".pr_agent.toml", ref=self.pr.head.sha + ).decoded_content return contents except Exception: return "" - + def publish_code_suggestions(self, code_suggestions: list) -> bool: """ Publishes code suggestions as comments on the PR. """ post_parameters_list = [] for suggestion in code_suggestions: - body = suggestion['body'] - relevant_file = suggestion['relevant_file'] - relevant_lines_start = suggestion['relevant_lines_start'] - relevant_lines_end = suggestion['relevant_lines_end'] + body = suggestion["body"] + relevant_file = suggestion["relevant_file"] + relevant_lines_start = suggestion["relevant_lines_start"] + relevant_lines_end = suggestion["relevant_lines_end"] if not relevant_lines_start or relevant_lines_start == -1: if get_settings().config.verbosity_level >= 2: logging.exception( - f"Failed to publish code suggestion, relevant_lines_start is {relevant_lines_start}") + f"Failed to publish code suggestion, relevant_lines_start is {relevant_lines_start}" + ) continue if relevant_lines_end < relevant_lines_start: if get_settings().config.verbosity_level >= 2: - logging.exception(f"Failed to publish code suggestion, " - f"relevant_lines_end is {relevant_lines_end} and " - f"relevant_lines_start is {relevant_lines_start}") + logging.exception( + f"Failed to publish code suggestion, " + f"relevant_lines_end is {relevant_lines_end} and " + f"relevant_lines_start is {relevant_lines_start}" + ) continue if relevant_lines_end > relevant_lines_start: @@ -75,8 +91,7 @@ class BitbucketProvider(GitProvider): "side": "RIGHT", } post_parameters_list.append(post_parameters) - - + try: self.publish_inline_comments(post_parameters_list) return True @@ -86,7 +101,12 @@ class BitbucketProvider(GitProvider): return False def is_supported(self, capability: str) -> bool: - if capability in ['get_issue_comments', 'create_inline_comment', 'publish_inline_comments', 'get_labels']: + if capability in [ + "get_issue_comments", + "create_inline_comment", + "publish_inline_comments", + "get_labels", + ]: return False return True @@ -99,57 +119,65 @@ class BitbucketProvider(GitProvider): def get_diff_files(self) -> list[FilePatchInfo]: diffs = self.pr.diffstat() - diff_split = ['diff --git%s' % x for x in self.pr.diff().split('diff --git') if x.strip()] - + diff_split = [ + "diff --git%s" % x for x in self.pr.diff().split("diff --git") if x.strip() + ] + diff_files = [] for index, diff in enumerate(diffs): - original_file_content_str = self._get_pr_file_content(diff.old.get_data('links')) - new_file_content_str = self._get_pr_file_content(diff.new.get_data('links')) - diff_files.append(FilePatchInfo(original_file_content_str, new_file_content_str, - diff_split[index], diff.new.path)) + original_file_content_str = self._get_pr_file_content( + diff.old.get_data("links") + ) + new_file_content_str = self._get_pr_file_content(diff.new.get_data("links")) + diff_files.append( + FilePatchInfo( + original_file_content_str, + new_file_content_str, + diff_split[index], + diff.new.path, + ) + ) return diff_files def publish_comment(self, pr_comment: str, is_temporary: bool = False): comment = self.pr.comment(pr_comment) if is_temporary: - self.temp_comments.append(comment['id']) + self.temp_comments.append(comment["id"]) def remove_initial_comment(self): try: for comment in self.temp_comments: - self.pr.delete(f'comments/{comment}') + self.pr.delete(f"comments/{comment}") except Exception as e: logging.exception(f"Failed to remove temp comments, error: {e}") - def publish_inline_comment(self, comment: str, from_line: int, to_line: int, file: str): - payload = json.dumps( { - "content": { - "raw": comment, - }, - "inline": { - "to": from_line, - "path": file - }, - }) + def publish_inline_comment( + self, comment: str, from_line: int, to_line: int, file: str + ): + payload = json.dumps( + { + "content": { + "raw": comment, + }, + "inline": {"to": from_line, "path": file}, + } + ) response = requests.request( - "POST", - self.bitbucket_comment_api_url, - data=payload, - headers=self.headers + "POST", self.bitbucket_comment_api_url, data=payload, headers=self.headers ) return response - - def publish_inline_comments(self, comments: list[dict]): for comment in comments: - self.publish_inline_comment(comment['body'], comment['start_line'], comment['line'], comment['path']) + self.publish_inline_comment( + comment["body"], comment["start_line"], comment["line"], comment["path"] + ) def get_title(self): return self.pr.title def get_languages(self): - languages = {self._get_repo().get_data('language'): 0} + languages = {self._get_repo().get_data("language"): 0} return languages def get_pr_branch(self): @@ -162,7 +190,9 @@ class BitbucketProvider(GitProvider): return 0 def get_issue_comments(self): - raise NotImplementedError("Bitbucket provider does not support issue comments yet") + raise NotImplementedError( + "Bitbucket provider does not support issue comments yet" + ) def add_eyes_reaction(self, issue_comment_id: int) -> Optional[int]: return True @@ -173,14 +203,16 @@ class BitbucketProvider(GitProvider): @staticmethod def _parse_pr_url(pr_url: str) -> Tuple[str, int]: parsed_url = urlparse(pr_url) - - if 'bitbucket.org' not in parsed_url.netloc: + + if "bitbucket.org" not in parsed_url.netloc: raise ValueError("The provided URL is not a valid Bitbucket URL") - path_parts = parsed_url.path.strip('/').split('/') - - if len(path_parts) < 4 or path_parts[2] != 'pull-requests': - raise ValueError("The provided URL does not appear to be a Bitbucket PR URL") + path_parts = parsed_url.path.strip("/").split("/") + + if len(path_parts) < 4 or path_parts[2] != "pull-requests": + raise ValueError( + "The provided URL does not appear to be a Bitbucket PR URL" + ) workspace_slug = path_parts[0] repo_slug = path_parts[1] @@ -193,7 +225,9 @@ class BitbucketProvider(GitProvider): def _get_repo(self): if self.repo is None: - self.repo = self.bitbucket_client.workspaces.get(self.workspace_slug).repositories.get(self.repo_slug) + self.repo = self.bitbucket_client.workspaces.get( + self.workspace_slug + ).repositories.get(self.repo_slug) return self.repo def _get_pr(self): @@ -204,3 +238,16 @@ class BitbucketProvider(GitProvider): def get_commit_messages(self): return "" # not implemented yet + + def publish_description(self, pr_title: str, pr_body: str): + pass + def create_inline_comment( + self, body: str, relevant_file: str, relevant_line_in_file: str + ): + pass + + def publish_labels(self, labels): + pass + + def get_labels(self): + pass diff --git a/pr_agent/secret_providers/__init__.py b/pr_agent/secret_providers/__init__.py new file mode 100644 index 00000000..1cc3ea7b --- /dev/null +++ b/pr_agent/secret_providers/__init__.py @@ -0,0 +1,16 @@ +from pr_agent.config_loader import get_settings + + +def get_secret_provider(): + try: + provider_id = get_settings().config.secret_provider + except AttributeError as e: + raise ValueError("secret_provider is a required attribute in the configuration file") from e + try: + if provider_id == 'google_cloud_storage': + from pr_agent.secret_providers.google_cloud_storage_secret_provider import GoogleCloudStorageSecretProvider + return GoogleCloudStorageSecretProvider() + else: + raise ValueError(f"Unknown secret provider: {provider_id}") + except Exception as e: + raise ValueError(f"Failed to initialize secret provider {provider_id}") from e diff --git a/pr_agent/secret_providers/google_cloud_storage_secret_provider.py b/pr_agent/secret_providers/google_cloud_storage_secret_provider.py new file mode 100644 index 00000000..18db5c4b --- /dev/null +++ b/pr_agent/secret_providers/google_cloud_storage_secret_provider.py @@ -0,0 +1,35 @@ +import ujson + +from google.cloud import storage + +from pr_agent.config_loader import get_settings +from pr_agent.git_providers.gitlab_provider import logger +from pr_agent.secret_providers.secret_provider import SecretProvider + + +class GoogleCloudStorageSecretProvider(SecretProvider): + def __init__(self): + try: + self.client = storage.Client.from_service_account_info(ujson.loads(get_settings().google_cloud_storage. + service_account)) + self.bucket_name = get_settings().google_cloud_storage.bucket_name + self.bucket = self.client.bucket(self.bucket_name) + except Exception as e: + logger.error(f"Failed to initialize Google Cloud Storage Secret Provider: {e}") + raise e + + def get_secret(self, secret_name: str) -> str: + try: + blob = self.bucket.blob(secret_name) + return blob.download_as_string() + except Exception as e: + logger.error(f"Failed to get secret {secret_name} from Google Cloud Storage: {e}") + return "" + + def store_secret(self, secret_name: str, secret_value: str): + try: + blob = self.bucket.blob(secret_name) + blob.upload_from_string(secret_value) + except Exception as e: + logger.error(f"Failed to store secret {secret_name} in Google Cloud Storage: {e}") + raise e diff --git a/pr_agent/secret_providers/secret_provider.py b/pr_agent/secret_providers/secret_provider.py new file mode 100644 index 00000000..df1e7780 --- /dev/null +++ b/pr_agent/secret_providers/secret_provider.py @@ -0,0 +1,12 @@ +from abc import ABC, abstractmethod + + +class SecretProvider(ABC): + + @abstractmethod + def get_secret(self, secret_name: str) -> str: + pass + + @abstractmethod + def store_secret(self, secret_name: str, secret_value: str): + pass diff --git a/pr_agent/servers/atlassian-connect.json b/pr_agent/servers/atlassian-connect.json new file mode 100644 index 00000000..1ff50865 --- /dev/null +++ b/pr_agent/servers/atlassian-connect.json @@ -0,0 +1,33 @@ +{ + "name": "CodiumAI PR-Agent", + "description": "CodiumAI PR-Agent", + "key": "app_key", + "vendor": { + "name": "CodiumAI", + "url": "https://codium.ai" + }, + "authentication": { + "type": "jwt" + }, + "baseUrl": "base_url", + "lifecycle": { + "installed": "/installed", + "uninstalled": "/uninstalled" + }, + "scopes": [ + "account", + "repository", + "pullrequest" + ], + "contexts": [ + "account" + ], + "modules": { + "webhooks": [ + { + "event": "*", + "url": "/webhook" + } + ] + } +} \ No newline at end of file diff --git a/pr_agent/servers/bitbucket_app.py b/pr_agent/servers/bitbucket_app.py new file mode 100644 index 00000000..cc6491d4 --- /dev/null +++ b/pr_agent/servers/bitbucket_app.py @@ -0,0 +1,139 @@ +import copy +import hashlib +import json +import logging +import os +import sys +import time + +import jwt +import requests +import uvicorn +from fastapi import APIRouter, FastAPI, Request, Response +from starlette.background import BackgroundTasks +from starlette.middleware import Middleware +from starlette.responses import JSONResponse +from starlette_context import context +from starlette_context.middleware import RawContextMiddleware + +from pr_agent.agent.pr_agent import PRAgent +from pr_agent.config_loader import get_settings, global_settings +from pr_agent.secret_providers import get_secret_provider + +logging.basicConfig(stream=sys.stdout, level=logging.INFO) +router = APIRouter() +secret_provider = get_secret_provider() + +async def get_bearer_token(shared_secret: str, client_key: str): + try: + now = int(time.time()) + url = "https://bitbucket.org/site/oauth2/access_token" + canonical_url = "GET&/site/oauth2/access_token&" + qsh = hashlib.sha256(canonical_url.encode("utf-8")).hexdigest() + app_key = get_settings().bitbucket.app_key + + payload = { + "iss": app_key, + "iat": now, + "exp": now + 240, + "qsh": qsh, + "sub": client_key, + } + token = jwt.encode(payload, shared_secret, algorithm="HS256") + payload = 'grant_type=urn%3Abitbucket%3Aoauth2%3Ajwt' + headers = { + 'Authorization': f'JWT {token}', + 'Content-Type': 'application/x-www-form-urlencoded' + } + response = requests.request("POST", url, headers=headers, data=payload) + bearer_token = response.json()["access_token"] + return bearer_token + except Exception as e: + logging.error(f"Failed to get bearer token: {e}") + raise e + +@router.get("/") +async def handle_manifest(request: Request, response: Response): + cur_dir = os.path.dirname(os.path.abspath(__file__)) + manifest = open(os.path.join(cur_dir, "atlassian-connect.json"), "rt").read() + try: + manifest = manifest.replace("app_key", get_settings().bitbucket.app_key) + manifest = manifest.replace("base_url", get_settings().bitbucket.base_url) + except: + logging.error("Failed to replace api_key in Bitbucket manifest, trying to continue") + manifest_obj = json.loads(manifest) + return JSONResponse(manifest_obj) + +@router.post("/webhook") +async def handle_github_webhooks(background_tasks: BackgroundTasks, request: Request): + print(request.headers) + jwt_header = request.headers.get("authorization", None) + if jwt_header: + input_jwt = jwt_header.split(" ")[1] + data = await request.json() + print(data) + async def inner(): + try: + owner = data["data"]["repository"]["owner"]["username"] + secrets = json.loads(secret_provider.get_secret(owner)) + shared_secret = secrets["shared_secret"] + client_key = secrets["client_key"] + jwt.decode(input_jwt, shared_secret, audience=client_key, algorithms=["HS256"]) + bearer_token = await get_bearer_token(shared_secret, client_key) + context['bitbucket_bearer_token'] = bearer_token + context["settings"] = copy.deepcopy(global_settings) + event = data["event"] + agent = PRAgent() + if event == "pullrequest:created": + pr_url = data["data"]["pullrequest"]["links"]["html"]["href"] + await agent.handle_request(pr_url, "review") + elif event == "pullrequest:comment_created": + pr_url = data["data"]["pullrequest"]["links"]["html"]["href"] + comment_body = data["data"]["comment"]["content"]["raw"] + await agent.handle_request(pr_url, comment_body) + except Exception as e: + logging.error(f"Failed to handle webhook: {e}") + background_tasks.add_task(inner) + return "OK" + +@router.get("/webhook") +async def handle_github_webhooks(request: Request, response: Response): + return "Webhook server online!" + +@router.post("/installed") +async def handle_installed_webhooks(request: Request, response: Response): + try: + print(request.headers) + data = await request.json() + print(data) + shared_secret = data["sharedSecret"] + client_key = data["clientKey"] + username = data["principal"]["username"] + secrets = { + "shared_secret": shared_secret, + "client_key": client_key + } + secret_provider.store_secret(username, json.dumps(secrets)) + except Exception as e: + logging.error(f"Failed to register user: {e}") + return JSONResponse({"error": "Unable to register user"}, status_code=500) + +@router.post("/uninstalled") +async def handle_uninstalled_webhooks(request: Request, response: Response): + data = await request.json() + print(data) + + +def start(): + get_settings().set("CONFIG.PUBLISH_OUTPUT_PROGRESS", False) + get_settings().set("CONFIG.GIT_PROVIDER", "bitbucket") + get_settings().set("PR_DESCRIPTION.PUBLISH_DESCRIPTION_AS_COMMENT", True) + middleware = [Middleware(RawContextMiddleware)] + app = FastAPI(middleware=middleware) + app.include_router(router) + + uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", "3000"))) + + +if __name__ == '__main__': + start() diff --git a/pr_agent/servers/github_app.py b/pr_agent/servers/github_app.py index 7a7208d0..1d1c901c 100644 --- a/pr_agent/servers/github_app.py +++ b/pr_agent/servers/github_app.py @@ -167,4 +167,4 @@ def start(): if __name__ == '__main__': - start() \ No newline at end of file + start() diff --git a/pr_agent/settings/configuration.toml b/pr_agent/settings/configuration.toml index c2350526..f8abd555 100644 --- a/pr_agent/settings/configuration.toml +++ b/pr_agent/settings/configuration.toml @@ -11,6 +11,7 @@ ai_timeout=180 max_description_tokens = 500 max_commits_tokens = 500 litellm_debugger=false +secret_provider="google_cloud_storage" [pr_reviewer] # /review # require_focused_review=false diff --git a/pr_agent/tools/pr_description.py b/pr_agent/tools/pr_description.py index 440675fd..acd272bc 100644 --- a/pr_agent/tools/pr_description.py +++ b/pr_agent/tools/pr_description.py @@ -153,12 +153,6 @@ class PRDescription: # Initialization pr_types = [] - # Iterate over the dictionary items and append the key and value to 'markdown_text' in a markdown format - markdown_text = "" - for key, value in data.items(): - markdown_text += f"## {key}\n\n" - markdown_text += f"{value}\n\n" - # If the 'PR Type' key is present in the dictionary, split its value by comma and assign it to 'pr_types' if 'PR Type' in data: if type(data['PR Type']) == list: @@ -194,6 +188,8 @@ class PRDescription: if idx < len(data) - 1: pr_body += "\n___\n" + markdown_text = f"## Title\n\n{title}\n\n___\n{pr_body}" + if get_settings().config.verbosity_level >= 2: logging.info(f"title:\n{title}\n{pr_body}") diff --git a/pyproject.toml b/pyproject.toml index 9a945dca..8d429668 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,7 +44,9 @@ dependencies = [ "starlette-context==0.3.6", "litellm~=0.1.445", "PyYAML==6.0", - "boto3~=1.28.25" + "boto3~=1.28.25", + "google-cloud-storage==2.10.0", + "ujson==5.8.0" ] [project.urls] diff --git a/requirements.txt b/requirements.txt index 0bbb6f28..fe92a74b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -15,3 +15,5 @@ PyYAML==6.0 starlette-context==0.3.6 litellm~=0.1.445 boto3~=1.28.25 +google-cloud-storage==2.10.0 +ujson==5.8.0