467 lines
18 KiB
Python
467 lines
18 KiB
Python
from __future__ import annotations
|
||
|
||
import sys
|
||
from device_auth import DeviceStore
|
||
import html
|
||
import json
|
||
import os
|
||
import re
|
||
import shutil
|
||
import sqlite3
|
||
import subprocess
|
||
import threading
|
||
import time
|
||
import uuid
|
||
import zipfile
|
||
from contextlib import contextmanager
|
||
from pathlib import Path, PurePosixPath
|
||
from typing import TYPE_CHECKING, Annotated, Any, Iterator
|
||
from xml.etree import ElementTree
|
||
|
||
from fastapi import Depends, FastAPI, File, Form, Header, HTTPException, UploadFile
|
||
from fastapi.responses import FileResponse
|
||
from num2words import num2words
|
||
|
||
ROOT = Path(os.environ.get("AUDIOBOOK_SERVICE_ROOT", "/home/sevenhill/apps/aletheia-audiobook"))
|
||
SERVICE_ROOT = ROOT / "service"
|
||
DATA_ROOT = ROOT / "audiobook-jobs"
|
||
DATABASE = DATA_ROOT / "jobs.db3"
|
||
MAX_UPLOAD_BYTES = 200 * 1024 * 1024
|
||
SERVICE_ROOT.mkdir(parents=True, exist_ok=True)
|
||
DATA_ROOT.mkdir(parents=True, exist_ok=True)
|
||
|
||
app = FastAPI(title="Aletheia CMP Audiobook Service", version="2.0.0")
|
||
queue_condition = threading.Condition()
|
||
queued_jobs: list[str] = []
|
||
worker_state = {"activeJobId": None, "modelLoaded": False}
|
||
|
||
|
||
@contextmanager
|
||
def database() -> Iterator[sqlite3.Connection]:
|
||
connection = sqlite3.connect(DATABASE, timeout=30)
|
||
connection.row_factory = sqlite3.Row
|
||
try:
|
||
yield connection
|
||
connection.commit()
|
||
finally:
|
||
connection.close()
|
||
|
||
|
||
def initialize_database() -> None:
|
||
with database() as connection:
|
||
connection.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS jobs (
|
||
id TEXT PRIMARY KEY,
|
||
title TEXT NOT NULL,
|
||
author TEXT NOT NULL,
|
||
source_name TEXT NOT NULL,
|
||
input_path TEXT NOT NULL,
|
||
output_path TEXT NOT NULL,
|
||
status TEXT NOT NULL,
|
||
stage TEXT NOT NULL,
|
||
processed_characters INTEGER NOT NULL DEFAULT 0,
|
||
total_characters INTEGER NOT NULL DEFAULT 0,
|
||
current_chapter TEXT,
|
||
duration_ms INTEGER NOT NULL DEFAULT 0,
|
||
chapters_json TEXT NOT NULL DEFAULT '[]',
|
||
error_message TEXT,
|
||
cancel_requested INTEGER NOT NULL DEFAULT 0,
|
||
created_at INTEGER NOT NULL,
|
||
updated_at INTEGER NOT NULL
|
||
)
|
||
"""
|
||
)
|
||
if "owner_id" not in {row[1] for row in connection.execute("PRAGMA table_info(jobs)")}:
|
||
connection.execute("ALTER TABLE jobs ADD COLUMN owner_id TEXT")
|
||
connection.execute("CREATE INDEX IF NOT EXISTS idx_jobs_owner ON jobs(owner_id)")
|
||
connection.execute(
|
||
"UPDATE jobs SET status='queued', stage='queued' "
|
||
"WHERE status IN ('extracting','synthesizing','encoding')"
|
||
)
|
||
rows = connection.execute(
|
||
"SELECT id FROM jobs WHERE status='queued' ORDER BY created_at"
|
||
).fetchall()
|
||
queued_jobs.extend(row["id"] for row in rows)
|
||
|
||
|
||
devices = DeviceStore(DATABASE)
|
||
|
||
|
||
def require_token(authorization: Annotated[str | None, Header()] = None) -> str:
|
||
try:
|
||
if not authorization or not authorization.startswith("Bearer "):
|
||
raise PermissionError()
|
||
return devices.authenticate(authorization[7:])
|
||
except PermissionError:
|
||
raise HTTPException(status_code=401, detail="Invalid device credential")
|
||
|
||
|
||
@app.post("/v1/devices/enroll")
|
||
def enroll_device(payload: dict) -> dict:
|
||
try:
|
||
return devices.enroll(payload.get("code", ""), payload.get("label", "Android"))
|
||
except PermissionError:
|
||
raise HTTPException(status_code=403, detail="Invalid or expired pairing code")
|
||
|
||
|
||
def owned_job(job_id: str, owner: str) -> sqlite3.Row:
|
||
row = get_job(job_id)
|
||
if row["owner_id"] != owner:
|
||
raise HTTPException(status_code=404, detail="Audiobook job not found")
|
||
return row
|
||
|
||
|
||
def now_ms() -> int:
|
||
return int(time.time() * 1000)
|
||
|
||
|
||
def get_job(job_id: str) -> sqlite3.Row:
|
||
with database() as connection:
|
||
row = connection.execute("SELECT * FROM jobs WHERE id=?", (job_id,)).fetchone()
|
||
if row is None:
|
||
raise HTTPException(status_code=404, detail="Audiobook job not found")
|
||
return row
|
||
|
||
|
||
def job_payload(row: sqlite3.Row) -> dict:
|
||
total = row["total_characters"]
|
||
processed = row["processed_characters"]
|
||
progress = min(100, int(processed * 100 / total)) if total else 0
|
||
if row["status"] == "ready":
|
||
progress = 100
|
||
return {
|
||
"id": row["id"],
|
||
"status": row["status"],
|
||
"stage": row["stage"],
|
||
"progress": progress,
|
||
"processedCharacters": processed,
|
||
"totalCharacters": total,
|
||
"currentChapter": row["current_chapter"],
|
||
"durationMs": row["duration_ms"],
|
||
"chapters": json.loads(row["chapters_json"] or "[]"),
|
||
"error": row["error_message"],
|
||
"createdAt": row["created_at"],
|
||
"updatedAt": row["updated_at"],
|
||
}
|
||
|
||
|
||
@app.get("/health")
|
||
def health() -> dict:
|
||
from voicestudio_client import check_available
|
||
try:
|
||
check_available()
|
||
except Exception:
|
||
raise HTTPException(status_code=503, detail="VoiceStudio unavailable")
|
||
return {
|
||
"status": "ok",
|
||
"workerMode": "voicestudio",
|
||
"modelLoaded": worker_state["modelLoaded"],
|
||
}
|
||
|
||
|
||
@app.post("/v1/audiobooks")
|
||
async def create_audiobook(
|
||
book: Annotated[UploadFile, File()],
|
||
owner: Annotated[str, Depends(require_token)],
|
||
title: Annotated[str, Form()],
|
||
author: Annotated[str, Form()] = "",
|
||
) -> dict:
|
||
extension = Path(book.filename or "book.epub").suffix.lower()
|
||
if extension not in {".epub", ".fb2"}:
|
||
raise HTTPException(status_code=400, detail="Only EPUB and FB2 are supported")
|
||
job_id = uuid.uuid4().hex
|
||
job_root = DATA_ROOT / job_id
|
||
job_root.mkdir(parents=True)
|
||
input_path = job_root / f"source{extension}"
|
||
received = 0
|
||
try:
|
||
with input_path.open("wb") as output:
|
||
while chunk := await book.read(1024 * 1024):
|
||
received += len(chunk)
|
||
if received > MAX_UPLOAD_BYTES:
|
||
raise HTTPException(status_code=413, detail="Book file is too large")
|
||
output.write(chunk)
|
||
except Exception:
|
||
shutil.rmtree(job_root, ignore_errors=True)
|
||
raise
|
||
timestamp = now_ms()
|
||
with database() as connection:
|
||
connection.execute(
|
||
"""
|
||
INSERT INTO jobs (
|
||
id,title,author,source_name,input_path,output_path,status,stage,created_at,updated_at,owner_id
|
||
) VALUES (?,?,?,?,?,?,?,?,?,?,?)
|
||
""",
|
||
(
|
||
job_id,
|
||
title.strip() or Path(book.filename or "Книга").stem,
|
||
author.strip(),
|
||
book.filename or input_path.name,
|
||
str(input_path),
|
||
str(job_root / "audiobook.m4a"),
|
||
"queued",
|
||
"queued",
|
||
timestamp,
|
||
timestamp,
|
||
owner,
|
||
),
|
||
)
|
||
with queue_condition:
|
||
queued_jobs.append(job_id)
|
||
queue_condition.notify()
|
||
return job_payload(get_job(job_id))
|
||
|
||
|
||
@app.get("/v1/audiobooks/{job_id}")
|
||
def audiobook_status(job_id: str, owner: Annotated[str, Depends(require_token)]) -> dict:
|
||
return job_payload(owned_job(job_id, owner))
|
||
|
||
|
||
@app.get("/v1/audiobooks/{job_id}/file")
|
||
def audiobook_file(job_id: str, owner: Annotated[str, Depends(require_token)]) -> FileResponse:
|
||
row = owned_job(job_id, owner)
|
||
output = Path(row["output_path"])
|
||
if row["status"] != "ready" or not output.is_file():
|
||
raise HTTPException(status_code=409, detail="Audiobook is not ready")
|
||
safe_title = re.sub(r"[^0-9A-Za-zА-Яа-яЁё._ -]+", "_", row["title"]).strip() or "audiobook"
|
||
return FileResponse(output, media_type="audio/mp4", filename=f"{safe_title}.m4a")
|
||
|
||
|
||
@app.delete("/v1/audiobooks/{job_id}")
|
||
def cancel_or_delete(job_id: str, owner: Annotated[str, Depends(require_token)]) -> dict:
|
||
row = owned_job(job_id, owner)
|
||
if row["status"] == "queued":
|
||
with queue_condition:
|
||
queued_jobs[:] = [queued_id for queued_id in queued_jobs if queued_id != job_id]
|
||
delete_job(job_id, Path(row["input_path"]).parent)
|
||
return {"status": "deleted"}
|
||
if row["status"] in {"extracting", "synthesizing", "encoding"}:
|
||
with database() as connection:
|
||
connection.execute(
|
||
"UPDATE jobs SET cancel_requested=1, stage='cancelling', updated_at=? WHERE id=?",
|
||
(now_ms(), job_id),
|
||
)
|
||
return {"status": "cancelling"}
|
||
delete_job(job_id, Path(row["input_path"]).parent)
|
||
return {"status": "deleted"}
|
||
|
||
|
||
def delete_job(job_id: str, job_root: Path) -> None:
|
||
shutil.rmtree(job_root, ignore_errors=True)
|
||
with database() as connection:
|
||
connection.execute("DELETE FROM jobs WHERE id=?", (job_id,))
|
||
|
||
|
||
def finish_cancellation(job_id: str, job_root: Path) -> None:
|
||
update_job(job_id, status="cancelled", stage="cancelled")
|
||
delete_job(job_id, job_root)
|
||
|
||
|
||
def local_name(tag: str) -> str:
|
||
return tag.rsplit("}", 1)[-1].lower()
|
||
|
||
|
||
def clean_text(value: str) -> str:
|
||
value = html.unescape(re.sub(r"<[^>]+>", " ", value))
|
||
value = value.replace("\u00a0", " ")
|
||
return re.sub(r"\s+", " ", value).strip()
|
||
|
||
|
||
def extract_epub(path: Path) -> list[tuple[str, str]]:
|
||
sections: list[tuple[str, str]] = []
|
||
with zipfile.ZipFile(path) as archive:
|
||
container = ElementTree.fromstring(archive.read("META-INF/container.xml"))
|
||
rootfile = next(node for node in container.iter() if local_name(node.tag) == "rootfile")
|
||
opf_path = rootfile.attrib["full-path"]
|
||
opf = ElementTree.fromstring(archive.read(opf_path))
|
||
manifest = {
|
||
node.attrib["id"]: node.attrib["href"]
|
||
for node in opf.iter()
|
||
if local_name(node.tag) == "item" and "id" in node.attrib and "href" in node.attrib
|
||
}
|
||
spine = [
|
||
node.attrib["idref"]
|
||
for node in opf.iter()
|
||
if local_name(node.tag) == "itemref" and "idref" in node.attrib
|
||
]
|
||
base = PurePosixPath(opf_path).parent
|
||
for index, item_id in enumerate(spine, start=1):
|
||
href = manifest.get(item_id)
|
||
if not href:
|
||
continue
|
||
entry = str(base / href.split("#", 1)[0])
|
||
raw = archive.read(entry).decode("utf-8", errors="replace")
|
||
heading_match = re.search(r"<h[1-3]\b[^>]*>(.*?)</h[1-3]>", raw, re.I | re.S)
|
||
title = clean_text(heading_match.group(1)) if heading_match else f"Раздел {index}"
|
||
paragraphs = [clean_text(match) for match in re.findall(r"<p\b[^>]*>(.*?)</p>", raw, re.I | re.S)]
|
||
text = "\n".join(paragraph for paragraph in paragraphs if paragraph)
|
||
if text:
|
||
sections.append((title, text))
|
||
return sections
|
||
|
||
|
||
def element_text_without_nested_sections(element: ElementTree.Element) -> str:
|
||
parts: list[str] = []
|
||
if element.text:
|
||
parts.append(element.text)
|
||
for child in element:
|
||
if local_name(child.tag) != "section":
|
||
parts.append(" ".join(child.itertext()))
|
||
if child.tail:
|
||
parts.append(child.tail)
|
||
return clean_text(" ".join(parts))
|
||
|
||
|
||
def extract_fb2(path: Path) -> list[tuple[str, str]]:
|
||
root = ElementTree.parse(path).getroot()
|
||
sections: list[tuple[str, str]] = []
|
||
for section in root.iter():
|
||
if local_name(section.tag) != "section":
|
||
continue
|
||
title_element = next((child for child in section if local_name(child.tag) == "title"), None)
|
||
title = clean_text(" ".join(title_element.itertext())) if title_element is not None else ""
|
||
text = element_text_without_nested_sections(section)
|
||
if text:
|
||
sections.append((title or f"Раздел {len(sections) + 1}", text))
|
||
return sections
|
||
|
||
|
||
def normalize_numbers(text: str) -> str:
|
||
def replace(match: re.Match[str]) -> str:
|
||
value = match.group(0)
|
||
try:
|
||
return num2words(int(value), lang="ru")
|
||
except (ValueError, OverflowError):
|
||
return value
|
||
|
||
return re.sub(r"(?<![\w+])\d{1,12}(?!\w)", replace, text)
|
||
|
||
|
||
def split_chunks(sections: list[tuple[str, str]], limit: int = 520) -> list[dict]:
|
||
chunks: list[dict] = []
|
||
for chapter, text in sections:
|
||
sentences = re.split(r"(?<=[.!?…])\s+|\n+", text)
|
||
current = ""
|
||
for sentence in sentences:
|
||
sentence = normalize_numbers(sentence.strip())
|
||
if not sentence or not re.search(r"[А-Яа-яЁё]", sentence):
|
||
continue
|
||
if len(sentence) > limit:
|
||
pieces = re.split(r"(?<=[,;:])\s+", sentence)
|
||
else:
|
||
pieces = [sentence]
|
||
for piece in pieces:
|
||
if current and len(current) + 1 + len(piece) > limit:
|
||
chunks.append({"chapter": chapter, "text": current})
|
||
current = piece
|
||
else:
|
||
current = f"{current} {piece}".strip()
|
||
if current:
|
||
chunks.append({"chapter": chapter, "text": current})
|
||
return chunks
|
||
|
||
|
||
def update_job(job_id: str, **values: object) -> None:
|
||
values["updated_at"] = now_ms()
|
||
assignments = ",".join(f"{key}=?" for key in values)
|
||
with database() as connection:
|
||
connection.execute(
|
||
f"UPDATE jobs SET {assignments} WHERE id=?",
|
||
(*values.values(), job_id),
|
||
)
|
||
|
||
|
||
def cancellation_requested(job_id: str) -> bool:
|
||
with database() as connection:
|
||
row = connection.execute("SELECT cancel_requested FROM jobs WHERE id=?", (job_id,)).fetchone()
|
||
return row is None or bool(row["cancel_requested"])
|
||
|
||
|
||
def run_job(job_id: str) -> None:
|
||
from voicestudio_client import render_audiobook, split_sections
|
||
row = get_job(job_id)
|
||
job_root = Path(row["input_path"]).parent
|
||
try:
|
||
if cancellation_requested(job_id):
|
||
finish_cancellation(job_id, job_root)
|
||
return
|
||
update_job(job_id, status="extracting", stage="extracting", error_message=None)
|
||
source = Path(row["input_path"])
|
||
sections = extract_epub(source) if source.suffix.lower() == ".epub" else extract_fb2(source)
|
||
sections = split_sections(sections)
|
||
if not sections:
|
||
raise RuntimeError("В книге не найден текст для озвучивания")
|
||
total = sum(len(text) for _, text in sections)
|
||
update_job(job_id, status="synthesizing", stage="synthesizing", total_characters=total,
|
||
processed_characters=0, duration_ms=0, chapters_json="[]",
|
||
current_chapter="Проверка сохранённых фрагментов и озвучивание недостающих…")
|
||
completed = {}
|
||
verified = set()
|
||
|
||
def on_event(event):
|
||
if event.get("type") == "chapter":
|
||
index = int(event["index"])
|
||
completed[index] = event
|
||
duration = 0
|
||
markers = []
|
||
for _, chapter in sorted(completed.items()):
|
||
markers.append({"title": chapter.get("title", "Раздел"), "startMs": duration})
|
||
duration += round(float(chapter.get("duration_s", 0)) * 1000)
|
||
processed = sum(len(sections[i][1]) for i in completed if 0 <= i < len(sections))
|
||
label = event.get("title")
|
||
if event.get("stage") != "gpu_render":
|
||
verified.add(index)
|
||
label = f"Проверка готовых фрагментов: {len(verified)} из {len(sections)}"
|
||
update_job(job_id, processed_characters=processed,
|
||
current_chapter=label, duration_ms=duration,
|
||
chapters_json=json.dumps(markers, ensure_ascii=False))
|
||
elif event.get("type") == "started":
|
||
update_job(job_id, current_chapter="Проверка готовых фрагментов перед сборкой M4A…")
|
||
elif event.get("type") == "assembling":
|
||
update_job(job_id, status="encoding", stage="encoding", current_chapter=None)
|
||
|
||
result = render_audiobook(sections, row["title"], row["author"], Path(row["output_path"]),
|
||
on_event, lambda: cancellation_requested(job_id))
|
||
if cancellation_requested(job_id):
|
||
finish_cancellation(job_id, job_root)
|
||
return
|
||
update_job(job_id, status="ready", stage="ready", processed_characters=total,
|
||
duration_ms=round(float(result["duration_s"]) * 1000), current_chapter=None)
|
||
except Exception as error:
|
||
if cancellation_requested(job_id):
|
||
finish_cancellation(job_id, job_root)
|
||
else:
|
||
update_job(job_id, status="failed", stage="failed", error_message=f"{type(error).__name__}: {error}")
|
||
|
||
def worker_loop() -> None:
|
||
while True:
|
||
with queue_condition:
|
||
while not queued_jobs:
|
||
queue_condition.wait()
|
||
job_id = queued_jobs.pop(0)
|
||
worker_state["activeJobId"] = job_id
|
||
try:
|
||
env = dict(os.environ, AUDIOBOOK_WORKER_PROCESS="1")
|
||
process = subprocess.Popen([sys.executable, str(Path(__file__).with_name("worker_runner.py")), job_id], env=env)
|
||
while process.poll() is None:
|
||
try:
|
||
worker_state["modelLoaded"] = get_job(job_id)["status"] == "synthesizing"
|
||
except HTTPException:
|
||
pass
|
||
time.sleep(0.5)
|
||
if process.returncode:
|
||
try:
|
||
update_job(job_id, status="failed", stage="failed", error_message="Audiobook worker stopped unexpectedly")
|
||
except HTTPException:
|
||
pass
|
||
finally:
|
||
worker_state["activeJobId"] = None
|
||
worker_state["modelLoaded"] = False
|
||
|
||
|
||
if os.environ.get("AUDIOBOOK_WORKER_PROCESS") != "1":
|
||
initialize_database()
|
||
if os.environ.get("AUDIOBOOK_DISABLE_WORKER") != "1":
|
||
threading.Thread(target=worker_loop, name="audiobook-worker", daemon=True).start()
|