handler.py

from typing import Any, Dict, List, Optional
import os
import google.generativeai as genai

# === Load environment variables ===
API_KEY = os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")
MODEL_NAME = os.getenv("GEMINI_MODEL", "gemini-2.5-pro")
TEMPERATURE = float(os.getenv("TEMPERATURE", "0.7"))
TOP_P = float(os.getenv("TOP_P", "0.95"))
MAX_TOKENS = int(os.getenv("MAX_OUTPUT_TOKENS", "1024"))
SYSTEM_PROMPT = os.getenv("SYSTEM_PROMPT", "You are a helpful assistant.")

genai.configure(api_key=API_KEY)

def _extract_text(resp: Any) -> str:
    if getattr(resp, "text", None):
        return resp.text
    try:
        for c in getattr(resp, "candidates", []):
            for p in getattr(c, "content", {}).get("parts", []):
                if p.get("text"):
                    return p["text"]
    except Exception:
        pass
    return ""

class EndpointHandler:
    def __init__(self, path: str = ""):
        print("[handler:init] Loading Gemini model...", flush=True)
        self.model = genai.GenerativeModel(
            MODEL_NAME,
            system_instruction=SYSTEM_PROMPT
        )
        print("[handler:init] Model ready ✅", flush=True)

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        print("[handler:call] Incoming data:", data, flush=True)

        # 1️⃣ Try flexible input parsing
        text = data.get("message") or data.get("inputs") or ""
        if isinstance(text, dict):
            text = text.get("content") or ""
        text = str(text).strip()

        if not text:
            return {"text": "(empty input)"}

        params = data.get("parameters", {})
        gen_cfg = {
            "temperature": float(params.get("temperature", TEMPERATURE)),
            "top_p": float(params.get("top_p", TOP_P)),
            "max_output_tokens": int(params.get("max_output_tokens", MAX_TOKENS))
        }

        history = data.get("history") or []

        try:
            if history:
                chat = self.model.start_chat(history=history)
                resp = chat.send_message(text, generation_config=gen_cfg)
            else:
                resp = self.model.generate_content(text, generation_config=gen_cfg)
            reply = _extract_text(resp)
            if not reply:
                reply = "(no response)"
            return {"text": reply}
        except Exception as e:
            print("[handler:error]", e, flush=True)
            return {"text": f"Error: {e}"}