前端界面
用 Streamlit 快速搭建界面
Streamlit 是最简单的方式,几行代码就能做出聊天界面。
基础界面
创建 app.py:
python
import streamlit as st
from rag import (
load_document,
split_documents,
create_vector_store,
load_vector_store,
answer_question
)
from config import config
import os
# 页面配置
st.set_page_config(
page_title="文档问答机器人",
page_icon="🤖",
layout="centered"
)
st.title("📄 文档问答机器人")
st.write("上传文档,然后提问")
# 侧边栏
with st.sidebar:
st.header("⚙️ 设置")
# 文档上传
uploaded_file = st.file_uploader(
"上传文档",
type=['pdf', 'docx', 'txt'],
help="支持 PDF、Word、TXT 格式"
)
# 处理上传的文档
if uploaded_file:
# 保存文件
file_path = f"data/{uploaded_file.name}"
os.makedirs("data", exist_ok=True)
with open(file_path, "wb") as f:
f.write(uploaded_file.getbuffer())
st.success(f"已上传: {uploaded_file.name}")
# 加载和处理文档
with st.spinner("处理文档中..."):
docs = load_document(file_path)
splits = split_documents(docs)
vectorstore = create_vector_store(splits)
st.success(f"处理完成!共 {len(splits)} 个文本块")
# 保存到 session state
st.session_state.vectorstore = vectorstore
# 主界面
if "vectorstore" not in st.session_state:
st.info("👈 请先在侧边栏上传文档")
else:
# 问题输入
question = st.text_input(
"❓ 你想问什么?",
placeholder="比如:这个文档主要讲了什么?",
height=100
)
if st.button("发送", type="primary"):
if question:
with st.spinner("思考中..."):
result = answer_question(
question,
st.session_state.vectorstore
)
# 显示答案
st.markdown("### 🤖 答案")
st.write(result['answer'])
# 显示来源
with st.expander("📚 查看来源"):
for i, source in enumerate(result['sources'], 1):
st.markdown(f"**来源 {i}**")
st.text(source['content'])
st.caption(f"元数据: {source['metadata']}")
st.divider()
else:
st.warning("请输入问题")运行:
bash
streamlit run app.py打开 http://localhost:8501 查看效果。
改进界面
1. 添加对话历史
python
# 初始化对话历史
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("你的问题"):
# 显示用户消息
st.session_state.messages.append({"role": "user", "content": prompt})
with st.chat_message("user"):
st.markdown(prompt)
# 生成回复
with st.chat_message("assistant"):
with st.spinner("思考中..."):
result = answer_question(
prompt,
st.session_state.vectorstore
)
st.markdown(result['answer'])
# 保存到历史
st.session_state.messages.append({
"role": "assistant",
"content": result['answer']
})2. 添加多文档支持
python
# 在侧边栏
uploaded_files = st.file_uploader(
"上传文档(可多选)",
type=['pdf', 'docx', 'txt'],
accept_multiple_files=True
)
if uploaded_files:
if st.button("处理文档"):
all_splits = []
progress_bar = st.progress(0)
for i, uploaded_file in enumerate(uploaded_files):
# 保存并处理
file_path = f"data/{uploaded_file.name}"
with open(file_path, "wb") as f:
f.write(uploaded_file.getbuffer())
docs = load_document(file_path)
splits = split_documents(docs)
all_splits.extend(splits)
progress_bar.progress((i + 1) / len(uploaded_files))
vectorstore = create_vector_store(all_splits)
st.session_state.vectorstore = vectorstore
st.success(f"处理完成!共 {len(all_splits)} 个文本块")3. 添加清空功能
python
# 在侧边栏
if st.button("🗑️ 清空对话"):
st.session_state.messages = []
st.rerun()
if st.button("🔄 重新上传文档"):
st.session_state.vectorstore = None
st.session_state.messages = []
st.rerun()4. 显示统计信息
python
# 在侧边栏
if "vectorstore" in st.session_state:
st.divider()
st.subheader("📊 统计")
# 这里需要根据实际情况调整
vectorstore = st.session_state.vectorstore
# 假设我们保存了文档信息
if "doc_info" in st.session_state:
st.metric("文档数量", st.session_state.doc_info['count'])
st.metric("文本块数量", st.session_state.doc_info['splits'])美化界面
自定义 CSS
python
st.markdown("""
<style>
.stTextInput > div > div > input {
background-color: #f0f2f6;
}
.stButton > button {
width: 100%;
border-radius: 20px;
}
</style>
""", unsafe_allow_html=True)添加 Logo 和描述
python
col1, col2, col3 = st.columns([1, 6, 1])
with col2:
st.image("logo.png", width=100)
st.title("📄 文档问答机器人")
st.caption("基于 RAG 技术的智能文档助手")完整代码示例
测试检查清单
上传测试文档后,检查:
- [ ] 文档能正常上传和解析
- [ ] 能回答文档中的问题
- [ ] 答案准确率 > 70%
- [ ] 来源显示正确
- [ ] 对话历史保存正常
- [ ] 清空功能工作正常
- [ ] 多文档上传正常
性能优化建议
1. 缓存向量化结果
python
@st.cache_resource
def get_vector_store(file_path):
"""缓存向量库"""
docs = load_document(file_path)
splits = split_documents(docs)
return create_vector_store(splits)2. 异步处理
python
import asyncio
async def async_answer_question(question, vectorstore):
# 使用异步版本的 API
pass3. 进度显示
python
import time
with st.spinner("处理中"):
for i in range(100):
time.sleep(0.01)
st.progress((i + 1) / 100)替代方案
如果 Streamlit 不能满足需求,可以考虑:
Chainlit(专为聊天设计)
python
import chainlit as cl
@cl.on_chat_start
async def start():
# 初始化
pass
@cl.on_message
async def main(message: str):
# 处理消息
passGradio
python
import gradio as gr
def chat(question, history):
# 处理逻辑
return answer, history
demo = gr.ChatInterface(chat)
demo.launch()下一步
界面做好了,测试一下有没有问题。