diff --git a/docker/Dockerfile b/docker/Dockerfile index 61ab74cf..4336cacc 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -2,13 +2,18 @@ FROM python:3.10 as base WORKDIR /app ADD pyproject.toml . -RUN pip install . && rm pyproject.toml +ADD requirements.txt . +RUN pip install . && rm pyproject.toml requirements.txt ENV PYTHONPATH=/app 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 fbfdffe9..a4036575 100644 --- a/pr_agent/git_providers/bitbucket_provider.py +++ b/pr_agent/git_providers/bitbucket_provider.py @@ -5,6 +5,7 @@ from urllib.parse import urlparse import requests from atlassian.bitbucket import Cloud +from starlette_context import context from ..algo.pr_processing import clip_tokens, find_line_number_of_relevant_line_in_file from ..config_loader import get_settings @@ -12,10 +13,18 @@ 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 @@ -27,37 +36,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: @@ -76,8 +92,7 @@ class BitbucketProvider(GitProvider): "side": "RIGHT", } post_parameters_list.append(post_parameters) - - + try: self.publish_inline_comments(post_parameters_list) return True @@ -88,6 +103,7 @@ class BitbucketProvider(GitProvider): def is_supported(self, capability: str) -> bool: if capability in ['get_issue_comments', 'publish_inline_comments', 'get_labels']: + return False return True @@ -100,28 +116,39 @@ 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}") + # funtion to create_inline_comment def create_inline_comment(self, body: str, relevant_file: str, relevant_line_in_file: str): position, absolute_position = find_line_number_of_relevant_line_in_file(self.get_diff_files(), relevant_file.strip('`'), relevant_line_in_file) @@ -146,14 +173,9 @@ class BitbucketProvider(GitProvider): }, }) 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: @@ -163,11 +185,12 @@ class BitbucketProvider(GitProvider): for comment in comments: self.publish_inline_comment(comment['body'],comment['position'], 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): @@ -180,7 +203,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 @@ -191,14 +216,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] @@ -211,7 +238,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): @@ -225,12 +254,12 @@ class BitbucketProvider(GitProvider): # bitbucket does not support labels def publish_description(self, pr_title: str, pr_body: str): - return "" + pass # bitbucket does not support labels def publish_labels(self, pr_types: list): - return "" + pass # bitbucket does not support labels def get_labels(self): - return "" + 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 1a0cd356..4ca84635 100644 --- a/pr_agent/tools/pr_description.py +++ b/pr_agent/tools/pr_description.py @@ -157,12 +157,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: @@ -198,6 +192,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..0e1289f4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,26 +26,10 @@ classifiers = [ "Operating System :: Independent", "Programming Language :: Python :: 3", ] +dynamic = ["dependencies"] -dependencies = [ - "dynaconf==3.1.12", - "fastapi==0.99.0", - "PyGithub==1.59.*", - "retry==0.9.2", - "openai==0.27.8", - "Jinja2==3.1.2", - "tiktoken==0.4.0", - "uvicorn==0.22.0", - "python-gitlab==3.15.0", - "pytest~=7.4.0", - "aiohttp~=3.8.4", - "atlassian-python-api==3.39.0", - "GitPython~=3.1.32", - "starlette-context==0.3.6", - "litellm~=0.1.445", - "PyYAML==6.0", - "boto3~=1.28.25" -] +[tool.setuptools.dynamic] +dependencies = {file = ["requirements.txt"]} [project.urls] "Homepage" = "https://github.com/Codium-ai/pr-agent" 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