Update handler.py

handler.py

@@ -57,28 +57,34 @@ def tools_to_system_text(tools: List[dict], max_tools=12, max_props=6) -> str:
     return "\n".join(lines)
 
 def build_schema_from_tools(tools: List[dict]) -> dict:
-    # strict per-request schema (name + args schema)
     from copy import deepcopy
     tool_variants, defs = [], {}
-    for t in tools:
-        f = t.get("function", {})
-        name = f.get("name")
-        params = f.get("parameters") or {"type":"object","properties":{},"additionalProperties":True}
-        if not isinstance(name, str): 
+    for t in tools or []:
+        f = t.get("function", t) if isinstance(t, dict) else {}
+        name = f.get("name") or f.get("api_call") or f.get("api_name")
+        if not isinstance(name, str) or not name:
             continue
+        params = f.get("parameters") or {"type":"object","properties":{},"additionalProperties":True}
+        if isinstance(params, list):  # allow list form
+            props = {}
+            for p in params:
+                if isinstance(p, dict) and "name" in p:
+                    nm = p["name"]; pd = {k:v for k,v in p.items() if k!="name"}
+                    props[nm] = pd
+            if props:
+                params = {"type":"object","properties":props}
         defs[f"{name}_args"] = deepcopy(params)
         tool_variants.append({
-            "type": "object",
-            "properties": {"name": {"const": name}, "arguments": {"$ref": f"#/$defs/{name}_args"}},
-            "required": ["name","arguments"],
+            "type":"object",
+            "properties":{"name":{"const":name},"arguments":{"$ref": f"#/$defs/{name}_args"}},
+            "required":["name","arguments"],
             "additionalProperties": False
         })
+
     return {
-        "$schema": "https://json-schema.org/draft/2020-12/schema",
-        "oneOf": [
-            {"type":"object","properties":{"function_call":{"oneOf": tool_variants}},
-             "required":["function_call"],"additionalProperties": False},
-            {"type":"object","properties":{"tool_calls":{"type":"array","minItems":1,"items":{"oneOf": tool_variants}}},
+        "$schema":"https://json-schema.org/draft/2020-12/schema",
+        "oneOf":[
+            {"type":"object","properties":{"tool_calls":{"type":"array","minItems":1,"items":{"oneOf":tool_variants}}},
              "required":["tool_calls"],"additionalProperties": False},
             {"type":"object","properties":{"final_answer":{"type":"string","minLength":1}},
              "required":["final_answer"],"additionalProperties": False}
@@ -86,6 +92,7 @@ def build_schema_from_tools(tools: List[dict]) -> dict:
         "$defs": defs
     }
 
+
 class EndpointHandler:
     def __init__(self, path: str = ""):
         model_id = path or os.getenv("MODEL_ID", ".")
@@ -166,22 +173,48 @@ class EndpointHandler:
 
         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)
+SYS_INSTR = (
+    "You're a tool-calling assistant. "
+    "Return ONLY valid JSON for your answer, with this exact shape:\n"
+    "{\"tool_calls\": [{\"name\": \"<function_name>\", \"arguments\": {<key>: <value>, ...}}, ...]}\n"
+    "No prose. No explanations. JSON only."
+)
+
+def _flat_tool(self, t: dict):
+    """Accept both {'function':{...}} and flat {'name':...,'parameters':...}."""
+    f = t.get("function", t) if isinstance(t, dict) else {}
+    name = f.get("name") or f.get("api_call") or f.get("api_name") or ""
+    params = f.get("parameters") or {}
+    # Normalize parameter names for signature display
+    prop_names = []
+    if isinstance(params, dict):
+        props = params.get("properties")
+        if isinstance(props, dict):
+            prop_names = list(props.keys())[:12]
+        elif isinstance(props, list):
+            prop_names = [p.get("name","") for p in props if isinstance(p,dict)][:12]
+    return name, params, prop_names
+
+def _render_tools_signature(self, tools: List[dict]) -> str:
+    lines = []
+    for t in tools[:12]:
+        name, _, pnames = self._flat_tool(t)
+        if not name: 
+            continue
+        lines.append(f"- {name}({', '.join(pnames)})" if pnames else f"- {name}()")
+    return "\n".join(lines) if lines else "- (tools omitted)"
+
+def _encode_messages(self, user_text: str, tools: List[dict]):
+    # Build the exact same prompt you used for training
+    sig = self._render_tools_signature(tools)
+    prompt = (
+        "<|system|>\n" + SYS_INSTR + "\n\n"
+        "<|tools|>\n" + sig + "\n\n"
+        "<|user|>\n" + user_text + "\n\n"
+        "<|assistant|>\n"
+    )
+    toks = self.tokenizer(prompt, 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)
@@ -192,12 +225,12 @@ class EndpointHandler:
 
         # Remember last user text for the guard’s heuristics
         user_text = ""
-        for m in reversed(msgs):
+        for m in reversed(messages):
             if m.get("role") == "user":
                 user_text = m.get("content", "")
                 break
 
-        input_ids = self._encode_messages(msgs)
+        input_ids = self._encode_messages(user_text, tools)
 
         gen_kwargs = dict(
             input_ids=input_ids,