89 lines
2.9 KiB
Python
89 lines
2.9 KiB
Python
import base64
|
|
import json
|
|
import os
|
|
from typing import Annotated
|
|
|
|
import httpx
|
|
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
|
|
|
|
|
|
app = FastAPI(title="Gemini image proxy", docs_url=None, redoc_url=None)
|
|
|
|
GEMINI_MODEL = os.getenv("GEMINI_MODEL", "gemini-3.6-flash")
|
|
GEMINI_API_URL = "https://generativelanguage.googleapis.com/v1beta/models"
|
|
MAX_IMAGE_BYTES = 10 * 1024 * 1024
|
|
ALLOWED_TYPES = {"image/jpeg", "image/png", "image/webp"}
|
|
|
|
|
|
@app.get("/health")
|
|
async def health() -> dict[str, str]:
|
|
return {"status": "ok"}
|
|
|
|
|
|
@app.post("/gemini")
|
|
async def recognize(
|
|
image: Annotated[UploadFile, File(...)],
|
|
prompt: Annotated[str, Form(...)],
|
|
generation_config: Annotated[str, Form()] = "{}",
|
|
api_key_override: Annotated[str | None, Form()] = None,
|
|
) -> dict:
|
|
if image.content_type not in ALLOWED_TYPES:
|
|
raise HTTPException(status_code=415, detail="Unsupported image type")
|
|
|
|
image_data = await image.read(MAX_IMAGE_BYTES + 1)
|
|
if len(image_data) > MAX_IMAGE_BYTES:
|
|
raise HTTPException(status_code=413, detail="Image is too large")
|
|
|
|
api_key = api_key_override or os.getenv("GEMINI_API_KEY")
|
|
if not api_key:
|
|
raise HTTPException(status_code=503, detail="Gemini is not configured")
|
|
|
|
try:
|
|
config = json.loads(generation_config)
|
|
except json.JSONDecodeError as exc:
|
|
raise HTTPException(status_code=422, detail="Invalid generation_config") from exc
|
|
if not isinstance(config, dict):
|
|
raise HTTPException(status_code=422, detail="generation_config must be an object")
|
|
|
|
payload = {
|
|
"contents": [{
|
|
"parts": [
|
|
{"text": prompt},
|
|
{
|
|
"inline_data": {
|
|
"mime_type": image.content_type,
|
|
"data": base64.b64encode(image_data).decode("ascii"),
|
|
}
|
|
},
|
|
]
|
|
}],
|
|
"generationConfig": config,
|
|
}
|
|
|
|
url = f"{GEMINI_API_URL}/{GEMINI_MODEL}:generateContent"
|
|
try:
|
|
async with httpx.AsyncClient(timeout=180) as client:
|
|
response = await client.post(
|
|
url,
|
|
params={"key": api_key},
|
|
json=payload,
|
|
)
|
|
except httpx.HTTPError as exc:
|
|
raise HTTPException(status_code=503, detail="Gemini unavailable") from exc
|
|
|
|
if response.status_code != 200:
|
|
try:
|
|
provider_error = response.json().get("error", {}).get("message")
|
|
except ValueError:
|
|
provider_error = None
|
|
detail = provider_error or "Gemini request failed"
|
|
raise HTTPException(status_code=502, detail=detail)
|
|
|
|
data = response.json()
|
|
try:
|
|
candidate = data["candidates"][0]
|
|
text = candidate["content"]["parts"][0]["text"]
|
|
except (KeyError, IndexError, TypeError) as exc:
|
|
raise HTTPException(status_code=502, detail="Invalid Gemini response") from exc
|
|
|
|
return {"text": text, "usage": data.get("usageMetadata", {})} |