app.py

import streamlit as st
from dotenv import load_dotenv
import os
from document_processor import DocumentProcessor
from vector_db import VectorDB
from rag_chain import RAGChain

# 加载环境变量
load_dotenv()

# 全局变量
if "vector_db" not in st.session_state:
    st.session_state.vector_db = VectorDB()
    st.session_state.vector_db.init_db()
    
if "document_processor" not in st.session_state:
    st.session_state.document_processor = DocumentProcessor()
    
if "rag_chain" not in st.session_state:
    st.session_state.rag_chain = RAGChain(st.session_state.vector_db)

# 设置页面配置
st.set_page_config(
    page_title="AutoRAG - 零代码RAG系统",
    page_icon="🤖",
    layout="wide"
)

# 页面标题
st.title("🤖 AutoRAG - 零代码RAG系统")
st.markdown("无需编程知识，上传文档即可自动构建完整的RAG系统")

# 侧边栏导航
page = st.sidebar.selectbox(
    "选择功能",
    ["首页", "文档上传", "智能问答", "系统设置"]
)

# 首页
if page == "首页":
    st.header("欢迎使用AutoRAG")
    st.markdown("""
    **AutoRAG** 是一个零代码的RAG（检索增强生成）系统，帮助您轻松构建基于文档的智能问答系统。
    
    ### 核心功能
    - 📤 **文档上传**：支持多种格式文档上传（PDF、Word、TXT等）
    - 📊 **自动分类**：智能分类上传的文档
    - 🔍 **高效检索**：基于向量数据库的快速检索
    - 🤖 **智能生成**：结合LLM生成准确回答
    - 📋 **问答模板**：提供多种预设问答模板
    
    ### 开始使用
    1. 点击左侧"文档上传"上传您的文档
    2. 系统自动构建向量数据库
    3. 进入"智能问答"开始提问
    """)
    
    # 显示数据库状态
    doc_count = st.session_state.vector_db.get_document_count()
    st.metric("数据库中文档数量", doc_count)

# 文档上传页面
elif page == "文档上传":
    st.header("📤 文档上传")
    st.markdown("上传您的文档，系统将自动处理并构建向量数据库")
    
    # 上传组件
    uploaded_files = st.file_uploader(
        "选择要上传的文档",
        accept_multiple_files=True,
        type=["pdf", "docx", "txt", "md"]
    )
    
    if uploaded_files:
        st.success(f"已上传 {len(uploaded_files)} 个文件")
        for file in uploaded_files:
            st.write(f"- {file.name} ({file.size / 1024:.2f} KB)")
        
        if st.button("开始处理文档", type="primary"):
            with st.spinner("正在处理文档..."):
                # 处理文档
                split_docs = st.session_state.document_processor.process_files(uploaded_files)
                
                # 添加到向量数据库
                st.session_state.vector_db.add_documents(split_docs)
                
                st.success(f"文档处理完成！共处理 {len(split_docs)} 个文档块")
                st.info(f"当前数据库中文档数量: {st.session_state.vector_db.get_document_count()}")
    
    # 显示当前数据库状态
    doc_count = st.session_state.vector_db.get_document_count()
    if doc_count > 0:
        st.markdown("---")
        st.subheader("当前数据库状态")
        st.metric("文档块数量", doc_count)
        
        if st.button("清空数据库", type="secondary"):
            if st.checkbox("确认清空所有数据"):
                st.session_state.vector_db.delete_all_documents()
                st.success("数据库已清空")

# 智能问答页面
elif page == "智能问答":
    st.header("🤖 智能问答")
    st.markdown("基于您上传的文档进行智能问答")
    
    # 问答模板选择
    template = st.selectbox(
        "选择问答模板",
        ["通用问答", "技术文档", "产品手册", "法律文档", "自定义"]
    )
    
    # 问题输入
    question = st.text_area("请输入您的问题", height=100)
    
    if st.button("获取答案", type="primary"):
        if question:
            # 检查数据库是否有文档
            if st.session_state.vector_db.get_document_count() == 0:
                st.warning("数据库为空，请先上传文档")
            else:
                with st.spinner("正在生成答案..."):
                    # 调用RAG链生成回答
                    try:
                        answer = st.session_state.rag_chain.invoke(question)
                        st.info(answer)
                    except Exception as e:
                        st.error(f"生成答案失败: {e}")
        else:
            st.warning("请输入问题")

# 系统设置页面
elif page == "系统设置":
    st.header("⚙️ 系统设置")
    
    # API密钥设置
    st.subheader("API密钥设置")
    openai_key = st.text_input("OpenAI API Key", type="password", value=os.getenv("OPENAI_API_KEY", ""))
    anthropic_key = st.text_input("Anthropic API Key", type="password", value=os.getenv("ANTHROPIC_API_KEY", ""))
    
    # 模型设置
    st.subheader("模型设置")
    llm_model = st.selectbox("选择LLM模型", ["gpt-3.5-turbo", "gpt-4", "claude-3-sonnet-20240229"])
    embedding_model = st.selectbox("选择嵌入模型", ["all-MiniLM-L6-v2", "text-embedding-ada-002"])
    
    if st.button("保存设置", type="primary"):
        # 保存API密钥到环境变量
        if openai_key:
            os.environ["OPENAI_API_KEY"] = openai_key
        if anthropic_key:
            os.environ["ANTHROPIC_API_KEY"] = anthropic_key
        
        # 更新向量数据库嵌入模型
        st.session_state.vector_db.update_embedding_model(embedding_model)
        
        # 更新RAG链LLM模型
        st.session_state.rag_chain.update_llm_model(llm_model)
        
        st.success("设置已保存！")