hunjing/backend/app/services/lorebook.py

"""世谱服务：注入编排 + CRUD。

注入编排（build_lorebook_injection）：查条目 → 匹配 → 预算 → 分组为 before/after 文本，
IO 在此层完成；纯匹配逻辑在 app.modules.lorebook。
"""

from __future__ import annotations

import json
from collections import defaultdict

from sqlalchemy.ext.asyncio import AsyncSession

from app.core.errors import AppError
from app.db.session import AsyncSessionLocal
from app.models.lorebook import Lorebook as LorebookOrm
from app.models.lorebook import LorebookEntry as LorebookEntryOrm
from app.modules.lorebook import (
    LorebookBook,
    LorebookEntry,
    build_injections,
    match_entries,
    select_within_budget,
)
from app.repositories import lorebook as lorebook_repo
from app.repositories.json_utils import dumps
from app.schemas.lorebook import (
    LorebookCreateRequest,
    LorebookEntryCreateRequest,
    LorebookEntryResponse,
    LorebookEntryUpdateRequest,
    LorebookResponse,
    LorebookUpdateRequest,
)

# ==================== 注入编排（供 send_message 调用）====================


async def build_lorebook_injection(
    session: AsyncSession,
    character_id: str,
    owner_scope: str,
    scan_text: str,
) -> tuple[str | None, str | None]:
    """扫描对话文本，返回世谱注入的 (before_char, after_char)；无命中返回 (None, None)。

    每本书各自按其 token_budget 裁剪，再跨书合并按 position 分组。
    """
    books = await lorebook_repo.list_active_books(session, character_id, owner_scope)
    if not books:
        return None, None

    entries = await lorebook_repo.list_entries_for_books(session, [b.id for b in books])
    if not entries:
        return None, None

    entries_by_book: dict[str, list[LorebookEntryOrm]] = defaultdict(list)
    for entry in entries:
        entries_by_book[entry.lorebook_id].append(entry)

    kept_all = []
    for book_orm in books:
        module_book = _to_module_book(book_orm, entries_by_book.get(book_orm.id, []))
        matched = match_entries(module_book, scan_text)
        kept_all.extend(
            select_within_budget(matched, token_budget=module_book.token_budget)
        )

    injections = build_injections(kept_all)
    return injections.get("before_char"), injections.get("after_char")


def _to_module_book(book_orm: LorebookOrm, entries_orm: list[LorebookEntryOrm]) -> LorebookBook:
    return LorebookBook(
        id=book_orm.id,
        name=book_orm.name,
        description=book_orm.description,
        entries=[_to_module_entry(e) for e in entries_orm],
        scan_depth=book_orm.scan_depth,
        token_budget=book_orm.token_budget,
        recursive_scanning=book_orm.recursive_scanning,
        scope=book_orm.scope,
    )


def _to_module_entry(entry_orm: LorebookEntryOrm) -> LorebookEntry:
    return LorebookEntry(
        id=entry_orm.id,
        keys=json.loads(entry_orm.keys_json or "[]"),
        content=entry_orm.content,
        secondary_keys=json.loads(entry_orm.secondary_keys_json or "[]"),
        constant=entry_orm.constant,
        selective=entry_orm.selective,
        enabled=entry_orm.enabled,
        case_sensitive=entry_orm.case_sensitive,
        use_regex=entry_orm.use_regex,
        position=entry_orm.position,
        insertion_order=entry_orm.insertion_order,
        priority=entry_orm.priority,
        comment=entry_orm.comment,
        bound_entity_id=entry_orm.bound_entity_id,
    )


# ==================== CRUD（供 API 调用）====================


async def create_lorebook(
    session: AsyncSession, payload: LorebookCreateRequest, owner_scope: str
) -> LorebookResponse:
    book = LorebookOrm(
        character_id=payload.character_id,
        name=payload.name,
        description=payload.description,
        scan_depth=payload.scan_depth,
        token_budget=payload.token_budget,
        recursive_scanning=payload.recursive_scanning,
        scope=payload.scope,
        owner_scope=owner_scope,
    )
    session.add(book)
    await session.flush()
    response = _book_response(book, [])
    await session.commit()
    return response


async def list_lorebooks(
    *,
    character_id: str | None,
    owner_scope: str,
    scope: str | None = None,
) -> list[LorebookResponse]:
    async with AsyncSessionLocal() as session:
        books = await lorebook_repo.list_books(session, character_id, owner_scope, scope=scope)
        return [_book_response(b, []) for b in books]


