Files
document-haness/build/lib/claridoc/providers/antigravity.py
T

93 lines
3.5 KiB
Python

from __future__ import annotations
import asyncio
import importlib.util
import inspect
import os
import threading
from contextlib import contextmanager
from pathlib import Path
from typing import Any, Iterator
from claridoc.providers.base import Provider, ProviderError, ProviderRequest, ProviderResponse, ProviderUnavailable
_CWD_LOCK = threading.Lock()
class AntigravityProvider(Provider):
"""Programmatic adapter for the Google Antigravity Python SDK."""
def generate(self, request: ProviderRequest) -> ProviderResponse:
try:
from google.antigravity import Agent, LocalAgentConfig # type: ignore[import-not-found]
except (ImportError, ModuleNotFoundError) as exc:
raise ProviderUnavailable(
"Google Antigravity SDK is not installed; install the optional 'antigravity' extra"
) from exc
config_values = self.spec.options.get("config", {})
if not isinstance(config_values, dict):
raise ProviderError("antigravity options.config must be an object")
if self.spec.model and "model" not in config_values:
config_values = {**config_values, "model": self.spec.model}
async def invoke() -> str:
try:
config = LocalAgentConfig(**config_values)
except TypeError as exc:
raise ProviderError(f"invalid Antigravity LocalAgentConfig options: {exc}") from exc
async with Agent(config) as agent:
response = await asyncio.wait_for(
agent.chat(request.prompt), timeout=self.spec.timeout_seconds
)
text_value = response.text()
if inspect.isawaitable(text_value):
text_value = await text_value
return str(text_value).strip()
# LocalAgentConfig operates on the current local environment. Serialize
# temporary cwd changes so concurrent threads cannot cross-contaminate runs.
try:
asyncio.get_running_loop()
except RuntimeError:
pass
else:
raise ProviderError("Antigravity provider must be called outside an active asyncio loop")
with _temporary_cwd(request.workdir):
try:
text = asyncio.run(invoke())
except (TimeoutError, asyncio.TimeoutError) as exc:
raise ProviderError(f"Antigravity timed out after {self.spec.timeout_seconds}s") from exc
except ProviderError:
raise
except Exception as exc:
raise ProviderError(f"Antigravity invocation failed: {exc}") from exc
if not text:
raise ProviderUnavailable("Antigravity returned an empty response")
return ProviderResponse(text=text, provider=self.name, model=self.spec.model, metadata={"mode": "sdk"})
def check(self) -> dict[str, Any]:
try:
available = importlib.util.find_spec("google.antigravity") is not None
except (ImportError, ModuleNotFoundError, ValueError):
available = False
return {
"provider": self.name,
"available": available,
"mode": "google-antigravity SDK",
"note": "Credentials and local agent access are verified only by a live invocation.",
}
@contextmanager
def _temporary_cwd(path: Path) -> Iterator[None]:
with _CWD_LOCK:
old = Path.cwd()
os.chdir(path)
try:
yield
finally:
os.chdir(old)