Update app.py

app.py

@@ -6,6 +6,8 @@ import pandas as pd
 import time
 from cnocr import CnOcr
 import numpy as np
+import openai
+from llama_index import GPTVectorStoreIndex, SimpleDirectoryReader, Prompt
 
 ocr = CnOcr()  # 初始化ocr模型
 history_max_len = 500  # 机器人记忆的最大长度
@@ -13,22 +15,12 @@ all_max_len = 2000  # 输入的最大长度
 
 
 def get_text_emb(open_ai_key, text):
-    url = 'https://api.openai.com/v1/embeddings'
-    headers = {
-        'Content-Type': 'application/json',
-        'Authorization': 'Bearer ' + open_ai_key
-    }
-    data = {
-        "model": "text-embedding-ada-002",
-        "input": text
-    }
-    result = requests.post(url=url,
-                           data=json.dumps(data),
-                           headers=headers
-                           )
-    if result.status_code != 200:
-        raise Exception('API请求出错，状态码为：' + str(result.status_code) + '，错误信息为：' + result.json())
-    return result.json()['data'][0]['embedding']
+    openai.api_key = open_ai_key
+    response = openai.Embedding.create(
+        input=text,
+        model="text-embedding-ada-002"
+    )
+    return response['data'][0]['embedding']
 
 
 def doc_index_self(open_ai_key, doc):  # 文档向量化
@@ -37,10 +29,10 @@ def doc_index_self(open_ai_key, doc):  # 文档向量化
     for text in texts:
         emb_list.append(get_text_emb(open_ai_key, text))
     return texts, emb_list, gr.Textbox.update(visible=True), gr.Button.update(visible=True), gr.Markdown.update(
-        value="""操作说明 step 3：PDF解析提交成功！ 🙋 可以开始对话啦~"""), gr.Chatbot.update(visible=True)
+        value="""操作说明 step 3：PDF解析提交成功！ 🙋 可以开始对话啦~"""), gr.Chatbot.update(visible=True), 1
 
 
-def get_response(open_ai_key, msg, bot, doc_text_list, doc_embeddings):  # 获取机器人回复
+def get_response_by_self(open_ai_key, msg, bot, doc_text_list, doc_embeddings):  # 获取机器人回复
     now_len = len(msg)  # 当前输入的长度
     his_bg = -1  # 历史记录的起始位置
     for i in range(len(bot) - 1, -1, -1):  # 从后往前遍历历史记录
@@ -96,29 +88,43 @@ def get_response(open_ai_key, msg, bot, doc_text_list, doc_embeddings):  # 获
         messages.append({"role": "user", "content": his[0]})  # 加入用户的历史记录
         messages.append({"role": "assistant", "content": his[1]})  # 加入机器人的历史记录
     messages.append({"role": "user", "content": msg})  # 加入用户的当前输入
+    openai.api_key = open_ai_key
+    chat_completion = openai.ChatCompletion.create(model="gpt-3.5-turbo", messages=messages)  # 获取机器人的回复
+    res = chat_completion.choices[0].message.content  # 获取机器人的回复
+    bot.append([msg, res])  # 加入历史记录
+    return bot[max(0, len(bot) - 3):]  # 返回最近3轮的历史记录
+
 
-    url = 'https://api.openai.com/v1/chat/completions'
-
-    data = {
-        "model": "gpt-3.5-turbo",
-        "messages": messages
-    }
-    print("data = \n", data)
-
-    headers = {
-        'Content-Type': 'application/json',
-        'Authorization': 'Bearer ' + open_ai_key
-    }
-    result = requests.post(url=url,
-                           data=json.dumps(data),
-                           headers=headers
-                           )
-    print("result = \n", result.json())
-    res = str(result.json()['choices'][0]['message']['content']).strip()
+def get_response_by_llama_index(open_ai_key, msg, bot, query_engine):  # 获取机器人回复
+    openai.api_key = open_ai_key
+    template = (
+        "你是一个有用的助手，可以使用文章内容准确地回答问题。使用提供的文章来生成你的答案，但避免逐字复制文章。尽可能使用自己的话。准确、有用、简洁、清晰。文章内容如下： \n"
+        "---------------------\n"
+        "{context_str}"
+        "\n---------------------\n"
+        "{query_str}\n"
+        "请基于文章内容回答用户的问题。\n"
+    )  # 定义模板
+    query_str = "历史对话如下：\n"
+    for his in bot:  # 遍历历史记录
+        query_str += "用户：" + his[0] + "\n"  # 加入用户的历史记录
+        query_str += "机器人：" + his[1] + "\n"  # 加入机器人的历史记录
+    query_str += "用户：" + msg + "\n"  # 加入用户的当前输入
+    qa_template = Prompt(template)  # 将模板转换成Prompt对象
+    query_engine = query_engine.as_query_engine(text_qa_template=qa_template)  # 建立查询引擎
+    res = query_engine.query(msg)  # 获取回答
+    print(res)  # 显示回答
     bot.append([msg, res])  # 加入历史记录
     return bot[max(0, len(bot) - 3):]  # 返回最近3轮的历史记录
 
 
