wip

pr_agent/servers/gitlab_webhook.py

@@ -3,9 +3,10 @@
 from datetime import datetime
 
 import uvicorn
-from fastapi import APIRouter, FastAPI, Request, status
+from fastapi import APIRouter, FastAPI, Request, status, HTTPException
 from fastapi.encoders import jsonable_encoder
 from fastapi.responses import JSONResponse
+from starlette.middleware.base import BaseHTTPMiddleware
 from starlette.background import BackgroundTasks
 from starlette.middleware import Middleware
 from starlette_context import context
@@ -196,14 +197,37 @@ async def root():
 if not gitlab_url:
     raise ValueError("GITLAB.URL is not set")
 get_settings().config.git_provider = "gitlab"
-middleware = [Middleware(RawContextMiddleware)]
-app = FastAPI(middleware=middleware)
+
+class IPWhitelistMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request: Request, call_next):
+        whitelisted_ips = get_settings().get("CONFIG.WHITELISTED_IPS", [])
+        client_ip = request.client.host
+
+        # strip port from client_ip if present.
+        if ":" in client_ip:
+            client_ip = client_ip.split(":")[0]
+
+        if client_ip not in whitelisted_ips and whitelisted_ips != []:
+            print (f"Client IP {client_ip} not in whitelisted IPs {whitelisted_ips}")
+            raise HTTPException(status_code=403, detail="Forbidden")
+        response = await call_next(request)
+        return response
+
+app = FastAPI()
 app.include_router(router)
+app.add_middleware(RawContextMiddleware)
+app.add_middleware(IPWhitelistMiddleware)
 
+# Custom exception handler for HTTPException
+@app.exception_handler(HTTPException)
+async def http_exception_handler(_: Request, exc: HTTPException):
+    return JSONResponse(
+        status_code=exc.status_code,
+        content={"message": exc.detail},
+    )
 
 def start():
     uvicorn.run(app, host="0.0.0.0", port=3000)
 
-
 if __name__ == '__main__':
     start()

pr_agent/settings/configuration.toml

@@ -1,131 +1,133 @@
 [config]
 # models
-model="gpt-4-turbo-2024-04-09"
-model_turbo="gpt-4o-2024-08-06"
-fallback_models=["gpt-4o-2024-05-13"]
+model = "gpt-4-turbo-2024-04-09"
+model_turbo = "gpt-4o-2024-08-06"
+fallback_models = ["gpt-4o-2024-05-13"]
 # CLI
-git_provider="github"
-publish_output=true
-publish_output_progress=true
-verbosity_level=0 # 0,1,2
-use_extra_bad_extensions=false
+git_provider = "github"
+publish_output = true
+publish_output_progress = true
+verbosity_level = 0              # 0,1,2
+use_extra_bad_extensions = false
 # Configurations
-use_wiki_settings_file=true
-use_repo_settings_file=true
-use_global_settings_file=true
-ai_timeout=120 # 2minutes
+use_wiki_settings_file = true
+use_repo_settings_file = true
+use_global_settings_file = true
+ai_timeout = 120                # 2minutes
 # token limits
 max_description_tokens = 500
 max_commits_tokens = 500
-max_model_tokens = 32000 # Limits the maximum number of tokens that can be used by any model, regardless of the model's default capabilities.
-custom_model_max_tokens=-1 # for models not in the default list
+max_model_tokens = 32000     # Limits the maximum number of tokens that can be used by any model, regardless of the model's default capabilities.
+custom_model_max_tokens = -1 # for models not in the default list
 # patch extension logic
-patch_extension_skip_types =[".md",".txt"]
-allow_dynamic_context=false
-max_extra_lines_before_dynamic_context = 10 # will try to include up to 10 extra lines before the hunk in the patch, until we reach an enclosing function or class
-patch_extra_lines_before = 3 # Number of extra lines (+3 default ones) to include before each hunk in the patch
-patch_extra_lines_after = 1 # Number of extra lines (+3 default ones) to include after each hunk in the patch
-secret_provider=""
-cli_mode=false
-ai_disclaimer_title=""  # Pro feature, title for a collapsible disclaimer to AI outputs
-ai_disclaimer=""  # Pro feature, full text for the AI disclaimer
-output_relevant_configurations=false
-large_patch_policy = "clip" # "clip", "skip"
-is_auto_command=false
+patch_extension_skip_types = [".md", ".txt"]
+allow_dynamic_context = false
+max_extra_lines_before_dynamic_context = 10  # will try to include up to 10 extra lines before the hunk in the patch, until we reach an enclosing function or class
+patch_extra_lines_before = 3                 # Number of extra lines (+3 default ones) to include before each hunk in the patch
+patch_extra_lines_after = 1                  # Number of extra lines (+3 default ones) to include after each hunk in the patch
+secret_provider = ""
+cli_mode = false
+ai_disclaimer_title = ""                     # Pro feature, title for a collapsible disclaimer to AI outputs
+ai_disclaimer = ""                           # Pro feature, full text for the AI disclaimer
+output_relevant_configurations = false
+large_patch_policy = "clip"                  # "clip", "skip"
+is_auto_command = false
 # seed
