NotesAgentic/backend/app/repository.py

"""SQLite Repository：笔记元数据、Block 与 FTS5 的读写。

向量（vec_blocks）不在这里处理，交给 Retrieval 基础设施层的 VectorStore（见
app/retrieval/vectorstore.py）。本层只负责 notes / blocks / blocks_fts 三张表的访问，
返回领域记录（NoteRecord / BlockHit / FtsHit），不负责业务编排。
"""

from __future__ import annotations

import json
import sqlite3
from contextlib import nullcontext
from dataclasses import dataclass, field
from datetime import datetime

from app.contracts import NoteBlock
from app.database.db import connect, transaction
from app.textutils import segment


def _iso(dt: datetime) -> str:
    return dt.isoformat()


def _parse_dt(value: str) -> datetime:
    return datetime.fromisoformat(value)


@dataclass
class NoteRecord:
    note_id: str
    title: str
    file_path: str
    folder: str
    tags: list[str]
    created_at: datetime
    updated_at: datetime
    blocks: list[NoteBlock] = field(default_factory=list)


@dataclass
class BlockHit:
    """检索时返回的完整 Block 上下文，用于组装 Citation 与 metadata 过滤。"""

    block_id: str
    note_id: str
    title: str
    file_path: str
    folder: str
    heading_path: list[str]
    content: str
    start_offset: int
    end_offset: int
    tags: list[str]
    created_at: datetime
    updated_at: datetime


@dataclass
class FtsHit:
    block_id: str
    note_id: str
    bm25: float


