streamlit_app.py

import streamlit as st
import requests
import json
import pandas as pd
from datetime import datetime

# API 基础 URL
API_BASE_URL = "http://localhost:6006"

# 页面配置
st.set_page_config(
    page_title="政府问答 RAG 系统",
    page_icon="🏛️",
    layout="wide"
)

# 侧边栏配置
st.sidebar.title("🏛️ 政府问答 RAG 系统")
page = st.sidebar.radio(
    "选择功能",
    ["📚 知识库管理", "📄 文档管理", "💬 智能问答", "🔍 文本向量化"]
)

# Token 配置（可以放在侧边栏或配置文件中）
TOKEN = st.sidebar.text_input("API Token", value="test_token", type="password")
MODEL = "bge-small-zh-v1.5"


# ==================== 知识库管理页面 ====================
if page == "📚 知识库管理":
    st.title("📚 知识库管理")

    # 创建知识库
    st.subheader("创建新知识库")
    col1, col2 = st.columns(2)

    with col1:
        kb_title = st.text_input("知识库名称", placeholder="例如：政府政策文档库")
    with col2:
        kb_category = st.text_input("知识库类别", placeholder="例如：政府文档")

    if st.button("创建知识库", type="primary"):
        if kb_title and kb_category:
            try:
                response = requests.post(
                    f"{API_BASE_URL}/v1/knowledge_base",
                    json={"title": kb_title, "category": kb_category}
                )
                if response.status_code == 200:
                    data = response.json()
                    st.success(f"✅ 知识库创建成功！ID: {data['knowledge_id']}")
                    st.json(data)
                else:
                    st.error(f"❌ 创建失败: {response.text}")
            except Exception as e:
                st.error(f"❌ 请求失败: {str(e)}")
        else:
            st.warning("⚠️ 请填写完整信息")

    st.divider()

    # 查询知识库
    st.subheader("查询知识库")
    col1, col2 = st.columns([3, 1])

    with col1:
        query_kb_id = st.number_input("知识库 ID", min_value=1, value=1, step=1)
    with col2:
        st.write("")  # 占位
        st.write("")  # 占位
        query_btn = st.button("查询")

    if query_btn:
        try:
            response = requests.get(
                f"{API_BASE_URL}/v1/knowledge_base",
                params={"knowledge_id": query_kb_id, "token": TOKEN}
            )
            if response.status_code == 200:
                data = response.json()
                st.success("✅ 查询成功")

                # 使用卡片展示
                col1, col2, col3 = st.columns(3)
                with col1:
                    st.metric("知识库 ID", data['knowledge_id'])
                with col2:
                    st.metric("名称", data['title'])
                with col3:
                    st.metric("类别", data['category'])

                st.json(data)
            else:
                st.error(f"❌ 查询失败: {response.text}")
        except Exception as e:
            st.error(f"❌ 请求失败: {str(e)}")

    st.divider()

    # 删除知识库
    st.subheader("删除知识库")
    st.warning("⚠️ 删除知识库前，请确保已删除其下的所有文档")

    delete_kb_id = st.number_input("要删除的知识库 ID", min_value=1, value=1, step=1, key="delete_kb")

    # 二次确认
    confirm = st.checkbox(f"✅ 我确认要删除知识库 ID: {delete_kb_id}", key="confirm_delete_kb")

    col1, col2, col3 = st.columns([1, 1, 3])
    with col1:
        delete_kb_btn = st.button("🗑️ 确认删除", type="primary", disabled=not confirm, key="delete_kb_btn")
    with col2:
        if st.button("取消"):
            st.rerun()

    if delete_kb_btn and confirm:
        try:
            response = requests.delete(
                f"{API_BASE_URL}/v1/knowledge_base",
                params={"knowledge_id": delete_kb_id, "token": TOKEN}
            )
            if response.status_code == 200:
                data = response.json()
                st.success(f"✅ 知识库 {delete_kb_id} 删除成功！")
                st.json(data)
            else:
                st.error(f"❌ 删除失败: {response.text}")
        except Exception as e:
            st.error(f"❌ 请求失败: {str(e)}")


