vent_agent_2.0/db/chat_store.py

"""会话与消息存取。

MySQL 是会话历史的**事实源**：图状态（checkpoint）丢了可以从这里补种，
反过来不行。所以每次对话结束都必须有一条 assistant 消息落库。

约定：``messages`` 里的 ``created_at`` 只到秒级且同秒可能重复，因此所有
「取最近 N 条」的查询都带上 ``rowid``/``id`` 的次级排序，保证分页稳定。
"""

from __future__ import annotations

import json
from typing import Any, Sequence

from db import db_conn
from obs import get_logger

logger = get_logger(__name__)


def _dumps(value: Any) -> str | None:
    """把结构化字段序列化成 JSON 字符串；None 保持 None。"""
    if value is None:
        return None
    if isinstance(value, str):
        return value
    return json.dumps(value, ensure_ascii=False)


def _loads(value: Any, default: Any) -> Any:
    """反序列化 JSON 字段；空值或坏数据回落默认值（不抛异常）。"""
    if value is None or value == "":
        return default
    if isinstance(value, (list, dict)):
        return value
    try:
        return json.loads(value)
    except (TypeError, ValueError):
        return default


# ============================================================
# 会话
# ============================================================

def create_session(user_name: str, mode: str = "full", title: str = "") -> dict[str, Any]:
    """新建会话并返回其完整行。会话 ID 即 LangGraph thread_id。"""
    session_id = db_conn.new_id()
    stamp = db_conn.fmt_dt()
    with db_conn.transaction():
        db_conn.execute(
            "INSERT INTO sessions (id, title, user_name, mode, pinned, created_at, updated_at)"
            " VALUES (?, ?, ?, ?, 0, ?, ?)",
            (session_id, title, user_name, mode, stamp, stamp),
        )
    logger.info("新建会话 session_id=%s user=%s mode=%s", session_id, user_name, mode)
    return get_session(session_id) or {}


def get_session(session_id: str) -> dict[str, Any] | None:
    """按 ID 取会话；不存在返回 None。"""
    row = db_conn.query_one("SELECT * FROM sessions WHERE id = ?", (session_id,))
    return _hydrate_session(row)


def get_session_for_user(session_id: str, user_name: str) -> dict[str, Any] | None:
    """取会话并校验归属。用于 400/404 判定前的存在性与越权检查。"""
    session = get_session(session_id)
    if session and session.get("user_name") != user_name:
        return None
    return session


def list_sessions(user_name: str, limit: int = 50, offset: int = 0) -> list[dict[str, Any]]:
    """列出某用户的会话，置顶优先、更新时间倒序。"""
    rows = db_conn.query(
        "SELECT * FROM sessions WHERE user_name = ?"
        " ORDER BY pinned DESC, updated_at DESC, id DESC LIMIT ? OFFSET ?",
        (user_name, limit, offset),
    )
    return [_hydrate_session(row) for row in rows]


def search_sessions(user_name: str, keyword: str, limit: int = 50) -> list[dict[str, Any]]:
    """按标题模糊搜索会话。"""
    rows = db_conn.query(
        "SELECT * FROM sessions WHERE user_name = ? AND title LIKE ?"
        " ORDER BY pinned DESC, updated_at DESC, id DESC LIMIT ?",
        (user_name, f"%{keyword}%", limit),
    )
    return [_hydrate_session(row) for row in rows]


def update_session_title(session_id: str, title: str) -> None:
    """更新会话标题。"""
    db_conn.execute(
        "UPDATE sessions SET title = ?, updated_at = ? WHERE id = ?",
        (title, db_conn.fmt_dt(), session_id),
    )


def update_session_mode(session_id: str, mode: str) -> None:
    """更新会话的权限模式（plan | full）。"""
    db_conn.execute(
        "UPDATE sessions SET mode = ?, updated_at = ? WHERE id = ?",
        (mode, db_conn.fmt_dt(), session_id),
    )


def update_session_pin(session_id: str, pinned: bool) -> None:
    """设置/取消会话置顶。"""
    db_conn.execute(
        "UPDATE sessions SET pinned = ?, updated_at = ? WHERE id = ?",
        (1 if pinned else 0, db_conn.fmt_dt(), session_id),
    )


