shudao-main/shudao-chat-py/tests/test_exam_response_sanitizer.py

import importlib.util
import json
import unittest
from pathlib import Path


CHAT_PATH = Path(__file__).resolve().parents[1] / "routers" / "chat.py"
spec = importlib.util.spec_from_file_location(
    "chat_under_test_exam", CHAT_PATH)
chat = importlib.util.module_from_spec(spec)
spec.loader.exec_module(chat)


def exam_payload(title="桩基础施工技术考核"):
    return {
        "title": title,
        "totalScore": 100,
        "totalQuestions": 1,
        "singleChoice": {
            "scorePerQuestion": 2,
            "totalScore": 2,
            "count": 1,
            "questions": [
                {
                    "text": "钻孔灌注桩清孔完成后应重点检查哪项指标？",
                    "options": [
                        {"key": "A", "text": "孔底沉渣厚度"},
                        {"key": "B", "text": "施工便道宽度"},
                        {"key": "C", "text": "钢筋棚颜色"},
                        {"key": "D", "text": "围挡广告内容"},
                    ],
                    "answer": "A",
                    "analysis": "孔底沉渣厚度直接影响桩端承载力。",
                }
            ],
        },
        "judge": {"scorePerQuestion": 3, "totalScore": 0, "count": 0, "questions": []},
        "multiple": {"scorePerQuestion": 5, "totalScore": 0, "count": 0, "questions": []},
        "short": {"scorePerQuestion": 10, "totalScore": 0, "count": 0, "questions": []},
    }


class ExamResponseSanitizerTests(unittest.TestCase):
    def test_removes_thinking_process_prefix(self):
        raw = "Thinking Process:\n\n1. Analyze the Request.\n\n" + \
            json.dumps(exam_payload(), ensure_ascii=False)

        cleaned = chat._sanitize_exam_response(raw)
        parsed = json.loads(cleaned)

        self.assertEqual(parsed["title"], "桩基础施工技术考核")
        self.assertIn("singleChoice", parsed)
        self.assertNotIn("Thinking Process", cleaned)

    def test_extracts_json_from_markdown_code_block(self):
        raw = "下面是生成结果：\n```json\n" + \
            json.dumps(exam_payload("桥梁考试"), ensure_ascii=False) + "\n```"

        cleaned = chat._sanitize_exam_response(raw)
        parsed = json.loads(cleaned)

        self.assertEqual(parsed["title"], "桥梁考试")

    def test_prefers_exam_payload_over_other_json_noise(self):
        raw = (
            "Thinking Process:\n"
            '{"note":"not exam"}\n'
            "Final Answer:\n"
            + json.dumps(exam_payload("最终试卷"), ensure_ascii=False)
        )

        cleaned = chat._sanitize_exam_response(raw)
        parsed = json.loads(cleaned)

        self.assertEqual(parsed["title"], "最终试卷")
        self.assertIn("singleChoice", parsed)

    def test_extracts_exam_payload_when_reasoning_contains_quotes_and_examples(self):
        raw = (
            'Thinking Process:\n'
            'The output must contain "title", "totalScore", "singleChoice".\n'
            'Use {"key": "A", "text": "..."} as the option shape example.\n'
            'Section example: {"scorePerQuestion": 2, "totalScore": 20, "count": 10, "questions": [...]}.\n\n'
            + json.dumps(exam_payload("带说明的最终试卷"), ensure_ascii=False)
        )

        cleaned = chat._sanitize_exam_response(raw)
        parsed = json.loads(cleaned)

        self.assertEqual(parsed["title"], "带说明的最终试卷")
        self.assertIn("singleChoice", parsed)
        self.assertFalse(cleaned.startswith("Thinking Process"))

    def test_extracts_trailing_exam_json_after_think_suffix(self):
        raw = (
            "Thinking Process:\n"
            'Use {"key": "A", "text": "..."} as example.\n'
            "</think>\n\n"
            + json.dumps(exam_payload("尾部试卷"), ensure_ascii=False)
        )

        cleaned = chat._sanitize_exam_response(raw)
        parsed = json.loads(cleaned)

        self.assertEqual(parsed["title"], "尾部试卷")
        self.assertEqual(parsed["totalQuestions"], 1)

    def test_repairs_unescaped_quotes_inside_string_values(self):
        payload = json.dumps(exam_payload("引号容错"), ensure_ascii=False)
        payload = payload.replace(
            "钻孔灌注桩清孔完成后应重点检查哪项指标？", '钻孔灌注桩必须实行"一炮三检"制度吗？')
        payload = payload.replace("孔底沉渣厚度直接影响桩端承载力。", '"一炮三检"是爆破作业的常见安全检查制度。')

        cleaned = chat._sanitize_exam_response(payload)
        parsed = json.loads(cleaned)

        self.assertEqual(parsed["title"], "引号容错")
        self.assertIn('"一炮三检"', parsed["singleChoice"]["questions"][0]["text"])


if __name__ == "__main__":
    unittest.main()