admin_panel.py

import streamlit as st
import os
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
import pandas as pd
import plotly.express as px
from datetime import datetime
import shutil


def show_admin_panel(vector_db):
    """
    管理面板：数据库统计、文档上传、数据库管理
    
    Args:
        vector_db: Chroma向量数据库实例
    """
    st.title("⚙️ Database Management Panel")
    st.caption("Monitor and manage your policy document database")
    
    # 创建标签页
    tab1, tab2, tab3, tab4 = st.tabs([
        "📊 Statistics", 
        "➕ Add Documents", 
        "🔍 Search Test",
        "🗑️ Database Management"
    ])
    
    # ========== TAB 1: 统计信息 ==========
    with tab1:
        st.subheader("Database Overview")
        
        try:
            # 获取所有文档
            all_data = vector_db.get()
            total_chunks = len(all_data['documents'])
            
            if total_chunks == 0:
                st.warning("⚠️ Database is empty. Please add documents in the 'Add Documents' tab.")
                return
            
            # 按来源分组统计
            source_stats = {}
            org_stats = {'IMF': 0, 'World Bank': 0, 'Other': 0}
            
            for metadata in all_data['metadatas']:
                if metadata and 'source' in metadata:
                    source = metadata['source'].split('\\')[-1]
                    source_stats[source] = source_stats.get(source, 0) + 1
                    
                    # 统计组织
                    if 'imf' in source.lower() or 'sdn' in source.lower():
                        org_stats['IMF'] += 1
                    elif 'world bank' in source.lower() or 'digital progress' in source.lower():
                        org_stats['World Bank'] += 1
                    else:
                        org_stats['Other'] += 1
            
            # 关键指标
            col1, col2, col3, col4 = st.columns(4)
            
            with col1:
                st.metric("Total Chunks", f"{total_chunks:,}")
            with col2:
                st.metric("Unique Documents", len(source_stats))
            with col3:
                st.metric("IMF Documents", sum(1 for k in source_stats.keys() if 'imf' in k.lower() or 'sdn' in k.lower()))
            with col4:
                st.metric("WB Documents", sum(1 for k in source_stats.keys() if 'world bank' in k.lower() or 'digital' in k.lower()))
            
            st.markdown("---")
            
            # 文档详情表格
            st.subheader("Document Breakdown")
            
            df = pd.DataFrame([
                {
                    "Document": k, 
                    "Chunks": v, 
                    "Organization": "🏦 IMF" if ("imf" in k.lower() or "sdn" in k.lower()) else 
                                   "🌍 World Bank" if ("world bank" in k.lower() or "digital" in k.lower()) else 
                                   "❓ Other",
                    "Percentage": f"{(v/total_chunks*100):.1f}%"
                }
                for k, v in source_stats.items()
            ]).sort_values("Chunks", ascending=False)
            
            st.dataframe(df, use_container_width=True, hide_index=True)
            
            # 可视化
            col1, col2 = st.columns(2)
            
            with col1:
                # 文档分布柱状图
                fig1 = px.bar(
                    df, 
                    x="Document", 
                    y="Chunks", 
                    color="Organization",
                    title="Document Coverage",
                    color_discrete_map={
                        "🏦 IMF": "#1f4788",
                        "🌍 World Bank": "#00ab51",
                        "❓ Other": "#cccccc"
                    }
                )
                fig1.update_layout(xaxis_tickangle=-45, height=400)
                st.plotly_chart(fig1, use_container_width=True)
            
            with col2:
                # 组织分布饼图
                org_df = pd.DataFrame([
                    {"Organization": k, "Chunks": v} 
                    for k, v in org_stats.items() if v > 0
                ])
                
                fig2 = px.pie(
                    org_df, 
                    values="Chunks", 
                    names="Organization",
                    title="Distribution by Organization",
                    color="Organization",
                    color_discrete_map={
                        "IMF": "#1f4788",
                        "World Bank": "#00ab51",
                        "Other": "#cccccc"
                    }
                )
                fig2.update_layout(height=400)
                st.plotly_chart(fig2, use_container_width=True)
            
            # 健康检查
            st.markdown("---")
            st.subheader("Database Health Check")
            
            health_col1, health_col2, health_col3 = st.columns(3)
            
            with health_col1:
                if total_chunks > 100:
                    st.success("✅ Sufficient data coverage")
                else:
                    st.warning("⚠️ Limited data - consider adding more documents")
            
            with health_col2:
                if len(source_stats) >= 2:
                    st.success("✅ Multiple sources available")
                else:
                    st.info("ℹ️ Single source detected")
            
            with health_col3:
                imf_ratio = org_stats['IMF'] / total_chunks if total_chunks > 0 else 0
                wb_ratio = org_stats['World Bank'] / total_chunks if total_chunks > 0 else 0
                
                if abs(imf_ratio - wb_ratio) < 0.3:
                    st.success("✅ Balanced org coverage")
                else:
                    st.warning("⚠️ Imbalanced coverage detected")
        
        except Exception as e:
            st.error(f"❌ Error loading database statistics: {str(e)}")
    
    # ========== TAB 2: 添加文档 ==========
    with tab2:
        st.subheader("Upload New Policy Documents")
        st.caption("Add PDF files from IMF, World Bank, or other international organizations")
        
        uploaded_files = st.file_uploader(
            "Choose PDF files",
            type="pdf",
            accept_multiple_files=True,
            help="Upload policy documents in PDF format"
        )
        
        if uploaded_files:
            st.write(f"**{len(uploaded_files)} file(s) selected:**")
            for file in uploaded_files:
                st.write(f"- {file.name} ({file.size / 1024:.1f} KB)")
            
            col1, col2 = st.columns([1, 3])
            
            with col1:
                process_button = st.button("📥 Process & Add to Database", type="primary")
            
            if process_button:
                progress_bar = st.progress(0)
                status_text = st.empty()
                
                try:
                    # 确保docs文件夹存在
                    if not os.path.exists("./docs"):
                        os.makedirs("./docs")
                    
                    # 保存文件
                    status_text.text("Saving uploaded files...")
                    for i, uploaded_file in enumerate(uploaded_files):
                        file_path = os.path.join("./docs", uploaded_file.name)
                        with open(file_path, "wb") as f:
                            f.write(uploaded_file.getvalue())
                        progress_bar.progress((i + 1) / (len(uploaded_files) * 2))
                    
                    # 处理文档
                    status_text.text("Processing documents and creating embeddings...")
                    
                    documents = []
                    for i, uploaded_file in enumerate(uploaded_files):
                        file_path = os.path.join("./docs", uploaded_file.name)
                        loader = PyPDFLoader(file_path)
                        documents.extend(loader.load())
                        progress_bar.progress(0.5 + (i + 1) / (len(uploaded_files) * 2))
                    
                    # 切分文档
                    text_splitter = RecursiveCharacterTextSplitter(
                        chunk_size=1000,
                        chunk_overlap=200,
                        separators=["\n\n", "\n", " ", ""],
                    )
                    chunks = text_splitter.split_documents(documents)
                    
                    # 添加到数据库
                    status_text.text("Adding to vector database...")
                    embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
                    
                    # 添加新文档到现有数据库
                    vector_db.add_documents(chunks)
                    
                    progress_bar.progress(1.0)
                    status_text.empty()
                    
                    st.success(f"✅ Successfully added {len(uploaded_files)} document(s) ({len(chunks)} chunks) to the database!")
                    st.info("💡 Refresh the Statistics tab to see updated metrics")
                    
                    # 自动刷新
                    st.rerun()
                
                except Exception as e:
                    st.error(f"❌ Error processing documents: {str(e)}")
                    status_text.empty()
                    progress_bar.empty()
    
    # ========== TAB 3: 搜索测试 ==========
    with tab3:
        st.subheader("Test Document Retrieval")
        st.caption("Verify that your documents are being retrieved correctly")
        
        test_query = st.text_input(
            "Enter a test query:",
            placeholder="e.g., climate finance, debt sustainability, AI risks"
        )
        
        k_value = st.slider("Number of results to retrieve:", 1, 20, 5)
        
        if test_query and st.button("🔍 Search"):
            with st.spinner("Searching..."):
                results = vector_db.similarity_search(test_query, k=k_value)
                
                if results:
                    st.success(f"Found {len(results)} relevant chunks")
                    
                    # 统计来源分布
                    source_dist = {}
                    for doc in results:
                        source = doc.metadata.get('source', 'Unknown').split('\\')[-1]
                        source_dist[source] = source_dist.get(source, 0) + 1
                    
                    st.write("**Source Distribution:**")
                    for source, count in source_dist.items():
                        org = "🏦 IMF" if ("imf" in source.lower() or "sdn" in source.lower()) else "🌍 World Bank"
                        st.write(f"{org} {source}: {count} result(s)")
                    
                    st.markdown("---")
                    
                    # 显示结果
                    for i, doc in enumerate(results, 1):
                        with st.expander(f"Result {i} - Page {doc.metadata.get('page', 'N/A')}"):
                            source = doc.metadata.get('source', 'Unknown').split('\\')[-1]
                            st.markdown(f"**Source:** {source}")
                            st.markdown(f"**Page:** {doc.metadata.get('page', 'N/A')}")
                            st.info(doc.page_content[:500] + "...")
                else:
                    st.warning("No results found for this query")
    
    # ========== TAB 4: 数据库管理 ==========
    with tab4:
        st.subheader("Database Management")
        st.warning("⚠️ **Danger Zone**: These actions are irreversible")
        
        st.markdown("---")
        
        # 导出数据库信息
        st.write("**Export Database Info**")
        if st.button("📤 Export Metadata to CSV"):
            try:
                all_data = vector_db.get()
                
                export_data = []
                for i, (doc, meta) in enumerate(zip(all_data['documents'], all_data['metadatas'])):
                    export_data.append({
                        'chunk_id': i,
                        'source': meta.get('source', 'Unknown').split('\\')[-1] if meta else 'Unknown',
                        'page': meta.get('page', 'N/A') if meta else 'N/A',
                        'content_preview': doc[:100] + '...' if len(doc) > 100 else doc
                    })
                
                df_export = pd.DataFrame(export_data)
                csv = df_export.to_csv(index=False).encode('utf-8')
                
                st.download_button(
                    label="📥 Download CSV",
                    data=csv,
                    file_name=f"database_metadata_{datetime.now().strftime('%Y%m%d')}.csv",
                    mime="text/csv"
                )
            except Exception as e:
                st.error(f"Error exporting data: {str(e)}")
        
        st.markdown("---")
        
        # 重建数据库
        st.write("**Rebuild Database**")
        st.info("This will re-process all PDFs in the docs/ folder")
        
        if st.button("🔄 Rebuild Database from docs/ folder"):
            with st.spinner("Rebuilding database..."):
                try:
                    # 删除旧数据库
                    if os.path.exists("./chroma_db"):
                        shutil.rmtree("./chroma_db")
                    
                    # 运行ingest.py的逻辑
                    from langchain_community.document_loaders import PyPDFLoader
                    from langchain_text_splitters import RecursiveCharacterTextSplitter
                    from langchain_openai import OpenAIEmbeddings
                    from langchain_community.vectorstores import Chroma
                    
                    documents = []
                    for file in os.listdir("./docs"):
                        if file.endswith(".pdf"):
                            file_path = os.path.join("./docs", file)
                            loader = PyPDFLoader(file_path)
                            documents.extend(loader.load())
                    
                    if not documents:
                        st.error("No PDF files found in docs/ folder")
                    else:
                        text_splitter = RecursiveCharacterTextSplitter(
                            chunk_size=1000,
                            chunk_overlap=200,
                            separators=["\n\n", "\n", " ", ""],
                        )
                        chunks = text_splitter.split_documents(documents)
                        
                        embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
                        vector_db_new = Chroma.from_documents(
                            documents=chunks, 
                            embedding=embeddings, 
                            persist_directory="./chroma_db"
                        )
                        
                        st.success(f"✅ Database rebuilt with {len(chunks)} chunks from {len(documents)} pages")
                        st.info("Please refresh the page to see the updated database")
                
                except Exception as e:
                    st.error(f"Error rebuilding database: {str(e)}")
        
        st.markdown("---")
        
        # 清空数据库
        st.write("**Clear Database**")
        st.error("⚠️ This will permanently delete all documents and embeddings!")
        
        confirm_clear = st.checkbox("I understand this action cannot be undone")
        
        if confirm_clear:
            if st.button("🗑️ Clear Entire Database", type="secondary"):
                try:
                    shutil.rmtree("./chroma_db")
                    os.makedirs("./chroma_db")
                    st.success("✅ Database cleared successfully")
                    st.info("Please re-run `python ingest.py` or use 'Rebuild Database' to add documents")
                    st.rerun()
                except Exception as e:
                    st.error(f"Error clearing database: {str(e)}")


# 测试函数
if __name__ == "__main__":
    print("Admin panel module loaded successfully")