final.py

import os
os.environ["OPENAI_API_KEY"] = "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"
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.document_loaders import PyPDFLoader
from langchain.vectorstores import Chroma
from pypinyin import lazy_pinyin
import gradio as gr
import openai
import random
# import logging
# logging.basicConfig(
#     filename='log/log.log',
#     level=logging.INFO,
#     format='%(asctime)s - %(levelname)s - %(message)s',
#     datefmt='%m/%d/%Y %H:%M:%S'
# )
embedding = OpenAIEmbeddings()
target_files = set()
topics = ["农业", "宗教与文化", "建筑业与制造业", "医疗卫生保健", "国家治理", "法律法规", "财政税收", "教育", "金融", "贸易", "宏观经济", "社会发展", "科学技术", "能源环保", "国际关系", "国防安全","不限主题"]

def get_path(target_string):
    folder_path = "./vector_data"
    all_vectors = os.listdir(folder_path)
    matching_files = [file for file in all_vectors if file.startswith(target_string)]
    for file in matching_files:
        file_path = os.path.join(folder_path, file)
        return file_path
    return ""


def extract_partial_message(res_message, response):
    for chunk in response:
        if len(chunk["choices"][0]["delta"]) != 0:
            res_message = res_message + chunk["choices"][0]["delta"]["content"]
            yield res_message


def format_messages(sys_prompt, history, message):
    history_openai_format = [{"role": "system", "content": sys_prompt}]
    for human, assistant in history:
        history_openai_format.append({"role": "user", "content": human})
        history_openai_format.append({"role": "assistant", "content": assistant})
    history_openai_format.append({"role": "user", "content": message})
    return history_openai_format

def get_domain(history, message):
    sys_prompt = """
    帮我根据用户的问题划分到以下几个类别,输出最匹配的一个类别：[宗教与文化, 农业, 建筑业与制造业, 医疗卫生保健, 国家治理, 法律法规, 财政税收, 教育, 金融, 贸易, 宏观经济, 社会发展, 科学技术, 能源环保, 国际关系, 国防安全]
    """
    history_openai_format = format_messages(sys_prompt, history, message)
    print("history_openai_format:", history_openai_format)
    # logging.info(f"history_openai_format: {history_openai_format}")
    response = openai.ChatCompletion.create(model="gpt-4", messages=history_openai_format, temperature=1.0, stream=False)
    domain = response['choices'][0]['message']['content']
    print("匹配领域：", domain)
    # logging.info(f"匹配领域: {domain}")
    return domain

