update LightZero RAG

.gitignore

@@ -1,2 +0,0 @@
-.env
-*bkp.py

README.md

@@ -56,7 +56,6 @@ QUESTION_LANG='cn' # The language of the question, currently available option is
 
 ```python
 
-# The difference between rag_demo.py and rag_demo_v0.py is that it can output the retrieved document chunks.
 if __name__ == "__main__":
     # Assuming documents are already present locally
     file_path = './documents/LightZero_README.zh.md'
@@ -91,9 +90,9 @@ if __name__ == "__main__":
 ```
 RAG/
 │
-├── rag_demo_v0.py         # RAG demonstration script without support for outputting retrieved document chunks.
 ├── rag_demo.py            # RAG demonstration script with support for outputting retrieved document chunks.
-├── app.py                 # Web-based interactive application built with Gradio and rag_demo.py.
+├── app_qa.py              # Web-based interactive application built with Gradio and rag_demo.py.
+├── app_mqa.py             # Web-based interactive application built with Gradio and rag_demo.py. Supports maintaining conversation history.
 ├── .env                   # Environment variable configuration file
 └── documents/             # Documents folder
     └── your_document.txt  # Context document
@@ -114,4 +113,4 @@ If you encounter any issues or require assistance, please submit a problem throu
 
 ## License
 
