Files
infra/stacks/books/app/library.py
T

483 lines
21 KiB
Python
Raw Blame History

This file contains invisible Unicode characters
This file contains invisible Unicode characters that are indistinguishable to humans but may be processed differently by a computer. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""Библиотека: разбор PDF/EPUB на главы, хранение в SQLite, гибридный поиск (FTS5 + эмбеддинги)."""
import datetime as dt
from html.parser import HTMLParser
import os
import re
import sqlite3
import threading
from pathlib import Path
import fitz # PyMuPDF: PDF и EPUB одним API
fitz.TOOLS.mupdf_display_errors(False) # шум про шрифты OCR-слоя в stderr — не ошибки разбора
import numpy as np
DATA = Path(os.environ.get("BOOKS_DATA", "/data"))
FILES = DATA / "files"
DB_PATH = DATA / "books.db"
EMB_MODEL = os.environ.get("EMB_MODEL", "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
CHUNK = 1200 # символов на фрагмент поиска
OVERLAP = 200
NO_TOC_PAGES = 20 # без оглавления — режем по 20 страниц
_local = threading.local()
_emb_lock = threading.Lock()
_emb = None
def db():
c = getattr(_local, "c", None)
if c is None:
DATA.mkdir(parents=True, exist_ok=True)
c = sqlite3.connect(DB_PATH, timeout=60, check_same_thread=False)
c.row_factory = sqlite3.Row
c.execute("PRAGMA journal_mode=WAL")
c.execute("PRAGMA foreign_keys=ON")
c.executescript(SCHEMA)
_local.c = c
return c
SCHEMA = """
CREATE TABLE IF NOT EXISTS book(id INTEGER PRIMARY KEY, title TEXT, author TEXT, lang TEXT, file TEXT, orig TEXT, fmt TEXT,
pages INTEGER, added TEXT, status TEXT, note TEXT, last_ch INTEGER);
CREATE TABLE IF NOT EXISTS chapter(id INTEGER PRIMARY KEY, book_id INTEGER REFERENCES book(id) ON DELETE CASCADE,
idx INTEGER, title TEXT, level INTEGER, p1 INTEGER, p2 INTEGER, text TEXT, chars INTEGER,
read INTEGER DEFAULT 0, read_at TEXT, UNIQUE(book_id, idx));
CREATE TABLE IF NOT EXISTS recap(chapter_id INTEGER PRIMARY KEY REFERENCES chapter(id) ON DELETE CASCADE,
md TEXT, ai_md TEXT, updated TEXT);
CREATE TABLE IF NOT EXISTS chunk(id INTEGER PRIMARY KEY, book_id INTEGER REFERENCES book(id) ON DELETE CASCADE,
chapter_id INTEGER REFERENCES chapter(id) ON DELETE CASCADE, kind TEXT, page INTEGER, text TEXT, emb BLOB);
CREATE INDEX IF NOT EXISTS chunk_book ON chunk(book_id, kind);
CREATE VIRTUAL TABLE IF NOT EXISTS chunk_fts USING fts5(text, content='chunk', content_rowid='id',
tokenize='unicode61 remove_diacritics 2');
CREATE TRIGGER IF NOT EXISTS chunk_ai AFTER INSERT ON chunk BEGIN
INSERT INTO chunk_fts(rowid, text) VALUES (new.id, new.text); END;
CREATE TRIGGER IF NOT EXISTS chunk_ad AFTER DELETE ON chunk BEGIN
INSERT INTO chunk_fts(chunk_fts, rowid, text) VALUES ('delete', old.id, old.text); END;
CREATE TABLE IF NOT EXISTS qa(id INTEGER PRIMARY KEY, book_id INTEGER REFERENCES book(id) ON DELETE CASCADE,
chapter_idx INTEGER, kind TEXT, q TEXT, a TEXT, ts TEXT);
"""
def now():
return dt.datetime.now().isoformat(timespec="seconds")
# ---------------------------------------------------------------- эмбеддинги
def embedder():
global _emb
with _emb_lock:
if _emb is None:
from fastembed import TextEmbedding
_emb = TextEmbedding(EMB_MODEL, cache_dir=str(DATA / "models"), threads=2)
return _emb
def embed(texts):
if not texts:
return np.zeros((0, 384), dtype=np.float32)
v = np.array(list(embedder().embed(texts, batch_size=32)), dtype=np.float32)
v /= np.linalg.norm(v, axis=1, keepdims=True) + 1e-9
return v
# ---------------------------------------------------------------- разбор книги
def detect_lang(text):
cyr = len(re.findall(r"[а-яёА-ЯЁ]", text[:20000]))
lat = len(re.findall(r"[a-zA-Z]", text[:20000]))
return "ru" if cyr > lat else "en"
def clean(t):
t = t.replace("­", "") # мягкие переносы
t = re.sub(r"(\w)-\n(\w)", r"\1\2", t) # перенос слова по слогам в PDF
t = re.sub(r"[ \t]+\n", "\n", t)
t = re.sub(r"\n{3,}", "\n\n", t)
return t.strip()
CHAPTER_RE = re.compile(r"^\s*(глава|chapter|часть|part|раздел)\s+([0-9IVXLC]+|[а-яё]+)\b", re.I)
MAX_AVG_PAGES = 45 # «главы» длиннее в среднем — это части книги, берём уровень глубже
def _num(tok):
tok = tok.upper()
if tok.isdigit():
return int(tok)
roman = {"I": 1, "V": 5, "X": 10, "L": 50, "C": 100}
if tok and all(ch in roman for ch in tok):
total = 0
for a, b in zip(tok, tok[1:] + " "):
total += -roman[a] if b in roman and roman[b] > roman[a] else roman[a]
return total
return None
def headings_from_pages(texts):
"""Без оглавления: страницы, которые начинаются с «Глава N» / «Chapter N».
Колонтитулы повторяют «Глава N» на каждой странице, а страницы оглавления перечисляют все главы подряд, поэтому:
берём первое появление каждого номера и самую длинную возрастающую по номеру цепочку (по порядку страниц)."""
found, seen = [], set()
for i, t in enumerate(texts):
head = [l.strip() for l in t.splitlines()[:4] if l.strip()]
for j, l in enumerate(head[:3]):
m = CHAPTER_RE.match(l)
if m and len(l) < 90:
kind, n = m.group(1).lower(), _num(m.group(2))
if kind not in ("глава", "chapter") or n is None or n in seen:
break
seen.add(n)
title = l if len(l) > 12 or j + 1 >= len(head) else f"{l}. {head[j + 1][:70]}"
title = re.sub(r"\.{2,}", ".", title)
found.append((n, i + 1, title))
break
# самая длинная возрастающая подпоследовательность по номеру главы (O(n²), глав немного)
best, prev = [1] * len(found), [-1] * len(found)
for a in range(len(found)):
for b in range(a):
if found[b][0] < found[a][0] and best[b] + 1 > best[a]:
best[a], prev[a] = best[b] + 1, b
if not found:
return []
k = max(range(len(found)), key=best.__getitem__)
chain = []
while k >= 0:
chain.append(found[k])
k = prev[k]
chain.reverse()
return [[1, title, page] for _, page, title in chain] if len(chain) >= 3 else []
def plan_chapters(toc, pages, texts=None):
"""Из оглавления — главы (верхний уровень, если он не из огромных «частей»). [(title, level, p1, p2)]."""
toc = [(lvl, (title or "").strip(), max(1, min(pages, p))) for lvl, title, p, *_ in toc if p and p > 0]
if not toc and texts:
toc = [tuple(x) for x in headings_from_pages(texts)]
if toc:
lvl1 = [t for t in toc if t[0] == 1]
lvl2 = [t for t in toc if t[0] <= 2]
use = lvl1 if len(lvl1) >= 3 and (pages / len(lvl1) <= MAX_AVG_PAGES or len(lvl2) <= len(lvl1)) else lvl2
use = sorted({(p, i): (lvl, title, p) for i, (lvl, title, p) in enumerate(use)}.items())
use = [v for _, v in use]
out = []
if use and use[0][2] > 1:
out.append(("Начало книги", 1, 1, use[0][2] - 1))
for i, (lvl, title, p) in enumerate(use):
p2 = use[i + 1][2] - 1 if i + 1 < len(use) else pages
out.append((title or f"Часть {i + 1}", lvl, p, max(p, p2)))
# слишком мелкие подряд (одна страница) оставляем как есть — это видно в оглавлении
return out
return [(f"Страницы {p}–{min(pages, p + NO_TOC_PAGES - 1)}", 1, p, min(pages, p + NO_TOC_PAGES - 1))
for p in range(1, pages + 1, NO_TOC_PAGES)]
# ---------------------------------------------------------------- EPUB: свой разбор (у PyMuPDF нет оглавления EPUB 3)
EPUB_PAGE = 2000 # у EPUB нет страниц — «страница» ≈ 2000 символов, для ссылок [с. N]
class _Text(HTMLParser):
BLOCK = {"p", "div", "br", "li", "h1", "h2", "h3", "h4", "h5", "h6", "tr", "pre", "blockquote", "section", "dt", "dd"}
def __init__(self):
super().__init__(convert_charrefs=True)
self.out, self.skip = [], 0
def handle_starttag(self, tag, attrs):
if tag in ("script", "style", "head"):
self.skip += 1
elif tag in self.BLOCK:
self.out.append("\n\n")
def handle_endtag(self, tag):
if tag in ("script", "style", "head"):
self.skip = max(0, self.skip - 1)
elif tag in self.BLOCK:
self.out.append("\n\n")
def handle_data(self, d):
if not self.skip:
self.out.append(d)
def html_text(raw):
t = _Text()
t.feed(raw)
return clean(re.sub(r"[ \t\r\f\v]*\n[ \t\r\f\v]*", "\n", re.sub(r"[ \t]+", " ", "".join(t.out))))
def parse_epub(path):
"""→ (title, author, [(глава, уровень, [тексты «страниц»])]). Главы — по оглавлению (nav/NCX), текст — по spine."""
import posixpath
import urllib.parse
import xml.etree.ElementTree as ET
import zipfile
z = zipfile.ZipFile(path)
tag = lambda e: e.tag.rsplit("}", 1)[-1]
cont = ET.fromstring(z.read("META-INF/container.xml"))
opf_path = next(e.get("full-path") for e in cont.iter() if tag(e) == "rootfile")
base = posixpath.dirname(opf_path)
norm = lambda href, rel=base: posixpath.normpath(posixpath.join(rel, urllib.parse.unquote(href.split("#")[0])))
opf = ET.fromstring(z.read(opf_path))
meta = {tag(e): (e.text or "").strip() for e in opf.iter() if tag(e) in ("title", "creator") and (e.text or "").strip()}
manifest, nav_href, ncx_href = {}, None, None
for e in opf.iter():
if tag(e) == "item":
manifest[e.get("id")] = e.get("href")
if "nav" in (e.get("properties") or "").split():
nav_href = e.get("href")
if e.get("media-type") == "application/x-dtbncx+xml":
ncx_href = e.get("href")
spine = [norm(manifest[e.get("idref")]) for e in opf.iter() if tag(e) == "itemref" and e.get("idref") in manifest]
titles = {} # файл spine → (название, уровень) первой записи оглавления
def add(href, title, lvl, rel):
f = norm(href, rel)
if title and f not in titles:
titles[f] = (title.strip(), lvl)
try:
if nav_href:
rel = posixpath.dirname(norm(nav_href))
nav = ET.fromstring(z.read(norm(nav_href)))
def walk(ol, lvl):
for li in ol:
if tag(li) != "li":
continue
a = next((x for x in li.iter() if tag(x) == "a"), None)
if a is not None and a.get("href"):
add(a.get("href"), "".join(a.itertext()), lvl, rel)
for sub in li:
if tag(sub) == "ol":
walk(sub, lvl + 1)
for n in nav.iter():
if tag(n) == "nav" and "toc" in (n.get("{http://www.idpf.org/2007/ops}type") or "toc"):
for ol in n:
if tag(ol) == "ol":
walk(ol, 1)
break
if not titles and ncx_href:
rel = posixpath.dirname(norm(ncx_href))
ncx = ET.fromstring(z.read(norm(ncx_href)))
def walkp(node, lvl):
for np_ in node:
if tag(np_) != "navPoint":
continue
label = next(("".join(x.itertext()) for x in np_.iter() if tag(x) == "text"), "")
src = next((x.get("src") for x in np_ if tag(x) == "content"), None)
if src:
add(src, label, lvl, rel)
walkp(np_, lvl + 1)
for n in ncx.iter():
if tag(n) == "navMap":
walkp(n, 1)
except (KeyError, ET.ParseError):
titles = {}
# верхний уровень, если его достаточно — иначе первые два
lvl1 = [f for f, (t, l) in titles.items() if l == 1]
keep_lvl = 1 if len(lvl1) >= 3 else 2
chapters = []
for f in spine:
try:
text = html_text(z.read(f).decode("utf-8", "replace"))
except KeyError:
continue
t = titles.get(f)
if not chapters or (t and t[1] <= keep_lvl) or (not titles and text):
first = next((l for l in text.splitlines() if l.strip()), "")[:80]
chapters.append([t[0] if t else (first or f"Часть {len(chapters) + 1}"), t[1] if t else 1, []])
if text:
chapters[-1][2].append(text)
out = []
for title, lvl, texts in chapters:
body = "\n\n".join(texts)
if not body.strip():
continue
pages, buf = [], ""
for para in re.split(r"\n\s*\n", body):
if buf and len(buf) + len(para) > EPUB_PAGE:
pages.append(buf)
buf = ""
buf = (buf + "\n\n" + para) if buf else para
if buf:
pages.append(buf)
out.append((title, lvl, pages))
return meta.get("title", ""), meta.get("creator", ""), out
def chunks_of(page_texts, p1):
"""Фрагменты для поиска с номером страницы начала."""
out = []
buf, start = "", p1
for i, t in enumerate(page_texts):
page = p1 + i
for para in re.split(r"\n\s*\n", t):
para = para.strip()
if not para:
continue
if not buf:
start = page
if len(buf) + len(para) > CHUNK and buf:
out.append((start, buf))
buf = buf[-OVERLAP:] + "\n" + para
start = page
else:
buf = (buf + "\n\n" + para) if buf else para
if buf.strip():
out.append((start, buf))
return out
def import_book(book_id):
"""Фоновая обработка загруженного файла: главы, текст, фрагменты, эмбеддинги."""
c = db()
b = c.execute("SELECT * FROM book WHERE id=?", (book_id,)).fetchone()
try:
if b["fmt"] == "epub":
et, ea, eps = parse_epub(FILES / b["file"])
texts, chapters = [], []
for ctitle, lvl, ptexts in eps:
chapters.append((ctitle, lvl, len(texts) + 1, len(texts) + len(ptexts)))
texts += ptexts
pages = len(texts)
meta = {"title": et, "author": ea}
else:
doc = fitz.open(FILES / b["file"])
if doc.needs_pass:
raise ValueError("файл защищён паролем")
pages = doc.page_count
texts = [clean(doc[i].get_text("text")) for i in range(pages)]
meta = doc.metadata or {}
chapters = plan_chapters(doc.get_toc(simple=True), pages, texts)
total = sum(len(t) for t in texts)
mt = (meta.get("title") or "").strip()
title = b["title"] or (mt if len(mt) > 2 and not mt.lower().startswith(("untitled", "microsoft")) else "") \
or Path(b["orig"] or b["file"]).stem
author = (meta.get("author") or "").strip()
note = ""
if total < pages * 80:
note = "Похоже на скан: текста почти нет (распознавание пока не поддерживается)."
if not pages or not chapters:
raise ValueError("в файле не нашлось текста")
c.execute("DELETE FROM chapter WHERE book_id=?", (book_id,))
all_chunks = []
for idx, (ctitle, lvl, p1, p2) in enumerate(chapters, 1):
ptexts = texts[p1 - 1:p2]
body = "\n\n".join(f"[[p{p1 + i}]]\n{t}" for i, t in enumerate(ptexts))
cur = c.execute("INSERT INTO chapter(book_id, idx, title, level, p1, p2, text, chars) VALUES (?,?,?,?,?,?,?,?)",
(book_id, idx, ctitle, lvl, p1, p2, body, sum(len(t) for t in ptexts)))
for page, text in chunks_of(ptexts, p1):
all_chunks.append((cur.lastrowid, page, text))
c.execute("UPDATE book SET title=?, author=?, lang=?, pages=?, note=?, status='indexing' WHERE id=?",
(title, author, detect_lang("".join(texts[:50])), pages, note, book_id))
c.commit()
for i in range(0, len(all_chunks), 256):
part = all_chunks[i:i + 256]
vecs = embed([t for _, _, t in part])
c.executemany("INSERT INTO chunk(book_id, chapter_id, kind, page, text, emb) VALUES (?,?,?,?,?,?)",
[(book_id, ch, "book", pg, t, v.tobytes()) for (ch, pg, t), v in zip(part, vecs)])
c.commit()
c.execute("UPDATE book SET status='ready' WHERE id=?", (book_id,))
c.commit()
except Exception as e: # отдаём ошибку в интерфейс, не роняем приложение
c.rollback()
c.execute("UPDATE book SET status='error', note=? WHERE id=?", (f"Не удалось разобрать: {e}", book_id))
c.commit()
def add_book(filename, data):
FILES.mkdir(parents=True, exist_ok=True)
ext = Path(filename).suffix.lower()
if ext not in (".pdf", ".epub"):
raise ValueError("Нужен PDF или EPUB")
c = db()
bid = c.execute("INSERT INTO book(orig, fmt, added, status) VALUES (?,?,?, 'parsing')",
(Path(filename).name, ext[1:], now())).lastrowid
name = f"{bid}{ext}"
(FILES / name).write_bytes(data)
c.execute("UPDATE book SET file=? WHERE id=?", (name, bid))
c.commit()
threading.Thread(target=import_book, args=(bid,), daemon=True).start()
return bid
def delete_book(bid):
c = db()
b = c.execute("SELECT file FROM book WHERE id=?", (bid,)).fetchone()
if b and b["file"]:
(FILES / b["file"]).unlink(missing_ok=True)
c.execute("DELETE FROM book WHERE id=?", (bid,))
c.commit()
# ---------------------------------------------------------------- рекапы в индекс
def index_recap(chapter_id):
c = db()
ch = c.execute("SELECT * FROM chapter WHERE id=?", (chapter_id,)).fetchone()
r = c.execute("SELECT md FROM recap WHERE chapter_id=?", (chapter_id,)).fetchone()
c.execute("DELETE FROM chunk WHERE chapter_id=? AND kind='recap'", (chapter_id,))
if r and (r["md"] or "").strip():
text = f"Рекап главы «{ch['title']}»:\n{r['md']}"
v = embed([text])[0]
c.execute("INSERT INTO chunk(book_id, chapter_id, kind, page, text, emb) VALUES (?,?,?,?,?,?)",
(ch["book_id"], chapter_id, "recap", ch["p1"], text, v.tobytes()))
c.commit()
# ---------------------------------------------------------------- поиск
def fts_query(q):
words = [w for w in re.findall(r"\w+", q.lower()) if len(w) >= 2]
out = []
for w in words:
stem = re.sub(r"(ами|ями|ов|ев|ей|ий|ый|ой|ая|яя|ое|ее|ые|ие|ам|ям|ах|ях|ом|ем|а|я|ы|и|у|ю|е|о)$", "", w) \
if re.fullmatch(r"[а-яё]{5,}", w) else w
out.append(f'"{stem}"*')
return " OR ".join(out)
def search(q, book_id=None, k=8):
"""Гибрид: BM25 (FTS5) + косинус эмбеддингов, слияние по рангам (RRF)."""
c = db()
where, args = ("AND c.book_id=?", [book_id]) if book_id else ("", [])
ranks = {}
fq = fts_query(q)
if fq:
rows = c.execute(f"SELECT c.id FROM chunk_fts f JOIN chunk c ON c.id=f.rowid WHERE chunk_fts MATCH ? {where} "
f"ORDER BY bm25(chunk_fts) LIMIT 40", [fq] + args).fetchall()
for r, row in enumerate(rows):
ranks[row["id"]] = ranks.get(row["id"], 0) + 1 / (60 + r)
rows = c.execute(f"SELECT c.id, c.emb FROM chunk c WHERE c.emb IS NOT NULL {where}", args).fetchall()
if rows:
qv = embed([q])[0]
m = np.frombuffer(b"".join(r["emb"] for r in rows), dtype=np.float32).reshape(len(rows), -1)
sims = m @ qv
for r, i in enumerate(np.argsort(-sims)[:40]):
ranks[rows[i]["id"]] = ranks.get(rows[i]["id"], 0) + 1 / (60 + r)
top = sorted(ranks, key=ranks.get, reverse=True)[:k]
if not top:
return []
res = c.execute(f"SELECT c.*, ch.idx AS ch_idx, ch.title AS ch_title, b.title AS book_title FROM chunk c "
f"JOIN chapter ch ON ch.id=c.chapter_id JOIN book b ON b.id=c.book_id "
f"WHERE c.id IN ({','.join('?' * len(top))})", top).fetchall()
order = {cid: i for i, cid in enumerate(top)}
return sorted(res, key=lambda r: order[r["id"]])
def find_selection(book_id, chapter_idx, text):
"""Окружение выделенного фрагмента в главе — для точного перевода/объяснения."""
c = db()
ch = c.execute("SELECT text FROM chapter WHERE book_id=? AND idx=?", (book_id, chapter_idx)).fetchone()
if not ch:
return ""
body = re.sub(r"\[\[p\d+\]\]\n", "", ch["text"])
probe = re.sub(r"\s+", " ", text.strip())[:60]
flat = re.sub(r"\s+", " ", body)
i = flat.find(probe)
if i < 0:
return ""
return flat[max(0, i - 1500): i + len(text) + 1500]