Add Antigravity 2.0 provider

This commit is contained in:
Jager Cooper 2026-05-31 13:54:10 -02:30
parent 3eb7e08538
commit 0344a07867
No known key found for this signature in database
9 changed files with 532 additions and 2 deletions

View file

@ -5,7 +5,7 @@
| Property | Details |
|-------|-------|
| Description | Server-side refresh-token adapters for subscription-backed inference providers. Clients call LiteLLM normally and LiteLLM refreshes provider OAuth tokens on the server. |
| Provider Routes on LiteLLM | `claude_max/` |
| Provider Routes on LiteLLM | `claude_max/`, `antigravity2/` |
| Supported Operations | `/chat/completions`, including streaming when the upstream provider supports it |
These providers are intended for private gateways where the LiteLLM server owns the subscription login and clients authenticate only to LiteLLM. Do not expose the server token files to client machines.
@ -53,7 +53,37 @@ export CLAUDE_MAX_API_BASE="https://api.anthropic.com/v1/messages"
## Antigravity 2.0
Antigravity 2.0 is not implemented by this provider page yet. The deprecated Cloud Code Assist / Gemini CLI `cloudcode-pa.googleapis.com/v1internal` gateway is intentionally not documented here because it is not the Antigravity 2.0 SDK/CLI runtime contract.
The `antigravity2/` provider uses Google's official Antigravity 2.0 Python SDK (`google-antigravity`) and its local agent runtime. This is intentionally separate from, and does not use, deprecated Gemini CLI, Cloud Code Assist, or internal gateway contracts.
Install the SDK in the LiteLLM server image because the PyPI wheels include the compiled Antigravity runtime binary:
```bash showLineNumbers title="Server dependency"
pip install google-antigravity
```
Authentication is managed on the LiteLLM server by the official runtime. The Antigravity CLI authenticates through the system keyring and falls back to Google Sign-In when no active session exists; the Python SDK can also use `GEMINI_API_KEY` / `ANTIGRAVITY2_API_KEY` or Vertex settings configured on the server. Clients should not send provider credentials.
```yaml showLineNumbers title="config.yaml"
model_list:
- model_name: frontier
litellm_params:
model: antigravity2/gemini-3.1-pro-preview
- model_name: antigravity-2-pro
litellm_params:
model: antigravity2/gemini-3.1-pro-preview
```
Optional server-side overrides:
```bash showLineNumbers title="Antigravity 2.0 overrides"
export ANTIGRAVITY2_APP_DATA_DIR="/var/lib/litellm/antigravity2"
export ANTIGRAVITY2_API_KEY="$GEMINI_API_KEY"
export ANTIGRAVITY2_VERTEX="false"
export ANTIGRAVITY2_PROJECT="my-gcp-project"
export ANTIGRAVITY2_LOCATION="us-central1"
```
For private inference-gateway usage, LiteLLM disables Antigravity built-in tools and subagents by default, so `antigravity2/` behaves like a chat inference provider rather than an autonomous file-editing/code-execution agent.
## LiteLLM virtual keys

View file

