# 导入必要的库和模块
import gradio as gr
import json
import random
from dotenv import load_dotenv
load_dotenv()  # 加载环境变量
from zhipuai import ZhipuAI  # 智谱AI的Python客户端
import os

# 设置文件路径和API密钥
file_path = 'tianji/prompt/yiyan_prompt/all_yiyan_prompt.json'
API_KEY = os.environ['ZHIPUAI_API_KEY']

# 读取包含不同场景提示词和示例对话的JSON文件
with open(file_path, 'r', encoding='utf-8') as file:
    json_data = json.load(file)

# 定义获取系统提示词的函数
def get_system_prompt_by_name(name):
    # ...

# 定义更改示例对话的函数
def change_example(name, cls_choose_value, chatbot):
    # ...

# 定义随机选择场景的函数
def random_button_click(chatbot):
    # ...

# 定义更改场景选择的函数
def cls_choose_change(idx):
    # ...

# 定义合并消息和聊天历史的函数
def combine_message_and_history(message, chat_history):
    # ...

# 定义生成回复的函数
def respond(system_prompt, message, chat_history):
    # ...

# 定义清除聊天历史的函数
def clear_history(chat_history):
    # ...

# 定义重新生成回复的函数
def regenerate(chat_history, system_prompt):
    # ...

# 使用Gradio创建Web界面
with gr.Blocks() as demo:
    # 定义界面状态
    chat_history = gr.State()
    now_json_data = gr.State(value=_get_id_json_id(0))
    now_name = gr.State()

    # 定义界面标题和描述
    gr.Markdown(TITLE)

    # 定义界面组件：单选按钮、下拉菜单、文本框、按钮等
    cls_choose = gr.Radio(...)
    input_example = gr.Dataset(...)
    dorpdown_name = gr.Dropdown(...)
    system_prompt = gr.TextArea(...)
    chatbot = gr.Chatbot(...)
    msg = gr.Textbox(...)
    submit = gr.Button(...)
    clear = gr.Button(...)
    regenerate = gr.Button(...)

    # 定义界面组件的布局
    with gr.Row():
        # ...

# 为界面组件设置事件处理函数
cls_choose.change(fn=cls_choose_change, inputs=cls_choose, outputs=[now_json_data, dorpdown_name])
dorpdown_name.change(fn=change_example, inputs=[dorpdown_name, now_json_data, chatbot], outputs=input_example)
input_example.click(fn=example_click, inputs=[input_example, dorpdown_name, now_json_data], outputs=[msg, system_prompt])
random_button.click(fn=random_button_click, inputs=chatbot, outputs=[cls_choose, now_json_data, dorpdown_name])

# 运行应用程序，用户可以通过界面进行交互
if __name__ == "__main__":
    demo.launch()

import os
import logging

import docx
import argparse

def argsParser():
    parser = argparse.ArgumentParser(
        description="该脚本能够将原始 .txt/.docx 转化为 .txt数据"
        "例如 `path`=liyi/ "
        "|-- liyi"
        "   |-- jingjiu"
        "       |-- *.txt"
        "       |-- ....."
        "   |-- songli"
        "       |-- *.docx"
        "       |-- ....."
        "将在 liyi/datasets 下生成处理后的 .txt 文件"
        "例如：python process_data.py \ "
        "--path liyi/"
    )
    parser.add_argument("--path", type=str, help="原始数据集目录")
    args = parser.parse_args()
    return args

log = logging.getLogger("myLogger")
log.setLevel(logging.DEBUG)

BASIC_FORMAT = "%(asctime)s %(levelname)-8s %(message)s"
formatter = logging.Formatter(BASIC_FORMAT)

chlr = logging.StreamHandler()  # console
chlr.setLevel(logging.DEBUG)
chlr.setFormatter(formatter)

log.addHandler(chlr)

def parser_docx(path):
    file = docx.Document(path)
    out = ""
    for para in file.paragraphs:
        text = para.text
        if text != "":
            out = out + text + "\n"
    return out

def parser_txt(path):
    out = ""
    with open(path, "r") as f:
        for line in f:
            line = line.strip()
            if line != "":
                out = out + line + "\n"
    return out

