#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""专题十五 · 步骤 5.1：PATCH 知识库 retrieval_model（命中测试用）。"""
from __future__ import annotations

import sys

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


def build_retrieval() -> dict:
    return {
        "search_method": "semantic_search",
        "reranking_enable": True,
        "reranking_mode": "reranking_model",
        "reranking_model": {
            "reranking_provider_name": C.XINFERENCE_PROVIDER,
            "reranking_model_name": C.RERANK_MODEL,
        },
        "top_k": 5,
        "score_threshold_enabled": False,
        "score_threshold": None,
    }


def main() -> int:
    d = Dify(C.DIFY_EMAIL, C.DIFY_PASSWORD)
    try:
        d.login()
        r = d.req("PATCH", "/datasets/%s" % C.DS_ID, body={"retrieval_model": build_retrieval()})
        print("HTTP", r.get("code"))
        j = r.get("json") or {}
        print("retrieval_model_dict:", j.get("retrieval_model_dict"))
        pp("patch dataset", r)
        return 0 if r.get("code") == "200" else 1
    finally:
        d.close()


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