chore: Create agent.py

add a ADK triaging agent to help triaging issues. This is an MVP. We can iterate further.
2026-03-30 10:57:20 -07:00 · 2025-06-04 21:25:02 -07:00
parent 78bfce4d89
commit cf689fd528
1 changed files with 114 additions and 0 deletions
@@ -0,0 +1,114 @@
+# 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 random
+import os
+import requests
+import time
+from google.adk import Agent
+from google.adk.tools.tool_context import ToolContext
+from google.genai import types
+
+
+# Read the PAT from the environment variable
+GITHUB_TOKEN = os.getenv("GITHUB_TOKEN")  # Ensure you've set this in your shell
+if not GITHUB_TOKEN:
+    raise ValueError("GITHUB_TOKEN environment variable not set")
+
+# Repository information
+OWNER = "google"
+REPO = "adk-python"
+
+# Base URL for the GitHub API
+BASE_URL = "https://api.github.com"
+
+# Headers including the Authorization header
+headers = {
+    "Authorization": f"token {GITHUB_TOKEN}",
+    "Accept": "application/vnd.github.v3+json"
+}
+
+
+def list_issues(per_page:int):
+    """
+    Generator to list all issues for the repository by handling pagination.
+
+    Args:
+      per_page: number of pages to return per page.
+
+    """
+    state="open"
+    # only process the 1st page for testing for now
+    page = 1
+    results = []
+    url = f"{BASE_URL}/repos/{OWNER}/{REPO}/issues"  # :contentReference[oaicite:16]{index=16}
+    # Warning: let's only handle max 10 issues at a time to avoid bad results
+    params = {
+        "state": state,
+        "per_page": per_page,
+        "page": page
+    }
+    response = requests.get(url, headers=headers, params=params)
+    response.raise_for_status()  # :contentReference[oaicite:17]{index=17}
+    issues = response.json()
+    if not issues:
+        return []
+    for issue in issues:
+        # Skip pull requests (issues API returns PRs as well)
+        if "pull_request" in issue:
+            continue
+        results.append(issue)
+    return results
+
+def add_label_to_issue(issue_number:str, label:str):
+    """
+    Add the specified label to the given issue number.
+
+    Args:
+      issue_number: issue number of the Github issue, in string foramt.
+      label: label to assign
+    """
+    url = f"{BASE_URL}/repos/{OWNER}/{REPO}/issues/{issue_number}/labels"
+    payload = [label]
+    response = requests.post(url, headers=headers, json=payload)
+    response.raise_for_status()
+    return response.json()
+
+root_agent = Agent(
+    model='gemini-2.5-pro-preview-05-06',
+    name='adk_triaging_assistant',
+    description=(
+        'Triage ADK issues.'
+    ),
+    instruction="""
+      You are a Github adk-python repo triaging bot. You will help get issues, and label them.
+      Here are the rules for labeling:
+      1. If it's about session, memory services, label it with "services"
+      2. If it's about UI/web, label it with "question"
+      3. If it's related to tools, label it with "tools"
+      4. In other cases, label it with "core".
+    """,
+    tools=[
+        list_issues,
+        add_label_to_issue,
+    ],
+    generate_content_config=types.GenerateContentConfig(
+        safety_settings=[
+            types.SafetySetting(  # avoid false alarm about rolling dice.
+                category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
+                threshold=types.HarmBlockThreshold.OFF,
+            ),
+        ]
+    ),
+)