agentgraph-connector-google 0.5.0__tar.gz

agentgraph_connector_google-0.5.0/PKG-INFO

@@ -0,0 +1,9 @@
+Metadata-Version: 2.4
+Name: agentgraph-connector-google
+Version: 0.5.0
+Summary: Google Docs, Sheets, Gmail, and Drive connectors for AgentGraph
+Requires-Python: >=3.12
+Requires-Dist: agentgraph-server<0.6,>=0.5.0
+Requires-Dist: google-api-python-client>=2.193.0
+Requires-Dist: google-auth>=2.49.1
+Requires-Dist: google-auth-oauthlib>=1.3.0

agentgraph_connector_google-0.5.0/agentgraph_connector_google/__init__.py

@@ -0,0 +1,11 @@
+from agentgraph_connector_google.gdocs import GoogleDocsConnector
+from agentgraph_connector_google.gdrive import DriveChangesConnector
+from agentgraph_connector_google.gmail import GmailConnector
+from agentgraph_connector_google.gsheets import GoogleSheetsConnector
+
+__all__ = [
+    "GoogleDocsConnector",
+    "GoogleSheetsConnector",
+    "GmailConnector",
+    "DriveChangesConnector",
+]

agentgraph_connector_google-0.5.0/agentgraph_connector_google/auth.py

