feat: add generative proxy with usage metrics
This commit is contained in:
vasily committed
commit
1f43b2d065
10 files changed
+1185
No files matched your search
+728
@@ -0,0 +1,728 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import hashlib
|
||||
import hmac
|
||||
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
|
||||
|
||||
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"))
|
||||
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"),
|
||||
},
|
||||
}
|
||||
ALLOWED_IMAGE_TYPES = {"image/jpeg", "image/png", "image/webp", "image/gif"}
|
||||
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 row["secret_hash"] and not secret_matches(secret, row["secret_hash"]):
|
||||
raise HTTPException(status_code=401, detail="Invalid client secret")
|
||||
|
||||
|
||||
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 ALLOWED_IMAGE_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 raw:
|
||||
encoded.append({"mime_type": mime, "b64": base64.b64encode(raw).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")
|
||||
hashed = secret_hash(client_secret) if client_secret else None
|
||||
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='Optional client secret'><button>Create 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>"""
|
||||
Reference in new issue
Block a user