780 lines
43 KiB
Python
780 lines
43 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
import hashlib
|
|
import hmac
|
|
import io
|
|
import json
|
|
import logging
|
|
import os
|
|
import time
|
|
import uuid
|
|
from contextlib import asynccontextmanager
|
|
from typing import Any
|
|
|
|
import aiosqlite
|
|
import httpx
|
|
from fastapi import FastAPI, File, Form, Header, HTTPException, Query, Request, UploadFile
|
|
from fastapi.responses import HTMLResponse
|
|
from PIL import Image, ImageOps, UnidentifiedImageError
|
|
|
|
LOG = logging.getLogger("phproxy")
|
|
DB_PATH = os.getenv("DB_PATH", "./data/proxy.db")
|
|
ADMIN_SECRET = os.getenv("ADMIN_SECRET", "")
|
|
MAX_UPLOAD_BYTES = int(os.getenv("MAX_UPLOAD_BYTES", str(50 * 1024 * 1024)))
|
|
MAX_IMAGES = int(os.getenv("MAX_IMAGES", "8"))
|
|
MAX_IMAGE_PIXELS = int(os.getenv("MAX_IMAGE_PIXELS", "25000000"))
|
|
HTTP_TIMEOUT_SECONDS = float(os.getenv("HTTP_TIMEOUT_SECONDS", "300"))
|
|
CONNECT_TIMEOUT_SECONDS = float(os.getenv("CONNECT_TIMEOUT_SECONDS", "15"))
|
|
MODEL_MAP = {
|
|
"nano-banana": {
|
|
"google": os.getenv("GOOGLE_MODEL_NANO_BANANA", "gemini-2.5-flash-image"),
|
|
"openrouter": os.getenv("OPENROUTER_MODEL_NANO_BANANA", "google/gemini-2.5-flash-image"),
|
|
},
|
|
}
|
|
IMAGE_FORMATS = {"JPEG": "image/jpeg", "PNG": "image/png", "WEBP": "image/webp"}
|
|
IMAGE_FORMAT_TYPES = {mime: image_format for image_format, mime in IMAGE_FORMATS.items()}
|
|
image_validation_semaphore = asyncio.Semaphore(1)
|
|
http_client: httpx.AsyncClient | None = None
|
|
|
|
|
|
def secret_hash(value: str) -> str:
|
|
return hashlib.sha256(value.encode("utf-8")).hexdigest()
|
|
|
|
|
|
def secret_matches(value: str | None, stored_hash: str | None) -> bool:
|
|
return bool(value and stored_hash and hmac.compare_digest(secret_hash(value), stored_hash))
|
|
|
|
|
|
@asynccontextmanager
|
|
async def db():
|
|
conn = await aiosqlite.connect(DB_PATH, timeout=10)
|
|
conn.row_factory = aiosqlite.Row
|
|
await conn.execute("PRAGMA foreign_keys=ON")
|
|
await conn.execute("PRAGMA busy_timeout=10000")
|
|
try:
|
|
yield conn
|
|
finally:
|
|
await conn.close()
|
|
|
|
|
|
async def init_db() -> None:
|
|
os.makedirs(os.path.dirname(os.path.abspath(DB_PATH)), exist_ok=True)
|
|
async with db() as conn:
|
|
await conn.execute("PRAGMA journal_mode=WAL")
|
|
await conn.executescript(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS clients (
|
|
client_id TEXT PRIMARY KEY,
|
|
name TEXT NOT NULL,
|
|
secret_hash TEXT,
|
|
is_active INTEGER NOT NULL DEFAULT 1 CHECK (is_active IN (0, 1)),
|
|
created_at REAL NOT NULL
|
|
);
|
|
CREATE TABLE IF NOT EXISTS provider_keys (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
client_id TEXT NOT NULL REFERENCES clients(client_id) ON DELETE CASCADE,
|
|
provider TEXT NOT NULL CHECK (provider IN ('google', 'openrouter')),
|
|
api_key TEXT NOT NULL,
|
|
priority INTEGER NOT NULL DEFAULT 1,
|
|
is_active INTEGER NOT NULL DEFAULT 1 CHECK (is_active IN (0, 1)),
|
|
fail_count INTEGER NOT NULL DEFAULT 0,
|
|
last_error_at REAL,
|
|
created_at REAL NOT NULL
|
|
);
|
|
CREATE INDEX IF NOT EXISTS provider_keys_route_idx
|
|
ON provider_keys(client_id, is_active, priority, id);
|
|
CREATE TABLE IF NOT EXISTS jobs (
|
|
job_id TEXT PRIMARY KEY,
|
|
client_id TEXT NOT NULL REFERENCES clients(client_id) ON DELETE CASCADE,
|
|
status TEXT NOT NULL CHECK (status IN ('queued', 'processing', 'completed', 'failed')),
|
|
request_json TEXT NOT NULL,
|
|
result_json TEXT,
|
|
error_json TEXT,
|
|
created_at REAL NOT NULL,
|
|
updated_at REAL NOT NULL
|
|
);
|
|
CREATE INDEX IF NOT EXISTS jobs_client_created_idx ON jobs(client_id, created_at);
|
|
CREATE TABLE IF NOT EXISTS generation_runs (
|
|
generation_id TEXT PRIMARY KEY,
|
|
client_id TEXT NOT NULL REFERENCES clients(client_id) ON DELETE CASCADE,
|
|
job_id TEXT,
|
|
status TEXT NOT NULL CHECK (status IN ('queued', 'processing', 'completed', 'failed')),
|
|
requested_provider TEXT,
|
|
provider_used TEXT,
|
|
requested_model TEXT NOT NULL,
|
|
model_used TEXT,
|
|
image_count INTEGER NOT NULL DEFAULT 0,
|
|
upload_bytes INTEGER NOT NULL DEFAULT 0,
|
|
request_intake_sec REAL,
|
|
upstream_sec REAL,
|
|
total_sec REAL,
|
|
failover_attempts INTEGER NOT NULL DEFAULT 0,
|
|
cost_usd REAL,
|
|
cost_source TEXT,
|
|
usage_json TEXT,
|
|
error_json TEXT,
|
|
created_at REAL NOT NULL,
|
|
started_at REAL,
|
|
completed_at REAL
|
|
);
|
|
CREATE INDEX IF NOT EXISTS generation_runs_created_idx
|
|
ON generation_runs(created_at DESC);
|
|
CREATE INDEX IF NOT EXISTS generation_runs_client_created_idx
|
|
ON generation_runs(client_id, created_at DESC);
|
|
"""
|
|
)
|
|
await conn.execute(
|
|
"UPDATE jobs SET status='failed', error_json=?, updated_at=? WHERE status IN ('queued','processing')",
|
|
(json.dumps({"message": "Proxy restarted while this job was running"}), time.time()),
|
|
)
|
|
now = time.time()
|
|
await conn.execute(
|
|
"UPDATE generation_runs SET status='failed',error_json=?,completed_at=?,"
|
|
"total_sec=MAX(0,?-created_at) WHERE status IN ('queued','processing')",
|
|
(json.dumps({"message": "Proxy restarted while this generation was running"}), now, now),
|
|
)
|
|
# Preserve basic timing from jobs created before execution history was added.
|
|
cur = await conn.execute(
|
|
"SELECT j.job_id,j.client_id,j.status,j.request_json,j.result_json,j.error_json,"
|
|
"j.created_at,j.updated_at FROM jobs j LEFT JOIN generation_runs r "
|
|
"ON r.generation_id=j.job_id WHERE r.generation_id IS NULL"
|
|
)
|
|
for row in await cur.fetchall():
|
|
try:
|
|
request_data = json.loads(row["request_json"])
|
|
result = json.loads(row["result_json"]) if row["result_json"] else {}
|
|
images = request_data.get("images", [])
|
|
upload_bytes = sum((len(image.get("b64", "")) * 3) // 4 for image in images)
|
|
await conn.execute(
|
|
"INSERT OR IGNORE INTO generation_runs(generation_id,client_id,job_id,status,"
|
|
"requested_provider,provider_used,requested_model,model_used,image_count,upload_bytes,"
|
|
"upstream_sec,total_sec,failover_attempts,cost_usd,cost_source,error_json,created_at,"
|
|
"started_at,completed_at) VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
|
|
(row["job_id"], row["client_id"], row["job_id"], row["status"],
|
|
request_data.get("preferred"), result.get("provider_used"),
|
|
request_data.get("model", "unknown"), result.get("model_used"), len(images), upload_bytes,
|
|
result.get("elapsed_sec"), max(0, row["updated_at"] - row["created_at"]),
|
|
result.get("failover_attempts", 0), result.get("cost_usd"), result.get("cost_source"),
|
|
row["error_json"], row["created_at"], row["created_at"], row["updated_at"]),
|
|
)
|
|
except (TypeError, ValueError, json.JSONDecodeError):
|
|
LOG.exception("Could not backfill generation history for job %s", row["job_id"])
|
|
await conn.commit()
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(_: FastAPI):
|
|
global http_client
|
|
if not ADMIN_SECRET:
|
|
raise RuntimeError("ADMIN_SECRET must be set before starting the service")
|
|
await init_db()
|
|
http_client = httpx.AsyncClient(
|
|
timeout=httpx.Timeout(HTTP_TIMEOUT_SECONDS, connect=CONNECT_TIMEOUT_SECONDS),
|
|
limits=httpx.Limits(max_keepalive_connections=50, max_connections=200),
|
|
follow_redirects=False,
|
|
)
|
|
yield
|
|
await http_client.aclose()
|
|
|
|
|
|
app = FastAPI(title="Generative AI Proxy", version="0.1.0", lifespan=lifespan)
|
|
|
|
|
|
@app.middleware("http")
|
|
async def content_length_limit(request: Request, call_next):
|
|
request.state.received_monotonic = time.monotonic()
|
|
request.state.received_at = time.time()
|
|
content_length = request.headers.get("content-length")
|
|
if content_length:
|
|
try:
|
|
if int(content_length) > MAX_UPLOAD_BYTES + 1024 * 1024:
|
|
return JSONResponse(status_code=413, content={"detail": "Request body too large"})
|
|
except ValueError:
|
|
return JSONResponse(status_code=400, content={"detail": "Invalid Content-Length"})
|
|
return await call_next(request)
|
|
|
|
|
|
from fastapi.responses import JSONResponse # noqa: E402
|
|
|
|
|
|
def resolve_model(model: str, provider: str) -> str:
|
|
if model in MODEL_MAP:
|
|
return MODEL_MAP[model][provider]
|
|
if provider == "google" and model.startswith("google/"):
|
|
return model.removeprefix("google/")
|
|
if provider == "openrouter" and "/" not in model:
|
|
return f"google/{model}"
|
|
return model
|
|
|
|
|
|
def auth_admin(value: str | None) -> None:
|
|
if not value or not hmac.compare_digest(value, ADMIN_SECRET):
|
|
raise HTTPException(status_code=401, detail="Unauthorized")
|
|
|
|
|
|
async def authorize_client(client_id: str, secret: str | None) -> None:
|
|
async with db() as conn:
|
|
cur = await conn.execute(
|
|
"SELECT secret_hash, is_active FROM clients WHERE client_id=?", (client_id,)
|
|
)
|
|
row = await cur.fetchone()
|
|
if not row or not row["is_active"]:
|
|
raise HTTPException(status_code=403, detail="Invalid or inactive client")
|
|
if not row["secret_hash"]:
|
|
raise HTTPException(status_code=403, detail="Client secret must be configured")
|
|
if not secret_matches(secret, row["secret_hash"]):
|
|
raise HTTPException(status_code=401, detail="Invalid client secret")
|
|
|
|
|
|
def sanitize_image(raw: bytes, declared_mime: str) -> tuple[bytes, str]:
|
|
"""Decode and re-encode accepted raster formats, stripping appended payloads/metadata."""
|
|
Image.MAX_IMAGE_PIXELS = MAX_IMAGE_PIXELS
|
|
try:
|
|
with Image.open(io.BytesIO(raw)) as source:
|
|
image_format = source.format
|
|
actual_mime = IMAGE_FORMATS.get(image_format or "")
|
|
if actual_mime is None or actual_mime != declared_mime:
|
|
raise HTTPException(status_code=415, detail="File content does not match an allowed image type")
|
|
width, height = source.size
|
|
if width <= 0 or height <= 0 or width * height > MAX_IMAGE_PIXELS:
|
|
raise HTTPException(status_code=413, detail="Image dimensions exceed configured limits")
|
|
source.load()
|
|
image = ImageOps.exif_transpose(source)
|
|
output = io.BytesIO()
|
|
if image_format == "JPEG":
|
|
if image.mode not in ("L", "RGB"):
|
|
image = image.convert("RGB")
|
|
image.save(output, format="JPEG", quality=95)
|
|
elif image_format == "PNG":
|
|
if image.mode not in ("1", "L", "LA", "P", "RGB", "RGBA"):
|
|
image = image.convert("RGBA" if "A" in image.getbands() else "RGB")
|
|
image.save(output, format="PNG")
|
|
else:
|
|
if image.mode not in ("RGB", "RGBA"):
|
|
image = image.convert("RGBA" if "A" in image.getbands() else "RGB")
|
|
image.save(output, format="WEBP", quality=95)
|
|
return output.getvalue(), actual_mime
|
|
except HTTPException:
|
|
raise
|
|
except (Image.DecompressionBombError, Image.DecompressionBombWarning) as exc:
|
|
raise HTTPException(status_code=413, detail="Image dimensions exceed configured limits") from exc
|
|
except (UnidentifiedImageError, OSError, ValueError) as exc:
|
|
raise HTTPException(status_code=415, detail="Uploaded file is not a valid supported image") from exc
|
|
|
|
|
|
async def read_images(images: list[UploadFile]) -> tuple[list[dict[str, str]], int]:
|
|
if len(images) > MAX_IMAGES:
|
|
raise HTTPException(status_code=413, detail=f"At most {MAX_IMAGES} images are allowed")
|
|
total = 0
|
|
encoded: list[dict[str, str]] = []
|
|
for image in images:
|
|
mime = (image.content_type or "").lower()
|
|
if mime not in IMAGE_FORMAT_TYPES:
|
|
raise HTTPException(status_code=415, detail=f"Unsupported image type: {mime or 'unknown'}")
|
|
chunks: list[bytes] = []
|
|
while chunk := await image.read(1024 * 1024):
|
|
total += len(chunk)
|
|
if total > MAX_UPLOAD_BYTES:
|
|
raise HTTPException(status_code=413, detail="Uploaded images exceed configured limit")
|
|
chunks.append(chunk)
|
|
raw = b"".join(chunks)
|
|
if not raw:
|
|
raise HTTPException(status_code=415, detail="Uploaded image is empty")
|
|
async with image_validation_semaphore:
|
|
sanitized, actual_mime = await asyncio.to_thread(sanitize_image, raw, mime)
|
|
if total - len(raw) + len(sanitized) > MAX_UPLOAD_BYTES:
|
|
raise HTTPException(status_code=413, detail="Validated images exceed configured upload limit")
|
|
total += len(sanitized) - len(raw)
|
|
encoded.append({"mime_type": actual_mime, "b64": base64.b64encode(sanitized).decode("ascii")})
|
|
await image.close()
|
|
return encoded, total
|
|
|
|
|
|
class UpstreamError(Exception):
|
|
def __init__(self, provider: str, status_code: int | None, retryable: bool):
|
|
self.provider = provider
|
|
self.status_code = status_code
|
|
self.retryable = retryable
|
|
|
|
|
|
async def post_upstream(provider: str, api_key: str, model: str, prompt: str,
|
|
images: list[dict[str, str]], params: dict[str, Any]) -> dict[str, Any]:
|
|
assert http_client is not None
|
|
resolved = resolve_model(model, provider)
|
|
try:
|
|
if provider == "google":
|
|
parts: list[dict[str, Any]] = [
|
|
{"inlineData": {"mimeType": image["mime_type"], "data": image["b64"]}}
|
|
for image in images
|
|
]
|
|
if prompt:
|
|
parts.append({"text": prompt})
|
|
config = dict(params.get("generationConfig", {}))
|
|
config.setdefault("responseModalities", params.get("responseModalities", ["TEXT", "IMAGE"]))
|
|
for source, target in (("temperature", "temperature"), ("topP", "topP"), ("seed", "seed")):
|
|
if source in params:
|
|
config[target] = params[source]
|
|
image_config = {}
|
|
if "aspect_ratio" in params:
|
|
image_config["aspectRatio"] = params["aspect_ratio"]
|
|
if "image_size" in params:
|
|
image_config["imageSize"] = params["image_size"]
|
|
if image_config:
|
|
config["imageConfig"] = image_config
|
|
if isinstance(params.get("thinking"), dict):
|
|
thinking = params["thinking"]
|
|
config["thinkingConfig"] = {
|
|
k: v for source, k in (("budget_tokens", "thinkingBudget"),
|
|
("include_thoughts", "includeThoughts"))
|
|
if (v := thinking.get(source)) is not None
|
|
}
|
|
payload: dict[str, Any] = {"contents": [{"role": "user", "parts": parts}],
|
|
"generationConfig": config}
|
|
if "safetySettings" in params:
|
|
payload["safetySettings"] = params["safetySettings"]
|
|
response = await http_client.post(
|
|
f"https://generativelanguage.googleapis.com/v1beta/models/{resolved}:generateContent",
|
|
params={"key": api_key}, json=payload,
|
|
)
|
|
else:
|
|
blocks: list[dict[str, Any]] = []
|
|
if prompt:
|
|
blocks.append({"type": "text", "text": prompt})
|
|
for image in images:
|
|
blocks.append({"type": "image_url", "image_url": {
|
|
"url": f"data:{image['mime_type']};base64,{image['b64']}"}})
|
|
payload = {"model": resolved, "messages": [{"role": "user", "content": blocks}],
|
|
"modalities": params.get("modalities", ["image", "text"])}
|
|
for key in ("temperature", "seed"):
|
|
if key in params:
|
|
payload[key] = params[key]
|
|
reasoning = params.get("reasoning")
|
|
if isinstance(reasoning, dict):
|
|
# OpenRouter expects an object here; boolean flags are invalid values.
|
|
payload["reasoning"] = reasoning
|
|
image_config = {k: params[source] for source, k in
|
|
(("aspect_ratio", "aspect_ratio"), ("image_size", "image_size"))
|
|
if source in params}
|
|
if image_config:
|
|
payload["image_config"] = image_config
|
|
if isinstance(params.get("extra_body"), dict):
|
|
payload.update(params["extra_body"])
|
|
# OpenRouter returns its billed amount in usage.cost when usage is requested.
|
|
payload.setdefault("usage", {"include": True})
|
|
response = await http_client.post(
|
|
"https://openrouter.ai/api/v1/chat/completions",
|
|
headers={"Authorization": f"Bearer {api_key}",
|
|
"X-OpenRouter-Metadata": "enabled"}, json=payload,
|
|
)
|
|
except httpx.TimeoutException as exc:
|
|
raise UpstreamError(provider, None, True) from exc
|
|
except httpx.RequestError as exc:
|
|
raise UpstreamError(provider, None, True) from exc
|
|
if response.status_code != 200:
|
|
raise UpstreamError(provider, response.status_code,
|
|
response.status_code == 429 or response.status_code >= 500)
|
|
try:
|
|
data = response.json()
|
|
result_images: list[dict[str, str]] = []
|
|
texts: list[str] = []
|
|
if provider == "google":
|
|
for candidate in data.get("candidates", []):
|
|
for part in candidate.get("content", {}).get("parts", []):
|
|
if "inlineData" in part:
|
|
inline = part["inlineData"]
|
|
result_images.append({"mime_type": inline.get("mimeType", "image/png"),
|
|
"b64_json": inline["data"]})
|
|
elif part.get("text") and not part.get("thought"):
|
|
texts.append(part["text"])
|
|
else:
|
|
message = data.get("choices", [{}])[0].get("message", {})
|
|
content = message.get("content") or ""
|
|
texts.append(content if isinstance(content, str) else "")
|
|
for item in message.get("images", []):
|
|
url = item.get("image_url", {}).get("url", "")
|
|
if url.startswith("data:"):
|
|
header, b64_data = url.split(",", 1)
|
|
result_images.append({"mime_type": header[5:].split(";")[0],
|
|
"b64_json": b64_data})
|
|
elif url:
|
|
result_images.append({"mime_type": "image/png", "url": url})
|
|
text = "\n".join(texts)
|
|
if not result_images and not text:
|
|
raise ValueError("empty content")
|
|
if provider == "openrouter":
|
|
usage = data.get("usage")
|
|
raw_cost = usage.get("cost") if isinstance(usage, dict) else None
|
|
cost_usd = (float(raw_cost) if isinstance(raw_cost, (int, float))
|
|
and not isinstance(raw_cost, bool) else None)
|
|
cost_source = "openrouter_usage" if cost_usd is not None else "unavailable"
|
|
upstream_request_id = data.get("id")
|
|
else:
|
|
usage = data.get("usageMetadata")
|
|
cost_usd = None
|
|
cost_source = "unavailable"
|
|
upstream_request_id = None
|
|
result = {"provider_used": provider, "model_used": data.get("model", resolved),
|
|
"images": result_images, "text": text,
|
|
"cost_usd": cost_usd, "cost_source": cost_source}
|
|
if isinstance(usage, dict):
|
|
result["usage"] = usage
|
|
if upstream_request_id:
|
|
result["upstream_request_id"] = upstream_request_id
|
|
if provider == "openrouter" and isinstance(data.get("openrouter_metadata"), dict):
|
|
result["upstream_metadata"] = data["openrouter_metadata"]
|
|
return result
|
|
except (ValueError, KeyError, IndexError, TypeError) as exc:
|
|
raise UpstreamError(provider, 502, False) from exc
|
|
|
|
|
|
async def execute_with_failover(client_id: str, preferred: str | None, model: str,
|
|
prompt: str, images: list[dict[str, str]],
|
|
params: dict[str, Any]) -> dict[str, Any]:
|
|
async with db() as conn:
|
|
cur = await conn.execute("SELECT is_active FROM clients WHERE client_id=?", (client_id,))
|
|
client = await cur.fetchone()
|
|
if not client or not client["is_active"]:
|
|
raise HTTPException(status_code=403, detail="Invalid or inactive client")
|
|
cur = await conn.execute(
|
|
"SELECT id, provider, api_key FROM provider_keys WHERE client_id=? AND is_active=1 "
|
|
"ORDER BY priority, id", (client_id,))
|
|
keys = await cur.fetchall()
|
|
if not keys:
|
|
raise HTTPException(status_code=400, detail="No active provider keys configured")
|
|
keys = sorted(keys, key=lambda key: (0 if preferred and key["provider"] == preferred else 1))
|
|
errors = []
|
|
for key in keys:
|
|
started = time.monotonic()
|
|
try:
|
|
result = await post_upstream(key["provider"], key["api_key"], model, prompt, images, params)
|
|
result["elapsed_sec"] = round(time.monotonic() - started, 2)
|
|
result["failover_attempts"] = len(errors)
|
|
return result
|
|
except UpstreamError as exc:
|
|
errors.append({"provider": key["provider"], "status_code": exc.status_code})
|
|
async with db() as conn:
|
|
await conn.execute("UPDATE provider_keys SET fail_count=fail_count+1,last_error_at=? WHERE id=?",
|
|
(time.time(), key["id"]))
|
|
await conn.commit()
|
|
if not exc.retryable:
|
|
break
|
|
raise HTTPException(status_code=502, detail={"message": "Upstream generation failed", "attempts": errors})
|
|
|
|
|
|
async def create_generation_run(generation_id: str, client_id: str, job_id: str | None,
|
|
request_data: dict[str, Any], status: str, image_count: int,
|
|
upload_bytes: int, request_intake_sec: float,
|
|
received_at: float) -> None:
|
|
async with db() as conn:
|
|
await conn.execute(
|
|
"INSERT INTO generation_runs(generation_id,client_id,job_id,status,requested_provider,"
|
|
"requested_model,image_count,upload_bytes,request_intake_sec,created_at,started_at) "
|
|
"VALUES(?,?,?,?,?,?,?,?,?,?,?)",
|
|
(generation_id, client_id, job_id, status, request_data.get("preferred"),
|
|
request_data["model"], image_count, upload_bytes, request_intake_sec,
|
|
received_at, time.time() if status == "processing" else None),
|
|
)
|
|
await conn.commit()
|
|
|
|
|
|
async def finish_generation_run(generation_id: str, status: str, completed_at: float,
|
|
total_sec: float, result: dict[str, Any] | None = None,
|
|
error: Any = None) -> None:
|
|
result = result or {}
|
|
async with db() as conn:
|
|
await conn.execute(
|
|
"UPDATE generation_runs SET status=?,provider_used=?,model_used=?,upstream_sec=?,"
|
|
"total_sec=?,failover_attempts=?,cost_usd=?,cost_source=?,usage_json=?,error_json=?,"
|
|
"completed_at=? WHERE generation_id=?",
|
|
(status, result.get("provider_used"), result.get("model_used"), result.get("elapsed_sec"),
|
|
total_sec, result.get("failover_attempts", 0), result.get("cost_usd"),
|
|
result.get("cost_source"), json.dumps(result.get("usage")) if result.get("usage") is not None else None,
|
|
json.dumps(error) if error is not None else None, completed_at, generation_id),
|
|
)
|
|
await conn.commit()
|
|
|
|
|
|
async def run_job(job_id: str, client_id: str, request_data: dict[str, Any],
|
|
run_metadata: dict[str, Any]) -> None:
|
|
started_at = time.time()
|
|
async with db() as conn:
|
|
await conn.execute("UPDATE jobs SET status='processing',updated_at=? WHERE job_id=?",
|
|
(started_at, job_id))
|
|
await conn.execute("UPDATE generation_runs SET status='processing',started_at=? WHERE generation_id=?",
|
|
(started_at, job_id))
|
|
await conn.commit()
|
|
try:
|
|
result = await execute_with_failover(**request_data)
|
|
completed_at = time.time()
|
|
result.update({"generation_id": job_id,
|
|
"request_intake_sec": run_metadata["request_intake_sec"],
|
|
"total_sec": round(completed_at - run_metadata["received_at"], 2)})
|
|
async with db() as conn:
|
|
await conn.execute("UPDATE jobs SET status='completed',result_json=?,updated_at=? WHERE job_id=?",
|
|
(json.dumps(result), completed_at, job_id))
|
|
await conn.commit()
|
|
await finish_generation_run(job_id, "completed", completed_at,
|
|
result["total_sec"], result=result)
|
|
except HTTPException as exc:
|
|
error = exc.detail
|
|
completed_at = time.time()
|
|
async with db() as conn:
|
|
await conn.execute("UPDATE jobs SET status='failed',error_json=?,updated_at=? WHERE job_id=?",
|
|
(json.dumps(error), completed_at, job_id))
|
|
await conn.commit()
|
|
await finish_generation_run(job_id, "failed", completed_at,
|
|
round(completed_at - run_metadata["received_at"], 2), error=error)
|
|
except Exception:
|
|
LOG.exception("Unexpected async job failure: %s", job_id)
|
|
completed_at = time.time()
|
|
error = {"message": "Internal generation error"}
|
|
async with db() as conn:
|
|
await conn.execute("UPDATE jobs SET status='failed',error_json=?,updated_at=? WHERE job_id=?",
|
|
(json.dumps(error), completed_at, job_id))
|
|
await conn.commit()
|
|
await finish_generation_run(job_id, "failed", completed_at,
|
|
round(completed_at - run_metadata["received_at"], 2), error=error)
|
|
|
|
|
|
@app.post("/v1/generate")
|
|
async def generate(request: Request, client_id: str = Form(...), prompt: str = Form(""),
|
|
model: str = Form("nano-banana"), provider: str | None = Form(None),
|
|
params: str = Form("{}"), async_mode: bool = Query(False),
|
|
images: list[UploadFile] = File(default=[]),
|
|
x_client_secret: str | None = Header(None)):
|
|
await authorize_client(client_id, x_client_secret)
|
|
if provider and provider.lower() not in {"google", "openrouter"}:
|
|
raise HTTPException(status_code=400, detail="provider must be google or openrouter")
|
|
try:
|
|
parsed_params = json.loads(params or "{}")
|
|
except json.JSONDecodeError as exc:
|
|
raise HTTPException(status_code=400, detail="params must be valid JSON") from exc
|
|
if not isinstance(parsed_params, dict):
|
|
raise HTTPException(status_code=400, detail="params must be a JSON object")
|
|
encoded, upload_bytes = await read_images(images)
|
|
request_data = {"client_id": client_id, "preferred": provider.lower() if provider else None,
|
|
"model": model, "prompt": prompt, "images": encoded, "params": parsed_params}
|
|
generation_id = str(uuid.uuid4())
|
|
received_at = request.state.received_at
|
|
request_intake_sec = round(time.monotonic() - request.state.received_monotonic, 2)
|
|
if not async_mode:
|
|
await create_generation_run(generation_id, client_id, None, request_data, "processing",
|
|
len(encoded), upload_bytes, request_intake_sec, received_at)
|
|
try:
|
|
result = await execute_with_failover(**request_data)
|
|
completed_at = time.time()
|
|
result.update({"generation_id": generation_id,
|
|
"request_intake_sec": request_intake_sec,
|
|
"total_sec": round(completed_at - received_at, 2)})
|
|
await finish_generation_run(generation_id, "completed", completed_at,
|
|
result["total_sec"], result=result)
|
|
return result
|
|
except HTTPException as exc:
|
|
completed_at = time.time()
|
|
await finish_generation_run(generation_id, "failed", completed_at,
|
|
round(completed_at - received_at, 2), error=exc.detail)
|
|
raise
|
|
except Exception:
|
|
LOG.exception("Unexpected generation failure: %s", generation_id)
|
|
completed_at = time.time()
|
|
error = {"message": "Internal generation error"}
|
|
await finish_generation_run(generation_id, "failed", completed_at,
|
|
round(completed_at - received_at, 2), error=error)
|
|
raise HTTPException(status_code=500, detail=error)
|
|
job_id = generation_id
|
|
now = time.time()
|
|
async with db() as conn:
|
|
await conn.execute("INSERT INTO jobs(job_id,client_id,status,request_json,created_at,updated_at) "
|
|
"VALUES(?,?, 'queued', ?, ?, ?)",
|
|
(job_id, client_id, json.dumps(request_data), now, now))
|
|
await conn.execute(
|
|
"INSERT INTO generation_runs(generation_id,client_id,job_id,status,requested_provider,"
|
|
"requested_model,image_count,upload_bytes,request_intake_sec,created_at) "
|
|
"VALUES(?,?,?,'queued',?,?,?,?,?,?)",
|
|
(job_id, client_id, job_id, request_data.get("preferred"), model, len(encoded),
|
|
upload_bytes, request_intake_sec, received_at),
|
|
)
|
|
await conn.commit()
|
|
asyncio.create_task(run_job(job_id, client_id, request_data,
|
|
{"received_at": received_at,
|
|
"request_intake_sec": request_intake_sec}))
|
|
return {"job_id": job_id, "status": "queued", "poll_url": f"/v1/jobs/{job_id}"}
|
|
|
|
|
|
@app.get("/v1/jobs/{job_id}")
|
|
async def get_job(job_id: str, client_id: str = Query(...),
|
|
x_client_secret: str | None = Header(None)):
|
|
await authorize_client(client_id, x_client_secret)
|
|
async with db() as conn:
|
|
cur = await conn.execute("SELECT * FROM jobs WHERE job_id=? AND client_id=?", (job_id, client_id))
|
|
row = await cur.fetchone()
|
|
if not row:
|
|
raise HTTPException(status_code=404, detail="Job not found")
|
|
response: dict[str, Any] = {"job_id": job_id, "status": row["status"],
|
|
"created_at": row["created_at"], "updated_at": row["updated_at"]}
|
|
if row["result_json"]:
|
|
response["result"] = json.loads(row["result_json"])
|
|
if row["error_json"]:
|
|
response["error"] = json.loads(row["error_json"])
|
|
return response
|
|
|
|
|
|
@app.get("/api/admin/state")
|
|
async def admin_state(x_admin_secret: str | None = Header(None)):
|
|
auth_admin(x_admin_secret)
|
|
async with db() as conn:
|
|
cur = await conn.execute("SELECT client_id,name,is_active,created_at FROM clients ORDER BY client_id")
|
|
clients = [dict(row) for row in await cur.fetchall()]
|
|
cur = await conn.execute("SELECT id,client_id,provider,substr(api_key,1,7)||'...' key_preview,"
|
|
"priority,is_active,fail_count FROM provider_keys ORDER BY client_id,priority,id")
|
|
keys = [dict(row) for row in await cur.fetchall()]
|
|
return {"clients": clients, "keys": keys}
|
|
|
|
|
|
@app.get("/api/admin/generations")
|
|
async def admin_generations(limit: int = Query(50, ge=1, le=200),
|
|
x_admin_secret: str | None = Header(None)):
|
|
auth_admin(x_admin_secret)
|
|
async with db() as conn:
|
|
cur = await conn.execute(
|
|
"SELECT COUNT(*) AS total,"
|
|
"SUM(CASE WHEN status='completed' THEN 1 ELSE 0 END) AS completed,"
|
|
"SUM(CASE WHEN status='failed' THEN 1 ELSE 0 END) AS failed,"
|
|
"COUNT(cost_usd) FILTER (WHERE status='completed') AS priced,"
|
|
"COALESCE(SUM(cost_usd) FILTER (WHERE status='completed'),0) AS total_cost_usd,"
|
|
"AVG(total_sec) FILTER (WHERE status='completed') AS avg_total_sec,"
|
|
"AVG(upstream_sec) FILTER (WHERE status='completed') AS avg_upstream_sec "
|
|
"FROM generation_runs"
|
|
)
|
|
summary = dict(await cur.fetchone())
|
|
cur = await conn.execute(
|
|
"SELECT client_id,COUNT(*) AS total,"
|
|
"SUM(CASE WHEN status='completed' THEN 1 ELSE 0 END) AS completed,"
|
|
"SUM(CASE WHEN status='failed' THEN 1 ELSE 0 END) AS failed,"
|
|
"COUNT(cost_usd) FILTER (WHERE status='completed') AS priced,"
|
|
"COALESCE(SUM(cost_usd) FILTER (WHERE status='completed'),0) AS total_cost_usd,"
|
|
"AVG(total_sec) FILTER (WHERE status='completed') AS avg_total_sec "
|
|
"FROM generation_runs GROUP BY client_id ORDER BY total_cost_usd DESC,client_id"
|
|
)
|
|
by_client = [dict(row) for row in await cur.fetchall()]
|
|
cur = await conn.execute(
|
|
"SELECT generation_id,client_id,job_id,status,requested_provider,provider_used,"
|
|
"requested_model,model_used,image_count,upload_bytes,request_intake_sec,upstream_sec,"
|
|
"total_sec,failover_attempts,cost_usd,cost_source,error_json,created_at,completed_at "
|
|
"FROM generation_runs ORDER BY created_at DESC LIMIT ?", (limit,)
|
|
)
|
|
generations = []
|
|
for row in await cur.fetchall():
|
|
generation = dict(row)
|
|
if generation["error_json"]:
|
|
try:
|
|
generation["error"] = json.loads(generation.pop("error_json"))
|
|
except json.JSONDecodeError:
|
|
generation["error"] = {"message": "Generation failed"}
|
|
generation.pop("error_json", None)
|
|
else:
|
|
generation.pop("error_json", None)
|
|
generations.append(generation)
|
|
return {"summary": summary, "by_client": by_client, "generations": generations}
|
|
|
|
|
|
@app.post("/api/admin/client")
|
|
async def add_client(client_id: str = Form(...), name: str = Form(...),
|
|
client_secret: str = Form(...), x_admin_secret: str | None = Header(None)):
|
|
auth_admin(x_admin_secret)
|
|
if not client_id.strip() or len(client_id) > 100 or not name.strip():
|
|
raise HTTPException(status_code=400, detail="client_id and name are required")
|
|
if len(client_secret) < 16 or len(client_secret) > 256:
|
|
raise HTTPException(status_code=400, detail="client_secret must be 16 to 256 characters")
|
|
hashed = secret_hash(client_secret)
|
|
async with db() as conn:
|
|
await conn.execute("INSERT INTO clients(client_id,name,secret_hash,created_at) VALUES(?,?,?,?) "
|
|
"ON CONFLICT(client_id) DO UPDATE SET name=excluded.name,secret_hash=excluded.secret_hash",
|
|
(client_id.strip(), name.strip(), hashed, time.time()))
|
|
await conn.commit()
|
|
return {"status": "ok", "client_id": client_id.strip()}
|
|
|
|
|
|
@app.post("/api/keys")
|
|
async def add_key(client_id: str = Form(...), provider: str = Form(...), api_key: str = Form(...),
|
|
priority: int = Form(1), auth_secret: str | None = Header(None)):
|
|
provider = provider.strip().lower()
|
|
if provider not in {"google", "openrouter"} or not api_key.strip():
|
|
raise HTTPException(status_code=400, detail="A supported provider and API key are required")
|
|
if priority < 1 or priority > 1000:
|
|
raise HTTPException(status_code=400, detail="priority must be between 1 and 1000")
|
|
auth_admin(auth_secret) if auth_secret and hmac.compare_digest(auth_secret, ADMIN_SECRET) else None
|
|
async with db() as conn:
|
|
cur = await conn.execute("SELECT secret_hash FROM clients WHERE client_id=?", (client_id,))
|
|
client = await cur.fetchone()
|
|
if not client:
|
|
raise HTTPException(status_code=404, detail="Client not found")
|
|
if not (auth_secret and hmac.compare_digest(auth_secret, ADMIN_SECRET)) and not secret_matches(auth_secret, client["secret_hash"]):
|
|
raise HTTPException(status_code=401, detail="Unauthorized")
|
|
await conn.execute("INSERT INTO provider_keys(client_id,provider,api_key,priority,created_at) VALUES(?,?,?,?,?)",
|
|
(client_id, provider, api_key.strip(), priority, time.time()))
|
|
await conn.commit()
|
|
return {"status": "ok"}
|
|
|
|
|
|
@app.delete("/api/keys/{key_id}")
|
|
async def delete_key(key_id: int, auth_secret: str | None = Header(None)):
|
|
async with db() as conn:
|
|
cur = await conn.execute("SELECT k.client_id,c.secret_hash FROM provider_keys k JOIN clients c "
|
|
"ON c.client_id=k.client_id WHERE k.id=?", (key_id,))
|
|
row = await cur.fetchone()
|
|
if not row:
|
|
raise HTTPException(status_code=404, detail="Key not found")
|
|
if not (auth_secret and hmac.compare_digest(auth_secret, ADMIN_SECRET)) and not secret_matches(auth_secret, row["secret_hash"]):
|
|
raise HTTPException(status_code=401, detail="Unauthorized")
|
|
await conn.execute("DELETE FROM provider_keys WHERE id=?", (key_id,))
|
|
await conn.commit()
|
|
return {"status": "deleted"}
|
|
|
|
|
|
@app.get("/healthz")
|
|
async def healthz():
|
|
return {"status": "ok"}
|
|
|
|
|
|
@app.get("/ui", response_class=HTMLResponse)
|
|
async def portal_ui():
|
|
return """<!doctype html><html><head><meta charset=utf-8><meta name=viewport content='width=device-width'>
|
|
<title>Proxy administration</title><style>
|
|
body{font:15px system-ui;max-width:1280px;margin:2rem auto;padding:0 1rem;background:#111;color:#eee}
|
|
input,select,button{padding:.6rem;margin:.25rem;background:#222;color:#eee;border:1px solid #555;border-radius:4px}
|
|
button{cursor:pointer;background:#2563eb}form{margin:1rem 0}.panel{background:#1b1b1b;border:1px solid #333;border-radius:8px;padding:1rem;margin:1rem 0}
|
|
.cards{display:grid;grid-template-columns:repeat(auto-fit,minmax(150px,1fr));gap:.75rem}.card{background:#222;padding:.8rem;border-radius:6px}
|
|
.value{font-size:1.3rem;font-weight:700}.muted{color:#aaa;font-size:.85rem}.table-wrap{overflow:auto}table{border-collapse:collapse;width:100%;font-size:.88rem}
|
|
th,td{text-align:left;padding:.55rem;border-bottom:1px solid #393939;vertical-align:top}th{position:sticky;top:0;background:#202020}
|
|
.ok{color:#70d89b}.bad{color:#ff8585}pre{white-space:pre-wrap;overflow-wrap:anywhere}
|
|
</style></head><body>
|
|
<h1>Generative proxy</h1><p class=muted>Admin secret is sent only to same-origin management APIs.</p>
|
|
<section class=panel><label>Admin secret <input id=s type=password autocomplete=current-password></label><button onclick=load()>Load / refresh</button></section>
|
|
<section class=panel><h2>Execution overview</h2><div id=summary class=cards></div><h3>By client</h3><div id=by-client></div></section>
|
|
<section class=panel><h2>Recent generations</h2><div class=table-wrap id=history></div></section>
|
|
<section class=panel><h2>Clients and provider keys</h2><pre id=o></pre></section>
|
|
<section class=panel><h2>Create or update client</h2><form id=c><input name=client_id placeholder='Client ID' required><input name=name placeholder=Name required><input name=client_secret placeholder='Client secret (16+ characters)' minlength=16 maxlength=256 required><button>Create or update client</button></form>
|
|
<h2>Add provider key</h2><form id=k><input name=client_id placeholder='Client ID' required><select name=provider><option>google</option><option>openrouter</option></select><input name=api_key placeholder='Provider API key' required><input name=priority type=number value=1 min=1><button>Add key</button></form></section>
|
|
<script>
|
|
const esc=v=>String(v??'').replace(/[&<>"']/g,c=>({'&':'&','<':'<','>':'>','"':'"',"'":'''}[c]));
|
|
const sec=v=>v==null?'—':Number(v).toFixed(2)+' s';
|
|
const money=v=>v==null?'—':'$'+Number(v).toFixed(6);
|
|
async function api(url,opts={}){opts.headers={...(opts.headers||{}),'x-admin-secret':document.querySelector('#s').value,'auth-secret':document.querySelector('#s').value};return fetch(url,opts)}
|
|
function table(headers,rows){return '<table><thead><tr>'+headers.map(x=>'<th>'+esc(x)+'</th>').join('')+'</tr></thead><tbody>'+rows.join('')+'</tbody></table>'}
|
|
async function load(){
|
|
const [sr,gr]=await Promise.all([api('/api/admin/state'),api('/api/admin/generations?limit=100')]);
|
|
const state=await sr.json();const stats=await gr.json();
|
|
if(!sr.ok||!gr.ok){document.querySelector('#o').textContent=JSON.stringify(state.detail||stats.detail||state||stats);return}
|
|
document.querySelector('#o').textContent=JSON.stringify(state,null,2);
|
|
const s=stats.summary;
|
|
document.querySelector('#summary').innerHTML=[['Generations',s.total],['Completed',s.completed||0],['Failed',s.failed||0],['Successful cost',money(s.total_cost_usd)],['Priced successes',s.priced],['Avg total (success)',sec(s.avg_total_sec)],['Avg upstream (success)',sec(s.avg_upstream_sec)]].map(([k,v])=>'<div class=card><div class=muted>'+esc(k)+'</div><div class=value>'+esc(v)+'</div></div>').join('');
|
|
document.querySelector('#by-client').innerHTML=table(['Client','Runs','Completed','Failed','Successful cost','Priced successes','Avg total (success)'],stats.by_client.map(x=>'<tr><td>'+esc(x.client_id)+'</td><td>'+esc(x.total)+'</td><td>'+esc(x.completed||0)+'</td><td>'+esc(x.failed||0)+'</td><td>'+esc(money(x.total_cost_usd))+'</td><td>'+esc(x.priced)+'</td><td>'+esc(sec(x.avg_total_sec))+'</td></tr>'));
|
|
document.querySelector('#history').innerHTML=table(['Time','Client','Status','Provider (requested → used)','Model','Images / upload','Intake','Upstream','Total','Failovers','Cost','Error'],stats.generations.map(x=>'<tr><td>'+esc(new Date(x.created_at*1000).toLocaleString())+'</td><td>'+esc(x.client_id)+'</td><td class="'+(x.status==='completed'?'ok':x.status==='failed'?'bad':'')+'">'+esc(x.status)+'</td><td>'+esc(x.requested_provider||'auto')+' → '+esc(x.provider_used||'—')+'</td><td>'+esc(x.model_used||x.requested_model)+'</td><td>'+esc(x.image_count)+' / '+esc((x.upload_bytes/1048576).toFixed(2))+' MiB</td><td>'+esc(sec(x.request_intake_sec))+'</td><td>'+esc(sec(x.upstream_sec))+'</td><td>'+esc(sec(x.total_sec))+'</td><td>'+esc(x.failover_attempts)+'</td><td>'+esc(money(x.cost_usd))+'<div class=muted>'+esc(x.cost_source||'')+'</div></td><td>'+esc(x.error?JSON.stringify(x.error):'')+'</td></tr>'));
|
|
}
|
|
document.querySelector('#c').onsubmit=async e=>{e.preventDefault();let r=await api('/api/admin/client',{method:'POST',body:new FormData(e.target)});document.querySelector('#o').textContent=await r.text();load()};
|
|
document.querySelector('#k').onsubmit=async e=>{e.preventDefault();let r=await api('/api/keys',{method:'POST',body:new FormData(e.target)});document.querySelector('#o').textContent=await r.text();load()};
|
|
</script></body></html>"""
|