def replace_note_metadata(
    *,
    conn: sqlite3.Connection,
    note_id: str,
    title: str,
    file_path: str,
    folder: str,
    tags: list[str],
    created_at: datetime,
    updated_at: datetime,
    blocks: list[NoteBlock],
) -> list[str]:
    """整体替换一条笔记的元数据、Block 与 FTS5 索引。

    不在此处开启/提交事务：由调用方（index_note）在同一连接上把「元数据 + 向量」包进
    单个事务，保证原子性。返回替换前的旧 block_id 列表，供调用方清理失效向量。
    """
    old_block_ids = [
        row["block_id"]
        for row in conn.execute("SELECT block_id FROM blocks WHERE note_id = ?", (note_id,))
    ]
    conn.execute(
        """
        INSERT INTO notes (note_id, title, file_path, folder, tags, created_at, updated_at)
        VALUES (?, ?, ?, ?, ?, ?, ?)
        ON CONFLICT(note_id) DO UPDATE SET
            title = excluded.title,
            file_path = excluded.file_path,
            folder = excluded.folder,
            tags = excluded.tags,
            updated_at = excluded.updated_at
        """,
        (note_id, title, file_path, folder, json.dumps(tags, ensure_ascii=False),
         _iso(created_at), _iso(updated_at)),
    )
    conn.execute("DELETE FROM blocks WHERE note_id = ?", (note_id,))
    conn.execute("DELETE FROM blocks_fts WHERE note_id = ?", (note_id,))
    for position, block in enumerate(blocks):
        conn.execute(
            """
            INSERT INTO blocks
                (block_id, note_id, heading_path, start_offset, end_offset,
                 content, content_hash, token_count, position)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (block.block_id, note_id, json.dumps(block.heading_path, ensure_ascii=False),
             block.start_offset, block.end_offset, block.content,
             block.content_hash, block.token_count, position),
        )
        # FTS5 存分词后的可检索文本；原文仍由 blocks.content 保留用于展示
        conn.execute(
            "INSERT INTO blocks_fts (block_id, note_id, heading_path, content) VALUES (?, ?, ?, ?)",
            (block.block_id, note_id, segment(" ".join(block.heading_path)), segment(block.content)),
        )
    return old_block_ids


def delete_note(
    note_id: str, *, conn: sqlite3.Connection | None = None
) -> list[str]:
    """删除笔记及其 Block、FTS5 索引；返回被删除的 block_id 供向量层清理。"""
    owns = conn is None
    conn = conn or connect()
    try:
        block_ids = [
            row["block_id"]
            for row in conn.execute("SELECT block_id FROM blocks WHERE note_id = ?", (note_id,))
        ]
        with transaction(conn) if owns else nullcontext():
            conn.execute("DELETE FROM blocks_fts WHERE note_id = ?", (note_id,))
            conn.execute("DELETE FROM notes WHERE note_id = ?", (note_id,))  # blocks 级联删除
        return block_ids
    finally:
        if owns:
            conn.close()


def get_note_record(note_id: str) -> NoteRecord | None:
    conn = connect()
    try:
        row = conn.execute("SELECT * FROM notes WHERE note_id = ?", (note_id,)).fetchone()
        if row is None:
            return None
        blocks = [
            _block_from_row(b)
            for b in conn.execute("SELECT * FROM blocks WHERE note_id = ? ORDER BY position", (note_id,))
        ]
        return NoteRecord(
            note_id=row["note_id"],
            title=row["title"],
            file_path=row["file_path"],
            folder=row["folder"],
            tags=json.loads(row["tags"] or "[]"),
            created_at=_parse_dt(row["created_at"]),
            updated_at=_parse_dt(row["updated_at"]),
            blocks=blocks,
        )
    finally:
        conn.close()


def list_note_summaries(
    *, limit: int = 50, offset: int = 0, folder: str | None = None, tag: str | None = None
) -> tuple[list, int]:
    conn = connect()
    try:
        where: list[str] = []
        params: list[str] = []
        if folder:
            where.append("folder = ?")
            params.append(folder)
        if tag:
            where.append("EXISTS (SELECT 1 FROM json_each(notes.tags) AS j WHERE j.value = ?)")
            params.append(tag)
        where_sql = ("WHERE " + " AND ".join(where)) if where else ""

        total = conn.execute(f"SELECT COUNT(*) FROM notes {where_sql}", params).fetchone()[0]
        rows = conn.execute(
            f"SELECT * FROM notes {where_sql} ORDER BY updated_at DESC LIMIT ? OFFSET ?",
            params + [limit, offset],
        ).fetchall()

        items = [
            {
                "note_id": r["note_id"],
                "title": r["title"],
                "file_path": r["file_path"],
                "tags": json.loads(r["tags"] or "[]"),
                "created_at": _parse_dt(r["created_at"]),
                "updated_at": _parse_dt(r["updated_at"]),
            }
            for r in rows
        ]
        return items, total
    finally:
        conn.close()


def fts_search(match: str, limit: int = 100) -> list[FtsHit]:
    conn = connect()
    try:
        rows = conn.execute(
            """
            SELECT block_id, note_id, bm25(blocks_fts) AS rank
            FROM blocks_fts
            WHERE blocks_fts MATCH ?
            ORDER BY rank
            LIMIT ?
            """,
            (match, limit),
        ).fetchall()
        return [FtsHit(block_id=r["block_id"], note_id=r["note_id"], bm25=r["rank"]) for r in rows]
    finally:
        conn.close()


def fts_search_page(
    *,
    match: str,
    limit: int,
    offset: int,
    folders: list[str],
    note_ids: list[str],
    tags: list[str],
    created_from: datetime | None,
    created_to: datetime | None,
    updated_from: datetime | None,
    updated_to: datetime | None,
) -> tuple[list[FtsHit], int]:
    """执行带元数据过滤的 FTS 精确分页，并返回过滤后的完整命中数。"""
    where = ["blocks_fts MATCH ?"]
    params: list[object] = [match]

    def add_in(column: str, values: list[str]) -> None:
        if not values:
            return
        placeholders = ",".join("?" * len(values))
        where.append(f"{column} IN ({placeholders})")
        params.extend(values)

    add_in("n.folder", folders)
    add_in("n.note_id", note_ids)
    if tags:
        placeholders = ",".join("?" * len(tags))
        where.append(
            f"EXISTS (SELECT 1 FROM json_each(n.tags) AS tag WHERE tag.value IN ({placeholders}))"
        )
        params.extend(tags)

    for column, lower, upper in (
        ("n.created_at", created_from, created_to),
        ("n.updated_at", updated_from, updated_to),
    ):
        if lower is not None:
            where.append(f"julianday({column}) >= julianday(?)")
            params.append(_iso(lower))
        if upper is not None:
            where.append(f"julianday({column}) <= julianday(?)")
            params.append(_iso(upper))

    from_sql = """
        FROM blocks_fts
        JOIN blocks AS b ON b.block_id = blocks_fts.block_id
        JOIN notes AS n ON n.note_id = b.note_id
    """
    where_sql = " AND ".join(where)

    conn = connect()
    try:
        total = conn.execute(
            f"SELECT COUNT(*) {from_sql} WHERE {where_sql}", params
        ).fetchone()[0]
        rows = conn.execute(
            f"""
            SELECT blocks_fts.block_id, blocks_fts.note_id, bm25(blocks_fts) AS rank
            {from_sql}
            WHERE {where_sql}
            ORDER BY rank
            LIMIT ? OFFSET ?
            """,
            [*params, limit, offset],
        ).fetchall()
        return (
            [FtsHit(block_id=row["block_id"], note_id=row["note_id"], bm25=row["rank"])
             for row in rows],
            total,
        )
    finally:
        conn.close()


def get_block_hits(block_ids: list[str]) -> list[BlockHit]:
    if not block_ids:
        return []
    conn = connect()
    try:
        placeholders = ",".join("?" * len(block_ids))
        rows = conn.execute(
            f"""
            SELECT b.block_id, b.note_id, b.heading_path, b.start_offset, b.end_offset, b.content,
                   n.title, n.file_path, n.folder, n.tags, n.created_at, n.updated_at
            FROM blocks b
            JOIN notes n ON n.note_id = b.note_id
            WHERE b.block_id IN ({placeholders})
            """,
            block_ids,
        ).fetchall()
        return [_block_hit_from_row(r) for r in rows]
    finally:
        conn.close()


def set_index_meta(
    kv: dict[str, str], *, conn: sqlite3.Connection | None = None
) -> None:
    """写入索引元信息；传入连接时加入调用方现有事务。"""
    owns = conn is None
    conn = conn or connect()
    try:
        with transaction(conn) if owns else nullcontext():
            for key, value in kv.items():
                conn.execute("INSERT OR REPLACE INTO index_meta (key, value) VALUES (?, ?)", (key, value))
    finally:
        if owns:
            conn.close()


def get_index_meta() -> dict[str, str]:
    conn = connect()
    try:
        return {r["key"]: r["value"] for r in conn.execute("SELECT key, value FROM index_meta")}
    finally:
        conn.close()


def clear_all() -> None:
    """清空元数据、Block 与 FTS5（重建索引用，向量由 VectorStore.clear 处理）。"""
    conn = connect()
    try:
        with transaction(conn):
            conn.execute("DELETE FROM blocks_fts")
            conn.execute("DELETE FROM blocks")
            conn.execute("DELETE FROM notes")
    finally:
        conn.close()


def stats() -> dict[str, int]:
    conn = connect()
    try:
        notes = conn.execute("SELECT COUNT(*) AS c FROM notes").fetchone()["c"]
        blocks = conn.execute("SELECT COUNT(*) AS c FROM blocks").fetchone()["c"]
        return {"notes": notes, "blocks": blocks}
    finally:
        conn.close()


def _block_from_row(row) -> NoteBlock:
    return NoteBlock(
        block_id=row["block_id"],
        note_id=row["note_id"],
        heading_path=json.loads(row["heading_path"] or "[]"),
        start_offset=row["start_offset"],
        end_offset=row["end_offset"],
        content=row["content"],
        content_hash=row["content_hash"],
        token_count=row["token_count"],
    )


def _block_hit_from_row(row) -> BlockHit:
    return BlockHit(
        block_id=row["block_id"],
        note_id=row["note_id"],
        title=row["title"],
        file_path=row["file_path"],
        folder=row["folder"],
        heading_path=json.loads(row["heading_path"] or "[]"),
        content=row["content"],
        start_offset=row["start_offset"],
        end_offset=row["end_offset"],
        tags=json.loads(row["tags"] or "[]"),
        created_at=_parse_dt(row["created_at"]),
        updated_at=_parse_dt(row["updated_at"]),
    )