@@ -0,0 +1,195 @@
+"""Google OAuth2 flow for all Google connectors (Docs, Sheets, Drive, Gmail)."""
+
+# pyright: reportUnknownMemberType=false, reportUnknownVariableType=false
+# pyright: reportUnknownArgumentType=false
+
+from __future__ import annotations
+
+import argparse
+import socket
+import webbrowser
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from typing import NoReturn
+from urllib.parse import parse_qs, urlparse
+
+import typer
+
+from agentgraph.auth.credentials import (
+    load_platform_accounts,
+    save_platform,
+    upsert_platform_account,
+)
+from agentgraph_connector_google.provider import (
+    GOOGLE_OAUTH_CLIENT_ID,
+    GOOGLE_OAUTH_CLIENT_SECRET,
+    GOOGLE_SCOPES,
+    GOOGLE_TOKEN_URI,
+    GoogleCredentials,
+    verify_google_auth,
+)
+
+
+class _AuthArgumentParser(argparse.ArgumentParser):
+    def error(self, message: str) -> NoReturn:
+        raise ValueError(message)
+
+
+def _validate_args(args: list[str]) -> None:
+    parser = _AuthArgumentParser(add_help=False, prog="agentgraph auth google")
+    parser.parse_args(args)
+
+
+def _find_free_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("", 0))
+        return s.getsockname()[1]
+
+
+def run_oauth_flow(
+    account_id: str | None = None,
+    add: bool = False,
+    args: list[str] | None = None,
+) -> None:
+    """Interactive OAuth2 browser flow. Stores credentials on completion."""
+    _validate_args(args or [])
+
+    existing_accounts = load_platform_accounts("google")
+    existing_data = None
+    if account_id is not None:
+        existing_data = next(
+            (item for item in existing_accounts if item.get("account_id") == account_id), None
+        )
+    elif existing_accounts and not add:
+        existing_data = existing_accounts[0]
+    existing = GoogleCredentials(**existing_data) if existing_data else None
+
+    if existing:
+        # Probe the saved refresh token before forcing the user back through
+        # the browser. If it still works, offer to skip.
+        raw_account_id = existing_data.get("account_id") if existing_data else None
+        verify_account_id = str(raw_account_id) if raw_account_id is not None else None
+        status, detail = (
+            verify_google_auth(verify_account_id)
+            if verify_account_id is not None
+            else verify_google_auth()
+        )
+        if status == "ok":
+            typer.echo(f"\nGoogle is already authenticated as {detail}.")
+            if not typer.confirm(
+                "Re-authenticate anyway (e.g. to switch accounts or grant new scopes)?",
+                default=False,
+            ):
+                typer.echo("Keeping existing credentials.")
+                return
+        else:
+            typer.echo(f"\nGoogle credentials need re-authentication: {detail or status}.")
+            typer.echo("Re-opening browser consent using AgentGraph's packaged OAuth client.")
+
+        typer.echo(
+            f"Re-authenticating as {existing.user_email or 'existing account'} with updated scopes."
+        )
+
+    client_id = GOOGLE_OAUTH_CLIENT_ID
+
+    port = _find_free_port()
+    redirect_uri = f"http://localhost:{port}"
+
+    from google_auth_oauthlib.flow import Flow  # type: ignore[import-untyped]
+
+    client_config = {
+        "installed": {
+            "client_id": client_id,
+            "client_secret": GOOGLE_OAUTH_CLIENT_SECRET,
+            "redirect_uris": [redirect_uri],
+            "auth_uri": "https://accounts.google.com/o/oauth2/auth",
+            "token_uri": GOOGLE_TOKEN_URI,
+        }
+    }
+
+    flow: Flow = Flow.from_client_config(
+        client_config,
+        scopes=GOOGLE_SCOPES,
+        redirect_uri=redirect_uri,
+    )
+
+    auth_url, _ = flow.authorization_url(
+        access_type="offline",
+        prompt="consent",
+    )
+
+    typer.echo("\nOpening browser for Google authorization...")
+    typer.echo(f"If the browser doesn't open, visit:\n  {auth_url}\n")
+    webbrowser.open(auth_url)
+
+    auth_code = _wait_for_callback(port)
+
+    flow.fetch_token(code=auth_code)
+    token = flow.credentials
+
+    user_email: str | None = None
+    display_name: str | None = None
+    try:
+        import requests  # type: ignore[import-untyped]
+
+        resp = requests.get(
+            "https://www.googleapis.com/oauth2/v2/userinfo",
+            headers={"Authorization": f"Bearer {token.token}"},
+            timeout=10,
+        )
+        if resp.ok:
+            info = resp.json()
+            user_email = info.get("email")
+            display_name = info.get("name")
+    except Exception:
+        pass
+
+    creds = GoogleCredentials(
+        client_id=client_id,
+        access_token=token.token or "",
+        refresh_token=token.refresh_token or "",
+        token_expiry=token.expiry,
+        user_email=user_email,
+        display_name=display_name,
+    )
+    resolved_account_id = account_id or (
+        user_email.lower() if user_email else f"google:{len(existing_accounts) + 1}"
+    )
+    if not add and len(existing_accounts) <= 1 and account_id is None:
+        save_platform(
+            "google", {**creds.model_dump(mode="json"), "account_id": resolved_account_id}
+        )
+    else:
+        upsert_platform_account("google", resolved_account_id, creds, make_default=True)
+    from agentgraph.config import CREDENTIALS_FILE
+
+    msg = f"Google credentials saved to {CREDENTIALS_FILE}"
+    if user_email:
+        msg += f" (authenticated as {user_email})"
+    typer.echo(msg)
+
+
+def _wait_for_callback(port: int) -> str:
+    """Spin up a one-shot HTTP server to capture the OAuth redirect code."""
+    auth_code: list[str] = []
+
+    class _Handler(BaseHTTPRequestHandler):
+        def do_GET(self) -> None:
+            params = parse_qs(urlparse(self.path).query)
+            code = params.get("code", [""])[0]
+            auth_code.append(code)
+            self.send_response(200)
+            self.end_headers()
+            self.wfile.write(
+                b"<html><body><h2>Authorization complete. You can close this tab.</h2></body></html>"
+            )
+
+        def log_message(self, format: str, *args: object) -> None:
+            pass
+
+    server = HTTPServer(("localhost", port), _Handler)
+    server.handle_request()
+    server.server_close()
+
+    if not auth_code or not auth_code[0]:
+        raise typer.Exit(code=1)
+    return auth_code[0]

agentgraph_connector_google-0.5.0/agentgraph_connector_google/gdocs.py

