mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
chore: refactor the ADK Triaging Agent to make the code easier to read
PiperOrigin-RevId: 776763061
This commit is contained in:
Xuan Yang
committed by
Copybara-Service
Copybara-Service
parent
ffa9b361db
commit
b6c7b5b64f
@@ -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
|
||||
|
||||
@@ -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],
|
||||
)
|
||||
|
||||
@@ -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 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)]
|
||||
)
|
||||
|
||||
final_response_text = ""
|
||||
async for event in runner.run_async(
|
||||
user_id=user_id,
|
||||
session_id=session_id,
|
||||
new_message=content,
|
||||
run_config=RunConfig(save_input_blobs_as_artifacts=False),
|
||||
):
|
||||
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_response_text += event.content.parts[0].text
|
||||
|
||||
return final_response_text
|
||||
|
||||
|
||||
async def main():
|
||||
app_name = "triage_app"
|
||||
user_id_1 = "triage_user"
|
||||
runner = InMemoryRunner(
|
||||
agent=agent.root_agent,
|
||||
app_name=app_name,
|
||||
app_name=APP_NAME,
|
||||
)
|
||||
session_11 = await runner.session_service.create_session(
|
||||
app_name=app_name, user_id=user_id_1
|
||||
session = await runner.session_service.create_session(
|
||||
user_id=USER_ID,
|
||||
app_name=APP_NAME,
|
||||
)
|
||||
|
||||
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 = []
|
||||
async for event in runner.run_async(
|
||||
user_id=user_id_1,
|
||||
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:
|
||||
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")
|
||||
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
|
||||
|
||||
event_name = os.getenv("EVENT_NAME")
|
||||
issue_number_str = os.getenv("ISSUE_NUMBER")
|
||||
|
||||
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']}."
|
||||
)
|
||||
await run_agent_prompt(session_11, prompt)
|
||||
else:
|
||||
print(
|
||||
f"No unlabelled issue details found for #{issue_number} or an error"
|
||||
" occurred. Skipping agent interaction."
|
||||
)
|
||||
|
||||
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
|
||||
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
|
||||
|
||||
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")
|
||||
|
||||
@@ -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"]
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user