diff --git a/tai.py b/tai.py index 90dc6c0..53029b5 100755 --- a/tai.py +++ b/tai.py @@ -881,6 +881,11 @@ def as_int(value, default=0): return default +def tool_name_clean(raw_name): + # Some backends leak harmony control tokens into the name ("shell<|channel|>commentary"). + return str(raw_name or "").split("<|", 1)[0].strip() + + def tool_args_text(raw_arguments): if not raw_arguments: return "{}" @@ -4344,7 +4349,7 @@ class ChatClient: calls = [] for call in message.get("tool_calls") or []: func = call.get("function") or {} - calls.append({"id": call.get("id") or uuid.uuid4().hex, "name": func.get("name") or "", "arguments": tool_args_text(func.get("arguments"))}) + calls.append({"id": call.get("id") or uuid.uuid4().hex, "name": tool_name_clean(func.get("name")), "arguments": tool_args_text(func.get("arguments"))}) return {"role": "assistant", "content": message.get("content") or "", "reasoning": message.get("reasoning") or message.get("reasoning_content") or "", "tool_calls": calls, "backend": label} def streaming(self, url, key, payload, label, sink, extra=None): @@ -4398,6 +4403,7 @@ class ChatClient: response.close() ordered = [calls[pos] for pos in sorted(calls)] for slot in ordered: + slot["name"] = tool_name_clean(slot["name"]) if not slot["id"]: slot["id"] = uuid.uuid4().hex return {"role": "assistant", "content": "".join(content_parts), "reasoning": "".join(reasoning_parts), "tool_calls": ordered, "backend": label} @@ -4509,7 +4515,7 @@ class ChatClient: def muse_reply(self, answer, label): calls = [] for call in answer.get("tool_calls") or []: - calls.append({"id": uuid.uuid4().hex, "name": call.get("name") or "", "arguments": tool_args_text(call.get("arguments"))}) + calls.append({"id": uuid.uuid4().hex, "name": tool_name_clean(call.get("name")), "arguments": tool_args_text(call.get("arguments"))}) return {"role": "assistant", "content": answer.get("content") or "", "reasoning": "", "tool_calls": calls, "backend": label} diff --git a/test_tai.py b/test_tai.py index e22f102..95619e3 100644 --- a/test_tai.py +++ b/test_tai.py @@ -3743,6 +3743,11 @@ class ToolArgumentsTests(unittest.TestCase): self.tmp.cleanup() os.environ.pop("TAI_HOME", None) + def test_tool_name_clean(self): + self.assertEqual(tai.tool_name_clean("shell<|channel|>commentary"), "shell") + self.assertEqual(tai.tool_name_clean(" read_file "), "read_file") + self.assertEqual(tai.tool_name_clean(None), "") + def test_tool_args_text(self): self.assertEqual(tai.tool_args_text(None), "{}") self.assertEqual(tai.tool_args_text(""), "{}")