if __name__ == "__main__":
    ARGS = argsParser()
    ori_data_path = ARGS.path

    data_dict = {}
    for sub_dir_name in os.listdir(ori_data_path):
        sub_dir_path = os.path.join(ori_data_path, sub_dir_name)
        data_dict.setdefault(sub_dir_path, {})
        samples = {}

        for sub_file_name in os.listdir(sub_dir_path):
            file_path = os.path.join(sub_dir_path, sub_file_name)

            sorted(file_path, reverse=True)
            if file_path.endswith(".docx"):
                samples.setdefault("docx", [])
                samples["docx"].append(sub_file_name)
            elif file_path.endswith(".txt"):
                samples.setdefault("txt", [])
                samples["txt"].append(sub_file_name)

        data_dict[sub_dir_path].setdefault("samples", samples)

    for datax, obj in data_dict.items():
        if "samples" in obj.keys():
            samples = obj["samples"]
            if "docx" in samples.keys():
                file_list = samples["docx"]
                file_list = sorted(
                    file_list, key=lambda file_path: int(file_path.split("-")[1][1:])
                )
                obj["samples"]["docx"] = file_list
            data_dict[datax] = obj

    docx_list = []
    txt_list = []
    for datax, obj in data_dict.items():
        if "samples" in obj.keys():
            samples = obj["samples"]
            if "docx" in samples.keys():
                docx_list.extend(os.path.join(datax, x) for x in samples["docx"])

            if "txt" in samples.keys():
                txt_list.extend(os.path.join(datax, x) for x in samples["txt"])

    data_dir = os.path.join(ori_data_path, "datasets")
    if not os.path.exists(data_dir):
        os.makedirs(data_dir)

    for ind, file in enumerate(docx_list):
        out_text = parser_docx(file)
        with open(os.path.join(data_dir, f"docx_{ind}.txt"), "w") as f:
            f.write(out_text)

    for ind, file in enumerate(txt_list):
        out_text = parser_txt(file)
        with open(os.path.join(data_dir, f"txt_{ind}.txt"), "w") as f:
            f.write(out_text)

# from metagpt.const import METAGPT_ROOT as TIANJI_PATH
class RQA_ST_Liyi_Chroma_Config:
    """
    检索问答增强（RQA）配置文件：
    基于Chroma检索数据库；
    基于Sentence-Transformer词向量模型构建的外挂礼仪（Liyi）知识库。
    """

    # 原始数据位置 online 设置为空
    ORIGIN_DATA = ""
    # 持久化数据库位置，例如 chroma/liyi/
    PERSIST_DIRECTORY = ""
    # Sentence-Transformer词向量模型权重位置
    HF_SENTENCE_TRANSFORMER_WEIGHT = (
        "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
    )

from langchain.document_loaders import DirectoryLoader, TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma

from . import RQA_ST_Liyi_Chroma_Config

if __name__ == "__main__":
    persist_directory = RQA_ST_Liyi_Chroma_Config.PERSIST_DIRECTORY
    data_directory = RQA_ST_Liyi_Chroma_Config.ORIGIN_DATA
    loader = DirectoryLoader(data_directory, glob="*.txt", loader_cls=TextLoader)

    text_splitter = RecursiveCharacterTextSplitter(chunk_size=3000, chunk_overlap=150)
    split_docs = text_splitter.split_documents(loader.load())

    embeddings = HuggingFaceEmbeddings(
        model_name="/root/weights/model/sentence-transformer"
    )
    vectordb = Chroma.from_documents(
        documents=split_docs, embedding=embeddings, persist_directory=persist_directory
    )
    vectordb.persist()

from langchain.llms.base import LLM
from typing import Any, List, Optional
from langchain.callbacks.manager import CallbackManagerForLLMRun
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import os

