1058 lines
40 KiB
Python
1058 lines
40 KiB
Python
"""agy provider adapter — Google Antigravity accounts behind ModelProviderPort.
|
|
|
|
``agy`` is Google's Antigravity: a Google-AI-subscription-billed route to the
|
|
Gemini 3 / Claude 4.5 / GPT-OSS families. There is no public API key; instead
|
|
each *account* is a Google OAuth identity whose refresh token is exchanged for
|
|
short-lived access tokens. This module implements that protocol directly, so
|
|
the platform never shells out to the CLI binary:
|
|
|
|
* :class:`AgyAuthClient` — PKCE OAuth in two shapes: **listener mode** (a
|
|
loopback HTTP server captures the redirect; used when the operator's browser
|
|
runs on the same machine) and **paste mode** (the operator copies the
|
|
callback URL from the address bar and hands it back; the headless/remote
|
|
case — the page after redirect never finishes loading, and that is expected).
|
|
* :class:`AgyModelAdapter` — the ``ModelProviderPort`` implementation. Converts
|
|
our message format to the ``/v1internal:generateContent`` envelope, calls
|
|
Google's cloudcode-pa backend, parses the Gemini-shaped response.
|
|
* :class:`AgyAccountPool` state inside the adapter — round-robin across the
|
|
logged-in accounts with automatic 429 benching: a rate-limited account is
|
|
skipped until its cooldown expires, both in-process (fast) and in the
|
|
database (so background workers respect a cooldown set by the API).
|
|
|
|
Nothing here ever returns or logs a refresh token: :func:`load_agy_accounts`
|
|
hands the adapter decrypted material in memory only, and admin endpoints see
|
|
email/project/status, never ciphertext.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import secrets
|
|
import time
|
|
import uuid
|
|
from collections.abc import Awaitable, Callable
|
|
from dataclasses import dataclass, field
|
|
from datetime import UTC, datetime, timedelta
|
|
from typing import Any
|
|
from urllib.parse import parse_qs, urlencode, urlparse
|
|
|
|
import httpx
|
|
|
|
from app.application.ports import ChatMessage, ModelHealth, ModelResponse
|
|
from app.core.config import settings
|
|
from app.core.crypto import decrypt
|
|
from app.core.errors import ProviderUnavailableError, ValidationError
|
|
from app.domain.ai_control.models import AiModel, ModelCredential
|
|
from app.domain.base import utcnow
|
|
|
|
logger = logging.getLogger("agy.adapter")
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Protocol constants (the public Antigravity CLI OAuth application and the
|
|
# cloudcode-pa backend it talks to).
|
|
# ---------------------------------------------------------------------------
|
|
|
|
OAUTH_AUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth"
|
|
OAUTH_TOKEN_URL = "https://oauth2.googleapis.com/token"
|
|
USERINFO_URL = "https://www.googleapis.com/oauth2/v1/userinfo?alt=json"
|
|
|
|
#: ``daily-cloudcode-pa`` first: it carries the free daily quota and 429s more
|
|
#: gracefully. ``cloudcode-pa`` is the fallback.
|
|
BASE_URLS: tuple[str, ...] = (
|
|
"https://daily-cloudcode-pa.sandbox.googleapis.com",
|
|
"https://cloudcode-pa.googleapis.com",
|
|
)
|
|
PATH_GENERATE = "/v1internal:generateContent"
|
|
PATH_MODELS = "/v1internal:fetchAvailableModels"
|
|
PATH_LOAD_PROJECT = "/v1internal:loadCodeAssist"
|
|
|
|
OAUTH_SCOPES: tuple[str, ...] = (
|
|
"https://www.googleapis.com/auth/cloud-platform",
|
|
"https://www.googleapis.com/auth/userinfo.email",
|
|
"https://www.googleapis.com/auth/userinfo.profile",
|
|
"https://www.googleapis.com/auth/cclog",
|
|
"https://www.googleapis.com/auth/experimentsandconfigs",
|
|
)
|
|
|
|
#: Same all-OFF matrix the Antigravity client sends. Without it Google applies
|
|
#: its defaults and refuses coding prompts it considers dangerous.
|
|
SAFETY_SETTINGS: list[dict[str, str]] = [
|
|
{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "OFF"},
|
|
{"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "OFF"},
|
|
{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "OFF"},
|
|
{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "OFF"},
|
|
{"category": "HARM_CATEGORY_CIVIC_INTEGRITY", "threshold": "OFF"},
|
|
]
|
|
|
|
#: How long a started login stays redeemable. Two OAuth round-trips plus a
|
|
#: manual paste must fit inside it comfortably.
|
|
CHALLENGE_TTL_SECONDS = 600
|
|
|
|
#: agy ``reasoning_effort`` → thinking budget (tokens), mirroring the Antigravity
|
|
#: client. An unknown value falls back to ``medium`` server-side, and the UI
|
|
#: only ever offers these five.
|
|
EFFORT_THINKING_BUDGETS: dict[str, int] = {
|
|
"minimal": 1024,
|
|
"low": 4096,
|
|
"medium": 8192,
|
|
"high": 16384,
|
|
"xhigh": 32768,
|
|
}
|
|
|
|
#: Where thinking is honoured. A request must never carry ``thinkingConfig`` to
|
|
#: a model that does not support it — that is a hard 400 from the backend.
|
|
_THINKING_MARKERS: tuple[str, ...] = (
|
|
"thinking",
|
|
"gemini-3",
|
|
"gemini-2.5-pro",
|
|
"gemini-2.5-flash",
|
|
)
|
|
|
|
|
|
def _supports_thinking(model_name: str) -> bool:
|
|
lowered = model_name.lower()
|
|
if "image" in lowered:
|
|
return False
|
|
return any(marker in lowered for marker in _THINKING_MARKERS)
|
|
|
|
_StatusFn = Callable[..., Awaitable[None]]
|
|
_LoaderFn = Callable[[], Awaitable[list[ModelCredential]]]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Small helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _aware(value: datetime | None) -> datetime | None:
|
|
"""SQLite hands back naive datetimes; comparisons need aware ones."""
|
|
if value is None:
|
|
return None
|
|
if value.tzinfo is None:
|
|
return value.replace(tzinfo=UTC)
|
|
return value
|
|
|
|
|
|
def _pkce_pair() -> tuple[str, str]:
|
|
verifier = base64.urlsafe_b64encode(secrets.token_bytes(32)).rstrip(b"=").decode()
|
|
digest = hashlib.sha256(verifier.encode()).digest()
|
|
challenge = base64.urlsafe_b64encode(digest).rstrip(b"=").decode()
|
|
return verifier, challenge
|
|
|
|
|
|
def _random_session_id() -> str:
|
|
return f"-{secrets.randbelow(9_000_000_000_000_000_000)}"
|
|
|
|
|
|
def _fallback_project_id() -> str:
|
|
"""Ad-hoc project when ``loadCodeAssist`` did not report one."""
|
|
adjectives = ("useful", "bright", "swift", "calm", "bold")
|
|
nouns = ("fuze", "wave", "spark", "flow", "core")
|
|
adj = secrets.choice(adjectives)
|
|
noun = secrets.choice(nouns)
|
|
return f"{adj}-{noun}-{uuid.uuid4().hex[:5].lower()}"
|
|
|
|
|
|
def _retry_after_seconds(value: str | None) -> int | None:
|
|
if not value:
|
|
return None
|
|
try:
|
|
seconds = int(value.strip())
|
|
except ValueError:
|
|
return None
|
|
return max(1, min(seconds, 4 * 3600))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Internal failure taxonomy (never crosses the port boundary)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _AgyRateLimited(Exception):
|
|
def __init__(self, retry_after: int | None) -> None:
|
|
super().__init__("agy account rate limited")
|
|
self.retry_after = retry_after
|
|
|
|
|
|
class _AgyAuthError(Exception):
|
|
"""Refresh token rejected / account not entitled — needs a fresh login."""
|
|
|
|
|
|
class _AgyTransientError(Exception):
|
|
"""5xx or transport failure on one backend URL — try the next account."""
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# OAuth login
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass(slots=True, kw_only=True)
|
|
class AgyLoginChallenge:
|
|
login_id: str
|
|
auth_url: str
|
|
redirect_uri: str
|
|
verifier: str
|
|
state: str
|
|
created_at: float
|
|
|
|
|
|
@dataclass(slots=True, kw_only=True)
|
|
class AgyLoginResult:
|
|
"""Decrypted account material, produced by a completed OAuth exchange."""
|
|
|
|
email: str
|
|
project_id: str
|
|
refresh_token: str
|
|
|
|
|
|
@dataclass(slots=True, kw_only=True)
|
|
class AgyAccountModelInfo:
|
|
"""Model catalogue item as reported by an agy account."""
|
|
|
|
id: str
|
|
display_name: str
|
|
supports_thinking: bool = False
|
|
supported_efforts: list[str] = field(default_factory=list)
|
|
default_effort: str | None = None
|
|
supports_images: bool = False
|
|
|
|
|
|
class AgyOAuthError(ValidationError):
|
|
"""The pasted URL or the token exchange was rejected by Google."""
|
|
|
|
|
|
class AgyAuthClient:
|
|
"""Google OAuth (PKCE, installed-app client) for agy accounts."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
client_id: str | None = None,
|
|
client_secret: str | None = None,
|
|
redirect_uri: str | None = None,
|
|
timeout: float = 30.0,
|
|
) -> None:
|
|
self.client_id = client_id or settings.agy_client_id
|
|
self.client_secret = client_secret or settings.agy_client_secret
|
|
self.redirect_uri = redirect_uri or settings.agy_redirect_uri
|
|
self.timeout = timeout
|
|
|
|
# -- flow -------------------------------------------------------------
|
|
|
|
def begin(self) -> AgyLoginChallenge:
|
|
"""Open a login: returns the URL the operator must open in a browser.
|
|
|
|
The verifier stays in the returned challenge — it is the proof that
|
|
whoever redeems the login is whoever started it, and it is never
|
|
serialised out of the process.
|
|
"""
|
|
verifier, challenge = _pkce_pair()
|
|
login_id = uuid.uuid4().hex
|
|
state = uuid.uuid4().hex
|
|
params = urlencode(
|
|
{
|
|
"client_id": self.client_id,
|
|
"response_type": "code",
|
|
"redirect_uri": self.redirect_uri,
|
|
"scope": " ".join(OAUTH_SCOPES),
|
|
"code_challenge": challenge,
|
|
"code_challenge_method": "S256",
|
|
"state": state,
|
|
"access_type": "offline",
|
|
"prompt": "consent",
|
|
}
|
|
)
|
|
return AgyLoginChallenge(
|
|
login_id=login_id,
|
|
auth_url=f"{OAUTH_AUTH_URL}?{params}",
|
|
redirect_uri=self.redirect_uri,
|
|
verifier=verifier,
|
|
state=state,
|
|
created_at=time.monotonic(),
|
|
)
|
|
|
|
async def exchange(self, challenge: AgyLoginChallenge, callback_url: str) -> AgyLoginResult:
|
|
"""Redeem a started login with the callback URL from the browser.
|
|
|
|
The operator may paste the full ``http://localhost:51121/oauth-callback?
|
|
…`` URL (it never finishes loading — that is the documented shape of
|
|
this flow, the parameters are already in the address bar).
|
|
"""
|
|
parsed = urlparse(callback_url.strip())
|
|
query = parse_qs(parsed.query)
|
|
if query.get("error"):
|
|
raise AgyOAuthError(f"Google خطای احراز هویت داد: {query['error'][0]}")
|
|
codes = query.get("code") or []
|
|
states = query.get("state") or []
|
|
if not codes or not states:
|
|
raise AgyOAuthError(
|
|
"URL بازگشت معتبر نیست؛ باید پارامترهای code و state را داشته باشد."
|
|
)
|
|
if states[0] != challenge.state:
|
|
raise AgyOAuthError("state بررسی نشد؛ این URL به این درخواست لاگین تعلق ندارد.")
|
|
return await self._exchange_code(challenge, codes[0])
|
|
|
|
async def _exchange_code(self, challenge: AgyLoginChallenge, code: str) -> AgyLoginResult:
|
|
tokens = await self._post_token(
|
|
{
|
|
"client_id": self.client_id,
|
|
"client_secret": self.client_secret,
|
|
"code": code,
|
|
"grant_type": "authorization_code",
|
|
"redirect_uri": challenge.redirect_uri,
|
|
"code_verifier": challenge.verifier,
|
|
}
|
|
)
|
|
refresh_token = tokens.get("refresh_token") or ""
|
|
if not refresh_token:
|
|
# Happens when the consent screen was already granted for these
|
|
# scopes: Google omits the refresh token. A re-run with
|
|
# ``prompt=consent`` (which begin() always sends) fixes it, so the
|
|
# operator message says exactly that.
|
|
raise AgyOAuthError(
|
|
"Google refresh_token برنگرداند. یک بار دیگر لاگین کنید و در صفحهٔ "
|
|
"Google حتماً «Block/Allow» را تأیید کنید."
|
|
)
|
|
email, project_id = await asyncio.gather(
|
|
self._fetch_email(tokens.get("access_token", "")),
|
|
self._fetch_project_id(tokens.get("access_token", "")),
|
|
)
|
|
return AgyLoginResult(
|
|
email=email, project_id=project_id, refresh_token=refresh_token
|
|
)
|
|
|
|
async def refresh(self, refresh_token: str) -> str:
|
|
tokens = await self._post_token(
|
|
{
|
|
"client_id": self.client_id,
|
|
"client_secret": self.client_secret,
|
|
"refresh_token": refresh_token,
|
|
"grant_type": "refresh_token",
|
|
}
|
|
)
|
|
access_token = tokens.get("access_token")
|
|
if not access_token:
|
|
raise _AgyAuthError("Google refresh را نپذیرفت (توکن منقضی یا باطل شده).")
|
|
return access_token
|
|
|
|
# -- Google endpoints --------------------------------------------------
|
|
|
|
async def _post_token(self, data: dict[str, str]) -> dict:
|
|
try:
|
|
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
|
response = await client.post(
|
|
OAUTH_TOKEN_URL,
|
|
data=data,
|
|
headers={"Content-Type": "application/x-www-form-urlencoded"},
|
|
)
|
|
except httpx.HTTPError as exc:
|
|
raise ProviderUnavailableError("اتصال به سرور OAuth گوگل برقرار نشد.") from exc
|
|
if response.status_code >= 400:
|
|
body = response.text[:400]
|
|
# A rejected refresh token is an account problem, not an outage.
|
|
if "invalid_grant" in body:
|
|
raise _AgyAuthError(f"Google refresh را نپذیرفت: {body}")
|
|
raise AgyOAuthError(f"تبادل توکن ناموفق بود ({response.status_code}): {body}")
|
|
return response.json()
|
|
|
|
async def _fetch_email(self, access_token: str) -> str:
|
|
try:
|
|
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
|
response = await client.get(
|
|
USERINFO_URL, headers={"Authorization": f"Bearer {access_token}"}
|
|
)
|
|
if response.is_success:
|
|
return str(response.json().get("email", ""))
|
|
except httpx.HTTPError:
|
|
logger.debug("agy userinfo fetch failed", exc_info=True)
|
|
return ""
|
|
|
|
async def _fetch_project_id(self, access_token: str) -> str:
|
|
for base in BASE_URLS:
|
|
try:
|
|
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
|
response = await client.post(
|
|
f"{base}{PATH_LOAD_PROJECT}",
|
|
headers=self._api_headers(access_token),
|
|
json={
|
|
"metadata": {
|
|
"ideType": "IDE_UNSPECIFIED",
|
|
"platform": "PLATFORM_UNSPECIFIED",
|
|
"pluginType": "GEMINI",
|
|
}
|
|
},
|
|
)
|
|
if response.is_success:
|
|
project = response.json().get("cloudaicompanionProject")
|
|
if isinstance(project, str) and project:
|
|
return project
|
|
if isinstance(project, dict) and project.get("id"):
|
|
return str(project["id"])
|
|
except httpx.HTTPError:
|
|
continue
|
|
return ""
|
|
|
|
def _api_headers(self, access_token: str) -> dict[str, str]:
|
|
return {
|
|
"Authorization": f"Bearer {access_token}",
|
|
"Content-Type": "application/json",
|
|
"User-Agent": settings.agy_user_agent,
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Pending-login store (process-local, bounded, TTL'd)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_CHALLENGES: dict[str, AgyLoginChallenge] = {}
|
|
_CHALLENGES_LOCK = asyncio.Lock()
|
|
|
|
|
|
def _prune_challenges(now: float) -> None:
|
|
expired = [
|
|
login_id
|
|
for login_id, challenge in _CHALLENGES.items()
|
|
if now - challenge.created_at > CHALLENGE_TTL_SECONDS
|
|
]
|
|
for login_id in expired:
|
|
_CHALLENGES.pop(login_id, None)
|
|
|
|
|
|
async def begin_login() -> tuple[str, str]:
|
|
"""Start a headless login. Returns ``(login_id, auth_url)``.
|
|
|
|
``auth_url`` is what the admin console shows the platform owner; the owner
|
|
opens it, signs in with Google, and the resulting callback URL (or, in
|
|
listener mode, the captured redirect) is redeemed against ``login_id``.
|
|
"""
|
|
challenge = AgyAuthClient().begin()
|
|
now = time.monotonic()
|
|
async with _CHALLENGES_LOCK:
|
|
_prune_challenges(now)
|
|
if len(_CHALLENGES) >= 32:
|
|
# Bound memory: pending logins are transient by nature.
|
|
oldest = min(_CHALLENGES, key=lambda k: _CHALLENGES[k].created_at)
|
|
_CHALLENGES.pop(oldest, None)
|
|
_CHALLENGES[challenge.login_id] = challenge
|
|
return challenge.login_id, challenge.auth_url
|
|
|
|
|
|
async def redeem_login(login_id: str, callback_url: str) -> AgyLoginResult:
|
|
challenge = _CHALLENGES.pop(login_id, None)
|
|
if challenge is None:
|
|
raise AgyOAuthError("این درخواست لاگین منقضی شده است؛ دوباره آغاز کنید.")
|
|
return await AgyAuthClient().exchange(challenge, callback_url)
|
|
|
|
|
|
async def redeem_login_by_waiting(
|
|
login_id: str, *, timeout_seconds: float = 300
|
|
) -> AgyLoginResult:
|
|
"""Listener mode: capture the redirect on the loopback callback port.
|
|
|
|
Used by the operator-side script (and by a same-machine deployment): Google
|
|
redirects the browser to ``redirect_uri``; this server catches it, closes,
|
|
and redeems the login without any copy-paste.
|
|
"""
|
|
challenge = _CHALLENGES.get(login_id)
|
|
if challenge is None:
|
|
raise AgyOAuthError("این درخواست لاگین منقضی شده است؛ دوباره آغاز کنید.")
|
|
callback_path = urlparse(challenge.redirect_uri).path or "/"
|
|
port = int(urlparse(challenge.redirect_uri).port or 51121)
|
|
loop = asyncio.get_running_loop()
|
|
future: asyncio.Future[str] = loop.create_future()
|
|
|
|
async def _handle(
|
|
reader: asyncio.StreamReader, writer: asyncio.StreamWriter
|
|
) -> None:
|
|
try:
|
|
request_line = await reader.readline()
|
|
parts = request_line.decode("latin-1").split(" ")
|
|
target = parts[1] if len(parts) > 1 else "/"
|
|
if urlparse(target).path != callback_path:
|
|
writer.write(b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\n\r\n")
|
|
await writer.drain()
|
|
return
|
|
writer.write(
|
|
b"HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=utf-8\r\n"
|
|
b"Content-Length: 46\r\nConnection: close\r\n\r\n"
|
|
b"<html><body>OK - you may close this tab.</body></html>"
|
|
)
|
|
await writer.drain()
|
|
if not future.done():
|
|
future.set_result(target)
|
|
except Exception: # noqa: BLE001 - the future carries the real error
|
|
if not future.done():
|
|
future.set_exception(AgyOAuthError("دریافت بازگشت OAuth ناموفق بود."))
|
|
finally:
|
|
writer.close()
|
|
|
|
server = await asyncio.start_server(_handle, "127.0.0.1", port)
|
|
try:
|
|
async with asyncio.timeout(timeout_seconds):
|
|
target = await future
|
|
except TimeoutError as exc:
|
|
raise AgyOAuthError("مهلت لاگین تمام شد؛ دوباره آغاز کنید.") from exc
|
|
finally:
|
|
server.close()
|
|
await server.wait_closed()
|
|
_CHALLENGES.pop(login_id, None)
|
|
redirect_uri = challenge.redirect_uri
|
|
return await AgyAuthClient(redirect_uri=redirect_uri).exchange(
|
|
challenge, f"http://127.0.0.1:{port}{target}"
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Account pool
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass(slots=True, kw_only=True)
|
|
class AgyAccount:
|
|
"""One usable account: decrypted material plus its rotation identity."""
|
|
|
|
credential_id: str
|
|
email: str
|
|
project_id: str
|
|
refresh_token: str
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class _PoolState:
|
|
#: Monotonic rotation cursor — spreads load across the account set.
|
|
cursor: int = 0
|
|
#: credential_id -> epoch seconds until which the account is benched.
|
|
benched_until: dict[str, float] = field(default_factory=dict)
|
|
#: credential_id -> (access token, expiry epoch) — avoids a Google
|
|
#: round-trip on every single completion.
|
|
tokens: dict[str, tuple[str, float]] = field(default_factory=dict)
|
|
|
|
|
|
def _account_of(credential: ModelCredential) -> AgyAccount:
|
|
"""Decrypt one credential row into usable login material.
|
|
|
|
The ciphertext is a JSON envelope written at login time; older key-shaped
|
|
credentials (empty email) are treated as unusable for agy.
|
|
"""
|
|
try:
|
|
payload = json.loads(decrypt(credential.encrypted_api_key))
|
|
return AgyAccount(
|
|
credential_id=credential.id,
|
|
email=str(payload.get("email") or credential.email or ""),
|
|
project_id=str(payload.get("project_id") or credential.project_id or ""),
|
|
refresh_token=str(payload.get("refresh_token") or ""),
|
|
)
|
|
except Exception:
|
|
return AgyAccount(
|
|
credential_id=credential.id,
|
|
email=credential.email,
|
|
project_id=credential.project_id,
|
|
refresh_token="",
|
|
)
|
|
|
|
|
|
class AgyModelAdapter:
|
|
"""ModelProviderPort over the agy account pool.
|
|
|
|
One instance per process (the adapter registry caches it), so the rotation
|
|
cursor and the in-memory bench are shared by every worker and every
|
|
request. Hard state (cooldown / error / last-used) is mirrored to the
|
|
database through the injected ``report`` callback so the background
|
|
workers see the same picture.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
loader: _LoaderFn | None = None,
|
|
report: _StatusFn | None = None,
|
|
timeout: float = 240.0,
|
|
) -> None:
|
|
self._loader = loader
|
|
self._report = report
|
|
self._timeout = timeout
|
|
self._state = _PoolState()
|
|
self._lock = asyncio.Lock()
|
|
|
|
# -- port ---------------------------------------------------------------
|
|
|
|
async def complete(
|
|
self,
|
|
*,
|
|
model: AiModel,
|
|
credential: ModelCredential | None = None,
|
|
credentials: list[ModelCredential] | None = None,
|
|
messages: list[ChatMessage],
|
|
system: str = "",
|
|
temperature: float = 0.2,
|
|
max_tokens: int = 2048,
|
|
) -> ModelResponse:
|
|
pool = credentials if credentials is not None else (
|
|
[credential] if credential is not None else await self._load()
|
|
)
|
|
accounts = [a for a in (_account_of(c) for c in pool) if a.refresh_token]
|
|
if not accounts:
|
|
raise ProviderUnavailableError(
|
|
"هیچ اکانت agy فعالی ثبت نشده است. مالک سامانه باید ابتدا از "
|
|
"کنسول مدیریت یک اکانت Google اضافه کند."
|
|
)
|
|
ordered = self._order_accounts(accounts)
|
|
failures: list[str] = []
|
|
for account in ordered:
|
|
try:
|
|
response = await self._complete_with(
|
|
account, model=model, messages=messages, system=system,
|
|
temperature=temperature, max_tokens=max_tokens,
|
|
)
|
|
except _AgyRateLimited as exc:
|
|
failures.append(f"{account.email or account.credential_id}: محدودیت نرخ")
|
|
await self._bench(account, exc.retry_after)
|
|
continue
|
|
except _AgyAuthError as exc:
|
|
failures.append(f"{account.email or account.credential_id}: {exc}")
|
|
await self._bench(account, 3600, status="error", error=str(exc))
|
|
continue
|
|
except _AgyTransientError as exc:
|
|
failures.append(f"{account.email or account.credential_id}: {exc}")
|
|
continue
|
|
await self._succeed(account)
|
|
return response
|
|
raise ProviderUnavailableError(
|
|
"همهٔ اکانتهای agy در دسترس نبودند: " + "; ".join(failures[:4])
|
|
)
|
|
|
|
async def health(self) -> ModelHealth:
|
|
accounts = [a for a in (_account_of(c) for c in await self._load()) if a.refresh_token]
|
|
if not accounts:
|
|
return ModelHealth(ok=False, detail="هیچ اکانت agy ثبت نشده است.")
|
|
try:
|
|
await self._access_token(accounts[0])
|
|
except _AgyAuthError as exc:
|
|
return ModelHealth(ok=False, detail=str(exc))
|
|
except Exception as exc: # noqa: BLE001 - health must not raise
|
|
return ModelHealth(ok=False, detail=str(exc))
|
|
return ModelHealth(ok=True, detail=f"{len(accounts)} اکانت فعال")
|
|
|
|
async def available_models(self) -> list[str]:
|
|
"""Models the logged-in subscription exposes (admin console helper)."""
|
|
accounts = [a for a in (_account_of(c) for c in await self._load()) if a.refresh_token]
|
|
if not accounts:
|
|
return []
|
|
items = await self._fetch_models(accounts[0])
|
|
return [m.id for m in items]
|
|
|
|
async def available_models_for(
|
|
self, credential: ModelCredential
|
|
) -> list[AgyAccountModelInfo]:
|
|
"""Live model catalog of exactly one account.
|
|
|
|
Called per account selection in the admin console, so the list is
|
|
fetched fresh from Google every time (the access token is cached, the
|
|
``fetchAvailableModels`` round-trip is not) and reflects what *this*
|
|
subscription is entitled to along with each model's supported efforts.
|
|
"""
|
|
account = _account_of(credential)
|
|
if not account.refresh_token:
|
|
raise ProviderUnavailableError(
|
|
"این اکانت اطلاعات ورود ندارد؛ یک بار دیگر لاگین کنید."
|
|
)
|
|
return await self._fetch_models(account)
|
|
|
|
async def _fetch_models(self, account: AgyAccount) -> list[AgyAccountModelInfo]:
|
|
token = await self._access_token(account)
|
|
for base in BASE_URLS:
|
|
try:
|
|
async with httpx.AsyncClient(timeout=30.0) as client:
|
|
response = await client.post(
|
|
f"{base}{PATH_MODELS}",
|
|
headers={
|
|
"Authorization": f"Bearer {token}",
|
|
"Content-Type": "application/json",
|
|
"User-Agent": settings.agy_user_agent,
|
|
},
|
|
json={},
|
|
)
|
|
if response.is_success:
|
|
models = response.json().get("models") or {}
|
|
return self._parse_model_catalog(models)
|
|
except httpx.HTTPError:
|
|
continue
|
|
return []
|
|
|
|
@classmethod
|
|
def _parse_model_catalog(
|
|
cls, raw_models: dict[str, Any]
|
|
) -> list[AgyAccountModelInfo]:
|
|
items: list[AgyAccountModelInfo] = []
|
|
for model_id, details in sorted(raw_models.items()):
|
|
if not isinstance(details, dict):
|
|
details = {}
|
|
# Omit internal non-generative / tab-completion preview models
|
|
if details.get("isInternal") and not details.get("displayName"):
|
|
continue
|
|
if model_id.startswith("tab_") or model_id.startswith("chat_"):
|
|
continue
|
|
|
|
display_name = str(details.get("displayName") or model_id)
|
|
supports_images = bool(details.get("supportsImages"))
|
|
supports_thinking = bool(details.get("supportsThinking"))
|
|
if not supports_thinking and _supports_thinking(model_id):
|
|
supports_thinking = True
|
|
|
|
supported_efforts = cls._supported_efforts_for(model_id, details)
|
|
default_effort = cls._default_effort_for(supported_efforts)
|
|
|
|
items.append(
|
|
AgyAccountModelInfo(
|
|
id=model_id,
|
|
display_name=display_name,
|
|
supports_thinking=supports_thinking,
|
|
supported_efforts=supported_efforts,
|
|
default_effort=default_effort,
|
|
supports_images=supports_images,
|
|
)
|
|
)
|
|
return items
|
|
|
|
@classmethod
|
|
def _supported_efforts_for(
|
|
cls, model_id: str, details: dict[str, Any]
|
|
) -> list[str]:
|
|
lowered = model_id.lower()
|
|
if "image" in lowered:
|
|
return []
|
|
if "claude" in lowered:
|
|
return []
|
|
|
|
supports_thinking = details.get("supportsThinking")
|
|
if supports_thinking is False:
|
|
return []
|
|
if supports_thinking is None and not _supports_thinking(model_id):
|
|
return []
|
|
|
|
if "gpt" in lowered:
|
|
return ["medium"]
|
|
if "pro" in lowered and "3.1" in lowered:
|
|
return ["low", "high"]
|
|
|
|
budget = details.get("thinkingBudget")
|
|
if isinstance(budget, int) and budget > 0:
|
|
if budget <= 4096:
|
|
return ["low"]
|
|
if budget <= 8192:
|
|
return ["medium"]
|
|
return ["high"]
|
|
|
|
return ["low", "medium", "high"]
|
|
|
|
@staticmethod
|
|
def _default_effort_for(supported: list[str]) -> str | None:
|
|
if not supported:
|
|
return None
|
|
if "high" in supported:
|
|
return "high"
|
|
if "medium" in supported:
|
|
return "medium"
|
|
return supported[0]
|
|
|
|
# -- internals ----------------------------------------------------------
|
|
|
|
async def _load(self) -> list[ModelCredential]:
|
|
if self._loader is None:
|
|
return []
|
|
return await self._loader()
|
|
|
|
async def _report_state(
|
|
self,
|
|
credential_id: str,
|
|
*,
|
|
status: str | None = None,
|
|
cooldown_until: datetime | None = None,
|
|
error: str | None = None,
|
|
used: bool = False,
|
|
) -> None:
|
|
if self._report is None:
|
|
return
|
|
try:
|
|
await self._report(
|
|
credential_id,
|
|
status=status,
|
|
cooldown_until=cooldown_until,
|
|
error=error,
|
|
used=used,
|
|
)
|
|
except Exception: # noqa: BLE001 - reporting must never kill a call
|
|
logger.debug("agy state report failed", exc_info=True)
|
|
|
|
def _order_accounts(self, accounts: list[AgyAccount]) -> list[AgyAccount]:
|
|
"""Round-robin, skipping accounts benched in this process.
|
|
|
|
The DB-shaped cooldown is already applied by the query that produced
|
|
the pool; this pass additionally honours recent 429s that have not been
|
|
persisted yet (or a reporter that is not wired).
|
|
"""
|
|
now = time.time()
|
|
usable = [a for a in accounts if self._state.benched_until.get(a.credential_id, 0) <= now]
|
|
if not usable:
|
|
# Everything is benched in memory — fall through to the DB's view
|
|
# rather than failing a request the pool could still serve.
|
|
usable = list(accounts)
|
|
n = len(usable)
|
|
self._state.cursor += 1
|
|
start = self._state.cursor % n
|
|
return usable[start:] + usable[:start]
|
|
|
|
async def _bench(
|
|
self,
|
|
account: AgyAccount,
|
|
retry_after: int | None,
|
|
*,
|
|
status: str = "cooldown",
|
|
error: str = "",
|
|
) -> None:
|
|
seconds = retry_after or settings.agy_429_cooldown_seconds
|
|
until = datetime.now(UTC) + timedelta(seconds=seconds)
|
|
async with self._lock:
|
|
self._state.benched_until[account.credential_id] = until.timestamp()
|
|
await self._report_state(
|
|
account.credential_id, status=status, cooldown_until=until, error=error
|
|
)
|
|
|
|
async def _succeed(self, account: AgyAccount) -> None:
|
|
async with self._lock:
|
|
self._state.benched_until.pop(account.credential_id, None)
|
|
await self._report_state(account.credential_id, status="active", used=True)
|
|
|
|
async def _access_token(self, account: AgyAccount) -> str:
|
|
"""Cached refresh: Google tokens live ~1h; one call serves the burst."""
|
|
cached = self._state.tokens.get(account.credential_id)
|
|
if cached and cached[1] > time.time() + 30:
|
|
return cached[0]
|
|
token = await AgyAuthClient().refresh(account.refresh_token)
|
|
async with self._lock:
|
|
self._state.tokens[account.credential_id] = (token, time.time() + 3300)
|
|
return token
|
|
|
|
async def _complete_with(
|
|
self,
|
|
account: AgyAccount,
|
|
*,
|
|
model: AiModel,
|
|
messages: list[ChatMessage],
|
|
system: str,
|
|
temperature: float,
|
|
max_tokens: int,
|
|
) -> ModelResponse:
|
|
token = await self._access_token(account)
|
|
payload = _build_generate_payload(
|
|
account, model.internal_name, messages, system, temperature, max_tokens,
|
|
reasoning_effort=model.reasoning_effort,
|
|
)
|
|
headers = {
|
|
"Authorization": f"Bearer {token}",
|
|
"Content-Type": "application/json",
|
|
"User-Agent": settings.agy_user_agent,
|
|
"Accept": "application/json",
|
|
}
|
|
started = time.perf_counter()
|
|
last_transient: Exception | None = None
|
|
retry_after: int | None = None
|
|
for base in BASE_URLS:
|
|
try:
|
|
async with httpx.AsyncClient(timeout=self._timeout) as client:
|
|
response = await client.post(
|
|
f"{base}{PATH_GENERATE}", json=payload, headers=headers
|
|
)
|
|
except httpx.HTTPError:
|
|
last_transient = _AgyTransientError(f"اتصال به {base} برقرار نشد.")
|
|
continue
|
|
if response.status_code == 429:
|
|
# Both backends share one quota, but a 429 can also be a
|
|
# backend hiccup — try the other base URL before benching the
|
|
# whole account.
|
|
retry_after = _retry_after_seconds(
|
|
response.headers.get("Retry-After")
|
|
) or retry_after
|
|
continue
|
|
if response.status_code in (401, 403):
|
|
raise _AgyAuthError(
|
|
f"اکانت {account.email or 'agy'} دسترسی ندارد یا نشست منقضی شده است."
|
|
)
|
|
if response.status_code >= 500:
|
|
last_transient = _AgyTransientError(f"{base} خطای {response.status_code} داد.")
|
|
continue
|
|
if response.status_code >= 400:
|
|
# Our request shape is wrong — another account will fail the
|
|
# same way, so surface it instead of burning the pool.
|
|
raise ProviderUnavailableError(
|
|
f"agy درخواست را رد کرد ({response.status_code}): "
|
|
f"{response.text[:300]}"
|
|
)
|
|
return _parse_generate_response(
|
|
response.json(), model=model, started=started
|
|
)
|
|
if retry_after is not None:
|
|
raise _AgyRateLimited(retry_after)
|
|
raise last_transient or _AgyTransientError("هیچ بکاند agy پاسخ نداد.")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Message / payload translation (portable, pure functions — unit-tested)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _build_generate_payload(
|
|
account: AgyAccount,
|
|
model_name: str,
|
|
messages: list[ChatMessage],
|
|
system: str,
|
|
temperature: float,
|
|
max_tokens: int,
|
|
reasoning_effort: str = "",
|
|
) -> dict:
|
|
"""Translate the port's OpenAI-shaped call into the Antigravity envelope."""
|
|
system_parts: list[str] = [system] if system else []
|
|
contents: list[dict] = []
|
|
for message in messages:
|
|
if message.role == "system":
|
|
system_parts.append(message.content)
|
|
continue
|
|
role = "model" if message.role == "assistant" else "user"
|
|
contents.append({"role": role, "parts": [{"text": message.content}]})
|
|
|
|
generation: dict = {
|
|
"temperature": temperature,
|
|
"maxOutputTokens": max_tokens,
|
|
}
|
|
# Reasoning effort is only sent to models that can think; the backend
|
|
# rejects ``thinkingConfig`` everywhere else with a 400.
|
|
if reasoning_effort and _supports_thinking(model_name):
|
|
budget = EFFORT_THINKING_BUDGETS.get(reasoning_effort, 8192)
|
|
generation["thinkingConfig"] = {
|
|
"thinkingBudget": budget,
|
|
"includeThoughts": True,
|
|
}
|
|
if max_tokens <= budget:
|
|
# The budget must fit inside the answer window or the thinking
|
|
# model has nowhere to emit its final answer.
|
|
generation["maxOutputTokens"] = budget + 8000
|
|
|
|
request: dict = {
|
|
"contents": contents,
|
|
"sessionId": _random_session_id(),
|
|
"generationConfig": generation,
|
|
"toolConfig": {"functionCallingConfig": {"mode": "VALIDATED"}},
|
|
"safetySettings": SAFETY_SETTINGS,
|
|
}
|
|
if system_parts:
|
|
request["systemInstruction"] = {
|
|
"role": "user",
|
|
"parts": [{"text": "\n\n".join(p for p in system_parts if p)}],
|
|
}
|
|
return {
|
|
"project": account.project_id or _fallback_project_id(),
|
|
"model": model_name,
|
|
"userAgent": "antigravity",
|
|
"requestId": f"agent-{uuid.uuid4()}",
|
|
"request": request,
|
|
}
|
|
|
|
|
|
def _parse_generate_response(
|
|
data: dict, *, model: AiModel, started: float
|
|
) -> ModelResponse:
|
|
candidates = data.get("candidates") or [{}]
|
|
candidate = candidates[0] if isinstance(candidates[0], dict) else {}
|
|
text_parts: list[str] = []
|
|
for part in (candidate.get("content") or {}).get("parts") or []:
|
|
if not isinstance(part, dict):
|
|
continue
|
|
text = part.get("text")
|
|
if text and not part.get("thought"):
|
|
text_parts.append(text)
|
|
usage = data.get("usageMetadata") or {}
|
|
latency_ms = int((time.perf_counter() - started) * 1000)
|
|
return ModelResponse(
|
|
content="".join(text_parts),
|
|
model_internal_name=model.internal_name,
|
|
input_tokens=int(usage.get("promptTokenCount", 0) or 0),
|
|
output_tokens=int(usage.get("candidatesTokenCount", 0) or 0),
|
|
latency_ms=latency_ms,
|
|
raw={
|
|
"finish_reason": candidate.get("finishReason", ""),
|
|
"thought_tokens": int(usage.get("thoughtsTokenCount", 0) or 0),
|
|
},
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Database wiring for the registry (each call opens its own short session, so
|
|
# the adapter never outlives or depends on an HTTP request's session).
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def load_agy_accounts() -> list[ModelCredential]:
|
|
"""Every usable agy credential row, ordered least-recently-used first."""
|
|
from app.core.database import SessionFactory
|
|
from app.infrastructure.repositories.ai_control import (
|
|
ModelCredentialRepository,
|
|
ProviderRepository,
|
|
)
|
|
|
|
async with SessionFactory() as session:
|
|
providers = await ProviderRepository(session).list()
|
|
provider = next(
|
|
(p for p in providers if p.type.value == "agy" and p.is_active), None
|
|
)
|
|
if provider is None:
|
|
return []
|
|
return await ModelCredentialRepository(session).list_active_for_provider(
|
|
provider.id
|
|
)
|
|
|
|
|
|
async def report_agy_state(
|
|
credential_id: str,
|
|
*,
|
|
status: str | None = None,
|
|
cooldown_until: datetime | None = None,
|
|
error: str | None = None,
|
|
used: bool = False,
|
|
) -> None:
|
|
"""Mirror rotation state into the DB for the workers' benefit."""
|
|
from app.core.database import SessionFactory
|
|
from app.infrastructure.repositories.ai_control import ModelCredentialRepository
|
|
|
|
async with SessionFactory() as session:
|
|
repo = ModelCredentialRepository(session)
|
|
if status is not None:
|
|
await repo.set_status(
|
|
credential_id,
|
|
status=status,
|
|
cooldown_until=cooldown_until,
|
|
error=error or "",
|
|
)
|
|
if used:
|
|
await repo.mark_used(credential_id, when=utcnow())
|
|
await session.commit()
|
|
|
|
|
|
__all__ = [
|
|
"AgyAccount",
|
|
"AgyAuthClient",
|
|
"AgyLoginChallenge",
|
|
"AgyLoginResult",
|
|
"AgyModelAdapter",
|
|
"AgyOAuthError",
|
|
"begin_login",
|
|
"load_agy_accounts",
|
|
"redeem_login",
|
|
"redeem_login_by_waiting",
|
|
"report_agy_state",
|
|
]
|