chore: refactor the ADK Triaging Agent to make the code easier to read

PiperOrigin-RevId: 776763061

.github/workflows/triage.yml

@@ -40,4 +40,5 @@ jobs:
           ISSUE_TITLE: ${{ github.event.issue.title }}
           ISSUE_BODY: ${{ github.event.issue.body }}
           ISSUE_COUNT_TO_PROCESS: '3' # Process 3 issues at a time on schedule
-        run: python contributing/samples/adk_triaging_agent/main.py
+          PYTHONPATH: contributing/samples
+        run: python -m adk_triaging_agent.main

contributing/samples/adk_triaging_agent/agent.py

@@ -12,26 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
+from typing import Any
 
+from adk_triaging_agent.settings import BOT_LABEL
+from adk_triaging_agent.settings import GITHUB_BASE_URL
+from adk_triaging_agent.settings import IS_INTERACTIVE
+from adk_triaging_agent.settings import OWNER
+from adk_triaging_agent.settings import REPO
+from adk_triaging_agent.utils import error_response
+from adk_triaging_agent.utils import get_request
+from adk_triaging_agent.utils import post_request
 from google.adk import Agent
 import requests
 
-GITHUB_TOKEN = os.getenv("GITHUB_TOKEN")
-if not GITHUB_TOKEN:
-  raise ValueError("GITHUB_TOKEN environment variable not set")
-
-OWNER = os.getenv("OWNER", "google")
-REPO = os.getenv("REPO", "adk-python")
-BOT_LABEL = os.getenv("BOT_LABEL", "bot_triaged")
-
-BASE_URL = "https://api.github.com"
-
-headers = {
-    "Authorization": f"token {GITHUB_TOKEN}",
-    "Accept": "application/vnd.github.v3+json",
-}
-
 ALLOWED_LABELS = [
     "documentation",
     "services",
@@ -45,24 +38,25 @@ ALLOWED_LABELS = [
     "web",
 ]
 
-
-def is_interactive():
-  return os.environ.get("INTERACTIVE", "1").lower() in ["true", "1"]
+APPROVAL_INSTRUCTION = (
+    "Do not ask for user approval for labeling! If you can't find appropriate"
+    " labels for the issue, do not label it."
+)
+if IS_INTERACTIVE:
+  APPROVAL_INSTRUCTION = "Only label them when the user approves the labeling!"
 
 
-def list_issues(issue_count: int):
-  """
-  Generator to list all issues for the repository by handling pagination.
+def list_unlabeled_issues(issue_count: int) -> dict[str, Any]:
+  """List most recent `issue_count` numer of unlabeled issues in the repo.
 
   Args:
     issue_count: number of issues to return
 
+  Returns:
+    The status of this request, with a list of issues when successful.
   """
+  url = f"{GITHUB_BASE_URL}/search/issues"
   query = f"repo:{OWNER}/{REPO} is:open is:issue no:label"
-
-  unlabelled_issues = []
-  url = f"{BASE_URL}/search/issues"
-
   params = {
       "q": query,
       "sort": "created",
@@ -70,57 +64,57 @@ def list_issues(issue_count: int):
       "per_page": issue_count,
       "page": 1,
   }
-  response = requests.get(url, headers=headers, params=params, timeout=60)
-  response.raise_for_status()
-  json_response = response.json()
-  issues = json_response.get("items", None)
-  if not issues:
-    return []
+
+  try:
+    response = get_request(url, params)
+  except requests.exceptions.RequestException as e:
+    return error_response(f"Error: {e}")
+  issues = response.get("items", None)
+
+  unlabeled_issues = []
   for issue in issues:
-    if not issue.get("labels", None) or len(issue["labels"]) == 0:
-      unlabelled_issues.append(issue)
-  return unlabelled_issues
+    if not issue.get("labels", None):
+      unlabeled_issues.append(issue)
+  return {"status": "success", "issues": unlabeled_issues}
 
 
-def add_label_to_issue(issue_number: str, label: str):
-  """
-  Add the specified label to the given issue number.
+def add_label_to_issue(issue_number: int, label: str) -> dict[str, Any]:
+  """Add the specified label to the given issue number.
 
   Args:
-    issue_number: issue number of the Github issue, in string foramt.
+    issue_number: issue number of the Github issue.
     label: label to assign
+
+  Returns:
+    The the status of this request, with the applied label when successful.
   """
   print(f"Attempting to add label '{label}' to issue #{issue_number}")
   if label not in ALLOWED_LABELS:
-    error_message = (
+    return error_response(
         f"Error: Label '{label}' is not an allowed label. Will not apply."
     )
-    print(error_message)
-    return {"status": "error", "message": error_message, "applied_label": None}
 
-  url = f"{BASE_URL}/repos/{OWNER}/{REPO}/issues/{issue_number}/labels"
+  url = f"{GITHUB_BASE_URL}/repos/{OWNER}/{REPO}/issues/{issue_number}/labels"
   payload = [label, BOT_LABEL]
-  response = requests.post(url, headers=headers, json=payload, timeout=60)
-  response.raise_for_status()
-  return response.json()
 
+  try:
+    response = post_request(url, payload)
+  except requests.exceptions.RequestException as e:
+    return error_response(f"Error: {e}")
+  return {
+      "status": "success",
+      "message": response,
+      "applied_label": label,
+  }
 
-approval_instruction = (
-    "Only label them when the user approves the labeling!"
-    if is_interactive()
-    else (
-        "Do not ask for user approval for labeling! If you can't find a"
-        " appropriate labels for the issue, do not label it."
-    )
-)
 
 root_agent = Agent(
-    model="gemini-2.5-pro-preview-05-06",
+    model="gemini-2.5-pro",
     name="adk_triaging_assistant",
     description="Triage ADK issues.",
     instruction=f"""
-      You are a Github adk-python repo triaging bot. You will help get issues, and recommend a label.
-      IMPORTANT: {approval_instruction}
+      You are a triaging bot for the Github {REPO} repo with the owner {OWNER}. You will help get issues, and recommend a label.
+      IMPORTANT: {APPROVAL_INSTRUCTION}
       Here are the rules for labeling:
       - If the user is asking about documentation-related questions, label it with "documentation".
       - If it's about session, memory services, label it with "services"
@@ -138,8 +132,5 @@ root_agent = Agent(
       - the issue summary in a few sentence
       - your label recommendation and justification
     """,
-    tools=[
-        list_issues,
-        add_label_to_issue,
-    ],
+    tools=[list_unlabeled_issues, add_label_to_issue],
 )

contributing/samples/adk_triaging_agent/main.py

@@ -13,48 +13,37 @@
 # limitations under the License.
 
 import asyncio
-import os
 import time
 
-import agent
-from dotenv import load_dotenv
+from adk_triaging_agent import agent
+from adk_triaging_agent.settings import EVENT_NAME
+from adk_triaging_agent.settings import GITHUB_BASE_URL
+from adk_triaging_agent.settings import ISSUE_BODY
+from adk_triaging_agent.settings import ISSUE_COUNT_TO_PROCESS
+from adk_triaging_agent.settings import ISSUE_NUMBER
+from adk_triaging_agent.settings import ISSUE_TITLE
+from adk_triaging_agent.settings import OWNER
+from adk_triaging_agent.settings import REPO
+from adk_triaging_agent.utils import get_request
+from adk_triaging_agent.utils import parse_number_string
 from google.adk.agents.run_config import RunConfig
 from google.adk.runners import InMemoryRunner
-from google.adk.sessions import Session
+from google.adk.runners import Runner
 from google.genai import types
 import requests
 
-load_dotenv(override=True)
-
-OWNER = os.getenv("OWNER", "google")
-REPO = os.getenv("REPO", "adk-python")
-GITHUB_TOKEN = os.getenv("GITHUB_TOKEN")
-BASE_URL = "https://api.github.com"
-headers = {
-    "Authorization": f"token {GITHUB_TOKEN}",
-    "Accept": "application/vnd.github.v3+json",
-}
-
-if not GITHUB_TOKEN:
-  print(
-      "Warning: GITHUB_TOKEN environment variable not set. API calls might"
-      " fail."
-  )
+APP_NAME = "adk_triage_app"
+USER_ID = "adk_triage_user"
 
 
 async def fetch_specific_issue_details(issue_number: int):
   """Fetches details for a single issue if it's unlabelled."""
-  if not GITHUB_TOKEN:
-    print("Cannot fetch issue details: GITHUB_TOKEN is not set.")
-    return None
-
-  url = f"{BASE_URL}/repos/{OWNER}/{REPO}/issues/{issue_number}"
+  url = f"{GITHUB_BASE_URL}/repos/{OWNER}/{REPO}/issues/{issue_number}"
   print(f"Fetching details for specific issue: {url}")
+
   try:
-    response = requests.get(url, headers=headers, timeout=60)
-    response.raise_for_status()
-    issue_data = response.json()
-    if not issue_data.get("labels") or len(issue_data["labels"]) == 0:
+    issue_data = get_request(url)
+    if not issue_data.get("labels", None):
       print(f"Issue #{issue_number} is unlabelled. Proceeding.")
       return {
           "number": issue_data["number"],
@@ -71,94 +60,91 @@ async def fetch_specific_issue_details(issue_number: int):
     return None
 
 
-async def main():
-  app_name = "triage_app"
-  user_id_1 = "triage_user"
-  runner = InMemoryRunner(
-      agent=agent.root_agent,
-      app_name=app_name,
-  )
-  session_11 = await runner.session_service.create_session(
-      app_name=app_name, user_id=user_id_1
+async def call_agent_async(
+    runner: Runner, user_id: str, session_id: str, prompt: str
+) -> str:
+  """Call the agent asynchronously with the user's prompt."""
+  content = types.Content(
+      role="user", parts=[types.Part.from_text(text=prompt)]
   )
 
-  async def run_agent_prompt(session: Session, prompt_text: str):
-    content = types.Content(
-        role="user", parts=[types.Part.from_text(text=prompt_text)]
-    )
-    print(f"\n>>>> Agent Prompt: {prompt_text}")
-    final_agent_response_parts = []
+  final_response_text = ""
   async for event in runner.run_async(
-        user_id=user_id_1,
-        session_id=session.id,
+      user_id=user_id,
+      session_id=session_id,
       new_message=content,
       run_config=RunConfig(save_input_blobs_as_artifacts=False),
   ):
-      if event.content.parts and event.content.parts[0].text:
+    if (
+        event.content
+        and event.content.parts
+        and hasattr(event.content.parts[0], "text")
+        and event.content.parts[0].text
+    ):
       print(f"** {event.author} (ADK): {event.content.parts[0].text}")
       if event.author == agent.root_agent.name:
-          final_agent_response_parts.append(event.content.parts[0].text)
-    print(f"<<<< Agent Final Output: {''.join(final_agent_response_parts)}\n")
+        final_response_text += event.content.parts[0].text
 
-  event_name = os.getenv("EVENT_NAME")
-  issue_number_str = os.getenv("ISSUE_NUMBER")
+  return final_response_text
 
-  if event_name == "issues" and issue_number_str:
-    print(f"EVENT: Processing specific issue due to '{event_name}' event.")
-    try:
-      issue_number = int(issue_number_str)
-      specific_issue = await fetch_specific_issue_details(issue_number)
 
-      if specific_issue:
-        prompt = (
-            f"A new GitHub issue #{specific_issue['number']} has been opened or"
-            f" reopened. Title: \"{specific_issue['title']}\"\nBody:"
-            f" \"{specific_issue['body']}\"\n\nBased on the rules, recommend an"
-            " appropriate label and its justification."
-            " Then, use the 'add_label_to_issue' tool to apply the label "
-            "directly to this issue."
-            f" The issue number is {specific_issue['number']}."
+async def main():
+  runner = InMemoryRunner(
+      agent=agent.root_agent,
+      app_name=APP_NAME,
   )
-        await run_agent_prompt(session_11, prompt)
-      else:
+  session = await runner.session_service.create_session(
+      user_id=USER_ID,
+      app_name=APP_NAME,
+  )
+
+  if EVENT_NAME == "issues" and ISSUE_NUMBER:
+    print(f"EVENT: Processing specific issue due to '{EVENT_NAME}' event.")
+    issue_number = parse_number_string(ISSUE_NUMBER)
+    if not issue_number:
+      print(f"Error: Invalid issue number received: {ISSUE_NUMBER}.")
+      return
+
+    specific_issue = await fetch_specific_issue_details(issue_number)
+    if specific_issue is None:
       print(
           f"No unlabelled issue details found for #{issue_number} or an error"
           " occurred. Skipping agent interaction."
       )
+      return
 
-    except ValueError:
-      print(f"Error: Invalid ISSUE_NUMBER received: {issue_number_str}")
-
-  else:
-    print(f"EVENT: Processing batch of issues (event: {event_name}).")
-    issue_count_str = os.getenv("ISSUE_COUNT_TO_PROCESS", "3")
-    try:
-      num_issues_to_process = int(issue_count_str)
-    except ValueError:
-      print(f"Warning: Invalid ISSUE_COUNT_TO_PROCESS. Defaulting to 3.")
-      num_issues_to_process = 3
-
+    issue_title = ISSUE_TITLE or specific_issue["title"]
+    issue_body = ISSUE_BODY or specific_issue["body"]
     prompt = (
-        f"List the first {num_issues_to_process} unlabelled open issues from"
-        f" the {OWNER}/{REPO} repository. For each issue, provide a summary,"
-        " recommend a label with justification, and then use the"
-        " 'add_label_to_issue' tool to apply the recommended label directly."
+        f"A new GitHub issue #{issue_number} has been opened or"
+        f' reopened. Title: "{issue_title}"\nBody:'
+        f' "{issue_body}"\n\nBased on the rules, recommend an'
+        " appropriate label and its justification."
+        " Then, use the 'add_label_to_issue' tool to apply the label "
+        "directly to this issue. Only label it, do not"
+        " process any other issues."
     )
-    await run_agent_prompt(session_11, prompt)
+  else:
+    print(f"EVENT: Processing batch of issues (event: {EVENT_NAME}).")
+    issue_count = parse_number_string(ISSUE_COUNT_TO_PROCESS, default_value=3)
+    prompt = f"Please triage the most recent {issue_count} issues."
+
+  response = await call_agent_async(runner, USER_ID, session.id, prompt)
+  print(f"<<<< Agent Final Output: {response}\n")
 
 
 if __name__ == "__main__":
   start_time = time.time()
   print(
-      "Script start time:",
-      time.strftime("%Y-%m-%d %H:%M:%S", time.gmtime(start_time)),
+      f"Start triaging {OWNER}/{REPO} issues at"
+      f" {time.strftime('%Y-%m-%d %H:%M:%S', time.gmtime(start_time))}"
   )
-  print("------------------------------------")
+  print("-" * 80)
   asyncio.run(main())
+  print("-" * 80)
   end_time = time.time()
-  print("------------------------------------")
   print(
-      "Script end time:",
-      time.strftime("%Y-%m-%d %H:%M:%S", time.gmtime(end_time)),
+      "Triaging finished at"
+      f" {time.strftime('%Y-%m-%d %H:%M:%S', time.gmtime(end_time))}",
   )
   print("Total script execution time:", f"{end_time - start_time:.2f} seconds")

contributing/samples/adk_triaging_agent/settings.py

@@ -0,0 +1,36 @@
+# 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.
+
+import os
+
+from dotenv import load_dotenv
+
+load_dotenv(override=True)
+
+GITHUB_BASE_URL = "https://api.github.com"
+
+GITHUB_TOKEN = os.getenv("GITHUB_TOKEN")
+if not GITHUB_TOKEN:
+  raise ValueError("GITHUB_TOKEN environment variable not set")
+
+OWNER = os.getenv("OWNER", "google")
+REPO = os.getenv("REPO", "adk-python")
+BOT_LABEL = os.getenv("BOT_LABEL", "bot_triaged")
+EVENT_NAME = os.getenv("EVENT_NAME")
+ISSUE_NUMBER = os.getenv("ISSUE_NUMBER")
+ISSUE_TITLE = os.getenv("ISSUE_TITLE")
+ISSUE_BODY = os.getenv("ISSUE_BODY")
+ISSUE_COUNT_TO_PROCESS = os.getenv("ISSUE_COUNT_TO_PROCESS")
+
+IS_INTERACTIVE = os.environ.get("INTERACTIVE", "1").lower() in ["true", "1"]

contributing/samples/adk_triaging_agent/utils.py

@@ -0,0 +1,55 @@
+# 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 typing import Any
+
+from adk_triaging_agent.settings import GITHUB_TOKEN
+import requests
+
+headers = {
+    "Authorization": f"token {GITHUB_TOKEN}",
+    "Accept": "application/vnd.github.v3+json",
+}
+
+
+def get_request(
+    url: str, params: dict[str, Any] | None = None
+) -> dict[str, Any]:
+  if params is None:
+    params = {}
+  response = requests.get(url, headers=headers, params=params, timeout=60)
+  response.raise_for_status()
+  return response.json()
+
+
+def post_request(url: str, payload: Any) -> dict[str, Any]:
+  response = requests.post(url, headers=headers, json=payload, timeout=60)
+  response.raise_for_status()
+  return response.json()
+
+
+def error_response(error_message: str) -> dict[str, Any]:
+  return {"status": "error", "message": error_message}
+
+
+def parse_number_string(number_str: str, default_value: int = 0) -> int:
+  """Parse a number from the given string."""
+  try:
+    return int(number_str)
+  except ValueError:
+    print(
+        f"Warning: Invalid number string: {number_str}. Defaulting to"
+        f" {default_value}."
+    )
+    return default_value