# ==================== 文档管理页面 ====================
elif page == "📄 文档管理":
    st.title("📄 文档管理")

    # 上传文档
    st.subheader("上传文档")

    col1, col2, col3 = st.columns(3)
    with col1:
        doc_kb_id = st.number_input("知识库 ID", min_value=1, value=1, step=1, key="upload_kb_id")
    with col2:
        doc_title = st.text_input("文档标题", placeholder="例如：政策文件2024")
    with col3:
        doc_category = st.text_input("文档类别", placeholder="例如：政策文件")

    uploaded_file = st.file_uploader("选择文档", type=['pdf', 'txt', 'docx'])

    if st.button("上传文档", type="primary"):
        if uploaded_file and doc_title and doc_category:
            try:
                files = {'file': (uploaded_file.name, uploaded_file, uploaded_file.type)}
                data = {
                    'knowledge_id': doc_kb_id,
                    'title': doc_title,
                    'category': doc_category,
                    'token': TOKEN
                }

                with st.spinner("上传中..."):
                    response = requests.post(
                        f"{API_BASE_URL}/v1/document",
                        files=files,
                        data=data
                    )

                if response.status_code == 200:
                    result = response.json()
                    st.success(f"✅ 文档上传成功！Document ID: {result['document_id']}")
                    st.json(result)
                else:
                    st.error(f"❌ 上传失败: {response.text}")
            except Exception as e:
                st.error(f"❌ 请求失败: {str(e)}")
        else:
            st.warning("⚠️ 请填写完整信息并选择文件")

    st.divider()

    # 查询文档
    st.subheader("查询文档")
    col1, col2 = st.columns([3, 1])

    with col1:
        query_doc_id = st.number_input("文档 ID", min_value=1, value=1, step=1)
    with col2:
        st.write("")
        st.write("")
        query_doc_btn = st.button("查询", key="query_doc")

    if query_doc_btn:
        try:
            response = requests.get(
                f"{API_BASE_URL}/v1/document",
                params={"document_id": query_doc_id, "token": TOKEN}
            )
            if response.status_code == 200:
                data = response.json()
                st.success("✅ 查询成功")

                col1, col2, col3, col4 = st.columns(4)
                with col1:
                    st.metric("文档 ID", data['document_id'])
                with col2:
                    st.metric("标题", data['title'])
                with col3:
                    st.metric("类别", data['category'])
                with col4:
                    st.metric("知识库 ID", data['knowledge_id'])

                st.json(data)
            else:
                st.error(f"❌ 查询失败: {response.text}")
        except Exception as e:
            st.error(f"❌ 请求失败: {str(e)}")

    st.divider()

    # 删除文档
    st.subheader("删除文档")

    delete_doc_id = st.number_input("要删除的文档 ID", min_value=1, value=1, step=1, key="delete_doc")

    # 二次确认
    confirm_doc = st.checkbox(f"✅ 我确认要删除文档 ID: {delete_doc_id}", key="confirm_delete_doc")

    col1, col2, col3 = st.columns([1, 1, 3])
    with col1:
        delete_doc_btn = st.button("🗑️ 确认删除", type="primary", disabled=not confirm_doc, key="delete_doc_btn")
    with col2:
        if st.button("取消", key="cancel_doc"):
            st.rerun()

    if delete_doc_btn and confirm_doc:
        try:
            response = requests.delete(
                f"{API_BASE_URL}/v1/document",
                params={"document_id": delete_doc_id, "token": TOKEN}
            )
            if response.status_code == 200:
                data = response.json()
                st.success(f"✅ 文档 {delete_doc_id} 删除成功！")
                st.json(data)
            else:
                st.error(f"❌ 删除失败: {response.text}")
        except Exception as e:
            st.error(f"❌ 请求失败: {str(e)}")


