Update handler.py

handler.py

@@ -1,30 +1,27 @@
-# handler.py — Hugging Face custom handler compatible with your local Flask shapes
-# - Env: GEMINI_API_KEY (or GOOGLE_API_KEY), GEMINI_MODEL, TEMPERATURE, TOP_P, MAX_OUTPUT_TOKENS, SYSTEM_PROMPT
-# - Input shapes supported:
-#   1) {"inputs": "hello", "parameters": {...}}
-#   2) {"message": "hello", "history": [...], "parameters": {...}}   # your Flask shape
-#   3) {"inputs": {"messages":[{"role":"user","content":"hello"}]}, "parameters": {...}}
+# handler.py — custom Hugging Face handler compatible with your Flask shape
+# Env needed in Endpoint Settings → Environment Variables:
+#   GEMINI_API_KEY=xxxxxxxxxxxxxxxxxxxxxxxxxxxx
+# Optional:
+#   GEMINI_MODEL=gemini-2.5-pro
+#   TEMPERATURE=0.7
+#   TOP_P=0.95
+#   MAX_OUTPUT_TOKENS=1024
+#   SYSTEM_PROMPT=You are a helpful assistant.
 
 from typing import Any, Dict, List, Optional
 import os, time
 import google.generativeai as genai
 
-# ---------- Config (env) ----------
 API_KEY = os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")
-MODEL_NAME = os.getenv("GEMINI_MODEL", "models/gemini-2.5-pro")  # keep your local default
+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.")
 
-GEN_CFG = {
-    "temperature": TEMPERATURE,
-    "top_p": TOP_P,
-    "max_output_tokens": MAX_TOKENS,
-}
+GEN_CFG = {"temperature": TEMPERATURE, "top_p": TOP_P, "max_output_tokens": MAX_TOKENS}
 
 def _extract_text(resp: Any) -> str:
-    # google.generativeai responses usually expose .text; keep a fallback
     if getattr(resp, "text", None):
         return resp.text
     try:
@@ -32,69 +29,69 @@ def _extract_text(resp: Any) -> str:
             content = getattr(c, "content", None)
             for p in getattr(content, "parts", []) or []:
                 t = getattr(p, "text", None)