async def get_lorebook(lorebook_id: str, owner_scope: str) -> LorebookResponse:
    async with AsyncSessionLocal() as session:
        book = await _require_book(session, lorebook_id, owner_scope)
        entries = await lorebook_repo.list_book_entries(session, lorebook_id)
        return _book_response(book, entries)


async def update_lorebook(
    lorebook_id: str, payload: LorebookUpdateRequest, owner_scope: str
) -> LorebookResponse:
    async with AsyncSessionLocal() as session:
        book = await _require_book(session, lorebook_id, owner_scope)
        for field in ("name", "description", "scan_depth", "token_budget", "recursive_scanning"):
            value = getattr(payload, field)
            if value is not None:
                setattr(book, field, value)
        await session.commit()
        entries = await lorebook_repo.list_book_entries(session, lorebook_id)
        return _book_response(book, entries)


async def delete_lorebook(lorebook_id: str, owner_scope: str) -> LorebookResponse:
    async with AsyncSessionLocal() as session:
        book = await _require_book(session, lorebook_id, owner_scope)
        book.status = "archived"
        await session.commit()
        return _book_response(book, [])


async def create_entry(
    session: AsyncSession,
    lorebook_id: str,
    payload: LorebookEntryCreateRequest,
    owner_scope: str,
) -> LorebookEntryResponse:
    await _require_book(session, lorebook_id, owner_scope)
    entry = LorebookEntryOrm(
        lorebook_id=lorebook_id,
        keys_json=dumps(payload.keys),
        content=payload.content,
        secondary_keys_json=dumps(payload.secondary_keys),
        constant=payload.constant,
        selective=payload.selective,
        enabled=payload.enabled,
        case_sensitive=payload.case_sensitive,
        use_regex=payload.use_regex,
        position=payload.position,
        insertion_order=payload.insertion_order,
        priority=payload.priority,
        comment=payload.comment,
        bound_entity_id=payload.bound_entity_id,
    )
    session.add(entry)
    await session.flush()
    response = _entry_response(entry)
    await session.commit()
    return response


async def update_entry(
    lorebook_id: str,
    entry_id: str,
    payload: LorebookEntryUpdateRequest,
    owner_scope: str,
) -> LorebookEntryResponse:
    async with AsyncSessionLocal() as session:
        await _require_book(session, lorebook_id, owner_scope)
        entry = await _require_entry(session, entry_id, lorebook_id)
        for key, value in payload.model_dump(exclude_unset=True).items():
            if key in ("keys", "secondary_keys"):
                setattr(entry, f"{key}_json", dumps(value))
            else:
                setattr(entry, key, value)
        await session.commit()
        return _entry_response(entry)


async def delete_entry(lorebook_id: str, entry_id: str, owner_scope: str) -> None:
    async with AsyncSessionLocal() as session:
        await _require_book(session, lorebook_id, owner_scope)
        entry = await _require_entry(session, entry_id, lorebook_id)
        await session.delete(entry)
        await session.commit()


async def _require_book(
    session: AsyncSession, lorebook_id: str, owner_scope: str
) -> LorebookOrm:
    book = await lorebook_repo.get_book(session, lorebook_id)
    if book is None or book.owner_scope != owner_scope or book.status == "archived":
        raise AppError("not_found", "世谱不存在。")
    return book


async def _require_entry(
    session: AsyncSession, entry_id: str, lorebook_id: str
) -> LorebookEntryOrm:
    entry = await lorebook_repo.get_entry(session, entry_id)
    if entry is None or entry.lorebook_id != lorebook_id:
        raise AppError("not_found", "条目不存在。")
    return entry


def _entry_response(entry: LorebookEntryOrm) -> LorebookEntryResponse:
    return LorebookEntryResponse(
        id=entry.id,
        lorebook_id=entry.lorebook_id,
        keys=json.loads(entry.keys_json or "[]"),
        content=entry.content,
        secondary_keys=json.loads(entry.secondary_keys_json or "[]"),
        constant=entry.constant,
        selective=entry.selective,
        enabled=entry.enabled,
        case_sensitive=entry.case_sensitive,
        use_regex=entry.use_regex,
        position=entry.position,
        insertion_order=entry.insertion_order,
        priority=entry.priority,
        comment=entry.comment,
        bound_entity_id=entry.bound_entity_id,
    )


def _book_response(
    book: LorebookOrm, entries: list[LorebookEntryOrm]
) -> LorebookResponse:
    return LorebookResponse(
        id=book.id,
        character_id=book.character_id,
        name=book.name,
        description=book.description,
        scope=book.scope,
        scan_depth=book.scan_depth,
        token_budget=book.token_budget,
        recursive_scanning=book.recursive_scanning,
        status=book.status,
        entries=[_entry_response(e) for e in entries],
    )