Update handler.py

handler.py

@@ -1,87 +1,152 @@
-# handler.py — warm-up + retry, always returns text or clear debug
+# 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": {...}}
 
-from typing import Any, Dict
-import os, time, socket, requests
+from typing import Any, Dict, List, Optional
+import os, time
 import google.generativeai as genai
 
-MODEL         = os.getenv("GEMINI_MODEL", "gemini-1.5-flash")
-TEMPERATURE   = float(os.getenv("TEMPERATURE", "0.7"))
-TOP_P         = float(os.getenv("TOP_P", "0.95"))
-MAX_TOKENS    = int(os.getenv("MAX_OUTPUT_TOKENS", "512"))
+# ---------- 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
+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,
+}
+
 def _extract_text(resp: Any) -> str:
-    if getattr(resp, "text", None): return resp.text
+    # google.generativeai responses usually expose .text; keep a fallback
+    if getattr(resp, "text", None):
+        return resp.text
     try:
         for c in getattr(resp, "candidates", []) or []:
             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.
+    """
+    out: List[Dict[str, Any]] = []
+    for m in history or []:
+        role = (m.get("role") or "").lower()
+        if "parts" in m:  # already Gemini-ish
+            parts = m.get("parts") or []
+            out.append({"role": role if role in ("user","model") else "user", "parts": parts})
+        else:
+            content = m.get("content", "")
+            if not isinstance(content, str):
+                try:
+                    content = str(content)
+                except Exception:
+                    content = ""
+            # map "assistant" -> "model"
+            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")
+    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 = ""):
-        self.api_key = os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")
-        if not self.api_key:
-            self.init_error = "Missing GEMINI_API_KEY in Endpoint → Settings → Environment Variables."
-            print("[handler] INIT ERROR:", self.init_error, flush=True)
+        if not API_KEY:
+            self._init_error = "Missing GEMINI_API_KEY/GOOGLE_API_KEY in Environment Variables."
+            print("[handler:init] ERROR:", self._init_error, flush=True)
             return
-        self.init_error = None
 
-        # Configure SDK
-        genai.configure(api_key=self.api_key)
-        self.model = genai.GenerativeModel(MODEL, system_instruction=SYSTEM_PROMPT)
-
-        # Warm-up: a tiny prompt so first real call isn’t cold
+        self._init_error = None
+        genai.configure(api_key=API_KEY)
+        self.model = genai.GenerativeModel(MODEL_NAME, system_instruction=SYSTEM_PROMPT)
+        # Warm-up (non-fatal if it fails)
         try:
             t0 = time.time()
-            resp = self.model.generate_content("ping", generation_config={"max_output_tokens": 4})
-            print("[handler] warmup ok in", round((time.time()-t0)*1000), "ms; text=",
-                  bool(_extract_text(resp)), flush=True)
+            _ = self.model.generate_content("ping", generation_config={"max_output_tokens": 4})
+            print("[handler:init] warm-up OK in", round((time.time()-t0)*1000), "ms", flush=True)
         except Exception as e:
-            # Non-fatal; we’ll still serve but log it
-            print("[handler] warmup failed:", repr(e), flush=True)
+            print("[handler:init] warm-up failed (non-fatal):", repr(e), flush=True)
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        if self.init_error:
-            return {"text":"", "debug":{"where":"init","error":self.init_error}}
-
-        # Normalize input
-        inputs = data.get("inputs")
-        if isinstance(inputs, str): user_text = inputs.strip()
-        elif isinstance(inputs, dict) and "messages" in inputs:
-            user_text = ""
-            for m in reversed(inputs["messages"] or []):
-                if (m.get("role") or "user").lower() == "user":
-                    user_text = str(m.get("content", "")).strip(); break
-        else: user_text = ""
+        # If init failed, surface that to the client
+        if getattr(self, "_init_error", None):
+            return {"text":"", "debug":{"where":"init", "error": self._init_error}}
+
+        # Parameters override (optional)
+        params = data.get("parameters") or {}
+        gen_cfg = {
+            "temperature": float(params.get("temperature", GEN_CFG["temperature"])),
+            "top_p": float(params.get("top_p", GEN_CFG["top_p"])),
+            "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":"handler","note":"Empty prompt"}}
+            return {"text":"", "debug":{"where":"input", "error":"Empty prompt"}}
 
-        gen_cfg = {"temperature":TEMPERATURE, "top_p":TOP_P, "max_output_tokens":MAX_TOKENS}
+        # Build history if provided (your local shape supported)
+        history = data.get("history") or []
+        gemini_history = _to_gemini_history(history)
 
-        # Retry with backoff (e.g., transient network, cold start)
-        last_exc = None
-        for attempt in range(3):
+        # Try chat w/ a tiny retry to dodge transient cold-start
+        last_err = None
+        for attempt in range(2):
             try:
-                t0 = time.time()
-                resp = self.model.generate_content(user_text, generation_config=gen_cfg)
-                dt = round((time.time()-t0)*1000)
+                if gemini_history:
+                    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, "debug":{"latency_ms": dt, "attempt": attempt+1}}
-                # no text -> include finish reasons if any
+                if txt:
+                    return {"text": txt}
+                # no text — return finish reasons if present
+                fin = []
                 try:
-                    fr = [getattr(c, "finish_reason", None) for c in (resp.candidates or [])]
+                    fin = [getattr(c, "finish_reason", None) for c in (resp.candidates or [])]
                 except Exception:
-                    fr = []
-                return {"text":"", "debug":{"where":"gemini_empty","finish_reasons":fr,"latency_ms":dt}}
-            # network/transient
+                    pass
+                return {"text":"", "debug":{"where":"empty_text", "finish_reasons": fin}}
             except Exception as e:
-                last_exc = repr(e)
-                time.sleep(0.4*(attempt+1))  # backoff
+                last_err = repr(e)
+                time.sleep(0.35)  # small backoff
 
-        return {"text":"", "debug":{"where":"gemini_exception","exception":last_exc}}
+        return {"text":"", "debug":{"where":"exception", "exception": last_err}}