def touch_session(session_id: str) -> None:
    """刷新会话的 updated_at，让最近活跃的会话排在列表前面。"""
    db_conn.execute(
        "UPDATE sessions SET updated_at = ? WHERE id = ?", (db_conn.fmt_dt(), session_id)
    )


def delete_session(session_id: str) -> None:
    """删除会话及其全部消息、用量、压缩存档。

    注意：调用方还须清除 LangGraph checkpoint 线程，否则图状态会残留并
    在下次复用同 ID 时「诈尸」。
    """
    with db_conn.transaction():
        db_conn.execute("DELETE FROM messages WHERE session_id = ?", (session_id,))
        db_conn.execute("DELETE FROM session_context_usage WHERE session_id = ?", (session_id,))
        db_conn.execute("DELETE FROM context_archives WHERE session_id = ?", (session_id,))
        db_conn.execute("DELETE FROM sessions WHERE id = ?", (session_id,))
    logger.info("删除会话 session_id=%s", session_id)


# ============================================================
# 消息
# ============================================================

def save_message(
    session_id: str,
    role: str,
    content: str,
    tool_calls: Any = None,
    citations: Any = None,
) -> str:
    """落一条消息，返回消息 ID。user 与 assistant 都走这里。"""
    message_id = db_conn.new_id()
    stamp = db_conn.fmt_dt()
    with db_conn.transaction():
        db_conn.execute(
            "INSERT INTO messages (id, session_id, msg_role, content, tool_calls, citations, created_at)"
            " VALUES (?, ?, ?, ?, ?, ?, ?)",
            (
                message_id,
                session_id,
                role,
                content or "",
                _dumps(tool_calls),
                _dumps(citations),
                stamp,
            ),
        )
        db_conn.execute(
            "UPDATE sessions SET updated_at = ? WHERE id = ?", (stamp, session_id)
        )
    return message_id


def get_messages(
    session_id: str, limit: int | None = None, offset: int = 0
) -> list[dict[str, Any]]:
    """取会话消息，按时间正序返回。

    ``limit`` 给定时取**最近** limit 条（先倒序截断再正序还原），这正是
    冷启动补种图状态需要的语义。
    """
    if limit is None:
        rows = db_conn.query(
            "SELECT * FROM messages WHERE session_id = ? ORDER BY created_at ASC, id ASC"
            " LIMIT ? OFFSET ?",
            (session_id, 100000, offset),
        )
    elif offset == 0:
        rows = db_conn.query(
            "SELECT * FROM (SELECT * FROM messages WHERE session_id = ?"
            " ORDER BY created_at DESC, id DESC LIMIT ?) AS recent"
            " ORDER BY created_at ASC, id ASC",
            (session_id, limit),
        )
    else:
        # 带偏移的分页：先按正序定位窗口
        rows = db_conn.query(
            "SELECT * FROM messages WHERE session_id = ?"
            " ORDER BY created_at ASC, id ASC LIMIT ? OFFSET ?",
            (session_id, limit, offset),
        )
    return [_hydrate_message(row) for row in rows]


def count_messages(session_id: str) -> int:
    """会话消息条数。"""
    return int(db_conn.scalar("SELECT COUNT(*) FROM messages WHERE session_id = ?", (session_id,), 0) or 0)


def delete_message(message_id: str) -> bool:
    """删除单条消息，返回是否命中。"""
    return db_conn.execute("DELETE FROM messages WHERE id = ?", (message_id,)) > 0


# ============================================================
# 上下文用量与压缩存档
# ============================================================

def save_context_usage(
    session_id: str,
    turn: int,
    input_tokens: int,
    output_tokens: int,
    cache_read_tokens: int = 0,
    usage_source: str = "provider",
    compacted: bool = False,
    compact_count: int = 0,
    snapshot: Any = None,
) -> str:
    """记录一轮对话的上下文用量。"""
    row_id = db_conn.new_id()
    total = int(input_tokens) + int(output_tokens)
    db_conn.execute(
        "INSERT INTO session_context_usage (id, session_id, turn, input_tokens, output_tokens,"
        " cache_read_tokens, total_tokens, usage_source, compacted, compact_count, snapshot, created_at)"
        " VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
        (
            row_id,
            session_id,
            int(turn),
            int(input_tokens),
            int(output_tokens),
            int(cache_read_tokens),
            total,
            usage_source,
            1 if compacted else 0,
            int(compact_count),
            _dumps(snapshot),
            db_conn.fmt_dt(),
        ),
    )
    return row_id


