doc-translator/src/worker/app.py

from fastapi import FastAPI
from pydantic import BaseModel
import fitz  # PyMuPDF
import os
import psycopg2
from datetime import datetime
import json

DB_DSN = os.getenv("DB_DSN", "postgresql://doc_user:doc_password@db:5432/doc_translator")

app = FastAPI()


class ProcessRequest(BaseModel):
    document_id: int
    original_path: str


def get_db_conn():
    return psycopg2.connect(DB_DSN)


@app.post("/process-document")
def process_document(req: ProcessRequest):
    doc_id = req.document_id
    path = req.original_path

    if not os.path.exists(path):
        _set_status(doc_id, "failed", f"File not found: {path}")
        return {"ok": False, "error": "file_not_found"}

    try:
        _set_status(doc_id, "processing", None)

        # 1. Открываем PDF
        pdf = fitz.open(path)
        page_count = pdf.page_count

        # Обновим page_count в documents
        _update_document_page_count(doc_id, page_count)

        # 2. Проходим по страницам и блокам текста
        translations = []
        for page_index in range(page_count):
            page = pdf[page_index]
            blocks = page.get_text("blocks")  # список блоков: (x0, y0, x1, y1, text, block_no, block_type, ...)
            segment_index = 0
            for b in blocks:
                x0, y0, x1, y1, text, *_rest = b
                if not text.strip():
                    continue
                segment_index += 1

                # TODO: здесь вызвать DeepSeek через Ollama
                translated = _translate_text(text)

                bbox = {"x0": x0, "y0": y0, "x1": x1, "y1": y1}

                translations.append({
                    "document_id": doc_id,
                    "page_number": page_index + 1,
                    "segment_index": segment_index,
                    "source_text": text,
                    "translated_text": translated,
                    "bbox": bbox,
                })

        # 3. Сохраняем сегменты в таблицу translations
        _insert_translations(translations)

        # 4. Собираем новый PDF с русским текстом
        output_path = _build_translated_pdf(path, pdf, translations, doc_id)

        # 5. Обновляем запись в documents
        _set_status(doc_id, "completed", None, output_path=output_path)

        return {"ok": True, "document_id": doc_id, "output_path": output_path}

    except Exception as e:
        _set_status(doc_id, "failed", str(e))
        return {"ok": False, "error": str(e)}


def _set_status(doc_id: int, status: str, error: str | None, output_path: str | None = None):
    conn = get_db_conn()
    cur = conn.cursor()
    cur.execute(
        """
        UPDATE documents
        SET status = %s,
            error_message = %s,
            output_path = COALESCE(%s, output_path),
            updated_at = %s
        WHERE id = %s
        """,
        (status, error, output_path, datetime.now(), doc_id),
    )
    conn.commit()
    cur.close()
    conn.close()


def _update_document_page_count(doc_id: int, page_count: int):
    conn = get_db_conn()
    cur = conn.cursor()
    cur.execute(
        "UPDATE documents SET page_count = %s WHERE id = %s",
        (page_count, doc_id),
    )
    conn.commit()
    cur.close()
    conn.close()


def _translate_text(text: str) -> str:
    """
    Заглушка перевода.
    Здесь надо:
    - взять активный промпт из таблицы prompts;
    - вызвать Ollama (DeepSeek) HTTP-запросом;
    - вернуть перевод.
    Пока просто возвращаем исходный текст.
    """
    return text  # TODO: заменить на реальный вызов LLM


def _insert_translations(rows: list[dict]):
    conn = get_db_conn()
    cur = conn.cursor()
    for r in rows:
        cur.execute(
            """
            INSERT INTO translations
            (document_id, page_number, segment_index, source_text, translated_text, bbox, model_name, created_at)
            VALUES (%s, %s, %s, %s, %s, %s::jsonb, %s, %s)
            """,
            (
                r["document_id"],
                r["page_number"],
                r["segment_index"],
                r["source_text"],
                r["translated_text"],
                json.dumps(r["bbox"]),
                "deepseek-r1:7b",
                datetime.now(),
            ),
        )
    conn.commit()
    cur.close()
    conn.close()


def _build_translated_pdf(original_path: str, pdf: fitz.Document, segments: list[dict], doc_id: int) -> str:
    """
    Строим новый PDF:
    - берем исходную страницу как фон;
    - поверх в тех же bbox'ах рисуем русский текст.
    Это упрощенный вариант: не идеальная типографика, но структура и картинки сохраняются.
    """
    import pathlib
    import json

    # Индексируем переводы по (page, segment_index) или лучше по bbox
    by_page: dict[int, list[dict]] = {}
    for s in segments:
        by_page.setdefault(s["page_number"], []).append(s)

    output_dir = os.getenv("OUTPUT_DIR", "/data/output")
    os.makedirs(output_dir, exist_ok=True)

    basename = os.path.basename(original_path)
    name, _ext = os.path.splitext(basename)
    output_path = os.path.join(output_dir, f"{name}_ru.pdf")

    # Создаем новый PDF
    out_doc = fitz.open()

    for page_index in range(pdf.page_count):
        src_page = pdf[page_index]
        rect = src_page.rect
        new_page = out_doc.new_page(width=rect.width, height=rect.height)

        # Рендерим исходную страницу как картинку (фон)
        pix = src_page.get_pixmap(dpi=150)
        img_rect = fitz.Rect(0, 0, rect.width, rect.height)
        new_page.insert_image(img_rect, pixmap=pix)

        # Наносим переведенный текст поверх
        page_number = page_index + 1
        for seg in by_page.get(page_number, []):
            bbox = seg["bbox"]
            r = fitz.Rect(bbox["x0"], bbox["y0"], bbox["x1"], bbox["y1"])
            new_page.insert_textbox(
                r,
                seg["translated_text"],
                fontsize=8,
                fontname="helv",
                align=fitz.TEXT_ALIGN_LEFT,
            )

    out_doc.save(output_path)
    out_doc.close()
    pdf.close()

    return output_path