"""EPUB ingestion into Richie DB.""" from __future__ import annotations import asyncio import hashlib import logging from dataclasses import dataclass from datetime import UTC, datetime from pathlib import Path from typing import TYPE_CHECKING import tiktoken from anyio import Path as AsyncPath from sqlalchemy import or_, select from python.ebook_search.epub_parse import parse_epub from python.ebook_search.protected_phrases.matching import index_chunk_phrase_mentions_for_book from python.orm.richie import EbookChapter, EbookChunk, EbookSource logger = logging.getLogger(__name__) DEFAULT_CHUNK_TOKENS = 700 DEFAULT_CHUNK_OVERLAP = 100 if TYPE_CHECKING: from sqlalchemy.ext.asyncio import AsyncSession from python.ebook_search.config import EbookSearchConfig from python.ebook_search.epub_parse import ParsedChapter @dataclass(frozen=True) class TextChunk: """A token-bounded chunk of text.""" text: str token_start: int token_count: int def chunk_text( text: str, *, chunk_tokens: int = DEFAULT_CHUNK_TOKENS, overlap_tokens: int = DEFAULT_CHUNK_OVERLAP, ) -> list[TextChunk]: """Split text into overlapping token chunks.""" if chunk_tokens <= 0: msg = "chunk_tokens must be positive" raise ValueError(msg) if overlap_tokens < 0 or overlap_tokens >= chunk_tokens: msg = "overlap_tokens must be non-negative and smaller than chunk_tokens" raise ValueError(msg) encoding = tiktoken.get_encoding("cl100k_base") tokens = encoding.encode(text) if not tokens: return [] chunks: list[TextChunk] = [] step = chunk_tokens - overlap_tokens for start in range(0, len(tokens), step): chunk = tokens[start : start + chunk_tokens] if not chunk: continue chunks.append( TextChunk( text=encoding.decode(chunk).strip(), token_start=start, token_count=len(chunk), ) ) if start + chunk_tokens >= len(tokens): break return [chunk for chunk in chunks if chunk.text] async def find_library_epubs(library_path: str) -> tuple[AsyncPath, list[AsyncPath] | None]: """Resolve one configured library path and collect its EPUB files asynchronously. Returns: tuple[Path, list[Path] | None]: The expanded path and its EPUB files, or ``None`` when the path is neither an EPUB file nor a directory. """ path = await AsyncPath(library_path).expanduser() if await path.is_file() and path.suffix.lower() == ".epub": return path, [path] if await path.is_dir(): return path, sorted([epub_path async for epub_path in path.rglob("*.epub")]) return path, None async def ingest_configured_paths(session: AsyncSession, config: EbookSearchConfig) -> int: """Ingest every EPUB found under configured library paths.""" count = 0 for library_path in config.library_paths: path, epub_paths = await find_library_epubs(library_path) logger.info(f"ebook_ingest_path_start {path=}") if epub_paths is None: logger.warning(f"ebook_ingest_path_missing {path=}") continue for epub_path in epub_paths: count += int(await ingest_file(session, epub_path, config)) logger.info(f"ebook_ingest_paths_complete {count=} configured_paths={len(config.library_paths)}") return count async def resolve_ingest_path(path: Path | AsyncPath) -> AsyncPath: """Expand and resolve an ingest path without blocking the event loop.""" expanded_path = await AsyncPath(path).expanduser() return await expanded_path.resolve() async def ingest_file(session: AsyncSession, path: Path | AsyncPath, config: EbookSearchConfig) -> bool: """Ingest one EPUB file. Return True when the database changed.""" try: resolved_path = await resolve_ingest_path(path) logger.info(f"ebook_ingest_file_start {resolved_path=}") file_hash = await sha256_file(resolved_path) existing = await find_existing_source(session, resolved_path, file_hash) if existing is not None and existing.file_sha256 == file_hash: stat = await resolved_path.stat() existing.file_path = str(resolved_path) existing.file_mtime = datetime.fromtimestamp(stat.st_mtime, tz=UTC) existing.file_size = stat.st_size await session.flush() logger.info(f"ebook_ingest_file_unchanged {existing.id=} {resolved_path=}") return False if existing is not None: logger.info(f"ebook_ingest_file_replacing {existing.id=} {resolved_path=}") await session.delete(existing) await session.flush() stat = await resolved_path.stat() parsed = await asyncio.to_thread(parse_epub, Path(resolved_path)) source = EbookSource( title=parsed.title, author=parsed.author, language=parsed.language, publisher=parsed.publisher, identifier=parsed.identifier, file_path=str(resolved_path), file_sha256=file_hash, file_mtime=datetime.fromtimestamp(stat.st_mtime, tz=UTC), file_size=stat.st_size, ) session.add(source) await session.flush() chunk_index = 0 for spine_index, parsed_chapter in enumerate(parsed.chapters): chapter = EbookChapter( source_id=source.id, spine_index=spine_index, title=parsed_chapter.title, href=parsed_chapter.href, ) session.add(chapter) await session.flush() chunk_index = add_chapter_chunks(session, source, chapter, parsed_chapter, chunk_index, config) mention_count = await index_chunk_phrase_mentions_for_book(session, source.id, config) await session.commit() logger.info( f"ebook_ingest_file_complete {source.id=} {resolved_path=} chapters={len(parsed.chapters)} {chunk_index=} " f"{mention_count=}" ) except Exception: await session.rollback() logger.exception(f"ebook_ingest_file_error {path=}") return False else: return True async def find_existing_source(session: AsyncSession, path: Path | AsyncPath, file_hash: str) -> EbookSource | None: """Find an existing source by canonical path or file hash.""" return await session.scalar( select(EbookSource).where(or_(EbookSource.file_path == str(path), EbookSource.file_sha256 == file_hash)) ) def add_chapter_chunks( session: AsyncSession, source: EbookSource, chapter: EbookChapter, parsed_chapter: ParsedChapter, chunk_index: int, config: EbookSearchConfig, ) -> int: """Add chunk rows for one parsed chapter and return the next chunk index.""" page_label = parsed_chapter.page_labels[0] if parsed_chapter.page_labels else None for text_chunk in chunk_text( parsed_chapter.text, chunk_tokens=config.chunk_tokens, overlap_tokens=config.chunk_overlap, ): session.add( EbookChunk( source_id=source.id, chapter_id=chapter.id, chunk_index=chunk_index, text=text_chunk.text, token_start=text_chunk.token_start, token_count=text_chunk.token_count, page_label=page_label, content_sha256=hashlib.sha256(text_chunk.text.encode()).hexdigest(), search_text=f"{source.title} {source.author or ''} {chapter.title or ''} {text_chunk.text}", ) ) chunk_index += 1 return chunk_index async def sha256_file(path: AsyncPath) -> str: """Calculate the SHA-256 digest for a file without blocking the event loop.""" digest = hashlib.sha256() async with await path.open("rb") as file: while block := await file.read(1024 * 1024): digest.update(block) return digest.hexdigest()