Skip to content

前端界面

用 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 技术的智能文档助手")

完整代码示例

点击查看完整 app.py

测试检查清单

上传测试文档后,检查:

  • [ ] 文档能正常上传和解析
  • [ ] 能回答文档中的问题
  • [ ] 答案准确率 > 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
    pass

3. 进度显示

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):
    # 处理消息
    pass

Gradio

python
import gradio as gr

def chat(question, history):
    # 处理逻辑
    return answer, history

demo = gr.ChatInterface(chat)
demo.launch()

下一步

界面做好了,测试一下有没有问题。

继续:测试优化 →


← 返回核心功能 | 返回项目一

最近更新

基于 Apache 2.0 许可发布