server.py

# app/server.py
# ================= 强制清除代理配置 (必须放在最前面) =================
import os
os.environ.pop("http_proxy", None)
os.environ.pop("https_proxy", None)
os.environ.pop("all_proxy", None)
# ===============================================================
import uvicorn
import shutil
from fastapi import FastAPI, HTTPException, Depends, UploadFile, File
from pydantic import BaseModel
from typing import List, Optional
from datetime import datetime
from sqlalchemy import create_engine, Column, Integer, String, Text, desc
from sqlalchemy.orm import declarative_base, sessionmaker, Session



# 引入你的 RAG 组件
from app.chains import fallback_chain, history_rewrite_chain
from app.graph import app_graph, get_retrievers
from app.retriever import add_single_file
from app.config import DATA_DIR

# [新增] 引入认证模块
from app.auth import (
    get_password_hash, 
    verify_password, 
    create_access_token, 
    get_current_user
)

app = FastAPI()

# ================= 1. MySQL 数据库配置（优化版） =================
# 【请务必修改】将 root:123456 替换为你本地 MySQL 的真实账号密码
SQLALCHEMY_DATABASE_URL = "mysql+pymysql://root:123456@localhost:3306/rag_db"

# [优化] 数据库连接池配置
engine = create_engine(
    SQLALCHEMY_DATABASE_URL,
    pool_size=10,              # 连接池大小
    max_overflow=20,           # 超出pool_size时最多创建的连接数
    pool_recycle=3600,         # 连接回收时间（1小时）
    pool_pre_ping=True,        # 使用前检查连接有效性
    echo=False                 # 生产环境关闭SQL日志
)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()

from contextlib import asynccontextmanager
from app.graph import get_retrievers  # 引入你的资源加载函数


# 1. 定义生命周期管理器
@asynccontextmanager
async def lifespan(app: FastAPI):
    # --- 启动时执行 ---
    print("🔥 [Startup] 正在预热系统，加载向量库和模型...")
    try:
        # 主动调用一次，触发 @lru_cache 加载数据
        retriever = get_retrievers()
        if retriever is None:
            print("⚠️  [Startup] 向量库初始化失败，但服务器将继续运行")
            print("💡 [Startup] 建议: 运行 python fix_chromadb.py 修复数据库")
        else:
            print("✅ [Startup] 系统预热完成，随时可以响应请求！")
    except Exception as e:
        error_msg = str(e)
        print(f"⚠️  [Startup] 预热失败: {error_msg[:200]}")
        if "Error in compaction" in error_msg or "hnsw" in error_msg.lower():
            print("💡 [Startup] 检测到 ChromaDB 数据库损坏")
            print("💡 [Startup] 建议: 运行 python fix_chromadb.py 修复数据库")
        print("⚠️  [Startup] 服务器将继续运行，但向量检索功能可能不可用")

    yield  # 服务运行中...

    # --- 关闭时执行 (可选) ---
    print("👋 [Shutdown] 服务器正在关闭...")


# 2. 注入 lifespan
app = FastAPI(lifespan=lifespan)


# --- 定义数据库模型（优化版） ---

class ChatHistory(Base):
    __tablename__ = "chat_history"
    id = Column(Integer, primary_key=True, index=True)
    username = Column(String(50), index=True, default="default")  # [新增] 用户隔离
    role = Column(String(10))  # 'user' 或 'ai'
    content = Column(Text)  # 对话内容
    time = Column(String(50))  # 时间戳


class User(Base):
    __tablename__ = "users"
    id = Column(Integer, primary_key=True, index=True)
    username = Column(String(50), unique=True, index=True)
    password = Column(String(255))  # [修改] 扩容以存储 bcrypt 加密后的密码


# 自动创建表结构 (如果不存在)
Base.metadata.create_all(bind=engine)


# 数据库依赖项
def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


# ================= 2. API 请求/响应结构 =================

class QueryRequest(BaseModel):
    question: str
    username: Optional[str] = "default"


class LoginRequest(BaseModel):
    username: str
    password: str


class LogoutRequest(BaseModel):
    username: str


class RegisterRequest(BaseModel):
    username: str
    password: str


class HistoryItem(BaseModel):
    role: str
    content: str
    time: str


# ================= 3. 接口逻辑 =================

