Get draft replies to emails that need an answer

from google.oauth2.credentials import Credentials
from googleapiclient.discovery import build
import base64
from email.message import EmailMessage
from clickruntime import on_trigger, report_fetched, credentials, state, ai

@on_trigger
def run(event):
    gmail = build("gmail", "v1", credentials=Credentials(credentials.connection("google", "work")), cache_discovery=False)
    msgs = gmail.users().messages().list(userId="me", q="in:inbox category:primary newer_than:1d -from:me").execute().get("messages", [])
    done = set(state.get("done") or [])
    new = [m for m in msgs if m["id"] not in done]
    report_fetched(len(new))
    if not new:
        return {"summary": "Nothing new"}
    drafted = 0
    for m in new:
        msg = gmail.users().messages().get(userId="me", id=m["id"], format="metadata",
                                           metadataHeaders=["From", "Subject", "Message-ID"]).execute()
        h = {x["name"]: x["value"] for x in msg["payload"]["headers"]}
        reply = ai.generate(f"If this email needs a personal reply, write a short, friendly one in first person. "
                            f"Otherwise answer NO_REPLY.\nFrom: {h['From']}\nSubject: {h['Subject']}\n{msg['snippet']}",
                            max_output_tokens=300).text.strip()
        done.add(m["id"])
        if reply == "NO_REPLY":
            continue
        draft = EmailMessage()
        draft["To"], draft["Subject"] = h["From"], "Re: " + h["Subject"]
        draft["In-Reply-To"] = draft["References"] = h.get("Message-ID", "")
        draft.set_content(reply)
        raw = base64.urlsafe_b64encode(draft.as_bytes()).decode()
        gmail.users().drafts().create(userId="me", body={"message": {"raw": raw, "threadId": msg["threadId"]}}).execute()
        drafted += 1
    state.set("done", sorted(done)[-1000:])
    return {"drafted": drafted}