tests/test_followup_routing.py

import ast
import re
import unittest
from pathlib import Path


def load_helpers():
    source = Path(__file__).resolve().parents[1] / "ui_server.py"
    tree = ast.parse(source.read_text(encoding="utf-8"))
    names = {
        "_is_confirmation_prompt",
        "_looks_like_data_parameter",
        "_is_data_request_followup",
        "_is_ocean_export_request",
        "_apply_ocean_export_defaults",
    }
    functions = [
        node for node in tree.body
        if isinstance(node, ast.FunctionDef) and node.name in names
    ]
    namespace = {
        "re": re,
        "_DATA_REQUEST_FOLLOWUP_ACTIONS": (
            "修改","调整","改动","变更","补充","补全","补上","增加","新增","添加","加上",
            "更换","换成","换为","改为","改成","改到","删掉","删除","去掉","移除","取消",
            "缩小","扩大","重新下载","重新导出","重下","重导","重试","改一下","再下载",
            "再导出","只要","仅要","只需","把","请把","改",
        ),
        "_DATA_REQUEST_FOLLOWUP_REFERENCES": (
            "刚才","之前","上一","上次","原来","上面","以上","这个","这些","该任务",
            "那个","此任务","下载","导出","结果","文件","任务","它",
        ),
        "_needs_ocean_mcp": lambda prompt: any(
            term in str(prompt).lower()
            for term in ("era5", "ocean", "sst", "cmems", "oisst")
        ),
    }
    exec(
        compile(ast.Module(body=functions, type_ignores=[]), str(source), "exec"),
        namespace,
    )
    return namespace


helpers = load_helpers()


class FollowupRoutingTests(unittest.TestCase):
    def test_short_confirmation_variants(self):
        confirm = helpers["_is_confirmation_prompt"]
        for value in ("确认", "好的", "继续", "OK", "yes"):
            self.assertTrue(confirm(value))
        self.assertFalse(confirm("确认一下ERA5数据范围"))

    def test_short_amendment_followups_are_detected(self):
        followup = helpers["_is_data_request_followup"]
        for value in (
            "把范围改成 120E–140E",
            "补充：2002年8月",
            "改成 2000年1月到2000年6月",
            "换成 CHL",
            "重新下载",
            "请调整刚才的变量，只要SST",
            "区域改为西北太平洋",
            "再补一个 csv 格式",
        ):
            self.assertTrue(followup(value), value)

    def test_new_questions_are_not_amendment_followups(self):
        followup = helpers["_is_data_request_followup"]
        for value in (
            "介绍一下柔鱼",
            "查询有哪些 Ocean 数据",
            "你好",
            "谢谢",
        ):
            self.assertFalse(followup(value), value)

    def test_missing_format_defaults_to_netcdf(self):
        apply_defaults = helpers["_apply_ocean_export_defaults"]
        prompt = "导出1998年1月4日ERA5 v10数据"
        routed = apply_defaults(prompt)
        self.assertIn("format=netcdf", routed)
        self.assertIn("mcp_marine_marine_export", routed)

    def test_explicit_format_is_preserved(self):
        apply_defaults = helpers["_apply_ocean_export_defaults"]
        prompt = "导出1998年1月4日ERA5 v10数据，格式csv"
        routed = apply_defaults(prompt)
        self.assertIn(prompt, routed)
        self.assertNotIn("format=netcdf", routed)
        self.assertIn("mcp_marine_marine_export", routed)


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