@ -132,6 +132,7 @@ LLM_CONFIG_NAMES = (
"DataRobotConfig",
"AnthropicConfig",
"ClaudeMaxConfig",
"Antigravity2Config",
"BedrockClaudePlatformConfig",
"AnthropicTextConfig",
"GroqSTTConfig",
@ -615,6 +616,10 @@ _LLM_CONFIGS_IMPORT_MAP = {
"DataRobotConfig": (".llms.datarobot.chat.transformation", "DataRobotConfig"),
"AnthropicConfig": (".llms.anthropic.chat.transformation", "AnthropicConfig"),
"ClaudeMaxConfig": (".llms.claude_max.chat.transformation", "ClaudeMaxConfig"),
"Antigravity2Config": (
".llms.antigravity2.chat.transformation",
"Antigravity2Config",
),
"BedrockClaudePlatformConfig": (
".llms.bedrock.claude_platform.transformation",
"BedrockClaudePlatformConfig",

View file

View file

@ -0,0 +1,324 @@
import asyncio
import os
import queue
import threading
from dataclasses import dataclass
from typing import Any, AsyncIterator, Iterator, List, Optional, Tuple, Union
import httpx
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
from litellm.types.utils import (
AllMessageValues,
GenericStreamingChunk,
ModelResponse,
Usage,
)
@dataclass
class Antigravity2SDK:
Agent: Any
LocalAgentConfig: Any
types: Any
class Antigravity2Config(BaseConfig):
"""LiteLLM provider for the official Antigravity 2.0 Python SDK.
Antigravity 2.0 is exposed by Google as a local agent runtime. Authentication
is owned by that runtime (system keyring / Google sign-in) or by Gemini API / Vertex
settings configured on the server, not by forwarding end-user API keys.
"""
@property
def custom_llm_provider(self) -> Optional[str]:
return "antigravity2"
def get_supported_openai_params(self, model: str) -> list:
return ["stream", "reasoning_effort"]
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
model: str,
drop_params: bool,
) -> dict:
if "stream" in non_default_params:
optional_params["stream"] = non_default_params["stream"]
if "reasoning_effort" in non_default_params:
optional_params["reasoning_effort"] = non_default_params["reasoning_effort"]
return optional_params
def validate_environment(
self,
headers: dict,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
) -> dict:
return headers
def transform_request(
self,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
headers: dict,
) -> dict:
return {"model": model, "messages": messages, **optional_params}
def transform_response(
self,
model: str,
raw_response: httpx.Response,
model_response: ModelResponse,
logging_obj: Any,
request_data: dict,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
encoding: Any,
api_key: Optional[str] = None,
json_mode: Optional[bool] = None,
) -> ModelResponse:
return model_response
def get_error_class(
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
) -> BaseLLMException:
return BaseLLMException(
status_code=status_code,
message=error_message,
headers=headers,
)
def _load_sdk(self) -> Antigravity2SDK:
try:
from google import antigravity as google_antigravity
from google.antigravity import types as ag_types
except ImportError as exc:
raise ImportError(
"The antigravity2 provider requires Google's official Antigravity 2.0 SDK. "
"Install it on the LiteLLM server with `pip install google-antigravity`."
) from exc
return Antigravity2SDK(
Agent=google_antigravity.Agent,
LocalAgentConfig=google_antigravity.LocalAgentConfig,
types=ag_types,
)
def _message_content_to_text(self, content: Any) -> str:
if content is None:
return ""
if isinstance(content, str):
return content
if isinstance(content, list):
parts: List[str] = []
for item in content:
if isinstance(item, dict):
if item.get("type") == "text":
parts.append(str(item.get("text", "")))
else:
parts.append(str(item))
else:
parts.append(str(item))
return "\n".join(part for part in parts if part)
return str(content)
def _split_system_and_prompt(self, messages: List[AllMessageValues]) -> Tuple[Optional[str], str]:
system_parts: List[str] = []
transcript_parts: List[str] = []
for message in messages:
role = str(message.get("role", "user"))
content = self._message_content_to_text(message.get("content"))
if role in {"system", "developer"}:
if content:
system_parts.append(content)
continue
transcript_parts.append(f"{role}: {content}")
prompt = "\n\n".join(transcript_parts).strip()
return ("\n\n".join(system_parts) or None, prompt)
def _thinking_level(self, sdk: Antigravity2SDK, optional_params: dict) -> Optional[Any]:
value = optional_params.get("reasoning_effort")
if value is None:
return None
mapping = {
"minimal": sdk.types.ThinkingLevel.MINIMAL,
"low": sdk.types.ThinkingLevel.LOW,
"medium": sdk.types.ThinkingLevel.MEDIUM,
"high": sdk.types.ThinkingLevel.HIGH,
}
return mapping.get(str(value).lower())
def _build_local_agent_config(
self,
sdk: Antigravity2SDK,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
api_key: Optional[str],
) -> Tuple[Any, str]:
system_instructions, prompt = self._split_system_and_prompt(messages)
thinking_level = self._thinking_level(sdk, optional_params)
generation = sdk.types.GenerationConfig(thinking_level=thinking_level)
model_entry = sdk.types.ModelEntry(name=model, api_key=api_key, generation=generation)
gemini_config = sdk.types.GeminiConfig(
api_key=api_key or os.getenv("ANTIGRAVITY2_API_KEY"),
vertex=os.getenv("ANTIGRAVITY2_VERTEX", "").lower() in {"1", "true", "yes"},
project=os.getenv("ANTIGRAVITY2_PROJECT"),
location=os.getenv("ANTIGRAVITY2_LOCATION"),
models=sdk.types.ModelConfig(default=model_entry),
)
capabilities = sdk.types.CapabilitiesConfig(
enabled_tools=sdk.types.BuiltinTools.none(),
enable_subagents=False,
)
kwargs = {
"system_instructions": system_instructions,
"capabilities": capabilities,
"policies": [],
"workspaces": [],
"gemini_config": gemini_config,
"model": model,
"api_key": api_key or os.getenv("ANTIGRAVITY2_API_KEY"),
}
app_data_dir = os.getenv("ANTIGRAVITY2_APP_DATA_DIR")
if app_data_dir:
kwargs["app_data_dir"] = app_data_dir
return sdk.LocalAgentConfig(**kwargs), prompt
def _usage_from_metadata(self, usage_metadata: Any) -> Usage:
if usage_metadata is None:
return Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0)
prompt_tokens = int(getattr(usage_metadata, "prompt_token_count", 0) or 0)
completion_tokens = int(getattr(usage_metadata, "candidates_token_count", 0) or 0)
total_tokens = int(getattr(usage_metadata, "total_token_count", 0) or (prompt_tokens + completion_tokens))
return Usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=total_tokens,
)
async def _achat_text(
self,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
api_key: Optional[str],
) -> Tuple[str, Usage]:
sdk = self._load_sdk()
config, prompt = self._build_local_agent_config(sdk, model, messages, optional_params, api_key)
async with sdk.Agent(config) as agent:
response = await agent.chat(prompt)
text = await response.text()
return text, self._usage_from_metadata(getattr(response, "usage_metadata", None))
async def _achat_stream(
self,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
api_key: Optional[str],
) -> AsyncIterator[GenericStreamingChunk]:
sdk = self._load_sdk()
config, prompt = self._build_local_agent_config(sdk, model, messages, optional_params, api_key)
async with sdk.Agent(config) as agent:
response = await agent.chat(prompt)
async for token in response:
yield {
"text": token,
"is_finished": False,
"finish_reason": "",
"usage": None,
}
yield {
"text": "",
"is_finished": True,
"finish_reason": "stop",
"usage": self._usage_from_metadata(getattr(response, "usage_metadata", None)),
}
def completion(
self,
model: str,
messages: List[AllMessageValues],
model_response: ModelResponse,
optional_params: dict,
api_key: Optional[str],
logging_obj: Any,
custom_llm_provider: str,
acompletion: bool = False,
) -> Any:
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
if optional_params.get("stream") is True:
stream = self._achat_stream(model, messages, optional_params, api_key)
if acompletion:
return CustomStreamWrapper(
stream, model=model, custom_llm_provider=custom_llm_provider, logging_obj=logging_obj
)
return CustomStreamWrapper(
self._sync_from_async_stream(stream),
model=model,
custom_llm_provider=custom_llm_provider,
logging_obj=logging_obj,
)
if acompletion:
return self._acompletion(model, messages, model_response, optional_params, api_key)
text, usage = self._run_sync(self._achat_text(model, messages, optional_params, api_key))
model_response.choices[0].message.content = text
model_response.usage = usage
return model_response
async def _acompletion(
self,
model: str,
messages: List[AllMessageValues],
model_response: ModelResponse,
optional_params: dict,
api_key: Optional[str],
) -> ModelResponse:
text, usage = await self._achat_text(model, messages, optional_params, api_key)
model_response.choices[0].message.content = text
model_response.usage = usage
return model_response
def _run_sync(self, coro: Any) -> Any:
try:
asyncio.get_running_loop()
except RuntimeError:
return asyncio.run(coro)
raise RuntimeError("Use litellm.acompletion() for antigravity2 calls from an active event loop")
def _sync_from_async_stream(
self, async_stream: AsyncIterator[GenericStreamingChunk]
) -> Iterator[GenericStreamingChunk]:
items: "queue.Queue[Any]" = queue.Queue()
sentinel = object()
async def produce() -> None:
try:
async for chunk in async_stream:
items.put(chunk)
except BaseException as exc:
items.put(exc)
finally:
items.put(sentinel)
thread = threading.Thread(target=lambda: asyncio.run(produce()), daemon=True)
thread.start()
while True:
item = items.get()
if item is sentinel:
break
if isinstance(item, BaseException):
raise item
yield item

