mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-06-06 20:15:17 +02:00
refactor(agents): extract shared Google OAuth helper from gmail connector
build_credentials/get_token_encryption are Google-OAuth helpers used by both the Gmail and Calendar connector tools. They lived inside gmail/tools/_helpers.py, forcing calendar -> gmail coupling. Move them to a neutral connector-level module (connectors/google_auth.py); gmail/_helpers.py re-exports them under the legacy private names so existing gmail tools are untouched, and calendar now imports the shared module directly.
This commit is contained in:
parent
0081b627e9
commit
490bb3c5c5
3 changed files with 76 additions and 52 deletions
|
|
@ -5,8 +5,8 @@ from langchain_core.tools import tool
|
|||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy.future import select
|
||||
|
||||
from app.agents.chat.multi_agent_chat.subagents.connectors.gmail.tools._helpers import (
|
||||
_build_credentials,
|
||||
from app.agents.chat.multi_agent_chat.subagents.connectors.google_auth import (
|
||||
build_credentials as _build_credentials,
|
||||
)
|
||||
from app.db import SearchSourceConnector, SearchSourceConnectorType
|
||||
|
||||
|
|
|
|||
|
|
@ -1,61 +1,26 @@
|
|||
"""Shared helpers for Gmail connector tools.
|
||||
"""Gmail-specific helpers for the Gmail connector tools.
|
||||
|
||||
Credential construction (``_build_credentials``) is also reused by the
|
||||
Calendar connector tools, since both are Google OAuth backed.
|
||||
Google OAuth credential construction lives in
|
||||
``app.agents.chat.multi_agent_chat.subagents.connectors.google_auth`` (shared
|
||||
with the Calendar connector). It is re-exported here under the legacy private
|
||||
names so the existing Gmail tools keep importing it from this module.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from app.db import SearchSourceConnector
|
||||
from app.agents.chat.multi_agent_chat.subagents.connectors.google_auth import (
|
||||
build_credentials as _build_credentials,
|
||||
get_token_encryption as _get_token_encryption,
|
||||
)
|
||||
|
||||
_token_encryption_cache: object | None = None
|
||||
|
||||
|
||||
def _get_token_encryption():
|
||||
global _token_encryption_cache
|
||||
if _token_encryption_cache is None:
|
||||
from app.config import config
|
||||
from app.utils.oauth_security import TokenEncryption
|
||||
|
||||
if not config.SECRET_KEY:
|
||||
raise RuntimeError("SECRET_KEY not configured for token decryption.")
|
||||
_token_encryption_cache = TokenEncryption(config.SECRET_KEY)
|
||||
return _token_encryption_cache
|
||||
|
||||
|
||||
def _build_credentials(connector: SearchSourceConnector):
|
||||
"""Build Google OAuth Credentials from a connector's stored config.
|
||||
|
||||
Handles both native OAuth connectors (with encrypted tokens) and
|
||||
Composio-backed connectors. Shared by Gmail and Calendar tools.
|
||||
"""
|
||||
from app.utils.google_credentials import COMPOSIO_GOOGLE_CONNECTOR_TYPES
|
||||
|
||||
if connector.connector_type in COMPOSIO_GOOGLE_CONNECTOR_TYPES:
|
||||
raise ValueError("Composio connectors must use Composio tool execution.")
|
||||
|
||||
from google.oauth2.credentials import Credentials
|
||||
|
||||
cfg = dict(connector.config)
|
||||
if cfg.get("_token_encrypted"):
|
||||
enc = _get_token_encryption()
|
||||
for key in ("token", "refresh_token", "client_secret"):
|
||||
if cfg.get(key):
|
||||
cfg[key] = enc.decrypt_token(cfg[key])
|
||||
|
||||
exp = (cfg.get("expiry") or "").replace("Z", "")
|
||||
return Credentials(
|
||||
token=cfg.get("token"),
|
||||
refresh_token=cfg.get("refresh_token"),
|
||||
token_uri=cfg.get("token_uri"),
|
||||
client_id=cfg.get("client_id"),
|
||||
client_secret=cfg.get("client_secret"),
|
||||
scopes=cfg.get("scopes", []),
|
||||
expiry=datetime.fromisoformat(exp) if exp else None,
|
||||
)
|
||||
__all__ = [
|
||||
"_build_credentials",
|
||||
"_format_gmail_summary",
|
||||
"_get_token_encryption",
|
||||
"_gmail_headers",
|
||||
]
|
||||
|
||||
|
||||
def _gmail_headers(message: dict[str, Any]) -> dict[str, str]:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,59 @@
|
|||
"""Google OAuth credential construction shared across Google connectors.
|
||||
|
||||
Both the Gmail and Calendar connector tools are Google OAuth backed and build
|
||||
``google.oauth2.credentials.Credentials`` from a stored ``SearchSourceConnector``
|
||||
the same way. This module is the single owner of that logic so neither connector
|
||||
has to import the other.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
from app.db import SearchSourceConnector
|
||||
|
||||
_token_encryption_cache: object | None = None
|
||||
|
||||
|
||||
def get_token_encryption():
|
||||
global _token_encryption_cache
|
||||
if _token_encryption_cache is None:
|
||||
from app.config import config
|
||||
from app.utils.oauth_security import TokenEncryption
|
||||
|
||||
if not config.SECRET_KEY:
|
||||
raise RuntimeError("SECRET_KEY not configured for token decryption.")
|
||||
_token_encryption_cache = TokenEncryption(config.SECRET_KEY)
|
||||
return _token_encryption_cache
|
||||
|
||||
|
||||
def build_credentials(connector: SearchSourceConnector):
|
||||
"""Build Google OAuth Credentials from a connector's stored config.
|
||||
|
||||
Handles both native OAuth connectors (with encrypted tokens) and
|
||||
Composio-backed connectors. Shared by Gmail and Calendar tools.
|
||||
"""
|
||||
from app.utils.google_credentials import COMPOSIO_GOOGLE_CONNECTOR_TYPES
|
||||
|
||||
if connector.connector_type in COMPOSIO_GOOGLE_CONNECTOR_TYPES:
|
||||
raise ValueError("Composio connectors must use Composio tool execution.")
|
||||
|
||||
from google.oauth2.credentials import Credentials
|
||||
|
||||
cfg = dict(connector.config)
|
||||
if cfg.get("_token_encrypted"):
|
||||
enc = get_token_encryption()
|
||||
for key in ("token", "refresh_token", "client_secret"):
|
||||
if cfg.get(key):
|
||||
cfg[key] = enc.decrypt_token(cfg[key])
|
||||
|
||||
exp = (cfg.get("expiry") or "").replace("Z", "")
|
||||
return Credentials(
|
||||
token=cfg.get("token"),
|
||||
refresh_token=cfg.get("refresh_token"),
|
||||
token_uri=cfg.get("token_uri"),
|
||||
client_id=cfg.get("client_id"),
|
||||
client_secret=cfg.get("client_secret"),
|
||||
scopes=cfg.get("scopes", []),
|
||||
expiry=datetime.fromisoformat(exp) if exp else None,
|
||||
)
|
||||
Loading…
Add table
Add a link
Reference in a new issue