import re
import requests
import json


def get_knowledge(query):
    # 定义请求的URL
    url = 'http://127.0.0.1:7861/chat/knowledge_base_chat'

    # 定义请求头
    headers = {
        'accept': 'application/json',
        'Content-Type': 'application/json'
    }

    # 定义请求的数据
    data = {
        "query": query,
        "knowledge_base_name": "My_cx",
        "top_k": 1,
        "score_threshold": 0.7,
        "history": [],
        "stream": False,
        "model_name": "fm9g_2b_hf_models",
        "temperature": 0.8,
        "max_tokens": 1,
        "prompt_name": "empty"
    }

    json_data = json.dumps(data)

    # 发送POST请求
    try:
        response = requests.post(url, headers=headers, data=json_data)
        response.raise_for_status()  # 确保请求成功

        # 尝试解析JSON
        if response.text:
            start = response.text.find('{')
            if start != -1:
                json_text = response.text[start:]
                response_json = json.loads(json_text)

                docs = response_json.get("docs", [])

                if docs and not any("未找到相关文档" in doc for doc in docs):
                    cleaned_docs = "<已知信息>\n"
                    for doc in docs:
                        # cleaned_doc = re.sub(r'出处 \[.*?\] \[.*?\]\(.*?\)', '', doc)
                        cleaned_doc = re.sub(r'\(.*?\)', '', doc)
                        cleaned_docs += cleaned_doc
                    cleaned_docs += "</已知信息>"
                    return cleaned_docs
                else:
                    return None

    except requests.exceptions.RequestException as e:
        # 如果请求失败，返回错误信息
        return f"请求失败: {e}"


if __name__ == '__main__':
    print(get_knowledge("队列"))
    print(get_knowledge("海湾战争"))