#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""专题十三 · 步骤 2：开启 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,
                "retrieval_model": {
                    "search_method": "semantic_search",
                    "reranking_enable": False,
                    "top_k": 5,
                    "score_threshold_enabled": False,
                },
            },
        )
        pp("hit-testing before rerank", r)
        records = r.get("json")
        if not isinstance(records, list):
            records = (records or {}).get("records") or []
        print("\n--- 前 5 条 ---")
        for i, rec in enumerate(records[:5], 1):
            seg = (rec.get("segment") or {}).get("content") or ""
            print(i, rec.get("score"), seg[:80].replace("\n", " "))
        return 0
    finally:
        d.close()


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