class InternLM_LLM(LLM):
    tokenizer: AutoTokenizer = None
    model: AutoModelForCausalLM = None

    def __init__(self, model_path: str):
        super().__init__()
        print("正在从本地加载模型...")
        self.tokenizer = AutoTokenizer.from_pretrained(
            model_path, trust_remote_code=True
        )
        self.model = (
            AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True)
            .to(torch.bfloat16)
            .cuda()
        )
        self.model = self.model.eval()
        print("完成本地模型的加载")

    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any
    ):
        system_prompt = """你是一名AI助手名为天机（SocialAI），也可称为来事儿AI。它能够处理中国传统人情世故的任务，例如如何敬酒、如何说好话、如何会来事儿等。
        """
        messages = [(system_prompt, "")]
        response, history = self.model.chat(self.tokenizer, prompt, history=messages)
        return response

    @property
    def _llm_type(self) -> str:
        return "InternLM"

class Zhipu_LLM(LLM):
    tokenizer: AutoTokenizer = None
    model: AutoModelForCausalLM = None
    client: Any = None

    def __init__(self):
        super().__init__()
        from zhipuai import ZhipuAI

        print("初始化模型...")
        self.client = ZhipuAI(api_key=os.environ.get("zhupuai_key"))
        print("完成模型初始化")

    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any
    ):
        system_prompt = """你是一名AI助手名为天机（SocialAI），也可称为来事儿AI。它能够处理中国传统人情世故的任务，例如如何敬酒、如何说好话、如何会来事儿等。
        你是一个信息抽取的知识库语料准备能手，你需要把我给你的文章做成几个知识点，这个知识点类似问答对的回答（陈述句的描述,不需要提问，比如：苹果是一种水果，可以吃和烹饪，而且是红色的，长在大树上），你不需要分1、2、3、4点， 只需要把相关的知识都划分成一个段落就好， ```The example is as follows, assuming I publish this article first: There is a very important etiquette in business banquets. If you ignore it, your customers will feel that you are very disrespectful. Everyone knows that discussing business at the dinner table often has a greater chance of success than sitting down in the office. Of course, wine is inseparable from this, so the etiquette of toasting in business banquets is very important. When toasting, pour the wine for the other party first, and then pour the wine for yourself. Hold the wine glass in your right hand and hold the bottom of the glass in your left hand. Our wine glass is a little lower than the other person's. If the other person is more modest and puts it lower than us, we can gently hold up the other person's wine glass with our left hand, which will show respect. In order to express our sincerity after drinking, we can use the gesture of toasting to tilt the cup slightly, with the mouth of the cup facing the other person. There is no need to turn the cup over directly, which will look very inelegant. What else should you pay attention to when toasting? We can leave a message to discuss it together. Your answer is a knowledgeable and calm reply, as follows as a whole: In business banquets, following etiquette is crucial to leaving a good impression on customers, and business negotiations at the dinner table usually have a higher success rate. During the toasting process, you should first pour the wine for the other person, and then pour the wine for yourself. The glass should be held with the right hand and the bottom with the left hand. When toasting, the wine glass should be kept lower than the other person's wine glass to show respect; if the other person's wine glass is lower, you can gently hold up the other person's wine glass with your left hand. After drinking, you should tilt the cup slightly with the mouth of the cup facing the other party in a toasting gesture. Avoid turning the cup over directly to maintain politeness and grace.``` 接下来你帮我解析新的知识，你只需要回复这个新的知识文章相关的内容就好，不要回复例子的内容！文章如下： ```Do you know the drinking process and toasting rhythm of a formal dinner? If you don’t know this video, be sure to like it and save it, because sooner or later you will use it for your last business drinking party. There are generally six stages, and most people are most likely to make mistakes in the second and fifth stages. Next, let’s talk about each one separately. First, the amount of common wine to drink at the beginning of the first stage of this drinking game depends on the home team. People in Shandong are very particular about taking turns to drink and toast together, while most places in Hebei are accustomed to drinking the first three drinks together. Different places have different rules. You can also leave a message to see what the rules are in your place. If everyone likes and follows enthusiastically enough, I can publish a special episode summarizing drinking customs across the country in the future. This second stage is when the host starts to toast. At this time, it is usually the host or host who takes the lead in toasting to each guest in turn, starting with the guest of honor. At this stage, the awareness of sequence and turn is very important. If you are a guest, do not return the favor to express your gratitude at this time, because it is not your time to show up yet. In the third stage, you can return the favor as a guest. You can take the lead and lead everyone to respond together first, and then respond separately. Then enter the fourth stage, drink the theme wine and the key wine, and highlight the theme based on the relationship between the lover and the theme. After drinking, the people at the table will understand why they are drinking this wine. Hey, this fifth stage is the free drinking stage. If you have a bad temper with someone, you can go over and have a drink with them. If you have unresolved issues with someone, you can use wine to ask for advice. If you don't like someone, you can also use wine to challenge them. Especially when you come with a mission, you must seize the time to implement the mission, because after this stage, you will no longer be free. In the sixth stage, that is, the whole house is finally full and it is almost time to disperse. The host will usually speak, and everyone will drink in front of the door, and the whole house will be filled with wine. Drinking this glass of wine means that the drinking party is officially over, and the following program can be eaten or vomited. In business banquets, following etiquette is crucial to leaving a good impression on customers. Business negotiations at the dinner table usually have a higher success rate. During the toasting process, you should first pour the wine for the other person, and then pour the wine for yourself. The glass should be held with the right hand and the bottom with the left hand. When toasting, the wine glass should be kept lower than the other person's wine glass to show respect; if the other person's wine glass is lower, you can gently hold up the other person's wine glass with your left hand. After drinking, you should tilt the cup slightly with the mouth of the cup facing the other party in a toasting gesture. Avoid turning the cup over directly to maintain politeness and grace.```
        """
        response = self.client.chat.completions.create(
            model="glm-4",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": prompt},
            ],
        )
        return response.choices[0].message.content

    @property
    def _llm_type(self) -> str:
        return "ZhipuLM"


