114 lines
3.9 KiB
Python
114 lines
3.9 KiB
Python
|
import streamlit as st
|
|||
|
|
from dotenv import load_dotenv
|
|||
|
|
from langchain_core.prompts import ChatPromptTemplate
|
|||
|
|
from langchain.prompts import MessagesPlaceholder
|
|||
|
|
from langchain_core.messages import AIMessage, HumanMessage
|
|||
|
|
from langchain.chains import create_retrieval_chain
|
|||
|
|
from langchain.chains.combine_documents import create_stuff_documents_chain
|
|||
|
|
from langchain_openai import ChatOpenAI
|
|||
|
|
from langchain.embeddings import HuggingFaceEmbeddings
|
|||
|
|
from langchain_chroma import Chroma
|
|||
|
|
|
|||
|
|
# 记住最近的6个聊天元素
|
|||
|
|
memory = 6
|
|||
|
|
|
|||
|
|
def init_chain():
|
|||
|
|
"""
|
|||
|
|
初始化聊天模型和检索链。
|
|||
|
|
|
|||
|
|
此函数设置环境变量，加载OpenAI模型，创建ChromaDB检索器，并编译用于问题回答的最终链。
|
|||
|
|
|
|||
|
|
返回：
|
|||
|
|
Runnable Sequence: 初始化的检索链。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
# 加载环境变量并初始化模型
|
|||
|
|
load_dotenv(".env")
|
|||
|
|
key = 1
|
|||
|
|
llm = ChatOpenAI(streaming=True,
|
|||
|
|
verbose=True,
|
|||
|
|
api_key=key,
|
|||
|
|
openai_api_base="http://localhost:8000/v1",
|
|||
|
|
model="Qwen1.5-32b-int4")
|
|||
|
|
|
|||
|
|
embedding = HuggingFaceEmbeddings(model_name = "/home/zhangxj/models/acge_text_embedding")
|
|||
|
|
# 通过加载ChromaDB创建一个数据库检索器
|
|||
|
|
retriever = Chroma(persist_directory="chroma-laws", embedding_function=embedding)
|
|||
|
|
retriever = retriever.as_retriever()
|
|||
|
|
|
|||
|
|
# 指令模板
|
|||
|
|
instruct_system_prompt = (
|
|||
|
|
"你是生命周期领域富有经验的专家。"
|
|||
|
|
"利用检索到的上下文来回答问题。如果上下文没有足够的信息，请说明。"
|
|||
|
|
"如果你不明白，请询问你不清楚的部分。"
|
|||
|
|
"尽量少用句子，但如果有必要可以完整引用文献资料。"
|
|||
|
|
"尽量保持对话有趣和机智。\n\n"
|
|||
|
|
"{context}")
|
|||
|
|
|
|||
|
|
instruct_prompt = ChatPromptTemplate.from_messages(
|
|||
|
|
[
|
|||
|
|
("system", instruct_system_prompt),
|
|||
|
|
MessagesPlaceholder("chat_history"),
|
|||
|
|
("human", "{input}"),
|
|||
|
|
])
|
|||
|
|
|
|||
|
|
# 编译最终链并返回
|
|||
|
|
qa_chain = create_stuff_documents_chain(llm, instruct_prompt)
|
|||
|
|
rag_chain = create_retrieval_chain(retriever, qa_chain)
|
|||
|
|
|
|||
|
|
return rag_chain
|
|||
|
|
|
|||
|
|
def user_in(uin, rag_chain, history):
|
|||
|
|
"""
|
|||
|
|
返回GPT生成的输出。
|
|||
|
|
|
|||
|
|
参数:
|
|||
|
|
uin (str): 用户输入。
|
|||
|
|
rag_chain (Runnable Sequence): llm链。
|
|||
|
|
history (list): 聊天历史列表。
|
|||
|
|
返回:
|
|||
|
|
str: 生成的响应
|
|||
|
|
"""
|
|||
|
|
result = rag_chain.invoke({"input": uin, "chat_history" : history})["answer"]
|
|||
|
|
|
|||
|
|
# 用新问题和响应更新聊天历史
|
|||
|
|
# history.extend([HumanMessage(content=uin),AIMessage(content=result)])
|
|||
|
|
return result
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main(memory=memory):
|
|||
|
|
st.title("LCA-GPT")
|
|||
|
|
|
|||
|
|
# 创建并存储在streamlit缓存中
|
|||
|
|
if "chat_history" not in st.session_state:
|
|||
|
|
st.session_state.chat_history = []
|
|||
|
|
|
|||
|
|
if "rag" not in st.session_state:
|
|||
|
|
st.session_state.rag = init_chain()
|
|||
|
|
|
|||
|
|
if "messages" not in st.session_state:
|
|||
|
|
st.session_state.messages = []
|
|||
|
|
|
|||
|
|
# 更新聊天消息
|
|||
|
|
for message in st.session_state.messages:
|
|||
|
|
with st.chat_message(message["role"]):
|
|||
|
|
st.markdown(message["content"])
|
|||
|
|
|
|||
|
|
if prompt := st.chat_input("请输入你的问题："):
|
|||
|
|
|
|||
|
|
with st.chat_message("user"):
|
|||
|
|
st.markdown(prompt)
|
|||
|
|
|
|||
|
|
st.session_state.messages.append({"role":"user", "content": prompt})
|
|||
|
|
response = user_in(prompt, st.session_state.rag, st.session_state.chat_history[-memory:])
|
|||
|
|
st.session_state.chat_history.extend([HumanMessage(content=prompt),AIMessage(content=response)])
|
|||
|
|
|
|||
|
|
with st.chat_message("assistant"):
|
|||
|
|
st.markdown(response)
|
|||
|
|
st.session_state.messages.append({"role" : "assistant", "content": response})
|
|||
|
|
print("历史:",st.session_state.chat_history[-memory:])
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|