-                if t:
-                    return t
+                if t: return t
     except Exception:
         pass
     return ""
 
 def _to_gemini_history(history: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
     """
-    Accepts your local history shape:
-      [{ "role":"user"|"model", "parts":[{"text":"..."}] }, ...]
-    …or an OpenAI-ish/messages list:
-      [{ "role":"user"|"assistant", "content":"..."}, ...]
-    Produces Gemini-compatible history.
+    Supports your Flask history: [{role:'user'|'model', parts:[{text:'...'}]}]
+    and OpenAI-style: [{role:'user'|'assistant', content:'...'}]
     """
     out: List[Dict[str, Any]] = []
     for m in history or []:
         role = (m.get("role") or "").lower()
-        if "parts" in m:  # already Gemini-ish
+        if "parts" in m:  # already Gemini-like
             parts = m.get("parts") or []
-            out.append({"role": role if role in ("user","model") else "user", "parts": parts})
+            role = role if role in ("user","model") else "user"
+            out.append({"role": role, "parts": parts})
         else:
             content = m.get("content", "")
             if not isinstance(content, str):
-                try:
-                    content = str(content)
-                except Exception:
-                    content = ""
-            # map "assistant" -> "model"
+                content = str(content)
             role = "model" if role == "assistant" else ("user" if role != "model" else "model")
             out.append({"role": role, "parts": [{"text": content}]})
     return out
 
-def _pick_user_text(payload: Dict[str, Any]) -> Optional[str]:
-    # 1) {"inputs": "text"}
-    if isinstance(payload.get("inputs"), str):
-        return payload["inputs"].strip()
-
-    # 2) {"message":"text"}
-    msg = payload.get("message")
-    if isinstance(msg, str) and msg.strip():
-        return msg.strip()
-
-    # 3) {"inputs":{"messages":[...]}} pick last user message
-    x = payload.get("inputs")
+def _pick_user_text(data: Dict[str, Any]) -> Optional[str]:
+    # Supports: {"message":"..."}, or {"inputs":"..."}, or {"inputs":{"messages":[...]}}
+    if isinstance(data.get("message"), str) and data["message"].strip():
+        return data["message"].strip()
+    if isinstance(data.get("inputs"), str) and data["inputs"].strip():
+        return data["inputs"].strip()
+    x = data.get("inputs")
     if isinstance(x, dict):
         msgs = x.get("messages") or []
         for m in reversed(msgs):
             if (m.get("role","user").lower()) == "user":
                 c = m.get("content","")
                 return c if isinstance(c,str) else str(c)
-
     return None
 
 class EndpointHandler:
     def __init__(self, path: str = ""):
         if not API_KEY:
-            self._init_error = "Missing GEMINI_API_KEY/GOOGLE_API_KEY in Environment Variables."
+            self._init_error = "Missing GEMINI_API_KEY/GOOGLE_API_KEY."
             print("[handler:init] ERROR:", self._init_error, flush=True)
             return
 
         self._init_error = None
         genai.configure(api_key=API_KEY)
-        self.model = genai.GenerativeModel(MODEL_NAME, system_instruction=SYSTEM_PROMPT)
+
+        # Try model name, then a prefixed fallback
+        try:
+            self.model = genai.GenerativeModel(MODEL_NAME, system_instruction=SYSTEM_PROMPT)
+        except Exception as e1:
+            alt = MODEL_NAME.replace("models/", "") if MODEL_NAME.startswith("models/") else f"models/{MODEL_NAME}"
+            try:
+                self.model = genai.GenerativeModel(alt, system_instruction=SYSTEM_PROMPT)
+                print(f"[handler:init] Fallback model name used: {alt}", flush=True)
+            except Exception as e2:
+                self._init_error = f"Model init failed: {repr(e1)} | fallback: {repr(e2)}"
+                print("[handler:init] ERROR:", self._init_error, flush=True)
+                return
+
         # Warm-up (non-fatal if it fails)
         try:
             t0 = time.time()
@@ -104,11 +101,10 @@ class EndpointHandler:
             print("[handler:init] warm-up failed (non-fatal):", repr(e), flush=True)
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        # If init failed, surface that to the client
         if getattr(self, "_init_error", None):
-            return {"text":"", "debug":{"where":"init", "error": self._init_error}}
+            return {"text": "", "debug": {"where": "init", "error": self._init_error}}
 
-        # Parameters override (optional)
+        # Per-request overrides
         params = data.get("parameters") or {}
         gen_cfg = {
             "temperature": float(params.get("temperature", GEN_CFG["temperature"])),
@@ -116,16 +112,13 @@ class EndpointHandler:
             "max_output_tokens": int(params.get("max_output_tokens", GEN_CFG["max_output_tokens"])),
         }
 
-        # Determine the user text
         user_text = _pick_user_text(data) or ""
         if not user_text:
-            return {"text":"", "debug":{"where":"input", "error":"Empty prompt"}}
+            return {"text": "", "debug": {"where": "input", "error": "Empty prompt"}}
 
-        # Build history if provided (your local shape supported)
-        history = data.get("history") or []
-        gemini_history = _to_gemini_history(history)
+        gemini_history = _to_gemini_history(data.get("history") or [])
 
-        # Try chat w/ a tiny retry to dodge transient cold-start
+        # Try once; tiny retry handles cold-start/transient
         last_err = None
         for attempt in range(2):
             try:
@@ -133,20 +126,18 @@ class EndpointHandler:
                     chat = self.model.start_chat(history=gemini_history)
                     resp = chat.send_message(user_text, generation_config=gen_cfg)
                 else:
-                    # stateless call
                     resp = self.model.generate_content(user_text, generation_config=gen_cfg)
                 txt = _extract_text(resp)
                 if txt:
                     return {"text": txt}
-                # no text — return finish reasons if present
                 fin = []
                 try:
                     fin = [getattr(c, "finish_reason", None) for c in (resp.candidates or [])]
                 except Exception:
                     pass
-                return {"text":"", "debug":{"where":"empty_text", "finish_reasons": fin}}
+                return {"text": "", "debug": {"where": "empty_text", "finish_reasons": fin}}
             except Exception as e:
                 last_err = repr(e)
-                time.sleep(0.35)  # small backoff
+                time.sleep(0.35)
 
-        return {"text":"", "debug":{"where":"exception", "exception": last_err}}
+        return {"text": "", "debug": {"where": "exception", "exception": last_err}}