u78649634_default/smart-org
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

افزودن ارائه‌دهندهٔ agy (Google Antigravity) با لاگین headless و پول چرخشی اکانت

Browse Source 
- Adapter اکانت‌محور با چرخش round-robin و بنچ‌شدن خودکار روی ۴۲۹/۴۰۱
- ورود ادمین از طریق URL بازگشتی گوگل (PKCE) و ذخیرهٔ رمزنگاری‌شدهٔ refresh token
- صفحهٔ مدیریت اکانت‌های agy در پنل ادمین به همراه تست و فعال/غیرفعال‌سازی
- مهاجرت Alembic b1c2d3e4f5a6 برای ستون‌های اکانت agy
- seed همیشه کاتالوگ agy را می‌سازد و فقط با AI_DEFAULT_PROVIDER=agy مسیریابی می‌شود
This commit is contained in:
mjalmoudy94
2026-10-13 08:20:27 +03:30
parent 54074f10c3
commit 3d1e98f4fb
25 changed files with 3073 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/app/infrastructure/adapters/agy.py
+904
View File
@@ -0,0 +1,904 @@
"""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 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
_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
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 []
token = await self._access_token(accounts[0])
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 sorted(models)
except httpx.HTTPError:
continue
return []
# -- 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
)
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,
) -> 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}]})
request: dict = {
"contents": contents,
"sessionId": _random_session_id(),
"generationConfig": {
"temperature": temperature,
"maxOutputTokens": max_tokens,
},
"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",
]
apps/api/app/infrastructure/adapters/mock.py
+1
View File
@@ -66,6 +66,7 @@ class MockModelProviderAdapter:
*,
model: AiModel,
credential=None,
credentials=None, # accepted for port parity; the mock has no pool
messages: list[ChatMessage],
system: str = "",
temperature: float = 0.2,
apps/api/app/infrastructure/adapters/openai_compatible.py
+2 -1
View File
@@ -34,7 +34,8 @@ class OpenAICompatibleModelAdapter:
self,
*,
model: AiModel,
credential: ModelCredential | None,
credential: ModelCredential | None = None,
credentials: list[ModelCredential] | None = None, # no pool here
messages: list[ChatMessage],
system: str = "",
temperature: float = 0.2,
apps/api/app/infrastructure/adapters/registry.py
+9 -1
View File
@@ -28,9 +28,17 @@ def get_model_provider():
base_url=settings.openai_compatible_base_url,
default_api_key=settings.openai_compatible_api_key,
)
if provider == "agy":
from app.infrastructure.adapters.agy import (
AgyModelAdapter,
load_agy_accounts,
report_agy_state,
)
return AgyModelAdapter(loader=load_agy_accounts, report=report_agy_state)
raise ProviderUnavailableError(
f"AI_DEFAULT_PROVIDER='{provider}' پشتیبانی نمی‌شود. "
"مقادیر مجاز: mock | openai_compatible"
"مقادیر مجاز: mock | openai_compatible | agy"
)