系列二 · 03 | 阶段三·查数据 + 权限红线：让 Agent 安全查询年假

def get_connection(db_path: str | None = None) -> sqlite3.Connection:
    """返回 SQLite 连接。生产换 MySQL/PostgreSQL 只改这里。"""
    db_path = db_path or os.getenv("LEAVE_DB_PATH") or str(DEFAULT_DB_PATH)
    Path(db_path).parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row
    return conn

def get_leave_balance(user_id: str) -> dict | None:
    """查询单个员工的假期余额。参数化查询，安全。"""
    with get_connection() as conn:
        cur = conn.execute(
            "SELECT user_id, name, annual_total, annual_used, sick_total, sick_used "
            "FROM leave_balances WHERE user_id = ?",   # ? 占位，杜绝 SQL 注入
            (user_id,),
        )
        row = cur.fetchone()
    if row is None:
        return None
    return {
        "user_id": row["user_id"], "name": row["name"],
        "annual_total": row["annual_total"], "annual_used": row["annual_used"],
        "annual_remaining": row["annual_total"] - row["annual_used"],
        "sick_total": row["sick_total"], "sick_used": row["sick_used"],
        "sick_remaining": row["sick_total"] - row["sick_used"],
    }

import contextvars
from fastapi import Header, HTTPException

# 关键点：该值由服务端注入，模型永远无法通过工具参数覆盖它——权限红线的根基。
user_id_ctx: contextvars.ContextVar[str] = contextvars.ContextVar("user_id", default="")

from app.auth import get_current_user, user_id_ctx

@app.post("/chat", response_model=ChatResponse)
async def chat(req: ChatRequest, user_id: str = Depends(get_current_user)):
    token = user_id_ctx.set(user_id)      # 服务端注入
    try:
        reply = await run_agent(req.message, user_id)
    finally:
        user_id_ctx.reset(token)          # 必须 reset
    return ChatResponse(reply=reply, user_id=user_id)

from langchain_core.tools import tool
from app.auth import user_id_ctx
from app.db.leave import get_leave_balance as fetch_leave_balance

@tool
def get_leave_balance() -> str:
    """查询**当前登录用户本人**的年假和病假余额。只能查本人，无法查他人。"""
    user_id = user_id_ctx.get()           # 来自服务端，不是模型传的
    if not user_id:
        return "未识别到登录用户，无法查询假期余额。"
    data = fetch_leave_balance(user_id)
    if data is None:
        return f"未找到用户 {user_id} 的假期记录，请联系 HR 核实。"
    return (f"【{data['name']}（{data['user_id']}）的假期余额】\n"
            f"年假：总共 {data['annual_total']} 天，已用 {data['annual_used']} 天，"
            f"剩余 {data['annual_remaining']} 天\n"
            f"病假：总共 {data['sick_total']} 天，已用 {data['sick_used']} 天，"
            f"剩余 {data['sick_remaining']} 天")

from app.agent.tools.leave_tool import get_leave_balance as get_leave_balance_tool

SYSTEM_PROMPT = """...（略）...
- 当用户询问"我的年假/病假还有多少"等**本人假期**问题时，调用 get_leave_balance 工具；
权限红线（务必遵守）：
- get_leave_balance 只能查**当前登录用户本人**，你无法也不得尝试查询他人假期；
  若用户想查别人的假期，请礼貌告知本助手仅支持查询本人。
- 你永远不要假设或编造某个具体员工（如"李四"）的假期数字。"""

tools: list[BaseTool] = [search_knowledge_base, get_leave_balance_tool]

[OK] DB 层：zhangsan 剩余 7 天；lisi 已用 0；nobody 返回 None
[OK] 权限红线：工具参数={}（无 user 字段，模型无法指定查谁）
[OK] 以 zhangsan 身份调用 -> 【张三（zhangsan）的假期余额】 ...（仅本人）
[OK] 尝试传入 user_id=lisi 被忽略，仍返回：【张三（zhangsan）的假期余额】
[OK] 未知用户 -> 未找到用户 ghost 的假期记录，请联系 HR 核实。
[OK] import app.main 成功（含 M2+M3 改动）

# 1) 灌样本数据
python scripts/seed_db.py

# 2) 起服务（另一个终端）
uvicorn app.main:app --reload

# 3) 用 Swagger 最省事：浏览器开 http://127.0.0.1:8000/docs
#    POST /chat → body: {"message":"我的年假还有几天？"} → X-User-Id 填 zhangsan → Execute