@app.post("/login")
def login_endpoint(request: LoginRequest, db: Session = Depends(get_db)):
    """
    用户登录接口
    - 使用明文密码验证
    - 返回 JWT Token
    """
    try:
        print(f"🔐 [Login] 收到登录请求: {request.username}")

        # 1. 查询用户
        user = db.query(User).filter(User.username == request.username).first()

        # 2. 验证账号密码（明文比较）
        if not user:
            return {"status": "error", "message": "用户不存在"}

        # 直接比较明文密码
        if request.password != user.password:
            return {"status": "error", "message": "密码错误"}

        # 3. [优化] 生成 JWT Token
        access_token = create_access_token(data={"sub": user.username})
        
        print(f"✅ [Login] {request.username} 验证通过")
        return {
            "status": "success",
            "message": "登录成功",
            "token": access_token,
            "username": user.username
        }
    except Exception as e:
        import traceback
        error_msg = f"登录失败: {str(e)}"
        print(f"❌ [Login] 错误: {error_msg}")
        traceback.print_exc()
        return {"status": "error", "message": error_msg}


@app.post("/register")
def register_endpoint(request: RegisterRequest, db: Session = Depends(get_db)):
    """
    用户注册接口
    - 密码明文存储（不加密）
    """
    try:
        print(f"📝 [Register] 收到注册请求: {request.username}")

        # 1. 检查用户是否已存在
        existing_user = db.query(User).filter(User.username == request.username).first()
        if existing_user:
            return {"status": "error", "message": "用户名已存在"}

        # 2. 直接存储明文密码
        new_user = User(username=request.username, password=request.password)
        
        db.add(new_user)
        db.commit()
        db.refresh(new_user)
        print(f"✅ [Register] 用户 {request.username} 注册成功")
        return {"status": "success", "message": "注册成功"}
    except Exception as e:
        db.rollback()
        import traceback
        error_msg = f"注册失败: {str(e)}"
        print(f"❌ [Register] 错误: {error_msg}")
        traceback.print_exc()
        return {"status": "error", "message": error_msg}


@app.post("/logout")
def logout_endpoint(
    request: LogoutRequest,
    current_user: str = Depends(get_current_user)  # [新增] 需要认证
):
    """
    用户登出接口（优化版）
    - 需要 Token 认证
    """
    print(f"👋 [Logout] 用户 {current_user} 已退出登录")
    return {"status": "success", "message": "已退出"}


# --- [优化] 清空历史记录接口（按用户隔离） ---
@app.post("/clear_history")
def clear_history_endpoint(
    current_user: str = Depends(get_current_user),  # [新增] 需要认证
    db: Session = Depends(get_db)
):
    """
    清空当前用户的聊天记录（优化版）
    - 只删除当前用户的历史记录
    - 需要 Token 认证
    """
    try:
        num_deleted = db.query(ChatHistory).filter(
            ChatHistory.username == current_user
        ).delete()
        db.commit()
        print(f"🧹 [Clear] 用户 {current_user} 清空历史记录，共删除 {num_deleted} 条")
        return {"status": "success", "message": f"已删除 {num_deleted} 条记录"}
    except Exception as e:
        db.rollback()
        return {"status": "error", "message": str(e)}


