File size: 1,347 Bytes
95a8a23 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 | import ast
import re
import unittest
from pathlib import Path
def load_sanitizer():
source = Path(__file__).resolve().parents[1] / "ui_server.py"
tree = ast.parse(source.read_text(encoding="utf-8"))
function = next(
node
for node in tree.body
if isinstance(node, ast.FunctionDef) and node.name == "_sanitize_final_answer"
)
namespace = {"re": re}
exec(compile(ast.Module(body=[function], type_ignores=[]), str(source), "exec"), namespace)
return namespace["_sanitize_final_answer"]
sanitize = load_sanitizer()
class OutputSanitizerTests(unittest.TestCase):
def test_rejects_literal_function_call_markup(self):
raw = '''我需要查询数据，让我调用工具。<function_calls>
<invoke name="mcp_marine_fisheries_inventory">
<parameter name="query">IATTC</parameter>
</invoke>
</function_calls>'''
self.assertEqual(sanitize(raw), "")
def test_keeps_final_answer_and_removes_trailing_markup(self):
raw = '''【FINAL】IATTC 数据已确认存在。
<function_calls><invoke name="x"></invoke></function_calls>'''
self.assertEqual(sanitize(raw), "IATTC 数据已确认存在。")
def test_keeps_normal_answer(self):
self.assertEqual(sanitize("【FINAL】正常回答"), "正常回答")
if __name__ == "__main__":
unittest.main()
|