#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""专题十三 · 步骤 8：开启 Rerank 后的命中测试（修改后）。"""
from __future__ import annotations

import sys

import _bootstrap  # noqa: F401
import _config as C
from dify_api import Dify, pp


QUERY = "法定假日加班的制度是什么，加班费怎么算？"


def main() -> int:
    d = Dify(C.DIFY_EMAIL, C.DIFY_PASSWORD)
    try:
        d.login()
        r = d.post("/datasets/%s/hit-testing" % C.DS_ID, body={"query": QUERY})
        pp("hit-testing after rerank", r)
        records = r.get("json")
        if not isinstance(records, list):
            records = (records or {}).get("records") or []
        print("\n--- 前 5 条（分数为 Rerank relevance_score）---")
        for i, rec in enumerate(records[:5], 1):
            seg = (rec.get("segment") or {}).get("content") or ""
            print(i, rec.get("score"), seg[:120].replace("\n", " "))
        return 0
    finally:
        d.close()


if __name__ == "__main__":
    sys.exit(main())