model.py

import json
import logging
import os

from google import genai
from google.genai import types

log = logging.getLogger("cardextractor")

# ponytail: key comes from the GEMINI_API_KEY Space secret; fail loud at startup if missing
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
MODEL = os.environ.get("GEMINI_MODEL", "gemma-4-26b-a4b-it")

PROMPT = """Extract the contact information from this business card.
Return ONLY JSON, no markdown, with exactly these keys:
{"name":"","company":"","designation":"","phone":"","email":"","website":"","address":""}"""


def extract_contact(image_path):
    with open(image_path, "rb") as f:
        image_bytes = f.read()

    # ponytail: let the image content-type ride on the extension; default jpeg
    ext = os.path.splitext(image_path)[1].lower().lstrip(".") or "jpeg"
    mime = "image/png" if ext == "png" else f"image/{ext}"

    # ponytail: Gemma models don't support response_mime_type=json, so parse
    # the JSON object out of the text ourselves (handles ```json fences too)
    resp = client.models.generate_content(
        model=MODEL,
        contents=[
            types.Part.from_bytes(data=image_bytes, mime_type=mime),
            PROMPT,
        ],
    )

    text = resp.text
    log.info("LLM response: %s", text)
    start = text.find("{")
    if start == -1:
        raise ValueError(f"no JSON in model output: {text!r}")
    data, _ = json.JSONDecoder().raw_decode(text[start:])
    return data