feat(providers): add Antigravity provider for unified Claude/Gemini access
Some checks failed
Semantic Release / release (push) Has been cancelled
Some checks failed
Semantic Release / release (push) Has been cancelled
Implements a new provider that uses Google's Antigravity unified gateway API to access Claude, Gemini, and other models through a single OAuth2-authenticated endpoint. Features: - OAuth2 token management with automatic refresh - Multi-account rotation for rate limit distribution - Support for Claude Opus/Sonnet 4.5 (with/without thinking) - Support for Gemini 2.5/3 models (Pro/Flash variants) - Thinking mode support with configurable tokens - Image processing support - Dual quota pool tracking (antigravity vs gemini-cli) - Gemini-style API request format Authentication: - Reads from ANTIGRAVITY_REFRESH_TOKEN env var (priority) - Falls back to ~/.config/opencode/antigravity-accounts.json - Automatic token refresh with retry logic - Rate limit tracking per account and quota pool Files added: - providers/antigravity.py - Main provider implementation - providers/antigravity_auth.py - OAuth token manager - providers/registries/antigravity.py - Model registry - conf/antigravity_models.json - Model definitions (11 models) - docs/antigravity_provider.md - Setup and usage docs - tests/test_antigravity_provider.py - Unit tests (14 pass) Integration: - Added to provider priority order after ZEN - Registered in server.py with auto-detection - ToS warning logged on first use
This commit is contained in:
576
providers/antigravity.py
Normal file
576
providers/antigravity.py
Normal file
@@ -0,0 +1,576 @@
|
||||
"""Antigravity model provider implementation.
|
||||
|
||||
Antigravity is Google's unified gateway API for accessing multiple AI models
|
||||
(Claude, Gemini, GPT-OSS) through a single Gemini-style interface.
|
||||
|
||||
WARNING: Using this provider may violate Google's Terms of Service.
|
||||
See docs/antigravity_provider.md for important information about risks.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import uuid
|
||||
from typing import TYPE_CHECKING, Any, ClassVar
|
||||
|
||||
import httpx
|
||||
|
||||
from utils.env import get_env
|
||||
from utils.image_utils import validate_image
|
||||
|
||||
from .antigravity_auth import AntigravityTokenManager
|
||||
from .base import ModelProvider
|
||||
from .registries.antigravity import AntigravityModelRegistry
|
||||
from .shared import ModelCapabilities, ModelResponse, ProviderType
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tools.models import ToolModelCategory
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Antigravity API endpoints
|
||||
PRODUCTION_ENDPOINT = "https://cloudcode-pa.googleapis.com"
|
||||
DAILY_ENDPOINT = "https://daily-cloudcode-pa.sandbox.googleapis.com"
|
||||
|
||||
# API path for content generation
|
||||
GENERATE_CONTENT_PATH = "/v1internal:generateContent"
|
||||
STREAM_GENERATE_CONTENT_PATH = "/v1internal:streamGenerateContent"
|
||||
|
||||
# ToS warning message
|
||||
TOS_WARNING = """
|
||||
================================================================================
|
||||
ANTIGRAVITY PROVIDER WARNING
|
||||
|
||||
Using the Antigravity provider may violate Google's Terms of Service.
|
||||
Users have reported account bans when using this approach.
|
||||
|
||||
HIGH-RISK SCENARIOS:
|
||||
- Fresh Google accounts have a very high chance of getting banned
|
||||
- New accounts with Pro/Ultra subscriptions are frequently flagged
|
||||
|
||||
By using this provider, you acknowledge and accept all associated risks.
|
||||
See docs/antigravity_provider.md for more information.
|
||||
================================================================================
|
||||
"""
|
||||
|
||||
|
||||
class AntigravityProvider(ModelProvider):
|
||||
"""Provider for Google's Antigravity unified gateway API.
|
||||
|
||||
Antigravity provides access to Claude, Gemini, and other models through a
|
||||
single Gemini-style API interface. This provider handles:
|
||||
|
||||
* OAuth2 authentication with refresh token management
|
||||
* Multi-account rotation on rate limits
|
||||
* Request transformation to Gemini-style format
|
||||
* Extended thinking support for compatible models
|
||||
* Image processing for vision requests
|
||||
|
||||
Configuration:
|
||||
ANTIGRAVITY_REFRESH_TOKEN: OAuth2 refresh token (env var override)
|
||||
ANTIGRAVITY_PROJECT_ID: Google Cloud project ID
|
||||
ANTIGRAVITY_BASE_URL: Custom endpoint URL (defaults to production)
|
||||
"""
|
||||
|
||||
# Model registry for capability metadata
|
||||
_registry: ClassVar[AntigravityModelRegistry | None] = None
|
||||
MODEL_CAPABILITIES: ClassVar[dict[str, ModelCapabilities]] = {}
|
||||
|
||||
# Thinking mode budgets (percentage of max_thinking_tokens)
|
||||
THINKING_BUDGETS = {
|
||||
"minimal": 0.005,
|
||||
"low": 0.08,
|
||||
"medium": 0.33,
|
||||
"high": 0.67,
|
||||
"max": 1.0,
|
||||
}
|
||||
|
||||
# Track whether warning has been shown this session
|
||||
_warning_shown: ClassVar[bool] = False
|
||||
|
||||
def __init__(self, api_key: str = "", **kwargs) -> None:
|
||||
"""Initialize Antigravity provider.
|
||||
|
||||
Args:
|
||||
api_key: Not used (authentication via OAuth2 refresh tokens)
|
||||
**kwargs: Additional configuration options
|
||||
"""
|
||||
# Initialize registry
|
||||
if AntigravityProvider._registry is None:
|
||||
AntigravityProvider._registry = AntigravityModelRegistry()
|
||||
AntigravityProvider.MODEL_CAPABILITIES = dict(self._registry.model_map)
|
||||
logger.info("Antigravity provider loaded %d models", len(self._registry.model_map))
|
||||
|
||||
super().__init__(api_key, **kwargs)
|
||||
|
||||
# Initialize token manager
|
||||
self._token_manager = AntigravityTokenManager()
|
||||
|
||||
# HTTP client for API requests
|
||||
self._http_client: httpx.Client | None = None
|
||||
|
||||
# Base URL configuration
|
||||
self._base_url = get_env("ANTIGRAVITY_BASE_URL") or PRODUCTION_ENDPOINT
|
||||
|
||||
# Show warning on first use
|
||||
if not AntigravityProvider._warning_shown:
|
||||
logger.warning(TOS_WARNING)
|
||||
AntigravityProvider._warning_shown = True
|
||||
|
||||
self._invalidate_capability_cache()
|
||||
|
||||
@property
|
||||
def http_client(self) -> httpx.Client:
|
||||
"""Lazy initialization of HTTP client."""
|
||||
if self._http_client is None:
|
||||
timeout = httpx.Timeout(
|
||||
connect=30.0,
|
||||
read=600.0, # 10 minutes for long responses
|
||||
write=60.0,
|
||||
pool=600.0,
|
||||
)
|
||||
self._http_client = httpx.Client(timeout=timeout, follow_redirects=True)
|
||||
return self._http_client
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Provider identity
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_provider_type(self) -> ProviderType:
|
||||
"""Return the provider type."""
|
||||
return ProviderType.ANTIGRAVITY
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Capability surface
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _lookup_capabilities(
|
||||
self,
|
||||
canonical_name: str,
|
||||
requested_name: str | None = None,
|
||||
) -> ModelCapabilities | None:
|
||||
"""Fetch capabilities from the registry."""
|
||||
if self._registry:
|
||||
return self._registry.get_capabilities(canonical_name)
|
||||
return None
|
||||
|
||||
def get_all_model_capabilities(self) -> dict[str, ModelCapabilities]:
|
||||
"""Return registry-backed capabilities."""
|
||||
if not self._registry:
|
||||
return {}
|
||||
return dict(self._registry.model_map)
|
||||
|
||||
def list_models(
|
||||
self,
|
||||
*,
|
||||
respect_restrictions: bool = True,
|
||||
include_aliases: bool = True,
|
||||
lowercase: bool = False,
|
||||
unique: bool = False,
|
||||
) -> list[str]:
|
||||
"""Return available Antigravity model names."""
|
||||
if not self._registry:
|
||||
return []
|
||||
|
||||
from utils.model_restrictions import get_restriction_service
|
||||
|
||||
restriction_service = get_restriction_service() if respect_restrictions else None
|
||||
allowed_configs: dict[str, ModelCapabilities] = {}
|
||||
|
||||
for model_name in self._registry.list_models():
|
||||
config = self._registry.resolve(model_name)
|
||||
if not config:
|
||||
continue
|
||||
|
||||
if restriction_service:
|
||||
if not restriction_service.is_allowed(self.get_provider_type(), model_name):
|
||||
continue
|
||||
|
||||
allowed_configs[model_name] = config
|
||||
|
||||
if not allowed_configs:
|
||||
return []
|
||||
|
||||
return ModelCapabilities.collect_model_names(
|
||||
allowed_configs,
|
||||
include_aliases=include_aliases,
|
||||
lowercase=lowercase,
|
||||
unique=unique,
|
||||
)
|
||||
|
||||
def _resolve_model_name(self, model_name: str) -> str:
|
||||
"""Resolve aliases to canonical model names."""
|
||||
if self._registry:
|
||||
config = self._registry.resolve(model_name)
|
||||
if config and config.model_name != model_name:
|
||||
logger.debug("Resolved Antigravity alias '%s' to '%s'", model_name, config.model_name)
|
||||
return config.model_name
|
||||
return model_name
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Content generation
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def generate_content(
|
||||
self,
|
||||
prompt: str,
|
||||
model_name: str,
|
||||
system_prompt: str | None = None,
|
||||
temperature: float = 0.7,
|
||||
max_output_tokens: int | None = None,
|
||||
thinking_mode: str = "medium",
|
||||
images: list[str] | None = None,
|
||||
**kwargs,
|
||||
) -> ModelResponse:
|
||||
"""Generate content using the Antigravity API.
|
||||
|
||||
Args:
|
||||
prompt: User prompt to send to the model
|
||||
model_name: Model identifier (e.g., 'claude-opus-4-5-thinking')
|
||||
system_prompt: Optional system instructions
|
||||
temperature: Sampling temperature (0.0-2.0)
|
||||
max_output_tokens: Maximum tokens in response
|
||||
thinking_mode: Thinking budget level for thinking models
|
||||
images: Optional list of image paths or data URLs
|
||||
**kwargs: Additional parameters
|
||||
|
||||
Returns:
|
||||
ModelResponse with generated content and metadata
|
||||
"""
|
||||
# Validate and resolve model name
|
||||
self.validate_parameters(model_name, temperature)
|
||||
capabilities = self.get_capabilities(model_name)
|
||||
resolved_model = self._resolve_model_name(model_name)
|
||||
|
||||
# Get authentication
|
||||
access_token, project_id, fingerprint_headers = self._token_manager.get_access_token(model=resolved_model)
|
||||
|
||||
# Build request
|
||||
request_body = self._build_request(
|
||||
prompt=prompt,
|
||||
model_name=resolved_model,
|
||||
project_id=project_id,
|
||||
system_prompt=system_prompt,
|
||||
temperature=temperature,
|
||||
max_output_tokens=max_output_tokens,
|
||||
thinking_mode=thinking_mode,
|
||||
images=images,
|
||||
capabilities=capabilities,
|
||||
)
|
||||
|
||||
# Build headers
|
||||
headers = {
|
||||
"Authorization": f"Bearer {access_token}",
|
||||
"Content-Type": "application/json",
|
||||
**fingerprint_headers,
|
||||
}
|
||||
|
||||
# Retry logic
|
||||
max_retries = 4
|
||||
retry_delays = [1.0, 3.0, 5.0, 8.0]
|
||||
attempt_counter = {"value": 0}
|
||||
|
||||
def _attempt() -> ModelResponse:
|
||||
attempt_counter["value"] += 1
|
||||
|
||||
url = f"{self._base_url}{GENERATE_CONTENT_PATH}"
|
||||
logger.debug("Antigravity request to %s for model %s", url, resolved_model)
|
||||
|
||||
response = self.http_client.post(
|
||||
url,
|
||||
json=request_body,
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
# Handle errors
|
||||
if response.status_code != 200:
|
||||
self._handle_error_response(response, resolved_model)
|
||||
|
||||
return self._parse_response(response.json(), resolved_model, capabilities, thinking_mode)
|
||||
|
||||
try:
|
||||
return self._run_with_retries(
|
||||
operation=_attempt,
|
||||
max_attempts=max_retries,
|
||||
delays=retry_delays,
|
||||
log_prefix=f"Antigravity API ({resolved_model})",
|
||||
)
|
||||
except Exception as exc:
|
||||
attempts = max(attempt_counter["value"], 1)
|
||||
error_msg = f"Antigravity API error for {resolved_model} after {attempts} attempt(s): {exc}"
|
||||
raise RuntimeError(error_msg) from exc
|
||||
|
||||
def _build_request(
|
||||
self,
|
||||
prompt: str,
|
||||
model_name: str,
|
||||
project_id: str,
|
||||
system_prompt: str | None = None,
|
||||
temperature: float = 0.7,
|
||||
max_output_tokens: int | None = None,
|
||||
thinking_mode: str = "medium",
|
||||
images: list[str] | None = None,
|
||||
capabilities: ModelCapabilities | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Build the Antigravity API request body."""
|
||||
# Build contents array (Gemini-style format)
|
||||
contents = []
|
||||
|
||||
# Build user message parts
|
||||
user_parts = []
|
||||
|
||||
# Add text prompt
|
||||
user_parts.append({"text": prompt})
|
||||
|
||||
# Add images if provided
|
||||
if images and capabilities and capabilities.supports_images:
|
||||
for image_path in images:
|
||||
try:
|
||||
image_part = self._process_image(image_path)
|
||||
if image_part:
|
||||
user_parts.append(image_part)
|
||||
except Exception as e:
|
||||
logger.warning("Failed to process image %s: %s", image_path, e)
|
||||
elif images:
|
||||
logger.warning("Model %s does not support images, ignoring %d image(s)", model_name, len(images))
|
||||
|
||||
contents.append({"role": "user", "parts": user_parts})
|
||||
|
||||
# Build generation config
|
||||
generation_config: dict[str, Any] = {
|
||||
"temperature": temperature,
|
||||
}
|
||||
|
||||
if max_output_tokens:
|
||||
generation_config["maxOutputTokens"] = max_output_tokens
|
||||
|
||||
# Add thinking config for thinking-capable models
|
||||
if capabilities and capabilities.supports_extended_thinking and thinking_mode in self.THINKING_BUDGETS:
|
||||
max_thinking = capabilities.max_thinking_tokens or 8192
|
||||
thinking_budget = int(max_thinking * self.THINKING_BUDGETS[thinking_mode])
|
||||
generation_config["thinkingConfig"] = {
|
||||
"thinkingBudget": thinking_budget,
|
||||
"includeThoughts": True,
|
||||
}
|
||||
# Ensure maxOutputTokens > thinkingBudget
|
||||
if not max_output_tokens or max_output_tokens <= thinking_budget:
|
||||
generation_config["maxOutputTokens"] = thinking_budget + 10000
|
||||
|
||||
# Build system instruction (must be object with parts, not string)
|
||||
system_instruction = None
|
||||
if system_prompt:
|
||||
system_instruction = {"parts": [{"text": system_prompt}]}
|
||||
|
||||
# Build full request
|
||||
request: dict[str, Any] = {
|
||||
"project": project_id,
|
||||
"model": model_name,
|
||||
"request": {
|
||||
"contents": contents,
|
||||
"generationConfig": generation_config,
|
||||
},
|
||||
"userAgent": "antigravity",
|
||||
"requestId": f"pal-{uuid.uuid4().hex[:12]}",
|
||||
}
|
||||
|
||||
if system_instruction:
|
||||
request["request"]["systemInstruction"] = system_instruction
|
||||
|
||||
return request
|
||||
|
||||
def _parse_response(
|
||||
self,
|
||||
data: dict[str, Any],
|
||||
model_name: str,
|
||||
capabilities: ModelCapabilities | None,
|
||||
thinking_mode: str,
|
||||
) -> ModelResponse:
|
||||
"""Parse the Antigravity API response."""
|
||||
response_data = data.get("response", data)
|
||||
|
||||
# Extract content from candidates
|
||||
content = ""
|
||||
thinking_content = ""
|
||||
candidates = response_data.get("candidates", [])
|
||||
|
||||
if candidates:
|
||||
candidate = candidates[0]
|
||||
candidate_content = candidate.get("content", {})
|
||||
parts = candidate_content.get("parts", [])
|
||||
|
||||
for part in parts:
|
||||
if part.get("thought"):
|
||||
# This is thinking content
|
||||
thinking_content += part.get("text", "")
|
||||
elif "thoughtSignature" in part and "text" in part:
|
||||
# Gemini thinking with signature
|
||||
thinking_content += part.get("text", "")
|
||||
elif "text" in part:
|
||||
content += part.get("text", "")
|
||||
|
||||
# Extract usage metadata
|
||||
usage_metadata = response_data.get("usageMetadata", {})
|
||||
usage = {
|
||||
"input_tokens": usage_metadata.get("promptTokenCount", 0),
|
||||
"output_tokens": usage_metadata.get("candidatesTokenCount", 0),
|
||||
"total_tokens": usage_metadata.get("totalTokenCount", 0),
|
||||
}
|
||||
|
||||
# Add thinking tokens if available
|
||||
if "thoughtsTokenCount" in usage_metadata:
|
||||
usage["thinking_tokens"] = usage_metadata["thoughtsTokenCount"]
|
||||
|
||||
# Extract finish reason
|
||||
finish_reason = "UNKNOWN"
|
||||
if candidates:
|
||||
finish_reason = candidates[0].get("finishReason", "STOP")
|
||||
|
||||
return ModelResponse(
|
||||
content=content,
|
||||
usage=usage,
|
||||
model_name=model_name,
|
||||
friendly_name="Antigravity",
|
||||
provider=ProviderType.ANTIGRAVITY,
|
||||
metadata={
|
||||
"thinking_mode": thinking_mode if capabilities and capabilities.supports_extended_thinking else None,
|
||||
"thinking_content": thinking_content if thinking_content else None,
|
||||
"finish_reason": finish_reason,
|
||||
"model_version": response_data.get("modelVersion"),
|
||||
"response_id": response_data.get("responseId"),
|
||||
"trace_id": data.get("traceId"),
|
||||
},
|
||||
)
|
||||
|
||||
def _handle_error_response(self, response: httpx.Response, model_name: str) -> None:
|
||||
"""Handle error responses from the API."""
|
||||
try:
|
||||
error_data = response.json()
|
||||
error = error_data.get("error", {})
|
||||
error_code = error.get("code", response.status_code)
|
||||
error_message = error.get("message", response.text)
|
||||
error_status = error.get("status", "UNKNOWN")
|
||||
|
||||
# Handle rate limiting
|
||||
if error_code == 429:
|
||||
# Extract retry delay from response
|
||||
retry_delay = 60.0
|
||||
details = error.get("details", [])
|
||||
for detail in details:
|
||||
if detail.get("@type", "").endswith("RetryInfo"):
|
||||
delay_str = detail.get("retryDelay", "60s")
|
||||
# Parse delay string (e.g., "3.957525076s")
|
||||
match = re.match(r"([\d.]+)s?", delay_str)
|
||||
if match:
|
||||
retry_delay = float(match.group(1))
|
||||
break
|
||||
|
||||
# Mark current account as rate limited
|
||||
self._token_manager.mark_rate_limited(model_name, retry_delay)
|
||||
|
||||
raise RuntimeError(f"Rate limited (429): {error_message}. Retry after {retry_delay}s")
|
||||
|
||||
# Re-raise with formatted message
|
||||
raise RuntimeError(f"Antigravity API error {error_code} ({error_status}): {error_message}")
|
||||
|
||||
except json.JSONDecodeError:
|
||||
raise RuntimeError(f"Antigravity API error {response.status_code}: {response.text}")
|
||||
|
||||
def _process_image(self, image_path: str) -> dict[str, Any] | None:
|
||||
"""Process an image for the Antigravity API."""
|
||||
try:
|
||||
image_bytes, mime_type = validate_image(image_path)
|
||||
|
||||
if image_path.startswith("data:"):
|
||||
# Extract base64 data from data URL
|
||||
_, data = image_path.split(",", 1)
|
||||
return {"inlineData": {"mimeType": mime_type, "data": data}}
|
||||
else:
|
||||
# Encode file bytes
|
||||
image_data = base64.b64encode(image_bytes).decode()
|
||||
return {"inlineData": {"mimeType": mime_type, "data": image_data}}
|
||||
|
||||
except ValueError as e:
|
||||
logger.warning("Image validation failed: %s", e)
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.error("Error processing image %s: %s", image_path, e)
|
||||
return None
|
||||
|
||||
def _is_error_retryable(self, error: Exception) -> bool:
|
||||
"""Determine if an error should be retried."""
|
||||
error_str = str(error).lower()
|
||||
|
||||
# Rate limits should be retried (we'll rotate accounts)
|
||||
if "429" in error_str or "rate limit" in error_str:
|
||||
return True
|
||||
|
||||
# Check for retryable indicators
|
||||
retryable_indicators = [
|
||||
"timeout",
|
||||
"connection",
|
||||
"network",
|
||||
"temporary",
|
||||
"unavailable",
|
||||
"retry",
|
||||
"500",
|
||||
"502",
|
||||
"503",
|
||||
"504",
|
||||
]
|
||||
|
||||
return any(indicator in error_str for indicator in retryable_indicators)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Preferred model selection
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_preferred_model(self, category: ToolModelCategory, allowed_models: list[str]) -> str | None:
|
||||
"""Get preferred model for a category from allowed models."""
|
||||
from tools.models import ToolModelCategory
|
||||
|
||||
if not allowed_models:
|
||||
return None
|
||||
|
||||
capability_map = self.get_all_model_capabilities()
|
||||
|
||||
def find_best(candidates: list[str]) -> str | None:
|
||||
return sorted(candidates, reverse=True)[0] if candidates else None
|
||||
|
||||
if category == ToolModelCategory.EXTENDED_REASONING:
|
||||
# Prefer thinking-capable models
|
||||
thinking_models = [
|
||||
m for m in allowed_models if m in capability_map and capability_map[m].supports_extended_thinking
|
||||
]
|
||||
if thinking_models:
|
||||
# Prefer Claude Opus, then Claude Sonnet, then Gemini
|
||||
opus = [m for m in thinking_models if "opus" in m]
|
||||
if opus:
|
||||
return find_best(opus)
|
||||
sonnet = [m for m in thinking_models if "sonnet" in m]
|
||||
if sonnet:
|
||||
return find_best(sonnet)
|
||||
return find_best(thinking_models)
|
||||
|
||||
elif category == ToolModelCategory.FAST_RESPONSE:
|
||||
# Prefer Flash models
|
||||
flash_models = [m for m in allowed_models if "flash" in m]
|
||||
if flash_models:
|
||||
return find_best(flash_models)
|
||||
|
||||
# Default: prefer models by capability rank
|
||||
return find_best(allowed_models)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Resource cleanup
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def close(self) -> None:
|
||||
"""Clean up resources."""
|
||||
if self._http_client:
|
||||
self._http_client.close()
|
||||
self._http_client = None
|
||||
if self._token_manager:
|
||||
self._token_manager.close()
|
||||
Reference in New Issue
Block a user