Upload app.py

app.py

@@ -1,4 +1,5 @@
 import json
+import re
 import time
 import uuid
 from typing import Optional
@@ -84,56 +85,136 @@ def list_models() -> str:
     return json.dumps(result)
 
 
+# ---------------------------------------------------------------------------
+# Tool call parsing (Hermes-style: <tool_call>{...}</tool_call>)
+# ---------------------------------------------------------------------------
+
+def _parse_tool_calls(text: str):
+    """
+    Detect and extract Hermes-style tool calls from model output.
+    Returns (tool_calls, remaining_content) where tool_calls is a list in
+    OpenAI format, or (None, text) if no tool calls are found.
+    """
+    pattern = r'<tool_call>(.*?)</tool_call>'
+    matches = re.findall(pattern, text, re.DOTALL)
+    if not matches:
+        return None, text
+
+    tool_calls = []
+    for match in matches:
+        try:
+            call = json.loads(match.strip())
+            tool_calls.append({
+                "id": f"call_{uuid.uuid4().hex[:24]}",
+                "type": "function",
+                "function": {
+                    "name": call.get("name", ""),
+                    "arguments": json.dumps(call.get("arguments", call.get("parameters", {}))),
+                },
+            })
+        except json.JSONDecodeError:
+            continue
+
+    if not tool_calls:
+        return None, text
+
+    # Strip tool call blocks and think tags from remaining content
+    remaining = re.sub(pattern, '', text, flags=re.DOTALL)
+    remaining = re.sub(r'<think>.*?</think>', '', remaining, flags=re.DOTALL).strip()
+    return tool_calls, remaining or None
+
+
 def chat_completions(
     messages_json: str,
     max_tokens: int = 512,
     temperature: float = 0.7,
     top_p: float = 0.9,
+    tools_json: str = "",
 ) -> str:
     """
     Non-streaming chat completions. Returns an OpenAI-compatible JSON string.
 
-    messages_json: JSON array of {role, content} objects,
-                   e.g. '[{"role":"user","content":"Hello"}]'
+    messages_json: JSON array of {role, content} objects
+    tools_json:    JSON array of OpenAI-format tool definitions (optional)
 
-    NOTE: Qwen3-14B supports thinking mode. Set enable_thinking=True in the
-          chat template call if you want chain-of-thought reasoning.
+    NOTE: Qwen3-14B supports thinking mode. enable_thinking is set to False
+          here for reliable tool call formatting.
     """
     try:
         messages = json.loads(messages_json)
     except json.JSONDecodeError as e:
         return json.dumps({"error": f"Invalid messages_json: {e}"})
 
+    tools = None
+    if tools_json:
+        try:
+            tools = json.loads(tools_json)
+        except json.JSONDecodeError:
+            pass
+
     try:
-        hf_messages = [{"role": m["role"], "content": m["content"]} for m in messages]
-        prompt = tokenizer.apply_chat_template(
-            hf_messages,
+        hf_messages = []
+        for m in messages:
+            role = m["role"]
+            # tool results come in as role=tool; map to role=tool with tool_call_id
+            if role == "tool":
+                hf_messages.append({
+                    "role": "tool",
+                    "content": m.get("content", ""),
+                    "tool_call_id": m.get("tool_call_id", ""),
+                })
+            elif role == "assistant" and m.get("tool_calls"):
+                hf_messages.append({
+                    "role": "assistant",
+                    "content": m.get("content") or "",
+                    "tool_calls": m["tool_calls"],
+                })
+            else:
+                hf_messages.append({"role": role, "content": m.get("content", "")})
+
+        template_kwargs = dict(
             tokenize=False,
             add_generation_prompt=True,
+            enable_thinking=False,
         )
+        if tools:
+            template_kwargs["tools"] = tools
+
+        prompt = tokenizer.apply_chat_template(hf_messages, **template_kwargs)
     except Exception as e:
         return json.dumps({"error": f"Prompt build failed: {e}"})
 
     gen_kwargs = dict(
         max_new_tokens=max_tokens,
-        temperature=temperature,
+        temperature=max(temperature, 0.01),
         top_p=top_p,
         do_sample=True,
         pad_token_id=tokenizer.eos_token_id,
     )
 
     try:
-        content = _generate_response(prompt, gen_kwargs)
+        raw = _generate_response(prompt, gen_kwargs)
     except Exception as e:
         return json.dumps({"error": f"Generation failed: {e}"})
 
     cid = f"chatcmpl-{uuid.uuid4().hex}"
+    tool_calls, content = _parse_tool_calls(raw)
+
+    if tool_calls:
+        message = {"role": "assistant", "content": content, "tool_calls": tool_calls}
+        finish_reason = "tool_calls"
+    else:
+        # Strip any stray think tags from plain responses
+        content = re.sub(r'<think>.*?</think>', '', raw, flags=re.DOTALL).strip()
+        message = {"role": "assistant", "content": content}
+        finish_reason = "stop"
+
     result = {
         "id": cid,
         "object": "chat.completion",
         "created": int(time.time()),
         "model": MODEL_ALIAS,
-        "choices": [{"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"}],
+        "choices": [{"index": 0, "message": message, "finish_reason": finish_reason}],
         "usage": {"prompt_tokens": -1, "completion_tokens": -1, "total_tokens": -1},
     }
     return json.dumps(result)
@@ -185,11 +266,12 @@ You can also chat directly below.
         _cc_max_tokens = gr.Number(label="max_tokens",   value=512)
         _cc_temp       = gr.Number(label="temperature",  value=0.7)
         _cc_top_p      = gr.Number(label="top_p",        value=0.9)
+        _cc_tools      = gr.Textbox(label="tools_json",  value="")
         _cc_out        = gr.Textbox(label="result")
         _cc_btn        = gr.Button("chat_completions")
         _cc_btn.click(
             fn=chat_completions,
-            inputs=[_cc_messages, _cc_max_tokens, _cc_temp, _cc_top_p],
+            inputs=[_cc_messages, _cc_max_tokens, _cc_temp, _cc_top_p, _cc_tools],
             outputs=[_cc_out],
             api_name="chat_completions",
         )