View file

@ -4477,6 +4477,20 @@ def completion( # type: ignore # noqa: PLR0915
logging_obj=logging,
)
elif custom_llm_provider == "antigravity2":
from litellm.llms.antigravity2.chat.transformation import Antigravity2Config
response = Antigravity2Config().completion(
model=model,
messages=messages,
model_response=model_response,
optional_params=optional_params,
api_key=api_key,
logging_obj=logging,
custom_llm_provider=custom_llm_provider,
acompletion=acompletion,
)
elif custom_llm_provider == "langgraph":
# LangGraph - Agent Runtime Provider
from litellm.llms.langgraph.chat.transformation import LangGraphConfig

View file

@ -3235,6 +3235,7 @@ class LlmProviders(str, Enum):
OPENAI = "openai"
CHATGPT = "chatgpt"
CLAUDE_MAX = "claude_max"
ANTIGRAVITY2 = "antigravity2"
OPENAI_LIKE = "openai_like" # embedding only
JINA_AI = "jina_ai"
XAI = "xai"

View file

@ -8261,6 +8261,7 @@ class ProviderConfigManager:
LlmProviders.GITHUB_COPILOT: (lambda: litellm.GithubCopilotConfig(), False),
LlmProviders.CHATGPT: (lambda: litellm.ChatGPTConfig(), False),
LlmProviders.CLAUDE_MAX: (lambda: litellm.ClaudeMaxConfig(), False),
LlmProviders.ANTIGRAVITY2: (lambda: litellm.Antigravity2Config(), False),
LlmProviders.GIGACHAT: (lambda: litellm.GigaChatConfig(), False),
LlmProviders.RAGFLOW: (lambda: litellm.RAGFlowConfig(), False),
LlmProviders.CUSTOM: (lambda: litellm.OpenAILikeChatConfig(), False),

