mirror of
https://github.com/usamashehab/antigravity-proxy.git
synced 2026-10-25 01:14:18 +00:00
usamashehab
3e664884af
- OpenAI /v1/chat/completions endpoint (streaming + non-streaming)
- Native Gemini passthrough /v1beta/models/{model}:generateContent endpoint
- Tool/function calling support with round-trip encoding
- Auto model name mapping (gemini-3.5-flash -> gemini-3.5-flash-low)
- OAuth token auto-refresh
- Zero pip dependencies (Python 3.10+ stdlib)
- systemd service file included
- Models: Gemini 3.1 Pro, 3 Flash, 3.5 Flash, 2.5 Pro/Flash, Claude Sonnet/Opus 4.6
1716 lines
67 KiB
Python
1716 lines
67 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Antigravity OpenAI-Compatible Proxy (Cloud Code Assist API)
|
|
|
|
Bridges Hermes Agent (OpenAI-compatible client) to Google Antigravity's Cloud
|
|
Code Assist API. Uses the OAuth token managed by the Antigravity CLI (agy) and
|
|
talks directly to https://cloudcode-pa.googleapis.com — no `agy` subprocess
|
|
needed, giving us native streaming + tool/function-call support.
|
|
|
|
Architecture:
|
|
1. OAuth token read from ~/.gemini/antigravity-cli/antigravity-oauth-token
|
|
Refreshed via https://oauth2.googleapis.com/token when expired.
|
|
2. Project ID discovered via v1internal:loadCodeAssist (cached).
|
|
3. Requests POSTed to v1internal:generateContent (or streamGenerateContent),
|
|
wrapped in an envelope: {project, model, request, requestType, ...}
|
|
4. Responses unwrapped from {"response": {<Gemini response>}} and transformed
|
|
to OpenAI chat.completion format.
|
|
|
|
Usage:
|
|
python3 antigravity_proxy.py [--port 8877] [--host 127.0.0.1]
|
|
|
|
Then configure Hermes:
|
|
hermes config set model.provider custom
|
|
hermes config set model.base_url http://127.0.0.1:8877/v1
|
|
hermes config set model.api_key antigravity
|
|
hermes config set model.default gemini-3.1-pro
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import re
|
|
import sys
|
|
import threading
|
|
import time
|
|
import traceback
|
|
import uuid
|
|
from datetime import datetime, timezone
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
from urllib.parse import urlparse
|
|
from urllib.request import Request, urlopen
|
|
from urllib.error import HTTPError, URLError
|
|
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
# Constants
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
|
|
# OAuth client credentials for Antigravity CLI (reverse-engineered from agy).
|
|
# These are the same values embedded in every Antigravity CLI binary — not
|
|
# personal secrets. Set them via env vars if you prefer, or see README.md
|
|
# for the default values to put in a .env file.
|
|
CLIENT_ID = os.environ.get(
|
|
"ANTIGRAVITY_CLIENT_ID",
|
|
"1071006060591-tmhssin2h21lcre2" "35vtolojh4g403ep.apps.googleusercontent.com",
|
|
)
|
|
CLIENT_SECRET = os.environ.get(
|
|
"ANTIGRAVITY_CLIENT_SECRET",
|
|
"GOCSPX-K58FWR486LdL" "J1mLB8sXC4z6qDAf",
|
|
)
|
|
|
|
TOKEN_FILE = os.path.expanduser("~/.gemini/antigravity-cli/antigravity-oauth-token")
|
|
OAUTH_TOKEN_URL = "https://oauth2.googleapis.com/token"
|
|
CLOUDCODE_BASE = "https://cloudcode-pa.googleapis.com"
|
|
LOAD_CODEASSIST_URL = f"{CLOUDCODE_BASE}/v1internal:loadCodeAssist"
|
|
GENERATE_CONTENT_URL = f"{CLOUDCODE_BASE}/v1internal:generateContent"
|
|
STREAM_GENERATE_CONTENT_URL = f"{CLOUDCODE_BASE}/v1internal:streamGenerateContent?alt=sse"
|
|
|
|
ANTIGRAVITY_USER_AGENT = (
|
|
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
|
|
"AppleWebKit/537.36 (KHTML, like Gecko) "
|
|
"Antigravity/1.0.14 Chrome/138.0.7204.235 Electron/37.3.1 Safari/537.36"
|
|
)
|
|
CLIENT_METADATA = json.dumps(
|
|
{"ideType": "ANTIGRAVITY", "platform": "MACOS", "pluginType": "GEMINI"},
|
|
separators=(",", ":"),
|
|
)
|
|
|
|
# Minimal system instruction — just enough for the Code Assist API to accept
|
|
# the request. The full Antigravity identity instruction causes the model to
|
|
# emit "unsupported version" warnings, so we keep it short.
|
|
DEFAULT_SYSTEM_INSTRUCTION = "You are a helpful AI assistant."
|
|
|
|
UPSTREAM_TIMEOUT = 300 # 5 minutes
|
|
TOKEN_REFRESH_SKEW = 120 # refresh this many seconds before actual expiry
|
|
|
|
# Model mapping: OpenAI-facing name -> Antigravity backend model name.
|
|
# Tested against cloudcode-pa.googleapis.com on 2026-06-30.
|
|
# Model mapping: OpenAI-facing name -> (backend model name, thinking level).
|
|
# Gemini 3 Pro models use the base name "gemini-3.1-pro-low" and control
|
|
# thinking compute via generationConfig.thinkingConfig.thinkingLevel.
|
|
MODEL_MAP = {
|
|
"gemini-3.1-pro": ("gemini-3.1-pro-low", "low"),
|
|
"gemini-3.1-pro-high": ("gemini-3.1-pro-low", "high"),
|
|
"gemini-3-flash": ("gemini-3-flash", None),
|
|
"gemini-3.5-flash": ("gemini-3.5-flash-low", None),
|
|
"gemini-2.5-pro": ("gemini-2.5-pro", None),
|
|
"gemini-2.5-flash": ("gemini-2.5-flash", None),
|
|
"claude-sonnet-4.6": ("claude-sonnet-4-6", None),
|
|
"claude-opus-4.6": ("claude-opus-4-6-thinking", None),
|
|
}
|
|
|
|
DEFAULT_MODEL = "gemini-2.5-flash"
|
|
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
# Token / OAuth management (thread-safe)
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
|
|
_token_lock = threading.Lock()
|
|
_cached_access_token: str = None
|
|
_cached_project_id: str = None
|
|
_project_lock = threading.Lock()
|
|
|
|
|
|
def _log(msg: str):
|
|
print(f"[antigravity-proxy] {msg}", file=sys.stderr, flush=True)
|
|
|
|
|
|
def _parse_expiry(expiry_raw):
|
|
"""Parse the 'expiry' field from the token file into a unix timestamp.
|
|
|
|
The Antigravity CLI stores expiry as an RFC3339 string like
|
|
"2026-06-30T12:55:03.123456789Z". Python stdlib datetime.fromisoformat
|
|
can't handle nanosecond precision, so we truncate to 6 digits.
|
|
"""
|
|
if not expiry_raw:
|
|
return 0.0
|
|
if isinstance(expiry_raw, (int, float)):
|
|
return float(expiry_raw)
|
|
s = str(expiry_raw).strip()
|
|
# Truncate sub-second precision to microseconds for fromisoformat.
|
|
try:
|
|
if "." in s:
|
|
head, tail = s.split(".", 1)
|
|
# Keep only the part up to 'Z' or timezone offset, truncate ns.
|
|
tz_part = ""
|
|
for i, ch in enumerate(tail):
|
|
if ch in "Z+-":
|
|
tz_part = tail[i:]
|
|
tail = tail[:i]
|
|
break
|
|
tail = tail[:6] # microseconds max
|
|
s = f"{head}.{tail}{tz_part}"
|
|
# fromisoformat in 3.12 handles 'Z' suffix.
|
|
dt = datetime.fromisoformat(s)
|
|
if dt.tzinfo is None:
|
|
dt = dt.replace(tzinfo=timezone.utc)
|
|
return dt.timestamp()
|
|
except Exception:
|
|
return 0.0
|
|
|
|
|
|
def _read_token_from_disk():
|
|
"""Read and parse the token file. Returns the raw dict or None."""
|
|
try:
|
|
with open(TOKEN_FILE, "r") as f:
|
|
return json.load(f)
|
|
except (OSError, json.JSONDecodeError) as e:
|
|
_log(f"WARNING: could not read token file {TOKEN_FILE}: {e}")
|
|
return None
|
|
|
|
|
|
def _write_token_to_disk(data: dict):
|
|
"""Persist updated token data back to disk (atomic-ish)."""
|
|
try:
|
|
tmp = TOKEN_FILE + ".tmp"
|
|
with open(tmp, "w") as f:
|
|
json.dump(data, f)
|
|
os.replace(tmp, TOKEN_FILE)
|
|
try:
|
|
os.chmod(TOKEN_FILE, 0o600)
|
|
except OSError:
|
|
pass
|
|
except OSError as e:
|
|
_log(f"WARNING: could not write token file: {e}")
|
|
|
|
|
|
def _refresh_access_token(refresh_token: str) -> dict:
|
|
"""Refresh the access token via Google's OAuth endpoint.
|
|
|
|
Returns the new token dict {access_token, refresh_token?, expiry, token_type}.
|
|
Raises RuntimeError on failure.
|
|
"""
|
|
body = json.dumps({
|
|
"client_id": CLIENT_ID,
|
|
"client_secret": CLIENT_SECRET,
|
|
"refresh_token": refresh_token,
|
|
"grant_type": "refresh_token",
|
|
}).encode()
|
|
req = Request(OAUTH_TOKEN_URL, data=body, method="POST")
|
|
req.add_header("Content-Type", "application/json")
|
|
try:
|
|
with urlopen(req, timeout=30) as resp:
|
|
payload = json.loads(resp.read().decode())
|
|
except HTTPError as e:
|
|
detail = ""
|
|
try:
|
|
detail = e.read().decode()
|
|
except Exception:
|
|
pass
|
|
raise RuntimeError(f"OAuth refresh failed (HTTP {e.code}): {detail}")
|
|
except URLError as e:
|
|
raise RuntimeError(f"OAuth refresh network error: {e}")
|
|
|
|
new_tok = {
|
|
"access_token": payload["access_token"],
|
|
"refresh_token": payload.get("refresh_token", refresh_token),
|
|
"token_type": payload.get("token_type", "Bearer"),
|
|
"expiry": payload.get("expires_in", 3600),
|
|
}
|
|
# Convert expires_in (seconds) to an RFC3339 expiry timestamp ~now.
|
|
expires_in = int(payload.get("expires_in", 3600))
|
|
exp_dt = datetime.now(timezone.utc)
|
|
from datetime import timedelta
|
|
exp_dt = exp_dt + timedelta(seconds=expires_in)
|
|
new_tok["expiry"] = exp_dt.strftime("%Y-%m-%dT%H:%M:%S.000000Z")
|
|
return new_tok
|
|
|
|
|
|
def get_access_token() -> str:
|
|
"""Return a valid access token, refreshing if necessary. Thread-safe."""
|
|
global _cached_access_token
|
|
with _token_lock:
|
|
data = _read_token_from_disk()
|
|
if not data:
|
|
raise RuntimeError(
|
|
f"No OAuth token found at {TOKEN_FILE}. "
|
|
"Run `agy` to authenticate first."
|
|
)
|
|
token_obj = data.get("token") or {}
|
|
access_token = token_obj.get("access_token")
|
|
expiry_ts = _parse_expiry(token_obj.get("expiry"))
|
|
now = time.time()
|
|
|
|
needs_refresh = (
|
|
not access_token
|
|
or expiry_ts == 0.0
|
|
or (expiry_ts - now) < TOKEN_REFRESH_SKEW
|
|
)
|
|
|
|
if needs_refresh:
|
|
refresh_token = token_obj.get("refresh_token")
|
|
if not refresh_token:
|
|
raise RuntimeError("No refresh_token available; re-run `agy` to authenticate.")
|
|
_log("Access token expired (or missing) — refreshing...")
|
|
new_tok = _refresh_access_token(refresh_token)
|
|
data["token"] = new_tok
|
|
_write_token_to_disk(data)
|
|
access_token = new_tok["access_token"]
|
|
|
|
_cached_access_token = access_token
|
|
return access_token
|
|
|
|
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
# Project discovery (cached)
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
|
|
def _load_code_assist(access_token: str) -> str:
|
|
"""Discover the cloudaicompanion project ID. Returns project id string.
|
|
|
|
NOTE: The metadata field (ideType/platform/pluginType) is rejected by the
|
|
API with INVALID_ARGUMENT for all known enum string values. Sending an
|
|
empty body {} works and returns the cloudaicompanionProject. The Client-Metadata
|
|
header still carries the ide/plugin info.
|
|
"""
|
|
body = json.dumps({}).encode()
|
|
req = Request(LOAD_CODEASSIST_URL, data=body, method="POST")
|
|
req.add_header("Authorization", f"Bearer {access_token}")
|
|
req.add_header("Content-Type", "application/json")
|
|
req.add_header("User-Agent", ANTIGRAVITY_USER_AGENT)
|
|
req.add_header("X-Goog-Api-Client", "google-cloud-sdk vscode_cloudshelleditor/0.1")
|
|
req.add_header("Client-Metadata", CLIENT_METADATA)
|
|
try:
|
|
with urlopen(req, timeout=60) as resp:
|
|
payload = json.loads(resp.read().decode())
|
|
except HTTPError as e:
|
|
detail = ""
|
|
try:
|
|
detail = e.read().decode()
|
|
except Exception:
|
|
pass
|
|
raise RuntimeError(f"loadCodeAssist failed (HTTP {e.code}): {detail}")
|
|
except URLError as e:
|
|
raise RuntimeError(f"loadCodeAssist network error: {e}")
|
|
|
|
project_id = (
|
|
payload.get("cloudaicompanionProject")
|
|
or payload.get("cloudaicompanion_project")
|
|
)
|
|
if not project_id:
|
|
# Some responses nest it differently; do a recursive search.
|
|
project_id = _deep_find(payload, "cloudaicompanionProject")
|
|
if not project_id:
|
|
raise RuntimeError(
|
|
f"loadCodeAssist did not return a project ID. Response: {payload}"
|
|
)
|
|
return project_id
|
|
|
|
|
|
def _deep_find(obj, key):
|
|
"""Recursively search for a key in nested dicts/lists. Returns first match."""
|
|
if isinstance(obj, dict):
|
|
if key in obj:
|
|
return obj[key]
|
|
for v in obj.values():
|
|
r = _deep_find(v, key)
|
|
if r is not None:
|
|
return r
|
|
elif isinstance(obj, list):
|
|
for item in obj:
|
|
r = _deep_find(item, key)
|
|
if r is not None:
|
|
return r
|
|
return None
|
|
|
|
|
|
def get_project_id() -> str:
|
|
"""Return cached project id, discovering it if needed. Thread-safe."""
|
|
global _cached_project_id
|
|
with _project_lock:
|
|
if _cached_project_id:
|
|
return _cached_project_id
|
|
token = get_access_token()
|
|
pid = _load_code_assist(token)
|
|
_cached_project_id = pid
|
|
_log(f"Discovered project ID: {pid}")
|
|
return pid
|
|
|
|
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
# OpenAI -> Gemini message transformation
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
|
|
def _content_to_text(content) -> str:
|
|
"""Normalize an OpenAI message 'content' field to plain text.
|
|
|
|
Handles str, list of content blocks [{type: text, text: ...}, ...], and None.
|
|
Non-text blocks (images, etc.) are represented as a placeholder marker so
|
|
the model knows something was there even if we can't inline binary data.
|
|
"""
|
|
if content is None:
|
|
return ""
|
|
if isinstance(content, str):
|
|
return content
|
|
if isinstance(content, list):
|
|
chunks = []
|
|
for block in content:
|
|
if not isinstance(block, dict):
|
|
chunks.append(str(block))
|
|
continue
|
|
btype = block.get("type", "text")
|
|
if btype == "text":
|
|
chunks.append(block.get("text", ""))
|
|
elif btype == "image_url":
|
|
url = ""
|
|
iu = block.get("image_url")
|
|
if isinstance(iu, dict):
|
|
url = iu.get("url", "")
|
|
elif isinstance(iu, str):
|
|
url = iu
|
|
# We pass the data URL through; Gemini's inlineData supports it.
|
|
if url.startswith("data:"):
|
|
# Best-effort: note its presence; full image passthrough
|
|
# would require parts[].inlineData. Keep as text marker.
|
|
chunks.append(f"[image: {url[:60]}...]")
|
|
else:
|
|
chunks.append(f"[image: {url}]")
|
|
elif btype == "input_text":
|
|
chunks.append(block.get("text", ""))
|
|
elif btype == "input_image":
|
|
chunks.append("[image provided]")
|
|
else:
|
|
chunks.append(f"[{btype}: {json.dumps(block.get('text', ''))[:100]}]")
|
|
return "\n".join(c for c in chunks if c)
|
|
return str(content)
|
|
|
|
|
|
def _openai_tools_to_gemini(tools: list) -> list:
|
|
"""Convert OpenAI tools array to Gemini functionDeclarations format.
|
|
|
|
OpenAI: [{"type":"function","function":{"name","description","parameters"}}]
|
|
Gemini: [{"functionDeclarations":[{"name","description","parameters"}]}]
|
|
"""
|
|
if not tools:
|
|
return []
|
|
decls = []
|
|
for tool in tools:
|
|
if not isinstance(tool, dict):
|
|
continue
|
|
func = tool.get("function") if tool.get("type") == "function" else tool
|
|
if not isinstance(func, dict):
|
|
continue
|
|
name = func.get("name")
|
|
if not name:
|
|
continue
|
|
decl = {
|
|
"name": name,
|
|
"description": func.get("description", ""),
|
|
}
|
|
params = func.get("parameters")
|
|
if isinstance(params, dict) and params:
|
|
decl["parameters"] = params
|
|
decls.append(decl)
|
|
if not decls:
|
|
return []
|
|
return [{"functionDeclarations": decls}]
|
|
|
|
|
|
def _tool_choice_to_gemini(tool_choice) -> str | None:
|
|
"""Convert OpenAI tool_choice to Gemini toolConfig.functionCallingConfig.mode.
|
|
|
|
Returns None if no transformation applies.
|
|
"""
|
|
if tool_choice is None:
|
|
return None
|
|
if tool_choice == "auto":
|
|
return "AUTO"
|
|
if tool_choice == "none":
|
|
return "NONE"
|
|
if tool_choice == "required":
|
|
return "ANY"
|
|
if isinstance(tool_choice, dict) and tool_choice.get("type") == "function":
|
|
# Specific function forced — Gemini uses ANY + allowed_function_names.
|
|
return "ANY"
|
|
return None
|
|
|
|
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
# Gemini-native passthrough (for invoice_reader.py / any Gemini SDK client)
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
|
|
# Matches /v1beta/models/{model}:generateContent
|
|
_GEMINI_PASSTHROUGH_RE = re.compile(
|
|
r"^/v1beta/models/(.+):generateContent$"
|
|
)
|
|
|
|
|
|
def _build_passthrough_envelope(model: str, gemini_body: dict) -> dict:
|
|
"""Wrap a native Gemini generateContent request in the Antigravity
|
|
envelope. The caller's body is used verbatim as ``request`` — no
|
|
message-format conversion, so inline_data / responseSchema / thinkingConfig
|
|
all pass straight through.
|
|
|
|
The friendly model name (e.g. ``gemini-3.5-flash``) is translated to the
|
|
Antigravity backend name (e.g. ``gemini-3.5-flash-low``) via MODEL_MAP so
|
|
callers can use the same name as Google's public API.
|
|
|
|
Safety settings are injected only when the caller did not supply its own.
|
|
"""
|
|
# Translate friendly model name to Antigravity backend name.
|
|
entry = MODEL_MAP.get(model)
|
|
backend_model = entry[0] if entry else model
|
|
thinking_level = entry[1] if entry else None
|
|
|
|
inner = dict(gemini_body) # shallow copy — don't mutate caller's dict
|
|
|
|
# Apply thinking level from MODEL_MAP if the caller didn't set one.
|
|
if thinking_level:
|
|
gc = inner.setdefault("generationConfig", {})
|
|
tc = gc.setdefault("thinkingConfig", {})
|
|
tc.setdefault("thinkingLevel", thinking_level)
|
|
|
|
if "safetySettings" not in inner:
|
|
inner["safetySettings"] = [
|
|
{"category": cat, "threshold": "BLOCK_NONE"}
|
|
for cat in (
|
|
"HARM_CATEGORY_HARASSMENT",
|
|
"HARM_CATEGORY_HATE_SPEECH",
|
|
"HARM_CATEGORY_SEXUALLY_EXPLICIT",
|
|
"HARM_CATEGORY_DANGEROUS_CONTENT",
|
|
)
|
|
]
|
|
|
|
project_id = get_project_id()
|
|
return {
|
|
"project": project_id,
|
|
"model": backend_model,
|
|
"request": inner,
|
|
"requestType": "agent",
|
|
"userAgent": "antigravity",
|
|
"requestId": f"agent-{uuid.uuid4().hex}",
|
|
}
|
|
|
|
|
|
def _build_gemini_request(body: dict, backend_model: str, thinking_level: str = None):
|
|
"""Transform an OpenAI chat completion body into a Gemini generateContent
|
|
request envelope.
|
|
|
|
Returns (envelope_dict, openai_model_name).
|
|
"""
|
|
messages = body.get("messages", [])
|
|
if not isinstance(messages, list):
|
|
messages = []
|
|
|
|
system_texts = []
|
|
contents = []
|
|
|
|
# Collect system instruction pieces and convert messages.
|
|
for msg in messages:
|
|
role = msg.get("role", "user")
|
|
|
|
if role == "system":
|
|
text = _content_to_text(msg.get("content"))
|
|
if text:
|
|
system_texts.append(text)
|
|
continue
|
|
|
|
if role == "tool":
|
|
# OpenAI tool result message -> Gemini functionResponse part.
|
|
tool_call_id = msg.get("tool_call_id", "")
|
|
content_text = _content_to_text(msg.get("content"))
|
|
# Try to parse JSON from the tool result; Gemini wants an object.
|
|
response_obj = {}
|
|
if content_text:
|
|
try:
|
|
parsed = json.loads(content_text)
|
|
response_obj = parsed if isinstance(parsed, dict) else {"result": parsed}
|
|
except (json.JSONDecodeError, TypeError):
|
|
response_obj = {"result": content_text}
|
|
# Derive a function name from tool_call_id if the tool name isn't given.
|
|
fn_name = msg.get("name") or tool_call_id or "tool_result"
|
|
# Decode the fc_id from the encoded tool_call_id (format: call_<fc_id>|<sig>).
|
|
fc_response_id = ""
|
|
if tool_call_id.startswith("call_"):
|
|
raw = tool_call_id[5:] # strip "call_"
|
|
if "|" in raw:
|
|
fc_response_id = raw.split("|", 1)[0]
|
|
else:
|
|
fc_response_id = raw
|
|
func_resp = {
|
|
"name": fn_name,
|
|
"response": response_obj,
|
|
}
|
|
if fc_response_id:
|
|
func_resp["id"] = fc_response_id
|
|
contents.append({
|
|
"role": "function",
|
|
"parts": [{"functionResponse": func_resp}],
|
|
})
|
|
continue
|
|
|
|
if role == "assistant":
|
|
parts = []
|
|
text = _content_to_text(msg.get("content"))
|
|
if text:
|
|
parts.append({"text": text})
|
|
# Prior tool calls from the assistant -> functionCall parts.
|
|
# Decode the encoded tool_call.id to recover the Gemini fc_id and thoughtSignature.
|
|
tool_calls = msg.get("tool_calls") or []
|
|
for tc in tool_calls:
|
|
if not isinstance(tc, dict):
|
|
continue
|
|
fn = tc.get("function") or {}
|
|
fn_name = fn.get("name", "")
|
|
args_str = fn.get("arguments", "{}")
|
|
try:
|
|
args = json.loads(args_str) if isinstance(args_str, str) else args_str
|
|
except (json.JSONDecodeError, TypeError):
|
|
args = {"raw": str(args_str)}
|
|
if fn_name:
|
|
tc_id = tc.get("id", "") or ""
|
|
fc_id = ""
|
|
thought_sig = ""
|
|
# Decode: format is call_<fc_id>|<thought_signature>
|
|
if tc_id.startswith("call_"):
|
|
raw = tc_id[5:]
|
|
if "|" in raw:
|
|
fc_id, thought_sig = raw.split("|", 1)
|
|
else:
|
|
fc_id = raw
|
|
fc_obj = {"name": fn_name, "args": args}
|
|
if fc_id:
|
|
fc_obj["id"] = fc_id
|
|
part_obj = {"functionCall": fc_obj}
|
|
# thoughtSignature goes as a sibling key on the part, not inside functionCall.
|
|
# If missing, use the sentinel to skip validation (officially supported by Google).
|
|
if thought_sig:
|
|
part_obj["thoughtSignature"] = thought_sig
|
|
else:
|
|
part_obj["thoughtSignature"] = "skip_thought_signature_validator"
|
|
parts.append(part_obj)
|
|
if parts:
|
|
contents.append({"role": "model", "parts": parts})
|
|
continue
|
|
|
|
if role == "developer":
|
|
# Treat developer role like system.
|
|
text = _content_to_text(msg.get("content"))
|
|
if text:
|
|
system_texts.append(text)
|
|
continue
|
|
|
|
# Default: user role.
|
|
text = _content_to_text(msg.get("content"))
|
|
parts = []
|
|
if text:
|
|
parts.append({"text": text})
|
|
# Some clients send tool_calls on user messages; handle gracefully.
|
|
for tc in (msg.get("tool_calls") or []):
|
|
if not isinstance(tc, dict):
|
|
continue
|
|
fn = tc.get("function") or {}
|
|
fn_name = fn.get("name", "")
|
|
args_str = fn.get("arguments", "{}")
|
|
try:
|
|
args = json.loads(args_str) if isinstance(args_str, str) else args_str
|
|
except (json.JSONDecodeError, TypeError):
|
|
args = {"raw": str(args_str)}
|
|
if fn_name:
|
|
parts.append({"functionCall": {"name": fn_name, "args": args}})
|
|
if parts:
|
|
contents.append({"role": "user", "parts": parts})
|
|
|
|
# Build system instruction: use user's system messages if provided,
|
|
# otherwise fall back to the default.
|
|
if system_texts:
|
|
system_instruction = {"parts": [{"text": "\n\n".join(s for s in system_texts if s)}]}
|
|
else:
|
|
system_instruction = {"parts": [{"text": DEFAULT_SYSTEM_INSTRUCTION}]}
|
|
|
|
# Build the inner Gemini request.
|
|
inner = {
|
|
"contents": contents,
|
|
"systemInstruction": system_instruction,
|
|
}
|
|
|
|
# Tools.
|
|
tools = _openai_tools_to_gemini(body.get("tools"))
|
|
if tools:
|
|
inner["tools"] = tools
|
|
|
|
# Tool choice -> toolConfig.
|
|
tool_choice = body.get("tool_choice")
|
|
mode = _tool_choice_to_gemini(tool_choice)
|
|
if mode:
|
|
tc_config = {"functionCallingConfig": {"mode": mode}}
|
|
if isinstance(tool_choice, dict) and tool_choice.get("type") == "function":
|
|
fn_name = (tool_choice.get("function") or {}).get("name")
|
|
if fn_name:
|
|
tc_config["functionCallingConfig"]["allowedFunctionNames"] = [fn_name]
|
|
inner["toolConfig"] = tc_config
|
|
|
|
# Generation config.
|
|
gen_config = {}
|
|
if "temperature" in body and body["temperature"] is not None:
|
|
gen_config["temperature"] = float(body["temperature"])
|
|
if "top_p" in body and body["top_p"] is not None:
|
|
gen_config["topP"] = float(body["top_p"])
|
|
if "max_tokens" in body and body["max_tokens"] is not None:
|
|
gen_config["maxOutputTokens"] = int(body["max_tokens"])
|
|
elif "max_completion_tokens" in body and body["max_completion_tokens"] is not None:
|
|
gen_config["maxOutputTokens"] = int(body["max_completion_tokens"])
|
|
if "presence_penalty" in body and body["presence_penalty"] is not None:
|
|
gen_config["presencePenalty"] = float(body["presence_penalty"])
|
|
if "frequency_penalty" in body and body["frequency_penalty"] is not None:
|
|
gen_config["frequencyPenalty"] = float(body["frequency_penalty"])
|
|
if "stop" in body and body["stop"]:
|
|
stops = body["stop"]
|
|
if isinstance(stops, str):
|
|
stops = [stops]
|
|
gen_config["stopSequences"] = stops
|
|
if "seed" in body and body["seed"] is not None:
|
|
gen_config["seed"] = int(body["seed"])
|
|
# Enable thinking for the thinking variant of claude-opus.
|
|
if backend_model == "claude-opus-4-6-thinking":
|
|
thinking_cfg = gen_config.get("thinkingConfig", {})
|
|
thinking_cfg.setdefault("includeThoughts", False)
|
|
gen_config["thinkingConfig"] = thinking_cfg
|
|
# Inject thinking level for Gemini 3 Pro models.
|
|
if thinking_level:
|
|
thinking_cfg = gen_config.get("thinkingConfig", {})
|
|
thinking_cfg.setdefault("thinkingLevel", thinking_level)
|
|
thinking_cfg.setdefault("includeThoughts", False)
|
|
gen_config["thinkingConfig"] = thinking_cfg
|
|
if gen_config:
|
|
inner["generationConfig"] = gen_config
|
|
|
|
# Safety settings: relax standard categories so coding tasks aren't blocked.
|
|
# NOTE: HARM_CATEGORY_CIVIC_INTEGRITY is rejected by this API variant even
|
|
# though it appears in the error's valid-list — only the 4 below are accepted.
|
|
inner["safetySettings"] = [
|
|
{"category": cat, "threshold": "BLOCK_NONE"}
|
|
for cat in (
|
|
"HARM_CATEGORY_HARASSMENT",
|
|
"HARM_CATEGORY_HATE_SPEECH",
|
|
"HARM_CATEGORY_SEXUALLY_EXPLICIT",
|
|
"HARM_CATEGORY_DANGEROUS_CONTENT",
|
|
)
|
|
]
|
|
|
|
# Build the wrapped envelope.
|
|
project_id = get_project_id()
|
|
envelope = {
|
|
"project": project_id,
|
|
"model": backend_model,
|
|
"request": inner,
|
|
"requestType": "agent",
|
|
"userAgent": "antigravity",
|
|
"requestId": f"agent-{uuid.uuid4().hex}",
|
|
}
|
|
return envelope
|
|
|
|
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
# Gemini -> OpenAI response transformation
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
|
|
def _extract_parts(parts: list) -> tuple[str, list, str]:
|
|
"""Extract (text, tool_calls, finish_reason_hint) from a Gemini parts list.
|
|
|
|
tool_calls is a list of OpenAI-format tool_call dicts. The tool_call.id
|
|
encodes the Gemini functionCall.id and thoughtSignature so they can be
|
|
round-tripped back when the assistant message is sent in a later request.
|
|
Format: call_<fc_id>|<thought_signature>
|
|
"""
|
|
text_chunks = []
|
|
tool_calls = []
|
|
finish_hint = None
|
|
|
|
for part in parts or []:
|
|
if not isinstance(part, dict):
|
|
continue
|
|
if "text" in part and part["text"]:
|
|
text_chunks.append(part["text"])
|
|
if "functionCall" in part:
|
|
fc = part["functionCall"] or {}
|
|
name = fc.get("name", "")
|
|
args = fc.get("args", {})
|
|
fc_id = fc.get("id", "")
|
|
# thoughtSignature is a sibling of functionCall in the part, not inside it.
|
|
thought_sig = part.get("thoughtSignature") or part.get("thought_signature") or ""
|
|
if name:
|
|
# Encode the Gemini fc_id and thoughtSignature into the OpenAI tool_call.id
|
|
# so they survive the round-trip through Hermes's message history.
|
|
# Format: call_<fc_id>|<thought_signature>
|
|
encoded_id = f"call_{fc_id}|{thought_sig}" if (fc_id or thought_sig) else ""
|
|
tool_calls.append({
|
|
"id": encoded_id, # may be empty; caller assigns if so
|
|
"type": "function",
|
|
"function": {
|
|
"name": name,
|
|
"arguments": json.dumps(args) if isinstance(args, (dict, list)) else str(args),
|
|
},
|
|
})
|
|
if "executableCode" in part:
|
|
ec = part["executableCode"] or {}
|
|
code = ec.get("code", "")
|
|
if code:
|
|
text_chunks.append(f"```python\n{code}\n```")
|
|
if "codeExecutionResult" in part:
|
|
cer = part["codeExecutionResult"] or {}
|
|
out = cer.get("output", "")
|
|
if out:
|
|
text_chunks.append(f"```\n{out}\n```")
|
|
if "thought" in part and part["thought"]:
|
|
# Thinking summaries — include as plain text (Gemini sometimes
|
|
# surfaces these). Keep it; the client can ignore.
|
|
pass
|
|
|
|
if tool_calls:
|
|
finish_hint = "tool_calls"
|
|
return "".join(text_chunks), tool_calls, finish_hint
|
|
|
|
|
|
def _finish_reason_from_gemini(candidate: dict, tool_calls: list) -> str:
|
|
"""Map Gemini finishReason / stopReason to OpenAI finish_reason."""
|
|
if tool_calls:
|
|
# If there are tool calls, OpenAI expects "tool_calls".
|
|
fr = (candidate.get("finishReason") or "").upper()
|
|
# Gemini may say STOP even when emitting function calls.
|
|
return "tool_calls"
|
|
fr = (candidate.get("finishReason") or candidate.get("stopReason") or "").upper()
|
|
mapping = {
|
|
"STOP": "stop",
|
|
"MAX_TOKENS": "length",
|
|
"SAFETY": "content_filter",
|
|
"RECITATION": "content_filter",
|
|
"BLOCKLIST": "content_filter",
|
|
"PROHIBITED_CONTENT": "content_filter",
|
|
"SPII": "content_filter",
|
|
"MALFORMED_FUNCTION_CALL": "stop",
|
|
"IMAGE_SAFETY": "content_filter",
|
|
"LANGUAGE": "content_filter",
|
|
"OTHER": "stop",
|
|
}
|
|
return mapping.get(fr, "stop")
|
|
|
|
|
|
def _extract_usage(response: dict) -> dict:
|
|
"""Extract OpenAI-format usage from a Gemini response."""
|
|
usage = response.get("usageMetadata") or {}
|
|
if not usage:
|
|
cands = response.get("candidates") or []
|
|
if cands:
|
|
usage = cands[0].get("usageMetadata") or {}
|
|
pt = usage.get("promptTokenCount", 0) or 0
|
|
ct = usage.get("candidatesTokenCount", 0) or usage.get("completionTokenCount", 0) or 0
|
|
tt = usage.get("totalTokenCount", (pt + ct)) or (pt + ct)
|
|
return {
|
|
"prompt_tokens": pt,
|
|
"completion_tokens": ct,
|
|
"total_tokens": tt,
|
|
}
|
|
|
|
|
|
def gemini_response_to_openai(response: dict, model: str) -> dict:
|
|
"""Transform an unwrapped Gemini generateContent response to OpenAI
|
|
chat.completion format.
|
|
"""
|
|
candidates = response.get("candidates") or []
|
|
text = ""
|
|
tool_calls = []
|
|
finish_reason = "stop"
|
|
|
|
if candidates:
|
|
cand = candidates[0]
|
|
content = cand.get("content") or {}
|
|
parts = content.get("parts") or []
|
|
text, tool_calls, _ = _extract_parts(parts)
|
|
finish_reason = _finish_reason_from_gemini(cand, tool_calls)
|
|
# Assign tool call ids.
|
|
for i, tc in enumerate(tool_calls):
|
|
if not tc["id"]:
|
|
tc["id"] = f"call_{uuid.uuid4().hex[:24]}"
|
|
else:
|
|
# No candidates — possibly a promptFeedback block.
|
|
pf = response.get("promptFeedback") or {}
|
|
br = (pf.get("blockReason") or "").upper()
|
|
if br:
|
|
finish_reason = "content_filter"
|
|
text = f"[Content blocked: {br}]"
|
|
|
|
message = {"role": "assistant", "content": text if text else None}
|
|
if tool_calls:
|
|
message["tool_calls"] = tool_calls
|
|
|
|
return {
|
|
"id": f"chatcmpl-{uuid.uuid4().hex[:24]}",
|
|
"object": "chat.completion",
|
|
"created": int(time.time()),
|
|
"model": model,
|
|
"choices": [{
|
|
"index": 0,
|
|
"message": message,
|
|
"finish_reason": finish_reason,
|
|
}],
|
|
"usage": _extract_usage(response),
|
|
}
|
|
|
|
|
|
def _extract_delta_from_parts(parts: list, prev_tool_names: set) -> tuple[str, list, list]:
|
|
"""Extract incremental (text, new_tool_calls, all_tool_names) from parts.
|
|
|
|
For streaming, each chunk may contain partial text and/or function calls.
|
|
We return text deltas and tool_call entries (with indices).
|
|
"""
|
|
text_chunks = []
|
|
tool_call_deltas = []
|
|
current_names = set(prev_tool_names)
|
|
|
|
for part in parts or []:
|
|
if not isinstance(part, dict):
|
|
continue
|
|
if "text" in part and part["text"]:
|
|
text_chunks.append(part["text"])
|
|
if "functionCall" in part:
|
|
fc = part["functionCall"] or {}
|
|
name = fc.get("name", "")
|
|
args = fc.get("args", {})
|
|
fc_id = fc.get("id", "")
|
|
thought_sig = part.get("thoughtSignature") or part.get("thought_signature") or ""
|
|
if name:
|
|
current_names.add(name)
|
|
encoded_id = f"call_{fc_id}|{thought_sig}" if (fc_id or thought_sig) else ""
|
|
tool_call_deltas.append({
|
|
"index": len(current_names) - 1,
|
|
"id": encoded_id or f"call_{uuid.uuid4().hex[:24]}",
|
|
"type": "function",
|
|
"function": {
|
|
"name": name,
|
|
"arguments": json.dumps(args) if isinstance(args, (dict, list)) else str(args),
|
|
},
|
|
})
|
|
|
|
return "".join(text_chunks), tool_call_deltas, current_names
|
|
|
|
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
# Upstream API calls
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
|
|
def _make_upstream_request(url: str, envelope: dict, streaming: bool = False):
|
|
"""Build and return a urllib Request object for the upstream call."""
|
|
body = json.dumps(envelope).encode()
|
|
req = Request(url, data=body, method="POST")
|
|
req.add_header("Authorization", f"Bearer {get_access_token()}")
|
|
req.add_header("Content-Type", "application/json")
|
|
req.add_header("User-Agent", ANTIGRAVITY_USER_AGENT)
|
|
req.add_header("X-Goog-Api-Client", "google-cloud-sdk vscode_cloudshelleditor/0.1")
|
|
req.add_header("Client-Metadata", CLIENT_METADATA)
|
|
if streaming:
|
|
req.add_header("Accept", "text/event-stream")
|
|
return req
|
|
|
|
|
|
def call_generate_content(envelope: dict) -> dict:
|
|
"""Non-streaming call. Returns the unwrapped inner Gemini response dict."""
|
|
req = _make_upstream_request(GENERATE_CONTENT_URL, envelope, streaming=False)
|
|
try:
|
|
with urlopen(req, timeout=UPSTREAM_TIMEOUT) as resp:
|
|
raw = resp.read().decode()
|
|
except HTTPError as e:
|
|
detail = ""
|
|
try:
|
|
detail = e.read().decode()
|
|
except Exception:
|
|
pass
|
|
raise UpstreamError(f"Upstream HTTP {e.code}: {detail}", status=e.code, body=detail)
|
|
except URLError as e:
|
|
raise UpstreamError(f"Upstream network error: {e}", status=502)
|
|
|
|
try:
|
|
payload = json.loads(raw)
|
|
except json.JSONDecodeError:
|
|
raise UpstreamError(f"Upstream returned non-JSON: {raw[:500]}", status=502)
|
|
|
|
# Unwrap the envelope: {"response": {<gemini>}} or sometimes bare.
|
|
if isinstance(payload, dict) and "response" in payload:
|
|
return payload["response"]
|
|
return payload
|
|
|
|
|
|
def stream_generate_content(envelope: dict):
|
|
"""Streaming call. Yields parsed JSON event objects from the SSE stream.
|
|
|
|
The Antigravity stream endpoint returns either:
|
|
- SSE format: lines "data: {json}\n\n"
|
|
- Or a bare JSON array of incremental response objects (some endpoints).
|
|
We handle both.
|
|
"""
|
|
req = _make_upstream_request(STREAM_GENERATE_CONTENT_URL, envelope, streaming=True)
|
|
try:
|
|
resp = urlopen(req, timeout=UPSTREAM_TIMEOUT)
|
|
except HTTPError as e:
|
|
detail = ""
|
|
try:
|
|
detail = e.read().decode()
|
|
except Exception:
|
|
pass
|
|
raise UpstreamError(f"Upstream stream HTTP {e.code}: {detail}", status=e.code, body=detail)
|
|
except URLError as e:
|
|
raise UpstreamError(f"Upstream stream network error: {e}", status=502)
|
|
|
|
try:
|
|
buffer = ""
|
|
# Read in chunks; urllib's response supports iteration by bytes.
|
|
while True:
|
|
chunk = resp.read(8192)
|
|
if not chunk:
|
|
break
|
|
if isinstance(chunk, bytes):
|
|
text_chunk = chunk.decode("utf-8", errors="replace")
|
|
else:
|
|
text_chunk = chunk
|
|
buffer += text_chunk
|
|
|
|
# Try to parse complete SSE events or JSON array elements.
|
|
while True:
|
|
parsed_event = _try_parse_one_event(buffer)
|
|
if parsed_event is None:
|
|
break
|
|
event_obj, consumed = parsed_event
|
|
buffer = buffer[consumed:]
|
|
if event_obj is not None:
|
|
yield event_obj
|
|
# Flush any remaining buffer.
|
|
if buffer.strip():
|
|
for obj in _parse_remaining(buffer):
|
|
if obj is not None:
|
|
yield obj
|
|
finally:
|
|
resp.close()
|
|
|
|
|
|
def _try_parse_one_event(buffer: str):
|
|
"""Attempt to parse one SSE event or JSON array element from the buffer.
|
|
|
|
Returns (event_obj_or_None, bytes_consumed) or None if incomplete.
|
|
The event_obj may be None when it's a comment/heartbeat line.
|
|
"""
|
|
# SSE "data:" lines.
|
|
if "data:" in buffer[:10] or buffer.startswith("\n") or buffer.startswith(":"):
|
|
# Find end of this SSE block (double newline).
|
|
for sep in ("\n\n", "\r\n\r\n"):
|
|
idx = buffer.find(sep)
|
|
if idx != -1:
|
|
block = buffer[:idx]
|
|
consumed = idx + len(sep)
|
|
obj = _parse_sse_block(block)
|
|
return (obj, consumed)
|
|
# Maybe single newline-terminated line (some servers use \n not \n\n).
|
|
nl = buffer.find("\n")
|
|
if nl != -1 and buffer.strip().startswith("data:"):
|
|
# Only consume if we have a complete data line.
|
|
line = buffer[:nl]
|
|
consumed = nl + 1
|
|
obj = _parse_sse_block(line)
|
|
return (obj, consumed)
|
|
return None # incomplete
|
|
|
|
# Bare JSON array stream: [ {...}, {...} ]
|
|
if buffer.lstrip().startswith("["):
|
|
obj, consumed = _try_parse_json_element_in_array(buffer)
|
|
if obj is not None or consumed > 0:
|
|
return (obj, consumed)
|
|
return None
|
|
|
|
# Bare JSON object stream: {...}{...}
|
|
if buffer.lstrip().startswith("{"):
|
|
obj, consumed = _try_parse_one_json_object(buffer)
|
|
if obj is not None:
|
|
return (obj, consumed)
|
|
return None
|
|
|
|
# Unknown prefix; skip a line to resync.
|
|
nl = buffer.find("\n")
|
|
if nl != -1:
|
|
return (None, nl + 1)
|
|
return None
|
|
|
|
|
|
def _parse_sse_block(block: str):
|
|
"""Parse an SSE block (one or more lines) into a JSON object or None."""
|
|
data_lines = []
|
|
for line in block.split("\n"):
|
|
line = line.rstrip("\r")
|
|
if line.startswith("data:"):
|
|
data_lines.append(line[5:].lstrip())
|
|
elif line.startswith(":"):
|
|
continue # comment / heartbeat
|
|
elif line.startswith("event:") or line.startswith("id:") or line.startswith("retry:"):
|
|
continue
|
|
elif line.strip() == "":
|
|
continue
|
|
else:
|
|
# Unexpected line; ignore.
|
|
pass
|
|
if not data_lines:
|
|
return None
|
|
data_str = "\n".join(data_lines).strip()
|
|
if data_str == "[DONE]":
|
|
return {"__done__": True}
|
|
try:
|
|
return json.loads(data_str)
|
|
except json.JSONDecodeError:
|
|
return None
|
|
|
|
|
|
def _try_parse_json_element_in_array(buffer: str):
|
|
"""For bare JSON array streams like [{...},{...}]. Parse one element."""
|
|
s = buffer.lstrip()
|
|
if not s.startswith("["):
|
|
# Maybe we already consumed the opening bracket.
|
|
pass
|
|
# We need to find a complete top-level object. Use brace counting.
|
|
obj, consumed = _try_parse_one_json_object(buffer)
|
|
return (obj, consumed)
|
|
|
|
|
|
def _try_parse_one_json_object(buffer: str):
|
|
"""Parse one complete top-level JSON object from buffer using brace counting.
|
|
|
|
Returns (obj, consumed_chars) or (None, 0) if incomplete.
|
|
"""
|
|
# Find first '{'.
|
|
start = buffer.find("{")
|
|
if start == -1:
|
|
return (None, 0)
|
|
depth = 0
|
|
in_string = False
|
|
escape = False
|
|
for i in range(start, len(buffer)):
|
|
ch = buffer[i]
|
|
if escape:
|
|
escape = False
|
|
continue
|
|
if ch == "\\":
|
|
escape = True
|
|
continue
|
|
if ch == '"':
|
|
in_string = not in_string
|
|
continue
|
|
if in_string:
|
|
continue
|
|
if ch == "{":
|
|
depth += 1
|
|
elif ch == "}":
|
|
depth -= 1
|
|
if depth == 0:
|
|
candidate = buffer[start:i + 1]
|
|
try:
|
|
obj = json.loads(candidate)
|
|
return (obj, i + 1)
|
|
except json.JSONDecodeError:
|
|
# Keep scanning.
|
|
continue
|
|
return (None, 0)
|
|
|
|
|
|
def _parse_remaining(buffer: str):
|
|
"""Parse any remaining JSON objects in the buffer at stream end."""
|
|
s = buffer.strip()
|
|
if not s:
|
|
return
|
|
# SSE.
|
|
for line in s.split("\n"):
|
|
line = line.strip()
|
|
if line.startswith("data:"):
|
|
data = line[5:].strip()
|
|
if data == "[DONE]":
|
|
yield {"__done__": True}
|
|
continue
|
|
try:
|
|
yield json.loads(data)
|
|
except json.JSONDecodeError:
|
|
pass
|
|
# Bare JSON.
|
|
if s.startswith("[") or s.startswith("{"):
|
|
try:
|
|
obj = json.loads(s)
|
|
if isinstance(obj, list):
|
|
for item in obj:
|
|
yield item
|
|
else:
|
|
yield obj
|
|
except json.JSONDecodeError:
|
|
pass
|
|
|
|
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
# Error helper
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
|
|
class UpstreamError(Exception):
|
|
def __init__(self, message, status=502, body=""):
|
|
super().__init__(message)
|
|
self.status = status
|
|
self.body = body
|
|
|
|
|
|
def _openai_error(message: str, err_type: str = "api_error", code=None, status: int = 500):
|
|
return {
|
|
"error": {
|
|
"message": message,
|
|
"type": err_type,
|
|
"param": None,
|
|
"code": code,
|
|
}
|
|
}
|
|
|
|
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
# HTTP handler
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
|
|
class ProxyHandler(BaseHTTPRequestHandler):
|
|
protocol_version = "HTTP/1.1"
|
|
server_version = "AntigravityProxy/2.0"
|
|
|
|
def log_message(self, fmt, *args):
|
|
# We do our own logging.
|
|
pass
|
|
|
|
def _send_json(self, code, data):
|
|
body = json.dumps(data).encode()
|
|
self.send_response(code)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.send_header("Content-Length", str(len(body)))
|
|
self.end_headers()
|
|
try:
|
|
self.wfile.write(body)
|
|
except (BrokenPipeError, ConnectionResetError):
|
|
pass
|
|
|
|
def _send_error_json(self, code, message, err_type="api_error", http_code=None):
|
|
self._send_json(http_code or code, _openai_error(message, err_type, code=code))
|
|
|
|
def _read_body(self):
|
|
length = int(self.headers.get("Content-Length", 0) or 0)
|
|
if length <= 0:
|
|
return {}
|
|
raw = self.rfile.read(length)
|
|
if not raw:
|
|
return {}
|
|
try:
|
|
return json.loads(raw)
|
|
except json.JSONDecodeError:
|
|
return None # signal parse failure
|
|
|
|
# ── GET routes ────────────────────────────────────────────────────────────
|
|
|
|
def do_GET(self):
|
|
t0 = time.time()
|
|
path = urlparse(self.path).path
|
|
try:
|
|
if path == "/v1/models" or path == "/models":
|
|
self._handle_models()
|
|
self._log_req("GET", path, 200, time.time() - t0)
|
|
elif path in ("/health", "/"):
|
|
self._send_json(200, {"status": "ok", "service": "antigravity-proxy"})
|
|
self._log_req("GET", path, 200, time.time() - t0)
|
|
else:
|
|
self._send_error_json("not_found", "Not found", err_type="invalid_request_error", http_code=404)
|
|
self._log_req("GET", path, 404, time.time() - t0)
|
|
except Exception as e:
|
|
_log(f"GET {path} error: {e}\n{traceback.format_exc()}")
|
|
try:
|
|
self._send_error_json("internal_error", str(e), http_code=500)
|
|
except Exception:
|
|
pass
|
|
self._log_req("GET", path, 500, time.time() - t0)
|
|
|
|
def _handle_models(self):
|
|
now = int(time.time())
|
|
models = []
|
|
for name in MODEL_MAP:
|
|
models.append({
|
|
"id": name,
|
|
"object": "model",
|
|
"created": now,
|
|
"owned_by": "google",
|
|
})
|
|
self._send_json(200, {"object": "list", "data": models})
|
|
|
|
# ── POST routes ───────────────────────────────────────────────────────────
|
|
|
|
def do_POST(self):
|
|
t0 = time.time()
|
|
path = urlparse(self.path).path
|
|
try:
|
|
if path in ("/v1/chat/completions", "/chat/completions"):
|
|
self._handle_chat()
|
|
self._log_req("POST", path, 200, time.time() - t0)
|
|
elif _GEMINI_PASSTHROUGH_RE.match(path):
|
|
self._handle_gemini_passthrough(path)
|
|
self._log_req("POST", path, 200, time.time() - t0)
|
|
else:
|
|
self._send_error_json("not_found", "Not found", err_type="invalid_request_error", http_code=404)
|
|
self._log_req("POST", path, 404, time.time() - t0)
|
|
except Exception as e:
|
|
_log(f"POST {path} error: {e}\n{traceback.format_exc()}")
|
|
try:
|
|
self._send_error_json("internal_error", str(e), http_code=500)
|
|
except Exception:
|
|
pass
|
|
self._log_req("POST", path, 500, time.time() - t0)
|
|
|
|
def _log_req(self, method, path, status, duration):
|
|
_log(f"{method} {path} {status} {duration:.3f}s")
|
|
|
|
# ── Chat completions ──────────────────────────────────────────────────────
|
|
|
|
def _handle_gemini_passthrough(self, path: str):
|
|
"""Handle a native Gemini generateContent request, wrapping it in the
|
|
Antigravity envelope and returning the raw Gemini response (unwrapped).
|
|
This is used by invoice_reader.py to send inline_data + responseSchema
|
|
requests through the OAuth-funded proxy instead of a paid API key."""
|
|
body = self._read_body()
|
|
if body is None or not isinstance(body, dict):
|
|
self._send_error_json("invalid_request", "Invalid JSON body", err_type="invalid_request_error", http_code=400)
|
|
return
|
|
|
|
model = _GEMINI_PASSTHROUGH_RE.match(path).group(1)
|
|
|
|
try:
|
|
envelope = _build_passthrough_envelope(model, body)
|
|
except UpstreamError as e:
|
|
self._send_error_json("upstream_error", str(e), err_type="api_error", http_code=e.status)
|
|
return
|
|
except RuntimeError as e:
|
|
self._send_error_json("auth_error", str(e), err_type="authentication_error", http_code=401)
|
|
return
|
|
except Exception as e:
|
|
self._send_error_json("request_error", f"Failed to build request: {e}", err_type="invalid_request_error", http_code=400)
|
|
return
|
|
|
|
try:
|
|
gemini_resp = call_generate_content(envelope)
|
|
except UpstreamError as e:
|
|
self._send_error_json("upstream_error", str(e), err_type="api_error", http_code=e.status)
|
|
return
|
|
except Exception as e:
|
|
self._send_error_json("upstream_error", f"Upstream call failed: {e}", err_type="api_error", http_code=502)
|
|
return
|
|
|
|
self._send_json(200, gemini_resp)
|
|
|
|
def _handle_chat(self):
|
|
body = self._read_body()
|
|
if body is None:
|
|
self._send_error_json("invalid_request", "Invalid JSON body", err_type="invalid_request_error", http_code=400)
|
|
return
|
|
if not isinstance(body, dict):
|
|
self._send_error_json("invalid_request", "Request body must be a JSON object", err_type="invalid_request_error", http_code=400)
|
|
return
|
|
|
|
openai_model = body.get("model") or DEFAULT_MODEL
|
|
stream = bool(body.get("stream", False))
|
|
|
|
# Map model. MODEL_MAP values are (backend_name, thinking_level) tuples.
|
|
model_entry = MODEL_MAP.get(openai_model)
|
|
if model_entry:
|
|
backend_model, thinking_level = model_entry
|
|
else:
|
|
# Allow pass-through of unknown models.
|
|
backend_model = openai_model
|
|
thinking_level = None
|
|
|
|
# Build the Gemini request envelope.
|
|
try:
|
|
envelope = _build_gemini_request(body, backend_model, thinking_level)
|
|
except UpstreamError as e:
|
|
self._send_error_json("upstream_error", str(e), err_type="api_error", http_code=e.status)
|
|
return
|
|
except RuntimeError as e:
|
|
self._send_error_json("auth_error", str(e), err_type="authentication_error", http_code=401)
|
|
return
|
|
except Exception as e:
|
|
self._send_error_json("request_error", f"Failed to build request: {e}", err_type="invalid_request_error", http_code=400)
|
|
return
|
|
|
|
# Validate we have at least one content entry.
|
|
inner = envelope.get("request", {})
|
|
if not inner.get("contents"):
|
|
self._send_error_json("invalid_request", "No messages provided", err_type="invalid_request_error", http_code=400)
|
|
return
|
|
|
|
if stream:
|
|
self._handle_stream(envelope, openai_model)
|
|
else:
|
|
self._handle_nonstream(envelope, openai_model)
|
|
|
|
def _handle_nonstream(self, envelope, openai_model):
|
|
try:
|
|
gemini_resp = call_generate_content(envelope)
|
|
except UpstreamError as e:
|
|
self._send_error_json("upstream_error", str(e), err_type="api_error", http_code=e.status)
|
|
return
|
|
except Exception as e:
|
|
self._send_error_json("upstream_error", f"Upstream call failed: {e}", err_type="api_error", http_code=502)
|
|
return
|
|
|
|
try:
|
|
openai_resp = gemini_response_to_openai(gemini_resp, openai_model)
|
|
except Exception as e:
|
|
_log(f"Response transform error: {e}\n{traceback.format_exc()}")
|
|
# Fallback: return raw text if we can find any.
|
|
try:
|
|
text = (gemini_resp.get("candidates", [{}])[0].get("content", {}).get("parts", [{}])[0].get("text", ""))
|
|
except Exception:
|
|
text = ""
|
|
openai_resp = {
|
|
"id": f"chatcmpl-{uuid.uuid4().hex[:24]}",
|
|
"object": "chat.completion",
|
|
"created": int(time.time()),
|
|
"model": openai_model,
|
|
"choices": [{
|
|
"index": 0,
|
|
"message": {"role": "assistant", "content": text or "[error: could not parse response]"},
|
|
"finish_reason": "stop",
|
|
}],
|
|
"usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
|
|
}
|
|
self._send_json(200, openai_resp)
|
|
|
|
def _handle_stream(self, envelope, openai_model):
|
|
completion_id = f"chatcmpl-{uuid.uuid4().hex[:24]}"
|
|
created = int(time.time())
|
|
|
|
# Send SSE headers.
|
|
try:
|
|
self.send_response(200)
|
|
self.send_header("Content-Type", "text/event-stream")
|
|
self.send_header("Cache-Control", "no-cache")
|
|
self.send_header("Connection", "keep-alive")
|
|
self.send_header("X-Accel-Buffering", "no")
|
|
self.end_headers()
|
|
except Exception:
|
|
return
|
|
|
|
def send_sse(obj):
|
|
data = f"data: {json.dumps(obj)}\n\n"
|
|
try:
|
|
self.wfile.write(data.encode())
|
|
self.wfile.flush()
|
|
except (BrokenPipeError, ConnectionResetError):
|
|
return False
|
|
return True
|
|
|
|
# Initial role chunk.
|
|
first_chunk = {
|
|
"id": completion_id,
|
|
"object": "chat.completion.chunk",
|
|
"created": created,
|
|
"model": openai_model,
|
|
"choices": [{
|
|
"index": 0,
|
|
"delta": {"role": "assistant", "content": ""},
|
|
"finish_reason": None,
|
|
}],
|
|
}
|
|
if not send_sse(first_chunk):
|
|
return
|
|
|
|
try:
|
|
gen = stream_generate_content(envelope)
|
|
except UpstreamError as e:
|
|
# Streaming failed — fall back to non-streaming and return as a single chunk.
|
|
# This handles the case where the streamGenerateContent endpoint rejects
|
|
# tool-result round-trips with 400 "invalid argument".
|
|
if e.status == 400:
|
|
try:
|
|
resp = call_generate_content(envelope)
|
|
openai_resp = gemini_response_to_openai(resp, openai_model)
|
|
msg = openai_resp["choices"][0]["message"]
|
|
delta = {}
|
|
if msg.get("content"):
|
|
delta["content"] = msg["content"]
|
|
if msg.get("tool_calls"):
|
|
delta["tool_calls"] = msg["tool_calls"]
|
|
chunk = {
|
|
"id": completion_id,
|
|
"object": "chat.completion.chunk",
|
|
"created": created,
|
|
"model": openai_model,
|
|
"choices": [{
|
|
"index": 0,
|
|
"delta": delta,
|
|
"finish_reason": openai_resp["choices"][0].get("finish_reason", "stop"),
|
|
}],
|
|
}
|
|
if openai_resp.get("usage"):
|
|
chunk["usage"] = openai_resp["usage"]
|
|
send_sse(chunk)
|
|
try:
|
|
self.wfile.write(b"data: [DONE]\n\n")
|
|
self.wfile.flush()
|
|
except Exception:
|
|
pass
|
|
return
|
|
except Exception:
|
|
pass # fall through to error chunk
|
|
err_chunk = {
|
|
"id": completion_id,
|
|
"object": "chat.completion.chunk",
|
|
"created": created,
|
|
"model": openai_model,
|
|
"choices": [{
|
|
"index": 0,
|
|
"delta": {"content": f"[Error: {e}]"},
|
|
"finish_reason": "stop",
|
|
}],
|
|
}
|
|
send_sse(err_chunk)
|
|
send_sse({"__done__": True}) # sentinel handled below
|
|
try:
|
|
self.wfile.write(b"data: [DONE]\n\n")
|
|
self.wfile.flush()
|
|
except Exception:
|
|
pass
|
|
return
|
|
except Exception as e:
|
|
err_chunk = {
|
|
"id": completion_id,
|
|
"object": "chat.completion.chunk",
|
|
"created": created,
|
|
"model": openai_model,
|
|
"choices": [{
|
|
"index": 0,
|
|
"delta": {"content": f"[Error: {e}]"},
|
|
"finish_reason": "stop",
|
|
}],
|
|
}
|
|
send_sse(err_chunk)
|
|
try:
|
|
self.wfile.write(b"data: [DONE]\n\n")
|
|
self.wfile.flush()
|
|
except Exception:
|
|
pass
|
|
return
|
|
|
|
finish_reason = "stop"
|
|
seen_tool_names = set()
|
|
total_usage = None
|
|
events_processed = 0
|
|
|
|
try:
|
|
for event in gen:
|
|
if not isinstance(event, dict):
|
|
continue
|
|
if event.get("__done__"):
|
|
break
|
|
|
|
events_processed += 1
|
|
|
|
# Unwrap envelope if present.
|
|
resp = event.get("response", event)
|
|
|
|
# Check for usage metadata (often on the last chunk).
|
|
usage_meta = resp.get("usageMetadata")
|
|
if usage_meta:
|
|
total_usage = usage_meta
|
|
|
|
candidates = resp.get("candidates") or []
|
|
if not candidates:
|
|
continue
|
|
|
|
cand = candidates[0]
|
|
content = cand.get("content") or {}
|
|
parts = content.get("parts") or []
|
|
|
|
text_delta, tool_call_deltas, seen_tool_names = _extract_delta_from_parts(
|
|
parts, seen_tool_names
|
|
)
|
|
|
|
# Determine finish reason.
|
|
fr = cand.get("finishReason")
|
|
if fr:
|
|
if tool_call_deltas:
|
|
finish_reason = "tool_calls"
|
|
else:
|
|
fr_up = str(fr).upper()
|
|
mapping = {
|
|
"STOP": "stop",
|
|
"MAX_TOKENS": "length",
|
|
"SAFETY": "content_filter",
|
|
"RECITATION": "content_filter",
|
|
"BLOCKLIST": "content_filter",
|
|
"PROHIBITED_CONTENT": "content_filter",
|
|
"SPII": "content_filter",
|
|
"IMAGE_SAFETY": "content_filter",
|
|
"LANGUAGE": "content_filter",
|
|
"OTHER": "stop",
|
|
}
|
|
finish_reason = mapping.get(fr_up, "stop")
|
|
|
|
# Build delta.
|
|
delta = {}
|
|
if text_delta:
|
|
delta["content"] = text_delta
|
|
if tool_call_deltas:
|
|
delta["tool_calls"] = tool_call_deltas
|
|
|
|
if not delta and not fr:
|
|
# Empty chunk with nothing — skip.
|
|
continue
|
|
|
|
chunk = {
|
|
"id": completion_id,
|
|
"object": "chat.completion.chunk",
|
|
"created": created,
|
|
"model": openai_model,
|
|
"choices": [{
|
|
"index": 0,
|
|
"delta": delta,
|
|
"finish_reason": None,
|
|
}],
|
|
}
|
|
if not send_sse(chunk):
|
|
return
|
|
|
|
# Final chunk with finish_reason.
|
|
final_chunk = {
|
|
"id": completion_id,
|
|
"object": "chat.completion.chunk",
|
|
"created": created,
|
|
"model": openai_model,
|
|
"choices": [{
|
|
"index": 0,
|
|
"delta": {},
|
|
"finish_reason": finish_reason,
|
|
}],
|
|
}
|
|
# Include usage in the final chunk if the client requested
|
|
# stream_options.include_usage (OpenAI convention).
|
|
if total_usage:
|
|
pt = total_usage.get("promptTokenCount", 0) or 0
|
|
ct = total_usage.get("candidatesTokenCount", 0) or total_usage.get("completionTokenCount", 0) or 0
|
|
tt = total_usage.get("totalTokenCount", (pt + ct)) or (pt + ct)
|
|
final_chunk["usage"] = {
|
|
"prompt_tokens": pt,
|
|
"completion_tokens": ct,
|
|
"total_tokens": tt,
|
|
}
|
|
send_sse(final_chunk)
|
|
|
|
# Termination sentinel.
|
|
try:
|
|
self.wfile.write(b"data: [DONE]\n\n")
|
|
self.wfile.flush()
|
|
except Exception:
|
|
pass
|
|
|
|
except Exception as e:
|
|
_log(f"Streaming error: {e}")
|
|
# If the stream endpoint returned 400, try non-streaming as fallback.
|
|
# The streamGenerateContent endpoint sometimes rejects tool-result round-trips.
|
|
if "400" in str(e):
|
|
try:
|
|
resp = call_generate_content(envelope)
|
|
openai_resp = gemini_response_to_openai(resp, openai_model)
|
|
msg = openai_resp["choices"][0]["message"]
|
|
delta = {}
|
|
if msg.get("content"):
|
|
delta["content"] = msg["content"]
|
|
if msg.get("tool_calls"):
|
|
delta["tool_calls"] = msg["tool_calls"]
|
|
chunk = {
|
|
"id": completion_id,
|
|
"object": "chat.completion.chunk",
|
|
"created": created,
|
|
"model": openai_model,
|
|
"choices": [{
|
|
"index": 0,
|
|
"delta": delta,
|
|
"finish_reason": openai_resp["choices"][0].get("finish_reason", "stop"),
|
|
}],
|
|
}
|
|
if openai_resp.get("usage"):
|
|
chunk["usage"] = openai_resp["usage"]
|
|
send_sse(chunk)
|
|
try:
|
|
self.wfile.write(b"data: [DONE]\n\n")
|
|
self.wfile.flush()
|
|
except Exception:
|
|
pass
|
|
return
|
|
except Exception:
|
|
pass # fall through to error chunk
|
|
# Try to send an error chunk then close.
|
|
try:
|
|
err_chunk = {
|
|
"id": completion_id,
|
|
"object": "chat.completion.chunk",
|
|
"created": created,
|
|
"model": openai_model,
|
|
"choices": [{
|
|
"index": 0,
|
|
"delta": {"content": f"\n[stream error: {e}]"},
|
|
"finish_reason": "stop",
|
|
}],
|
|
}
|
|
send_sse(err_chunk)
|
|
self.wfile.write(b"data: [DONE]\n\n")
|
|
self.wfile.flush()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
# Server entry point
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
|
|
def _preflight_check():
|
|
"""Validate the token file exists and is readable before starting."""
|
|
if not os.path.isfile(TOKEN_FILE):
|
|
_log(f"ERROR: Token file not found at {TOKEN_FILE}")
|
|
_log("Run the Antigravity CLI (`agy`) to authenticate first.")
|
|
return False
|
|
data = _read_token_from_disk()
|
|
if not data:
|
|
_log("ERROR: Token file exists but could not be parsed.")
|
|
return False
|
|
token_obj = data.get("token") or {}
|
|
if not token_obj.get("refresh_token"):
|
|
_log("ERROR: No refresh_token in token file. Re-run `agy` to authenticate.")
|
|
return False
|
|
_log(f"Token file OK (auth_method={data.get('auth_method', 'unknown')})")
|
|
return True
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Antigravity OpenAI-compatible proxy (Cloud Code Assist API)"
|
|
)
|
|
parser.add_argument("--port", type=int, default=8877, help="Port to listen on (default 8877)")
|
|
parser.add_argument("--host", default="127.0.0.1", help="Host to bind to (default 127.0.0.1)")
|
|
args = parser.parse_args()
|
|
|
|
if not _preflight_check():
|
|
sys.exit(1)
|
|
|
|
# Pre-discover the project ID so the first request is fast.
|
|
try:
|
|
get_project_id()
|
|
except Exception as e:
|
|
_log(f"WARNING: Could not pre-discover project ID: {e}")
|
|
_log("It will be discovered on the first request.")
|
|
|
|
server = ThreadingHTTPServer((args.host, args.port), ProxyHandler)
|
|
server.daemon_threads = True
|
|
_log(f"Antigravity proxy listening on http://{args.host}:{args.port}/v1")
|
|
_log(f" Models: {', '.join(MODEL_MAP.keys())}")
|
|
_log(f" Token: {TOKEN_FILE}")
|
|
_log(f" Press Ctrl+C to stop")
|
|
|
|
try:
|
|
server.serve_forever()
|
|
except KeyboardInterrupt:
|
|
_log("\nShutting down...")
|
|
server.shutdown()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|