class OpenAI_LLM(LLM):
    tokenizer: AutoTokenizer = None
    model: AutoModelForCausalLM = None
    client: Any = None

    def __init__(self, base_url="https://api.deepseek.com/v1"):
        super().__init__()
        from openai import OpenAI

        print("初始化模型...")
        self.client = OpenAI(
            api_key=os.environ.get("openai_key", None), base_url=base_url
        )
        print("完成模型初始化")

    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any
    ):
        system_prompt = """你是一名AI助手名为天机（SocialAI），也可称为来事儿AI。它能够处理中国传统人情世故的任务，例如如何敬酒、如何说好话、如何会来事儿等。
        你是一个信息抽取的知识库语料准备能手，你需要把我给你的文章做成几个知识点，这个知识点类似问答对的回答（陈述句的描述,不需要提问，比如：苹果是一种水果，可以吃和烹饪，而且是红色的，长在大树上），你不需要分1、2、3、4点， 只需要把相关的知识都划分成一个段落就好， ```The example is as follows, assuming I publish this article first: There is a very important etiquette in business banquets. If you ignore it, your customers will feel that you are very disrespectful. Everyone knows that discussing business at the dinner table often has a greater chance of success than sitting down in the office. Of course, wine is inseparable from this, so the etiquette of toasting in business banquets is very important. When toasting, pour the wine for the other party first, and then pour the wine for yourself. Hold the wine glass in your right hand and hold the bottom of the glass in your left hand. Our wine glass is a little lower than the other person's. If the other person is more modest and puts it lower than us, we can gently hold up the other person's wine glass with our left hand, which will show respect. In order to express our sincerity after drinking, we can use the gesture of toasting to tilt the cup slightly, with the mouth of the cup facing the other person. There is no need to turn the cup over directly, which will look very inelegant. What else should you pay attention to when toasting? We can leave a message to discuss it together. Your answer is a knowledgeable and calm reply, as follows as a whole: In business banquets, following etiquette is crucial to leaving a good impression on customers, and business negotiations at the dinner table usually have a higher success rate. During the toasting process, you should first pour the wine for the other person, and then pour the wine for yourself. The glass should be held with the right hand and the bottom with the left hand. When toasting, the wine glass should be kept lower than the other person's wine glass to show respect; if the other person's wine glass is lower, you can gently hold up the other person's wine glass with your left hand. After drinking, you should tilt the cup slightly with the mouth of the cup facing the other party in a toasting gesture. Avoid turning the cup over directly to maintain politeness and grace.``` 接下来你帮我解析新的知识，你只需要回复这个新的知识文章相关的内容就好，不要回复例子的内容！文章如下： ```Do you know the drinking process and toasting rhythm of a formal dinner? If you don’t know this video, be sure to like it and save it, because sooner or later you will use it for your last business drinking party. There are generally six stages, and most people are most likely to make mistakes in the second and fifth stages. Next, let’s talk about each one separately. First, the amount of common wine to drink at the beginning of the first stage of this drinking game depends on the home team. People in Shandong are very particular about taking turns to drink and toast together, while most places in Hebei are accustomed to drinking the first three drinks together. Different places have different rules. You can also leave a message to see what the rules are in your place. If everyone likes and follows enthusiastically enough, I can publish a special episode summarizing drinking customs across the country in the future. This second stage is when the host starts to toast. At this time, it is usually the host or host who takes the lead in toasting to each guest in turn, starting with the guest of honor. At this stage, the awareness of sequence and turn is very important. If you are a guest, do not return the favor to express your gratitude at this time, because it is not your time to show up yet. In the third stage, you can return the favor as a guest. You can take the lead and lead everyone to respond together first, and then respond separately. Then enter the fourth stage, drink the theme wine and the key wine, and highlight the theme based on the relationship between the lover and the theme. After drinking, the people at the table will understand why they are drinking this wine. Hey, this fifth stage is the free drinking stage. If you have a bad temper with someone, you can go over and have a drink with them. If you have unresolved issues with someone, you can use wine to ask for advice. If you don't like someone, you can also use wine to challenge them. Especially when you come with a mission, you must seize the time to implement the mission, because after this stage, you will no longer be free. In the sixth stage, that is, the whole house is finally full and it is almost time to disperse. The host will usually speak, and everyone will drink in front of the door, and the whole house will be filled with wine. Drinking this glass of wine means that the drinking party is officially over, and the following program can be eaten or vomited. In business banquets, following etiquette is crucial to leaving a good impression on customers. Business negotiations at the dinner table usually have a higher success rate. During the toasting process, you should first pour the wine for the other person, and then pour the wine for yourself. The glass should be held with the right hand and the bottom with the left hand. When toasting, the wine glass should be kept lower than the other person's wine glass to show respect; if the other person's wine glass is lower, you can gently hold up the other person's wine glass with your left hand. After drinking, you should tilt the cup slightly with the mouth of the cup facing the other party in a toasting gesture. Avoid turning the cup over directly to maintain politeness and grace.```
        """
        response = self.client.chat.completions.create(
            model="glm-4",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": prompt},
            ],
        )
        return response.choices[0].message.content

    @property
    def _llm_type(self) -> str:
        return "OpenAILM"

