playchat/backend/app/core/notes.py

"""ユーザーとAIが共同で編集する共有ノート。

システムプロンプトにはページの「タイトル一覧」だけを注入し、本文はモデルが
read_note ツールで必要なときだけ読み出す。書き込みは write_note ツールで行う。

例外として、ユーザーがセッションに紐付けたページ（SessionNotePage）は本文も
毎回そのままシステムプロンプトに載る。ツールの往復を待たずに確実に参照させたい
資料のための仕組みで、指定はセッション単位なので会話ごとに組み合わせを変えられる。
"""
import io
import logging
import re
import zipfile

from sqlalchemy.orm import Session as DBSession

from app.core.config import (
    NOTES_IMPORT_MAX_BYTES,
    NOTES_IMPORT_MAX_PAGES,
    NOTES_TOOL_CONTENT_LIMIT,
    SESSION_NOTES_CONTENT_LIMIT,
)
from app.db.models import NotePage, SessionNotePage

logger = logging.getLogger(__name__)


READ_NOTE_TOOL = {
    "type": "function",
    "function": {
        "name": "read_note",
        "description": "共有ノートの特定ページの本文を読み出す。システムプロンプトに提示されたタイトル一覧の中から、"
        "今の会話に関係しそうなページがあるときに使う。",
        "parameters": {
            "type": "object",
            "properties": {
                "title": {"type": "string", "description": "読み出すページのタイトル（タイトル一覧に載っているもの）"},
            },
            "required": ["title"],
        },
    },
}

WRITE_NOTE_TOOL = {
    "type": "function",
    "function": {
        "name": "write_note",
        "description": "共有ノートのページに書き込む。ユーザーからノートへの記録・追記・整理を頼まれたときに使う。"
        "存在しないタイトルを指定した場合は新しいページを作成する。"
        "既存ページを replace で上書きすると元の内容は失われるので、追記でよい場合は必ず append を使うこと。",
        "parameters": {
            "type": "object",
            "properties": {
                "title": {"type": "string", "description": "書き込む対象のページタイトル"},
                "content": {"type": "string", "description": "書き込む本文（Markdown可）"},
                "mode": {
                    "type": "string",
                    "enum": ["append", "replace"],
                    "description": "append=末尾に追記（既定） / replace=ページ全体を置き換え",
                },
            },
            "required": ["title", "content"],
        },
    },
}

NOTE_TOOLS = [READ_NOTE_TOOL, WRITE_NOTE_TOOL]


def _find_page(db: DBSession, title: str) -> NotePage | None:
    """タイトルでページを探す。完全一致 → 大文字小文字・前後空白を無視した一致 の順に試す。"""
    pages = db.query(NotePage).order_by(NotePage.position, NotePage.id).all()
    for page in pages:
        if page.title == title:
            return page
    normalized = title.strip().casefold()
    for page in pages:
        if page.title.strip().casefold() == normalized:
            return page
    return None


def next_position(db: DBSession) -> int:
    """新規ページを末尾に置くための position を返す。"""
    last = db.query(NotePage).order_by(NotePage.position.desc()).first()
    return (last.position + 1) if last else 0


def set_section(db: DBSession, page_id: int, section: str | None) -> NotePage | None:
    """ページの所属セクションを変更し、同じセクションのページ群の直後に連続配置し直す。

    タブバーはposition順に描画され、隣接して同じsectionが続く区間を1つのグループとして
    まとめて表示する（app.api.notes / フロントのグルーピング表示を参照）。そのため
    セクションを変更したページは、常に同じセクションの最後のページの直後に移動する。
    """
    page = db.get(NotePage, page_id)
    if not page:
        return None

    normalized = (section or "").strip() or None
    ordered = [p for p in db.query(NotePage).order_by(NotePage.position, NotePage.id).all() if p.id != page_id]

    if normalized is None:
        insert_at = len(ordered)
    else:
        matching = [i for i, p in enumerate(ordered) if p.section == normalized]
        insert_at = (matching[-1] + 1) if matching else len(ordered)

    ordered.insert(insert_at, page)
    page.section = normalized
    for i, p in enumerate(ordered):
        p.position = i
    db.commit()
    db.refresh(page)
    return page


