Google Drive

pip install httpx geniffy

uv add httpx geniffy

import httpx
from geniffy import BadRequestError, Geniffy, NotFoundError

geniffy = Geniffy()                               # reads GENIFFY_API_KEY
LABELS = {"channel": "google-drive"}
FIELDS = "id, name, mimeType, size, trashed"
XLSX = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
EXPORT = {                                        # Google's own files, as Drive exports them: every tab of a Sheet
    "application/vnd.google-apps.document": ("text/plain", ".txt"),
    "application/vnd.google-apps.presentation": ("text/plain", ".txt"),
    "application/vnd.google-apps.spreadsheet": (XLSX, ".xlsx"),
}
READS = (".pdf", ".docx", ".pptx", ".xlsx", ".txt", ".md", ".csv", ".html")    # and these, as they are


def drive(token: str) -> httpx.Client:
    """One user's Drive, with the access token from their Google sign-in."""
    return httpx.Client(base_url="https://www.googleapis.com/drive/v3", timeout=60,
                        headers={"Authorization": f"Bearer {token}"})


def keep(mem, api: httpx.Client, file: dict) -> bool:
    """Add one file under its Drive id. False when it holds nothing Geniffy can read."""
    if file["mimeType"] in EXPORT:
        kind, ext = EXPORT[file["mimeType"]]
        got, name = api.get(f"/files/{file['id']}/export", params={"mimeType": kind}), file["name"] + ext
    elif file["name"].lower().endswith(READS) and int(file.get("size", 0)) <= 25 * 1024 * 1024:
        got, name = api.get(f"/files/{file['id']}", params={"alt": "media"}), file["name"]
    else:
        return False                              # a folder, an image, a video, or over 25 MB
    try:
        mem.memories.add_file(got.raise_for_status().content, filename=name, title=file["name"],
                              external_id=f"drive:{file['id']}", labels=LABELS)
    except BadRequestError:                       # empty, or a scan with no text in it
        return False
    return True


def forget(mem, file_id: str) -> None:
    try:
        mem.sources.delete(external_id=f"drive:{file_id}")
    except NotFoundError:                         # never added: a folder, an image
        pass


def read_all(mem, api: httpx.Client) -> None:
    """Everything in the user's Drive; and what memory holds from Drive that is no longer there goes."""
    seen, params = set(), {"q": "trashed = false", "pageSize": 100, "fields": f"nextPageToken, files({FIELDS})"}
    while True:
        out = api.get("/files", params=params).raise_for_status().json()
        seen |= {f"drive:{file['id']}" for file in out["files"] if keep(mem, api, file)}
        if "nextPageToken" not in out:
            break
        params["pageToken"] = out["nextPageToken"]
    mem.sources.delete_labelled(LABELS, keep=seen)


def sync(user_id: str, token: str, page_token: str | None = None) -> str:
    """Bring one user's Drive into their memory. Returns the token to pass next time: pass None the
    first time and everything is read; after that, only what changed since the last run is fetched."""
    mem = geniffy.space(f"user_{user_id}")
    with drive(token) as api:
        if page_token is None:
            start = api.get("/changes/startPageToken").raise_for_status().json()["startPageToken"]
            read_all(mem, api)
            return start
        params = {"pageToken": page_token, "pageSize": 100,
                  "fields": f"nextPageToken, newStartPageToken, changes(changeType, fileId, removed, file({FIELDS}))"}
        while True:
            out = api.get("/changes", params=params).raise_for_status().json()
            for change in out["changes"]:
                if change.get("changeType") == "drive":
                    continue                      # a shared drive itself changed, not a file in it
                file = change.get("file")
                if change.get("removed") or not file or file.get("trashed") or not keep(mem, api, file):
                    forget(mem, change["fileId"])
            if "newStartPageToken" in out:
                return out["newStartPageToken"]
            params["pageToken"] = out["nextPageToken"]


def disconnect(user_id: str) -> int:
    """The user disconnected Drive: everything that came from it goes. Returns how many files."""
    return geniffy.space(f"user_{user_id}").sources.delete_labelled(LABELS)