@@ -0,0 +1,253 @@
+"""Google Docs connector."""
+
+# pyright: reportUnknownMemberType=false, reportUnknownVariableType=false
+# pyright: reportUnknownArgumentType=false, reportAttributeAccessIssue=false
+
+from __future__ import annotations
+
+import logging
+import re
+from datetime import UTC, datetime
+from typing import Any
+
+from googleapiclient.discovery import build  # type: ignore[import-untyped]
+from googleapiclient.errors import HttpError  # type: ignore[import-untyped]
+
+from agentgraph.connectors.base import (
+    BaseConnector,
+    ConnectorAccount,
+    EdgeRecord,
+    EntityBatch,
+    EntityRecord,
+    FetchPolicy,
+    PersonRecord,
+    ResourceType,
+    SourceReference,
+)
+from agentgraph.graph.upsert import upsert_batch
+from agentgraph_connector_google.provider import (
+    get_credentials as google_credentials,
+)
+from agentgraph_connector_google.provider import (
+    get_user_email,
+    list_google_accounts,
+    verify_google_auth,
+)
+
+logger = logging.getLogger(__name__)
+
+_GDOCS_URL_RE = re.compile(r"https://docs\.google\.com/document/d/(?P<doc_id>[a-zA-Z0-9_-]+)")
+
+# Staleness: re-fetch if doc hasn't been synced in the last 15 minutes
+_STALE_AFTER = 15 * 60
+
+
+def _build_drive_service(account_id: str | None = None) -> Any:
+    return build("drive", "v3", credentials=google_credentials(account_id))
+
+
+def _build_drive_service_for(account_id: str | None) -> Any:
+    return _build_drive_service(account_id) if account_id is not None else _build_drive_service()
+
+
+def _web_url(doc_id: str) -> str:
+    return f"https://docs.google.com/document/d/{doc_id}/view"
+
+
+def _download_url(doc_id: str) -> str:
+    return f"https://www.googleapis.com/drive/v3/files/{doc_id}/export?mimeType=text/html"
+
+
+def _metadata(
+    doc_id: str, account_id: str | None = None
+) -> dict[str, str | int | float | bool | None]:
+    meta: dict[str, str | int | float | bool | None] = {
+        "web_url": _web_url(doc_id),
+        "download_url": _download_url(doc_id),
+        "content_type": "text/html",
+        "mime_type": "text/html",
+    }
+    if account_id:
+        meta["account_id"] = account_id
+    return meta
+
+
+def _export_as_html(drive_service: Any, doc_id: str) -> str:
+    """Export a Google Doc as HTML via Drive."""
+    html: bytes = (
+        drive_service.files()
+        .export(
+            fileId=doc_id,
+            mimeType="text/html",
+        )
+        .execute()
+    )
+    return html.decode("utf-8", errors="replace")
+
+
+class GoogleDocsConnector(BaseConnector):
+    source = "gdocs"
+    fetch_policy = FetchPolicy(stale_after_seconds=_STALE_AFTER)
+    url_patterns = ["https://docs.google.com/document/*"]
+    auth_label = "google"
+    auth_description = (
+        "Google Docs: Document entities with full HTML body content and owner authorship."
+    )
+    onboard_prompt = "Set up Google?"
+
+    @classmethod
+    def run_auth_flow(
+        cls,
+        account_id: str | None = None,
+        add: bool = False,
+        args: list[str] | None = None,
+    ) -> None:
+        from agentgraph_connector_google.auth import run_oauth_flow
+
+        run_oauth_flow(account_id=account_id, add=add, args=args)
+
+    @classmethod
+    def get_authenticated_user(cls) -> str | None:
+        return get_user_email()
+
+    @classmethod
+    def list_accounts(cls) -> list[ConnectorAccount]:
+        return [
+            ConnectorAccount(
+                account_id=str(account["account_id"]),
+                label=str(account["label"]),
+                auth_group=cls.auth_label or cls.source,
+                source=cls.source,
+                user_id=account.get("email"),
+                email=account.get("email"),
+            )
+            for account in list_google_accounts()
+        ]
+
+    @classmethod
+    async def verify_auth(cls, account_id: str | None = None) -> tuple[str, str | None]:
+        import asyncio
+
+        return await asyncio.to_thread(verify_google_auth, account_id)
+
+    @classmethod
+    def current_user_ids(cls) -> list[str]:
+        return [str(account["email"]) for account in list_google_accounts() if account.get("email")]
+
+    def can_handle(self, url: str) -> bool:
+        return self.resolve_url(url) is not None
+
+    def resolve_url(self, url: str) -> SourceReference | None:
+        match = _GDOCS_URL_RE.match(url)
+        if match is None:
+            return None
+        return SourceReference(
+            source=self.source,
+            resource_type="document",
+            resource_id=match.group("doc_id"),
+        )
+
+    def entity_url(self, platform_entity_id: str) -> str | None:
+        return f"https://docs.google.com/document/d/{platform_entity_id}"
+
+    async def fetch(
+        self,
+        resource_type: ResourceType,
+        resource_id: str,
+        meta: dict[str, str] | None = None,
+        account_id: str | None = None,
+    ) -> EntityBatch:
+        last_sync = await self.last_synced_at(resource_id)
+        decision = self.fetch_policy.decide(last_sync)
+
+        if decision == FetchPolicy.FRESH:
+            logger.debug("gdocs/%s is fresh — updating last_accessed only", resource_id)
+            await _touch_last_accessed(resource_id)
+            return EntityBatch()
+
+        logger.info("Fetching Google Doc %s (policy=%s)", resource_id, decision)
+        selected_account_id = account_id or ((meta or {}).get("account_id") if meta else None)
+        batch = await _fetch_doc(resource_id, account_id=selected_account_id)
+        await upsert_batch(batch)
+        return batch
+
+    async def download(
+        self,
+        resource_type: ResourceType,
+        resource_id: str,
+        output_path: str | None = None,
+    ) -> dict[str, Any]:
+        from agentgraph_connector_google.gdrive import download_drive_file
+
+        return await download_drive_file(resource_id, output_path)
+
+
+async def _touch_last_accessed(doc_id: str) -> None:
+    from agentgraph.core.context import get_backend
+
+    await get_backend().touch_last_accessed("gdocs", doc_id)
+
+
+async def _fetch_doc(doc_id: str, account_id: str | None = None) -> EntityBatch:
+    import asyncio
+
+    loop = asyncio.get_event_loop()
+    drive_service = await loop.run_in_executor(None, _build_drive_service_for, account_id)
+
+    try:
+        file_meta: dict[str, Any] = await loop.run_in_executor(
+            None,
+            lambda: (
+                drive_service.files()
+                .get(
+                    fileId=doc_id,
+                    fields="name,owners",
+                )
+                .execute()
+            ),
+        )
+    except HttpError as exc:
+        if exc.resp.status == 404:
+            raise ValueError(f"Google Doc not found or not accessible: {doc_id}") from exc
+        raise
+
+    title: str = file_meta.get("name", "")
+    content = await loop.run_in_executor(None, _export_as_html, drive_service, doc_id)
+
+    persons: list[PersonRecord] = []
+    edges: list[EdgeRecord] = []
+
+    for owner in file_meta.get("owners", []):
+        email: str = owner.get("emailAddress", "")
+        name: str = owner.get("displayName", "")
+        if email:
+            persons.append(
+                PersonRecord(
+                    platform="gdocs",
+                    platform_user_id=email,
+                    canonical_email=email,
+                    display_name=name or None,
+                )
+            )
+            edges.append(
+                EdgeRecord(
+                    edge_type="authored",
+                    source_platform_user_id=email,
+                    target_platform_entity_id=doc_id,
+                    platform="gdocs",
+                )
+            )
+
+    entity = EntityRecord(
+        entity_type="Document",
+        platform="gdocs",
+        platform_entity_id=doc_id,
+        title=title,
+        content=content,
+        metadata=_metadata(doc_id, account_id),
+        updated_at=datetime.now(UTC),
+    )
+
+    batch = EntityBatch(entities=[entity], persons=persons, edges=edges)
+    batch.add_stubs_from(entity)
+    return batch