diff --git a/docs/my-website/docs/providers/server_managed_oauth_subscriptions.md b/docs/my-website/docs/providers/server_managed_oauth_subscriptions.md index f9e29630296..9c24890e587 100644 --- a/docs/my-website/docs/providers/server_managed_oauth_subscriptions.md +++ b/docs/my-website/docs/providers/server_managed_oauth_subscriptions.md @@ -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 diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index 5b573ac3ebd..01dc34baf67 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -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", diff --git a/litellm/llms/antigravity2/__init__.py b/litellm/llms/antigravity2/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/antigravity2/chat/__init__.py b/litellm/llms/antigravity2/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/antigravity2/chat/transformation.py b/litellm/llms/antigravity2/chat/transformation.py new file mode 100644 index 00000000000..5835d20b82b --- /dev/null +++ b/litellm/llms/antigravity2/chat/transformation.py @@ -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 diff --git a/litellm/main.py b/litellm/main.py index 8c884ddba01..a31cac243bd 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -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 diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 75c82b1476e..9cc6b8e9111 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -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" diff --git a/litellm/utils.py b/litellm/utils.py index cdfce3a1506..136e2db063b 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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), diff --git a/tests/test_litellm/llms/antigravity2/test_antigravity2.py b/tests/test_litellm/llms/antigravity2/test_antigravity2.py new file mode 100644 index 00000000000..31ad5d71f38 --- /dev/null +++ b/tests/test_litellm/llms/antigravity2/test_antigravity2.py @@ -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