# --- [优化] 文件上传接口（需要认证） ---
@app.post("/upload")
async def upload_file_endpoint(
    file: UploadFile = File(...),
    current_user: str = Depends(get_current_user)  # [新增] 需要认证
):
    """
    上传文件 -> 保存 -> 索引 -> 返回详细日志（优化版）
    - 需要 Token 认证
    """
    upload_logs = []

    def log(msg):
        print(msg)
        upload_logs.append(msg)

    try:
        if not os.path.exists(DATA_DIR):
            os.makedirs(DATA_DIR)

        file_path = os.path.join(DATA_DIR, file.filename)

        # 1. 保存文件
        log(f"📂 [Upload] 正在接收文件: {file.filename}...")
        with open(file_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
        log(f"✅ [Upload] 文件已保存至: {file_path}")

        # 2. 触发向量库更新
        log(f"⚙️ [Index] 开始构建向量索引...")
        try:
            # 获取 add_single_file 返回的内部日志
            retriever_logs = add_single_file(file_path)
            upload_logs.extend(retriever_logs)
        except ValueError as e:
            # 文件格式不支持或其他明确的错误
            err_msg = f"❌ [Index] 索引构建失败: {e}"
            log(err_msg)
            return {"status": "error", "message": str(e), "logs": upload_logs}
        except Exception as e:
            # ChromaDB 内部错误或其他未知错误
            err_msg = f"❌ [Index] 索引构建失败: {e}"
            log(err_msg)
            import traceback
            traceback.print_exc()
            
            # 如果是 ChromaDB 错误，提供更友好的提示
            error_str = str(e)
            if "Error in compaction" in error_str or "hnsw" in error_str.lower():
                user_message = "向量数据库内部错误，可能是数据库文件损坏。建议重启服务器或联系管理员。"
            else:
                user_message = str(e)
            
            return {"status": "error", "message": user_message, "logs": upload_logs}

        # 3. 清除缓存
        log(f"🔄 [Cache] 清除检索器缓存，应用最新数据...")
        get_retrievers.cache_clear()

        return {
            "status": "success",
            "message": f"文件 {file.filename} 上传成功",
            "logs": upload_logs  # [核心修改] 返回日志列表
        }

    except Exception as e:
        return {"status": "error", "message": str(e), "logs": upload_logs}


@app.post("/chat")
def chat_endpoint(
    request: QueryRequest,
    current_user: str = Depends(get_current_user),  # [新增] 需要认证
    db: Session = Depends(get_db)
):
    """
    智能问答接口（优化版）
    - 需要 Token 认证
    - 按用户隔离历史记录
    """
    raw_question = request.question
    current_time = datetime.now().strftime("%H:%M")

    # 初始化日志收集器
    server_logs = []

    def log(msg):
        print(msg)
        server_logs.append(msg)

    log(f"📨 [Server] 用户 {current_user} 请求: {raw_question}")

    try:
        final_question = raw_question

        # --- Step 1: [优化] 获取当前用户的最近历史记录 ---
        recent_records = db.query(ChatHistory).filter(
            ChatHistory.username == current_user
        ).order_by(desc(ChatHistory.id)).limit(6).all()
        recent_records = recent_records[::-1]

        # 构建历史对话字符串
        history_str = ""
        if recent_records:
            history_parts = []
            for record in recent_records:
                role_label = "用户" if record.role == "user" else "助手"
                history_parts.append(f"{role_label}: {record.content}")
            history_str = "\n".join(history_parts)

        # --- Step 2: 执行 RAG 流程 (LangGraph) ---
        log(f"🚀 [执行] 使用问题: '{final_question}' 进入检索...")

        inputs = {
            "question": final_question,
            "chat_history": history_str,
            "logs": []
        }

        graph_result = app_graph.invoke(inputs)
        answer = graph_result["generation"]
        server_logs.extend(graph_result.get("logs", []))

        # --- Step 3: 智能兜底检测 ---
        refusal_keywords = ["无法从现有资料中找到答案", "无法从现有上下文中", "No answer found"]
        if any(k in answer for k in refusal_keywords):
            log(f"   ⚠️ [结果] RAG 知识库无相关文档")
            log(f"   🔄 [兜底] 切换策略 -> 调用通用大模型回答...")
            answer = fallback_chain.invoke({
                "question": final_question,
                "chat_history": history_str
            })
        else:
            log("   ✅ [结果] 知识库检索成功")

        # --- Step 4: [优化] 存入数据库（带用户名） ---
        user_record = ChatHistory(
            username=current_user,
            role="user",
            content=raw_question,
            time=current_time
        )
        db.add(user_record)
        ai_record = ChatHistory(
            username=current_user,
            role="ai",
            content=answer,
            time=current_time
        )
        db.add(ai_record)
        db.commit()

        return {
            "answer": answer,
            "logs": server_logs
        }

    except Exception as e:
        import traceback
        traceback.print_exc()
        return {
            "answer": f"服务器内部错误: {str(e)}",
            "logs": server_logs + [f"❌ Error: {str(e)}"]
        }


@app.get("/get_history", response_model=List[HistoryItem])
def get_history(
    current_user: str = Depends(get_current_user),  # [新增] 需要认证
    db: Session = Depends(get_db)
):
    """
    获取历史记录接口（优化版）
    - 只返回当前用户的历史记录
    - 需要 Token 认证
    """
    records = db.query(ChatHistory).filter(
        ChatHistory.username == current_user
    ).order_by(ChatHistory.id).all()
    return records


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)