Update handler.py

handler.py

@@ -1,31 +1,52 @@
-import os
+# handler.py — HF-compliant, stateless Gemini proxy
+
 from typing import Any, Dict, List, Union
+import os
 import google.generativeai as genai
 
-DEFAULT_MODEL = os.getenv("GEMINI_MODEL", "models/gemini-2.5-pro")
-DEFAULT_TEMPERATURE = float(os.getenv("TEMPERATURE", "0.7"))
-DEFAULT_TOP_P = float(os.getenv("TOP_P", "0.95"))
-DEFAULT_MAX_OUTPUT = int(os.getenv("MAX_OUTPUT_TOKENS", "1024"))
-DEFAULT_CANDIDATE_COUNT = int(os.getenv("CANDIDATE_COUNT", "1"))
-DEFAULT_SYSTEM_PROMPT = os.getenv("SYSTEM_PROMPT", "You are a helpful assistant.")
-USE_HISTORY = os.getenv("USE_HISTORY", "true").lower() in {"1", "true", "yes"}
+# Config via HF Endpoint → Settings → Environment Variables
+MODEL          = os.getenv("GEMINI_MODEL", "gemini-1.5-flash")   # safe 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", "512"))
+SYSTEM_PROMPT  = os.getenv("SYSTEM_PROMPT", "You are a helpful assistant.")
+
+def _extract_text(resp: Any) -> str:
+    # 1) standard property
+    if getattr(resp, "text", None):
+        return resp.text
+    # 2) candidates/parts
+    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
+    except Exception:
+        pass
+    return ""
+
+def _last_user_from_messages(msgs: List[Dict[str, Any]]) -> str:
+    for m in reversed(msgs or []):
+        if (m.get("role") or "user").lower() == "user":
+            return str(m.get("content", "")).strip()
+    return ""
 
 class EndpointHandler:
-    def __init__(self, model_dir: str, *args, **kwargs):
+    def __init__(self, path: str = ""):
         api_key = os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")
         if not api_key:
-            self._init_error = "Missing GEMINI_API_KEY in Endpoint Environment Variables."
+            self._init_error = "Missing GEMINI_API_KEY/GOOGLE_API_KEY in Endpoint Environment Variables."
+            print("[handler:init] ERROR:", self._init_error, flush=True)
             return
         self._init_error = None
         genai.configure(api_key=api_key)
 
-        # Proper system instruction
-        self.model = genai.GenerativeModel(
-            DEFAULT_MODEL,
-            system_instruction=DEFAULT_SYSTEM_PROMPT
-        )
+        # Proper system prompt
+        self.model = genai.GenerativeModel(MODEL, system_instruction=SYSTEM_PROMPT)
 
-        # Slightly relaxed safety (optional)
+        # Optional: slightly relaxed safety to avoid silent blocks of normal prompts
         self.safety_settings = None
         try:
             from google.generativeai.types import HarmBlockThreshold, HarmCategory
@@ -35,104 +56,64 @@ class EndpointHandler:
                 HarmCategory.HARM_CATEGORY_SEXUAL:      HarmBlockThreshold.BLOCK_ONLY_HIGH,
                 HarmCategory.HARM_CATEGORY_DANGEROUS:   HarmBlockThreshold.BLOCK_ONLY_HIGH,
             }