-seed=-1 # set positive value to fix the seed (and ensure temperature=0)
-temperature=0.2
+seed = -1         # set positive value to fix the seed (and ensure temperature=0)
+temperature = 0.2
+# server
+server_whitelisted_ips = []
 
 [pr_reviewer] # /review #
 # enable/disable features
-require_score_review=false
-require_tests_review=true
-require_estimate_effort_to_review=true
-require_can_be_split_review=false
-require_security_review=true
-extra_issue_links=false
+require_score_review = false
+require_tests_review = true
+require_estimate_effort_to_review = true
+require_can_be_split_review = false
+require_security_review = true
+extra_issue_links = false
 # soc2
-require_soc2_ticket=false
-soc2_ticket_prompt="Does the PR description include a link to ticket in a project management system (e.g., Jira, Asana, Trello, etc.) ?"
+require_soc2_ticket = false
+soc2_ticket_prompt = "Does the PR description include a link to ticket in a project management system (e.g., Jira, Asana, Trello, etc.) ?"
 # general options
-num_code_suggestions=0
+num_code_suggestions = 0
 inline_code_comments = false
-ask_and_reflect=false
+ask_and_reflect = false
 #automatic_review=true
-persistent_comment=true
+persistent_comment = true
 extra_instructions = ""
 final_update_message = true
 # review labels
-enable_review_labels_security=true
-enable_review_labels_effort=true
+enable_review_labels_security = true
+enable_review_labels_effort = true
 # specific configurations for incremental review (/review -i)
-require_all_thresholds_for_incremental_review=false
-minimal_commits_for_incremental_review=0
-minimal_minutes_for_incremental_review=0
-enable_help_text=false # Determines whether to include help text in the PR review. Enabled by default.
+require_all_thresholds_for_incremental_review = false
+minimal_commits_for_incremental_review = 0
+minimal_minutes_for_incremental_review = 0
+enable_help_text = false                              # Determines whether to include help text in the PR review. Enabled by default.
 # auto approval
-enable_auto_approval=false
-maximal_review_effort=5
+enable_auto_approval = false
+maximal_review_effort = 5
 
 
 [pr_description] # /describe #
-publish_labels=true
-add_original_user_description=true
-generate_ai_title=false
-use_bullet_points=true
+publish_labels = true
+add_original_user_description = true
+generate_ai_title = false
+use_bullet_points = true
 extra_instructions = ""
-enable_pr_type=true
+enable_pr_type = true
 final_update_message = true
-enable_help_text=false
-enable_help_comment=true
+enable_help_text = false
+enable_help_comment = true
 # describe as comment
-publish_description_as_comment=false
-publish_description_as_comment_persistent=true
+publish_description_as_comment = false
+publish_description_as_comment_persistent = true
 ## changes walkthrough section
-enable_semantic_files_types=true
-collapsible_file_list='adaptive' # true, false, 'adaptive'
-inline_file_summary=false # false, true, 'table'
+enable_semantic_files_types = true
+collapsible_file_list = 'adaptive' # true, false, 'adaptive'
+inline_file_summary = false        # false, true, 'table'
 # markers
-use_description_markers=false
-include_generated_by_header=true
+use_description_markers = false
+include_generated_by_header = true
 # large pr mode 💎
-enable_large_pr_handling=true
-max_ai_calls=4
-async_ai_calls=true
-mention_extra_files=true
+enable_large_pr_handling = true
+max_ai_calls = 4
+async_ai_calls = true
+mention_extra_files = true
 #custom_labels = ['Bug fix', 'Tests', 'Bug fix with tests', 'Enhancement', 'Documentation', 'Other']
 
 [pr_questions] # /ask #
-enable_help_text=false
+enable_help_text = false
 
 
 [pr_code_suggestions] # /improve #
-max_context_tokens=14000
-num_code_suggestions=4
+max_context_tokens = 14000
+num_code_suggestions = 4
 commitable_code_suggestions = false
 extra_instructions = ""
 rank_suggestions = false
-enable_help_text=false
-persistent_comment=true
-max_history_len=4
+enable_help_text = false
+persistent_comment = true
+max_history_len = 4
 # enable to apply suggestion 💎
-apply_suggestions_checkbox=true
+apply_suggestions_checkbox = true
 # suggestions scoring
-self_reflect_on_suggestions=true
-suggestions_score_threshold=0 # [0-10]. highly recommend not to set this value above 8, since above it may clip highly relevant suggestions
+self_reflect_on_suggestions = true
+suggestions_score_threshold = 0    # [0-10]. highly recommend not to set this value above 8, since above it may clip highly relevant suggestions
 # params for '/improve --extended' mode
-auto_extended_mode=true
-num_code_suggestions_per_chunk=4
+auto_extended_mode = true
+num_code_suggestions_per_chunk = 4
 max_number_of_calls = 3
 parallel_calls = true
 rank_extended_suggestions = false
 final_clip_factor = 0.8
 # self-review checkbox
