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 a GPT-4o 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
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. 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"))
        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]:
        found = []
        for p in HTR_ERROR_PATTERNS:
            if p["htr"] in text:
                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:
            if p["modern"].lower() in 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 = []

        # 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['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)
            )

        # Alertas de patrones HTR detectados
        if htr_errors:
            hints = "\n".join(
                f"  • '{p['htr']}' puede ser '{p['gt']}': {p['context']} (ej: {p['example']})"
                for p in htr_errors
            )
            sections.append(f"POSIBLES ERRORES HTR DETECTADOS EN ESTE TEXTO:\n{hints}")

        # 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(
                f"ALERTA — GRAFÍAS QUE NO DEBEN MODERNIZARSE:\n{warns}"
            )

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

        return (
            f"{context_block}\n\n"
            f"TEXTO 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, # usa el modelo que llega, no el de .env
            temperature=0.1,        # baja temperatura: reproducibilidad
            max_tokens=1024,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user",   "content": user_prompt},
            ],
        )
        return response.choices[0].message.content.strip().strip('"').strip("'")