Upload 2 files

app.py

@@ -0,0 +1,1123 @@
+import streamlit as st
+import time as _time
+_BOOT = _time.time()
+import json
+import time
+import logging
+import hashlib
+from datetime import datetime
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
+logger = logging.getLogger(__name__)
+
+def _perf(label):
+    logger.info(f"[PERF] {label}: {_time.time()-_BOOT:.2f}s")
+
+_perf("stdlib imports done")
+
+# numpy 延迟导入
+_np_module = None
+def _get_np():
+    global _np_module
+    if _np_module is None:
+        import numpy
+        _np_module = numpy
+        _perf("numpy loaded")
+    return _np_module
+
+# =========================
+# 1. 页面配置 & 样式注入
+# =========================
+st.set_page_config(page_title="RAG 知识库助手 v3 (HF+Supabase)", page_icon="🛡️", layout="wide")
+_perf("page_config done")
+
+
+def inject_custom_css():
+    st.markdown("""
+        <style>
+            [data-testid="stSidebarContent"] { padding-top: 1.5rem !important; }
+            [data-testid="stVerticalBlock"] > div { gap: 0.8rem !important; }
+            [data-testid="stFileUploader"] section > div { display: none; }
+            [data-testid="stFileUploaderDropzoneInstructions"] { display: none !important; }
+            [data-testid="stFileUploader"] section::before {
+                content: "拖拽文档至此";
+                color: #555; font-size: 14px; display: block; margin-bottom: 10px;
+            }
+            [data-testid="stFileUploader"] section::after {
+                content: "支持格式：TXT, PDF, DOCX";
+                color: #888; font-size: 12px; display: block; margin-top: 5px;
+            }
+            [data-testid="stFileUploader"] button { font-size: 0 !important; }
+            [data-testid="stFileUploader"] button::after {
+                content: "选择文件";
+                font-size: 14px !important;
+            }
+        </style>
+    """, unsafe_allow_html=True)
+
+
+inject_custom_css()
+st.title("🛡️ 智能知识库助手 v3")
+_perf("CSS + title done")
+
+# =========================
+# 1.5 Supabase 客户端初始化
+# =========================
+from supabase import create_client
+
+SUPABASE_URL = st.secrets.get("SUPABASE_URL", "")
+SUPABASE_KEY = st.secrets.get("SUPABASE_KEY", "")  # service_role key (后端使用)
+STORAGE_BUCKET = "rag-files"
+
+if not SUPABASE_URL or not SUPABASE_KEY:
+    st.error("⚠️ 未配置 SUPABASE_URL 或 SUPABASE_KEY，请在 Secrets 中设置。")
+    st.stop()
+
+
+@st.cache_resource
+def _get_supabase():
+    return create_client(SUPABASE_URL, SUPABASE_KEY)
+
+
+def _sb():
+    """快捷获取 Supabase 客户端。"""
+    return _get_supabase()
+
+
+_perf("supabase client ready")
+
+# =========================
+# 2. 用户管理（Supabase users 表）
+# =========================
+MAX_LOGIN_ATTEMPTS = 10
+
+
+def _hash_password(password):
+    return hashlib.sha256(password.encode("utf-8")).hexdigest()
+
+
+def _load_users():
+    """从 Supabase users 表加载所有用户，返回 {username: {password_hash, role, created_at}}。"""
+    try:
+        resp = _sb().table("users").select("*").execute()
+        users = {}
+        for row in resp.data:
+            users[row["username"]] = {
+                "password_hash": row["password_hash"],
+                "role": row["role"],
+                "created_at": row["created_at"][:16] if row.get("created_at") else "未知",
+            }
+        return users
+    except Exception as e:
+        logger.error(f"加载用户表失败: {e}")
+        return {}
+
+
+def _ensure_admin():
+    """首次运行时，从 secrets 创建管理员（如果 users 表为空）。"""
+    users = _load_users()
+    if users:
+        return
+    admin_user = st.secrets.get("ADMIN_USER", "admin")
+    admin_pass = st.secrets.get("ADMIN_PASSWORD", "")
+    if not admin_pass:
+        return
+    try:
+        _sb().table("users").upsert({
+            "username": admin_user,
+            "password_hash": _hash_password(admin_pass),
+            "role": "admin",
+        }).execute()
+        logger.info(f"初始管理员 {admin_user} 已创建")
+    except Exception as e:
+        logger.error(f"创建管理员失败: {e}")
+
+
+_ensure_admin()
+
+
+def _save_user(username, password_hash, role="user"):
+    """新增或更新单个用户。"""
+    _sb().table("users").upsert({
+        "username": username,
+        "password_hash": password_hash,
+        "role": role,
+    }).execute()
+
+
+def _delete_user_db(username):
+    """删除用户记录。"""
+    _sb().table("users").delete().eq("username", username).execute()
+
+
+def _get_invite_code():
+    """从 app_meta 表读取邀请码。"""
+    try:
+        resp = _sb().table("app_meta").select("value").eq("key", "invite_code").execute()
+        if resp.data:
+            return resp.data[0]["value"]
+    except Exception:
+        pass
+    return st.secrets.get("INVITE_CODE", "")
+
+
+def _set_invite_code(new_code):
+    _sb().table("app_meta").upsert({"key": "invite_code", "value": new_code}).execute()
+
+
+def register_user(username, password, invite_code):
+    if not username or not password:
+        return False, "用户名和密码不能为空"
+    if len(username) < 2 or len(username) > 20:
+        return False, "用户名长度需要 2-20 个字符"
+    if len(password) < 4:
+        return False, "密码至少 4 个字符"
+    if username.startswith("__"):
+        return False, "用户名不能以 __ 开头"
+
+    correct_code = _get_invite_code()
+    if not correct_code:
+        return False, "邀请码未配置，请联系管理员"
+    if invite_code != correct_code:
+        return False, "邀请码错误"
+
+    users = _load_users()
+    if username in users:
+        return False, "用户名已存在"
+
+    try:
+        _save_user(username, _hash_password(password), "user")
+        logger.info(f"新用户注册: {username}")
+        return True, "注册成功，请登录"
+    except Exception as e:
+        logger.error(f"注册失败: {e}")
+        return False, f"注册失败: {e}"
+
+
+def verify_user(username, password):
+    users = _load_users()
+    user_info = users.get(username)
+    if not user_info or not isinstance(user_info, dict):
+        return False, None
+    if user_info.get("password_hash") != _hash_password(password):
+        return False, None
+    return True, user_info.get("role", "user")
+
+
+# --- 认证 UI ---
+if "login_attempts" not in st.session_state:
+    st.session_state.login_attempts = 0
+if "current_user" not in st.session_state:
+    st.session_state.current_user = None
+if "current_role" not in st.session_state:
+    st.session_state.current_role = None
+if "auth_mode" not in st.session_state:
+    st.session_state.auth_mode = "login"
+
+with st.sidebar:
+    with st.expander("🔑 账号"):
+        if st.session_state.login_attempts >= MAX_LOGIN_ATTEMPTS:
+            st.error("🚫 尝试次数过多，请刷新页面后重试。")
+            st.stop()
+
+        users_data = _load_users()
+        if not users_data:
+            st.error("⚠️ 未配置管理员，请在 secrets 中设置 ADMIN_USER 和 ADMIN_PASSWORD。")
+            st.stop()
+
+        auth_mode = st.radio(
+            "操作", ["登录", "注册"], horizontal=True,
+            label_visibility="collapsed", key="auth_radio",
+        )
+
+        if auth_mode == "登录":
+            input_username = st.text_input("用户名", key="login_user")
+            input_password = st.text_input("密码", type="password", key="login_pass")
+
+            if input_username == "" or input_password == "":
+                st.stop()
+
+            ok, role = verify_user(input_username, input_password)
+            if not ok:
+                st.session_state.login_attempts += 1
+                remaining = MAX_LOGIN_ATTEMPTS - st.session_state.login_attempts
+                st.warning(f"⚠️ 用户名或密码错误（剩余 {remaining} 次）")
+                st.stop()
+            else:
+                st.session_state.login_attempts = 0
+                st.session_state.current_user = input_username
+                st.session_state.current_role = role
+                role_label = "管理员" if role == "admin" else "普通用户"
+                st.success(f"✅ {input_username}（{role_label}）")
+
+        else:  # 注册
+            reg_user = st.text_input("用户名", key="reg_user")
+            reg_pass = st.text_input("密码", type="password", key="reg_pass")
+            reg_pass2 = st.text_input("确认密码", type="password", key="reg_pass2")
+            reg_code = st.text_input("邀请码", type="password", key="reg_code")
+
+            if st.button("注册", use_container_width=True, key="btn_register"):
+                if reg_pass != reg_pass2:
+                    st.error("两次密码不一致")
+                else:
+                    ok, msg = register_user(reg_user, reg_pass, reg_code)
+                    if ok:
+                        st.success(f"✅ {msg}")
+                        time.sleep(1)
+                        st.rerun()
+                    else:
+                        st.error(f"❌ {msg}")
+            st.stop()
+
+CURRENT_USER = st.session_state.current_user
+IS_ADMIN = st.session_state.current_role == "admin"
+_perf("auth done")
+
+# =========================
+# 3. 安全配置与 Embedding 策略
+# =========================
+TAVILY_KEY = st.secrets.get("TAVILY_API_KEY", "")
+DS_API_KEY = st.secrets.get("DEEPSEEK_API_KEY", "")
+BAIDU_TOKEN = st.secrets.get("BAIDU_BEARER_TOKEN", "")
+BAIDU_APP_ID = st.secrets.get("BAIDU_APP_ID", "")
+OR_KEY = st.secrets.get("OPENROUTER_API_KEY", "")
+
+
+@st.cache_resource
+def _get_embedding_client():
+    from openai import OpenAI
+    if BAIDU_TOKEN and BAIDU_APP_ID:
+        return OpenAI(
+            api_key=BAIDU_TOKEN,
+            base_url="https://qianfan.baidubce.com/v2",
+            default_headers={"appid": BAIDU_APP_ID},
+        ), "bge-large-zh"
+    if OR_KEY:
+        return OpenAI(
+            api_key=OR_KEY,
+            base_url="https://openrouter.ai/api/v1",
+        ), "BAAI/bge-small-zh"
+    return None, None
+
+
+def _api_encode(texts):
+    np = _get_np()
+    client, model = _get_embedding_client()
+    if client is None:
+        return None
+    try:
+        batch_size = 32
+        all_vecs = []
+        for i in range(0, len(texts), batch_size):
+            batch = texts[i:i + batch_size]
+            resp = client.embeddings.create(model=model, input=batch)
+            all_vecs.extend([np.array(item.embedding) for item in resp.data])
+        return all_vecs
+    except Exception as e:
+        logger.warning(f"API embedding 失败，回退到本地模型: {e}")
+        return None
+
+
+def _get_local_model():
+    if "_local_emb_model" not in st.session_state:
+        try:
+            with st.spinner("API 不可用，正在加载本地向量模型（仅首次）..."):
+                from sentence_transformers import SentenceTransformer
+                st.session_state._local_emb_model = SentenceTransformer("BAAI/bge-small-zh")
+        except ImportError:
+            logger.error("sentence-transformers 未安装，本地模型不可用")
+            return None
+    return st.session_state.get("_local_emb_model")
+
+
+def encode_texts(texts):
+    if not texts:
+        return []
+    if isinstance(texts, str):
+        texts = [texts]
+    result = _api_encode(texts)
+    if result is not None:
+        return result
+    model = _get_local_model()
+    if model is None:
+        st.error("❌ Embedding 服务不可用：API 调用失败且本地模型未安装。请检查 API Key 配置。")
+        return []
+    return list(model.encode(texts))
+
+
+def encode_query(text):
+    vecs = encode_texts([text])
+    return vecs[0]
+
+
+# =========================
+# 4. Supabase 索引管理（替代本地文件）
+# =========================
+
+def _load_library(scope):
+    """从 Supabase documents 表加载指定 scope 的所有文档切片。
+    返回 (docs, embeddings, sources)。"""
+    np = _get_np()
+    try:
+        resp = _sb().table("documents").select(
+            "content, embedding, source_file"
+        ).eq("scope", scope).execute()
+
+        docs = []
+        embeddings = []
+        sources = []
+        for row in resp.data:
+            docs.append(row["content"])
+            embeddings.append(np.array(row["embedding"]))
+            sources.append(row["source_file"])
+        return docs, embeddings, sources
+    except Exception as e:
+        logger.error(f"加载索引失败 [scope={scope}]: {e}")
+        return [], [], []
+
+
+def _save_chunks_to_db(scope, chunks, vectors, source_file):
+    """将新切片批量写入 Supabase documents 表。"""
+    rows = []
+    for content, vec, src in zip(chunks, vectors, [source_file] * len(chunks)):
+        rows.append({
+            "scope": scope,
+            "source_file": src,
+            "content": content,
+            "embedding": vec.tolist() if hasattr(vec, 'tolist') else list(vec),
+        })
+    # Supabase 批量插入（每次最多 500 行）
+    batch_size = 500
+    for i in range(0, len(rows), batch_size):
+        _sb().table("documents").insert(rows[i:i + batch_size]).execute()
+
+
+def _delete_chunks_by_file(scope, filename):
+    """删除指定 scope + filename 的所有切片。"""
+    _sb().table("documents").delete().eq("scope", scope).eq("source_file", filename).execute()
+
+
+def _clear_all_chunks(scope):
+    """清空指定 scope 的所有文档切片。"""
+    _sb().table("documents").delete().eq("scope", scope).execute()
+
+
+def _count_chunks(scope):
+    """返回指定 scope 的切片数量。"""
+    try:
+        resp = _sb().table("documents").select("id", count="exact").eq("scope", scope).execute()
+        return resp.count or 0
+    except Exception:
+        return 0
+
+
+# --- 原始文件管理（Supabase Storage + uploaded_files 表）---
+
+def _save_uploaded_file_to_storage(scope, uploaded_file):
+    """上传原始文件到 Supabase Storage，并记录元数据。"""
+    storage_path = f"{scope}/{uploaded_file.name}"
+    uploaded_file.seek(0)
+    file_bytes = uploaded_file.read()
+
+    # 上传到 Storage（存在则覆盖）
+    try:
+        _sb().storage.from_(STORAGE_BUCKET).upload(
+            storage_path, file_bytes,
+            file_options={"content-type": "application/octet-stream", "upsert": "true"}
+        )
+    except Exception as e:
+        # supabase-py 某些版本 upsert 需要先删再传
+        logger.warning(f"Storage upload fallback: {e}")
+        try:
+            _sb().storage.from_(STORAGE_BUCKET).remove([storage_path])
+        except Exception:
+            pass
+        _sb().storage.from_(STORAGE_BUCKET).upload(
+            storage_path, file_bytes,
+            file_options={"content-type": "application/octet-stream"}
+        )
+
+    # 记录元数据到 uploaded_files 表
+    _sb().table("uploaded_files").upsert({
+        "scope": scope,
+        "filename": uploaded_file.name,
+        "file_size": len(file_bytes),
+        "storage_path": storage_path,
+    }, on_conflict="scope,filename").execute()
+
+
+def _list_uploaded_files_db(scope):
+    """列出某个 scope 已上传的文件。返回 [(filename, size_str, storage_path), ...]。"""
+    try:
+        resp = _sb().table("uploaded_files").select(
+            "filename, file_size, storage_path"
+        ).eq("scope", scope).order("filename").execute()
+
+        result = []
+        for row in resp.data:
+            size = row.get("file_size", 0) or 0
+            if size < 1024:
+                size_str = f"{size}B"
+            elif size < 1048576:
+                size_str = f"{size / 1024:.1f}KB"
+            else:
+                size_str = f"{size / 1048576:.1f}MB"
+            result.append((row["filename"], size_str, row.get("storage_path", "")))
+        return result
+    except Exception as e:
+        logger.error(f"列出文件失败 [scope={scope}]: {e}")
+        return []
+
+
+def _delete_uploaded_file_from_storage(scope, filename):
+    """删除 Storage 中的文件和 uploaded_files 表记录。"""
+    storage_path = f"{scope}/{filename}"
+    try:
+        _sb().storage.from_(STORAGE_BUCKET).remove([storage_path])
+    except Exception as e:
+        logger.warning(f"Storage 删除失败: {e}")
+    try:
+        _sb().table("uploaded_files").delete().eq("scope", scope).eq("filename", filename).execute()
+    except Exception as e:
+        logger.warning(f"uploaded_files 记录删除失败: {e}")
+
+
+def _clear_uploaded_files_storage(scope):
+    """清空某个 scope 的所有上传文件。"""
+    files = _list_uploaded_files_db(scope)
+    paths = [f["storage_path"] for f in files] if files else []
+    # 从 uploaded_files 表拿 storage_path
+    try:
+        resp = _sb().table("uploaded_files").select("storage_path").eq("scope", scope).execute()
+        paths = [row["storage_path"] for row in resp.data]
+        if paths:
+            _sb().storage.from_(STORAGE_BUCKET).remove(paths)
+        _sb().table("uploaded_files").delete().eq("scope", scope).execute()
+    except Exception as e:
+        logger.warning(f"清空文件失败 [scope={scope}]: {e}")
+
+
+# --- 初始化 session_state 中的缓存 ---
+def _init_library(key_prefix, scope):
+    """加载 Supabase 中的索引到 session_state。"""
+    docs_key = f"{key_prefix}_docs"
+    emb_key = f"{key_prefix}_embeddings"
+    src_key = f"{key_prefix}_sources"
+    loaded_key = f"{key_prefix}_loaded"
+
+    if docs_key not in st.session_state or not st.session_state.get(loaded_key):
+        docs, embeddings, sources = _load_library(scope)
+        st.session_state[docs_key] = docs
+        st.session_state[emb_key] = embeddings
+        st.session_state[src_key] = sources
+        st.session_state[loaded_key] = True
+
+
+def _refresh_library(key_prefix, scope):
+    """强制从 Supabase 重新加载索引到 session_state。"""
+    docs, embeddings, sources = _load_library(scope)
+    st.session_state[f"{key_prefix}_docs"] = docs
+    st.session_state[f"{key_prefix}_embeddings"] = embeddings
+    st.session_state[f"{key_prefix}_sources"] = sources
+
+
+_perf("before init_library")
+PUBLIC_SCOPE = "public"
+_init_library("public", PUBLIC_SCOPE)
+PRIVATE_SCOPE = CURRENT_USER  # 私有库 scope = 用户名
+_init_library("private", PRIVATE_SCOPE)
+_perf("init_library done")
+
+
+def _get_embeddings_np(key_prefix):
+    np = _get_np()
+    np_key = f"{key_prefix}_embeddings_np"
+    ver_key = f"{key_prefix}_emb_version"
+    emb_key = f"{key_prefix}_embeddings"
+    emb_list = st.session_state.get(emb_key, [])
+    current_ver = id(emb_list)
+    if np_key not in st.session_state or st.session_state.get(ver_key) != current_ver:
+        if emb_list:
+            st.session_state[np_key] = np.array(emb_list)
+        else:
+            st.session_state[np_key] = np.array([])
+        st.session_state[ver_key] = current_ver
+    return st.session_state[np_key]
+
+
+# =========================
+# 5. 缓存 LLM 客户端
+# =========================
+@st.cache_resource
+def get_or_client():
+    from openai import OpenAI
+    return OpenAI(api_key=OR_KEY, base_url="https://openrouter.ai/api/v1")
+
+
+@st.cache_resource
+def get_ds_client():
+    from openai import OpenAI
+    return OpenAI(api_key=DS_API_KEY, base_url="https://api.deepseek.com")
+
+
+@st.cache_resource
+def get_baidu_client():
+    from openai import OpenAI
+    return OpenAI(
+        api_key=BAIDU_TOKEN,
+        base_url="https://qianfan.baidubce.com/v2",
+        default_headers={"appid": BAIDU_APP_ID},
+    )
+
+
+# =========================
+# 6. 实用功能函数
+# =========================
+_text_splitter_cache = None
+def _get_text_splitter():
+    global _text_splitter_cache
+    if _text_splitter_cache is None:
+        from langchain_text_splitters import RecursiveCharacterTextSplitter
+        _text_splitter_cache = RecursiveCharacterTextSplitter(chunk_size=400, chunk_overlap=50)
+        _perf("text_splitter loaded")
+    return _text_splitter_cache
+
+SYSTEM_PROMPT = (
+    "你是一个专业的知识问答助手。请基于提供的参考资料回答用户问题。"
+    "如果资料中没有相关信息，请诚实说明。回答要准确、有条理、简洁。"
+    "不要编造不在资料中的信息。"
+)
+
+
+def web_search(query):
+    if not TAVILY_KEY:
+        return "⚠️ 未配置搜索 Key"
+    from tavily import TavilyClient
+    tavily = TavilyClient(api_key=TAVILY_KEY)
+    current_year = datetime.now().year
+    try:
+        search_result = tavily.search(
+            query=f"{current_year}年 {query}",
+            search_depth="advanced",
+            max_results=3,
+        )
+        results = [
+            f"来源: {r.get('url')}\n内容: {r.get('content', '')[:700]}"
+            for r in search_result["results"]
+        ]
+        return "\n\n".join(results)[:2500]
+    except Exception as e:
+        logger.error(f"联网搜索异常: {e}")
+        return f"联网搜索异常：{str(e)}"
+
+
+def estimate_tokens(text):
+    if not text:
+        return 0
+    zh_count = sum(1 for c in text if "\u4e00" <= c <= "\u9fff")
+    return int(zh_count * 1.5 + (len(text) - zh_count) * 0.4)
+
+
+def extract_text(file):
+    fname = file.name.lower()
+    text = ""
+    try:
+        if fname.endswith(".txt"):
+            text = file.read().decode("utf-8", errors="ignore")
+        elif fname.endswith(".pdf"):
+            import io
+            file.seek(0)
+            pdf_bytes = file.read()
+            if len(pdf_bytes) < 100:
+                raise ValueError("PDF 文件过小，可能已损坏")
+            if not pdf_bytes[:5] == b"%PDF-":
+                raise ValueError("不是有效的 PDF 文件（缺少 %PDF- 头）")
+            pdf_stream = io.BytesIO(pdf_bytes)
+            pages_text = []
+            import pdfplumber
+            with pdfplumber.open(pdf_stream) as pdf:
+                for i, page in enumerate(pdf.pages):
+                    try:
+                        page_text = page.extract_text() or ""
+                        pages_text.append(page_text)
+                    except Exception as page_err:
+                        logger.warning(f"PDF 第{i+1}页解析失败: {page_err}")
+                        pages_text.append("")
+            text = "\n".join(pages_text)
+        elif fname.endswith(".docx"):
+            from docx import Document
+            doc = Document(file)
+            text = "\n".join(para.text for para in doc.paragraphs)
+    except Exception as e:
+        logger.error(f"文件解析失败 [{file.name}]: {e}", exc_info=True)
+        st.error(f"解析失败: {e}")
+    return text
+
+
+def process_upload(uploaded_files, target_prefix, scope):
+    """处理上传文件：解析 → 切片 → 编码 → 写入 Supabase。"""
+    if not uploaded_files:
+        return False
+    file_fingerprint = str(sorted((f.name, f.size) for f in uploaded_files))
+    fp_key = f"_last_upload_fp_{target_prefix}"
+    if file_fingerprint == st.session_state.get(fp_key):
+        return False
+
+    try:
+        all_new_chunks = []
+        all_new_sources = []
+        with st.spinner("正在自动解析文档并更新索引..."):
+            for f in uploaded_files:
+                try:
+                    # 保存原始文件到 Supabase Storage
+                    f.seek(0)
+                    _save_uploaded_file_to_storage(scope, f)
+
+                    # 解析文本
+                    f.seek(0)
+                    raw_text = extract_text(f)
+                    if not raw_text.strip():
+                        st.warning(f"文件 {f.name} 内容为空，已跳过。")
+                        continue
+                    chunks = _get_text_splitter().split_text(raw_text)
+                    all_new_chunks.extend(chunks)
+                    all_new_sources.extend([f.name] * len(chunks))
+                except Exception as file_err:
+                    logger.error(f"文件 {f.name} 处理失败: {file_err}", exc_info=True)
+                    st.warning(f"⚠️ 文件 {f.name} 处理失败：{str(file_err)[:100]}，已跳过。")
+
+            if all_new_chunks:
+                # 分批编码
+                batch_size = 64
+                all_vecs = []
+                for i in range(0, len(all_new_chunks), batch_size):
+                    batch = all_new_chunks[i:i + batch_size]
+                    all_vecs.extend(encode_texts(batch))
+
+                # 按 source_file 分组写入 Supabase
+                file_groups = {}
+                for chunk, vec, src in zip(all_new_chunks, all_vecs, all_new_sources):
+                    file_groups.setdefault(src, ([], []))
+                    file_groups[src][0].append(chunk)
+                    file_groups[src][1].append(vec)
+
+                for src_file, (chunks, vecs) in file_groups.items():
+                    _save_chunks_to_db(scope, chunks, vecs, src_file)
+
+                # 刷新 session_state 缓存
+                _refresh_library(target_prefix, scope)
+
+                st.session_state[fp_key] = file_fingerprint
+                # 递增上传组件 key
+                ukey = f"_upload_ver_{target_prefix}"
+                st.session_state[ukey] = st.session_state.get(ukey, 0) + 1
+                st.success(f"自动导入 {len(all_new_chunks)} 个知识切片！")
+                time.sleep(1)
+                st.rerun()
+            else:
+                st.error("解析失败，未发现有效文字内容。")
+    except Exception as e:
+        logger.error(f"上传处理异常: {e}", exc_info=True)
+        st.error(f"❌ 上传处理出错：{str(e)[:200]}")
+    return False
+
+
+# =========================
+# 7. 侧边栏 UI & 逻辑
+# =========================
+model_mapping = {
+    "⭐ Step-3.5 (首选)": "stepfun/step-3.5-flash:free",
+    "🌐 OR-Auto (避堵)": "openrouter/free",
+    "🧠 GLM-4.5 (推理)": "z-ai/glm-4.5-air:free",
+    "🔥 Gemma-3-27B (旗舰)": "google/gemma-3-27b-it:free",
+    "🐋 Nemotron (120B)": "nvidia/nemotron-3-super-120b-a12b:free",
+    "⚡ Trinity-L (极速)": "arcee-ai/trinity-large-preview:free",
+    "💭 Liquid-Think (思维链)": "liquid/lfm-2.5-1.2b-thinking:free",
+    "🏎️ Liquid-Ins (1.0s)": "liquid/lfm-2.5-1.2b-instruct:free",
+    "⚖️ Gemma-3-12B (平衡)": "google/gemma-3-12b-it:free",
+    "💎 Gemma-3n-e4b (稳)": "google/gemma-3n-e4b-it:free",
+    "🤖 Nemotron-Nano (混)": "nvidia/nemotron-3-nano-30b-a3b:free",
+    "📉 Trinity-M (1.8s)": "arcee-ai/trinity-mini:free",
+    "🍃 Nemotron-9B": "nvidia/nemotron-nano-9b-v2:free",
+    "🪶 Gemma-3-4B": "google/gemma-3-4b-it:free",
+    "🫧 Gemma-3n-e2b": "google/gemma-3n-e2b-it:free",
+    "📷 Nemotron-VL": "nvidia/nemotron-nano-12b-v2-vl:free",
+    "🛡️ DeepSeek (官方)": "deepseek-chat",
+    "🏢 百度文心 (官方)": "ernie-3.5-8k",
+}
+
+_perf("before sidebar UI")
+with st.sidebar:
+    pub_chunk_count = len(st.session_state.get("public_docs", []))
+    with st.expander(f"📚 公共知识库（{pub_chunk_count} 切片）"):
+        st.caption("所有人可搜索")
+
+        # 文件列表
+        pub_file_list = _list_uploaded_files_db(PUBLIC_SCOPE)
+        if pub_file_list:
+            st.caption(f"📎 已上传 {len(pub_file_list)} 个文件：")
+            for fname, size_str, _ in pub_file_list:
+                if IS_ADMIN:
+                    col_name, col_del = st.columns([4, 1])
+                    col_name.text(f"📄 {fname} ({size_str})")
+                    if col_del.button("🗑", key=f"delpub_{fname}", help=f"删除 {fname}"):
+                        _delete_chunks_by_file(PUBLIC_SCOPE, fname)
+                        _delete_uploaded_file_from_storage(PUBLIC_SCOPE, fname)
+                        _refresh_library("public", PUBLIC_SCOPE)
+                        st.success(f"已删除 {fname}")
+                        time.sleep(0.5)
+                        st.rerun()
+                else:
+                    st.text(f"📄 {fname} ({size_str})")
+
+        if IS_ADMIN:
+            pub_upload_key = f"upload_public_{st.session_state.get('_upload_ver_public', 0)}"
+            pub_files = st.file_uploader(
+                "上传到公共库",
+                type=["txt", "pdf", "docx"],
+                accept_multiple_files=True,
+                label_visibility="collapsed",
+                key=pub_upload_key,
+            )
+            if pub_files:
+                process_upload(pub_files, "public", PUBLIC_SCOPE)
+
+            if pub_chunk_count > 0 and len(pub_file_list) >= 2:
+                if st.button("🗑️ 清空公共库", use_container_width=True, type="secondary", key="clear_pub"):
+                    _clear_all_chunks(PUBLIC_SCOPE)
+                    _clear_uploaded_files_storage(PUBLIC_SCOPE)
+                    _refresh_library("public", PUBLIC_SCOPE)
+                    st.success("公共知识库已清空。")
+                    time.sleep(0.5)
+                    st.rerun()
+        else:
+            st.caption("*仅管理员可维护公共库*")
+
+    # --- 私有知识库 ---
+    priv_chunk_count = len(st.session_state.get("private_docs", []))
+    with st.expander(f"🔒 我的私有库（{priv_chunk_count} 切片）"):
+        st.caption(f"用户：{CURRENT_USER}，仅自己可见")
+
+        priv_file_list = _list_uploaded_files_db(PRIVATE_SCOPE)
+        if priv_file_list:
+            st.caption(f"📎 已上传 {len(priv_file_list)} 个文件：")
+            for fname, size_str, _ in priv_file_list:
+                col_name, col_del = st.columns([4, 1])
+                col_name.text(f"📄 {fname} ({size_str})")
+                if col_del.button("🗑", key=f"delpriv_{fname}", help=f"删除 {fname}"):
+                    _delete_chunks_by_file(PRIVATE_SCOPE, fname)
+                    _delete_uploaded_file_from_storage(PRIVATE_SCOPE, fname)
+                    _refresh_library("private", PRIVATE_SCOPE)
+                    st.success(f"已删除 {fname}")
+                    time.sleep(0.5)
+                    st.rerun()
+
+        priv_upload_key = f"upload_private_{st.session_state.get('_upload_ver_private', 0)}"
+        priv_files = st.file_uploader(
+            "上传到私有库",
+            type=["txt", "pdf", "docx"],
+            accept_multiple_files=True,
+            label_visibility="collapsed",
+            key=priv_upload_key,
+        )
+        if priv_files:
+            process_upload(priv_files, "private", PRIVATE_SCOPE)
+
+        if priv_chunk_count > 0 and len(priv_file_list) >= 2:
+            if st.button("🗑️ 清空我的私有库", use_container_width=True, type="secondary", key="clear_priv"):
+                _clear_all_chunks(PRIVATE_SCOPE)
+                _clear_uploaded_files_storage(PRIVATE_SCOPE)
+                _refresh_library("private", PRIVATE_SCOPE)
+                st.success("私有知识库已清空。")
+                time.sleep(0.5)
+                st.rerun()
+
+    # --- 模型设置 ---
+    with st.expander("⚙️ 模型设置"):
+        selected_display_name = st.selectbox(
+            "模型", list(model_mapping.keys()), index=0, label_visibility="collapsed"
+        )
+
+        web_on = st.toggle("🌐 联网增强", value=False)
+
+        c1, c2 = st.columns(2)
+        with c1:
+            ui_top_k = st.number_input("Top-K", 1, 15, 5)
+        with c2:
+            ui_threshold = st.number_input("阈值", 0.0, 1.0, 0.25, step=0.05)
+
+    # --- 修改密码 ---
+    with st.expander("🔐 修改密码"):
+        old_pass = st.text_input("当前密码", type="password", key="self_old_pass")
+        new_pass1 = st.text_input("新密码", type="password", key="self_new_pass1")
+        new_pass2 = st.text_input("确认新密码", type="password", key="self_new_pass2")
+        if st.button("✅ 确认修改", key="btn_change_pass"):
+            ok, _ = verify_user(CURRENT_USER, old_pass)
+            if not ok:
+                st.error("当前密码错误")
+            elif len(new_pass1) < 4:
+                st.error("新密码至少 4 个字符")
+            elif new_pass1 != new_pass2:
+                st.error("两次新密码不一致")
+            else:
+                _save_user(CURRENT_USER, _hash_password(new_pass1),
+                           st.session_state.current_role)
+                st.success("密码修改成功，请重新登录")
+                time.sleep(1)
+                st.rerun()
+
+    # --- 管理员面板 ---
+    if IS_ADMIN:
+        with st.expander("👥 用户管理"):
+            all_users = _load_users()
+            user_list = [(u, info) for u, info in all_users.items() if isinstance(info, dict)]
+
+            st.caption(f"共 **{len(user_list)}** 个用户")
+            for uname, uinfo in user_list:
+                role_tag = "👑" if uinfo.get("role") == "admin" else "👤"
+                created = uinfo.get("created_at", "未知")
+                st.text(f"{role_tag} {uname}（{created}）")
+
+            deletable = [u for u, _ in user_list if u != CURRENT_USER]
+            if deletable:
+                del_target = st.selectbox("选择要删除的用户", deletable, key="del_user_select")
+                if st.button("❌ 删除该用户", key="btn_del_user"):
+                    _delete_user_db(del_target)
+                    # 清除该用户的私有库
+                    _clear_all_chunks(del_target)
+                    _clear_uploaded_files_storage(del_target)
+                    st.success(f"用户 {del_target} 已删除")
+                    time.sleep(0.5)
+                    st.rerun()
+
+            resetable = [u for u, _ in user_list if u != CURRENT_USER]
+            if resetable:
+                reset_target = st.selectbox("选择要重置密码的用户", resetable, key="reset_user_select")
+                new_pass = st.text_input("新密码", type="password", key="reset_new_pass")
+                if st.button("🔄 重置密码", key="btn_reset_pass"):
+                    if len(new_pass) < 4:
+                        st.error("密码至少 4 个字符")
+                    else:
+                        target_role = all_users[reset_target].get("role", "user")
+                        _save_user(reset_target, _hash_password(new_pass), target_role)
+                        st.success(f"用户 {reset_target} 密码已重置")
+                        time.sleep(0.5)
+                        st.rerun()
+
+        with st.expander("📩 邀请码管理"):
+            current_code = _get_invite_code()
+            st.text(f"当前邀请码：{current_code if current_code else '未设置'}")
+            new_code = st.text_input("新邀请码", key="new_invite_code")
+            if st.button("✏️ 更新邀请码", key="btn_update_code"):
+                if new_code.strip():
+                    _set_invite_code(new_code.strip())
+                    st.success("邀请码已更新")
+                    time.sleep(0.5)
+                    st.rerun()
+                else:
+                    st.error("邀请码不能为空")
+
+        with st.expander("🛠️ 数据库概览"):
+            st.caption("Supabase 数据统计")
+            try:
+                pub_cnt = _count_chunks(PUBLIC_SCOPE)
+                st.text(f"📚 公共库切片数: {pub_cnt}")
+
+                # 统计所有 scope
+                resp = _sb().rpc("", {}).execute() if False else None  # placeholder
+                # 简单统计各用户私有库
+                for uname, _ in user_list:
+                    cnt = _count_chunks(uname)
+                    if cnt > 0:
+                        st.text(f"🔒 {uname} 私有库: {cnt} 切片")
+            except Exception as e:
+                st.warning(f"统计失败: {e}")
+
+            st.divider()
+            st.caption("📋 用户列表")
+            display_users = {}
+            for k, v in all_users.items():
+                if isinstance(v, dict) and "password_hash" in v:
+                    v_copy = dict(v)
+                    v_copy["password_hash"] = v_copy["password_hash"][:8] + "..."
+                    display_users[k] = v_copy
+                else:
+                    display_users[k] = v
+            st.json(display_users)
+
+    # --- 清空聊天记录 ---
+    with st.expander("🧹 清空聊天记录"):
+        st.caption("清空后不可恢复")
+        if st.button("确认清空", use_container_width=True, type="secondary", key="btn_clear_chat"):
+            st.session_state.messages = []
+            st.rerun()
+
+
+# =========================
+# 8. 核心搜索逻辑（合并公共库 + 私有库）
+# =========================
+def _cosine_scores(query_vec, matrix):
+    np = _get_np()
+    query_norm = np.linalg.norm(query_vec)
+    if query_norm < 1e-10:
+        return np.zeros(matrix.shape[0])
+    mat_norms = np.linalg.norm(matrix, axis=1)
+    mat_norms = np.maximum(mat_norms, 1e-10)
+    return (matrix @ query_vec) / (mat_norms * query_norm)
+
+
+def search_local(query, top_k, threshold):
+    query_vec = encode_query(query)
+    all_results = []
+
+    pub_docs = st.session_state.get("public_docs", [])
+    pub_np = _get_embeddings_np("public")
+    if pub_docs and pub_np.size > 0:
+        scores = _cosine_scores(query_vec, pub_np)
+        for i, s in enumerate(scores):
+            if s > threshold:
+                all_results.append((float(s), pub_docs[i]))
+
+    priv_docs = st.session_state.get("private_docs", [])
+    priv_np = _get_embeddings_np("private")
+    if priv_docs and priv_np.size > 0:
+        scores = _cosine_scores(query_vec, priv_np)
+        for i, s in enumerate(scores):
+            if s > threshold:
+                all_results.append((float(s), priv_docs[i]))
+
+    all_results.sort(key=lambda x: x[0], reverse=True)
+    return [doc for _, doc in all_results[:top_k]]
+
+
+# =========================
+# 9. LLM 回答逻辑
+# =========================
+def llm_answer(query, context_docs, selected_display_name, web_enabled):
+    all_context = ""
+    curr_time = datetime.now().strftime("%Y-%m-%d %H:%M")
+
+    if context_docs:
+        all_context += "【知识库资料】：\n" + "\n".join(context_docs) + "\n"
+
+    if web_enabled:
+        search_res = web_search(query)
+        all_context += f"\n【互联网实时资料】：\n{search_res}"
+
+    prompt_content = f"当前时间：{curr_time}\n\n参考资料：\n{all_context[:6500]}\n\n用户问题：{query}"
+    input_tokens = estimate_tokens(prompt_content)
+
+    or_client = get_or_client()
+    ds_client = get_ds_client()
+    baidu_client = get_baidu_client()
+
+    special_clients = {"deepseek-chat": ds_client, "ernie-3.5-8k": baidu_client}
+    selected_id = model_mapping[selected_display_name]
+
+    retry_queue = []
+    retry_queue.append(
+        (special_clients.get(selected_id, or_client), selected_id, f"首选-{selected_display_name}")
+    )
+
+    if selected_id != "stepfun/step-3.5-flash:free":
+        retry_queue.append((or_client, "stepfun/step-3.5-flash:free", "⚡ 快速备选-Step3.5"))
+
+    if selected_id != "openrouter/free":
+        retry_queue.append((or_client, "openrouter/free", "OR-Auto 免费避堵"))
+
+    paid_backups = [
+        ("deepseek-chat", "🛡️ DeepSeek 官方", ds_client),
+        ("ernie-3.5-8k", "🏢 百度文心", baidu_client),
+    ]
+    for p_id, p_label, p_client in paid_backups:
+        if selected_id != p_id:
+            retry_queue.append((p_client, p_id, f"💰 收费兜底-{p_label}"))
+
+    messages = [
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {"role": "user", "content": prompt_content},
+    ]
+
+    for idx, (client, m_id, label) in enumerate(retry_queue):
+        logger.info(f"[{CURRENT_USER}] 尝试链路: {label}")
+        try:
+            extra_h = (
+                {"HTTP-Referer": "https://streamlit.io", "X-Title": "RAG_v3"}
+                if client is or_client
+                else None
+            )
+            response = client.chat.completions.create(
+                model=m_id,
+                messages=messages,
+                stream=True,
+                extra_headers=extra_h,
+                timeout=25,
+            )
+
+            full_text = ""
+            has_content = False
+            for chunk in response:
+                if chunk.choices and chunk.choices[0].delta.content:
+                    content = chunk.choices[0].delta.content
+                    full_text += content
+                    has_content = True
+                    yield content
+
+            if has_content:
+                st.session_state["last_meta"] = (
+                    f"🟢 {label} | 📊 ~{input_tokens}/{estimate_tokens(full_text)} Tokens"
+                )
+                return
+
+        except Exception as e:
+            err_msg = str(e)
+            logger.warning(f"{label} 失败: {err_msg[:100]}")
+            if "429" in err_msg:
+                st.toast(f"{label} 拥堵，切换备选...", icon="⏳")
+                time.sleep(1.5)
+            continue
+
+    yield "❌ 抱歉，所有免费和收费线路均暂时不可用。"
+
+
+# =========================
+# 10. 聊天渲染
+# =========================
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+for m in st.session_state.messages:
+    with st.chat_message(m["role"]):
+        st.markdown(m["content"])
+        if "meta" in m:
+            st.caption(m["meta"])
+
+if q := st.chat_input("输入问题...", key="chat_input_v3"):
+    st.session_state.messages.append({"role": "user", "content": q})
+    with st.chat_message("user"):
+        st.markdown(q)
+
+    with st.chat_message("assistant"):
+        relevant_docs = search_local(q, ui_top_k, ui_threshold)
+        container = st.empty()
+        container.markdown("*🤔 正在组织语言...*")
+
+        if web_on:
+            with st.status("🌐 正在抓取实时网络数据...", expanded=False) as s:
+                time.sleep(0.1)
+                s.update(label="✅ 网络资料已就绪", state="complete")
+
+        try:
+            full_response = container.write_stream(
+                llm_answer(q, relevant_docs, selected_display_name, web_on)
+            )
+            meta_info = st.session_state.get("last_meta", "")
+            st.caption(meta_info)
+            st.session_state.messages.append(
+                {"role": "assistant", "content": full_response, "meta": meta_info}
+            )
+        except Exception as e:
+            logger.error(f"模型调用异常: {e}")
+            container.error(f"❌ 抱歉，连接模型时出错了: {str(e)}")
+
+_perf("script execution complete")

requirements.txt

@@ -1,3 +1,8 @@
-altair
-pandas
-streamlit
+streamlit
+numpy
+openai
+langchain-text-splitters
+python-docx
+tavily-python
+pdfplumber
+supabase