-demand_code_suggestions_self_review=false # add a checkbox for the author to self-review the code suggestions
-code_suggestions_self_review_text= "**Author self-review**: I have reviewed the PR code suggestions, and addressed the relevant ones."
-approve_pr_on_self_review=false # Pro feature. if true, the PR will be auto-approved after the author clicks on the self-review checkbox
+demand_code_suggestions_self_review = false                                                                                             # add a checkbox for the author to self-review the code suggestions
+code_suggestions_self_review_text = "**Author self-review**: I have reviewed the PR code suggestions, and addressed the relevant ones."
+approve_pr_on_self_review = false                                                                                                       # Pro feature. if true, the PR will be auto-approved after the author clicks on the self-review checkbox
 # Suggestion impact
-publish_post_process_suggestion_impact=true
+publish_post_process_suggestion_impact = true
 
 [pr_custom_prompt] # /custom_prompt #
 prompt = """\
@@ -134,45 +136,47 @@ The code suggestions should focus only on the following:
 - ...
 ...
 """
-suggestions_score_threshold=0
-num_code_suggestions_per_chunk=4
-self_reflect_on_custom_suggestions=true
-enable_help_text=false
+suggestions_score_threshold = 0
+num_code_suggestions_per_chunk = 4
+self_reflect_on_custom_suggestions = true
+enable_help_text = false
 
 
 [pr_add_docs] # /add_docs #
 extra_instructions = ""
-docs_style = "Sphinx" # "Google Style with Args, Returns, Attributes...etc", "Numpy Style", "Sphinx Style", "PEP257", "reStructuredText"
-file = ""              # in case there are several components with the same name, you can specify the relevant file
-class_name = ""        # in case there are several methods with the same name in the same file, you can specify the relevant class name
+docs_style = "Sphinx"   # "Google Style with Args, Returns, Attributes...etc", "Numpy Style", "Sphinx Style", "PEP257", "reStructuredText"
+file = ""               # in case there are several components with the same name, you can specify the relevant file
+class_name = ""         # in case there are several methods with the same name in the same file, you can specify the relevant class name
 
 [pr_update_changelog] # /update_changelog #
-push_changelog_changes=false
+push_changelog_changes = false
 extra_instructions = ""
 
 [pr_analyze] # /analyze #
-enable_help_text=true
+enable_help_text = true
 
 [pr_test] # /test #
 extra_instructions = ""
-testing_framework = "" # specify the testing framework you want to use
-num_tests=3            # number of tests to generate. max 5.
-avoid_mocks=true       # if true, the generated tests will prefer to use real objects instead of mocks
-file = ""              # in case there are several components with the same name, you can specify the relevant file
-class_name = ""        # in case there are several methods with the same name in the same file, you can specify the relevant class name
-enable_help_text=false
+testing_framework = ""   # specify the testing framework you want to use
+num_tests = 3            # number of tests to generate. max 5.
+avoid_mocks = true       # if true, the generated tests will prefer to use real objects instead of mocks
+file = ""                # in case there are several components with the same name, you can specify the relevant file
+class_name = ""          # in case there are several methods with the same name in the same file, you can specify the relevant class name
+enable_help_text = false
 
 [pr_improve_component] # /improve_component #
-num_code_suggestions=4
+num_code_suggestions = 4
 extra_instructions = ""
-file = ""              # in case there are several components with the same name, you can specify the relevant file
-class_name = ""        # in case there are several methods with the same name in the same file, you can specify the relevant class name
+file = ""                # in case there are several components with the same name, you can specify the relevant file
+class_name = ""          # in case there are several methods with the same name in the same file, you can specify the relevant class name
 
 [checks] # /checks (pro feature) #
-enable_auto_checks_feedback=true
-excluded_checks_list=["lint"] # list of checks to exclude, for example: ["check1", "check2"]
-persistent_comment=true
-enable_help_text=true
+enable_auto_checks_feedback = true
+excluded_checks_list = [
+    "lint",
+] # list of checks to exclude, for example: ["check1", "check2"]
+persistent_comment = true
+enable_help_text = true
 final_update_message = false
 
 [pr_help] # /help #
@@ -211,10 +215,7 @@ push_trigger_ignore_merge_commits = true
 push_trigger_wait_for_initial_review = true
 push_trigger_pending_tasks_backlog = true
 push_trigger_pending_tasks_ttl = 300
-push_commands = [
-    "/describe",
-    "/review --pr_reviewer.num_code_suggestions=0",
-]
+push_commands = ["/describe", "/review --pr_reviewer.num_code_suggestions=0"]
 ignore_pr_title = []
 ignore_bot_pr = true
 
@@ -226,10 +227,7 @@ pr_commands = [
     "/improve",
 ]
 handle_push_trigger = false
-push_commands = [
-    "/describe",
-    "/review --pr_reviewer.num_code_suggestions=0",
-]
+push_commands = ["/describe", "/review --pr_reviewer.num_code_suggestions=0"]
 
 [bitbucket_app]
 pr_commands = [