text_to_solver_program/modeling/extraction.py

"""Stufen 2 und 5 der Pipeline: Modellspezifikation aus dem Text extrahieren und nach der Prüfung korrigieren."""
from pathlib import Path
from typing import List

from core.llm import LLMSession
from core.model_schema import (
    AssumptionReview,
    CheckIssue,
    CheckSeverity,
    OptimizationModelSpecification,
    TaskUnderstanding,
)

ROUTING_PROMPT_FILE = Path(__file__).parent / "Routing-und-Extraktions-Prompt.txt"


def text_to_model(session: LLMSession, task: str,
                  understanding: TaskUnderstanding | None = None,
                  data_prompt: str = "") -> OptimizationModelSpecification:
    inhalt = task if not data_prompt else f"{task}\n\n{data_prompt}"
    if understanding is not None:
        # Die Verstehensstufe hat Teilfragen und Zahlen schon erfasst – die Extraktion muss den Text
        # nicht noch einmal danach absuchen, sondern baut darauf auf.
        inhalt += ("\n\nBEREITS ERFASSTES VERSTÄNDNIS (Teilfragen und Zahlen; übernimm die Werte aus facts "
                   "in extracted_values, statt sie neu zu suchen):\n" + understanding.model_dump_json(indent=2))
    return session.create(
        OptimizationModelSpecification,
        [
            {"role": "system", "content": ROUTING_PROMPT_FILE.read_text(encoding="utf-8")},
            {"role": "user", "content": inhalt},
        ],
        step="Extraktion",
    )


def revise_model(
    session: LLMSession, task: str, spec: OptimizationModelSpecification, review: AssumptionReview,
    issues: List[CheckIssue] = (), data_prompt: str = "",
) -> OptimizationModelSpecification:
    errors = [i for i in issues if i.severity == CheckSeverity.ERROR]
    check_text = "\n".join(f"- [{i.code}] {i.message}" for i in errors) or "keine"
    # Warnungen lösen für sich keine Korrekturrunde aus; findet aber ohnehin eine statt,
    # sollen sie mit behoben werden – sonst erreicht z. B. unerklärter Fachjargon das LLM nie.
    hinweise = [i for i in issues if i.severity == CheckSeverity.WARNING]
    hinweis_text = "\n".join(f"- [{i.code}] {i.message}" for i in hinweise) or "keine"
    aufgabe = f"{task}\n\n{data_prompt}" if data_prompt else task   # Datenprofil wie bei der Extraktion
    return session.create(
        OptimizationModelSpecification,
        [
            {"role": "system", "content": ROUTING_PROMPT_FILE.read_text(encoding="utf-8")},
            {
                "role": "user",
                "content": (
                    f"AUFGABENTEXT:\n{aufgabe}\n\n"
                    f"BISHERIGE MODELLSPEZIFIKATION:\n{spec.model_dump_json(indent=2)}\n\n"
                    f"PRÜFBERICHT:\n{review.model_dump_json(indent=2)}\n\n"
                    f"FEHLER DER AUTOMATISCHEN KONSISTENZPRÜFUNG:\n{check_text}\n\n"
                    f"HINWEISE DER AUTOMATISCHEN PRÜFUNG (kein Fehler, aber zu verbessern):\n{hinweis_text}\n\n"
                    "Überarbeite die Modellspezifikation: Setze alle Korrekturen aus dem Prüfbericht um, "
                    "behebe alle Fehler der Konsistenzprüfung, passe die betroffenen Restriktionen, Variablen "
                    "und Parameter an und weise fehlende Annahmen in open_assumptions aus. "
                    "Arbeite auch die Hinweise ab, soweit das Modell dadurch nicht schlechter wird – "
                    "unerklärte Fachbegriffe gehören in executive_summary, nicht weggelassen. "
                    "Behalte korrekte Teile unverändert bei."
                ),
            },
        ],
        step="Korrektur",
    )