app/models/db.py

"""SQLite 元数据 + LangGraph checkpoint 共享 db.

用 sqlite3 stdlib 即可, 暂不引入 SQLAlchemy (减少依赖).
WAL 模式开启, 允许 LangGraph checkpoint 与文档读并发.
"""
from __future__ import annotations

import json
import logging
import sqlite3
import threading
import time
from contextlib import contextmanager
from pathlib import Path
from typing import Any, Iterator

from app.config import settings
from app.core.paths import sqlite_dir

logger = logging.getLogger(__name__)

_local = threading.local()


def _connect() -> sqlite3.Connection:
    """每线程一个连接 (sqlite3 本身不线程安全)."""
    conn = sqlite3.connect(
        str(settings.sqlite_db_path),
        check_same_thread=False,
        isolation_level=None,  # autocommit, 手动 BEGIN
        timeout=30.0,
    )
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA synchronous=NORMAL")
    conn.execute("PRAGMA foreign_keys=ON")
    return conn


def get_conn() -> sqlite3.Connection:
    """获取当前线程的连接."""
    conn = getattr(_local, "conn", None)
    if conn is None:
        conn = _connect()
        _local.conn = conn
    return conn


@contextmanager
def transaction() -> Iterator[sqlite3.Connection]:
    """显式事务 wrapper."""
    conn = get_conn()
    conn.execute("BEGIN")
    try:
        yield conn
    except Exception:
        conn.execute("ROLLBACK")
        raise
    else:
        conn.execute("COMMIT")


SCHEMA = """
CREATE TABLE IF NOT EXISTS documents (
    id              TEXT PRIMARY KEY,
    filename        TEXT NOT NULL,
    mime            TEXT,
    size            INTEGER NOT NULL,
    page_count      INTEGER,
    chunk_count     INTEGER DEFAULT 0,
    status          TEXT NOT NULL DEFAULT 'uploading',
    progress        INTEGER DEFAULT 0,         -- 0-100, 用于前端进度条
    progress_label  TEXT,                       -- "正在分块 (12/48)..." 等
    error           TEXT,
    sha256          TEXT NOT NULL,
    parser          TEXT,
    chunk_size      INTEGER,
    chunk_overlap   INTEGER,
    semantic_chunking      INTEGER DEFAULT 0,
    contextual_retrieval   INTEGER DEFAULT 0,
    meta_json       TEXT,
    created_at      REAL NOT NULL,
    updated_at      REAL NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_documents_sha ON documents(sha256);
CREATE INDEX IF NOT EXISTS idx_documents_status ON documents(status);

CREATE TABLE IF NOT EXISTS chunks (
    id              TEXT PRIMARY KEY,
    doc_id          TEXT NOT NULL,
    parent_id       TEXT,
    chunk_index     INTEGER NOT NULL,
    text            TEXT NOT NULL,
    token_count     INTEGER,
    page_no         INTEGER,
    heading         TEXT,
    context_prefix  TEXT,
    created_at      REAL NOT NULL,
    FOREIGN KEY (doc_id) REFERENCES documents(id) ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS idx_chunks_doc ON chunks(doc_id);
CREATE INDEX IF NOT EXISTS idx_chunks_parent ON chunks(parent_id);

CREATE TABLE IF NOT EXISTS sessions (
    id              TEXT PRIMARY KEY,
    user_id         TEXT NOT NULL DEFAULT 'default',
    title           TEXT,
    message_count   INTEGER DEFAULT 0,
    created_at      REAL NOT NULL,
    updated_at      REAL NOT NULL
);

CREATE TABLE IF NOT EXISTS messages (
    id              TEXT PRIMARY KEY,
    session_id      TEXT NOT NULL,
    role            TEXT NOT NULL,
    content         TEXT NOT NULL,
    citations_json  TEXT,
    tool_calls_json TEXT,
    created_at      REAL NOT NULL,
    FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS idx_messages_session ON messages(session_id, created_at);
"""


