MapleGu/ChuanhuChatGPT

1import os2import logging3 4from llama_index import GPTSimpleVectorIndex5from llama_index import download_loader6from llama_index import (7    Document,8    LLMPredictor,9    PromptHelper,10    QuestionAnswerPrompt,11    RefinePrompt,12)13from langchain.llms import OpenAI14import colorama15 16 17from presets import *18from utils import *19 20 21def get_documents(file_src):22    documents = []23    index_name = ""24    logging.debug("Loading documents...")25    logging.debug(f"file_src: {file_src}")26    for file in file_src:27        logging.debug(f"file: {file.name}")28        index_name += file.name29        if os.path.splitext(file.name)[1] == ".pdf":30            logging.debug("Loading PDF...")31            CJKPDFReader = download_loader("CJKPDFReader")32            loader = CJKPDFReader()33            documents += loader.load_data(file=file.name)34        elif os.path.splitext(file.name)[1] == ".docx":35            logging.debug("Loading DOCX...")36            DocxReader = download_loader("DocxReader")37            loader = DocxReader()38            documents += loader.load_data(file=file.name)39        elif os.path.splitext(file.name)[1] == ".epub":40            logging.debug("Loading EPUB...")41            EpubReader = download_loader("EpubReader")42            loader = EpubReader()43            documents += loader.load_data(file=file.name)44        else:45            logging.debug("Loading text file...")46            with open(file.name, "r", encoding="utf-8") as f:47                text = add_space(f.read())48                documents += [Document(text)]49    index_name = sha1sum(index_name)50    return documents, index_name51 52 53def construct_index(54    api_key,55    file_src,56    max_input_size=4096,57    num_outputs=1,58    max_chunk_overlap=20,59    chunk_size_limit=600,60    embedding_limit=None,61    separator=" ",62    num_children=10,63    max_keywords_per_chunk=10,64):65    os.environ["OPENAI_API_KEY"] = api_key66    chunk_size_limit = None if chunk_size_limit == 0 else chunk_size_limit67    embedding_limit = None if embedding_limit == 0 else embedding_limit68    separator = " " if separator == "" else separator69 70    llm_predictor = LLMPredictor(71        llm=OpenAI(model_name="gpt-3.5-turbo-0301", openai_api_key=api_key)72    )73    prompt_helper = PromptHelper(74        max_input_size,75        num_outputs,76        max_chunk_overlap,77        embedding_limit,78        chunk_size_limit,79        separator=separator,80    )81    documents, index_name = get_documents(file_src)82    if os.path.exists(f"./index/{index_name}.json"):83        logging.info("找到了缓存的索引文件，加载中……")84        return GPTSimpleVectorIndex.load_from_disk(f"./index/{index_name}.json")85    else:86        try:87            logging.debug("构建索引中……")88            index = GPTSimpleVectorIndex(89                documents, llm_predictor=llm_predictor, prompt_helper=prompt_helper90            )91            os.makedirs("./index", exist_ok=True)92            index.save_to_disk(f"./index/{index_name}.json")93            return index94        except Exception as e:95            print(e)96            return None97 98 99def chat_ai(100    api_key,101    index,102    question,103    context,104    chatbot,105):106    os.environ["OPENAI_API_KEY"] = api_key107 108    logging.info(f"Question: {question}")109 110    response, chatbot_display, status_text = ask_ai(111        api_key,112        index,113        question,114        replace_today(PROMPT_TEMPLATE),115        REFINE_TEMPLATE,116        SIM_K,117        INDEX_QUERY_TEMPRATURE,118        context,119    )120    if response is None:121        status_text = "查询失败，请换个问法试试"122        return context, chatbot123    response = response124 125    context.append({"role": "user", "content": question})126    context.append({"role": "assistant", "content": response})127    chatbot.append((question, chatbot_display))128 129    os.environ["OPENAI_API_KEY"] = ""130    return context, chatbot, status_text131 132 133def ask_ai(134    api_key,135    index,136    question,137    prompt_tmpl,138    refine_tmpl,139    sim_k=1,140    temprature=0,141    prefix_messages=[],142):143    os.environ["OPENAI_API_KEY"] = api_key144 145    logging.debug("Index file found")146    logging.debug("Querying index...")147    llm_predictor = LLMPredictor(148        llm=OpenAI(149            temperature=temprature,150            model_name="gpt-3.5-turbo-0301",151            prefix_messages=prefix_messages,152        )153    )154 155    response = None  # Initialize response variable to avoid UnboundLocalError156    qa_prompt = QuestionAnswerPrompt(prompt_tmpl)157    rf_prompt = RefinePrompt(refine_tmpl)158    response = index.query(159        question,160        llm_predictor=llm_predictor,161        similarity_top_k=sim_k,162        text_qa_template=qa_prompt,163        refine_template=rf_prompt,164        response_mode="compact",165    )166 167    if response is not None:168        logging.info(f"Response: {response}")169        ret_text = response.response170        nodes = []171        for index, node in enumerate(response.source_nodes):172            brief = node.source_text[:25].replace("\n", "")173            nodes.append(174                f"<details><summary>[{index+1}]\t{brief}...</summary><p>{node.source_text}</p></details>"175            )176        new_response = ret_text + "\n----------\n" + "\n\n".join(nodes)177        logging.info(178            f"Response: {colorama.Fore.BLUE}{ret_text}{colorama.Style.RESET_ALL}"179        )180        os.environ["OPENAI_API_KEY"] = ""181        return ret_text, new_response, f"查询消耗了{llm_predictor.last_token_usage} tokens"182    else:183        logging.warning("No response found, returning None")184        os.environ["OPENAI_API_KEY"] = ""185        return None186 187 188def add_space(text):189    punctuations = {"，": "， ", "。": "。 ", "？": "？ ", "！": "！ ", "：": "： ", "；": "； "}190    for cn_punc, en_punc in punctuations.items():191        text = text.replace(cn_punc, en_punc)192    return text193