def reorder(db: DBSession, page_id: int, section: str | None, order: list[int]) -> bool:
    """タブのドラッグ&ドロップによる並べ替えを反映する。

    order にはドラッグ後の全ページIDを新しい順序で渡す（フロント側で計算する）。
    page_id のページには先に section を設定し、その後 order の順に position を振り直す。
    こうすることで、別セクションのタブの前後にドロップした場合はそのセクションへの
    移動を兼ね、同じセクション内でのドロップなら単純な並べ替えになる。

    order が現在の全ページIDの集合と一致しない場合は何もせず False を返す。
    """
    pages = {p.id: p for p in db.query(NotePage).all()}
    if page_id not in pages or set(order) != set(pages.keys()):
        return False

    pages[page_id].section = (section or "").strip() or None
    for i, pid in enumerate(order):
        pages[pid].position = i
    db.commit()
    return True


def rename_section(db: DBSession, old_name: str, new_name: str | None) -> int:
    """あるセクションに属する全ページのセクション名を一括で変更する。new_name が空ならグループ解除になる。

    戻り値: 変更したページ数。
    """
    normalized_new = (new_name or "").strip() or None
    pages = db.query(NotePage).filter(NotePage.section == old_name).order_by(NotePage.position, NotePage.id).all()
    for page in pages:
        page.section = normalized_new
    db.commit()
    return len(pages)


# 取り込めるテキストファイルの拡張子。これ以外はZIP内でも単体でも無視する
_TEXT_SUFFIXES = (".md", ".markdown", ".txt", ".text")
# ダウンロードしたZIPのファイル名に付く並び順の通し番号（01_タイトル.md）
_EXPORT_PREFIX = re.compile(r"^\d{2,}_")


def _decode(raw: bytes) -> str:
    """取り込んだファイルを文字列にする。UTF-8優先、日本語Windows由来のCP932も拾う。"""
    for encoding in ("utf-8-sig", "utf-8", "cp932"):
        try:
            return raw.decode(encoding)
        except UnicodeDecodeError:
            continue
    return raw.decode("utf-8", errors="replace")


def _title_from_filename(filename: str) -> str:
    """ファイル名からページタイトルを作る。拡張子とエクスポート時の通し番号を落とす。"""
    stem = filename.rsplit("/", 1)[-1].rsplit("\\", 1)[-1]
    for suffix in _TEXT_SUFFIXES:
        if stem.lower().endswith(suffix):
            stem = stem[: -len(suffix)]
            break
    return _EXPORT_PREFIX.sub("", stem).strip() or "無題のページ"


def _unique_title(db: DBSession, title: str, taken: set[str]) -> str:
    """既存ページや同時取り込み分とタイトルがぶつかったら「(2)」を付けてずらす。

    上書きせずに増やすのは、エクスポートしたZIPを既存のノートに投げ戻したときに
    元のページを壊さないため。要らない方は後からタブの × で消せる。
    """
    existing = {p.title for p in db.query(NotePage.title).all()} | taken
    if title not in existing:
        return title
    n = 2
    while f"{title} ({n})" in existing:
        n += 1
    return f"{title} ({n})"


def _create_imported_page(db: DBSession, title: str, content: str, section: str | None) -> NotePage:
    page = NotePage(title=title, content=content, position=next_position(db))
    db.add(page)
    db.commit()
    db.refresh(page)
    if section:
        # 同じセクションの末尾に並べ直す（連続配置の前提を保つ）
        page = set_section(db, page.id, section) or page
    return page