def get_context_usage(session_id: str, limit: int = 50) -> dict[str, Any]:
    """取会话的用量明细（倒序）与会话累计压缩次数。"""
    rows = db_conn.query(
        "SELECT * FROM session_context_usage WHERE session_id = ?"
        " ORDER BY created_at DESC, id DESC LIMIT ?",
        (session_id, limit),
    )
    history = []
    for row in rows:
        item = dict(row)
        item["snapshot"] = _loads(item.get("snapshot"), {})
        item["compacted"] = bool(item.get("compacted"))
        history.append(item)
    compact_count = int(
        db_conn.scalar(
            "SELECT MAX(compact_count) FROM session_context_usage WHERE session_id = ?",
            (session_id,),
            0,
        )
        or 0
    )
    latest = history[0] if history else None
    return {
        "session_id": session_id,
        "compact_count": compact_count,
        "latest": latest,
        "history": history,
    }


def save_context_archive(
    session_id: str, turn: int, file_path: str, tokens: int = 0
) -> str:
    """登记一次「压缩前原文」的落盘位置，供事后追溯。"""
    row_id = db_conn.new_id()
    db_conn.execute(
        "INSERT INTO context_archives (id, session_id, turn, file_path, tokens, created_at)"
        " VALUES (?, ?, ?, ?, ?, ?)",
        (row_id, session_id, int(turn), file_path, int(tokens), db_conn.fmt_dt()),
    )
    return row_id


def list_context_archives(session_id: str) -> list[dict[str, Any]]:
    """列出会话的压缩存档索引。"""
    return db_conn.query(
        "SELECT * FROM context_archives WHERE session_id = ? ORDER BY created_at ASC",
        (session_id,),
    )


# ============================================================
# 用户偏好
# ============================================================

def save_preference(user_name: str, key: str, value: str) -> None:
    """写入/覆盖一条用户偏好（先查后写，避免方言相关的 upsert 语法）。"""
    existing = db_conn.query_one(
        "SELECT id FROM user_preferences WHERE user_name = ? AND pref_key = ?",
        (user_name, key),
    )
    stamp = db_conn.fmt_dt()
    if existing:
        db_conn.execute(
            "UPDATE user_preferences SET pref_value = ?, updated_at = ? WHERE id = ?",
            (value, stamp, existing["id"]),
        )
    else:
        db_conn.execute(
            "INSERT INTO user_preferences (id, user_name, pref_key, pref_value, updated_at)"
            " VALUES (?, ?, ?, ?, ?)",
            (db_conn.new_id(), user_name, key, value, stamp),
        )


def list_preferences(user_name: str) -> list[dict[str, Any]]:
    """列出某用户的全部偏好。"""
    return db_conn.query(
        "SELECT pref_key, pref_value, updated_at FROM user_preferences"
        " WHERE user_name = ? ORDER BY pref_key ASC",
        (user_name,),
    )


def delete_preference(user_name: str, key: str) -> bool:
    """删除一条偏好，返回是否命中。"""
    return db_conn.execute(
        "DELETE FROM user_preferences WHERE user_name = ? AND pref_key = ?",
        (user_name, key),
    ) > 0


def preferences_as_text(user_name: str, limit: int = 20) -> str:
    """把偏好渲染成可注入提示词的一行文本；无偏好返回空串。"""
    rows = list_preferences(user_name)[:limit]
    if not rows:
        return ""
    pairs = [f"{row['pref_key']}={row['pref_value']}" for row in rows]
    return "；".join(pairs)


# ============================================================
# 行归一化
# ============================================================

def _hydrate_session(row: dict[str, Any] | None) -> dict[str, Any] | None:
    if not row:
        return None
    session = dict(row)
    session["pinned"] = bool(session.get("pinned"))
    return session


def _hydrate_message(row: dict[str, Any]) -> dict[str, Any]:
    message = dict(row)
    message["tool_calls"] = _loads(message.get("tool_calls"), [])
    message["citations"] = _loads(message.get("citations"), [])
    return message