1145 lines
60 KiB
Python
1145 lines
60 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
import binascii
|
|
import hashlib
|
|
import hmac
|
|
import io
|
|
import json
|
|
import logging
|
|
import os
|
|
import re
|
|
import sqlite3
|
|
import time
|
|
import uuid
|
|
from contextlib import asynccontextmanager
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from urllib.parse import quote
|
|
|
|
import aiosqlite
|
|
import httpx
|
|
from fastapi import FastAPI, File, Form, Header, HTTPException, Query, Request, UploadFile
|
|
from fastapi.responses import FileResponse, HTMLResponse, Response
|
|
from PIL import Image, ImageOps, UnidentifiedImageError
|
|
|
|
LOG = logging.getLogger("phproxy")
|
|
DB_PATH = os.getenv("DB_PATH", "./data/proxy.db")
|
|
OUTPUT_DIR = Path(os.getenv("OUTPUT_DIR", "./data/generated")).resolve()
|
|
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"))
|
|
IMAGE_LINK_TTL_SECONDS = int(os.getenv("IMAGE_LINK_TTL_SECONDS", "3600"))
|
|
OUTPUT_WEBP_QUALITY = int(os.getenv("OUTPUT_WEBP_QUALITY", "90"))
|
|
OUTPUT_RETENTION_SECONDS = int(os.getenv("OUTPUT_RETENTION_SECONDS", "86400"))
|
|
OUTPUT_CLEANUP_INTERVAL_SECONDS = 300
|
|
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))
|
|
|
|
|
|
def make_image_download_token(job_id: str, client_id: str, image_index: int) -> str:
|
|
claims = json.dumps(
|
|
{"job_id": job_id, "client_id": client_id, "image_index": image_index,
|
|
"expires_at": int(time.time()) + IMAGE_LINK_TTL_SECONDS},
|
|
separators=(",", ":"),
|
|
).encode("utf-8")
|
|
payload = base64.urlsafe_b64encode(claims).rstrip(b"=").decode("ascii")
|
|
signature = hmac.new(ADMIN_SECRET.encode("utf-8"), payload.encode("ascii"), hashlib.sha256).hexdigest()
|
|
return f"{payload}.{signature}"
|
|
|
|
|
|
def signed_image_download_url(job_id: str, client_id: str, image_index: int) -> str:
|
|
token = make_image_download_token(job_id, client_id, image_index)
|
|
return f"/v1/jobs/{job_id}/images/{image_index}?token={quote(token, safe='')}"
|
|
|
|
|
|
def verify_image_download_token(token: str, job_id: str, image_index: int) -> str:
|
|
try:
|
|
payload, supplied_signature = token.split(".", 1)
|
|
expected_signature = hmac.new(
|
|
ADMIN_SECRET.encode("utf-8"), payload.encode("ascii"), hashlib.sha256
|
|
).hexdigest()
|
|
if not hmac.compare_digest(supplied_signature, expected_signature):
|
|
raise ValueError("invalid signature")
|
|
raw_claims = base64.urlsafe_b64decode(payload + "=" * (-len(payload) % 4))
|
|
claims = json.loads(raw_claims)
|
|
if (claims.get("job_id") != job_id or claims.get("image_index") != image_index
|
|
or not isinstance(claims.get("client_id"), str)):
|
|
raise ValueError("token scope mismatch")
|
|
if not isinstance(claims.get("expires_at"), int) or claims["expires_at"] < int(time.time()):
|
|
raise HTTPException(status_code=410, detail="Image link has expired")
|
|
return claims["client_id"]
|
|
except HTTPException:
|
|
raise
|
|
except (ValueError, TypeError, UnicodeDecodeError, json.JSONDecodeError) as exc:
|
|
raise HTTPException(status_code=401, detail="Invalid image link") from exc
|
|
|
|
|
|
def output_image_path(job_id: str, image_index: int, mime_type: str) -> Path:
|
|
image_format = IMAGE_FORMAT_TYPES.get(mime_type)
|
|
if image_format is None or image_index < 0:
|
|
raise ValueError("Unsupported generated image type or index")
|
|
job_key = uuid.UUID(job_id).hex
|
|
return OUTPUT_DIR / f"{job_key}-{image_index}.{image_format.lower()}"
|
|
|
|
|
|
def convert_image_to_webp(content: bytes, mime_type: str) -> bytes:
|
|
if not 1 <= OUTPUT_WEBP_QUALITY <= 100:
|
|
raise ValueError("OUTPUT_WEBP_QUALITY must be between 1 and 100")
|
|
Image.MAX_IMAGE_PIXELS = MAX_IMAGE_PIXELS
|
|
try:
|
|
with Image.open(io.BytesIO(content)) as source:
|
|
actual_mime = IMAGE_FORMATS.get(source.format or "")
|
|
if actual_mime is None or actual_mime != mime_type:
|
|
raise ValueError("Generated image content does not match its MIME type")
|
|
width, height = source.size
|
|
if width <= 0 or height <= 0 or width * height > MAX_IMAGE_PIXELS:
|
|
raise ValueError("Generated image dimensions exceed configured limits")
|
|
source.load()
|
|
image = ImageOps.exif_transpose(source)
|
|
has_alpha = "A" in image.getbands() or "transparency" in image.info
|
|
image = image.convert("RGBA" if has_alpha else "RGB")
|
|
output = io.BytesIO()
|
|
image.save(output, format="WEBP", quality=OUTPUT_WEBP_QUALITY, method=4)
|
|
return output.getvalue()
|
|
except (Image.DecompressionBombError, Image.DecompressionBombWarning) as exc:
|
|
raise ValueError("Generated image dimensions exceed configured limits") from exc
|
|
except (UnidentifiedImageError, OSError) as exc:
|
|
raise ValueError("Provider returned an invalid image") from exc
|
|
|
|
|
|
def write_output_file(path: Path, content: bytes) -> None:
|
|
temp_path = OUTPUT_DIR / f".{path.name}.{uuid.uuid4().hex}.tmp"
|
|
with temp_path.open("xb") as image_file:
|
|
os.chmod(temp_path, 0o600)
|
|
image_file.write(content)
|
|
image_file.flush()
|
|
os.fsync(image_file.fileno())
|
|
os.replace(temp_path, path)
|
|
os.chmod(path, 0o600)
|
|
|
|
|
|
def cleanup_expired_output_files(now: float | None = None) -> int:
|
|
"""Delete generated image files once they have been stored for 24 hours."""
|
|
cutoff = (time.time() if now is None else now) - OUTPUT_RETENTION_SECONDS
|
|
generated_name = re.compile(r"[0-9a-f]{32}-[0-9]+\.(?:webp|png|jpe?g)")
|
|
removed = 0
|
|
try:
|
|
paths = list(OUTPUT_DIR.iterdir())
|
|
except FileNotFoundError:
|
|
return 0
|
|
for path in paths:
|
|
if not generated_name.fullmatch(path.name):
|
|
continue
|
|
try:
|
|
if path.is_file() and path.stat().st_mtime <= cutoff:
|
|
path.unlink()
|
|
removed += 1
|
|
except FileNotFoundError:
|
|
continue
|
|
except OSError:
|
|
LOG.exception("Could not remove expired output image %s", path.name)
|
|
return removed
|
|
|
|
|
|
async def output_cleanup_loop() -> None:
|
|
while True:
|
|
await asyncio.sleep(OUTPUT_CLEANUP_INTERVAL_SECONDS)
|
|
try:
|
|
removed = await asyncio.to_thread(cleanup_expired_output_files)
|
|
if removed:
|
|
LOG.info("Removed %s expired generated image(s)", removed)
|
|
except Exception:
|
|
LOG.exception("Generated image cleanup failed")
|
|
|
|
|
|
def persist_result_images(job_id: str, result: dict[str, Any]) -> tuple[dict[str, Any], bool, list[Path]]:
|
|
"""Store async image outputs as WebP and return compact result metadata."""
|
|
compact_result = dict(result)
|
|
compact_images = []
|
|
written_paths: list[Path] = []
|
|
superseded_paths: list[Path] = []
|
|
changed = False
|
|
try:
|
|
for image_index, image in enumerate(result.get("images", [])):
|
|
image_metadata = dict(image)
|
|
encoded = image_metadata.pop("b64_json", None)
|
|
if encoded:
|
|
mime_type = image_metadata.get("mime_type")
|
|
source_content = base64.b64decode(encoded, validate=True)
|
|
webp_content = convert_image_to_webp(source_content, mime_type)
|
|
path = output_image_path(job_id, image_index, "image/webp")
|
|
write_output_file(path, webp_content)
|
|
written_paths.append(path)
|
|
image_metadata["mime_type"] = "image/webp"
|
|
changed = True
|
|
elif image_metadata.get("mime_type") in IMAGE_FORMAT_TYPES:
|
|
old_path = output_image_path(job_id, image_index, image_metadata["mime_type"])
|
|
if image_metadata["mime_type"] != "image/webp" and old_path.is_file():
|
|
source_content = old_path.read_bytes()
|
|
webp_content = convert_image_to_webp(source_content, image_metadata["mime_type"])
|
|
path = output_image_path(job_id, image_index, "image/webp")
|
|
write_output_file(path, webp_content)
|
|
written_paths.append(path)
|
|
superseded_paths.append(old_path)
|
|
image_metadata["mime_type"] = "image/webp"
|
|
changed = True
|
|
compact_images.append(image_metadata)
|
|
compact_result["images"] = compact_images
|
|
return compact_result, changed, superseded_paths
|
|
except Exception:
|
|
for path in written_paths:
|
|
path.unlink(missing_ok=True)
|
|
raise
|
|
|
|
|
|
def compact_sqlite_database() -> None:
|
|
"""Checkpoint and compact SQLite without an outstanding async cursor."""
|
|
conn = sqlite3.connect(DB_PATH, timeout=30)
|
|
try:
|
|
checkpoint = conn.execute("PRAGMA wal_checkpoint(TRUNCATE)")
|
|
checkpoint.fetchall()
|
|
checkpoint.close()
|
|
vacuum = conn.execute("VACUUM")
|
|
vacuum.close()
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
@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)
|
|
OUTPUT_DIR.mkdir(parents=True, exist_ok=True, mode=0o700)
|
|
os.chmod(OUTPUT_DIR, 0o700)
|
|
migrated_result_count = 0
|
|
cleared_request_count = 0
|
|
superseded_paths: list[Path] = []
|
|
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()
|
|
|
|
cur = await conn.execute(
|
|
"SELECT job_id,result_json,request_json FROM jobs "
|
|
"WHERE status IN ('completed','failed')"
|
|
)
|
|
for row in await cur.fetchall():
|
|
if row["result_json"]:
|
|
try:
|
|
result = json.loads(row["result_json"])
|
|
compact_result, changed, old_paths = await asyncio.to_thread(
|
|
persist_result_images, row["job_id"], result
|
|
)
|
|
superseded_paths.extend(old_paths)
|
|
if changed:
|
|
await conn.execute(
|
|
"UPDATE jobs SET result_json=? WHERE job_id=?",
|
|
(json.dumps(compact_result), row["job_id"]),
|
|
)
|
|
migrated_result_count += 1
|
|
except Exception:
|
|
# Keep the legacy JSON available if an individual historic image cannot be migrated.
|
|
LOG.exception("Could not migrate stored output images for job %s", row["job_id"])
|
|
if row["request_json"] and row["request_json"] != "{}":
|
|
await conn.execute("UPDATE jobs SET request_json='{}' WHERE job_id=?", (row["job_id"],))
|
|
cleared_request_count += 1
|
|
await conn.commit()
|
|
|
|
# Remove old PNG/JPEG outputs only after their WebP replacements are committed.
|
|
for old_path in superseded_paths:
|
|
try:
|
|
old_path.unlink(missing_ok=True)
|
|
except OSError:
|
|
LOG.exception("Could not remove superseded output image %s", old_path.name)
|
|
|
|
async with db() as conn:
|
|
cur = await conn.execute("PRAGMA freelist_count")
|
|
free_page_count = (await cur.fetchone())[0]
|
|
if migrated_result_count or cleared_request_count or free_page_count > 1000:
|
|
# Reclaim SQLite pages freed by removing Base64 images and completed request payloads.
|
|
try:
|
|
await asyncio.to_thread(compact_sqlite_database)
|
|
except sqlite3.Error:
|
|
LOG.exception("Could not compact the SQLite database; service will continue")
|
|
if migrated_result_count or cleared_request_count:
|
|
LOG.info("Migrated %s job result(s) to files and cleared %s completed request payload(s)",
|
|
migrated_result_count, cleared_request_count)
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(_: FastAPI):
|
|
global http_client
|
|
cleanup_task: asyncio.Task | None = None
|
|
if not ADMIN_SECRET:
|
|
raise RuntimeError("ADMIN_SECRET must be set before starting the service")
|
|
if OUTPUT_RETENTION_SECONDS <= 0:
|
|
raise RuntimeError("OUTPUT_RETENTION_SECONDS must be greater than zero")
|
|
await init_db()
|
|
removed = await asyncio.to_thread(cleanup_expired_output_files)
|
|
if removed:
|
|
LOG.info("Removed %s expired generated image(s) during startup", removed)
|
|
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,
|
|
)
|
|
cleanup_task = asyncio.create_task(output_cleanup_loop())
|
|
try:
|
|
yield
|
|
finally:
|
|
cleanup_task.cancel()
|
|
try:
|
|
await cleanup_task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
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)})
|
|
compact_result, _, superseded_paths = await asyncio.to_thread(
|
|
persist_result_images, job_id, result
|
|
)
|
|
async with db() as conn:
|
|
await conn.execute(
|
|
"UPDATE jobs SET status='completed',result_json=?,request_json='{}',updated_at=? WHERE job_id=?",
|
|
(json.dumps(compact_result), completed_at, job_id),
|
|
)
|
|
await conn.commit()
|
|
for old_path in superseded_paths:
|
|
try:
|
|
old_path.unlink(missing_ok=True)
|
|
except OSError:
|
|
LOG.exception("Could not remove superseded output image %s", old_path.name)
|
|
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',request_json='{}',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',request_json='{}',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(...),
|
|
include_result: bool = Query(True),
|
|
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"]:
|
|
result = json.loads(row["result_json"])
|
|
if include_result:
|
|
legacy_result = dict(result)
|
|
legacy_images = []
|
|
for index, image in enumerate(result.get("images", [])):
|
|
image_data = dict(image)
|
|
if not image_data.get("b64_json") and image_data.get("mime_type") in IMAGE_FORMAT_TYPES:
|
|
try:
|
|
path = output_image_path(job_id, index, image_data["mime_type"])
|
|
if path.is_file():
|
|
content = await asyncio.to_thread(path.read_bytes)
|
|
image_data["b64_json"] = base64.b64encode(content).decode("ascii")
|
|
except (ValueError, OSError):
|
|
LOG.exception("Could not read stored output image for job %s, image %s",
|
|
job_id, index)
|
|
legacy_images.append(image_data)
|
|
legacy_result["images"] = legacy_images
|
|
response["result"] = legacy_result
|
|
else:
|
|
# Keep status polling small; image bytes are fetched through the binary route.
|
|
result_metadata = {key: value for key, value in result.items() if key != "images"}
|
|
result_metadata["images"] = []
|
|
for index, image in enumerate(result.get("images", [])):
|
|
image_metadata = {key: value for key, value in image.items()
|
|
if key not in {"b64_json", "url"}}
|
|
image_is_local = bool(image.get("b64_json"))
|
|
if not image_is_local and image.get("mime_type") in IMAGE_FORMAT_TYPES:
|
|
try:
|
|
path = output_image_path(job_id, index, image.get("mime_type"))
|
|
image_is_local = path.is_file()
|
|
except ValueError:
|
|
LOG.warning("Stored output image has unsupported MIME type for job %s, image %s",
|
|
job_id, index)
|
|
if image_is_local:
|
|
image_metadata["download_url"] = signed_image_download_url(
|
|
job_id, client_id, index
|
|
)
|
|
elif image.get("url"):
|
|
image_metadata["url"] = image["url"]
|
|
result_metadata["images"].append(image_metadata)
|
|
response["result"] = result_metadata
|
|
if row["error_json"]:
|
|
response["error"] = json.loads(row["error_json"])
|
|
return response
|
|
|
|
|
|
@app.get("/v1/jobs/{job_id}/images/{image_index}")
|
|
async def get_job_image(job_id: str, image_index: int, client_id: str | None = Query(None),
|
|
token: str | None = Query(None),
|
|
x_client_secret: str | None = Header(None)):
|
|
"""Return raw image bytes via a short-lived signed link or client-secret auth."""
|
|
if token:
|
|
token_client_id = verify_image_download_token(token, job_id, image_index)
|
|
if client_id and client_id != token_client_id:
|
|
raise HTTPException(status_code=401, detail="Invalid image link")
|
|
client_id = token_client_id
|
|
elif client_id:
|
|
await authorize_client(client_id, x_client_secret)
|
|
else:
|
|
raise HTTPException(status_code=401, detail="An image link or client authentication is required")
|
|
|
|
async with db() as conn:
|
|
cur = await conn.execute(
|
|
"SELECT j.status,j.result_json,c.is_active FROM jobs j "
|
|
"JOIN clients c ON c.client_id=j.client_id WHERE j.job_id=? AND j.client_id=?",
|
|
(job_id, client_id),
|
|
)
|
|
row = await cur.fetchone()
|
|
if not row or not row["is_active"]:
|
|
raise HTTPException(status_code=404, detail="Job not found")
|
|
if row["status"] != "completed" or not row["result_json"]:
|
|
raise HTTPException(status_code=409, detail="Job has no completed result")
|
|
images = json.loads(row["result_json"]).get("images", [])
|
|
if image_index < 0 or image_index >= len(images):
|
|
raise HTTPException(status_code=404, detail="Image not found")
|
|
image = images[image_index]
|
|
mime_type = image.get("mime_type")
|
|
if mime_type in IMAGE_FORMAT_TYPES:
|
|
path = output_image_path(job_id, image_index, mime_type)
|
|
if path.is_file():
|
|
extension = IMAGE_FORMAT_TYPES[mime_type].lower()
|
|
return FileResponse(
|
|
path,
|
|
media_type=mime_type,
|
|
filename=f"generated-{image_index}.{extension}",
|
|
content_disposition_type="inline",
|
|
headers={
|
|
"Cache-Control": "private, no-store",
|
|
"X-Content-Type-Options": "nosniff",
|
|
"Referrer-Policy": "no-referrer",
|
|
},
|
|
)
|
|
|
|
# Keep legacy completed jobs accessible if a file migration was not possible.
|
|
encoded = image.get("b64_json")
|
|
if encoded:
|
|
try:
|
|
content = base64.b64decode(encoded, validate=True)
|
|
except (ValueError, TypeError, binascii.Error) as exc:
|
|
LOG.error("Stored output image is invalid for job %s, image %s", job_id, image_index)
|
|
raise HTTPException(status_code=500, detail="Stored image is invalid") from exc
|
|
return Response(
|
|
content=content,
|
|
media_type=mime_type if mime_type in IMAGE_FORMAT_TYPES else "application/octet-stream",
|
|
headers={
|
|
"Content-Disposition": f'inline; filename="generated-{image_index}"',
|
|
"Cache-Control": "private, no-store",
|
|
"X-Content-Type-Options": "nosniff",
|
|
"Referrer-Policy": "no-referrer",
|
|
},
|
|
)
|
|
if mime_type in IMAGE_FORMAT_TYPES:
|
|
raise HTTPException(status_code=410, detail="Image is no longer available")
|
|
raise HTTPException(status_code=409, detail="Image is available only from its upstream URL")
|
|
|
|
|
|
@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>"""
|