"""Framework-agnostic translation runner. Runs the OCR -> translate -> render pipeline in a worker thread (so its internal ``asyncio.run`` works) and streams per-phase progress through a queue. Knows nothing about Gradio, so it can be unit-tested in isolation. """ from __future__ import annotations import logging import queue import tempfile import threading import traceback import uuid from collections.abc import Callable, Iterator from dataclasses import dataclass from pathlib import Path from typing import Any logger = logging.getLogger(__name__) @dataclass class TranslationRequest: """One translation job, with the provider already resolved to its key.""" pdf_path: str | None provider: str # Phase-2 provider key (e.g. "openrouter") api_key: str model: str | None src_lang: str tgt_lang: str font: str pages: list[int] | None @dataclass(frozen=True) class Progress: """A per-phase progress update streamed while the pipeline runs.""" frac: float msg: str @dataclass(frozen=True) class Result: """The terminal outcome of a run. ``data`` carries the step's payload for the stepped flow (e.g. the parsed dict, or ``{"translated": ..., "out_path": ...}``). ``out_path`` is kept for the legacy one-shot ``stream_translation`` path. """ status: str # "ok" | "invalid" | "error" out_path: str | None = None detail: str = "" data: Any = None def validate(req: TranslationRequest) -> str | None: """Return a user-facing error message if the request can't run, else None.""" if not req.pdf_path: return "Vui lòng tải lên một file PDF." if not req.api_key or not req.api_key.strip(): return "Thiếu API key — nhập API key của provider ở thanh bên." if not req.src_lang or not req.tgt_lang: return "Chọn ngôn ngữ nguồn và ngôn ngữ đích." return None def list_models(provider: str, api_key: str) -> list[str]: """Fetch model ids from a provider's OpenAI-compatible ``GET /models`` endpoint. ``provider`` is the resolved key (e.g. "deepseek"). Returns a sorted list of model ids, or [] on any failure (bad key, no endpoint, non-OpenAI response) — the UI then falls back to free-text entry. """ import httpx from pdf2zh.translation.config import provider_base_url if not api_key or not api_key.strip(): return [] try: base = provider_base_url(provider).rstrip("/") resp = httpx.get( f"{base}/models", headers={"Authorization": f"Bearer {api_key.strip()}"}, timeout=15, verify=False, ) resp.raise_for_status() data = resp.json().get("data", []) return sorted({m["id"] for m in data if isinstance(m, dict) and m.get("id")}) except Exception: # noqa: BLE001 — listing is best-effort; fall back to manual logger.warning("list_models failed for provider %s", provider, exc_info=True) return [] def stream_translation(req: TranslationRequest) -> Iterator[Progress | Result]: """Yield ``Progress`` updates while translating, then one terminal ``Result``.""" # Imported lazily so the lightweight bits above (dataclasses, validate) stay # importable without the heavy ML stack (torch, surya, ...) that e2e pulls in. from pdf2zh.e2e import run_pipeline q: queue.Queue = queue.Queue() work_dir = Path(tempfile.gettempdir()) / f"pdf2zh_{uuid.uuid4().hex}" def on_progress(frac: float, msg: str) -> None: q.put(Progress(frac, msg)) def worker() -> None: try: out = run_pipeline( pdf_path=req.pdf_path, src_lang=req.src_lang, tgt_lang=req.tgt_lang, provider=req.provider, api_key=req.api_key, model=req.model, pages=req.pages, font=req.font, work_dir=work_dir, progress=on_progress, ) q.put(Result("ok", out_path=out)) except ValueError as exc: # user-facing input error q.put(Result("invalid", detail=str(exc))) except Exception as exc: # noqa: BLE001 — surface anything else to the UI logger.exception("pipeline failed") tail = "".join(traceback.format_exc().splitlines(keepends=True)[-6:]) q.put( Result("error", detail=f"{type(exc).__name__}: {exc}\n```\n{tail}\n```") ) threading.Thread(target=worker, daemon=True).start() while True: item = q.get() yield item if isinstance(item, Result): return # --------------------------------------------------------------------------- # # Stepped flow — run one phase in a worker thread and stream its progress. # --------------------------------------------------------------------------- # def _stream( fn: Callable[[Callable[[float, str], None]], Any], ) -> Iterator[Progress | Result]: """Run ``fn(progress_cb)`` in a worker thread; stream Progress then a Result. ``fn`` receives a ``progress(frac, msg)`` callback and returns the payload placed on ``Result.data``. A ValueError becomes an ``invalid`` result (user-facing input error); anything else becomes an ``error`` result. """ q: queue.Queue = queue.Queue() def on_progress(frac: float, msg: str) -> None: q.put(Progress(frac, msg)) def worker() -> None: try: data = fn(on_progress) q.put(Result("ok", data=data)) except ValueError as exc: q.put(Result("invalid", detail=str(exc))) except Exception as exc: # noqa: BLE001 — surface anything else to the UI logger.exception("step failed") tail = "".join(traceback.format_exc().splitlines(keepends=True)[-6:]) q.put( Result("error", detail=f"{type(exc).__name__}: {exc}\n```\n{tail}\n```") ) threading.Thread(target=worker, daemon=True).start() while True: item = q.get() yield item if isinstance(item, Result): return def stream_parse( pdf_path: str, pages: list[int] | None, work_dir: str | Path ) -> Iterator[Progress | Result]: """Phase 1 — parse. ``Result.data`` is the parsed doc dict.""" from pdf2zh.e2e import run_parse return _stream(lambda p: run_parse(pdf_path, pages, work_dir, p)) def stream_translate_render( pdf_path: str, parsed: dict, src_lang: str, tgt_lang: str, provider: str, api_key: str, model: str | None, pages: list[int] | None, font: str, work_dir: str | Path, ) -> Iterator[Progress | Result]: """Phase 2 + 3 — translate the (edited) parsed doc, then render. ``Result.data`` is ``{"translated": dict, "out_path": str}``. """ from pdf2zh.e2e import run_render, run_translate def fn(p: Callable[[float, str], None]) -> dict: translated = run_translate( parsed, src_lang, tgt_lang, provider, api_key, model, work_dir, p ) out_path = run_render(pdf_path, translated, pages, font, work_dir, p) return {"translated": translated, "out_path": out_path} return _stream(fn) def stream_render( pdf_path: str, translated: dict, pages: list[int] | None, font: str, work_dir: str | Path, ) -> Iterator[Progress | Result]: """Phase 3 only — re-render the (edited) translated doc. ``Result.data`` is ``{"out_path": str}``.""" from pdf2zh.e2e import run_render def fn(p: Callable[[float, str], None]) -> dict: return {"out_path": run_render(pdf_path, translated, pages, font, work_dir, p)} return _stream(fn)