def init_db() -> None:
    """启动时调用. 幂等."""
    sqlite_dir()
    conn = get_conn()
    conn.executescript(SCHEMA)
    # 在线迁移: 给旧 documents 表加新列 (progress / parser / chunk_size 等).
    # 失败无害 (列已存在).
    _MIGRATIONS = [
        "ALTER TABLE documents ADD COLUMN progress INTEGER DEFAULT 0",
        "ALTER TABLE documents ADD COLUMN progress_label TEXT",
        "ALTER TABLE documents ADD COLUMN parser TEXT",
        "ALTER TABLE documents ADD COLUMN chunk_size INTEGER",
        "ALTER TABLE documents ADD COLUMN chunk_overlap INTEGER",
        "ALTER TABLE documents ADD COLUMN semantic_chunking INTEGER DEFAULT 0",
        "ALTER TABLE documents ADD COLUMN contextual_retrieval INTEGER DEFAULT 0",
    ]
    for sql in _MIGRATIONS:
        try:
            conn.execute(sql)
        except Exception:  # noqa: BLE001
            pass  # 列已存在, 跳过
    logger.info("DB schema ensured at %s", settings.sqlite_db_path)


# ========== Document CRUD ==========
def doc_insert(doc: dict[str, Any]) -> None:
    """doc 必须含: id, filename, size, sha256, created_at, updated_at."""
    conn = get_conn()
    conn.execute(
        """
        INSERT INTO documents (id, filename, mime, size, page_count, chunk_count,
                               status, progress, progress_label, error, sha256,
                               parser, chunk_size, chunk_overlap,
                               semantic_chunking, contextual_retrieval,
                               meta_json, created_at, updated_at)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        """,
        (
            doc["id"], doc["filename"], doc.get("mime"), doc["size"],
            doc.get("page_count"), doc.get("chunk_count", 0),
            doc.get("status", "uploading"),
            doc.get("progress", 0), doc.get("progress_label"),
            doc.get("error"),
            doc["sha256"], doc.get("parser"),
            doc.get("chunk_size"), doc.get("chunk_overlap"),
            int(bool(doc.get("semantic_chunking", False))),
            int(bool(doc.get("contextual_retrieval", False))),
            json.dumps(doc.get("meta", {}), ensure_ascii=False),
            doc["created_at"], doc.get("updated_at", doc["created_at"]),
        ),
    )


def doc_update_status(
    doc_id: str,
    status: str,
    *,
    error: str | None = None,
    chunk_count: int | None = None,
    page_count: int | None = None,
    progress: int | None = None,
    progress_label: str | None = None,
) -> None:
    conn = get_conn()
    sets = ["status = ?", "updated_at = ?"]
    params: list[Any] = [status, time.time()]
    if error is not None:
        sets.append("error = ?")
        params.append(error)
    if chunk_count is not None:
        sets.append("chunk_count = ?")
        params.append(chunk_count)
    if page_count is not None:
        sets.append("page_count = ?")
        params.append(page_count)
    if progress is not None:
        sets.append("progress = ?")
        params.append(max(0, min(100, int(progress))))
    if progress_label is not None:
        sets.append("progress_label = ?")
        params.append(progress_label)
    params.append(doc_id)
    conn.execute(f"UPDATE documents SET {', '.join(sets)} WHERE id = ?", params)


def doc_get(doc_id: str) -> dict[str, Any] | None:
    row = get_conn().execute("SELECT * FROM documents WHERE id = ?", (doc_id,)).fetchone()
    return _row_to_doc(row) if row else None


def doc_find_by_sha256(sha: str) -> dict[str, Any] | None:
    row = get_conn().execute(
        "SELECT * FROM documents WHERE sha256 = ? AND status = 'ready' ORDER BY created_at DESC LIMIT 1",
        (sha,),
    ).fetchone()
    return _row_to_doc(row) if row else None


def doc_list(limit: int = 100, offset: int = 0) -> list[dict[str, Any]]:
    rows = get_conn().execute(
        "SELECT * FROM documents ORDER BY created_at DESC LIMIT ? OFFSET ?",
        (limit, offset),
    ).fetchall()
    return [_row_to_doc(r) for r in rows]


def doc_delete(doc_id: str) -> bool:
    cur = get_conn().execute("DELETE FROM documents WHERE id = ?", (doc_id,))
    return cur.rowcount > 0


def _row_to_doc(row: sqlite3.Row) -> dict[str, Any]:
    d = dict(row)
    if d.get("meta_json"):
        try:
            d["meta"] = json.loads(d["meta_json"])
        except json.JSONDecodeError:
            d["meta"] = {}
    d.pop("meta_json", None)
    # 拆出 boolean 字段, 方便前端直接使用
    d["semantic_chunking"] = bool(d.get("semantic_chunking"))
    d["contextual_retrieval"] = bool(d.get("contextual_retrieval"))
    return d


