adk-python/contributing/samples/adk_pr_triaging_agent/agent.py

# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pathlib import Path
from typing import Any

from adk_pr_triaging_agent.settings import BOT_LABEL
from adk_pr_triaging_agent.settings import GITHUB_BASE_URL
from adk_pr_triaging_agent.settings import IS_INTERACTIVE
from adk_pr_triaging_agent.settings import OWNER
from adk_pr_triaging_agent.settings import REPO
from adk_pr_triaging_agent.utils import error_response
from adk_pr_triaging_agent.utils import get_diff
from adk_pr_triaging_agent.utils import post_request
from adk_pr_triaging_agent.utils import read_file
from adk_pr_triaging_agent.utils import run_graphql_query
from google.adk import Agent
import requests

LABEL_TO_OWNER = {
    "documentation": "polong-lin",
    "services": "DeanChensj",
    "tools": "seanzhou1023",
    "eval": "ankursharmas",
    "live": "hangfei",
    "models": "genquan9",
    "tracing": "Jacksunwei",
    "core": "Jacksunwei",
    "web": "wyf7107",
}

CONTRIBUTING_MD = read_file(
    Path(__file__).resolve().parents[3] / "CONTRIBUTING.md"
)

APPROVAL_INSTRUCTION = (
    "Do not ask for user approval for labeling or commenting! If you can't find"
    " appropriate labels for the PR, do not label it."
)
if IS_INTERACTIVE:
  APPROVAL_INSTRUCTION = (
      "Only label or comment when the user approves the labeling or commenting!"
  )


def get_pull_request_details(pr_number: int) -> str:
  """Get the details of the specified pull request.

  Args:
    pr_number: number of the Github pull request.

  Returns:
    The status of this request, with the details when successful.
  """
  print(f"Fetching details for PR #{pr_number} from {OWNER}/{REPO}")
  query = """
    query($owner: String!, $repo: String!, $prNumber: Int!) {
      repository(owner: $owner, name: $repo) {
        pullRequest(number: $prNumber) {
          id
          title
          body
          author {
            login
          }
          labels(last: 10) {
            nodes {
              name
            }
          }
          files(last: 50) {
            nodes {
              path
            }
          }
          comments(last: 50) {
            nodes {
              id
              body
              createdAt
              author {
                login
              }
            }
          }
          commits(last: 50) {
            nodes {
              commit {
                url
                message
              }
            }
          }
          statusCheckRollup {
            state
            contexts(last: 20) {
              nodes {
                ... on StatusContext {
                  context
                  state
                  targetUrl
                }
                ... on CheckRun {
                  name
                  status
                  conclusion
                  detailsUrl
                }
              }
            }
          }
        }
      }
    }
  """
  variables = {"owner": OWNER, "repo": REPO, "prNumber": pr_number}
  url = f"{GITHUB_BASE_URL}/repos/{OWNER}/{REPO}/pulls/{pr_number}"

  try:
    response = run_graphql_query(query, variables)
    if "errors" in response:
      return error_response(str(response["errors"]))

    pr = response.get("data", {}).get("repository", {}).get("pullRequest")
    if not pr:
      return error_response(f"Pull Request #{pr_number} not found.")

    # Filter out main merge commits.
    original_commits = pr.get("commits", {}).get("nodes", {})
    if original_commits:
      filtered_commits = [
          commit_node
          for commit_node in original_commits
          if not commit_node["commit"]["message"].startswith(
              "Merge branch 'main' into"
          )
      ]
      pr["commits"]["nodes"] = filtered_commits

    # Get diff of the PR and truncate it to avoid exceeding the maximum tokens.
    pr["diff"] = get_diff(url)[:10000]

    return {"status": "success", "pull_request": pr}
  except requests.exceptions.RequestException as e:
    return error_response(str(e))


def add_label_and_reviewer_to_pr(pr_number: int, label: str) -> dict[str, Any]:
  """Adds a specified label and requests a review from a mapped reviewer on a PR.

  Args:
      pr_number: the number of the Github pull request
      label: the label to add

  Returns:
      The the status of this request, with the applied label and assigned
      reviewer when successful.
  """
  print(f"Attempting to add label '{label}' and a reviewer to PR #{pr_number}")
  if label not in LABEL_TO_OWNER:
    return error_response(
        f"Error: Label '{label}' is not an allowed label. Will not apply."
    )

  # Pull Request is a special issue in Github, so we can use issue url for PR.
  label_url = (
      f"{GITHUB_BASE_URL}/repos/{OWNER}/{REPO}/issues/{pr_number}/labels"
  )
  label_payload = [label, BOT_LABEL]

  try:
    response = post_request(label_url, label_payload)
  except requests.exceptions.RequestException as e:
    return error_response(f"Error: {e}")

  owner = LABEL_TO_OWNER.get(label, None)
  if not owner:
    return {
        "status": "warning",
        "message": (
            f"{response}\n\nLabel '{label}' does not have an owner. Will not"
            " assign."
        ),
        "applied_label": label,
    }
  reviewer_url = f"{GITHUB_BASE_URL}/repos/{OWNER}/{REPO}/pulls/{pr_number}/requested_reviewers"
  reviewer_payload = {"reviewers": [owner]}
  try:
    post_request(reviewer_url, reviewer_payload)
  except requests.exceptions.RequestException as e:
    return {
        "status": "warning",
        "message": f"Reviewer not assigned: {e}",
        "applied_label": label,
    }

  return {
      "status": "success",
      "applied_label": label,
      "assigned_reviewer": owner,
  }