View file

@ -0,0 +1,155 @@
import asyncio
from types import SimpleNamespace
from litellm.llms.antigravity2.chat.transformation import Antigravity2Config, Antigravity2SDK
from litellm.types.utils import LlmProviders, ModelResponse
from litellm.utils import ProviderConfigManager
class _Obj:
def __init__(self, **kwargs):
self.__dict__.update(kwargs)
class _FakeBuiltinTools:
@classmethod
def none(cls):
return []
class _FakeTypes:
class ThinkingLevel:
MINIMAL = "minimal"
LOW = "low"
MEDIUM = "medium"
HIGH = "high"
GenerationConfig = _Obj
ModelEntry = _Obj
ModelConfig = _Obj
GeminiConfig = _Obj
CapabilitiesConfig = _Obj
BuiltinTools = _FakeBuiltinTools
class _FakeLocalAgentConfig(_Obj):
pass
class _FakeResponse:
usage_metadata = SimpleNamespace(
prompt_token_count=3,
candidates_token_count=2,
total_token_count=5,
)
async def text(self):
return "hello from ag2"
def __aiter__(self):
async def _tokens():
yield "hel"
yield "lo"
return _tokens()
class _FakeAgent:
last_config = None
last_prompt = None
def __init__(self, config):
_FakeAgent.last_config = config
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc, tb):
return None
async def chat(self, prompt):
_FakeAgent.last_prompt = prompt
return _FakeResponse()
class _FakeAntigravity2Config(Antigravity2Config):
def _load_sdk(self):
return Antigravity2SDK(
Agent=_FakeAgent,
LocalAgentConfig=_FakeLocalAgentConfig,
types=_FakeTypes,
)
def test_antigravity2_provider_config_registration():
config = ProviderConfigManager.get_provider_chat_config("gemini-3.1-pro-preview", LlmProviders.ANTIGRAVITY2)
assert isinstance(config, Antigravity2Config)
def test_build_local_agent_config_uses_official_sdk_contract_and_disables_tools(monkeypatch):
monkeypatch.setenv("ANTIGRAVITY2_APP_DATA_DIR", "/srv/ag2")
config = _FakeAntigravity2Config()
sdk = config._load_sdk()
local_config, prompt = config._build_local_agent_config(
sdk=sdk,
model="gemini-3.1-pro-preview",
messages=[
{"role": "system", "content": "be brief"},
{"role": "user", "content": "hi"},
],
optional_params={"reasoning_effort": "high"},
api_key=None,
)
assert local_config.system_instructions == "be brief"
assert prompt == "user: hi"
assert local_config.model == "gemini-3.1-pro-preview"
assert local_config.capabilities.enabled_tools == []
assert local_config.capabilities.enable_subagents is False
assert local_config.app_data_dir == "/srv/ag2"
assert local_config.gemini_config.models.default.generation.thinking_level == "high"
def test_antigravity2_async_completion_maps_text_and_usage():
config = _FakeAntigravity2Config()
model_response = ModelResponse()
result = asyncio.run(
config._acompletion(
model="gemini-3.1-pro-preview",
messages=[{"role": "user", "content": "say hello"}],
model_response=model_response,
optional_params={},
api_key=None,
)
)
assert _FakeAgent.last_prompt == "user: say hello"
assert result.choices[0].message.content == "hello from ag2"
assert result.usage.prompt_tokens == 3
assert result.usage.completion_tokens == 2
assert result.usage.total_tokens == 5
def test_antigravity2_stream_yields_generic_chunks():
config = _FakeAntigravity2Config()
async def collect():
return [
chunk
async for chunk in config._achat_stream(
model="gemini-3.1-pro-preview",
messages=[{"role": "user", "content": "stream"}],
optional_params={},
api_key=None,
)
]
chunks = asyncio.run(collect())
assert [chunk["text"] for chunk in chunks] == ["hel", "lo", ""]
assert chunks[-1]["is_finished"] is True
assert chunks[-1]["finish_reason"] == "stop"
assert chunks[-1]["usage"].total_tokens == 5