# ========== Chunk CRUD (用于上下文回查, 实际检索走 Chroma) ==========
def chunk_insert_bulk(chunks: list[dict[str, Any]]) -> None:
    if not chunks:
        return
    conn = get_conn()
    with transaction():
        conn.executemany(
            """
            INSERT OR REPLACE INTO chunks
            (id, doc_id, parent_id, chunk_index, text, token_count, page_no, heading, context_prefix, created_at)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            [
                (
                    c["id"], c["doc_id"], c.get("parent_id"), c["chunk_index"],
                    c["text"], c.get("token_count"), c.get("page_no"),
                    c.get("heading"), c.get("context_prefix"), c.get("created_at", time.time()),
                )
                for c in chunks
            ],
        )


def chunk_get_by_doc(doc_id: str, limit: int = 200) -> list[dict[str, Any]]:
    rows = get_conn().execute(
        "SELECT * FROM chunks WHERE doc_id = ? ORDER BY chunk_index LIMIT ?",
        (doc_id, limit),
    ).fetchall()
    return [dict(r) for r in rows]


def chunk_delete_by_doc(doc_id: str) -> int:
    cur = get_conn().execute("DELETE FROM chunks WHERE doc_id = ?", (doc_id,))
    return cur.rowcount


# ========== Session / Message CRUD (阶段 3 实际使用) ==========
def session_upsert(session_id: str, user_id: str = "default", title: str | None = None) -> None:
    """插入或更新 session 行.

    关键: ON CONFLICT 时, title 用 COALESCE(sessions.title, excluded.title):
    - 原 title 已设: 保留原 title (不覆盖)
    - 原 title 为空: 用新 title (auto_title 首次写)
    这样调用方传 title=None 时不会清掉已有标题, 传有效 title 时只在原空时写入.
    """
    now = time.time()
    conn = get_conn()
    conn.execute(
        """
        INSERT INTO sessions (id, user_id, title, message_count, created_at, updated_at)
        VALUES (?, ?, ?, 0, ?, ?)
        ON CONFLICT(id) DO UPDATE SET
            title = COALESCE(sessions.title, excluded.title),
            updated_at = excluded.updated_at
        """,
        (session_id, user_id, title, now, now),
    )


def session_list(limit: int = 50) -> list[dict[str, Any]]:
    rows = get_conn().execute(
        "SELECT * FROM sessions ORDER BY updated_at DESC LIMIT ?", (limit,)
    ).fetchall()
    return [dict(r) for r in rows]


def session_get(session_id: str) -> dict[str, Any] | None:
    row = get_conn().execute("SELECT * FROM sessions WHERE id = ?", (session_id,)).fetchone()
    return dict(row) if row else None


def session_delete(session_id: str) -> bool:
    cur = get_conn().execute("DELETE FROM sessions WHERE id = ?", (session_id,))
    return cur.rowcount > 0


def message_insert(msg: dict[str, Any]) -> None:
    now = msg.get("created_at", time.time())
    conn = get_conn()
    with transaction():
        conn.execute(
            """
            INSERT INTO messages (id, session_id, role, content, citations_json, tool_calls_json, created_at)
            VALUES (?, ?, ?, ?, ?, ?, ?)
            """,
            (
                msg["id"], msg["session_id"], msg["role"], msg["content"],
                json.dumps(msg.get("citations", []), ensure_ascii=False),
                json.dumps(msg.get("tool_calls", []), ensure_ascii=False),
                now,
            ),
        )
        # 更新 session 计数
        conn.execute(
            "UPDATE sessions SET message_count = message_count + 1, updated_at = ? WHERE id = ?",
            (now, msg["session_id"]),
        )


def message_list_by_session(session_id: str, limit: int = 200) -> list[dict[str, Any]]:
    rows = get_conn().execute(
        "SELECT * FROM messages WHERE session_id = ? ORDER BY created_at LIMIT ?",
        (session_id, limit),
    ).fetchall()
    out = []
    for r in rows:
        d = dict(r)
        for k in ("citations_json", "tool_calls_json"):
            v = d.pop(k, None)
            if v:
                try:
                    d[k.removesuffix("_json")] = json.loads(v)
                except json.JSONDecodeError:
                    d[k.removesuffix("_json")] = []
        out.append(d)
    return out