-All code in this repository is compliant with [Apache License 2.0](https://www.apache.org/licenses/LICENSE-2.0).
+All code in this repository is compliant with [Apache License 2.0](https://www.apache.org/licenses/LICENSE-2.0).

README_zh.md

@@ -43,7 +43,6 @@ QUESTION_LANG='cn' # 问题语言，目前可选值为 'cn'
 
 ```python
 
-# rag_demo.py 相对 rag_demo_v0.py 的不同之处在于可以输出检索到的文档块。
 if __name__ == "__main__":
     # 假设文档已存在于本地
     file_path = './documents/LightZero_README.zh.md'
@@ -78,9 +77,9 @@ if __name__ == "__main__":
 ```
 RAG/
 │
-├── rag_demo_v0.py         # RAG 演示脚本，不支持输出检索到的文档块。
 ├── rag_demo.py            # RAG 演示脚本，支持输出检索到的文档块。
-├── app.py                 # 基于 Gradio 和 rag_demo.py 构建的网页交互式应用。
+├── app_qa.py              # 基于 Gradio 和 rag_demo.py 构建的网页交互式应用。
+├── app_mqa.py             # 基于 Gradio 和 rag_demo.py 构建的网页交互式应用。支持保持对话历史。
 ├── .env                   # 环境变量配置文件
 └── documents/             # 文档文件夹
     └── your_document.txt  # 上下文文档

app.py

@@ -1,16 +1,3 @@
-"""
-这段代码的整体功能是创建一个Gradio应用，用户可以在其中输入问题，应用会使用Retrieval-Augmented Generation (RAG)模型来寻找答案并将结果显示在界面上。
-其中，检索到的上下文会在Markdown文档中高亮显示，帮助用户理解答案的来源。应用界面分为两部分：顶部是问答区，底部展示了RAG模型参考的上下文。
-
-结构概述：
-- 导入必要的库和函数。
-- 设置环境变量和全局变量。
-- 加载和处理Markdown文档。
-- 定义处理用户问题并返回答案和高亮显示上下文的函数。
-- 使用Gradio构建用户界面，包括Markdown、输入框、按钮和输出框。
-- 启动Gradio应用并设置为可以分享。
-"""
-
 import os
 
 import gradio as gr
@@ -22,7 +9,6 @@ from rag_demo import load_and_split_document, create_vector_store, setup_rag_cha
 # 环境设置
 load_dotenv()  # 加载环境变量
 QUESTION_LANG = os.getenv("QUESTION_LANG")  # 从环境变量获取 QUESTION_LANG
-
 assert QUESTION_LANG in ['cn', 'en'], QUESTION_LANG
 
 if QUESTION_LANG == "cn":
@@ -31,8 +17,8 @@ if QUESTION_LANG == "cn":
     <div align="center">
         <img src="https://raw.githubusercontent.com/puyuan1996/RAG/main/assets/banner.svg" width="80%" height="20%" alt="Banner Image">
     </div>
-    <h2 style="text-align: center; color: black;"><a href="https://github.com/puyuan1996/RAG"> 🎭LightZero RAG Demo</a></h2>
-    <h4 align="center"> 📢说明：请您在下面的"问题"框中输入任何关于 LightZero 的问题，然后点击"提交"按钮。右侧"回答"框中会显示 RAG 模型给出的回答。在QA栏的下方会给出参考文档（检索得到的 context 用黄色高亮显示）。</h4>
+    <h2 style="text-align: center; color: black;"><a href="https://github.com/puyuan1996/RAG"> LightZero RAG Demo</a></h2>
+    <h4 align="center"> 📢说明：请您在下面的"问题（Q）"框中输入任何关于 LightZero 的问题，然后点击"提交"按钮。右侧"回答（A）"框中会显示 RAG 模型给出的回答。在 QA 栏的下方会给出参考文档（其中检索得到的相关文段会用黄色高亮显示）。</h4>
     <h4 align="center"> 如果你喜欢这个项目，请给我们在 GitHub 点个 star ✨ 。我们将会持续保持更新。  </h4>
     <strong><h5 align="center">注意：算法模型的输出可能包含一定的随机性。相关结果不代表任何开发者和相关 AI 服务的态度和意见。本项目开发者不对生成结果作任何保证，仅供参考。<h5></strong>
     """
@@ -47,55 +33,101 @@ if QUESTION_LANG == "cn":
 
 # 路径变量，方便之后的文件使用
 file_path = './documents/LightZero_README.zh.md'
-chunks = load_and_split_document(file_path)
-retriever = create_vector_store(chunks)
-# rag_chain = setup_rag_chain(model_name="gpt-4")
-rag_chain = setup_rag_chain(model_name="gpt-3.5-turbo")
 
 # 加载原始Markdown文档
 loader = TextLoader(file_path)
 orig_documents = loader.load()
 
+# 存储对话历史
+conversation_history = []
+
 
-def rag_answer(question):
-    retrieved_documents, answer = execute_query(retriever, rag_chain, question)
-    # Highlight the context in the document
-    context = [retrieved_documents[i].page_content for i in range(len(retrieved_documents))]
-    highlighted_document = orig_documents[0].page_content
-    for i in range(len(context)):
-        highlighted_document = highlighted_document.replace(context[i], f"<mark>{context[i]}</mark>")
+def rag_answer(question, model_name, temperature, embedding_model, k):
+    """
+    处理用户问题并返回答案和高亮显示的上下文
+
+    :param question: 用户输入的问题
+    :param model_name: 使用的语言模型名称
+    :param temperature: 生成答案时使用的温度参数
+    :param embedding_model: 使用的嵌入模型
+    :param k: 检索到的文档块数量
+    :return: 模型生成的答案和高亮显示上下文的Markdown文本
+    """
+    try:
+        chunks = load_and_split_document(file_path, chunk_size=5000, chunk_overlap=500)
+        retriever = create_vector_store(chunks, model=embedding_model, k=k)
+        rag_chain = setup_rag_chain(model_name=model_name, temperature=temperature)
+
+        # 将问题添加到对话历史中
+        conversation_history.append(("User", question))
+
+        # 将对话历史转换为字符串
+        history_str = "\n".join([f"{role}: {text}" for role, text in conversation_history])
+
+        retrieved_documents, answer = execute_query(retriever, rag_chain, history_str, model_name=model_name,
+                                                    temperature=temperature)
+        # 在文档中高亮显示上下文
+        context = [retrieved_documents[i].page_content for i in range(len(retrieved_documents))]
+        highlighted_document = orig_documents[0].page_content
+        for i in range(len(context)):
+            highlighted_document = highlighted_document.replace(context[i], f"<mark>{context[i]}</mark>")
+
+        # 将回答添加到���话历史中
+        conversation_history.append(("Assistant", answer))
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        return "处理您的问题时出现错误，请稍后再试。", ""
     return answer, highlighted_document
 
-"""
-在下面的代码中，gr.Blocks构建了Gradio的界面布局，gr.Textbox用于创建文本输入框，gr.Button创建了一个按钮，gr.Markdown则用于显示Markdown格式的内容。
-gr_submit.click是一个事件处理器，当用户点击提交按钮时，它会调用rag_answer函数，并将输入和输出的组件关联起来。
-代码中的rag_answer函数负责接收用户的问题，使用RAG模型检索和生成答案，并将检索到的文本段落在Markdown原文中高亮显示。
-该函数返回模型生成的答案和高亮显示上下文的Markdown文本。
-"""
-with gr.Blocks(title=title, theme='ParityError/Interstellar') as rag_demo:
-    gr.Markdown(title_markdown)
-
-    with gr.Row():
-        with gr.Column():
-            inputs = gr.Textbox(
-                placeholder="请您输入任何关于 LightZero 的问题。",
-                label="问题 (Q)")  # 设置输出框，包括答案和高亮显示参考文档
-            gr_submit = gr.Button('提交')
-
-        outputs_answer = gr.Textbox(placeholder="当你点击提交按钮后，这里会显示 RAG 模型给出的回答。",
-                                    label="回答 (A)")
-    with gr.Row():
-        # placeholder="当你点击提交按钮后，这里会显示参考的文档，其中检索得到的与问题最相关的 context 用高亮显示。"
-        outputs_context = gr.Markdown(label="参考的文档，检索得到的 context 用高亮显示 (C)")
-
-    gr.Markdown(tos_markdown)
-
-    gr_submit.click(
-        rag_answer,
-        inputs=inputs,
-        outputs=[outputs_answer, outputs_context],
-    )
+
+def clear_context():
+    """
+    清除对话历史
+    """
+    global conversation_history
+    conversation_history = []
+    return "", ""
+
 
 if __name__ == "__main__":
-    # 启动界面，设置为可以分享。如果分享公网链接失败，可以在本地执行 ngrok http 7860 将本地端口映射到公网
-    rag_demo.launch(share=True)
+    with gr.Blocks(title=title, theme='ParityError/Interstellar') as rag_demo:
+        gr.Markdown(title_markdown)
+
+        with gr.Row():
+            with gr.Column():
+                inputs = gr.Textbox(
+                    placeholder="请您输入任何关于 LightZero 的问题。",
+                    label="问题 (Q)")
+                model_name = gr.Dropdown(
+                    choices=['kimi', 'abab6-chat', 'glm-4', 'gpt-3.5-turbo', 'gpt-4', 'gpt-4-turbo', 'azure_gpt-4', 'azure_gpt-35-turbo-16k', 'azure_gpt-35-turbo'],
+                    # value='azure_gpt-4',
+                    value='kimi',
+                    label="选择语言模型")
+                temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.01, step=0.01, label="温度参数")
+                embedding_model = gr.Dropdown(
+                    choices=['HuggingFace', 'TensorflowHub', 'OpenAI'],
+                    value='OpenAI',
+                    label="选择嵌入模型")
+                k = gr.Slider(minimum=1, maximum=10, value=5, step=1, label="检索到的文档块数量")
+                with gr.Row():
+                    gr_submit = gr.Button('提交')
+                    gr_clear = gr.Button('清除上下文')
+
+            outputs_answer = gr.Textbox(placeholder="当你点击提交按钮后，这里会显示 RAG 模型给出的回答。",
+                                        label="回答 (A)")
+        with gr.Row():
+            outputs_context = gr.Markdown(label="参考的文档，检索得到的 context 用高亮显示 (C)")
+
+        gr.Markdown(tos_markdown)
+
+        gr_submit.click(
+            rag_answer,
+            inputs=[inputs, model_name, temperature, embedding_model, k],
+            outputs=[outputs_answer, outputs_context],
+        )
+        gr_clear.click(clear_context, outputs=[outputs_answer, outputs_context])
+
+    concurrency = int(os.environ.get('CONCURRENCY', os.cpu_count()))
+    favicon_path = os.path.join(os.path.dirname(__file__), 'assets', 'avatar.png')
+    rag_demo.queue().launch(max_threads=concurrency, favicon_path=favicon_path, share=True)
+    

app_mqa.py

@@ -0,0 +1,132 @@
+import os
+
+import gradio as gr
+from dotenv import load_dotenv
+from langchain.document_loaders import TextLoader
+
+from rag_demo import load_and_split_document, create_vector_store, setup_rag_chain, execute_query
+
+# 环境设置
+load_dotenv()  # 加载环境变量
+QUESTION_LANG = os.getenv("QUESTION_LANG")  # 从环境变量获取 QUESTION_LANG
+assert QUESTION_LANG in ['cn', 'en'], QUESTION_LANG
+
+if QUESTION_LANG == "cn":
+    title = "LightZero RAG Demo"
+    title_markdown = """
+    <div align="center">
+        <img src="https://raw.githubusercontent.com/puyuan1996/RAG/main/assets/banner.svg" width="80%" height="20%" alt="Banner Image">
+    </div>
+    <h2 style="text-align: center; color: black;"><a href="https://github.com/puyuan1996/RAG"> LightZero RAG Demo</a></h2>
+    <h4 align="center"> 📢说明：请您在下面的"问题（Q）"框中输入任何关于 LightZero 的问题，然后点击"提交"按钮。右侧"回答（A）"框中会显示 RAG 模型给出的回答。在 QA 栏的下方会给出参考文档（其中检索得到的相关文段会用黄色高亮显示）。</h4>
+    <h4 align="center"> 如果你喜欢这个项目，请给我们在 GitHub 点个 star ✨ 。我们将会持续保持更新。  </h4>
+    <strong><h5 align="center">注意：算法模型的输出可能包含一定的随机性。相关结果不代表任何开发者和相关 AI 服务的态度和意见。本项目开发者不对生成结果作任何保证，仅供参考。<h5></strong>
+    """
+    tos_markdown = """
+    ### 使用条款
+    玩家使用本服务须同意以下条款：
+    该服务是一项探索性研究预览版，仅供非商业用途。它仅提供有限的安全措施，并可能生成令人反感的内容。不得将其用于任何非法、有害、暴力、种族主义等目的。
+    如果您的游玩体验有不佳之处，请发送邮件至 opendilab@pjlab.org.cn ！ 我们将删除相关信息，并不断改进这个项目。
+    为了获得最佳体验，请使用台式电脑，因为移动设备可能会影响可视化效果。
+    **版权所有 2024 OpenDILab。**
+    """
+
+# 路径变量，方便之后的文件使用
+file_path = './documents/LightZero_README.zh.md'
+
+# 加载原始Markdown文档
+loader = TextLoader(file_path)
+orig_documents = loader.load()
+
+# 存储对话历史
+conversation_history = []
+
+
+def rag_answer(question, model_name, temperature, embedding_model, k):
+    """
+    处理用户问题并返回答案和高亮显示的上下文
+
+    :param question: 用户输入的问题
+    :param model_name: 使用的语言模型名称
+    :param temperature: 生成答案时使用的温度参数
+    :param embedding_model: 使用的嵌入模型
+    :param k: 检索到的文档块数量
+    :return: 模型生成的答案和高亮显示上下文的Markdown文本
+    """
+    try:
+        chunks = load_and_split_document(file_path, chunk_size=5000, chunk_overlap=500)
+        retriever = create_vector_store(chunks, model=embedding_model, k=k)
+        rag_chain = setup_rag_chain(model_name=model_name, temperature=temperature)
+
+        # 将问题添加到对话历史中
+        conversation_history.append(("User", question))
+
+        # 将对话历史转换为字符串
+        history_str = "\n".join([f"{role}: {text}" for role, text in conversation_history])
+
+        retrieved_documents, answer = execute_query(retriever, rag_chain, history_str, model_name=model_name,
+                                                    temperature=temperature)
+        # 在文档中高亮显示上下文
+        context = [retrieved_documents[i].page_content for i in range(len(retrieved_documents))]
+        highlighted_document = orig_documents[0].page_content
+        for i in range(len(context)):
+            highlighted_document = highlighted_document.replace(context[i], f"<mark>{context[i]}</mark>")
+
+        # 将回答添加到对话历史中
+        conversation_history.append(("Assistant", answer))
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        return "处理您的问题时出现错误，请稍后再试。", ""
+    return answer, highlighted_document
+
+
+def clear_context():
+    """
+    清除对话历史
+    """
+    global conversation_history
+    conversation_history = []
+    return "", ""
+
+
+if __name__ == "__main__":
+    with gr.Blocks(title=title, theme='ParityError/Interstellar') as rag_demo:
+        gr.Markdown(title_markdown)
+
+        with gr.Row():
+            with gr.Column():
+                inputs = gr.Textbox(
+                    placeholder="请您输入任何关于 LightZero 的问题。",
+                    label="问题 (Q)")
+                model_name = gr.Dropdown(
+                    choices=['kimi', 'abab6-chat', 'glm-4', 'gpt-3.5-turbo', 'gpt-4', 'gpt-4-turbo', 'azure_gpt-4', 'azure_gpt-35-turbo-16k', 'azure_gpt-35-turbo'],
+                    # value='azure_gpt-4',
+                    value='kimi',
+                    label="选择语言模型")
+                temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.01, step=0.01, label="温度参数")
+                embedding_model = gr.Dropdown(
+                    choices=['HuggingFace', 'TensorflowHub', 'OpenAI'],
+                    value='OpenAI',
+                    label="选择嵌入模型")
+                k = gr.Slider(minimum=1, maximum=10, value=5, step=1, label="检索到的文档块数量")
+                with gr.Row():
+                    gr_submit = gr.Button('提交')
+                    gr_clear = gr.Button('清除上下文')
+
+            outputs_answer = gr.Textbox(placeholder="当你点击提交按钮后，这里会显示 RAG 模型给出的回答。",
+                                        label="回答 (A)")
+        with gr.Row():
+            outputs_context = gr.Markdown(label="参考的文档，检索得到的 context 用高亮显示 (C)")
+
+        gr.Markdown(tos_markdown)
+
+        gr_submit.click(
+            rag_answer,
+            inputs=[inputs, model_name, temperature, embedding_model, k],
+            outputs=[outputs_answer, outputs_context],
+        )
+        gr_clear.click(clear_context, outputs=[outputs_answer, outputs_context])
+
+    concurrency = int(os.environ.get('CONCURRENCY', os.cpu_count()))
+    favicon_path = os.path.join(os.path.dirname(__file__), 'assets', 'avatar.png')
+    rag_demo.queue().launch(max_threads=concurrency, favicon_path=favicon_path, share=True)

app_qa.py

@@ -0,0 +1,106 @@
+import os
+
+import gradio as gr
+from dotenv import load_dotenv
+from langchain.document_loaders import TextLoader
+
+from rag_demo import load_and_split_document, create_vector_store, setup_rag_chain, execute_query
+
+# 环境设置
+load_dotenv()  # 加载环境变量
+QUESTION_LANG = os.getenv("QUESTION_LANG")  # 从环境变量获取 QUESTION_LANG
+assert QUESTION_LANG in ['cn', 'en'], QUESTION_LANG
+
+if QUESTION_LANG == "cn":
+    title = "LightZero RAG Demo"
+    title_markdown = """
+    <div align="center">
+        <img src="https://raw.githubusercontent.com/puyuan1996/RAG/main/assets/banner.svg" width="80%" height="20%" alt="Banner Image">
+    </div>
+    <h2 style="text-align: center; color: black;"><a href="https://github.com/puyuan1996/RAG"> LightZero RAG Demo</a></h2>
+    <h4 align="center"> 📢说明：请您在下面的"问题（Q）"框中输入任何关于 LightZero 的问题，然后点击"提交"按钮。右侧"回答（A）"框中会显示 RAG 模型给出的回答。在 QA 栏的下方会给出参考文档（其中检索得到的相关文段会用黄色高亮显示）。</h4>
+    <h4 align="center"> 如果你喜欢这个项目，请给我们在 GitHub 点个 star ✨ 。我们将会持续保持更新。  </h4>
+    <strong><h5 align="center">注意：算法模型的输出可能包含一定的随机性。相关结果不代表任何开发者和相关 AI 服务的态度和意见。本项目开发者不对生成结果作任何保证，仅供参考。<h5></strong>
+    """
+    tos_markdown = """
+    ### 使用条款
+    玩家使用本服务须同意以下条款：
+    该服务是一项探索性研究预览版，仅供非商业用途。它仅提供有限的安全措施，并可能生成令人反感的内容。不得将其用于任何非法、有害、暴力、种族主义等目的。
+    如果您的游玩体验有不佳之处，请发送邮件至 opendilab@pjlab.org.cn ！ 我们将删除相关信息，并不断改进这个项目。
+    为了获得最佳体验，请使用台式电脑，因为移动设备可能会影响可视化效果。
+    **版权所有 2024 OpenDILab。**
+    """
+
+# 路径变量，方便之后的文件使用
+file_path = './documents/LightZero_README.zh.md'
+
+# 加载原始Markdown文档
+loader = TextLoader(file_path)
+orig_documents = loader.load()
+
+def rag_answer(question, model_name, temperature, embedding_model, k):
+    """
+    处理用户问题并返回答案和高亮显示的上下文
+
+    :param question: 用户输入的问题
+    :param model_name: 使用的语言模型名称
+    :param temperature: 生成答案时使用的温度参数
+    :param embedding_model: 使用的嵌入模型
+    :param k: 检索到的文档块数量
+    :return: 模型生成的答案和高亮显示上下文的Markdown文本
+    """
+    try:
+        chunks = load_and_split_document(file_path, chunk_size=5000, chunk_overlap=500)
+        retriever = create_vector_store(chunks, model=embedding_model, k=k)
+        rag_chain = setup_rag_chain(model_name=model_name, temperature=temperature)
+
+        retrieved_documents, answer = execute_query(retriever, rag_chain, question, model_name=model_name, temperature=temperature)
+        # 在文档中高亮显示上下文
+        context = [retrieved_documents[i].page_content for i in range(len(retrieved_documents))]
+        highlighted_document = orig_documents[0].page_content
+        for i in range(len(context)):
+            highlighted_document = highlighted_document.replace(context[i], f"<mark>{context[i]}</mark>")
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        return "处理您的问题时出现错误，请稍后再试。", ""
+    return answer, highlighted_document
+
+
+if __name__ == "__main__":
+    with gr.Blocks(title=title, theme='ParityError/Interstellar') as rag_demo:
+        gr.Markdown(title_markdown)
+
+        with gr.Row():
+            with gr.Column():
+                inputs = gr.Textbox(
+                    placeholder="请您输入任何关于 LightZero 的问题。",
+                    label="问题 (Q)")
+                model_name = gr.Dropdown(
+                    choices=['kimi', 'abab6-chat', 'glm-4', 'gpt-3.5-turbo', 'gpt-4', 'gpt-4-turbo', 'azure_gpt-4', 'azure_gpt-35-turbo-16k', 'azure_gpt-35-turbo'],
+                    # value='azure_gpt-4',
+                    value='kimi',
+                    label="选择语言模型")
+                temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.01, step=0.01, label="温度参数")
+                embedding_model = gr.Dropdown(
+                    choices=['HuggingFace', 'TensorflowHub', 'OpenAI'],
+                    value='OpenAI',
+                    label="选择嵌入模型")
+                k = gr.Slider(minimum=1, maximum=10, value=5, step=1, label="检索到的文档块数量")
+                gr_submit = gr.Button('提交')
+
+            outputs_answer = gr.Textbox(placeholder="当你点击提交按钮后，这里会显示 RAG 模型给出的回答。",
+                                        label="回答 (A)")
+        with gr.Row():
+            outputs_context = gr.Markdown(label="参考的文档，检索得到的 context 用高亮显示 (C)")
+
+        gr.Markdown(tos_markdown)
+
+        gr_submit.click(
+            rag_answer,
+            inputs=[inputs, model_name, temperature, embedding_model, k],
+            outputs=[outputs_answer, outputs_context],
+        )
+
+    concurrency = int(os.environ.get('CONCURRENCY', os.cpu_count()))
+    favicon_path = os.path.join(os.path.dirname(__file__), 'assets', 'avatar.png')
+    rag_demo.queue().launch(max_threads=concurrency, favicon_path=favicon_path, share=True)

rag_demo.py

@@ -2,24 +2,34 @@
 参考博客：https://mp.weixin.qq.com/s/RUdZjQMSlVOfHfhErSNXnA
 """
 # 导入必要的库与模块
+import json
 import os
 import textwrap
 
+import requests
 from dotenv import load_dotenv
 from langchain.chat_models import ChatOpenAI
 from langchain.document_loaders import TextLoader
-from langchain.embeddings import OpenAIEmbeddings
+from langchain.embeddings import OpenAIEmbeddings, HuggingFaceEmbeddings, TensorflowHubEmbeddings
 from langchain.prompts import ChatPromptTemplate
 from langchain.schema.output_parser import StrOutputParser
-from langchain.schema.runnable import RunnablePassthrough
 from langchain.text_splitter import CharacterTextSplitter
 from langchain.vectorstores import Weaviate
 from weaviate import Client
 from weaviate.embedded import EmbeddedOptions
+from zhipuai import ZhipuAI
+from openai import AzureOpenAI
 
 # 环境设置与文档下载
 load_dotenv()  # 加载环境变量
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")  # 从环境变量获取 OpenAI API 密钥
+MIMIMAX_API_KEY = os.getenv("MIMIMAX_API_KEY")
+MIMIMAX_GROUP_ID = os.getenv("MIMIMAX_GROUP_ID")
+ZHIPUAI_API_KEY = os.getenv("ZHIPUAI_API_KEY")
+KIMI_OPENAI_API_KEY = os.getenv("KIMI_OPENAI_API_KEY")
+
+AZURE_OPENAI_KEY = os.getenv("AZURE_OPENAI_KEY")
+AZURE_ENDPOINT = os.getenv("AZURE_ENDPOINT")
 
 # 确保 OPENAI_API_KEY 被正确设置
 if not OPENAI_API_KEY:
@@ -37,79 +47,231 @@ def load_and_split_document(file_path, chunk_size=500, chunk_overlap=50):
 
 
 # 向量存储建立
-def create_vector_store(chunks, model="OpenAI"):
+def create_vector_store(chunks, model="OpenAI", k=4):
     """将文档块转换为向量并存储到 Weaviate 中"""
     client = Client(embedded_options=EmbeddedOptions())
-    embedding_model = OpenAIEmbeddings() if model == "OpenAI" else None  # 可以根据需要替换为其他嵌入模型
+
+    if model == "OpenAI":
+        embedding_model = OpenAIEmbeddings()
+    elif model == "HuggingFace":
+        embedding_model = HuggingFaceEmbeddings()
+    elif model == "TensorflowHub":
+        embedding_model = TensorflowHubEmbeddings()
+    else:
+        raise ValueError(f"Unsupported embedding model: {model}")
+
     vectorstore = Weaviate.from_documents(
         client=client,
         documents=chunks,
         embedding=embedding_model,
         by_text=False
     )
-    return vectorstore.as_retriever()
+    return vectorstore.as_retriever(search_kwargs={'k': k})
 
 
-# 定义检索增强生成流程
 def setup_rag_chain(model_name="gpt-4", temperature=0):
     """设置检索增强生成流程"""
-    prompt_template = """You are an assistant for question-answering tasks. 
-    Use your knowledge to answer the question if the provided context is not relevant. 
-    Otherwise, use the context to inform your answer. 
-    Question: {question} 
-    Context: {context} 
-    Answer:
-    """
-    prompt = ChatPromptTemplate.from_template(prompt_template)
-    llm = ChatOpenAI(model_name=model_name, temperature=temperature)
-    # 创建 RAG 链，参考 https://python.langchain.com/docs/expression_language/
-    rag_chain = (
-            prompt
-            | llm
-            | StrOutputParser()
-    )
+    if model_name.startswith("gpt"):
+        # 如果是以gpt开头的模型,使用原来的逻辑
+        prompt_template = """您是一个用于问答任务的专业助手。
+        在处理问答任务时,请根据所提供的[上下文信息]给出回答。
+        如果[上下文信息]与[问题]不相关,那么请运用您的知识库为提问者提供准确的答复。
+        请确保回答内容的质量, 包括相关性、准确性和可读性。
+        [问题]: {question} 
+        [上下文信息]: {context} 
+        [回答]:
+        """
+        prompt = ChatPromptTemplate.from_template(prompt_template)
+        llm = ChatOpenAI(model_name=model_name, temperature=temperature)
+        # 创建 RAG 链,参考 https://python.langchain.com/docs/expression_language/
+        rag_chain = (
+                prompt
+                | llm
+                | StrOutputParser()
+        )
+    else:
+        # 如果不是以gpt开头的模型,返回None
+        rag_chain = None
     return rag_chain
 
 
 # 执行查询并打印结果
-def execute_query(retriever, rag_chain, query):
-    """执行查询并返回结果及检索到的文档块"""
+def execute_query(retriever, rag_chain, query, model_name="gpt-4", temperature=0):
+    """
+    执行查询并返回结果及检索到的文档块
+
+    参数:
+    retriever: 文档检索器对象
+    rag_chain: 检索增强生成链对象,如果为None则不使用RAG链
+    query: 查询问题
+    model_name: 使用的语言模型名称,默认为"gpt-4"
+    temperature: 生成温度,默认为0
+
+    返回:
+    retrieved_documents: 检索到的文档块列表
+    response_text: 生成的回答文本
+    """
+    # 使用检索器检索相关文档块
     retrieved_documents = retriever.invoke(query)
-    rag_chain_response = rag_chain.invoke({"context": retrieved_documents, "question": query})
-    return retrieved_documents, rag_chain_response
+
+    if rag_chain is not None:
+        # 如果有RAG链,则使用RAG链生成回答
+        rag_chain_response = rag_chain.invoke({"context": retrieved_documents, "question": query})
+        response_text = rag_chain_response
+    else:
+        # 如果没有RAG链,则将检索到的文档块和查询问题按照指定格式输入给语言模型
+        if model_name == "kimi":
+            # 对于有检索能力的模型,使用不同的模板
+            prompt_template = """您是一个用于问答任务的专业助手。
+            在处理问答任务时，请根据所提供的【上下文信息】和【你的知识库和检索到的相关文档】给出回答。
+            请确保回答内容的质量，包括相关性、准确性和可读性。
+           【问题】: {question} 
+           【上下文信息】: {context} 
+           【回答】:
+            """
+        else:
+            prompt_template = """您是一个用于问答任务的专业助手。
+            在处理问答任务时，请根据所提供的【上下文信息】给出回答。
+            如果【上下文信息】与【问题】不相关,那么请运用您的知识库为提问者提供准确的答复。
+            请确保回答内容的质量，包括相关性、准确性和可读性。
+           【问题】: {question} 
+           【上下文信息】: {context} 
+           【回答】:
+            """
+
+        context = '\n'.join(
+            [f'**Document {i}**: ' + retrieved_documents[i].page_content for i in range(len(retrieved_documents))])
+        prompt = prompt_template.format(question=query, context=context)
+        response_text = execute_query_no_rag(model_name=model_name, temperature=temperature, query=prompt)
+    return retrieved_documents, response_text
 
 
-# 执行无 RAG 链的查询
 def execute_query_no_rag(model_name="gpt-4", temperature=0, query=""):
     """执行无 RAG 链的查询"""
-    llm = ChatOpenAI(model_name=model_name, temperature=temperature)
-    response = llm.invoke(query)
-    return response.content
+    if model_name.startswith("gpt"):
+        # 如果是以gpt开头的模型,使用原来的逻辑
+        llm = ChatOpenAI(model_name=model_name, temperature=temperature)
+        response = llm.invoke(query)
+        return response.content
+    elif model_name.startswith("azure_gpt"):
+        client = AzureOpenAI(
+            azure_endpoint=AZURE_ENDPOINT,
+            api_key=AZURE_OPENAI_KEY,
+            api_version="2024-02-15-preview"
+        )
+        message_text = [{"role": "user", "content": query}, ]
+        completion = client.chat.completions.create(
+            model=model_name[6:],  # model_name = 'azure_gpt-4', 'azure_gpt-35-turbo-16k', 'azure_gpt-35-turbo'
+            messages=message_text,
+            temperature=temperature,
+            top_p=0.95,
+            frequency_penalty=0,
+            presence_penalty=0,
+            stop=None
+        )
+        return completion.choices[0].message.content
+    elif model_name == 'abab6-chat':
+        # 如果是'abab6-chat'模型,使用专门的API调用方式
+        url = "https://api.minimax.chat/v1/text/chatcompletion_pro?GroupId=" + MIMIMAX_GROUP_ID
+        headers = {"Content-Type": "application/json", "Authorization": "Bearer " + MIMIMAX_API_KEY}
+        payload = {
+            "bot_setting": [
+                {
+                    "bot_name": "MM智能助理",
+                    "content": "MM智能助理是一款由MiniMax自研的,没有调用其他产品的接口的大型语言模型。MiniMax是一家中国科技公司,一直致力于进行大模型相关的研究。",
+                }
+            ],
+            "messages": [{"sender_type": "USER", "sender_name": "小明", "text": query}],
+            "reply_constraints": {"sender_type": "BOT", "sender_name": "MM智能助理"},
+            "model": model_name,
+            "tokens_to_generate": 1034,
+            "temperature": temperature,
+            "top_p": 0.9,
+        }
+
+        response = requests.request("POST", url, headers=headers, json=payload)
+        # 将 JSON 字符串解析为字典
+        response_dict = json.loads(response.text)
+        # 提取 'reply' 键对应的值
+        return response_dict['reply']
+
+    elif model_name == 'glm-4':
+        # 如果是'glm-4'模型,使用专门的API调用方式
+        client = ZhipuAI(api_key=ZHIPUAI_API_KEY)  # 填写您自己的APIKey
+        response = client.chat.completions.create(
+            model=model_name,  # 填写需要调用的模型名称
+            messages=[{"role": "user", "content": query}]
+        )
+        return response.choices[0].message.content
+    elif model_name == 'kimi':
+        # 如果是'kimi'模型,使用专门的API调用方式
+        from openai import OpenAI
+        client = OpenAI(
+            api_key=KIMI_OPENAI_API_KEY,
+            base_url="https://api.moonshot.cn/v1",
+        )
+        messages = [
+            {
+                "role": "system",
+                "content": "你是 Kimi，由 Moonshot AI 提供的人工智能助手，你更擅长中文和英文的对话。你会为用户提供安全，有帮助，准确的回答。同时，你会拒绝一切涉及恐怖主义，种族歧视，黄色暴力等问题的回答。Moonshot AI 为专有名词，不可翻译成其他语言。",
+            },
+            {"role": "user",
+             "content": query},
+        ]
+        completion = client.chat.completions.create(
+            model="moonshot-v1-128k",
+            messages=messages,
+            temperature=0.01,
+            top_p=1.0,
+            n=1,  # 为每条输入消息生成多少个结果
+            stream=False  # 流式输出
+        )
+        return completion.choices[0].message.content
+    else:
+        # 如果模型不支持,抛出异常
+        raise ValueError(f"Unsupported model: {model_name}")
 
 
-# rag_demo.py 相对 rag_demo_v0.py 的不同之处在于可以输出检索到的文档块。
 if __name__ == "__main__":
     # 假设文档已存在于本地
     file_path = './documents/LightZero_README.zh.md'
+    # model_name = "glm-4"  # model_name=['abab6-chat', 'glm-4', 'gpt-3.5-turbo', 'gpt-4', 'gpt-4-turbo', 'azure_gpt-4', 'azure_gpt-35-turbo-16k', 'azure_gpt-35-turbo']
+    model_name = 'azure_gpt-4'
+    temperature = 0.01
+    # embedding_model = 'HuggingFace'  # embedding_model=['HuggingFace', 'TensorflowHub', 'OpenAI']
+    embedding_model = 'OpenAI'  # embedding_model=['HuggingFace', 'TensorflowHub', 'OpenAI']
 
     # 加载和分割文档
-    chunks = load_and_split_document(file_path)
+    chunks = load_and_split_document(file_path, chunk_size=5000, chunk_overlap=500)
 
     # 创建向量存储
-    retriever = create_vector_store(chunks)
+    retriever = create_vector_store(chunks, model=embedding_model, k=5)
 
     # 设置 RAG 流程
-    rag_chain = setup_rag_chain()
+    rag_chain = setup_rag_chain(model_name=model_name, temperature=temperature)
 
     # 提出问题并获取答案
-    query = "请问 LightZero 里面实现的 AlphaZero 算法支持在 Atari 环境上运行吗？请详细解释原因"
-    # query = "请详细解释 MCTS 算法的原理，并给出带有详细中文注释的 Python 代码示例"
+    query = ("GitHub - opendilab/LightZero: [NeurIPS 2023 Spotlight] LightZero: A Unified Benchmark for Monte Carl  请根据这个仓库回答下面的问题：（1）请简要介绍一下 LightZero （2）请详细介绍 LightZero 的框架结构。 （3）请给出安装 LightZero，运行他们的示例代码的详细步骤 （4）- 请问 LightZero 具体支持什么任务（tasks/environments）? （5）请问 LightZero 具体支持什么算法?（6）请问 LightZero 具体支持什么算法，各自支持在哪些任务上运行? （7）请问 LightZero 里面实现的 MuZero 算法支持在 Atari 任务上运行吗？（8）请问 LightZero 里面实现的 AlphaZero 算法支持在 Atari 任务上运行吗？（9）LightZero 支持哪些算法? 各自的优缺点是什么? 我应该如何根据任务特点进行选择呢？（10）请结合 LightZero 中的代码介绍他们是如何实现 MCTS 的。（11）请问对这个仓库提出详细的改进建议")
+    """
+    （1）请简要介绍一下 LightZero 
+    （2）请详细介绍 LightZero 的框架结构。 
+    （3）请给出安装 LightZero，运行他们的示例代码的详细步骤 
+    （4）请问 LightZero 具体支持什么任务（tasks/environments）? 
+    （5）请问 LightZero 具体支持什么算法?
+    （6）请问 LightZero 具体支持什么算法，各自支持在哪些任务上运行? 
+    （7）请问 LightZero 里面实现的 MuZero 算法支持在 Atari 任务上运行吗？
+    （8）请问 LightZero 里面实现的 AlphaZero 算法支持在 Atari 任务上运行吗？
+    （9）LightZero 支持哪些算法? 各自的优缺点是什么? 我应该如何根据任务特点进行选择呢？
+    （10）请结合 LightZero 中的代码介绍他们是如何实现 MCTS 的。
+    （11）请问对这个仓库提出详细的改进建议。
+    """
 
     # 使用 RAG 链获取参考的文档与答案
-    retrieved_documents, result_with_rag = execute_query(retriever, rag_chain, query)
+    retrieved_documents, result_with_rag = execute_query(retriever, rag_chain, query, model_name=model_name,
+                                                         temperature=temperature)
 
     # 不使用 RAG 链获取答案
-    result_without_rag = execute_query_no_rag(query=query)
+    result_without_rag = execute_query_no_rag(model_name=model_name, query=query, temperature=temperature)
 
     # 打印并对比两种方法的结果
     # 使用textwrap.fill来自动分段文本，width参数可以根据你的屏幕宽度进行调整

rag_demo_v0.py

@@ -1,136 +0,0 @@
-"""
-参考博客：https://mp.weixin.qq.com/s/RUdZjQMSlVOfHfhErSNXnA
-"""
-# 导入必要的库与模块
-import os
-import textwrap
-
-from dotenv import load_dotenv
-from langchain.chat_models import ChatOpenAI
-from langchain.document_loaders import TextLoader
-from langchain.embeddings import OpenAIEmbeddings
-from langchain.prompts import ChatPromptTemplate
-from langchain.schema.output_parser import StrOutputParser
-from langchain.schema.runnable import RunnablePassthrough
-from langchain.text_splitter import CharacterTextSplitter
-from langchain.vectorstores import Weaviate
-from weaviate import Client
-from weaviate.embedded import EmbeddedOptions
-
-# 环境设置与文档下载
-load_dotenv()  # 加载环境变量
-OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")  # 从环境变量获取 OpenAI API 密钥
-
-# 确保 OPENAI_API_KEY 被正确设置
-if not OPENAI_API_KEY:
-    raise ValueError("OpenAI API Key not found in the environment variables.")
-
-
-# 文档加载与分割
-def load_and_split_document(file_path, chunk_size=500, chunk_overlap=50):
-    """加载文档并分割成小块"""
-    loader = TextLoader(file_path)
-    documents = loader.load()
-    text_splitter = CharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
-    chunks = text_splitter.split_documents(documents)
-    return chunks
-
-
-# 向量存储建立
-def create_vector_store(chunks, model="OpenAI"):
-    """将文档块转换为向量并存储到 Weaviate 中"""
-    client = Client(embedded_options=EmbeddedOptions())
-    embedding_model = OpenAIEmbeddings() if model == "OpenAI" else None  # 可以根据需要替换为其他嵌入模型
-    vectorstore = Weaviate.from_documents(
-        client=client,
-        documents=chunks,
-        embedding=embedding_model,
-        by_text=False
-    )
-    return vectorstore.as_retriever()
-
-
-# 定义检索增强生成流程
-def setup_rag_chain_v0(retriever, model_name="gpt-4", temperature=0):
-    """设置检索增强生成流程"""
-    prompt_template = """You are an assistant for question-answering tasks. 
-    Use your knowledge to answer the question if the provided context is not relevant. 
-    Otherwise, use the context to inform your answer.
-    Question: {question} 
-    Context: {context} 
-    Answer:
-    """
-    prompt = ChatPromptTemplate.from_template(prompt_template)
-    llm = ChatOpenAI(model_name=model_name, temperature=temperature)
-    # 创建 RAG 链，参考 https://python.langchain.com/docs/expression_language/
-    rag_chain = (
-            {"context": retriever, "question": RunnablePassthrough()}
-            | prompt
-            | llm
-            | StrOutputParser()
-    )
-    return rag_chain
-
-
-# 执行查询并打印结果
-def execute_query_v0(rag_chain, query):
-    """执行查询并返回结果"""
-    return rag_chain.invoke(query)
-
-
-# 执行无 RAG 链的查询
-def execute_query_no_rag(model_name="gpt-4", temperature=0, query=""):
-    """执行无 RAG 链的查询"""
-    llm = ChatOpenAI(model_name=model_name, temperature=temperature)
-    response = llm.invoke(query)
-    return response.content
-
-
-# rag_demo.py 相对 rag_demo_v0.py 的不同之处在于可以输出检索到的文档块。
-if __name__ == "__main__":
-    # 下载并保存文档到本地（这里被注释掉了，因为已经假设文档存在于本地）
-    # url = "https://raw.githubusercontent.com/langchain-ai/langchain/master/docs/docs/modules/state_of_the_union.txt"
-    # res = requests.get(url)
-    # with open("state_of_the_union.txt", "w") as f:
-    #     f.write(res.text)
-
-    # 假设文档已存在于本地
-    # file_path = './documents/state_of_the_union.txt'
-    file_path = './documents/LightZero_README.zh.md'
-
-    # 加载和分割文档
-    chunks = load_and_split_document(file_path)
-
-    # 创建向量存储
-    retriever = create_vector_store(chunks)
-
-    # 设置 RAG 流程
-    rag_chain = setup_rag_chain_v0(retriever)
-
-    # 提出问题并获取答案
-    # query = "请你分别用中英文简介 LightZero"
-    # query = "请你用英文简介 LightZero"
-    query = "请你用中文简介 LightZero"
-    # query = "请问 LightZero 支持哪些环境和算法，应该如何快速上手使用？"
-    # query = "请问 LightZero 里面实现的 MuZero 算法支持在 Atari 环境上运行吗？"
-    # query = "请问 LightZero 里面实现的 AlphaZero 算法支持在 Atari 环境上运行吗？请详细解释原因"
-    # query = "请详细解释 MCTS 算法的原理，并给出带有详细中文注释的 Python 代码示例"
-
-    # 使用 RAG 链获取答案
-    result_with_rag = execute_query_v0(rag_chain, query)
-
-    # 不使用 RAG 链获取答案
-    result_without_rag = execute_query_no_rag(query=query)
-
-    # 打印并对比两种方法的结果
-    # 使用textwrap.fill来自动分段文本，width参数可以根据你的屏幕宽度进行调整
-    wrapped_result_with_rag = textwrap.fill(result_with_rag, width=80)
-    wrapped_result_without_rag = textwrap.fill(result_without_rag, width=80)
-
-    # 打印自动分段后的文本
-    print("="*40)
-    print(f"我的问题是:\n{query}")
-    print("="*40)
-    print(f"Result with RAG:\n{wrapped_result_with_rag}")
-    print("="*40)
-    print(f"Result without RAG:\n{wrapped_result_without_rag}")

requirements.txt

@@ -5,3 +5,4 @@ weaviate-client
 requests
 python-dotenv
 tiktoken
+sentence-transformers