import tianji.utils.knowledge_tool as knowledgetool
from tianji.agents.knowledges.config import AGENT_KNOWLEDGE_PATH, AGENT_EMBEDDING_PATH
from dotenv import load_dotenv

load_dotenv()

# KNOWLEDGE_PATH = r"D:\1-wsl\TIANJI\Tianji\tianji\knowledges\04-Wishes\knowledges.txt"
# SAVE_PATH = r"D:\1-wsl\TIANJI\Tianji\temp"

# doclist = knowledgetool.get_docs_list_query_openai(query_str="春节",loader_file_path=KNOWLEDGE_PATH, \
#                                   persist_directory = SAVE_PATH,k_num=5)

doclist = knowledgetool.get_docs_list_query_zhipuai(
    query_str="春节",
    loader_file_path=AGENT_KNOWLEDGE_PATH.WISHES.path(),
    persist_directory=AGENT_EMBEDDING_PATH.WISHES.path(filename="zhipuai"),
    k_num=5,
)

if doclist is not []:
    print(doclist)
else:
    print("doclist is [] !")

from tianji.knowledges.RAG.demo import model_center

if __name__ == "__main__":
    model = model_center()
    question = "如何给长辈敬酒？"
    chat_history = []
    _, response = model.qa_chain_self_answer(question, chat_history)
    print(response)