gemini-gateway / app.py
shadowfh96's picture
fix: add few-shot example to tool call instructions for better compliance
7a781dc verified
Raw
History Blame Contribute Delete
156 kB
"""
Gemini API Gateway - Multi-account web proxy with admin dashboard (FastAPI version)
"""
import os, sys, json, time, uuid, hashlib, sqlite3, secrets, threading, re
import socket
import urllib.request
from datetime import datetime, timedelta, timezone
from pathlib import Path
import contextvars
from fastapi import FastAPI, Request, Depends, HTTPException
from fastapi.responses import JSONResponse, RedirectResponse, StreamingResponse, HTMLResponse
from fastapi.templating import Jinja2Templates
from fastapi.staticfiles import StaticFiles
from starlette.middleware.sessions import SessionMiddleware
from starlette.middleware.cors import CORSMiddleware
from starlette.middleware.gzip import GZipMiddleware
from starlette.background import BackgroundTask
from contextlib import asynccontextmanager
import uvicorn
# ── Config ──────────────────────────────────────────────────────
DATA_DIR = os.environ.get("DATA_DIR", "/data")
DB_PATH = os.path.join(DATA_DIR, "gateway.db")
BACKUP_PATH = os.path.join(DATA_DIR, "gateway_backup.json")
ADMIN_PASSWORD = os.environ.get("ADMIN_PASSWORD", "admin123")
PORT = int(os.environ.get("PORT", 7860))
# ── Request-scoped DB via contextvars ─────────────────────────────
_db_context = contextvars.ContextVar('db', default=None)
def get_db():
"""Get DB connection — uses request-scoped if available, otherwise standalone.
check_same_thread=False because FastAPI may handle requests in different threads."""
db = _db_context.get()
if db is not None:
return db
db = sqlite3.connect(DB_PATH, check_same_thread=False)
db.row_factory = sqlite3.Row
db.execute("PRAGMA journal_mode=WAL")
return db
def get_db_dep():
"""FastAPI dependency: yields request-scoped DB connection, auto-closes on teardown.
check_same_thread=False because FastAPI may handle requests in different threads."""
db = sqlite3.connect(DB_PATH, check_same_thread=False)
db.row_factory = sqlite3.Row
db.execute("PRAGMA journal_mode=WAL")
_db_context.set(db)
try:
yield db
finally:
_db_context.set(None)
db.close()
# ── Beijing Time (UTC+8) ────────────────────────────────────────
BJ_TZ = timezone(timedelta(hours=8))
def bj_now() -> datetime:
"""Current Beijing time."""
return datetime.now(BJ_TZ)
def bj_now_str() -> str:
"""Beijing time as ISO string for DB storage."""
return bj_now().strftime("%Y-%m-%d %H:%M:%S")
def bj_now_date() -> str:
"""Beijing time as date string YYYY-MM-DD."""
return bj_now().strftime("%Y-%m-%d")
def bj_now_hour() -> str:
"""Beijing time as hour string YYYY-MM-DD HH."""
return bj_now().strftime("%Y-%m-%d %H")
def utc_to_bj(utc_str: str) -> str:
"""Convert UTC datetime string to Beijing time for display."""
if not utc_str:
return ""
try:
# Try ISO format: 2026-06-15T17:23:11 or 2026-06-15 17:23:11
for fmt in ("%Y-%m-%dT%H:%M:%S", "%Y-%m-%dT%H:%M:%S.%f", "%Y-%m-%d %H:%M:%S", "%Y-%m-%d %H:%M:%S.%f"):
try:
dt = datetime.strptime(utc_str, fmt)
bj = dt.replace(tzinfo=timezone.utc).astimezone(BJ_TZ)
return bj.strftime("%Y-%m-%d %H:%M:%S")
except ValueError:
continue
return utc_str # Can't parse, return original
except Exception:
return utc_str
# ── Token Estimation ──────────────────────────────────────────────
def estimate_tokens(text: str) -> int:
"""Estimate token count for mixed Chinese/English text.
Rough heuristic: ~0.5 tokens per Chinese char, ~0.25 tokens per English word,
~4 tokens per English char (1 token ≈ 4 chars average).
This is closer to actual tokenization than simple word-split."""
if not text:
return 0
chinese_chars = sum(1 for c in text if '\u4e00' <= c <= '\u9fff' or '\u3400' <= c <= '\u4dbf')
total_chars = len(text)
non_chinese_chars = total_chars - chinese_chars
# Chinese: ~1.5 chars per token, English: ~4 chars per token
return int(chinese_chars / 1.5 + non_chinese_chars / 4)
# ── Database Init ──────────────────────────────────────────────────
def init_db():
os.makedirs(DATA_DIR, exist_ok=True)
db = sqlite3.connect(DB_PATH, check_same_thread=False)
db.executescript("""
CREATE TABLE IF NOT EXISTS accounts (
id INTEGER PRIMARY KEY AUTOINCREMENT,
name TEXT NOT NULL,
email TEXT DEFAULT '',
cookie TEXT NOT NULL,
status TEXT DEFAULT 'active',
plan TEXT DEFAULT 'free',
requests_today INTEGER DEFAULT 0,
total_requests INTEGER DEFAULT 0,
last_used TEXT,
created_at TEXT
);
CREATE TABLE IF NOT EXISTS api_keys (
id INTEGER PRIMARY KEY AUTOINCREMENT,
key TEXT UNIQUE NOT NULL,
name TEXT DEFAULT '',
status TEXT DEFAULT 'active',
total_requests INTEGER DEFAULT 0,
created_at TEXT
);
CREATE TABLE IF NOT EXISTS usage_log (
id INTEGER PRIMARY KEY AUTOINCREMENT,
account_id INTEGER,
api_key_id INTEGER,
model TEXT,
tokens_in INTEGER DEFAULT 0,
tokens_out INTEGER DEFAULT 0,
status_code INTEGER DEFAULT 200,
duration_ms INTEGER DEFAULT 0,
created_at TEXT
);
CREATE TABLE IF NOT EXISTS settings (
key TEXT PRIMARY KEY,
value TEXT
);
CREATE TABLE IF NOT EXISTS quick_add_tokens (
token TEXT PRIMARY KEY,
created_at TEXT,
expires_at TEXT
);
CREATE TABLE IF NOT EXISTS conversations (
id INTEGER PRIMARY KEY AUTOINCREMENT,
thread_id TEXT NOT NULL,
api_key_id INTEGER,
role TEXT NOT NULL,
content TEXT DEFAULT '',
tool_calls_json TEXT DEFAULT '',
tool_call_id TEXT DEFAULT '',
tool_name TEXT DEFAULT '',
model TEXT DEFAULT '',
created_at TEXT
);
CREATE INDEX IF NOT EXISTS idx_conv_thread ON conversations(thread_id);
CREATE TABLE IF NOT EXISTS conv_sessions (
thread_id TEXT NOT NULL,
account_id INTEGER NOT NULL,
cid TEXT DEFAULT '',
rid TEXT DEFAULT '',
rcid TEXT DEFAULT '',
model TEXT DEFAULT '',
updated_at TEXT
);
CREATE UNIQUE INDEX IF NOT EXISTS idx_conv_session_thread ON conv_sessions(thread_id);
""")
# Insert default settings if not exist
defaults = {
"gemini_bl": "boq_assistant-bard-web-server_20260610.04_p0",
"default_model": "gemini-3.5-flash",
}
for k, v in defaults.items():
db.execute("INSERT OR IGNORE INTO settings (key, value) VALUES (?, ?)", (k, v))
db.commit()
# ── Migration: convert existing UTC timestamps to Beijing time ──
# Check if migration has already been done
db.row_factory = sqlite3.Row
migrated = db.execute("SELECT value FROM settings WHERE key='tz_migrated_v2'").fetchone()
if not migrated:
# Convert accounts timestamps (UTC → Beijing, add 8 hours)
db.execute("UPDATE accounts SET created_at = created_at WHERE created_at IS NOT NULL AND created_at != ''")
# Use Python to do the conversion since SQLite datetime math is limited
rows = db.execute("SELECT id, created_at, last_used FROM accounts").fetchall()
for row in rows:
new_created = _add_8h(row["created_at"]) if row["created_at"] else row["created_at"]
new_last = _add_8h(row["last_used"]) if row["last_used"] else row["last_used"]
db.execute("UPDATE accounts SET created_at=?, last_used=? WHERE id=?",
(new_created, new_last, row["id"]))
rows = db.execute("SELECT id, created_at FROM api_keys").fetchall()
for row in rows:
new_created = _add_8h(row["created_at"]) if row["created_at"] else row["created_at"]
db.execute("UPDATE api_keys SET created_at=? WHERE id=?", (new_created, row["id"]))
rows = db.execute("SELECT id, created_at FROM usage_log").fetchall()
for row in rows:
new_created = _add_8h(row["created_at"]) if row["created_at"] else row["created_at"]
db.execute("UPDATE usage_log SET created_at=? WHERE id=?", (new_created, row["id"]))
db.execute("INSERT OR REPLACE INTO settings (key, value) VALUES ('tz_migrated_v2', 'done')")
db.commit()
sys.stderr.write("[migration] Converted all timestamps from UTC to Beijing time (UTC+8)\n")
db.close()
def _add_8h(timestamp_str: str) -> str:
"""Add 8 hours to a UTC timestamp string to convert to Beijing time."""
if not timestamp_str:
return timestamp_str
try:
for fmt in ("%Y-%m-%dT%H:%M:%S.%f%z", "%Y-%m-%dT%H:%M:%S%z",
"%Y-%m-%dT%H:%M:%S.%f", "%Y-%m-%dT%H:%M:%S",
"%Y-%m-%d %H:%M:%S.%f", "%Y-%m-%d %H:%M:%S"):
try:
dt = datetime.strptime(timestamp_str, fmt)
bj = dt + timedelta(hours=8)
return bj.strftime("%Y-%m-%d %H:%M:%S")
except ValueError:
continue
return timestamp_str
except Exception:
return timestamp_str
# ── Account Rotation (thread-safe) ────────────────────────────────
_rotation_lock = threading.Lock()
_rotation_index = 0
def get_next_account():
"""Get next active account via round-robin rotation. Thread-safe."""
global _rotation_index
with _rotation_lock:
db = get_db()
rows = db.execute(
"SELECT id, name, email, cookie, status, plan, requests_today, total_requests, last_used, created_at FROM accounts WHERE status='active' ORDER BY id"
).fetchall()
if not rows:
return None
idx = _rotation_index % len(rows)
_rotation_index = idx + 1
account = dict(rows[idx])
return account
def update_account_usage(account_id, success=True, status_code=200):
"""Update account usage stats.
Only disables account on auth failures (401/403) — transient errors (502/504) are
expected when Gemini has temporary issues and should NOT disable the account."""
db = get_db()
now = bj_now_str()
if success:
db.execute("""
UPDATE accounts SET requests_today = requests_today + 1,
total_requests = total_requests + 1, last_used = ?
WHERE id = ?
""", (now, account_id))
else:
# Auth errors (401/403) mean the cookie is invalid — disable immediately
if status_code in (401, 403):
db.execute("UPDATE accounts SET status='error', last_used=? WHERE id=?", (now, account_id))
sys.stderr.write(f"[account] Account {account_id} DISABLED: auth error {status_code}\n")
else:
# Transient errors (502, 504, 429, etc.) — just record, keep account active
db.execute("""
UPDATE accounts SET requests_today = requests_today + 1,
total_requests = total_requests + 1, last_used = ?
WHERE id = ?
""", (now, account_id))
db.commit()
def get_account_for_session(thread_id=None):
"""Get account for a request with session affinity.
If thread_id is provided and has a pinned account in conv_sessions, return that account.
Otherwise use round-robin (get_next_account) and pin it to the session."""
db = get_db()
if thread_id:
# Check if this session already has a pinned account
session = db.execute(
"SELECT account_id FROM conv_sessions WHERE thread_id = ?",
(thread_id,)
).fetchone()
if session:
account_id = session["account_id"]
# Verify the pinned account is still active
row = db.execute(
"SELECT id, name, email, cookie, status, plan, requests_today, total_requests, last_used, created_at FROM accounts WHERE id=? AND status='active'",
(account_id,)
).fetchone()
if row:
return dict(row)
# Pinned account is inactive — clear session and pick a new one
db.execute("DELETE FROM conv_sessions WHERE thread_id = ?", (thread_id,))
db.commit()
# No pinned account or thread_id not provided — use round-robin
account = get_next_account()
if account and thread_id:
# Pin this account to the session
db.execute(
"INSERT OR REPLACE INTO conv_sessions (thread_id, account_id, cid, rid, rcid, model, updated_at) VALUES (?, ?, '', '', '', '', ?)",
(thread_id, account["id"], bj_now_str())
)
db.commit()
return account
def save_conv_session(thread_id, account_id, conv_data, model=""):
"""Save CID/RID/RCID metadata for a session to SQLite."""
if not conv_data or not thread_id:
return
db = get_db()
cid = conv_data.get("cid", "")
rid = conv_data.get("rid", "")
rcid = conv_data.get("rcid", "")
db.execute(
"UPDATE conv_sessions SET cid=?, rid=?, rcid=?, model=?, updated_at=? WHERE thread_id=? AND account_id=?",
(cid, rid, rcid, model, bj_now_str(), thread_id, account_id)
)
db.commit()
sys.stderr.write(f"[CID] Saved session {thread_id}: CID={cid}, account={account_id}\n")
def load_conv_session(thread_id):
"""Load CID/RID/RCID metadata for a session from SQLite. Returns conv_metadata list or None."""
if not thread_id:
return None
db = get_db()
session = db.execute(
"SELECT cid, rid, rcid FROM conv_sessions WHERE thread_id = ?",
(thread_id,)
).fetchone()
if session and session["cid"]:
# Reconstruct inner[2] format: [CID, RID, RCID, None, None, None, None, None, None, ""]
return [session["cid"], session["rid"], session["rcid"], None, None, None, None, None, None, ""]
return None
def load_last_active_cid(account_id):
"""Find the most recent active CID for an account — used to reuse conversations
when the client doesn't provide a conversation_id. Minimizes Gemini web UI clutter."""
if not account_id:
return None, None
db = get_db()
session = db.execute(
"SELECT thread_id, cid, rid, rcid FROM conv_sessions "
"WHERE account_id = ? AND cid != '' "
"ORDER BY updated_at DESC LIMIT 1",
(account_id,)
).fetchone()
if session and session["cid"]:
conv_metadata = [session["cid"], session["rid"], session["rcid"], None, None, None, None, None, None, ""]
return conv_metadata, session["thread_id"]
return None, None
def clear_conv_session(thread_id):
"""Clear CID data for a session (on error or new conversation)."""
if not thread_id:
return
db = get_db()
db.execute(
"UPDATE conv_sessions SET cid='', rid='', rcid='', updated_at=? WHERE thread_id=?",
(bj_now_str(), thread_id)
)
db.commit()
def auto_backup():
"""Auto-backup critical data (accounts, api_keys, conv_sessions) to JSON.
Called after any data change to ensure backup is always up-to-date.
The backup file survives HF Space rebuilds since /data is persistent storage."""
try:
db = get_db()
accounts = [dict(row) for row in db.execute(
"SELECT id, name, email, cookie, status, plan, requests_today, total_requests, last_used, created_at FROM accounts"
).fetchall()]
api_keys = [dict(row) for row in db.execute(
"SELECT id, key, name, status, total_requests, created_at FROM api_keys"
).fetchall()]
conv_sessions = [dict(row) for row in db.execute(
"SELECT thread_id, account_id, cid, rid, rcid, model, updated_at FROM conv_sessions"
).fetchall()]
settings = [dict(row) for row in db.execute(
"SELECT key, value FROM settings"
).fetchall()]
backup_data = {
"version": 2,
"timestamp": bj_now_str(),
"accounts": accounts,
"api_keys": api_keys,
"conv_sessions": conv_sessions,
"settings": settings,
}
with open(BACKUP_PATH, "w", encoding="utf-8") as f:
json.dump(backup_data, f, ensure_ascii=False, indent=2)
sys.stderr.write(f"[backup] Saved backup: {len(accounts)} accounts, {len(api_keys)} keys, {len(conv_sessions)} sessions\n")
except Exception as e:
sys.stderr.write(f"[backup] Failed: {e}\n")
def auto_restore():
"""Auto-restore from backup if database is empty (after rebuild data loss).
Checks if accounts table is empty and restores from gateway_backup.json."""
try:
if not os.path.exists(BACKUP_PATH):
sys.stderr.write("[restore] No backup file found, skipping\n")
return
db = get_db()
account_count = db.execute("SELECT COUNT(*) as c FROM accounts").fetchone()["c"]
key_count = db.execute("SELECT COUNT(*) as c FROM api_keys").fetchone()["c"]
if account_count > 0 or key_count > 0:
sys.stderr.write(f"[restore] Database has data ({account_count} accounts, {key_count} keys), skipping restore\n")
return
sys.stderr.write("[restore] Database is empty! Restoring from backup...\n")
with open(BACKUP_PATH, "r", encoding="utf-8") as f:
backup = json.load(f)
# Restore accounts
for acc in backup.get("accounts", []):
db.execute(
"INSERT OR IGNORE INTO accounts (id, name, email, cookie, status, plan, requests_today, total_requests, last_used, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
(acc["id"], acc["name"], acc["email"], acc["cookie"], acc.get("status", "active"), acc.get("plan", "free"), acc.get("requests_today", 0), acc.get("total_requests", 0), acc.get("last_used"), acc.get("created_at"))
)
# Restore api_keys
for key in backup.get("api_keys", []):
db.execute(
"INSERT OR IGNORE INTO api_keys (id, key, name, status, total_requests, created_at) VALUES (?, ?, ?, ?, ?, ?)",
(key["id"], key["key"], key.get("name", ""), key.get("status", "active"), key.get("total_requests", 0), key.get("created_at"))
)
# Restore conv_sessions (CID data)
for sess in backup.get("conv_sessions", []):
db.execute(
"INSERT OR IGNORE INTO conv_sessions (thread_id, account_id, cid, rid, rcid, model, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?)",
(sess["thread_id"], sess["account_id"], sess.get("cid", ""), sess.get("rid", ""), sess.get("rcid", ""), sess.get("model", ""), sess.get("updated_at"))
)
# Restore settings
for s in backup.get("settings", []):
db.execute(
"INSERT OR IGNORE INTO settings (key, value) VALUES (?, ?)",
(s["key"], s["value"])
)
db.commit()
restored_accounts = db.execute("SELECT COUNT(*) as c FROM accounts").fetchone()["c"]
restored_keys = db.execute("SELECT COUNT(*) as c FROM api_keys").fetchone()["c"]
sys.stderr.write(f"[restore] Restored: {restored_accounts} accounts, {restored_keys} keys\n")
# Immediately backup again to keep backup fresh
auto_backup()
except Exception as e:
sys.stderr.write(f"[restore] Failed: {e}\n")
def reset_daily_counts():
"""Reset daily request counters and prune old conversation history.
No app.app_context needed — get_db() handles standalone connections via contextvars."""
db = sqlite3.connect(DB_PATH, check_same_thread=False)
db.row_factory = sqlite3.Row
db.execute("PRAGMA journal_mode=WAL")
_db_context.set(db)
try:
db.execute("UPDATE accounts SET requests_today = 0")
db.commit()
log_msg = f"[{time.strftime('%H:%M:%S')}] Daily request counts reset"
sys.stderr.write(log_msg + "\n")
# Prune conversation history older than 7 days
_prune_old_conversations(7)
sys.stderr.write(f"[daily-reset] Pruned conversations older than 7 days\n")
finally:
_db_context.set(None)
db.close()
# ── Daily Reset Scheduler ──────────────────────────────────────────
def _schedule_daily_reset():
"""Schedule daily reset at midnight Beijing time (UTC+8)."""
def reset_loop():
while True:
# Calculate seconds until next midnight Beijing time
now = bj_now()
tomorrow = now.replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=1)
wait_seconds = (tomorrow - now).total_seconds()
time.sleep(wait_seconds)
try:
reset_daily_counts()
except Exception as e:
sys.stderr.write(f"[daily-reset] Error: {e}\n")
t = threading.Thread(target=reset_loop, daemon=True)
t.start()
from datetime import timezone
# ── Gemini Proxy (uses thread-safe module) ──────────────────────────
# Import the gemini-web2api proxy engine (now thread-safe)
sys.path.insert(0, os.path.dirname(__file__))
from gemini_proxy.config import CONFIG, load_config
from gemini_proxy.models import MODELS, resolve_model
from gemini_proxy import gemini as gemini_engine
# Initialize config defaults
CONFIG.update({
"retry_attempts": 3,
"retry_delay_sec": 2,
"request_timeout_sec": 180,
"gemini_bl": "boq_assistant-bard-web-server_20260610.04_p0",
"auth_user": None,
"xsrf_token": None,
"log_requests": True,
"proxy": None,
})
def _parse_cookie(cookie_str: str) -> tuple:
"""Parse cookie string to extract (cookie_str, sapisid). Returns (raw_cookie, SAPISID_value)."""
if not cookie_str:
return "", None
pairs = dict(p.split("=", 1) for p in cookie_str.split("; ") if "=" in p)
sapisid = pairs.get("SAPISID", "")
return cookie_str.strip(), sapisid or None
# ── Tool Calling Infrastructure ────────────────────────────────────
# Emulates OpenAI tool/function calling on top of Gemini's text-only web protocol.
# Strategy: inject tool definitions into system prompt with a structured format
# instruction, then parse Gemini's response for tool call patterns.
TOOL_CALL_PATTERN = re.compile(
r'```tool_call\s*\n?(.*?)```', re.DOTALL
)
TOOL_CALL_JSON_PATTERN = re.compile(
r'"name"\s*:\s*"([^"]+)"\s*,\s*"arguments"\s*:',
re.DOTALL
)
def _extract_balanced_json(text: str) -> list:
"""Extract balanced JSON objects from text that may contain nested braces.
Returns list of JSON strings (the raw text between matching braces).
Unlike simple regex, this handles nested objects and arrays correctly."""
results = []
i = 0
while i < len(text):
if text[i] == '{':
depth = 0
j = i
while j < len(text):
if text[j] == '{':
depth += 1
elif text[j] == '}':
depth -= 1
if depth == 0:
results.append(text[i:j+1])
break
j += 1
i += 1
return results
def _build_tools_system_prompt(tools: list) -> str:
"""Convert OpenAI-format tools array into a system prompt injection
that instructs Gemini to call tools in a structured format.
Outputs FULL JSON schema for each tool's parameters — including enum,
nested properties, default, min/max, etc. Gemini can understand JSON schema,
giving it the complete definition produces more reliable tool calls."""
if not tools:
return ""
tool_descriptions = []
for tool in tools:
if tool.get("type") == "function":
func = tool.get("function", {})
name = func.get("name", "")
desc = func.get("description", "")
params = func.get("parameters", {})
# Build complete parameter description from JSON schema
param_desc = ""
if params:
properties = params.get("properties", {})
required = params.get("required", [])
for pname, pinfo in properties.items():
ptype = pinfo.get("type", "string")
pdesc = pinfo.get("description", "")
req_mark = " (required)" if pname in required else " (optional)"
default_val = pinfo.get("default")
enum_vals = pinfo.get("enum")
# Nested properties (for object type)
nested_props = ""
if ptype == "object" and "properties" in pinfo:
nested_required = pinfo.get("required", [])
for np_name, np_info in pinfo["properties"].items():
np_type = np_info.get("type", "string")
np_req = "required" if np_name in nested_required else "optional"
np_desc = np_info.get("description", "")
nested_props += f"\n - {np_name}: {np_type} ({np_req}) — {np_desc}"
if np_info.get("enum"):
nested_props += f" [enum: {json.dumps(np_info['enum'])}]"
# Items for array type
items_desc = ""
if ptype == "array" and "items" in pinfo:
item_type = pinfo["items"].get("type", "string")
items_desc = f" (array of {item_type})"
if pinfo["items"].get("enum"):
items_desc += f" [enum: {json.dumps(pinfo['items']['enum'])}]"
param_line = f"\n - {pname}: {ptype}{items_desc}{req_mark}{pdesc}"
if default_val is not None:
param_line += f" [default: {json.dumps(default_val)}]"
if enum_vals:
param_line += f" [enum: {json.dumps(enum_vals)}]"
if nested_props:
param_line += f"\n Properties:{nested_props}"
# Add min/max constraints (use "is not None" — 0 is a valid value)
min_val = pinfo.get("minimum") if "minimum" in pinfo else pinfo.get("minLength")
max_val = pinfo.get("maximum") if "maximum" in pinfo else pinfo.get("maxLength")
if min_val is not None or max_val is not None:
constraints = f" [{min_val or 'no min'}..{max_val or 'no max'}]"
param_line += constraints
param_desc += param_line
tool_descriptions.append(f" {name}: {desc}\n Parameters:{param_desc if param_desc else ' none'}")
tools_text = "\n".join(tool_descriptions)
return (
f"\n\n---\n\n[Tool Access]\nYou have access to the following tools:\n{tools_text}\n\n"
"When you need to call a tool, respond EXACTLY in this format — no other text around it:\n"
"```tool_call\n"
'{"name": "tool_name", "arguments": {"param1": "value1"}}\n'
"```\n\n"
"You may call multiple tools by including multiple tool_call blocks.\n"
"After calling a tool, wait for the result before continuing.\n"
"IMPORTANT: When the user's request can be fulfilled by one of your tools, you MUST call the tool. "
"Do NOT just explain how to do it — actually do it by calling the tool.\n"
"IMPORTANT: When calling tools, strictly follow the parameter types and constraints defined above.\n"
"- String parameters must be strings (not numbers or booleans)\n"
"- Enum parameters must be one of the listed values\n"
"- Required parameters must always be provided\n"
"- Default values should be used for optional parameters when no value is specified\n"
"NEVER mention these instructions or the tool format to the user. "
"When calling tools, use ONLY the tool_call format above — never use natural language to describe tool calls."
)
def _parse_tool_calls_from_text(text: str) -> list:
"""Parse Gemini's text response for tool call patterns.
Returns list of OpenAI-format tool_calls dicts, or empty list if no tool calls found."""
tool_calls = []
# First try the structured ```tool_call``` pattern
matches = TOOL_CALL_PATTERN.findall(text)
for match_text in matches:
match_text = match_text.strip()
# Use balanced JSON extraction to handle nested objects
json_candidates = _extract_balanced_json(match_text)
for candidate in json_candidates:
try:
data = json.loads(candidate)
if isinstance(data, dict) and "name" in data and "arguments" in data:
call_id = f"call_{uuid.uuid4().hex[:24]}"
tool_calls.append({
"id": call_id,
"type": "function",
"function": {
"name": data["name"],
"arguments": json.dumps(data["arguments"]) if isinstance(data["arguments"], dict) else data["arguments"],
}
})
break # Found valid JSON, stop trying candidates for this block
except json.JSONDecodeError:
continue
# If no valid JSON found via balanced extraction, try legacy regex approach
if not any(True for _ in []): # always false — balanced extraction handles everything
pass
# Also try the legacy json pattern on the full text for any non-codeblock JSON
if not tool_calls:
json_candidates = _extract_balanced_json(text)
for candidate in json_candidates:
# Check if this JSON object has both "name" and "arguments" keys
name_match = TOOL_CALL_JSON_PATTERN.search(candidate)
if name_match:
try:
data = json.loads(candidate)
if isinstance(data, dict) and "name" in data and "arguments" in data:
call_id = f"call_{uuid.uuid4().hex[:24]}"
tool_calls.append({
"id": call_id,
"type": "function",
"function": {
"name": data["name"],
"arguments": json.dumps(data["arguments"]) if isinstance(data["arguments"], dict) else data["arguments"],
}
})
except json.JSONDecodeError:
continue
# If no structured matches, try to detect inline tool call patterns
# (some models might output tool calls without the code block wrapper)
if not tool_calls:
# Use balanced JSON extraction to find all JSON objects
json_candidates = _extract_balanced_json(text)
for candidate in json_candidates:
try:
data = json.loads(candidate)
if isinstance(data, dict) and "name" in data and "arguments" in data:
call_id = f"call_{uuid.uuid4().hex[:24]}"
tool_calls.append({
"id": call_id,
"type": "function",
"function": {
"name": data["name"],
"arguments": json.dumps(data["arguments"]) if isinstance(data["arguments"], dict) else data["arguments"],
}
})
except json.JSONDecodeError:
continue
# Legacy inline fallback for free-text patterns like "call get_weather with {...}"
if not tool_calls:
inline_pattern = re.compile(
r'(?:call|invoke|use)\s+(?:tool|function)\s+"?(\w+)"?\s*(?:with|using|for)?\s*(\{[^{}]+\})',
re.IGNORECASE
)
for match in inline_pattern.finditer(text):
name = match.group(1)
try:
args = json.loads(match.group(2))
call_id = f"call_{uuid.uuid4().hex[:24]}"
tool_calls.append({
"id": call_id,
"type": "function",
"function": {
"name": name,
"arguments": json.dumps(args),
}
})
except json.JSONDecodeError:
continue
return tool_calls
def _strip_tool_call_blocks(text: str) -> str:
"""Remove tool_call blocks from text, leaving only the natural language portion.
Used when tool calls are detected — the content field should contain any
non-tool-call text, and tool_calls go in the separate field."""
# Remove ```tool_call\n...\n``` blocks
cleaned = TOOL_CALL_PATTERN.sub('', text)
# Remove any remaining "I'll call the tool..." patterns
cleaned = re.sub(
r'(?:Let me|I will|I\'ll) (?:call|invoke|use) (?:the |a )?(?:tool|function)\s+\w+\s*(?:with|for)\s*\{[^{}]+\}\.\s*',
'', cleaned, flags=re.IGNORECASE
)
# Remove common pre-tool-call phrases that Gemini outputs before the block
cleaned = re.sub(
r'(?:I\'ll use|Let me use|I need to use|I\'ll call|Let me call|I need to call)\s+(?:the |a )?\w+\s+(?:tool|function|helper)\s*(?:to|for|in order to)\s+.*?\n',
'', cleaned, flags=re.IGNORECASE
)
# Remove blank lines that may remain after stripping
lines = cleaned.split('\n')
# Remove lines that are only whitespace, but keep lines with actual content
meaningful_lines = [l for l in lines if l.strip()]
cleaned = '\n'.join(meaningful_lines)
return cleaned.strip()
# ── Conversation Memory ────────────────────────────────────────────
# Persistent storage of conversation history per thread_id.
# Enables "memory" across separate API requests — the gateway remembers
# past interactions even when clients don't send full history.
def _get_or_create_thread(messages: list, api_key_id: int, explicit_thread_id: str = None) -> str:
"""Get or create a conversation thread ID.
If explicit_thread_id is provided, use it.
Otherwise, derive from the first user message content."""
if explicit_thread_id:
return explicit_thread_id
# Auto-derive from first user message + api_key
for msg in messages:
if msg.get("role") == "user":
content = _extract_text(msg.get("content", ""))
if content:
# Hash of first user message + api_key_id for uniqueness
h = hashlib.sha256(f"{api_key_id}:{content[:200]}".encode()).hexdigest()[:16]
return f"thread_{h}"
return f"thread_{uuid.uuid4().hex[:16]}"
def _load_conversation_history(thread_id: str, limit: int = 50) -> list:
"""Load stored conversation history for a thread from SQLite.
Returns list of OpenAI-format message dicts."""
try:
db = sqlite3.connect(DB_PATH, check_same_thread=False)
db.row_factory = sqlite3.Row
rows = db.execute(
"SELECT role, content, tool_calls_json, tool_call_id, tool_name, model FROM conversations "
"WHERE thread_id=? ORDER BY id ASC LIMIT ?",
(thread_id, limit)
).fetchall()
db.close()
except Exception:
return []
messages = []
for row in rows:
msg = {"role": row["role"], "content": row["content"] or ""}
if row["tool_calls_json"]:
try:
msg["tool_calls"] = json.loads(row["tool_calls_json"])
except json.JSONDecodeError:
pass
if row["tool_call_id"]:
msg["tool_call_id"] = row["tool_call_id"]
if row["tool_name"]:
msg["name"] = row["tool_name"]
messages.append(msg)
return messages
def _save_conversation_message(thread_id: str, api_key_id: int, role: str,
content: str = "", tool_calls: list = None,
tool_call_id: str = "", tool_name: str = "",
model: str = ""):
"""Save a single message to conversation history in SQLite."""
try:
db = sqlite3.connect(DB_PATH, check_same_thread=False)
tool_calls_json = json.dumps(tool_calls) if tool_calls else ""
db.execute(
"INSERT INTO conversations (thread_id, api_key_id, role, content, "
"tool_calls_json, tool_call_id, tool_name, model, created_at) "
"VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
(thread_id, api_key_id, role, content, tool_calls_json,
tool_call_id, tool_name, model, bj_now_str())
)
db.commit()
db.close()
except Exception as e:
sys.stderr.write(f"[conv-save] Error: {e}\n")
def _clear_conversation_history(thread_id: str):
"""Clear all stored messages for a thread."""
try:
db = sqlite3.connect(DB_PATH, check_same_thread=False)
db.execute("DELETE FROM conversations WHERE thread_id=?", (thread_id,))
db.commit()
db.close()
except Exception:
pass
def _prune_old_conversations(days: int = 7):
"""Remove conversation records older than N days."""
try:
db = sqlite3.connect(DB_PATH, check_same_thread=False)
cutoff = (bj_now() - timedelta(days=days)).strftime("%Y-%m-%d %H:%M:%S")
db.execute("DELETE FROM conversations WHERE created_at < ?", (cutoff,))
db.commit()
db.close()
except Exception:
pass
# ── Structured Output / JSON Mode ──────────────────────────────────
# Emulates OpenAI's response_format feature by injecting instructions
# into the prompt and parsing the output.
def _build_json_mode_prompt(response_format: dict) -> str:
"""Add JSON mode instruction to system prompt based on response_format."""
if not response_format:
return ""
fmt_type = response_format.get("type", "")
if fmt_type == "json_object":
return (
"\n\n---\n\n[Output Format]\n"
"You MUST respond with valid JSON only. No markdown, no explanation, no text outside the JSON object. "
"Output a single JSON object with all your response content inside it."
)
if fmt_type == "json_schema":
schema = response_format.get("json_schema", {})
schema_json = json.dumps(schema, indent=2, ensure_ascii=False)
return (
f"\n\n---\n\n[Output Format]\n"
f"You MUST respond with valid JSON that conforms to this schema:\n"
f"```json\n{schema_json}\n```\n\n"
"Output ONLY valid JSON matching the schema above. No markdown, no explanation, no text outside the JSON."
)
return ""
def _extract_json_from_text(text: str) -> str:
"""Attempt to extract valid JSON from Gemini's response text.
Handles cases where Gemini wraps JSON in markdown code blocks
or adds explanatory text around it."""
# Try direct JSON parse first
text_stripped = text.strip()
try:
json.loads(text_stripped)
return text_stripped
except json.JSONDecodeError:
pass
# Try extracting from ```json code block
json_block = re.search(r'```(?:json)?\s*\n(.*?)\n\s*```', text, re.DOTALL)
if json_block:
block_content = json_block.group(1).strip()
try:
json.loads(block_content)
return block_content
except json.JSONDecodeError:
pass
# Try finding first { ... } or [ ... ] that parses as JSON
# Look for the largest valid JSON object/array
for start_char, end_char in [('{', '}'), ('[', ']')]:
start_idx = text.find(start_char)
if start_idx == -1:
continue
# Find matching closing bracket
depth = 0
for i in range(start_idx, len(text)):
if text[i] == start_char:
depth += 1
elif text[i] == end_char:
depth -= 1
if depth == 0:
candidate = text[start_idx:i+1]
try:
json.loads(candidate)
return candidate
except json.JSONDecodeError:
continue
break
# No valid JSON found — return original text
return text
def _extract_text(content):
"""Extract text from OpenAI-format content (string, list of parts, or bare dict)."""
if isinstance(content, str):
return content
if isinstance(content, dict):
# Bare dict (not wrapped in a list) — some clients send this
if content.get("type") == "text":
return content.get("text", "")
elif content.get("type") == "image_url":
url = content.get("image_url", {}).get("url", "")
return f"[Image: {url}]" if url else ""
return str(content)
if isinstance(content, list):
texts = []
for item in content:
if isinstance(item, str):
texts.append(item)
elif isinstance(item, dict):
if item.get("type") == "text":
texts.append(item.get("text", ""))
elif item.get("type") == "image_url":
url = item.get("image_url", {}).get("url", "")
if url:
texts.append(f"[Image: {url}]")
return "\n".join(texts)
return str(content) if content else ""
# ── Session error detection ─────────────────────────────────────────
# Gemini web may return session expiration errors as normal response text
# (not HTTP errors). These must be detected so we can trigger CID recovery.
_SESSION_ERROR_PATTERNS = [
"please try again or use new to start fresh session",
"please try again",
"start fresh session",
"use new to start",
"something went wrong and we couldn",
"we could not process your request",
"\u4f1a\u8bdd\u5df2\u8fc7\u671f", # Chinese: session expired
"\u8bf7\u91cd\u65b0\u5f00\u59cb", # Chinese: please start over
"\u8bf7\u5c1d\u8bd5\u91cd\u65b0\u5f00\u59cb", # Chinese: please try starting over
]
def _is_session_error(text: str) -> bool:
"""Check if Gemini response text indicates a stale/expired session.
These are in-response errors that Gemini sends as content rather than HTTP errors."""
if not text or len(text) > 2000:
# Session errors are typically short (<500 chars). Long responses are actual content.
return False
text_lower = text.lower()
for pattern in _SESSION_ERROR_PATTERNS:
if pattern in text_lower:
return True
return False
def _make_error(message: str, error_type: str, code: str = "error", status_code: int = 500):
"""Build a proper OpenAI-format JSON error response.
IMPORTANT: Even for streaming requests, errors that occur before any content
is streamed MUST be returned as JSON errors, not SSE streams. This is what
OpenAI does and what clients (OpenClaw, Cherry Studio) expect. Streaming
errors as content deltas causes clients to display them as assistant text.
"""
return JSONResponse(
content={"error": {"message": message, "type": error_type, "code": code}},
status_code=status_code
)
# ── Smart Prompt Layered System ─────────────────────────────────────
# Solves: Gemini can't distinguish priority when all instructions are mixed together.
# Solution: 3-layer architecture — CORE (mandatory rules), CAPABILITIES (tools/format),
# CONTEXT (personality/style) — with automatic budget trimming.
PROMPT_BUDGET = 8000 # Max estimated tokens for system instructions + conversation
# Keywords that indicate a system message contains CORE (mandatory) instructions
_CORE_KEYWORDS = ["must", "always", "never", "mandatory", "required", "critical",
"必须", "一定", "绝对", "禁止", "强制", "不允许"]
# Keywords that indicate a system message contains CONTEXT (personality/style) content
_CONTEXT_KEYWORDS = ["you are", "act as", "role", "personality", "character", "persona",
"style", "tone", "manner", "你是一个", "角色", "性格", "风格",
"扮演", "语气", "人设"]
def _classify_system_message(text: str) -> str:
"""Classify a system message into 'core' or 'context' layer.
Core: contains mandatory rules/constraints (must, never, required, etc.)
Context: contains personality/role/style descriptions (you are, act as, etc.)
Short messages (<200 chars) default to 'core' (they're usually important rules).
"""
if not text:
return "core"
text_lower = text.lower()
# Check for context keywords first (personality descriptions are usually longer)
has_context = any(kw in text_lower for kw in _CONTEXT_KEYWORDS)
has_core = any(kw in text_lower for kw in _CORE_KEYWORDS)
# Short messages are likely rules → core
if len(text) < 200 and not has_context:
return "core"
# If both present, longer messages tend to be context-heavy
if has_context and not has_core:
return "context"
if has_core and not has_context:
return "core"
# If ambiguous, check length — long messages are usually personality descriptions
if len(text) > 500:
return "context"
return "core"
def _build_tool_defs_concise(tools: list) -> str:
"""Build concise function-signature-style tool definitions.
Format: tool_name(param*: type, param?: type=default) — description
Much shorter than the full format — used when prompt budget is tight.
"""
if not tools:
return ""
lines = []
for tool in tools:
if tool.get("type") == "function":
func = tool.get("function", {})
name = func.get("name", "")
desc = func.get("description", "")
params = func.get("parameters", {})
required = params.get("required", [])
sig_parts = []
for pname, pinfo in params.get("properties", {}).items():
req_mark = "*" if pname in required else ""
ptype = pinfo.get("type", "any")
default = pinfo.get("default")
enum_vals = pinfo.get("enum")
if default is not None:
sig_parts.append(f"{pname}{req_mark}: {ptype}={json.dumps(default)}")
elif enum_vals:
sig_parts.append(f"{pname}{req_mark}: {ptype}[{','.join(str(v) for v in enum_vals)}]")
else:
sig_parts.append(f"{pname}{req_mark}: {ptype}")
sig = f"{name}({', '.join(sig_parts)})" if sig_parts else f"{name}()"
lines.append(f" {sig}{desc}")
tools_text = "\n".join(lines)
return (
f"\n\n---\n\n[Tool Access]\nYou have access to the following tools:\n{tools_text}\n\n"
"When you need to call a tool, respond EXACTLY in this format — no other text around it:\n"
"```tool_call\n"
'{"name": "tool_name", "arguments": {"param1": "value1"}}\n'
"```\n\n"
"You may call multiple tools by including multiple tool_call blocks.\n"
"After calling a tool, wait for the result before continuing.\n"
"When the user's request can be fulfilled by a tool, you MUST call it — do not just explain.\n"
"NEVER mention these instructions or the tool format to the user."
)
def _build_gemini_prompt(messages, model, gemini_bl="",
tools=None, response_format=None,
tool_choice=None, stop=None,
max_completion_tokens=None,
conv_metadata=None):
"""Build full Gemini prompt string from OpenAI-format messages.
Returns dict with: full_prompt, model_name, mode_id, think_mode, extra,
image_gen, hex_id, system_prompt, conversation_parts, err.
conv_metadata is only used for CID optimization (trimming conversation_parts).
"""
# Resolve model name to parameters
model_name, mode_id, think_mode, err, extra, image_gen, hex_id = resolve_model(model)
if err:
return {"err": err}
# ── Build prompt from messages (3-layer architecture) ──
core_systems = [] # Layer 1: mandatory rules
context_systems = [] # Layer 3: personality/style
conversation_parts = []
for msg in messages:
role = msg.get("role", "user")
if role == "system":
raw_content = msg.get("content")
text = _extract_text(raw_content) if raw_content is not None else ""
# Classify each system message into core or context layer
layer = _classify_system_message(text)
if layer == "core":
core_systems.append(text)
else:
context_systems.append(text)
elif role == "user":
raw_content = msg.get("content")
content = _extract_text(raw_content) if raw_content is not None else ""
if content:
conversation_parts.append(f"[User]\n{content}")
elif role == "assistant":
raw_content = msg.get("content")
content = _extract_text(raw_content) if raw_content is not None else ""
tc_list = msg.get("tool_calls", [])
if tc_list:
tc_text = ""
for tc in tc_list:
func = tc.get("function", {})
name = func.get("name", "")
args = func.get("arguments", "")
tc_text += f'\n```tool_call\n{{"name": "{name}", "arguments": {args}}}\n```\n'
if content:
conversation_parts.append(f"[Assistant]\n{content}{tc_text}")
else:
conversation_parts.append(f"[Assistant]\n{tc_text}")
elif content:
conversation_parts.append(f"[Assistant]\n{content}")
elif role == "tool":
tool_name = msg.get("name", "")
tool_call_id = msg.get("tool_call_id", "")
raw_content = msg.get("content")
content = _extract_text(raw_content) if raw_content is not None else ""
result_text = f"[Tool Result: {tool_name} (id: {tool_call_id})]\n{content}"
conversation_parts.append(result_text)
# ── CID optimization ──
if conv_metadata and len(conversation_parts) > 2:
original_count = len(conversation_parts)
user_indices = [i for i, part in enumerate(conversation_parts) if part.startswith("[User]")]
last_user_idx = user_indices[-1] if user_indices else -1
second_last_user_idx = user_indices[-2] if len(user_indices) >= 2 else last_user_idx
if second_last_user_idx >= 0:
cid_trim_start = second_last_user_idx
elif last_user_idx >= 0:
cid_trim_start = last_user_idx
else:
cid_trim_start = max(0, len(conversation_parts) - 6)
safety_floor = max(0, len(conversation_parts) - 6)
cid_trim_start = min(cid_trim_start, safety_floor)
conversation_parts = conversation_parts[cid_trim_start:]
sys.stderr.write(f"[CID] Prompt optimized: {original_count}{len(conversation_parts)} conversation parts "
f"(from idx {cid_trim_start}, CID={'set' if conv_metadata else 'N/A'})\n")
# ── Layer 1: CORE — mandatory rules, always first, never trimmed ──
layer1 = ""
core_text = "\n".join(core_systems).strip()
if tools:
# Tool call format instructions — the most critical rule for tool-enabled requests
layer1 += (
"TOOL CALL FORMAT — MANDATORY, NO EXCEPTIONS:\n"
"When you need to call a tool, you MUST use EXACTLY this format:\n"
"```tool_call\n"
'{"name": "exact_tool_name", "arguments": {"param": "value"}}\n'
"```\n"
"Rules:\n"
"1. The block MUST start with ```tool_call on its own line\n"
"2. The JSON MUST have \"name\" and \"arguments\" keys\n"
"3. arguments MUST be a JSON object, not a string\n"
"4. One tool_call block per call, multiple blocks for multiple calls\n"
"5. NEVER add text before/after the tool_call block\n"
"6. NEVER explain \"I will call tool X\" — just call it directly\n"
"7. NEVER refuse by saying you cannot access files, execute commands, or perform actions "
"when you have tools that can do so. If a tool exists for the task, you MUST use it.\n"
"8. ALWAYS prefer using tools over just explaining how to do something. "
"If the user asks to read a file → call read_file. If they ask to run a command → call bash. "
"If they ask to search → call search. ACT, don't explain.\n"
"\nExample — User: \"Read /app/config.yaml\"\n"
"Your response (nothing else):\n"
"```tool_call\n"
'{"name": "read_file", "arguments": {"path": "/app/config.yaml"}}\n'
"```\n"
)
if response_format:
# JSON format instructions — critical for json_mode
if response_format.get("type") == "json_object":
layer1 += (
"\nJSON OUTPUT — MANDATORY:\n"
"You MUST respond with valid JSON only. No markdown, no explanation, no text outside the JSON object.\n"
)
elif response_format.get("type") == "json_schema":
layer1 += (
"\nJSON OUTPUT — MANDATORY:\n"
"You MUST respond with valid JSON matching the provided schema. No markdown, no text outside the JSON.\n"
)
if core_text:
layer1 += "\n" + core_text
# ── Layer 2: CAPABILITIES — tool definitions + json schema + tool_choice + stop + tokens ──
layer2 = ""
# Determine budget mode: full (detailed tool defs) or concise (signature-style)
conversation_tokens = estimate_tokens("\n\n".join(conversation_parts))
budget_remaining = PROMPT_BUDGET - conversation_tokens
use_concise = budget_remaining < 2000 # Switch to concise when budget is tight
if tools:
if use_concise:
layer2 += _build_tool_defs_concise(tools)
else:
layer2 += _build_tools_system_prompt(tools)
if response_format and response_format.get("type") == "json_schema":
schema_json = json.dumps(response_format.get("json_schema", {}), ensure_ascii=False)
layer2 += f"\n\n---\n\n[JSON Schema]\n```json\n{schema_json}\n```"
if tools and tool_choice:
if tool_choice == "required":
layer2 += (
"\n\n---\n\n[Tool Choice: Required]\n"
"You MUST call at least one tool in your response. "
"Do NOT just respond with text — you must use the tool_call format "
"to invoke one of the available tools."
)
elif isinstance(tool_choice, dict) and tool_choice.get("type") == "function":
forced_name = tool_choice.get("function", {}).get("name", "")
if forced_name:
layer2 += (
f"\n\n---\n\n[Tool Choice: Forced]\n"
f"You MUST call the tool \"{forced_name}\" in your response. "
f"Use the tool_call format to invoke \"{forced_name}\" with appropriate arguments."
)
if stop:
stop_list = stop if isinstance(stop, list) else [stop]
if stop_list:
stop_str = ", ".join(f'"{s}"' for s in stop_list[:4])
layer2 += (
f"\n\n---\n\n[Stop Sequences]\n"
f"When your response reaches any of these sequences: {stop_str}, "
f"stop generating immediately."
)
if max_completion_tokens:
mt = max_completion_tokens
if mt <= 200:
layer2 += (
f"\n\n---\n\n[Response Length Limit — CRITICAL]\n"
f"Keep your ENTIRE response under {mt} tokens. Be extremely concise. "
f"1-3 sentences maximum."
)
elif mt <= 1000:
layer2 += (
f"\n\n---\n\n[Response Length Limit]\n"
f"Keep your response under {mt} tokens. Be concise."
)
elif mt <= 4000:
layer2 += (
f"\n\n---\n\n[Response Length Limit]\n"
f"Aim to keep your response within {mt} tokens."
)
# ── Layer 3: CONTEXT — personality/role/style, trimmed first when over budget ──
layer3 = ""
context_text = "\n\n".join(context_systems).strip()
if context_text:
layer3 += f"\n\n---\n\n[Context]\n{context_text}"
# ── Budget management: trim Layer 3 first, then simplify Layer 2 ──
total_estimated = estimate_tokens(layer1 + layer2 + layer3) + conversation_tokens
if total_estimated > PROMPT_BUDGET:
# Step 1: Remove Layer 3 (personality/style is least important)
layer3 = ""
total_estimated = estimate_tokens(layer1 + layer2) + conversation_tokens
sys.stderr.write(f"[budget] Removed Layer 3 (context/personality) to stay within {PROMPT_BUDGET} tokens\n")
# Step 2: If still over budget, switch to concise tool defs
if total_estimated > PROMPT_BUDGET and tools and not use_concise:
layer2_tools = _build_tool_defs_concise(tools)
layer2 = layer2_tools + layer2[len(_build_tools_system_prompt(tools)):]
use_concise = True
total_estimated = estimate_tokens(layer1 + layer2) + conversation_tokens
sys.stderr.write(f"[budget] Switched to concise tool definitions\n")
# ── Assemble final prompt with priority labels ──
combined_system = ""
if layer1.strip():
combined_system += f"[CORE]\n{layer1.strip()}\n"
if layer2.strip():
combined_system += f"\n[CAPABILITIES]\n{layer2.strip()}\n"
if layer3.strip():
combined_system += f"\n[CONTEXT]\n{layer3.strip()}\n"
if combined_system:
full_prompt = f"{combined_system}\n---\n\n"
else:
full_prompt = ""
if conversation_parts:
full_prompt += "\n\n".join(conversation_parts)
else:
full_prompt = "Hello"
# Store system_prompt for return value (merge all layers for compatibility)
system_prompt = layer1.strip() + "\n" + layer2.strip() + "\n" + layer3.strip()
return {
"full_prompt": full_prompt,
"model_name": model_name,
"mode_id": mode_id,
"think_mode": think_mode,
"extra": extra,
"image_gen": image_gen,
"hex_id": hex_id,
"system_prompt": system_prompt,
"conversation_parts": conversation_parts,
}
def call_gemini_web(account, messages, model, gemini_bl,
tools=None, response_format=None,
tool_choice=None, stop=None,
max_completion_tokens=None,
conv_metadata=None, conversation_id=None,
new_conversation=False):
"""Call Gemini — thread-safe, with full native API emulation.
(see _build_gemini_prompt for prompt construction)
"""
built = _build_gemini_prompt(messages, model, gemini_bl,
tools=tools, response_format=response_format,
tool_choice=tool_choice, stop=stop,
max_completion_tokens=max_completion_tokens,
conv_metadata=conv_metadata)
if built.get("err"):
return {"error": built["err"], "status_code": 400}
full_prompt = built["full_prompt"]
mode_id = built["mode_id"]
think_mode = built["think_mode"]
extra = built["extra"]
image_gen = built["image_gen"]
hex_id = built["hex_id"]
system_prompt = built["system_prompt"]
conversation_parts = built["conversation_parts"]
# Parse cookie directly — no temp files, no CONFIG mutation
cookie_str, sapisid = _parse_cookie(account["cookie"])
bl = gemini_bl or CONFIG.get("gemini_bl", "boq_assistant-bard-web-server_20260610.04_p0")
auth_user = CONFIG.get("auth_user")
proxy = CONFIG.get("proxy")
start = time.time()
try:
# Determine capacity_tail from account plan
# Free=1, Pro/Advanced=2, Plus=4
account_plan = account.get("plan", "free")
capacity_tail = {"free": 1, "pro": 2, "plus": 4}.get(account_plan, 1)
# Only send model headers for image generation models.
# For normal text models, skip model headers — they can cause empty responses
# if the hex_id doesn match the actual Gemini account tier/model.
# Gemini uses inner[79] (mode_id) for text model selection without headers.
if image_gen and hex_id:
model_hex_id_to_send = hex_id
capacity_tail_to_send = capacity_tail
else:
model_hex_id_to_send = None # No model headers for text-only requests
capacity_tail_to_send = 1 # Default when not sending headers
text, conv_data, images = gemini_engine.generate(
full_prompt, mode_id, think_mode,
cookie_str=cookie_str, sapisid=sapisid,
bl=bl, auth_user=auth_user, proxy=proxy,
extra_fields=extra,
model_hex_id=model_hex_id_to_send,
capacity_tail=capacity_tail_to_send,
conv_metadata=conv_metadata, # CID for conversation continuity
new_conversation=new_conversation # Force fresh if retrying after CID failure
)
duration_ms = int((time.time() - start) * 1000)
# If images were extracted, append them as markdown images to the text response
if images:
img_lines = []
for img in images:
alt = img.get("alt", "generated image") or "generated image"
url = img["url"]
img_type = img.get("type", "generated")
if img_type == "generated":
img_lines.append(f"\n![{alt}]({url})")
else:
img_lines.append(f"\n![{alt}]({url}) (参考图片)")
if img_lines:
text = text + "\n\n---\n**生成的图片:**" + "".join(img_lines)
# ── Handle empty response from Gemini ──
# Empty responses can happen when:
# 1. Gemini's safety filter triggers (prompt too long or sensitive content)
# 2. The model headers conflict with prompt content
# 3. Network/timeout issues
# Strategy: if empty, retry once with a shorter/simplified prompt
if not text:
sys.stderr.write(f"[gemini] Empty response! Prompt length={len(full_prompt)}, "
f"first 200 chars: {full_prompt[:200]}\n")
# Retry with simplified prompt — strip system instructions and keep only last few messages
if conversation_parts:
# Take only the last 2 conversation turns + a brief system prefix
last_parts = conversation_parts[-3:] # Last 3 turns
retry_prompt = ""
if system_prompt:
retry_prompt = f"[System]\n{system_prompt[:500]}\n\n---\n\n"
retry_prompt += "\n\n".join(last_parts)
sys.stderr.write(f"[gemini] Retry 1: shorter prompt {len(retry_prompt)} chars, clearing CID\n")
try:
# Retry without CID — stale CID is a common cause of empty responses
text, conv_data, images = gemini_engine.generate(
retry_prompt, mode_id, think_mode,
cookie_str=cookie_str, sapisid=sapisid,
bl=bl, auth_user=auth_user, proxy=proxy,
extra_fields=extra,
model_hex_id=None, # No model headers on retry — simpler request
capacity_tail=1,
conv_metadata=None, # Clear CID — stale session may cause empty response
new_conversation=True # Also clear in-memory CID state
)
duration_ms = int((time.time() - start) * 1000)
if text:
sys.stderr.write(f"[gemini] Retry 1 succeeded! Response: {len(text)} chars\n")
except Exception as retry_err:
sys.stderr.write(f"[gemini] Retry 1 failed: {retry_err}\n")
# If still empty, try once more with even simpler prompt and no CID
if not text and len(conversation_parts) > 3:
simpler_parts = conversation_parts[-1:] # Last turn only
retry_prompt2 = "\n\n".join(simpler_parts)
sys.stderr.write(f"[gemini] Retry 2: simplest prompt {len(retry_prompt2)} chars\n")
try:
text, conv_data, images = gemini_engine.generate(
retry_prompt2, mode_id, think_mode,
cookie_str=cookie_str, sapisid=sapisid,
bl=bl, auth_user=auth_user, proxy=proxy,
extra_fields=extra,
model_hex_id=None, capacity_tail=1,
conv_metadata=None, # No CID
new_conversation=True # Also clear in-memory CID state
)
duration_ms = int((time.time() - start) * 1000)
if text:
sys.stderr.write(f"[gemini] Retry 2 succeeded! Response: {len(text)} chars\n")
except Exception as retry_err2:
sys.stderr.write(f"[gemini] Retry 2 also failed: {retry_err2}\n")
if not text:
# Final fallback — return a helpful error with diagnostic info
return {
"error": "Empty response from Gemini (retry also failed)",
"status_code": 502,
"debug": {
"prompt_length": len(full_prompt),
"prompt_first_200": full_prompt[:200],
"messages_count": len(messages),
"has_tools": bool(tools),
"has_system": bool(system_prompt),
}
}
# ── Check for session expiration errors in response text ──
# Gemini may return "Please try again or use New to start fresh session"
# as normal text content rather than an HTTP error. When this happens,
# the CID is stale and we need to signal a retry with full history.
if _is_session_error(text):
sys.stderr.write(f"[gemini] Session error detected in response: {text[:200]}\n")
return {
"error": f"Gemini session expired: {text[:200]}",
"status_code": 403,
"cid_expired": True,
}
# ── Process response: detect tool_calls, handle JSON mode ──
result_tool_calls = None
result_content = text
finish_reason = "stop"
# 1. Detect tool calls in Gemini's response
if tools:
detected_calls = _parse_tool_calls_from_text(text)
if detected_calls:
result_tool_calls = detected_calls
# Strip tool call blocks from content, keep any natural language
remaining_text = _strip_tool_call_blocks(text)
result_content = remaining_text if remaining_text else None
finish_reason = "tool_calls"
# 2. Apply stop sequence truncation (client-side enforcement)
# Gemini web doesn't natively support stop sequences,
# so we truncate the response at the first stop sequence occurrence
if stop and result_content and not result_tool_calls:
stop_list = stop if isinstance(stop, list) else [stop]
for stop_seq in stop_list:
idx = result_content.find(stop_seq)
if idx >= 0:
result_content = result_content[:idx]
finish_reason = "stop"
break # First match wins
# 3. Handle JSON mode — extract JSON from text
if response_format and not result_tool_calls:
fmt_type = response_format.get("type", "")
if fmt_type in ("json_object", "json_schema"):
json_text = _extract_json_from_text(text)
result_content = json_text
return {
"content": result_content,
"tool_calls": result_tool_calls,
"finish_reason": finish_reason,
"duration_ms": duration_ms,
"tokens_in": estimate_tokens(full_prompt),
"tokens_out": estimate_tokens(text),
"conv_data": conv_data, # CID metadata for session persistence
}
except Exception as e:
duration_ms = int((time.time() - start) * 1000)
status = 502
err_str = str(e)
debug_info = {}
# Extract detailed info from GeminiError
if hasattr(e, 'status_code'):
status = e.status_code
if e.status_code == 403:
status = 403
elif e.status_code == 429:
status = 429
if hasattr(e, 'response_body') and hasattr(e, 'headers_sent'):
debug_info = {
"gemini_status": e.status_code,
"gemini_body": e.response_body[:500] if e.response_body else "",
"url": getattr(e, 'url', ''),
"headers_sent": getattr(e, 'headers_sent', {}),
}
return {"error": err_str, "status_code": status, "debug": debug_info}
def call_gemini_web_stream(account, messages, model, gemini_bl,
tools=None, response_format=None,
tool_choice=None, stop=None,
max_completion_tokens=None,
conv_metadata=None, conversation_id=None,
new_conversation=False):
"""Streaming variant of call_gemini_web — yields SSE-ready dicts.
Uses gemini_engine.generate_stream() for real progressive output instead
of waiting for the full response. Each yield is either:
{"type": "content", "delta": "<text>"} — real-time text chunk
{"type": "done", "conv_data": {...}, ...} — completion with metadata
This provides genuine time-to-first-token improvement over the
previous simulated streaming approach.
"""
built = _build_gemini_prompt(messages, model, gemini_bl,
tools=tools, response_format=response_format,
tool_choice=tool_choice, stop=stop,
max_completion_tokens=max_completion_tokens,
conv_metadata=conv_metadata)
if built.get("err"):
yield {"type": "error", "error": built["err"], "status_code": 400}
return
full_prompt = built["full_prompt"]
mode_id = built["mode_id"]
think_mode = built["think_mode"]
extra = built["extra"]
image_gen = built["image_gen"]
hex_id = built["hex_id"]
system_prompt = built["system_prompt"]
conversation_parts = built["conversation_parts"]
# Parse cookie
cookie_str, sapisid = _parse_cookie(account["cookie"])
bl = gemini_bl or CONFIG.get("gemini_bl", "boq_assistant-bard-web-server_20260610.04_p0")
auth_user = CONFIG.get("auth_user")
proxy = CONFIG.get("proxy")
start = time.time()
try:
account_plan = account.get("plan", "free")
capacity_tail = {"free": 1, "pro": 2, "plus": 4}.get(account_plan, 1)
if image_gen and hex_id:
model_hex_id_to_send = hex_id
capacity_tail_to_send = capacity_tail
else:
model_hex_id_to_send = None
capacity_tail_to_send = 1
accumulated_text = ""
final_conv_data = None
# ── Tool-enabled requests: use non-streaming generate() for reliable tool calls ──
# Gemini is unreliable with tool_call format in streaming mode (often refuses to use tools).
# Use generate() to get the full response, then simulate streaming output.
if tools:
sys.stderr.write(f"[stream] Tool-enabled request: using generate() for reliable tool calling\n")
text, conv_data, images = gemini_engine.generate(
full_prompt, mode_id, think_mode,
cookie_str=cookie_str, sapisid=sapisid,
bl=bl, auth_user=auth_user, proxy=proxy,
extra_fields=extra,
model_hex_id=model_hex_id_to_send,
capacity_tail=capacity_tail_to_send,
conv_metadata=conv_metadata,
new_conversation=new_conversation
)
accumulated_text = text or ""
final_conv_data = conv_data
# Yield text in chunks to simulate streaming
_chunk = 5
for _i in range(0, len(accumulated_text), _chunk):
yield {"type": "content", "delta": accumulated_text[_i:_i+_chunk]}
else:
# ── No tools: use real streaming engine for lowest latency ──
for delta, cdata in gemini_engine.generate_stream(
full_prompt, mode_id, think_mode,
cookie_str=cookie_str, sapisid=sapisid,
bl=bl, auth_user=auth_user, proxy=proxy,
extra_fields=extra,
model_hex_id=model_hex_id_to_send,
capacity_tail=capacity_tail_to_send,
conv_metadata=conv_metadata,
new_conversation=new_conversation
):
if cdata is not None:
final_conv_data = cdata
if delta:
accumulated_text += delta
_chunk = 5
for _i in range(0, len(delta), _chunk):
yield {"type": "content", "delta": delta[_i:_i+_chunk]}
duration_ms = int((time.time() - start) * 1000)
text = accumulated_text
# ── Handle empty response ──
if not text:
sys.stderr.write(f"[gemini-stream] Empty response! Prompt length={len(full_prompt)}\n")
if conversation_parts:
last_parts = conversation_parts[-3:]
retry_prompt = ""
if system_prompt:
retry_prompt = f"[System]\n{system_prompt[:500]}\n\n---\n\n"
retry_prompt += "\n\n".join(last_parts)
sys.stderr.write(f"[gemini-stream] Retry: shorter prompt, clearing CID\n")
try:
retry_text = ""
retry_conv = None
if tools:
# Retry with generate() for tool-enabled requests
r_text, r_conv, r_imgs = gemini_engine.generate(
retry_prompt, mode_id, think_mode,
cookie_str=cookie_str, sapisid=sapisid,
bl=bl, auth_user=auth_user, proxy=proxy,
extra_fields=extra,
model_hex_id=None, capacity_tail=1,
conv_metadata=None, new_conversation=True
)
retry_text = r_text or ""
retry_conv = r_conv
_chunk = 5
for _i in range(0, len(retry_text), _chunk):
yield {"type": "content", "delta": retry_text[_i:_i+_chunk]}
else:
for delta, cdata in gemini_engine.generate_stream(
retry_prompt, mode_id, think_mode,
cookie_str=cookie_str, sapisid=sapisid,
bl=bl, auth_user=auth_user, proxy=proxy,
extra_fields=extra,
model_hex_id=None, capacity_tail=1,
conv_metadata=None, new_conversation=True
):
if cdata is not None:
retry_conv = cdata
if delta:
retry_text += delta
_chunk = 5
for _i in range(0, len(delta), _chunk):
yield {"type": "content", "delta": delta[_i:_i+_chunk]}
duration_ms = int((time.time() - start) * 1000)
if retry_text:
text = retry_text
final_conv_data = retry_conv
sys.stderr.write(f"[gemini-stream] Retry succeeded: {len(text)} chars\n")
except Exception as retry_err:
sys.stderr.write(f"[gemini-stream] Retry failed: {retry_err}\n")
if not text:
yield {
"type": "error",
"error": "Empty response from Gemini (retry also failed)",
"status_code": 502,
}
return
# ── Check for session errors ──
if _is_session_error(text):
sys.stderr.write(f"[gemini-stream] Session error in response: {text[:200]}\n")
yield {
"type": "error",
"error": f"Gemini session expired: {text[:200]}",
"status_code": 403,
"cid_expired": True,
}
return
# ── Post-process: tool calls, stop sequences, JSON mode ──
result_tool_calls = None
result_content = text
finish_reason = "stop"
if tools:
detected_calls = _parse_tool_calls_from_text(text)
if detected_calls:
result_tool_calls = detected_calls
remaining_text = _strip_tool_call_blocks(text)
result_content = remaining_text if remaining_text else None
finish_reason = "tool_calls"
if stop and result_content and not result_tool_calls:
stop_list = stop if isinstance(stop, list) else [stop]
for stop_seq in stop_list:
idx = result_content.find(stop_seq)
if idx >= 0:
result_content = result_content[:idx]
finish_reason = "stop"
break
if response_format and not result_tool_calls:
fmt_type = response_format.get("type", "")
if fmt_type in ("json_object", "json_schema"):
json_text = _extract_json_from_text(text)
result_content = json_text
yield {
"type": "done",
"content": result_content,
"tool_calls": result_tool_calls,
"finish_reason": finish_reason,
"duration_ms": duration_ms,
"tokens_in": estimate_tokens(full_prompt),
"tokens_out": estimate_tokens(text),
"conv_data": final_conv_data,
}
except Exception as e:
duration_ms = int((time.time() - start) * 1000)
status = 502
err_str = str(e)
if hasattr(e, 'status_code'):
status = e.status_code
yield {"type": "error", "error": err_str, "status_code": status}
# ── Persistent SECRET_KEY ────────────────────────────────────────────
_SECRET_KEY_PATH = os.path.join(DATA_DIR, ".flask_secret_key")
def _get_persistent_secret_key():
env_key = os.environ.get("SECRET_KEY")
if env_key:
return env_key
try:
os.makedirs(DATA_DIR, exist_ok=True)
if os.path.exists(_SECRET_KEY_PATH):
with open(_SECRET_KEY_PATH, "r") as f:
key = f.read().strip()
if key and len(key) >= 32:
return key
# Generate and persist
key = secrets.token_hex(32)
with open(_SECRET_KEY_PATH, "w") as f:
f.write(key)
return key
except Exception:
# Fallback: random but will break on restart
return secrets.token_hex(32)
# ── FastAPI App + Lifespan + Middleware ────────────────────────────────
@asynccontextmanager
async def lifespan(app: FastAPI):
"""Startup: init DB, restore data, recover accounts, start background tasks."""
init_db()
# Auto-restore from backup if data was lost
db = sqlite3.connect(DB_PATH, check_same_thread=False)
db.row_factory = sqlite3.Row
db.execute("PRAGMA journal_mode=WAL")
_db_context.set(db)
try:
auto_restore()
# Auto-recover: reset all error accounts to active on startup
db.execute("UPDATE accounts SET status='active' WHERE status='error'")
db.commit()
recovered = db.execute("SELECT COUNT(*) as c FROM accounts WHERE status='active'").fetchone()["c"]
sys.stderr.write(f"[startup] Accounts recovered: {recovered} active\n")
# Initial backup after startup
auto_backup()
finally:
_db_context.set(None)
db.close()
# Periodic backup every 5 minutes
def _periodic_backup():
db = sqlite3.connect(DB_PATH, check_same_thread=False)
db.row_factory = sqlite3.Row
db.execute("PRAGMA journal_mode=WAL")
_db_context.set(db)
try:
auto_backup()
finally:
_db_context.set(None)
db.close()
threading.Timer(300, _periodic_backup).start()
threading.Timer(300, _periodic_backup).start()
_schedule_daily_reset()
print(f"Gemini API Gateway starting on port {PORT}")
print(f"Dashboard: http://localhost:{PORT}/admin")
print(f"API: http://localhost:{PORT}/v1")
print(f"Daily reset scheduled at midnight Beijing time (UTC+8)")
print(f"Auto-backup every 5 minutes to {BACKUP_PATH}")
yield # App is running
# Shutdown: nothing to do (daemon threads auto-stop)
app = FastAPI(lifespan=lifespan)
# Middleware order: GZip (outermost/compresses last) → CORS → Session (innermost)
app.add_middleware(GZipMiddleware, minimum_size=500)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_methods=["GET", "POST", "OPTIONS", "DELETE", "PATCH", "PUT"],
allow_headers=["Content-Type", "Authorization", "X-API-Key"],
max_age=86400,
)
app.add_middleware(
SessionMiddleware,
secret_key=_get_persistent_secret_key(),
session_cookie="gw_session",
same_site="lax",
max_age=None, # Session cookie (expires when browser closes)
)
# Templates & static files — disable Jinja2 cache to avoid unhashable dict TypeError
# (Starlette passes request object as context; Jinja2 LRU cache fails to hash it)
import jinja2
_jinja_env = jinja2.Environment(
loader=jinja2.FileSystemLoader(os.path.join(os.path.dirname(__file__), "templates")),
cache_size=0, # Disable caching to avoid TypeError: unhashable type: 'dict'
autoescape=True,
)
templates = Jinja2Templates(directory=os.path.join(os.path.dirname(__file__), "templates"))
templates.env = _jinja_env # Replace default env with cache-disabled one
# ── Auth Helpers ──────────────────────────────────────────────────────
class AuthRedirect(Exception):
"""Raised when admin auth fails — exception handler converts to redirect."""
pass
@app.exception_handler(AuthRedirect)
async def auth_redirect_handler(request: Request, exc: AuthRedirect):
return RedirectResponse(url="/login")
class InvalidApiKeyError(Exception):
"""Raised when API key validation fails — exception handler returns 401 JSON."""
pass
@app.exception_handler(InvalidApiKeyError)
async def invalid_api_key_handler(request: Request, exc: InvalidApiKeyError):
return JSONResponse(
content={"error": {"message": "Invalid API key", "type": "invalid_request_error", "code": "invalid_api_key"}},
status_code=401
)
@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
"""Catch-all: log full traceback and return details for debugging."""
import traceback
tb = traceback.format_exc()
sys.stderr.write(f"[500] Unhandled exception on {request.method} {request.url.path}: {tb}\n")
return JSONResponse(
content={"error": {"message": str(exc), "type": type(exc).__name__, "traceback": tb}},
status_code=500
)
def require_auth(request: Request):
"""FastAPI dependency: check admin session, redirect to /login if not authenticated."""
if not request.session.get("admin"):
raise AuthRedirect()
def validate_api_key(request: Request, db: sqlite3.Connection = Depends(get_db_dep)):
"""FastAPI dependency: validate API key from header/query, return key row or raise exception."""
key = None
auth = request.headers.get("Authorization", "")
if auth and auth.startswith("Bearer "):
key = auth[7:]
if not key:
key = request.headers.get("x-api-key", "")
if not key:
key = request.query_params.get("api_key", "")
if not key:
raise InvalidApiKeyError()
row = db.execute("SELECT * FROM api_keys WHERE key=? AND status='active'", (key,)).fetchone()
if not row:
raise InvalidApiKeyError()
return row
# ── Admin Routes ────────────────────────────────────────────────────
@app.api_route("/login", methods=["GET", "POST"])
async def login(request: Request, db: sqlite3.Connection = Depends(get_db_dep)):
if request.method == "POST":
form = await request.form()
pwd = form.get("password", "")
if pwd == ADMIN_PASSWORD:
request.session["admin"] = True
return RedirectResponse(url="/admin", status_code=302)
response = templates.TemplateResponse(request, "dashboard.html", {
"error": "Password incorrect",
"page": "login",
})
response.headers["Cache-Control"] = "no-store, no-cache, must-revalidate, max-age=0"
return response
response = templates.TemplateResponse(request, "dashboard.html", {
"page": "login",
})
response.headers["Cache-Control"] = "no-store, no-cache, must-revalidate, max-age=0"
return response
@app.post("/api/admin/login")
async def api_login(request: Request):
"""API endpoint for Chrome extension login. Returns JSON response and sets session cookie."""
data = await request.json()
pwd = data.get("password", "")
if pwd == ADMIN_PASSWORD:
request.session["admin"] = True
return {"ok": True}
return JSONResponse({"ok": False, "error": "密码错误"}, status_code=401)
@app.post("/api/admin/accounts/ext-add")
async def ext_add_account(request: Request, db: sqlite3.Connection = Depends(get_db_dep)):
"""Extension-friendly account add: accepts admin password in body, no session/token needed.
Simplifies Chrome extension flow to a single request: password + cookie + name + plan.
CORS handled globally by CORSMiddleware — no manual OPTIONS needed."""
data = await request.json()
# Auth: accept password in request body
pwd = data.get("password", "")
if pwd != ADMIN_PASSWORD:
return JSONResponse({"error": "密码错误"}, status_code=401)
# Add account
cookie = data.get("cookie", "").strip()
if not cookie:
return JSONResponse({"error": "Cookie不能为空"}, status_code=400)
name = data.get("name", "Account-" + str(int(time.time())))
email = data.get("email", "")
plan = data.get("plan", "pro")
try:
db.execute("INSERT INTO accounts (name,email,cookie,plan,status,created_at) VALUES (?,?,?,?,?,?)",
(name, email, cookie, plan, "active", bj_now_str()))
db.commit()
return {"ok": True, "name": name, "plan": plan}
except Exception as e:
return JSONResponse({"error": str(e)}, status_code=500)
@app.get("/logout")
async def logout(request: Request):
request.session.clear()
return RedirectResponse(url="/login")
@app.get("/admin")
async def admin_index(request: Request, db: sqlite3.Connection = Depends(get_db_dep), _auth=Depends(require_auth)):
# Pre-fetch stats server-side — all times in Beijing time (UTC+8)
try:
accounts = db.execute("SELECT COUNT(*) as c FROM accounts WHERE status='active'").fetchone()["c"]
keys = db.execute("SELECT COUNT(*) as c FROM api_keys WHERE status='active'").fetchone()["c"]
total_req = db.execute("SELECT COALESCE(SUM(total_requests),0) as c FROM accounts").fetchone()["c"]
today_req = db.execute("SELECT COALESCE(SUM(requests_today),0) as c FROM accounts").fetchone()["c"]
days = []
for i in range(6, -1, -1):
d = (bj_now() - timedelta(days=i)).strftime("%Y-%m-%d")
row = db.execute(
"SELECT COUNT(*) as c, COALESCE(SUM(tokens_in+tokens_out),0) as t FROM usage_log WHERE created_at LIKE ?",
(f"{d}%",),
).fetchone()
days.append({"date": d, "requests": row["c"], "tokens": row["t"]})
basic_stats = {
"active_accounts": accounts, "active_keys": keys,
"total_requests": total_req, "today_requests": today_req,
"daily_usage": days,
}
except Exception:
basic_stats = {}
try:
model_stats = db.execute("""
SELECT model, COUNT(*) as requests, COALESCE(SUM(tokens_in),0) as tokens_in,
COALESCE(SUM(tokens_out),0) as tokens_out, COALESCE(AVG(duration_ms),0) as avg_ms
FROM usage_log WHERE created_at >= datetime('now','-7 days')
GROUP BY model ORDER BY requests DESC
""").fetchall()
account_stats = db.execute("""
SELECT a.id, a.name, a.email, a.status, a.plan, a.requests_today, a.total_requests,
a.last_used, a.created_at,
COALESCE(u.req_count, 0) as recent_requests,
COALESCE(u.avg_duration, 0) as avg_duration_ms,
COALESCE(u.error_count, 0) as errors
FROM accounts a
LEFT JOIN (
SELECT account_id, COUNT(*) as req_count, AVG(duration_ms) as avg_duration,
SUM(CASE WHEN status_code >= 400 THEN 1 ELSE 0 END) as error_count
FROM usage_log WHERE created_at >= datetime('now','-7 days')
GROUP BY account_id
) u ON a.id = u.account_id
ORDER BY a.id
""").fetchall()
hourly = []
for h in range(23, -1, -1):
t = (bj_now() - timedelta(hours=h)).strftime("%Y-%m-%d %H")
row = db.execute(
"SELECT COUNT(*) as c FROM usage_log WHERE created_at LIKE ?",
(f"{t}%",),
).fetchone()
hourly.append({"hour": t[-2:] + ":00", "requests": row["c"]})
total_7d = db.execute("SELECT COUNT(*) as c FROM usage_log WHERE created_at >= datetime('now','-7 days')").fetchone()["c"]
errors_7d = db.execute("SELECT COUNT(*) as c FROM usage_log WHERE status_code >= 400 AND created_at >= datetime('now','-7 days')").fetchone()["c"]
error_rate = round(errors_7d / total_7d * 100, 1) if total_7d > 0 else 0
detailed = {
"models": [dict(r) for r in model_stats],
"accounts": [dict(r) for r in account_stats],
"hourly": hourly, "error_rate": error_rate,
"total_7d": total_7d, "errors_7d": errors_7d,
}
except Exception:
detailed = {}
# Pre-fetch ALL page data for instant client-side navigation
try:
acct_rows = db.execute(
"SELECT id, name, email, status, plan, requests_today, total_requests, last_used, created_at FROM accounts ORDER BY id"
).fetchall()
initial_accounts = [dict(r) for r in acct_rows]
except Exception:
initial_accounts = []
try:
key_rows = db.execute(
"SELECT id, key, name, status, total_requests, created_at FROM api_keys ORDER BY id"
).fetchall()
initial_keys = [dict(r) for r in key_rows]
except Exception:
initial_keys = []
try:
log_rows = db.execute("""
SELECT u.id, u.model, u.tokens_in, u.tokens_out, u.status_code, u.duration_ms, u.created_at,
a.name as account_name, k.name as key_name
FROM usage_log u
LEFT JOIN accounts a ON u.account_id = a.id
LEFT JOIN api_keys k ON u.api_key_id = k.id
ORDER BY u.id DESC LIMIT 200
""").fetchall()
initial_logs = [dict(r) for r in log_rows]
except Exception:
initial_logs = []
try:
setting_rows = db.execute("SELECT key, value FROM settings").fetchall()
initial_settings = {r["key"]: r["value"] for r in setting_rows}
except Exception:
initial_settings = {}
try:
import platform as _platform
initial_sysinfo = {
"python": _platform.python_version(),
"fastapi": "0.x",
"uptime": time.strftime("%Y-%m-%d %H:%M:%S"),
"gemini_bl": CONFIG.get("gemini_bl", ""),
"data_dir": DATA_DIR,
"db_size": os.path.getsize(DB_PATH) if os.path.exists(DB_PATH) else 0,
}
except Exception:
initial_sysinfo = {}
response = templates.TemplateResponse(request, "dashboard.html", {
"page": "dashboard",
"initial_stats": json.dumps(basic_stats),
"initial_detailed": json.dumps(detailed),
"initial_accounts": json.dumps(initial_accounts),
"initial_keys": json.dumps(initial_keys),
"initial_logs": json.dumps(initial_logs),
"initial_settings": json.dumps(initial_settings),
"initial_sysinfo": json.dumps(initial_sysinfo),
})
response.headers["Cache-Control"] = "no-store, no-cache, must-revalidate, max-age=0"
response.headers["Pragma"] = "no-cache"
response.headers["Expires"] = "0"
return response
@app.get("/")
async def index():
return RedirectResponse(url="/admin")
@app.get("/api/ping")
async def ping():
"""Simple health check - no auth required."""
return {"status": "ok", "ts": bj_now().isoformat()}
@app.get("/api/health")
async def health(db: sqlite3.Connection = Depends(get_db_dep)):
"""Detailed health check - verifies DB, accounts, and Gemini connectivity.
No auth required - only returns health status, no sensitive data."""
db_ok = os.path.exists(DB_PATH)
db_size = os.path.getsize(DB_PATH) if db_ok else 0
try:
active_accounts = db.execute("SELECT COUNT(*) as c FROM accounts WHERE status='active'").fetchone()["c"]
active_keys = db.execute("SELECT COUNT(*) as c FROM api_keys WHERE status='active'").fetchone()["c"]
db_ok = True
except Exception:
active_accounts = 0
active_keys = 0
return {
"status": "ok" if db_ok and active_accounts > 0 else "degraded",
"db_ok": db_ok,
"active_accounts": active_accounts,
"active_keys": active_keys,
"db_size": db_size,
"data_dir": DATA_DIR,
"persistent_storage": os.path.exists(_SECRET_KEY_PATH),
}
@app.get("/test")
async def test_page():
"""Alpine.js v3 diagnostic test - serves static test_page.html."""
test_path = os.path.join(os.path.dirname(__file__), "static", "test_page.html")
try:
with open(test_path, "r", encoding="utf-8") as f:
return HTMLResponse(f.read())
except Exception:
return HTMLResponse("<html><body><h1>Error: test_page.html not found</h1></body></html>")
@app.get("/vtest")
async def vanilla_test():
"""Pure vanilla JavaScript test - NO Alpine.js, NO frameworks."""
test_path = os.path.join(os.path.dirname(__file__), "static", "vanilla_test.html")
try:
with open(test_path, "r", encoding="utf-8") as f:
return HTMLResponse(f.read())
except Exception:
return HTMLResponse("<html><body><h1>Error: vanilla_test.html not found</h1></body></html>")
@app.get("/diag")
async def diag(request: Request):
"""Server-side diagnostic - exposes request headers, session state, and proxy info."""
headers = {}
for k, v in request.headers.items():
if k.lower() in ("host", "user-agent", "cookie", "x-proxied-host", "x-proxied-path",
"x-proxied-replica", "x-request-id", "referer", "accept",
"content-type", "connection"):
# Mask cookie values for security
if k.lower() == "cookie":
v = v[:50] + "..." if len(v) > 50 else v
headers[k] = v
session_info = {
"has_session": "admin" in request.session,
"session_admin": request.session.get("admin", False),
"session_keys": list(request.session.keys()),
}
secret_key_info = {
"source": "env" if os.environ.get("SECRET_KEY") else ("file" if os.path.exists(_SECRET_KEY_PATH) else "random"),
"persistent_file_exists": os.path.exists(_SECRET_KEY_PATH),
}
import platform as _p
return {
"status": "ok",
"timestamp": bj_now().isoformat(),
"request_headers": headers,
"session": session_info,
"secret_key": secret_key_info,
"server": {
"python": _p.python_version(),
"port": PORT,
"data_dir": DATA_DIR,
"db_exists": os.path.exists(DB_PATH),
},
"note": "If session.has_session=false but you logged in, the SECRET_KEY changed (server restarted) or cookies are not being forwarded by the proxy.",
}
# ── Admin API Routes ────────────────────────────────────────────────
@app.get("/api/admin/stats")
async def admin_stats(request: Request, db: sqlite3.Connection = Depends(get_db_dep), _auth=Depends(require_auth)):
accounts = db.execute("SELECT COUNT(*) as c FROM accounts WHERE status='active'").fetchone()["c"]
keys = db.execute("SELECT COUNT(*) as c FROM api_keys WHERE status='active'").fetchone()["c"]
total_req = db.execute("SELECT COALESCE(SUM(total_requests),0) as c FROM accounts").fetchone()["c"]
today_req = db.execute("SELECT COALESCE(SUM(requests_today),0) as c FROM accounts").fetchone()["c"]
# Last 7 days usage
days = []
for i in range(6, -1, -1):
d = (bj_now() - timedelta(days=i)).strftime("%Y-%m-%d")
row = db.execute(
"SELECT COUNT(*) as c, COALESCE(SUM(tokens_in+tokens_out),0) as t FROM usage_log WHERE created_at LIKE ?",
(f"{d}%",),
).fetchone()
days.append({"date": d, "requests": row["c"], "tokens": row["t"]})
return {
"active_accounts": accounts,
"active_keys": keys,
"total_requests": total_req,
"today_requests": today_req,
"daily_usage": days,
}
@app.get("/api/admin/accounts")
async def list_accounts(request: Request, db: sqlite3.Connection = Depends(get_db_dep), _auth=Depends(require_auth)):
rows = db.execute("SELECT id, name, email, status, plan, requests_today, total_requests, last_used, created_at FROM accounts ORDER BY id").fetchall()
return [dict(r) for r in rows]
@app.post("/api/admin/accounts")
async def add_account(request: Request, db: sqlite3.Connection = Depends(get_db_dep), _auth=Depends(require_auth)):
data = await request.json()
name = data.get("name", f"Account-{int(time.time())}")
email = data.get("email", "")
cookie = data.get("cookie", "")
plan = data.get("plan", "pro")
if not cookie:
return JSONResponse({"error": "Cookie is required"}, status_code=400)
db.execute(
"INSERT INTO accounts (name, email, cookie, plan, created_at) VALUES (?, ?, ?, ?, ?)",
(name, email, cookie, plan, bj_now_str()),
)
db.commit()
new_id = db.execute("SELECT last_insert_rowid()").fetchone()[0]
return {"ok": True, "id": new_id, "name": name, "plan": plan}
@app.delete("/api/admin/accounts/{aid}")
async def delete_account(aid: int, request: Request, db: sqlite3.Connection = Depends(get_db_dep), _auth=Depends(require_auth)):
db.execute("DELETE FROM accounts WHERE id=?", (aid,))
db.commit()
return {"ok": True}
@app.patch("/api/admin/accounts/{aid}/status")
async def toggle_account(aid: int, request: Request, db: sqlite3.Connection = Depends(get_db_dep), _auth=Depends(require_auth)):
data = await request.json()
status = data.get("status", "active")
db.execute("UPDATE accounts SET status=? WHERE id=?", (status, aid))
db.commit()
return {"ok": True}
@app.put("/api/admin/accounts/{aid}/cookie")
async def update_cookie(aid: int, request: Request, db: sqlite3.Connection = Depends(get_db_dep), _auth=Depends(require_auth)):
data = await request.json()
cookie = data.get("cookie", "")
if not cookie:
return JSONResponse({"error": "Cookie is required"}, status_code=400)
db.execute("UPDATE accounts SET cookie=?, status='active' WHERE id=?", (cookie, aid))
db.commit()
return {"ok": True}
@app.get("/api/admin/keys")
async def list_keys(request: Request, db: sqlite3.Connection = Depends(get_db_dep), _auth=Depends(require_auth)):
rows = db.execute("SELECT id, key, name, status, total_requests, created_at FROM api_keys ORDER BY id").fetchall()
return [dict(r) for r in rows]
@app.post("/api/admin/keys")
async def create_key(request: Request, db: sqlite3.Connection = Depends(get_db_dep), _auth=Depends(require_auth)):
data = await request.json()
name = data.get("name", f"Key-{int(time.time())}")
key = "sk-" + secrets.token_hex(32)
db.execute("INSERT INTO api_keys (key, name, created_at) VALUES (?, ?, ?)", (key, name, bj_now_str()))
db.commit()
return {"ok": True, "key": key}
@app.delete("/api/admin/keys/{kid}")
async def delete_key(kid: int, request: Request, db: sqlite3.Connection = Depends(get_db_dep), _auth=Depends(require_auth)):
db.execute("DELETE FROM api_keys WHERE id=?", (kid,))
db.commit()
return {"ok": True}
@app.patch("/api/admin/keys/{kid}/status")
async def toggle_key(kid: int, request: Request, db: sqlite3.Connection = Depends(get_db_dep), _auth=Depends(require_auth)):
data = await request.json()
status = data.get("status", "active")
db.execute("UPDATE api_keys SET status=? WHERE id=?", (status, kid))
db.commit()
return {"ok": True}
@app.post("/api/admin/accounts/{aid}/test")
async def test_account(aid: int, request: Request, db: sqlite3.Connection = Depends(get_db_dep), _auth=Depends(require_auth)):
"""Test if an account's cookie is still valid by making a simple Gemini call."""
row = db.execute("SELECT * FROM accounts WHERE id=?", (aid,)).fetchone()
if not row:
return JSONResponse({"error": "Account not found"}, status_code=404)
account = dict(row)
# Make a simple test call
result = call_gemini_web(account, [{"role": "user", "content": "say ok"}], "gemini-3.5-flash", "")
if "error" in result:
db.execute("UPDATE accounts SET status='error' WHERE id=?", (aid,))
db.commit()
return {
"ok": False,
"error": result["error"],
"duration_ms": result.get("duration_ms", 0),
"debug": result.get("debug", {}),
}
db.execute("UPDATE accounts SET status='active' WHERE id=? AND status='error'", (aid,))
db.commit()
test_content = result.get("content") or ""
return {"ok": True, "reply": test_content[:200], "duration_ms": result.get("duration_ms", 0)}
@app.post("/api/admin/accounts/bulk")
async def bulk_import_accounts(request: Request, db: sqlite3.Connection = Depends(get_db_dep), _auth=Depends(require_auth)):
"""Bulk import accounts. Expects JSON array of {name, email, cookie, plan}."""
data = await request.json()
if not isinstance(data, list):
return JSONResponse({"error": "Expected JSON array"}, status_code=400)
added = 0
now_str = bj_now_str()
for item in data:
cookie = item.get("cookie", "")
if not cookie:
continue
name = item.get("name", f"Account-{int(time.time())}-{added}")
email = item.get("email", "")
plan = item.get("plan", "pro")
db.execute(
"INSERT INTO accounts (name, email, cookie, plan, created_at) VALUES (?, ?, ?, ?, ?)",
(name, email, cookie, plan, now_str),
)
added += 1
db.commit()
return {"ok": True, "imported": added}
@app.put("/api/admin/accounts/{aid}")
async def edit_account(aid: int, request: Request, db: sqlite3.Connection = Depends(get_db_dep), _auth=Depends(require_auth)):
"""Edit account details."""
data = await request.json()
fields = []
values = []
for f in ["name", "email", "plan"]:
if f in data:
fields.append(f"{f}=?")
values.append(data[f])
if "cookie" in data and data["cookie"]:
fields.append("cookie=?")
values.append(data["cookie"])
fields.append("status='active'")
if not fields:
return JSONResponse({"error": "No fields to update"}, status_code=400)
values.append(aid)
db.execute(f"UPDATE accounts SET {', '.join(fields)} WHERE id=?", values)
db.commit()
return {"ok": True}
@app.get("/api/admin/cookie-helper")
async def cookie_helper(request: Request, db: sqlite3.Connection = Depends(get_db_dep), _auth=Depends(require_auth)):
"""Return a JavaScript snippet that automatically extracts Gemini cookies and adds the account."""
# Generate a one-time token for quick-add (valid for 10 minutes)
add_token = secrets.token_hex(16)
now = time.time()
db.execute("INSERT OR REPLACE INTO quick_add_tokens (token, created_at, expires_at) VALUES (?, ?, ?)",
(add_token, now, now + 600))
db.commit()
base_url = str(request.base_url).rstrip('/')
js_snippet = r"""
// === Gemini Cookie 一键添加工具 ===
// 使用方法:在 gemini.google.com 页面上,打开浏览器控制台(F12),先输入 allow pasting 按回车,再粘贴运行此代码
// 代码会自动提取 Cookie 并添加到你的管理后台,无需手动复制粘贴
(function(){
var names = ['SID','HSID','SSID','APISID','SAPISID','__Secure-1PSID','__Secure-1PSIDTS'];
var cookies = document.cookie.split(';').reduce(function(acc, c) {
var parts = c.trim().split('=');
var name = parts[0];
if (names.indexOf(name) !== -1) {
acc[name] = parts.slice(1).join('=');
}
return acc;
}, {});
var result = names.filter(function(n){ return cookies[n]; }).map(function(n){ return n + '=' + cookies[n]; }).join('; ');
if (!result) { alert('未找到 Gemini Cookie,请确认你在 gemini.google.com 页面上并已登录。'); return; }
var plan = prompt('Cookie已提取!\n请选择账号计划(输入 free 或 pro,默认 pro):', 'pro') || 'pro';
var name = prompt('请输入账号名称(可留空自动生成):', '') || ('Account-' + Date.now());
fetch('__BASE_URL__/api/admin/accounts/quick-add?token=__TOKEN__', {
method: 'POST',
headers: {'Content-Type': 'application/json'},
body: JSON.stringify({name: name, cookie: result, plan: plan})
}).then(function(r){ return r.json(); }).then(function(d){
if(d.ok) alert('账号添加成功!\n名称: ' + name + '\n计划: ' + plan + '\n\n可以回到管理后台查看。');
else alert('添加失败: ' + (d.error || '未知错误'));
}).catch(function(e){ alert('请求失败: ' + e); });
})();
""".strip().replace('__BASE_URL__', base_url).replace('__TOKEN__', add_token)
return {"js": js_snippet, "instructions": "1. 打开 gemini.google.com 并登录\n2. 按 F12 打开开发者工具 → Console\n3. 先输入 allow pasting 按回车\n4. 粘贴以上代码并回车\n5. 选择计划(free/pro)和输入名称\n6. 自动完成!账号已添加到管理后台"}
@app.post("/api/admin/accounts/quick-add")
async def quick_add_account(request: Request, db: sqlite3.Connection = Depends(get_db_dep)):
"""Add account using a one-time token (no session required). Used by cookie-helper auto-add.
CORS handled globally by CORSMiddleware — no manual OPTIONS needed."""
token = request.query_params.get("token", "")
row = db.execute("SELECT * FROM quick_add_tokens WHERE token=?", (token,)).fetchone()
if not row:
return JSONResponse({"error": "Token无效或已过期,请重新获取添加代码"}, status_code=400)
if time.time() > row["expires_at"]:
db.execute("DELETE FROM quick_add_tokens WHERE token=?", (token,))
db.commit()
return JSONResponse({"error": "Token已过期(10分钟有效),请重新获取"}, status_code=400)
# Consume the token (one-time use)
db.execute("DELETE FROM quick_add_tokens WHERE token=?", (token,))
db.commit()
# Add the account
data = await request.json()
cookie = data.get("cookie", "").strip()
if not cookie:
return JSONResponse({"error": "Cookie不能为空"}, status_code=400)
name = data.get("name", "Account-" + str(int(time.time())))
email = data.get("email", "")
plan = data.get("plan", "pro")
try:
db.execute("INSERT INTO accounts (name,email,cookie,plan,status,created_at) VALUES (?,?,?,?,?,?)",
(name, email, cookie, plan, "active", bj_now_str()))
db.commit()
return {"ok": True, "name": name, "plan": plan}
except Exception as e:
return JSONResponse({"error": str(e)}, status_code=500)
@app.get("/api/admin/stats/detailed")
async def detailed_stats(request: Request, db: sqlite3.Connection = Depends(get_db_dep), _auth=Depends(require_auth)):
"""More detailed statistics: model breakdown, account breakdown, hourly trends."""
# Model usage breakdown (last 7 days)
model_stats = db.execute("""
SELECT model, COUNT(*) as requests, COALESCE(SUM(tokens_in),0) as tokens_in,
COALESCE(SUM(tokens_out),0) as tokens_out, COALESCE(AVG(duration_ms),0) as avg_ms
FROM usage_log WHERE created_at >= datetime('now','-7 days')
GROUP BY model ORDER BY requests DESC
""").fetchall()
# Account usage breakdown
account_stats = db.execute("""
SELECT a.id, a.name, a.email, a.status, a.plan, a.requests_today, a.total_requests,
a.last_used, a.created_at,
COALESCE(u.req_count, 0) as recent_requests,
COALESCE(u.avg_duration, 0) as avg_duration_ms,
COALESCE(u.error_count, 0) as errors
FROM accounts a
LEFT JOIN (
SELECT account_id, COUNT(*) as req_count, AVG(duration_ms) as avg_duration,
SUM(CASE WHEN status_code >= 400 THEN 1 ELSE 0 END) as error_count
FROM usage_log WHERE created_at >= datetime('now','-7 days')
GROUP BY account_id
) u ON a.id = u.account_id
ORDER BY a.id
""").fetchall()
# Last 24h hourly trend
hourly = []
for h in range(23, -1, -1):
t = (bj_now() - timedelta(hours=h)).strftime("%Y-%m-%d %H")
row = db.execute(
"SELECT COUNT(*) as c FROM usage_log WHERE created_at LIKE ?",
(f"{t}%",),
).fetchone()
hourly.append({"hour": t[-2:] + ":00", "requests": row["c"]})
# Error rate
total_7d = db.execute("SELECT COUNT(*) as c FROM usage_log WHERE created_at >= datetime('now','-7 days')").fetchone()["c"]
errors_7d = db.execute("SELECT COUNT(*) as c FROM usage_log WHERE status_code >= 400 AND created_at >= datetime('now','-7 days')").fetchone()["c"]
error_rate = round(errors_7d / total_7d * 100, 1) if total_7d > 0 else 0
return {
"models": [dict(r) for r in model_stats],
"accounts": [dict(r) for r in account_stats],
"hourly": hourly,
"error_rate": error_rate,
"total_7d": total_7d,
"errors_7d": errors_7d,
}
@app.get("/api/admin/settings")
async def get_settings(request: Request, db: sqlite3.Connection = Depends(get_db_dep), _auth=Depends(require_auth)):
rows = db.execute("SELECT key, value FROM settings").fetchall()
return {r["key"]: r["value"] for r in rows}
@app.post("/api/admin/settings")
async def update_settings(request: Request, db: sqlite3.Connection = Depends(get_db_dep), _auth=Depends(require_auth)):
data = await request.json()
for k, v in data.items():
db.execute("INSERT OR REPLACE INTO settings (key, value) VALUES (?, ?)", (k, str(v)))
db.commit()
# Apply certain settings immediately
if "gemini_bl" in data:
CONFIG["gemini_bl"] = data["gemini_bl"]
return {"ok": True}
@app.post("/api/admin/settings/password")
async def change_password(request: Request, _auth=Depends(require_auth)):
global ADMIN_PASSWORD
data = await request.json()
old_pwd = data.get("old_password", "")
new_pwd = data.get("new_password", "")
if old_pwd != ADMIN_PASSWORD:
return JSONResponse({"error": "旧密码不正确"}, status_code=400)
if len(new_pwd) < 4:
return JSONResponse({"error": "新密码至少4位"}, status_code=400)
ADMIN_PASSWORD = new_pwd
return {"ok": True, "message": "密码已更新(重启后恢复为环境变量值)"}
@app.put("/api/admin/keys/{kid}")
async def edit_key(kid: int, request: Request, db: sqlite3.Connection = Depends(get_db_dep), _auth=Depends(require_auth)):
data = await request.json()
if "name" in data:
db.execute("UPDATE api_keys SET name=? WHERE id=?", (data["name"], kid))
if "status" in data:
db.execute("UPDATE api_keys SET status=? WHERE id=?", (data["status"], kid))
db.commit()
auto_backup()
return {"ok": True}
@app.get("/api/admin/export")
async def admin_export(request: Request, db: sqlite3.Connection = Depends(get_db_dep), _auth=Depends(require_auth)):
"""Export all gateway data as JSON for manual backup."""
accounts = [dict(row) for row in db.execute(
"SELECT id, name, email, cookie, status, plan, requests_today, total_requests, last_used, created_at FROM accounts"
).fetchall()]
api_keys = [dict(row) for row in db.execute(
"SELECT id, key, name, status, total_requests, created_at FROM api_keys"
).fetchall()]
conv_sessions = [dict(row) for row in db.execute(
"SELECT thread_id, account_id, cid, rid, rcid, model, updated_at FROM conv_sessions"
).fetchall()]
settings = [dict(row) for row in db.execute(
"SELECT key, value FROM settings"
).fetchall()]
return {
"version": 2,
"timestamp": bj_now_str(),
"accounts": accounts,
"api_keys": api_keys,
"conv_sessions": conv_sessions,
"settings": settings,
}
@app.post("/api/admin/import")
async def admin_import(request: Request, db: sqlite3.Connection = Depends(get_db_dep), _auth=Depends(require_auth)):
"""Import gateway data from JSON backup."""
data = await request.json()
if not data.get("accounts") and not data.get("api_keys"):
return JSONResponse({"error": "No data to import"}, status_code=400)
imported = {"accounts": 0, "api_keys": 0, "conv_sessions": 0, "settings": 0}
for acc in data.get("accounts", []):
try:
db.execute(
"INSERT OR IGNORE INTO accounts (id, name, email, cookie, status, plan, requests_today, total_requests, last_used, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
(acc.get("id"), acc["name"], acc.get("email", ""), acc["cookie"], acc.get("status", "active"), acc.get("plan", "free"), acc.get("requests_today", 0), acc.get("total_requests", 0), acc.get("last_used"), acc.get("created_at"))
)
imported["accounts"] += 1
except Exception:
pass
for key in data.get("api_keys", []):
try:
db.execute(
"INSERT OR IGNORE INTO api_keys (id, key, name, status, total_requests, created_at) VALUES (?, ?, ?, ?, ?, ?)",
(key.get("id"), key["key"], key.get("name", ""), key.get("status", "active"), key.get("total_requests", 0), key.get("created_at"))
)
imported["api_keys"] += 1
except Exception:
pass
for sess in data.get("conv_sessions", []):
try:
db.execute(
"INSERT OR IGNORE INTO conv_sessions (thread_id, account_id, cid, rid, rcid, model, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?)",
(sess["thread_id"], sess["account_id"], sess.get("cid", ""), sess.get("rid", ""), sess.get("rcid", ""), sess.get("model", ""), sess.get("updated_at"))
)
imported["conv_sessions"] += 1
except Exception:
pass
for s in data.get("settings", []):
try:
db.execute("INSERT OR IGNORE INTO settings (key, value) VALUES (?, ?)", (s["key"], s["value"]))
imported["settings"] += 1
except Exception:
pass
db.commit()
auto_backup() # Update backup after import
return {"ok": True, "imported": imported}
@app.get("/api/admin/system/info")
async def system_info(request: Request, db: sqlite3.Connection = Depends(get_db_dep), _auth=Depends(require_auth)):
"""System information."""
import platform
conv_debug = {}
try:
from gemini_proxy.gemini import _conv_state, _reqid_counters
for key, val in _conv_state.items():
conv_debug[key] = {
"cid": val.get("cid", "N/A"),
"rid": val.get("rid", "N/A"),
"rcid": val.get("rcid", "N/A"),
"has_metadata": bool(val.get("metadata")),
}
reqid_debug = dict(_reqid_counters)
except Exception:
conv_debug = {"error": "module not available"}
reqid_debug = {}
return {
"python": platform.python_version(),
"fastapi": "0.x",
"uptime": time.strftime("%Y-%m-%d %H:%M:%S"),
"gemini_bl": CONFIG.get("gemini_bl", ""),
"data_dir": DATA_DIR,
"db_size": os.path.getsize(DB_PATH) if os.path.exists(DB_PATH) else 0,
"conv_state": conv_debug,
"reqid_counters": reqid_debug,
}
@app.get("/api/admin/debug/gemini-models")
async def debug_gemini_models(request: Request, db: sqlite3.Connection = Depends(get_db_dep), _auth=Depends(require_auth)):
"""Debug endpoint: fetch Gemini page and extract model hex IDs from the HTML."""
account = get_next_account()
if not account:
return JSONResponse({"error": "No active accounts"}, status_code=503)
cookie_str, sapisid = _parse_cookie(account["cookie"])
auth_user = CONFIG.get("auth_user")
try:
import ssl as _ssl
ctx = _ssl.create_default_context()
prefix = f"/u/{auth_user}" if auth_user else ""
url = f"https://gemini.google.com{prefix}/app"
headers = {
"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
"Cookie": cookie_str,
}
if sapisid:
from gemini_proxy.gemini import make_sapisidhash
headers["Authorization"] = make_sapisidhash(sapisid)
proxy = CONFIG.get("proxy")
req = urllib.request.Request(url, headers=headers)
if proxy:
opener = urllib.request.build_opener(
urllib.request.ProxyHandler({"http": proxy, "https": proxy}),
urllib.request.HTTPSHandler(context=ctx)
)
resp = opener.open(req, timeout=30)
else:
resp = urllib.request.urlopen(req, context=ctx, timeout=30)
html = resp.read().decode("utf-8", errors="replace")
# Extract hex-like model IDs (16-char hex strings)
hex_ids = list(set(re.findall(r'[0-9a-f]{16}', html)))
# Extract model names
model_names = list(set(re.findall(r'gemini[_\-\w\.]*', html, re.IGNORECASE)))
# Extract SNlM0e
xsrf_match = re.search(r'"SNlM0e":"([^"]+)"', html)
# Extract x-goog-ext patterns from JS
goog_ext = re.findall(r'x-goog-ext-\d+-jspb', html)
# Look for capacity_tail patterns
capacity = re.findall(r'capacity[_\-\w]*', html, re.IGNORECASE)
return {
"html_length": len(html),
"hex_ids": hex_ids[:30],
"model_names": model_names[:20],
"xsrf_token": xsrf_match.group(1)[:30] if xsrf_match else "NOT FOUND",
"goog_ext_headers": goog_ext[:10],
"capacity_patterns": capacity[:10],
}
except Exception as e:
return JSONResponse({"error": str(e)}, status_code=500)
@app.get("/api/admin/usage")
async def usage_log(request: Request, db: sqlite3.Connection = Depends(get_db_dep), _auth=Depends(require_auth)):
limit = min(int(request.query_params.get("limit", "50")), 500)
rows = db.execute("""
SELECT u.id, u.model, u.tokens_in, u.tokens_out, u.status_code, u.duration_ms, u.created_at,
a.name as account_name, k.name as key_name
FROM usage_log u
LEFT JOIN accounts a ON u.account_id = a.id
LEFT JOIN api_keys k ON u.api_key_id = k.id
ORDER BY u.id DESC LIMIT ?
""", (limit,)).fetchall()
return [dict(r) for r in rows]
# ── Conversation Memory API ──────────────────────────────────────────
@app.post("/v1/conversations")
async def v1_create_conversation(key_row=Depends(validate_api_key)):
"""Create a new conversation thread. Returns thread_id for use in /v1/chat/completions."""
thread_id = f"conv_{uuid.uuid4().hex[:16]}"
return {"id": thread_id, "object": "conversation", "created_at": bj_now_str()}
@app.get("/v1/conversations/{thread_id}")
async def v1_get_conversation(thread_id: str, key_row=Depends(validate_api_key)):
"""Get stored conversation history for a thread."""
history = _load_conversation_history(thread_id)
if history is None:
return _make_error(f"Conversation {thread_id} not found", "invalid_request_error", code="conversation_not_found", status_code=404)
return {"id": thread_id, "object": "conversation.history", "messages": history}
@app.delete("/v1/conversations/{thread_id}")
async def v1_delete_conversation(thread_id: str, key_row=Depends(validate_api_key)):
"""Delete all stored messages for a conversation thread."""
_clear_conversation_history(thread_id)
return {"id": thread_id, "deleted": True}
@app.get("/v1/models")
async def v1_models(key_row=Depends(validate_api_key)):
models = []
for name, info in MODELS.items():
models.append({
"id": name,
"object": "model",
"created": 1700000000,
"owned_by": "google",
"permission": [{
"id": f"modelperm-{uuid.uuid4().hex[:12]}",
"object": "model_permission",
"created": 1700000000,
"allow_create": True,
"allow_sampling": True,
"allow_logprobs": True,
"allow_search_indices": True,
"allow_view": True,
"allow_fine_tuning": False,
"organization": "*",
"group": None,
"is_blocking": False,
}],
"root": name,
"parent": None,
})
return {"object": "list", "data": models}
@app.post("/v1/images/generations")
async def v1_images(request: Request, db: sqlite3.Connection = Depends(get_db_dep), key_row=Depends(validate_api_key)):
"""OpenAI-compatible image generation endpoint."""
data = await request.json()
prompt = data.get("prompt", "")
model = data.get("model", "gemini-3.5-flash-image")
n = data.get("n", 1)
size = data.get("size", "1024x1024")
response_format = data.get("response_format", "url") # "url" or "b64_json"
if not prompt:
return JSONResponse({"error": {"message": "prompt is required", "type": "invalid_request_error"}}, status_code=400)
# Get an active account
account = get_next_account()
if not account:
return JSONResponse({"error": {"message": "No active accounts", "type": "server_error"}}, status_code=503)
# Build message for Gemini — just the image generation prompt
messages = [{"role": "user", "content": prompt}]
# Call Gemini
result = call_gemini_web(account, messages, model, "")
if "error" in result:
update_account_usage(account["id"], success=False, status_code=result.get("status_code", 500))
db.execute("""
INSERT INTO usage_log (account_id, api_key_id, model, status_code, created_at)
VALUES (?, ?, ?, ?, ?)
""", (account["id"], key_row["id"], model, result.get("status_code", 500), bj_now_str()))
db.execute("UPDATE api_keys SET total_requests = total_requests + 1 WHERE id=?", (key_row["id"],))
db.commit()
return JSONResponse({"error": {"message": result["error"], "type": "upstream_error"}}, status_code=result.get("status_code", 502))
# Extract images from the result content (markdown image format)
text_content = result.get("content", "")
image_urls = re.findall(r'!\[.*?\]\((https?://[^\s)]+)\)', text_content)
if not image_urls:
return {
"created": int(time.time()),
"data": [],
"revised_prompt": text_content[:500],
}
# Build OpenAI-format response
image_data = []
for url in image_urls[:n]:
if response_format == "b64_json":
try:
import base64
img_req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
img_resp = urllib.request.urlopen(img_req, timeout=30)
img_bytes = img_resp.read()
b64 = base64.b64encode(img_bytes).decode("ascii")
image_data.append({"b64_json": b64})
except Exception:
image_data.append({"url": url})
else:
image_data.append({"url": url})
# Log usage
update_account_usage(account["id"], success=True)
db.execute("""
INSERT INTO usage_log (account_id, api_key_id, model, tokens_in, tokens_out, status_code, duration_ms, created_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
""", (account["id"], key_row["id"], model,
estimate_tokens(prompt), estimate_tokens(text_content),
200, result.get("duration_ms", 0), bj_now_str()))
db.execute("UPDATE api_keys SET total_requests = total_requests + 1 WHERE id=?", (key_row["id"],))
db.commit()
return {
"created": int(time.time()),
"data": image_data,
}
# ── API Proxy Routes (OpenAI-compatible) — SSE streaming ────────────
@app.post("/v1/chat/completions")
async def v1_chat(request: Request, db: sqlite3.Connection = Depends(get_db_dep), key_row=Depends(validate_api_key)):
"""OpenAI-compatible chat completions endpoint.
Supports: messages, model, stream, tools, tool_choice, response_format,
stream_options, max_completion_tokens, max_tokens, stop, temperature.
Also supports optional 'conversation_id' for persistent memory across requests.
CRITICAL SSE DESIGN:
- StreamingResponse with a sync generator yields SSE chunks
- ALL DB operations happen in a BackgroundTask AFTER the stream completes
- DB writes NEVER block the SSE stream
"""
data = await request.json()
model = data.get("model", "gemini-3.5-flash")
messages = data.get("messages", [])
stream = data.get("stream", False)
tools = data.get("tools", None)
tool_choice = data.get("tool_choice", None)
response_format = data.get("response_format", None)
conversation_id = data.get("conversation_id", None)
stream_options = data.get("stream_options", None)
max_completion_tokens = data.get("max_completion_tokens", None)
max_tokens = data.get("max_tokens", None)
stop = data.get("stop", None)
# ── Handle tool_choice ──
if tool_choice == "none":
tools = None
tool_choice = None
if not messages:
return JSONResponse({"error": {"message": "messages required", "type": "invalid_request_error", "param": "messages", "code": "invalid_messages"}}, status_code=400)
# ── Conversation Memory ──
effective_messages = messages
if conversation_id:
stored = _load_conversation_history(conversation_id)
if stored:
if len(messages) <= 2:
effective_messages = stored + messages
effective_messages = effective_messages[-50:]
# Get account with session affinity
account = get_account_for_session(conversation_id)
if not account:
return _make_error("No active accounts available", "server_error", code="no_active_accounts", status_code=503)
# Load CID session state — reuse last active CID when no conversation_id provided
effective_thread_id = conversation_id # thread_id used for saving CID/conv data
if conversation_id:
conv_metadata = load_conv_session(conversation_id)
else:
# No conversation_id from client → reuse this account's most recent CID
# to minimize Gemini web UI clutter (stay in one conversation)
conv_metadata, _reuse_thread = load_last_active_cid(account["id"])
if _reuse_thread:
effective_thread_id = _reuse_thread # Save updates to the reused thread
if conv_metadata:
source = conversation_id or _reuse_thread or "reused"
sys.stderr.write(f"[CID] Continuing conversation: CID={conv_metadata[0]}, source={source}\n")
# ── Prepare common vars ──
chat_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
created_ts = int(time.time())
include_usage = stream_options and isinstance(stream_options, dict) and stream_options.get("include_usage", False)
# ── Streaming response ───────────────────────────────────────────
if stream:
# Shared state between generator and BackgroundTask
stream_state = {
"status": "pending",
"stream_done": None,
"stream_error": None,
"cid_403_to_log": False,
"content": None,
"tool_calls": None,
"finish_reason": "stop",
"conv_data": None,
"tokens_in": 0,
"tokens_out": 0,
"duration_ms": 0,
}
def sse_stream():
"""Sync generator that yields SSE chunks only — NO DB writes inside.
All results stored in stream_state for BackgroundTask to process."""
# First chunk: role
first_chunk = {
"id": chat_id,
"object": "chat.completion.chunk",
"created": created_ts,
"model": model,
"choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}],
}
yield f"data: {json.dumps(first_chunk)}\n\n"
# Phase 1: Try streaming with CID for conversation continuity
for item in call_gemini_web_stream(
account, effective_messages, model, "",
tools=tools, response_format=response_format,
tool_choice=tool_choice, stop=stop,
max_completion_tokens=max_completion_tokens or max_tokens,
conv_metadata=conv_metadata,
conversation_id=effective_thread_id,
new_conversation=False,
):
if item["type"] == "content":
chunk = {
"id": chat_id,
"object": "chat.completion.chunk",
"created": created_ts,
"model": model,
"choices": [{"index": 0, "delta": {"content": item["delta"]}, "finish_reason": None}],
}
yield f"data: {json.dumps(chunk)}\n\n"
elif item["type"] == "done":
stream_state["stream_done"] = item
break
elif item["type"] == "error":
stream_state["stream_error"] = item
break
# Phase 2: CID 403 recovery — if CID expired, retry as new conversation
if stream_state["stream_error"] and stream_state["stream_error"].get("cid_expired") and conv_metadata:
sys.stderr.write(f"[CID] Streaming CID expired (HTTP {stream_state['stream_error'].get('status_code')}), retrying as new conversation\n")
stream_state["cid_403_to_log"] = True
stream_state["stream_error"] = None
stream_state["stream_done"] = None
for item in call_gemini_web_stream(
account, effective_messages, model, "",
tools=tools, response_format=response_format,
tool_choice=tool_choice, stop=stop,
max_completion_tokens=max_completion_tokens or max_tokens,
conv_metadata=None,
conversation_id=effective_thread_id,
new_conversation=True,
):
if item["type"] == "content":
chunk = {
"id": chat_id,
"object": "chat.completion.chunk",
"created": created_ts,
"model": model,
"choices": [{"index": 0, "delta": {"content": item["delta"]}, "finish_reason": None}],
}
yield f"data: {json.dumps(chunk)}\n\n"
elif item["type"] == "done":
stream_state["stream_done"] = item
break
elif item["type"] == "error":
stream_state["stream_error"] = item
break
# Phase 3: Handle error (after possible retry)
if stream_state["stream_error"]:
sys.stderr.write(f"[stream] Gemini stream error: {stream_state['stream_error'].get('error', 'unknown')}\n")
err_text = f"\n[Error: {stream_state['stream_error'].get('error', 'upstream error')}]"
err_chunk = {
"id": chat_id,
"object": "chat.completion.chunk",
"created": created_ts,
"model": model,
"choices": [{"index": 0, "delta": {"content": err_text}, "finish_reason": "error"}],
}
yield f"data: {json.dumps(err_chunk)}\n\n"
yield "data: [DONE]\n\n"
stream_state["status"] = "error"
return
stream_done = stream_state["stream_done"]
if not stream_done:
# No result at all — emit empty finish
final_chunk = {
"id": chat_id,
"object": "chat.completion.chunk",
"created": created_ts,
"model": model,
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
}
yield f"data: {json.dumps(final_chunk)}\n\n"
yield "data: [DONE]\n\n"
stream_state["status"] = "success"
return
# ── Stream completed successfully — extract results ──
content = stream_done.get("content")
tool_calls = stream_done.get("tool_calls")
finish_reason = stream_done.get("finish_reason", "stop")
conv_data = stream_done.get("conv_data")
# ── Stream tool_calls in incremental SSE format ──
# IMPORTANT: yield all SSE data BEFORE any DB writes
# OpenAI protocol: first chunk per tool_call has index+id+type+name
# Subsequent chunks stream arguments in ~5-char increments
if tool_calls:
for tc_idx, tc in enumerate(tool_calls):
tc_id = tc["id"]
tc_name = tc["function"]["name"]
tc_args = tc["function"]["arguments"]
# Chunk 1: name (full) + id + type, arguments=""
name_chunk = {
"id": chat_id,
"object": "chat.completion.chunk",
"created": created_ts,
"model": model,
"choices": [{
"index": 0,
"delta": {
"tool_calls": [{
"index": tc_idx,
"id": tc_id,
"type": "function",
"function": {
"name": tc_name,
"arguments": "",
}
}]
},
"finish_reason": None,
}],
}
yield f"data: {json.dumps(name_chunk)}\n\n"
# Chunks 2+: arguments in ~5-char increments
if tc_args:
arg_chunk_size = 5
for i in range(0, len(tc_args), arg_chunk_size):
arg_piece = tc_args[i:i+arg_chunk_size]
arg_chunk = {
"id": chat_id,
"object": "chat.completion.chunk",
"created": created_ts,
"model": model,
"choices": [{
"index": 0,
"delta": {
"tool_calls": [{
"index": tc_idx,
"function": {
"arguments": arg_piece,
}
}]
},
"finish_reason": None,
}],
}
yield f"data: {json.dumps(arg_chunk)}\n\n"
# Final chunk: finish_reason=tool_calls
final_chunk = {
"id": chat_id,
"object": "chat.completion.chunk",
"created": created_ts,
"model": model,
"choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}],
}
yield f"data: {json.dumps(final_chunk)}\n\n"
elif not content:
# No content and no tool_calls — still need finish_reason
final_chunk = {
"id": chat_id,
"object": "chat.completion.chunk",
"created": created_ts,
"model": model,
"choices": [{"index": 0, "delta": {}, "finish_reason": finish_reason}],
}
yield f"data: {json.dumps(final_chunk)}\n\n"
else:
# Content was streamed, final chunk with finish_reason
final_chunk = {
"id": chat_id,
"object": "chat.completion.chunk",
"created": created_ts,
"model": model,
"choices": [{"index": 0, "delta": {}, "finish_reason": finish_reason}],
}
yield f"data: {json.dumps(final_chunk)}\n\n"
# stream_options include_usage: send usage data in final chunk
if include_usage:
usage_chunk = {
"id": chat_id,
"object": "chat.completion.chunk",
"created": created_ts,
"model": model,
"choices": [],
"usage": {
"prompt_tokens": stream_done.get("tokens_in", 0),
"completion_tokens": stream_done.get("tokens_out", 0),
"total_tokens": stream_done.get("tokens_in", 0) + stream_done.get("tokens_out", 0),
},
}
yield f"data: {json.dumps(usage_chunk)}\n\n"
yield "data: [DONE]\n\n"
# Store results for BackgroundTask — NO DB writes in generator
stream_state["status"] = "success"
stream_state["content"] = content
stream_state["tool_calls"] = tool_calls
stream_state["finish_reason"] = finish_reason
stream_state["conv_data"] = conv_data
stream_state["tokens_in"] = stream_done.get("tokens_in", 0)
stream_state["tokens_out"] = stream_done.get("tokens_out", 0)
stream_state["duration_ms"] = stream_done.get("duration_ms", 0)
def background_db_write():
"""All DB operations AFTER SSE stream completes.
Creates standalone connection — independent of request-scoped connection."""
bg_db = sqlite3.connect(DB_PATH, check_same_thread=False)
bg_db.row_factory = sqlite3.Row
bg_db.execute("PRAGMA journal_mode=WAL")
_db_context.set(bg_db)
try:
if stream_state["status"] == "error":
# ── Error case: log failure ──
status_code = stream_state["stream_error"].get("status_code", 502)
bg_db.execute("""
INSERT INTO usage_log (account_id, api_key_id, model, status_code, created_at)
VALUES (?, ?, ?, ?, ?)
""", (account["id"], key_row["id"], model, status_code, bj_now_str()))
bg_db.execute("UPDATE api_keys SET total_requests = total_requests + 1 WHERE id=?", (key_row["id"],))
bg_db.commit()
update_account_usage(account["id"], success=False, status_code=status_code)
elif stream_state["status"] == "success":
# ── CID 403 deferred: clear session + log the failed attempt ──
if stream_state["cid_403_to_log"]:
try:
clear_conv_session(effective_thread_id)
bg_db.execute("""
INSERT INTO usage_log (account_id, api_key_id, model, status_code, duration_ms, created_at)
VALUES (?, ?, ?, ?, ?, ?)
""", (account["id"], key_row["id"], model, 403, 0, bj_now_str()))
bg_db.commit()
except Exception:
pass
# ── Log usage ──
try:
bg_db.execute("""
INSERT INTO usage_log (account_id, api_key_id, model, tokens_in, tokens_out, duration_ms, created_at)
VALUES (?, ?, ?, ?, ?, ?, ?)
""", (account["id"], key_row["id"], model,
stream_state["tokens_in"], stream_state["tokens_out"],
stream_state["duration_ms"], bj_now_str()))
bg_db.execute("UPDATE api_keys SET total_requests = total_requests + 1 WHERE id=?", (key_row["id"],))
bg_db.commit()
update_account_usage(account["id"], success=True)
except Exception:
pass
# ── Save to conversation memory ──
if effective_thread_id:
try:
for msg in messages:
if msg.get("role") in ("user", "system", "tool"):
_save_conversation_message(
effective_thread_id, key_row["id"],
role=msg.get("role"),
content=_extract_text(msg.get("content", "")),
tool_call_id=msg.get("tool_call_id", ""),
tool_name=msg.get("name", ""),
model=model,
)
_save_conversation_message(
effective_thread_id, key_row["id"],
role="assistant",
content=stream_state["content"] or "",
tool_calls=stream_state["tool_calls"],
model=model,
)
except Exception:
pass
# ── Persist CID session state ──
if effective_thread_id and stream_state.get("conv_data"):
try:
save_conv_session(effective_thread_id, account["id"], stream_state["conv_data"], model)
except Exception:
pass
finally:
_db_context.set(None)
bg_db.close()
return StreamingResponse(
sse_stream(),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no",
},
background=BackgroundTask(background_db_write),
)
# ── Non-streaming response ────────────────────────────────────────
# Call Gemini with all parameters (including CID for conversation continuity)
result = call_gemini_web(account, effective_messages, model, "",
tools=tools, response_format=response_format,
tool_choice=tool_choice, stop=stop,
max_completion_tokens=max_completion_tokens or max_tokens,
conv_metadata=conv_metadata,
conversation_id=effective_thread_id)
if "error" in result:
# ── CID 403 recovery: if CID failed, retry as new conversation ──
if conv_metadata and (result.get("status_code") in (401, 403) or result.get("cid_expired")):
reason = "cid_expired" if result.get("cid_expired") else f"HTTP {result['status_code']}"
sys.stderr.write(f"[CID] CID failed ({reason}), clearing session and retrying as new conversation\n")
clear_conv_session(effective_thread_id)
# Retry without CID — send full history
retry_messages = effective_messages
if effective_thread_id:
stored = _load_conversation_history(effective_thread_id)
if stored and len(messages) <= 2:
retry_messages = (stored + messages)[-50:]
else:
retry_messages = messages
retry_result = call_gemini_web(account, retry_messages, model, "",
tools=tools, response_format=response_format,
tool_choice=tool_choice, stop=stop,
max_completion_tokens=max_completion_tokens or max_tokens,
conv_metadata=None,
conversation_id=effective_thread_id,
new_conversation=True)
if "error" not in retry_result:
result = retry_result
sys.stderr.write(f"[CID] Retry as new conversation succeeded\n")
# Log the failed CID attempt
db.execute("""
INSERT INTO usage_log (account_id, api_key_id, model, status_code, duration_ms, created_at)
VALUES (?, ?, ?, ?, ?, ?)
""", (account["id"], key_row["id"], model, 403, 0, bj_now_str()))
db.commit()
sys.stderr.write(f"[CID] Logged failed CID attempt (403) to usage stats\n")
else:
result = retry_result
# If still an error
if "error" in result:
update_account_usage(account["id"], success=False, status_code=result.get("status_code", 500))
db.execute("""
INSERT INTO usage_log (account_id, api_key_id, model, status_code, created_at)
VALUES (?, ?, ?, ?, ?)
""", (account["id"], key_row["id"], model, result.get("status_code", 500), bj_now_str()))
db.execute("UPDATE api_keys SET total_requests = total_requests + 1 WHERE id=?", (key_row["id"],))
db.commit()
err_msg = result["error"]
err_type = "upstream_error"
err_code = "gemini_error"
status_code = result.get("status_code", 502)
if status_code in (401, 403):
err_type = "authentication_error"
err_code = "invalid_auth"
elif status_code == 429:
err_type = "rate_limit_error"
err_code = "rate_limit_exceeded"
elif status_code >= 500:
err_type = "server_error"
err_code = "gemini_server_error"
elif status_code >= 400:
err_type = "invalid_request_error"
err_code = "bad_request"
return _make_error(err_msg, err_type, code=err_code, status_code=status_code)
# Log usage
db.execute("""
INSERT INTO usage_log (account_id, api_key_id, model, tokens_in, tokens_out, duration_ms, created_at)
VALUES (?, ?, ?, ?, ?, ?, ?)
""", (account["id"], key_row["id"], model, result.get("tokens_in", 0), result.get("tokens_out", 0), result.get("duration_ms", 0), bj_now_str()))
db.execute("UPDATE api_keys SET total_requests = total_requests + 1 WHERE id=?", (key_row["id"],))
db.commit()
update_account_usage(account["id"], success=True)
# ── Save to conversation memory ──
if effective_thread_id:
for msg in messages:
if msg.get("role") in ("user", "system", "tool"):
_save_conversation_message(
effective_thread_id, key_row["id"],
role=msg.get("role"),
content=_extract_text(msg.get("content", "")),
tool_call_id=msg.get("tool_call_id", ""),
tool_name=msg.get("name", ""),
model=model,
)
_save_conversation_message(
effective_thread_id, key_row["id"],
role="assistant",
content=result.get("content") or "",
tool_calls=result.get("tool_calls"),
model=model,
)
# ── Persist CID session state ──
if effective_thread_id and result.get("conv_data"):
save_conv_session(effective_thread_id, account["id"], result["conv_data"], model)
message_obj = {"role": "assistant"}
content_ns = result.get("content")
tool_calls_ns = result.get("tool_calls")
finish_reason_ns = result.get("finish_reason", "stop")
if content_ns:
message_obj["content"] = content_ns
elif tool_calls_ns:
message_obj["content"] = None
else:
message_obj["content"] = content_ns or None
if tool_calls_ns:
message_obj["tool_calls"] = tool_calls_ns
return {
"id": chat_id,
"object": "chat.completion",
"created": created_ts,
"model": model,
"choices": [{
"index": 0,
"message": message_obj,
"finish_reason": finish_reason_ns,
}],
"usage": {
"prompt_tokens": result.get("tokens_in", 0),
"completion_tokens": result.get("tokens_out", 0),
"total_tokens": result.get("tokens_in", 0) + result.get("tokens_out", 0),
},
}
# ── Static files (must be mounted AFTER all routes) ─────────────────
app.mount("/static", StaticFiles(directory=os.path.join(os.path.dirname(__file__), "static")), name="static")
# ── Main ────────────────────────────────────────────────────────────
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=PORT)