def add_comment_to_pr(pr_number: int, comment: str) -> dict[str, Any]:
  """Add the specified comment to the given PR number.

  Args:
    pr_number: the number of the Github pull request
    comment: the comment to add

  Returns:
    The the status of this request, with the applied comment when successful.
  """
  print(f"Attempting to add comment '{comment}' to issue #{pr_number}")

  # Pull Request is a special issue in Github, so we can use issue url for PR.
  url = f"{GITHUB_BASE_URL}/repos/{OWNER}/{REPO}/issues/{pr_number}/comments"
  payload = {"body": comment}

  try:
    post_request(url, payload)
  except requests.exceptions.RequestException as e:
    return error_response(f"Error: {e}")
  return {
      "status": "success",
      "added_comment": comment,
  }


root_agent = Agent(
    model="gemini-2.5-pro",
    name="adk_pr_triaging_assistant",
    description="Triage ADK pull requests.",
    instruction=f"""
      # 1. Identity
      You are a Pull Request (PR) triaging bot for the Github {REPO} repo with the owner {OWNER}.

      # 2. Responsibilities
      Your core responsibility includes:
      - Get the pull request details.
      - Add a label to the pull request.
      - Assign a reviewer to the pull request.
      - Check if the pull request is following the contribution guidelines.
      - Add a comment to the pull request if it's not following the guidelines.

      **IMPORTANT: {APPROVAL_INSTRUCTION}**

      # 3. Guidelines & Rules
      Here are the rules for labeling:
      - If the PR is about documentations, label it with "documentation".
      - If it's about session, memory, artifacts services, label it with "services"
      - If it's about UI/web, label it with "web"
      - If it's related to tools, label it with "tools"
      - If it's about agent evalaution, then label it with "eval".
      - If it's about streaming/live, label it with "live".
      - If it's about model support(non-Gemini, like Litellm, Ollama, OpenAI models), label it with "models".
      - If it's about tracing, label it with "tracing".
      - If it's agent orchestration, agent definition, label it with "core".
      - If you can't find a appropriate labels for the PR, follow the previous instruction that starts with "IMPORTANT:".

      Here is the contribution guidelines:
      `{CONTRIBUTING_MD}`

      Here are the guidelines for checking if the PR is following the guidelines:
      - The "statusCheckRollup" in the pull request details may help you to identify if the PR is following some of the guidelines (e.g. CLA compliance).

      Here are the guidelines for the comment:
      - **Be Polite and Helpful:** Start with a friendly tone.
      - **Be Specific:** Clearly list only the sections from the contribution guidelines that are still missing.
      - **Address the Author:** Mention the PR author by their username (e.g., `@username`).
      - **Provide Context:** Explain *why* the information or action is needed.
      - **Do not be repetitive:** If you have already commented on an PR asking for information, do not comment again unless new information has been added and it's still incomplete.
      - **Identify yourself:** Include a bolded note (e.g. "Response from ADK Triaging Agent") in your comment to indicate this comment was added by an ADK Answering Agent.

      **Example Comment for a PR:**
      > **Response from ADK Triaging Agent**
      >
      > Hello @[pr-author-username], thank you for creating this PR!
      >
      > This PR is a bug fix, could you please associate the github issue with this PR? If there is no existing issue, could you please create one?
      >
      > In addition, could you please provide logs or screenshot after the fix is applied?
      >
      > This information will help reviewers to review your PR more efficiently. Thanks!

      # 4. Steps
      When you are given a PR, here are the steps you should take:
      - Call the `get_pull_request_details` tool to get the details of the PR.
      - Skip the PR (i.e. do not label or comment) if the PR is closed or is labeled with "{BOT_LABEL}" or "google-contributior".
      - Check if the PR is following the contribution guidelines.
        - If it's not following the guidelines, recommend or add a comment to the PR that points to the contribution guidelines (https://github.com/google/adk-python/blob/main/CONTRIBUTING.md).
        - If it's following the guidelines, recommend or add a label to the PR.

      # 5. Output
      Present the followings in an easy to read format highlighting PR number and your label.
      - The PR summary in a few sentence
      - The label you recommended or added with the justification
      - The owner of the label if you assigned a reviewer to the PR
      - The comment you recommended or added to the PR with the justification
    """,
    tools=[
        get_pull_request_details,
        add_label_and_reviewer_to_pr,
        add_comment_to_pr,
    ],
)