Fix lint errors

.github/actions/summarizer/gemini_summarize_pr.py

@@ -1,15 +1,18 @@
+"""Summarizes the pull request using the Gemini model and adds summary to a PR comment."""
+
 import json
 import os
+import requests
 
 from github import Github
-import requests
 import vertexai
 from vertexai.generative_models import GenerationConfig, GenerativeModel
 
+GEMINI_MODEL = "gemini-1.5-flash-001"
 
-def get_pr_number() -> str:
-    event_path = os.getenv("GITHUB_EVENT_PATH", "")
 
+def get_pr_number(event_path: str) -> str:
+    """Retrieves the pull request number from GitHub event data."""
     # Load event data
     with open(event_path, "r", encoding="utf-8") as f:
         event_data = json.load(f)
@@ -26,22 +29,53 @@ def get_pr_number() -> str:
     raise ValueError("Unable to determine pull request number from event data.")
 
 
-def call_gemini(
-    pull_request_content: str, model_id: str = "gemini-1.5-flash-001"
+def get_pr_content(pr: Github.PullRequest) -> str:
+    """Returns the content of the pull request as a string."""
+    pr_content = f"""
+    Title: {pr.title}
+    Pull Request Description: {pr.body}
+
+    --- Files Changed ---
+    """
+
+    for file in pr.get_files():
+        pr_content += f"File name: {file.filename}\n\n"
+
+        # Attempt to fetch raw content if patch is not available
+        if file.patch is None:
+            try:
+                raw_content = requests.get(file.raw_url, timeout=10).text
+                pr_content += f"Raw File Content:\n`\n{raw_content}\n`\n\n"
+            except requests.exceptions.RequestException:
+                pr_content += "Unable to fetch raw file content.\n\n"
+        else:  # Use patch if available
+            pr_content += f"Code Diff:\n{file.patch}\n\n"
+
+    return pr_content
+
+
+def summarize_pr_gemini(
+    pull_request_content: str,
+    project_id: str,
+    location: str = "us-central1",
+    model_id: str = GEMINI_MODEL,
 ) -> str:
-    project_id = os.getenv("GOOGLE_CLOUD_PROJECT_ID")
-    vertexai.init(project=project_id, location="us-central1")
+    """Calls the Gemini model to summarize the pull request content."""
+    vertexai.init(project=project_id, location=location)
 
     model = GenerativeModel(
         model_id,
         system_instruction=[
             "You are an expert software engineer, proficient in Generative AI, Git and GitHub.",
         ],
-        generation_config=GenerationConfig(temperature=0.0),
+        generation_config=GenerationConfig(temperature=0.0, max_output_tokens=8192),
     )
 
     prompt = [
-        "The following is the content of a GitHub Pull Request for a repository focused on Generative AI with Google Cloud. This content includes the Pull Request title, Pull Request description, a list of all of the files changed with the file name, the code diff and the raw file content. Your task is to output a summary of the Pull Request in Markdown format.",
+        "The following is the content of a GitHub Pull Request for a repository focused on Generative AI with Google Cloud."
+        "This content includes the Pull Request title, Pull Request description, "
+        "a list of all of the files changed with the file name, the code diff and the raw file content."
+        "Your task is to output a summary of the Pull Request in Markdown format.",
         "Content:",
         pull_request_content,
         "Summary:",
@@ -54,34 +88,8 @@ def call_gemini(
     return response.text.replace("## Pull Request Summary", "")
 
 
-def summarize_pr(token, repo_name, pr_number):
-    """Summarizes the pull request using the Gemini model and updates/creates a comment."""
-    g = Github(token)
-    repo = g.get_repo(repo_name)
-    pr = repo.get_pull(pr_number)
-
-    pr_content = f"""
-    Title: {pr.title}
-    Pull Request Description: {pr.body}
-    
-    --- Files Changed ---
-    """
-
-    for file in pr.get_files():
-        pr_content += f"File name: {file.filename}\n\n"
-
-        # Attempt to fetch raw content if patch is not available
-        if file.patch is None:
-            try:
-                raw_content = requests.get(file.raw_url).text
-                pr_content += f"Raw File Content:\n`\n{raw_content}\n`\n\n"
-            except requests.exceptions.RequestException:
-                pr_content += "Unable to fetch raw file content.\n\n"
-        else:  # Use patch if available
-            pr_content += f"Code Diff:\n{file.patch}\n\n"
-
-    summary = call_gemini(pr_content)
-
+def add_pr_comment(pr: Github.PullRequest, summary: str) -> None:
+    """Comments on the pull request with the provided summary."""
     comment_header = "## Pull Request Summary from [Gemini ✨](https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/overview)"
     comment_body = (
         f"{comment_header}\n{summary}\n---\nGenerated at `{pr.get_commits()[-1].sha}`"
@@ -98,13 +106,21 @@ def summarize_pr(token, repo_name, pr_number):
     pr.create_issue_comment(comment_body)
 
 
-def main():
+def main() -> None:
+    """Summarizes the pull request using the Gemini model and adds summary to a PR comment."""
+
     # Get GitHub token and repository details
     repo_name = os.getenv("GITHUB_REPOSITORY")
     token = os.getenv("GITHUB_TOKEN")
-    pr_number = get_pr_number()
+    pr_number = get_pr_number(os.getenv("GITHUB_EVENT_PATH", ""))
+
+    g = Github(token)
+    repo = g.get_repo(repo_name)
+    pr = repo.get_pull(pr_number)
 
-    summarize_pr(token, repo_name, pr_number)
+    pr_content = get_pr_content(pr)
+    summary = summarize_pr_gemini(pr_content, os.getenv("GOOGLE_CLOUD_PROJECT_ID", ""))
+    add_pr_comment(pr, summary)
 
 
 if __name__ == "__main__":