def import_files(db: DBSession, files: list[tuple[str, bytes]]) -> list[NotePage]:
    """アップロードされたファイルをノートのページとして取り込む。

    files は (ファイル名, 中身) のリスト。.md / .txt はそのまま1ページになり、.zip は
    中のテキストファイルを1つずつページにする（ZIP内のフォルダ名がセクションになる）。
    エクスポートで出したZIPをそのまま投げ戻せば、セクションと並び順ごと復元できる。

    取り込めないファイル（拡張子が対象外・大きすぎる）は黙って飛ばす。作成したページを返す。
    """
    created: list[NotePage] = []
    taken: set[str] = set()

    def add(filename: str, raw: bytes, section: str | None) -> None:
        if len(created) >= NOTES_IMPORT_MAX_PAGES or len(raw) > NOTES_IMPORT_MAX_BYTES:
            logger.info("Skipped note import: %s (%d bytes)", filename, len(raw))
            return
        title = _unique_title(db, _title_from_filename(filename), taken)
        taken.add(title)
        created.append(_create_imported_page(db, title, _decode(raw), section))

    for filename, raw in files:
        if filename.lower().endswith(".zip"):
            try:
                with zipfile.ZipFile(io.BytesIO(raw)) as archive:
                    for info in archive.infolist():
                        if info.is_dir() or not info.filename.lower().endswith(_TEXT_SUFFIXES):
                            continue
                        if info.file_size > NOTES_IMPORT_MAX_BYTES:
                            logger.info("Skipped zip entry: %s (%d bytes)", info.filename, info.file_size)
                            continue
                        # ZIP内のフォルダ名をセクションにする（ネストは最初の階層だけ見る）
                        parts = info.filename.replace("\\", "/").split("/")
                        section = parts[0] if len(parts) > 1 and parts[0] not in ("", ".", "..") else None
                        add(info.filename, archive.read(info), section)
            except zipfile.BadZipFile:
                logger.warning("Not a valid zip: %s", filename)
        elif filename.lower().endswith(_TEXT_SUFFIXES):
            add(filename, raw, None)
        else:
            logger.info("Skipped unsupported file: %s", filename)

    return created


def get_attached_page_ids(db: DBSession, session_id: str) -> list[int]:
    """セッションに紐付けられた（本文をプロンプトに載せる）ページIDを position 順で返す。"""
    rows = (
        db.query(SessionNotePage.page_id)
        .join(NotePage, NotePage.id == SessionNotePage.page_id)
        .filter(SessionNotePage.session_id == session_id)
        .order_by(NotePage.position, NotePage.id)
        .all()
    )
    return [row[0] for row in rows]


def set_attached_pages(db: DBSession, session_id: str, page_ids: list[int]) -> list[int]:
    """セッションに紐付けるページを page_ids で置き換える。存在しないIDは黙って捨てる。

    戻り値は実際に紐付いたページIDを position 順で並べたもの。
    """
    valid = {p.id for p in db.query(NotePage.id).filter(NotePage.id.in_(page_ids)).all()} if page_ids else set()

    existing = db.query(SessionNotePage).filter(SessionNotePage.session_id == session_id).all()
    for link in existing:
        if link.page_id not in valid:
            db.delete(link)
    already = {link.page_id for link in existing if link.page_id in valid}
    for page_id in valid - already:
        db.add(SessionNotePage(session_id=session_id, page_id=page_id))
    db.commit()
    return get_attached_page_ids(db, session_id)


def get_session_notes_context(db: DBSession, session_id: str) -> str | None:
    """セッションに紐付けられたページの本文を、システムプロンプト用のテキストとして返す。

    1ページも紐付いていなければ None。1ページあたり SESSION_NOTES_CONTENT_LIMIT 文字で切る。
    """
    pages = (
        db.query(NotePage)
        .join(SessionNotePage, SessionNotePage.page_id == NotePage.id)
        .filter(SessionNotePage.session_id == session_id)
        .order_by(NotePage.position, NotePage.id)
        .all()
    )
    if not pages:
        return None

    lines = [
        "【この会話で参照するノート】",
        "ユーザーがこの会話に紐付けたページです。本文をそのまま載せているので、"
        "これらのページに read_note を使う必要はありません。",
        "",
    ]
    for page in pages:
        content = page.content.strip() or "(空)"
        if len(content) > SESSION_NOTES_CONTENT_LIMIT:
            content = content[:SESSION_NOTES_CONTENT_LIMIT] + "\n…(以下省略。全文は read_note で読めます)"
        lines += [f"## {page.title}", content, ""]
    return "\n".join(lines).rstrip()


