Torbjørn Lindahl
5add230d4c
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
57 lines
2.2 KiB
Python
57 lines
2.2 KiB
Python
"""Antigravity model registry for managing model configurations and aliases."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from ..shared import ModelCapabilities, ProviderType
|
|
from .base import CAPABILITY_FIELD_NAMES, CapabilityModelRegistry
|
|
|
|
|
|
class AntigravityModelRegistry(CapabilityModelRegistry):
|
|
"""Capability registry backed by ``conf/antigravity_models.json``.
|
|
|
|
The Antigravity provider accesses Claude, Gemini, and other models through
|
|
Google's unified gateway API. Models are split across two quota pools:
|
|
|
|
* **antigravity**: Claude models and Gemini 3 variants
|
|
* **gemini-cli**: Gemini 2.5 and preview models using Gemini CLI quota
|
|
"""
|
|
|
|
# Extra fields beyond standard ModelCapabilities
|
|
EXTRA_KEYS = {"quota_pool"}
|
|
|
|
def __init__(self, config_path: str | None = None) -> None:
|
|
super().__init__(
|
|
env_var_name="ANTIGRAVITY_MODELS_CONFIG_PATH",
|
|
default_filename="antigravity_models.json",
|
|
provider=ProviderType.ANTIGRAVITY,
|
|
friendly_prefix="Antigravity ({model})",
|
|
config_path=config_path,
|
|
)
|
|
|
|
def _extra_keys(self) -> set[str]:
|
|
"""Include quota_pool as a valid entry key."""
|
|
return self.EXTRA_KEYS
|
|
|
|
def _finalise_entry(self, entry: dict) -> tuple[ModelCapabilities, dict]:
|
|
"""Build ModelCapabilities and extract extras like quota_pool."""
|
|
# Extract extras before building capabilities
|
|
extras = {}
|
|
quota_pool = entry.pop("quota_pool", None)
|
|
if quota_pool:
|
|
extras["quota_pool"] = quota_pool
|
|
|
|
# Build capabilities from remaining fields
|
|
filtered = {k: v for k, v in entry.items() if k in CAPABILITY_FIELD_NAMES}
|
|
filtered.setdefault("provider", ProviderType.ANTIGRAVITY)
|
|
filtered.setdefault("friendly_name", f"Antigravity ({entry.get('model_name', 'unknown')})")
|
|
capability = ModelCapabilities(**filtered)
|
|
|
|
return capability, extras
|
|
|
|
def get_quota_pool(self, model_name: str) -> str:
|
|
"""Return the quota pool for a model ('antigravity' or 'gemini-cli')."""
|
|
extras = self.get_entry(model_name)
|
|
if extras and "quota_pool" in extras:
|
|
return extras["quota_pool"]
|
|
return "antigravity" # Default to antigravity quota pool
|