Blame for rag_corrector.py

"""
rag_corrector.py
────────────────
Núcleo del sistema RAG.
  1. Detecta posibles errores HTR y grafías modernas en el texto de entrada.
  2. Recupera ejemplos similares del vector store (few-shot dinámico).
  3. Construye el prompt con reglas, ejemplos y alertas.
  4. Llama al LLM y devuelve el texto corregido + trazabilidad.

Uso:
    from rag_corrector import RAGCorrector
    corrector = RAGCorrector(vector_store)
    result = corrector.correct("texto htr aqui")
    print(result["corrected"])
    print(result["prompt"])          # para depuración
"""

import os
import re
from typing import List, Dict, Tuple
from openai import OpenAI
from dotenv import load_dotenv

from knowledge_base import HTR_ERROR_PATTERNS, GRAFIA_PATTERNS

load_dotenv()

MODEL = os.getenv("OPENAI_MODEL", "gpt-4o")
TOP_K = int(os.getenv("TOP_K", 5))

SYSTEM_PROMPT = """Eres un corrector especializado en documentos notariales y judiciales \
españoles del siglo XVI (castellano antiguo).
Tu ÚNICA tarea es corregir los errores introducidos por el proceso automático de \
reconocimiento de texto manuscrito (HTR). NO debes modernizar el texto bajo ninguna \
circunstancia.
REGLAS ABSOLUTAS — incumplirlas invalida la corrección:
1. Conserva SIEMPRE las grafías propias del s.XVI:
   fizo, fazer, hazer, merçed, vezino, mesmo, çibdad, escriuano,
   dho (=dicho), q̃ (=que), nro (=nuestro), vn/vna, etc.
2. NO conviertas f→h inicial latina (fizo ≠ hizo, fazer ≠ hacer).
3. Conserva abreviaturas y tildes voladas (q̃, nro, dho, sr).
4. Corrige SOLO lo que claramente sea un error HTR (grafema confundido visualmente).
5. Si no estás seguro de si algo es error HTR o forma s.XVI válida → conserva el original.
6. Presta especial atención a los errores marcados como [CRÍTICO] — son los de mayor \
impacto en la comprensión del texto.
7. NO corrijas ningún patrón marcado con ⚠ NO_CORREGIR — esas formas son válidas en s.XVI.
8. Responde ÚNICAMENTE con el texto corregido. Sin explicaciones, sin comillas."""