-        except Exception:
-            pass
-
-    # ---- Helpers ----
-    def _to_gemini_history(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
-        # Only user/model roles are valid for chat history.
-        out = []
-        for m in messages:
-            role = (m.get("role") or "user").lower()
-            if role == "assistant":
-                role = "model"
-            if role not in ("user", "model"):
-                role = "user"
-            out.append({"role": role, "parts": [{"text": str(m.get("content", ""))}]})
-        return out
-
-    def _extract_text(self, resp: Any) -> Dict[str, Any]:
-        # 1) Standard
-        if getattr(resp, "text", None):
-            return {"text": resp.text}
-        # 2) Candidates/parts
-        try:
-            cands = getattr(resp, "candidates", None) or []
-            for c in cands:
-                if getattr(c, "content", None) and getattr(c.content, "parts", None):
-                    for p in c.content.parts:
-                        t = getattr(p, "text", None)
-                        if t:
-                            return {"text": t}
-        except Exception:
-            pass
-        # 3) Diagnostics fallback
-        diag = {}
-        try:
-            if cands:
-                fr = [getattr(c, "finish_reason", None) for c in cands]
-                diag["finish_reasons"] = fr
-        except Exception:
-            pass
-        return {"text": "I couldn’t generate a response.", "debug": diag or {"note": "empty model text"}}
-
-    def _gen_cfg(self, payload: Dict[str, Any]) -> Dict[str, Any]:
-        params = payload.get("parameters") or {}
-        return {
-            "temperature": float(params.get("temperature", DEFAULT_TEMPERATURE)),
-            "top_p": float(params.get("top_p", DEFAULT_TOP_P)),
-            "max_output_tokens": int(params.get("max_output_tokens", DEFAULT_MAX_OUTPUT)),
-            "candidate_count": int(params.get("candidate_count", DEFAULT_CANDIDATE_COUNT)),
-        }
-
-    # ---- Main entry ----
-    def __call__(self, data: Union[Dict[str, Any], List[Dict[str, Any]]]):
+        except Exception as e:
+            print("[handler:init] safety settings skipped:", repr(e), flush=True)
+
+        print(f"[handler:init] OK MODEL={MODEL}", flush=True)
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
         if self._init_error:
-            return {"error": self._init_error}
+            return {"text": "", "debug": {"error": self._init_error}}
+
         try:
-            if isinstance(data, list):
-                return [self._handle_one(d) for d in data]
-            return self._handle_one(data)
+            print("[handler:call] payload:", data, flush=True)
+            # HF guarantees top-level `inputs`
+            inputs = data.get("inputs")
+
+            # Accept either:
+            # A) {"inputs":"plain text"}
+            # B) {"inputs":{"messages":[{"role":"user","content":"..."}]}}
+            # (compat) Also accept top-level "messages" if present.
+            if isinstance(inputs, str):
+                user_text = inputs.strip()
+            elif isinstance(inputs, dict) and "messages" in inputs:
+                user_text = _last_user_from_messages(inputs.get("messages"))
+            elif "messages" in data:
+                user_text = _last_user_from_messages(data.get("messages"))
+            else:
+                user_text = ""
+
+            if not user_text:
+                return {"text": "", "debug": {"note": "Empty prompt."}}
+
+            gen_cfg = {
+                "temperature": TEMPERATURE,
+                "top_p": TOP_P,
+                "max_output_tokens": MAX_TOKENS,
+            }
+
+            print("[handler:call] generate_content:", repr(user_text[:160]), flush=True)
+            resp = self.model.generate_content(
+                user_text,
+                generation_config=gen_cfg,
+                safety_settings=self.safety_settings
+            )
+            print("[handler:call] raw resp:", repr(resp), flush=True)
+
+            text = _extract_text(resp)
+            if text:
+                return {"text": text}
+
+            # Diagnostics if empty
+            debug = {}
+            try:
+                fr = [getattr(c, "finish_reason", None) for c in (resp.candidates or [])]
+                if fr:
+                    debug["finish_reasons"] = fr
+            except Exception:
+                pass
+            return {"text": "", "debug": debug or {"note": "Empty model text"}}
+
         except Exception as e:
-            return {"error": str(e)}
-
-    def _handle_one(self, payload: Dict[str, Any]) -> Dict[str, Any]:
-        cfg = self._gen_cfg(payload)
-        data_inputs = payload.get("inputs")
-
-        # A) Chat: {"inputs":{"messages":[...]}}
-        if isinstance(data_inputs, dict) and "messages" in data_inputs:
-            msgs = data_inputs["messages"] or []
-            # Stateless mode if USE_HISTORY=false
-            if not USE_HISTORY:
-                # use only the last user message
-                last_user = next((m for m in reversed(msgs) if (m.get("role") or "user").lower() == "user"), None)
-                text = (last_user or {}).get("content", "") if last_user else ""
-                resp = self.model.generate_content(text, generation_config=cfg, safety_settings=self.safety_settings)
-                return self._extract_text(resp)
-
-            # With history
-            last_user = next((m for m in reversed(msgs) if (m.get("role") or "user").lower() == "user"), None)
-            user_text = (last_user or {}).get("content", "")
-            history_msgs = [m for m in msgs if m is not last_user]
-            chat = self.model.start_chat(history=self._to_gemini_history(history_msgs))
-            resp = chat.send_message(user_text, generation_config=cfg, safety_settings=self.safety_settings)
-            return self._extract_text(resp)
-
-        # B) Plain text: {"inputs":"..."}
-        if isinstance(data_inputs, str):
-            prompt = data_inputs.strip()
-            if not prompt:
-                return {"text": "Empty prompt."}
-            resp = self.model.generate_content(prompt, generation_config=cfg, safety_settings=self.safety_settings)
-            return self._extract_text(resp)
-
-        # C) Fallbacks
-        if "messages" in (payload or {}):
-            msgs = payload["messages"] or []
-            last_user = next((m for m in reversed(msgs) if (m.get("role") or "user").lower() == "user"), None)
-            text = (last_user or {}).get("content", "")
-            resp = self.model.generate_content(text, generation_config=cfg, safety_settings=self.safety_settings)
-            return self._extract_text(resp)
-
-        return {"text": "Empty prompt."}
+            print("[handler:call] EXC:", repr(e), flush=True)
+            return {"text": "", "debug": {"exception": str(e)}}