def echo(message, history, flag1, flag2):
    global target_files, topics
    print("flag1:", flag1)
    print("flag2:", flag2) 
    print("history:", history)
    print("message:", message)
    # logging.info(f"flag1: {flag1}")
    # logging.info(f"flag2: {flag2}")
    # logging.info(f"history: {history}")
    # logging.info(f"message: {message}")
    if len(flag1) == 0:  # 不进行研报问答&研报生成
        target_files.clear()
        history.clear()
        if flag2 not in [None, 16]:
            domain = topics[flag2]
            message = f"{domain}领域相关内容"
        elif flag2 in [None, 16]:
            message = f"{topics[random.randint(0, 16)]}领域相关内容" if message == "" else message
            domain = get_domain(history, message)
        persist_vector_path = get_path("".join(lazy_pinyin(domain)))
        
        db = Chroma(persist_directory=persist_vector_path, embedding_function=embedding)
        docs = db.similarity_search_with_score(query=message, k=5)
        contents = [doc[0] for doc in docs]
        relevance = "  ".join(doc.page_content for doc in contents)
        source = [doc.metadata for doc in contents]
        for item in source:
            target_files.add(item['source'] )
        print("研报搜索结果：", target_files)
        # logging.info(f"研报搜索结果: {target_files}")

        sys_prompt = """
        你是一个研报助手，根据这篇文章：{}
        来回复用户的问题生成总结，你需要严格按照这种格式回复：以上文章总结了*,主要观点是*。你只能回复中文。
        如果这篇文章无法回答用户的问题，你必须根据你的知识面直接编造内容进行回答，不能回复不知道，不能说这段文字没有提供等话语。
        """
        sys_prompt = sys_prompt.format(relevance)
        history_openai_format = format_messages(sys_prompt, history, message)
        print("history_openai_format:", history_openai_format)
        # logging.info(f"history_openai_format: {history_openai_format}")
        response = openai.ChatCompletion.create(model="gpt-4", messages=history_openai_format, temperature=1.0, stream=True)
        partial_message = "搜索结果来自以下研报:\n" + '\n'.join(i for i in target_files) + '\n\n'
        for result_message in extract_partial_message(partial_message, response):
            yield result_message
        

    elif flag1 == ['研报问答']:
        print("target_files:", target_files)
        # logging.info(f"target_files: {target_files}")
        QA_pages = []
        if not target_files:
            yield "请取消选中研报问答，先进行研报检索，再进行问答。"
        else:
            for item in target_files:
                loader = PyPDFLoader(item)
                QA_pages.extend(loader.load_and_split())
            text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
            documents = text_splitter.split_documents(QA_pages)
            db = Chroma.from_documents(documents, OpenAIEmbeddings())

            docs = db.similarity_search_with_score(query=message, k=3)
            contents = [doc[0] for doc in docs]
            relevance = "  ".join(doc.page_content for doc in contents)

            sys_prompt = """
            你是一个研报助手，根据这篇文章：{}
            来回复用户的问题，如果这篇文章无法回答用户的问题，你必须根据你的知识面来编造进行专业的回答，
            不能回复不知道，不能回复这篇文章不能回答的这种话语，你只能回复中文。
            """
            sys_prompt = sys_prompt.format(relevance)
            history_openai_format = format_messages(sys_prompt, history, message)
            print("history_openai_format:", history_openai_format)
            # logging.info(f"history_openai_format: {history_openai_format}")
            response = openai.ChatCompletion.create(model="gpt-4", messages=history_openai_format, temperature=1.0, stream=True)
            for result_message in extract_partial_message("", response):
                yield result_message

    elif flag1 == ['研报生成']:
        target_files.clear()
        sys_prompt = """
        你是一个研报助手，请根据用户的要求回复问题。
        """
        history_openai_format = format_messages(sys_prompt, history, message)
        print("history_openai_format:", history_openai_format)
        # logging.info(f"history_openai_format: {history_openai_format}")
        response = openai.ChatCompletion.create(model="gpt-4", messages=history_openai_format, temperature=1.0, stream=True)
        for result_message in extract_partial_message("", response):
            yield result_message

    elif len(flag1) == 2:
        yield "请选中一个选项，进行相关问答。"


demo = gr.ChatInterface(
        echo, 
        chatbot=gr.Chatbot(height=430, label="ChatReport"),
        textbox=gr.Textbox(placeholder="请输入问题", container=False, scale=7),
        title="研报助手",
        description="清芬院研报助手",
        theme="soft",
        additional_inputs=[
            # gr.Radio(["研报问答", "研报生成"], type="index", label = "function"),
            # gr.Checkbox(label = "研报问答"),
            # gr.Checkbox(label = "研报生成"),
            gr.CheckboxGroup(["研报问答", "研报生成"], label="Function"),
            gr.Dropdown(topics, type="index"),
            # gr.Button(value="Run").click(echo, inputs=["", "", [], None], outputs=[""])
            # btn.click(combine, inputs=[txt, txt_2], outputs=[txt_3])
            # gr.Blocks()
        ],
        # retry_btn="retry",
        undo_btn="清空输入框",
        clear_btn="清空聊天记录"
        ).queue()


if __name__ == "__main__":
    demo.launch(share=True)
    





'''
flag1: ['研报问答']
flag2: None
history: []
message: gg
target_files: set()
'''