+def get_response(open_ai_key, msg, bot, doc_text_list, doc_embeddings, query_engine, index_type):  # 获取机器人回复
+    if index_type == 1:
+        return get_response_by_self(open_ai_key, msg, bot, doc_text_list, doc_embeddings)
+    else:
+        return get_response_by_llama_index(open_ai_key, msg, bot, query_engine)
+
+
 def up_file(files):  # 上传文件
     doc_text_list = []  # 用于存储文档
     for idx, file in enumerate(files):  # 遍历文件
@@ -155,10 +161,26 @@ def up_file(files):  # 上传文件
     doc_text_list = [str(text).strip() for text in doc_text_list if len(str(text).strip()) > 0]  # 去除空格
     print(doc_text_list)
     return gr.Textbox.update(value='\n'.join(doc_text_list), visible=True), gr.Button.update(
+        visible=True), gr.Button.update(
         visible=True), gr.Markdown.update(
         value="操作说明 step 2：确认PDF解析结果（可修正），点击“建立索引”，随后进行对话")
 
 
+def doc_index_llama(open_ai_key, txt):  # 建立索引
+    # 根据时间戳新建目录，保存txt文件
+    path = str(time.time())
+    import os
+    os.mkdir(path)
+    with open(path + '/doc.txt', mode='w', encoding='utf-8') as f:
+        f.write(txt)
+    openai.api_key = open_ai_key  # 设置OpenAI API Key
+    documents = SimpleDirectoryReader(path).load_data()  # 读取文档
+    index = GPTVectorStoreIndex.from_documents(documents)  # 建立索引
+    query_engine = index.as_query_engine()  # 建立查询引擎
+    return query_engine, gr.Textbox.update(visible=True), gr.Button.update(visible=True), gr.Markdown.update(
+        value="""操作说明 step 3：PDF解析提交成功！ 🙋 可以开始对话啦~"""), gr.Chatbot.update(visible=True), 0
+
+
 with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
@@ -166,21 +188,27 @@ with gr.Blocks() as demo:
             file = gr.File(file_types=['.pdf'], label='点击上传PDF，进行解析(支持多文档、表格、OCR)',
                            file_count='multiple')  # 支持多文档、表格、OCR
             txt = gr.Textbox(label='PDF解析结果', visible=False)  # PDF解析结果
-            index_self_bu = gr.Button(value='建立索引（by self）', visible=False)  #
-            index_llama_bu = gr.Button(value='建立索引（by llama_index）', visible=False)  #
+            with gr.Row():
+                index_llama_bu = gr.Button(value='建立索引（by llama_index）', visible=False)  # 建立索引（by llama_index）
+                index_self_bu = gr.Button(value='建立索引（by self）', visible=False)  # 建立索引（by self）
             doc_text_state = gr.State([])  # 存储PDF解析结果
             doc_emb_state = gr.State([])  # 存储PDF解析结果的embedding
+            query_engine = gr.State([])  # 存储查询引擎
+            index_type = gr.State([])  # 存储索引类型
         with gr.Column():
             md = gr.Markdown("""操作说明 step 1：点击左侧区域，上传PDF，进行解析""")  # 操作说明
             chat_bot = gr.Chatbot(visible=False)  # 聊天机器人
             msg_txt = gr.Textbox(label='消息框', placeholder='输入消息，点击发送', visible=False)  # 消息框
-            with gr.Row():
-                chat_bu = gr.Button(value='发送', visible=False)  # 发送按钮
+            chat_bu = gr.Button(value='发送', visible=False)  # 发送按钮
 
-    file.change(up_file, [file], [txt, index_self_bu, md])  # 上传文件
+    file.change(up_file, [file], [txt, index_self_bu, index_llama_bu, md])  # 上传文件
     index_self_bu.click(doc_index_self, [open_ai_key, txt],
-                        [doc_text_state, doc_emb_state, msg_txt, chat_bu, md, chat_bot])  # 提交解析结果
-    chat_bu.click(get_response, [open_ai_key, msg_txt, chat_bot, doc_text_state, doc_emb_state], [chat_bot])  # 发送消息
+                        [doc_text_state, doc_emb_state, msg_txt, chat_bu, md, chat_bot, index_type])  # 提交解析结果
+    index_llama_bu.click(doc_index_llama, [open_ai_key, txt],
+                         [query_engine, msg_txt, chat_bu, md, chat_bot, index_type])  # 提交解析结果
+    chat_bu.click(get_response,
+                  [open_ai_key, msg_txt, chat_bot, doc_text_state, doc_emb_state, query_engine, index_type],
+                  [chat_bot])  # 发送消息
 
 if __name__ == "__main__":
     demo.queue().launch()