Create app.py

app.py

@@ -0,0 +1,388 @@
+"""
+Gemini CLI → OpenAI-Compatible API Proxy
+Ultra-fast, reliable, with full streaming support.
+Deploy on HuggingFace Spaces (Docker SDK, port 7860).
+"""
+
+import os
+import json
+import time
+import asyncio
+import logging
+from uuid import uuid4
+from datetime import datetime, timezone
+from contextlib import asynccontextmanager
+from typing import AsyncIterator, Any
+
+import httpx
+from fastapi import FastAPI, Request, HTTPException, Depends
+from fastapi.responses import StreamingResponse, JSONResponse
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+from fastapi.middleware.cors import CORSMiddleware
+from google.oauth2.credentials import Credentials
+from google.auth.transport.requests import Request as GoogleAuthRequest
+
+# ────────────────────── Logging ──────────────────────
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+log = logging.getLogger("gemini-proxy")
+
+# ────────────────────── Config ───────────────────────
+AUTH_PASSWORD  = os.environ.get("GEMINI_AUTH_PASSWORD", "")
+CREDS_JSON     = os.environ.get("GEMINI_CREDENTIALS", "{}")
+CLIENT_ID      = os.environ.get("GEMINI_CLIENT_ID", "")
+CLIENT_SECRET  = os.environ.get("GEMINI_CLIENT_SECRET", "")
+API_BASE       = os.environ.get("GEMINI_API_BASE", "https://cloudcode-pa.googleapis.com")
+
+# ────────────────────── Globals ──────────────────────
+_http: httpx.AsyncClient | None = None
+_creds: Credentials | None = None
+_lock = asyncio.Lock()
+_sec = HTTPBearer(auto_error=False)
+
+MODELS = [
+    "gemini-2.5-pro",
+    "gemini-2.5-flash",
+    "gemini-2.0-flash",
+    "gemini-2.5-pro-search",
+    "gemini-2.5-flash-search",
+    "gemini-2.5-pro-nothinking",
+    "gemini-2.5-flash-nothinking",
+    "gemini-2.5-pro-maxthinking",
+    "gemini-2.5-flash-maxthinking",
+]
+
+# ════════════════════ APP LIFESPAN ═══════════════════
+
+@asynccontextmanager
+async def lifespan(_app: FastAPI):
+    global _http
+    _http = httpx.AsyncClient(
+        timeout=httpx.Timeout(connect=10, read=300, write=30, pool=10),
+        limits=httpx.Limits(max_connections=100, max_keepalive_connections=20),
+        http2=True,
+    )
+    log.info("Proxy ready — %d models", len(MODELS))
+    yield
+    await _http.aclose()
+
+app = FastAPI(title="Gemini OpenAI Proxy", lifespan=lifespan)
+app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
+
+
+# ════════════════════ AUTH ═══════════════════════════
+
+def _auth(c: HTTPAuthorizationCredentials = Depends(_sec)):
+    if not AUTH_PASSWORD:
+        raise HTTPException(500, "Server missing GEMINI_AUTH_PASSWORD")
+    if not c or c.credentials != AUTH_PASSWORD:
+        raise HTTPException(401, "Invalid Bearer token")
+
+
+# ════════════════════ TOKEN ══════════════════════════
+
+async def _token() -> str:
+    global _creds
+    async with _lock:
+        if _creds and _creds.valid and not _creds.expired:
+            return _creds.token
+        return await asyncio.to_thread(_refresh)
+
+
+def _refresh() -> str:
+    global _creds
+    d = json.loads(CREDS_JSON)
+
+    cid  = d.get("client_id") or CLIENT_ID
+    csec = d.get("client_secret") or CLIENT_SECRET
+    rtok = d.get("refresh_token")
+    atok = d.get("access_token") or d.get("token")
+
+    missing = []
+    if not cid:  missing.append("GEMINI_CLIENT_ID")
+    if not csec: missing.append("GEMINI_CLIENT_SECRET")
+    if not rtok: missing.append("refresh_token")
+    if missing:
+        raise HTTPException(500, f"Missing: {', '.join(missing)}")
+
+    exp = None
+    if "expiry_date" in d:
+        exp = datetime.fromtimestamp(d["expiry_date"] / 1000, tz=timezone.utc)
+
+    c = Credentials(
+        token=atok, refresh_token=rtok,
+        token_uri=d.get("token_uri", "https://oauth2.googleapis.com/token"),
+        client_id=cid, client_secret=csec, expiry=exp,
+    )
+    if not c.valid or c.expired:
+        c.refresh(GoogleAuthRequest())
+        log.info("Token refreshed → expires %s", c.expiry)
+    _creds = c
+    return c.token
+
+
+# ════════════════════ GEMINI HELPERS ═════════════════
+
+def _parse_model(model: str):
+    """Returns (base_model, use_search, thinking_budget)."""
+    search = model.endswith("-search")
+    no_think = model.endswith("-nothinking")
+    max_think = model.endswith("-maxthinking")
+
+    base = (model.removesuffix("-search")
+                 .removesuffix("-nothinking")
+                 .removesuffix("-maxthinking"))
+
+    budget = None
+    if no_think:  budget = 0
+    if max_think: budget = 24576
+
+    return base, search, budget
+
+
+def _to_gemini(messages: list, search: bool, budget, **kw) -> dict:
+    """OpenAI messages → Gemini request body."""
+    contents = []
+    sys_parts = []
+
+    for m in messages:
+        role, text = m.get("role", "user"), m.get("content", "")
+        if role == "system":
+            sys_parts.append({"text": text})
+        else:
+            contents.append({
+                "role": "user" if role == "user" else "model",
+                "parts": [{"text": text}],
+            })
+
+    body: dict[str, Any] = {"contents": contents}
+    if sys_parts:
+        body["systemInstruction"] = {"parts": sys_parts}
+
+    gc: dict[str, Any] = {}
+    if kw.get("temperature") is not None: gc["temperature"] = kw["temperature"]
+    if kw.get("max_tokens"):              gc["maxOutputTokens"] = kw["max_tokens"]
+    if kw.get("top_p") is not None:       gc["topP"] = kw["top_p"]
+    if kw.get("stop"):
+        gc["stopSequences"] = kw["stop"] if isinstance(kw["stop"], list) else [kw["stop"]]
+    if budget is not None:
+        gc["thinkingConfig"] = {"thinkingBudget": budget}
+    if gc:
+        body["generationConfig"] = gc
+
+    if search:
+        body["tools"] = [{"googleSearch": {}}]
+
+    return body
+
+
+# ─── Stream parser: handles both SSE and JSON-array ──
+
+async def _gemini_stream(url: str, headers: dict, body: dict) -> AsyncIterator[dict]:
+    """Yields individual Gemini response objects from a stream."""
+    sse_url = url + "?alt=sse"
+
+    async with _http.stream("POST", sse_url, json=body, headers=headers) as r:
+        if r.status_code != 200:
+            err = (await r.aread()).decode(errors="replace")
+            raise HTTPException(r.status_code, f"Gemini: {err[:500]}")
+
+        ct = r.headers.get("content-type", "")
+
+        if "text/event-stream" in ct:
+            # ── SSE mode (fast, line-by-line) ──
+            async for line in r.aiter_lines():
+                if not line.startswith("data:"):
+                    continue
+                payload = line[5:].strip()
+                if not payload or payload == "[DONE]":
+                    continue
+                try:
+                    yield json.loads(payload)
+                except json.JSONDecodeError:
+                    continue
+        else:
+            # ── JSON-array fallback ──
+            buf = ""
+            async for chunk in r.aiter_text():
+                buf += chunk
+                while True:
+                    buf = buf.lstrip(" \t\n\r,[")
+                    if not buf or buf[0] != "{":
+                        # also strip trailing ] at end of array
+                        buf = buf.lstrip("]")
+                        break
+                    # find matching }
+                    depth = 0
+                    in_s = 0   # 1 = inside string
+                    esc = 0    # 1 = next char is escaped
+                    found = -1
+                    for i, c in enumerate(buf):
+                        if esc:
+                            esc = 0; continue
+                        if c == "\\" and in_s:
+                            esc = 1; continue
+                        if c == '"':
+                            in_s ^= 1; continue
+                        if in_s:
+                            continue
+                        if c == "{": depth += 1
+                        elif c == "}":
+                            depth -= 1
+                            if depth == 0:
+                                found = i; break
+                    if found < 0:
+                        break  # incomplete, need more data
+                    try:
+                        yield json.loads(buf[:found + 1])
+                    except json.JSONDecodeError:
+                        pass
+                    buf = buf[found + 1:]
+
+
+def _text(obj: dict) -> str:
+    """Extract non-thought text from Gemini response."""
+    parts = obj.get("candidates", [{}])[0].get("content", {}).get("parts", [])
+    return "".join(p.get("text", "") for p in parts if not p.get("thought"))
+
+
+def _usage(obj: dict) -> dict:
+    m = obj.get("usageMetadata", {})
+    return {
+        "prompt_tokens":     m.get("promptTokenCount", 0),
+        "completion_tokens": m.get("candidatesTokenCount", 0),
+        "total_tokens":      m.get("totalTokenCount", 0),
+    }
+
+
+_FINISH_MAP = {"STOP": "stop", "MAX_TOKENS": "length", "SAFETY": "content_filter"}
+
+def _finish(obj: dict) -> str | None:
+    r = obj.get("candidates", [{}])[0].get("finishReason")
+    return _FINISH_MAP.get(r)
+
+
+# ════════════════════ ROUTES ═════════════════════════
+
+@app.get("/")
+async def health():
+    return {"status": "ok", "service": "gemini-openai-proxy", "models": len(MODELS)}
+
+
+@app.get("/v1/models")
+async def list_models(_=Depends(_auth)):
+    return {
+        "object": "list",
+        "data": [{"id": m, "object": "model", "owned_by": "google", "created": 0} for m in MODELS],
+    }
+
+
+@app.post("/v1/chat/completions")
+async def chat(request: Request, _=Depends(_auth)):
+    body = await request.json()
+    model   = body.get("model", "gemini-2.5-pro")
+    msgs    = body.get("messages", [])
+    stream  = body.get("stream", False)
+
+    if not msgs:
+        raise HTTPException(400, "messages required")
+
+    base, search, budget = _parse_model(model)
+    gemini_body = _to_gemini(
+        msgs, search, budget,
+        temperature=body.get("temperature"),
+        max_tokens=body.get("max_tokens") or body.get("max_completion_tokens"),
+        top_p=body.get("top_p"),
+        stop=body.get("stop"),
+    )
+
+    tok = await _token()
+    hdrs = {"Authorization": f"Bearer {tok}", "Content-Type": "application/json"}
+
+    if stream:
+        url = f"{API_BASE}/v1/models/{base}:streamGenerateContent"
+        return StreamingResponse(
+            _sse_stream(url, hdrs, gemini_body, model),
+            media_type="text/event-stream",
+            headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
+        )
+
+    # ── Non-streaming ──
+    url = f"{API_BASE}/v1/models/{base}:generateContent"
+    return await _non_stream(url, hdrs, gemini_body, model)
+
+
+# ─── Streaming response ─────────────────────────────
+
+async def _sse_stream(url, hdrs, body, model):
+    cid = f"chatcmpl-{uuid4().hex[:24]}"
+    ts  = int(time.time())
+
+    # role chunk
+    yield _sse({"id": cid, "object": "chat.completion.chunk", "created": ts, "model": model,
+                "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}]})
+
+    try:
+        async for obj in _gemini_stream(url, hdrs, body):
+            txt = _text(obj)
+            fin = _finish(obj)
+
+            if txt:
+                yield _sse({"id": cid, "object": "chat.completion.chunk", "created": ts, "model": model,
+                            "choices": [{"index": 0, "delta": {"content": txt}, "finish_reason": None}]})
+            if fin:
+                yield _sse({"id": cid, "object": "chat.completion.chunk", "created": ts, "model": model,
+                            "choices": [{"index": 0, "delta": {}, "finish_reason": fin}]})
+    except HTTPException as e:
+        # Send error as SSE event so client knows what happened
+        yield _sse({"error": {"message": e.detail, "code": e.status_code}})
+
+    yield "data: [DONE]\n\n"
+
+
+# ─── Non-streaming response ─────────────────────────
+
+async def _non_stream(url, hdrs, body, model):
+    # Retry once on 401 (expired token)
+    for attempt in range(2):
+        r = await _http.post(url, json=body, headers=hdrs)
+        if r.status_code == 401 and attempt == 0:
+            global _creds
+            async with _lock:
+                _creds = None
+            tok = await _token()
+            hdrs["Authorization"] = f"Bearer {tok}"
+            continue
+        break
+
+    if r.status_code != 200:
+        raise HTTPException(r.status_code, f"Gemini: {r.text[:500]}")
+
+    data = r.json()
+
+    # Handle error in body
+    if "error" in data:
+        e = data["error"]
+        raise HTTPException(e.get("code", 500), e.get("message", "Unknown"))
+
+    # Gemini may return list or dict
+    if isinstance(data, list):
+        full_text = "".join(_text(item) for item in data)
+        usg = next((_usage(i) for i in data if _usage(i).get("total_tokens")), _usage({}))
+        fin = next((_finish(i) for i in data if _finish(i)), "stop")
+    else:
+        full_text = _text(data)
+        usg = _usage(data)
+        fin = _finish(data) or "stop"
+
+    return JSONResponse({
+        "id": f"chatcmpl-{uuid4().hex[:24]}",
+        "object": "chat.completion",
+        "created": int(time.time()),
+        "model": model,
+        "choices": [{"index": 0, "message": {"role": "assistant", "content": full_text}, "finish_reason": fin}],
+        "usage": usg,
+    })
+
+
+def _sse(obj: dict) -> str:
+    return f"data: {json.dumps(obj, ensure_ascii=False)}\n\n"