class RAGCorrector:
    def __init__(self, vector_store):
        self.vs = vector_store
        self.client = OpenAI(
            api_key=os.getenv("OPENAI_API_KEY"),
            base_url=os.getenv("OPENAI_BASE_URL", "https://api.x.ai/v1"),
        )

    # ── API pública ──────────────────────────────────────────────────────────

    def correct(self, htr_text: str, top_k: int = TOP_K, model: str = None) -> Dict:
        """
        Corrige un texto HTR usando RAG.
        Retorna dict con:
          corrected     : str   — texto corregido
          prompt        : str   — prompt completo enviado al LLM
          retrieved     : list  — documentos recuperados del vector store
          htr_errors    : list  — patrones HTR detectados
          grafia_warns  : list  — grafías modernas detectadas (alertas)
          model         : str   — modelo usado
        """
        retrieved    = self.vs.retrieve(htr_text, k=top_k)
        htr_errors   = self._detect_htr_errors(htr_text)
        grafia_warns = self._detect_grafias(htr_text)

        prompt = self._build_prompt(htr_text, retrieved, htr_errors, grafia_warns)
        corrected = self._call_llm(prompt, model=model or MODEL)

        return {
            "corrected":    corrected,
            "prompt":       prompt,
            "retrieved":    retrieved,
            "htr_errors":   htr_errors,
            "grafia_warns": grafia_warns,
            "model":        model or MODEL,
        }

    # ── Detección de patrones ────────────────────────────────────────────────

    def _detect_htr_errors(self, text: str) -> List[Dict]:
        """Detecta patrones de error HTR en el texto usando word boundaries."""
        found = []
        for p in HTR_ERROR_PATTERNS:
            htr_val = p.get("htr", "")
            # Ignorar patrones descriptivos (fusion_palabras, omision_*, etc.)
            if "_" in htr_val or len(htr_val) > 10:
                continue
            pattern = r'\b' + re.escape(htr_val) + r'\b'
            if re.search(pattern, text, re.IGNORECASE):
                found.append(p)
        return found

    def _detect_grafias(self, text: str) -> List[Dict]:
        """Detecta formas modernas que NO deberían modernizarse."""
        found = []
        lower = text.lower()
        for p in GRAFIA_PATTERNS:
            pattern = r'\b' + re.escape(p["modern"].lower()) + r'\b'
            if re.search(pattern, lower):
                found.append(p)
        return found

    # ── Constructor de prompt ────────────────────────────────────────────────

    def _build_prompt(
        self,
        htr_text:     str,
        retrieved:    List[Dict],
        htr_errors:   List[Dict],
        grafia_warns: List[Dict],
    ) -> str:

        sections = []

        # ── Bloque 1: Few-shot dinámico (ejemplos recuperados) ────────────────
        if retrieved:
            examples = []
            for i, doc in enumerate(retrieved, 1):
                corr = "; ".join(doc["corrections"]) if doc["corrections"] else "—"
                examples.append(
                    f"Ejemplo {i} [{doc['type']}, {doc['region']}, {doc['date']}]"
                    f" caligrafia: {doc.get('caligrafia', '?')}"
                    f" (similitud={doc['score']}):\n"
                    f"  HTR:  \"{doc['htr']}\"\n"
                    f"  GT:   \"{doc['gt']}\"\n"
                    f"  Correcciones aplicadas: {corr}"
                )
            sections.append(
                "EJEMPLOS DEL CORPUS (similares al texto a corregir):\n"
                + "\n\n".join(examples)
            )

        # ── Bloque 2: Alertas de errores HTR detectados ───────────────────────
        if htr_errors:
            # Ordenar: críticos primero
            SEV_ORDER = {"alta": 0, "media": 1, "baja": 2}
            sorted_errors = sorted(
                htr_errors,
                key=lambda p: SEV_ORDER.get(p.get("severity", "media"), 1)
            )

            hints = []
            for p in sorted_errors:
                sev      = p.get("severity", "media").upper()
                htr_c    = p.get("htr", "")
                gt_c     = p.get("gt", "")
                context  = p.get("context", "")
                tendencia= p.get("tendencia", "")
                warning  = p.get("warning")

                # Ejemplos: hasta 3 de la lista completa
                ejemplos_raw = p.get("examples", p.get("example", ""))
                if isinstance(ejemplos_raw, list):
                    ejemplos_str = "; ".join(ejemplos_raw[:3])
                else:
                    ejemplos_str = ejemplos_raw

                linea = (
                    f"  • [{sev}] '{htr_c}' puede ser '{gt_c}'\n"
                    f"    Causa    : {context}"
                )
                if tendencia:
                    linea += f"\n    Tendencia: {tendencia}"
                if ejemplos_str:
                    linea += f"\n    Ejemplos : {ejemplos_str}"
                if warning:
                    linea += f"\n    ⚠ {warning}"

                hints.append(linea)

            sections.append(
                "POSIBLES ERRORES HTR DETECTADOS EN ESTE TEXTO:\n"
                + "\n\n".join(hints)
            )

        # ── Bloque 3: Alertas de grafías modernas ─────────────────────────────
        if grafia_warns:
            warns = "\n".join(
                f"  • '{p['modern']}' → mantener como '{p['ancient']}': {p['rule']}"
                for p in grafia_warns
            )
            sections.append(
                "ALERTA — GRAFÍAS QUE NO DEBEN MODERNIZARSE:\n" + warns
            )

        context_block = "\n\n".join(sections)

        return (
            f"{context_block}\n\nTEXTO HTR A CORREGIR:\n\"{htr_text}\""
            if context_block
            else f"TEXTO HTR A CORREGIR:\n\"{htr_text}\""
        )

    # ── Llamada al LLM ───────────────────────────────────────────────────────

    def _call_llm(self, user_prompt: str, model: str = MODEL) -> str:
        response = self.client.chat.completions.create(
            model=model,
            temperature=0.1,
            max_tokens=1024,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user",   "content": user_prompt},
            ],
        )
        return response.choices[0].message.content.strip().strip('"').strip("'")