def get_notes_context(db: DBSession, session_id: str | None = None) -> str | None:
    """共有ノートのタイトル一覧をシステムプロンプト用のテキストとして返す。セクションごとに見出しを付ける。

    session_id を渡すと、そのセッションに紐付いた（本文が別途プロンプトに載る）ページに
    印を付ける。同じ本文を read_note で読み直させないため。
    """
    pages = db.query(NotePage).order_by(NotePage.position, NotePage.id).all()
    attached = set(get_attached_page_ids(db, session_id)) if session_id else set()

    lines = [
        "【共有ノート】",
        "ユーザーとあなたが共同で編集できるノートです。以下のページがあります"
        "（セクションはページをグループ化する見出しで、あれば併記します）。",
        "",
    ]
    if pages:
        current_section = "__unset__"  # どのページのsectionとも一致しない番兵
        for page in pages:
            if page.section != current_section:
                current_section = page.section
                if current_section:
                    lines.append(f"## {current_section}")
            if page.id in attached:
                lines.append(f"- {page.title} — （本文は「この会話で参照するノート」に全文あり）")
                continue
            preview = " ".join(page.content.split())[:40]
            suffix = f" — {preview}…" if preview else " — (空)"
            lines.append(f"- {page.title}{suffix}")
    else:
        lines.append("- (まだページがありません)")
    lines += [
        "",
        "本文はここには含まれていません。内容が必要なときは read_note ツールで読み出してください。",
        "ノートへの記録・追記を頼まれたときは write_note ツールで書き込んでください"
        "（新しいタイトルを指定すればページが新規作成されます）。",
        "ユーザーが明示的に頼んでいないのに勝手に書き込まないこと。",
    ]
    return "\n".join(lines)


def read_note(db: DBSession, title: str) -> str:
    """ページ本文をツール結果用のテキストとして返す。"""
    page = _find_page(db, title)
    if not page:
        available = ", ".join(p.title for p in db.query(NotePage).order_by(NotePage.position, NotePage.id).all())
        return f"「{title}」というページはありません。現在のページ: {available or '(なし)'}"
    if not page.content.strip():
        return f"ページ「{page.title}」は空です。"
    content = page.content
    if len(content) > NOTES_TOOL_CONTENT_LIMIT:
        content = content[:NOTES_TOOL_CONTENT_LIMIT] + "\n…(以下省略)"
    return f"ページ「{page.title}」の本文:\n\n{content}"


_LIST_ITEM = re.compile(r"^\s*([-*+]|\d+[.)])\s")


def _append(existing: str, addition: str) -> str:
    """既存本文に追記する。箇条書き同士をつなぐ場合だけ空行を挟まない。"""
    if not existing.strip():
        return addition
    body = existing.rstrip()
    last_line = body.rsplit("\n", 1)[-1]
    first_line = addition.lstrip("\n").split("\n", 1)[0]
    separator = "\n" if _LIST_ITEM.match(last_line) and _LIST_ITEM.match(first_line) else "\n\n"
    return f"{body}{separator}{addition}"


def write_note(db: DBSession, title: str, content: str, mode: str = "append") -> str:
    """ページに書き込む。存在しなければ新規作成する。"""
    if not title.strip():
        return "タイトルが空のため書き込めませんでした。"

    page = _find_page(db, title)
    if page is None:
        page = NotePage(title=title.strip(), content=content, position=next_position(db))
        db.add(page)
        db.commit()
        return f"新しいページ「{page.title}」を作成し、本文を書き込みました。"

    if mode == "replace":
        page.content = content
        result = f"ページ「{page.title}」の内容を置き換えました。"
    else:
        page.content = _append(page.content, content)
        result = f"ページ「{page.title}」に追記しました。"
    db.commit()
    return result


def execute_note_tool(db: DBSession, name: str, args: dict) -> tuple[str, bool]:
    """ノート系ツールを実行する。戻り値: (ツール結果テキスト, ノートを変更したか)。"""
    try:
        if name == "read_note":
            return read_note(db, args.get("title", "")), False
        if name == "write_note":
            result = write_note(
                db,
                args.get("title", ""),
                args.get("content", ""),
                args.get("mode", "append"),
            )
            return result, True
    except Exception as e:
        db.rollback()
        logger.exception("Note tool %s failed", name)
        return f"ノート操作に失敗しました: {e}", False
    return f"未知のツールです: {name}", False