# ==================== 智能问答页面 ====================
elif page == "💬 智能问答":
    st.title("💬 智能问答")

    # 选择知识库
    chat_kb_id = st.number_input(
        "选择知识库 ID",
        min_value=1,
        value=1,
        step=1,
        help="选择要查询的知识库"
    )

    # 初始化对话历史
    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)

        # 调用 API
        with st.chat_message("assistant"):
            with st.spinner("思考中..."):
                try:
                    # 构造消息格式
                    messages = [{"role": msg["role"], "content": msg["content"]}
                               for msg in st.session_state.messages]

                    response = requests.post(
                        f"{API_BASE_URL}/chat",
                        json={
                            "knowledge_id": chat_kb_id,
                            "message": messages
                        }
                    )

                    if response.status_code == 200:
                        data = response.json()
                        assistant_message = data['message'][-1]['content']
                        st.markdown(assistant_message)

                        # 添加助手消息
                        st.session_state.messages.append({
                            "role": "assistant",
                            "content": assistant_message
                        })

                        # 显示处理时间
                        st.caption(f"⏱️ 处理时间: {data['processing_time']:.2f}s")
                    else:
                        error_msg = f"❌ 请求失败: {response.text}"
                        st.error(error_msg)
                        st.session_state.messages.append({
                            "role": "assistant",
                            "content": error_msg
                        })
                except Exception as e:
                    error_msg = f"❌ 请求失败: {str(e)}"
                    st.error(error_msg)
                    st.session_state.messages.append({
                        "role": "assistant",
                        "content": error_msg
                    })

    # 清除对话按钮
    if st.button("🗑️ 清除对话历史"):
        st.session_state.messages = []
        st.rerun()


# ==================== 文本向量化页面 ====================
elif page == "🔍 文本向量化":
    st.title("🔍 文本向量化")

    st.write("将文本转换为向量表示，用于语义相似度计算")

    # 单文本向量化
    st.subheader("单文本向量化")
    text_input = st.text_area("输入文本", placeholder="例如：人工智能的发展趋势", height=100)

    if st.button("生成向量", type="primary"):
        if text_input:
            try:
                with st.spinner("生成中..."):
                    response = requests.post(
                        f"{API_BASE_URL}/v1/embedding",
                        json={
                            "text": text_input,
                            "token": TOKEN,
                            "model": MODEL
                        }
                    )

                if response.status_code == 200:
                    data = response.json()
                    vector = data['vector'][0]

                    st.success("✅ 向量生成成功")

                    col1, col2 = st.columns(2)
                    with col1:
                        st.metric("向量维度", len(vector))
                    with col2:
                        st.metric("处理时间", f"{data['processing_time']:.4f}s")

                    # 显示向量的前10个值
                    st.subheader("向量预览（前10个维度）")
                    preview_df = pd.DataFrame({
                        "维度": range(1, 11),
                        "值": vector[:10]
                    })
                    st.dataframe(preview_df, use_container_width=True)

                    # 完整向量（可展开）
                    with st.expander("查看完整向量"):
                        st.json(vector)
                else:
                    st.error(f"❌ 生成失败: {response.text}")
            except Exception as e:
                st.error(f"❌ 请求失败: {str(e)}")
        else:
            st.warning("⚠️ 请输入文本")

    st.divider()

    # 批量文本向量化
    st.subheader("批量文本向量化")
    batch_text = st.text_area(
        "输入多个文本（每行一个）",
        placeholder="人工智能\n机器学习\n深度学习",
        height=150
    )

    if st.button("批量生成向量", type="primary", key="batch"):
        if batch_text:
            texts = [line.strip() for line in batch_text.split('\n') if line.strip()]

            try:
                with st.spinner("批量生成中..."):
                    response = requests.post(
                        f"{API_BASE_URL}/v1/embedding",
                        json={
                            "text": texts,
                            "token": TOKEN,
                            "model": MODEL
                        }
                    )

                if response.status_code == 200:
                    data = response.json()
                    vectors = data['vector']

                    st.success(f"✅ 成功生成 {len(vectors)} 个向量")

                    col1, col2, col3 = st.columns(3)
                    with col1:
                        st.metric("文本数量", len(texts))
                    with col2:
                        st.metric("向量维度", len(vectors[0]))
                    with col3:
                        st.metric("处理时间", f"{data['processing_time']:.4f}s")

                    # 显示每个文本的向量预览
                    for i, (text, vector) in enumerate(zip(texts, vectors)):
                        with st.expander(f"文本 {i+1}: {text}"):
                            st.write(f"向量前5个值: {vector[:5]}")
                else:
                    st.error(f"❌ 生成失败: {response.text}")
            except Exception as e:
                st.error(f"❌ 请求失败: {str(e)}")
        else:
            st.warning("⚠️ 请输入文本")


# 页脚
st.sidebar.divider()
st.sidebar.caption("🏛️ 政府问答 RAG 系统 v1.0")
st.sidebar.caption(f"⏰ {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")