Spaces:
Runtime error
Runtime error
| """ | |
| 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") | |
| else: | |
| img_lines.append(f"\n (参考图片)") | |
| 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 ──────────────────────────────── | |
| 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 | |
| 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 | |
| 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 | |
| ) | |
| 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 ──────────────────────────────────────────────────── | |
| 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 | |
| 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) | |
| 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) | |
| async def logout(request: Request): | |
| request.session.clear() | |
| return RedirectResponse(url="/login") | |
| 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 | |
| async def index(): | |
| return RedirectResponse(url="/admin") | |
| async def ping(): | |
| """Simple health check - no auth required.""" | |
| return {"status": "ok", "ts": bj_now().isoformat()} | |
| 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), | |
| } | |
| 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>") | |
| 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>") | |
| 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 ──────────────────────────────────────────────── | |
| 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, | |
| } | |
| 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] | |
| 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} | |
| 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} | |
| 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} | |
| 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} | |
| 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] | |
| 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} | |
| 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} | |
| 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} | |
| 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)} | |
| 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} | |
| 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} | |
| 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. 自动完成!账号已添加到管理后台"} | |
| 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) | |
| 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, | |
| } | |
| 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} | |
| 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} | |
| 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": "密码已更新(重启后恢复为环境变量值)"} | |
| 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} | |
| 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, | |
| } | |
| 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} | |
| 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, | |
| } | |
| 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) | |
| 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 ────────────────────────────────────────── | |
| 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()} | |
| 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} | |
| 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} | |
| 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} | |
| 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 ──────────── | |
| 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) | |