text_to_solver_program/understanding/figures.py

"""Abbildungen in Aufgabentexten: vom Text trennen, als Bild an die Verstehensstufe geben, Übertragung einsetzen.

docling exportiert Abbildungen als `![Image](data:image/png;base64,…)` (ImageRefMode.EMBEDDED) oder als Verweis
auf eine Bilddatei (REFERENCED). Beides bliebe sonst Zeichensalat im Text jeder Stufe: kein Modell sähe den Graphen,
jeder Aufruf bezahlte das Base64 erneut, und die Zahlenprüfung liefe über Base64-Zeichen. Deshalb bekommt nur die
Verstehensstufe das Bild; ab dann arbeitet die Pipeline mit der Übertragung in Text.
"""
import base64
import re
from pathlib import Path
from typing import List, Tuple

from core.model_schema import FigureInput, FigureTranscription

MAX_FIGURES = 6
MAX_BYTES = 5 * 1024 * 1024
IMAGE_TYPES = {".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".gif": "image/gif", ".webp": "image/webp"}
# ![alt](data:image/png;base64,…) oder ![alt](pfad/bild.png); Base64 darf Zeilenumbrüche enthalten
IMAGE_PATTERN = re.compile(r"!\[([^\]]*)\]\(\s*(?:data:(image/[\w.+-]+);base64,([A-Za-z0-9+/=\s]+?)|([^)\s]+))\s*\)")
# docling schreibt als Alternativtext meist nur „Image“ – der trägt nichts zur Aufgabe bei
GENERIC_ALT = {"", "image", "bild", "figure", "abbildung", "picture"}


class FigureError(ValueError):
    """Abbildungen, die nicht verarbeitet werden können (zu viele, zu groß, kaputtes Base64)."""


def split_figures(task: str, basis_dir: Path | str | None = None) -> Tuple[str, List[FigureInput]]:
    """Abbildungen aus dem Aufgabentext lösen; im Text bleibt ein Platzhalter „[Abbildung N]“.

    Verweise auf Dateien werden relativ zu `basis_dir` (dem Verzeichnis der Aufgabendatei) aufgelöst. Nicht
    vorhandene Dateien und Verweise ohne Bild-Endung (etwa Webadressen) bleiben unverändert im Text.
    """
    basis = Path(basis_dir) if basis_dir else Path.cwd()
    figures: List[FigureInput] = []

    def ersetze(treffer: re.Match) -> str:
        alt, mime, daten, pfad = treffer.group(1).strip(), treffer.group(2), treffer.group(3), treffer.group(4)
        if daten is not None:
            daten = re.sub(r"\s", "", daten)
            try:
                roh = base64.b64decode(daten, validate=True)
            except ValueError as exc:
                raise FigureError(f"Eingebettetes Bild {len(figures) + 1} ist kein gültiges Base64: {exc}") from exc
            quelle = "eingebettet"
        else:
            datei = (basis / pfad) if not Path(pfad).is_absolute() else Path(pfad)
            mime = IMAGE_TYPES.get(datei.suffix.lower())
            if mime is None or not datei.is_file():
                return treffer.group(0)
            roh = datei.read_bytes()
            daten = base64.b64encode(roh).decode("ascii")
            quelle = str(datei)
        if len(roh) > MAX_BYTES:
            raise FigureError(f"Abbildung {len(figures) + 1} ist {len(roh) // 1024} KB groß – erlaubt sind "
                              f"{MAX_BYTES // 1024 // 1024} MB.")
        kennung = f"Abbildung {len(figures) + 1}"
        figures.append(FigureInput(id=kennung, alt=alt, mime=mime, data=daten, source=quelle))
        return f"[{kennung}]" if alt.lower() in GENERIC_ALT else f"[{kennung}: {alt}]"

    text = IMAGE_PATTERN.sub(ersetze, task)
    if len(figures) > MAX_FIGURES:
        raise FigureError(f"Die Aufgabe enthält {len(figures)} Abbildungen – verarbeitet werden höchstens {MAX_FIGURES}.")
    return text, figures


def canonical_figure_id(kennung: str) -> str:
    """„1“, „Abb. 1“, „Abbildung 1“ → „Abbildung 1“ – beobachtet: eine Lesung schrieb nur „1“."""
    nummer = re.search(r"\d+", kennung or "")
    return f"Abbildung {nummer.group(0)}" if nummer else (kennung or "").strip()


def normalize_figure_ids(transcriptions: List[FigureTranscription]) -> List[FigureTranscription]:
    return [t.model_copy(update={"id": canonical_figure_id(t.id)}) for t in transcriptions]


def image_parts(figures: List[FigureInput]) -> list:
    """Nachrichtenteile im OpenAI-Format: je Abbildung ihre Kennung als Text, dann das Bild als Data-URI."""
    teile = []
    for figur in figures:
        teile.append({"type": "text", "text": f"{figur.id}:"})
        teile.append({"type": "image_url", "image_url": {"url": f"data:{figur.mime};base64,{figur.data}"}})
    return teile


def with_transcriptions(text: str, transcriptions: List[FigureTranscription]) -> str:
    """Aufgabentext für alle weiteren Stufen: Platzhalter bleiben, die Übertragungen folgen am Ende."""
    if not transcriptions:
        return text
    bloecke = []
    for t in transcriptions:
        block = f"[{t.id}] {t.shows}\n{t.transcription.strip()}"
        if t.uncertain:
            block += "\nUnsicher gelesen: " + "; ".join(t.uncertain)
        bloecke.append(block)
    return (text.rstrip() + "\n\nABBILDUNGEN (aus dem Bild in Text übertragen; das Original liegt dem Report bei):\n\n"
            + "\n\n".join(bloecke))