Update handler.py

handler.py

@@ -135,48 +135,87 @@ class EndpointHandler:
         # best-effort: return canonicalized even if schema still complains
         return obj
 
-    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        # Accept either {"messages":[...], "tools":[...]} or {"inputs": "..."}
-        messages = data.get("messages")
-        tools = data.get("tools") or data.get("functions") or []
-        temperature = float(data.get("temperature", 0.0))
-        max_new = int(data.get("max_new_tokens", 192))
-
+    def _unpack(self, data: Dict[str, Any]):
+        """Normalize payload coming from IE:
+        - accept top-level or inputs-nested messages/tools
+        - accept parameters both top-level and nested
+        """
+        body = data.get("inputs", data)  # if no "inputs", body == data
+        params = data.get("parameters") or {}
+        # pull messages/tools from body if dict
+        messages = None
+        tools = None
+        if isinstance(body, dict):
+            messages = body.get("messages")
+            tools = body.get("tools") or body.get("functions")
+        # allow top-level fallbacks
+        if messages is None:
+            messages = data.get("messages")
+        if tools is None:
+            tools = data.get("tools") or data.get("functions") or []
+
+        # if still no messages, treat body as raw text
         if not messages:
-            text = data.get("inputs") or data.get("text") or ""
-            messages = [{"role": "user", "content": text}]
+            raw = body if isinstance(body, str) else data.get("text", "")
+            messages = [{"role": "user", "content": str(raw)}]
+
+        # generation params (support both locations)
+        temperature = float(params.get("temperature", data.get("temperature", 0.0)))
+        max_new = int(params.get("max_new_tokens", data.get("max_new_tokens", 192)))
+        top_p = float(params.get("top_p", data.get("top_p", 1.0)))
+
+        return messages, tools, temperature, max_new, top_p
+
+    def _encode_messages(self, msgs: List[dict]):
+        # Try chat template; fallback to a simple role-tagged prompt
+        try:
+            return self.tokenizer.apply_chat_template(
+                msgs, add_generation_prompt=True, return_tensors="pt"
+            ).to(self.model.device)
+        except Exception:
+            lines = []
+            for m in msgs:
+                role = m.get("role", "user")
+                content = m.get("content", "")
+                lines.append(f"{role}: {content}")
+            lines.append("assistant:")
+            prompt_text = "\n".join(lines)
+            toks = self.tokenizer(prompt_text, return_tensors="pt")
+            return toks["input_ids"].to(self.model.device)
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        messages, tools, temperature, max_new, top_p = self._unpack(data)
 
-        # Build a fresh system message from tools; prepend to any user/system provided
-        sys_text = tools_to_system_text(tools) if isinstance(tools, list) and tools else None
-        msgs = []
-        if sys_text:
-            msgs.append({"role": "system", "content": sys_text})
-        msgs.extend(messages)
+        # Build a system message from tools; prepend to conversation
+        sys_text = tools_to_system_text(tools) if tools else None
+        msgs = [{"role": "system", "content": sys_text}] + messages if sys_text else messages
 
-        # The last user content is used by the guard to infer heuristics (e.g., "top 5")
+        # Remember last user text for the guard’s heuristics
         user_text = ""
-        for m in msgs:
+        for m in reversed(msgs):
             if m.get("role") == "user":
                 user_text = m.get("content", "")
+                break
 
-        prompt = self.tokenizer.apply_chat_template(
-            msgs, add_generation_prompt=True, return_tensors="pt"
-        ).to(self.model.device)
+        input_ids = self._encode_messages(msgs)
+
+        gen_kwargs = dict(
+            input_ids=input_ids,
+            max_new_tokens=max_new,
+            eos_token_id=self.tokenizer.eos_token_id,
+        )
+        if temperature > 0:
+            gen_kwargs.update(do_sample=True, temperature=temperature, top_p=top_p)
+        else:
+            gen_kwargs.update(do_sample=False)
 
         with torch.inference_mode():
-            out = self.model.generate(
-                input_ids=prompt,
-                max_new_tokens=max_new,
-                do_sample=temperature > 0,
-                temperature=temperature if temperature > 0 else None,
-                eos_token_id=self.tokenizer.eos_token_id,
-            )
-        raw = self.tokenizer.decode(out[0][prompt.shape[-1]:], skip_special_tokens=True).strip()
+            out = self.model.generate(**gen_kwargs)
 
+        raw = self.tokenizer.decode(out[0][input_ids.shape[-1]:], skip_special_tokens=True).strip()
         guarded = self._apply_guard(user_text, tools, raw)
 
-        # Return both for convenience
         return {
-            "generated_text": raw,   # string (for quick cURL)
-            "envelope": guarded      # dict with {"tool_calls":[...]} | {"function_call":...} | {"final_answer":...}
+            "generated_text": raw,
+            "envelope": guarded
         }