IPSExtractor/application/normalizer.py

"""Normalizador LLM: adapta un JSON fuente al esquema canónico de rutas.

Usa la API de OpenAI (Chat Completions) para transformar el JSON devuelto por
el extractor a un objeto que respete exactamente las rutas/campos canónicos
seleccionados por el usuario (incluyendo listas items/*).
"""
from __future__ import annotations

import logging
import time
from typing import Any, Dict, List, Optional

from .config import settings
from .prompts import PromptBuilder
from .llm_client import LLMClient

logger = logging.getLogger('application.normalizer')
_llm_client = LLMClient()


def _ensure_dict_at_path(root: Dict[str, Any], key: str) -> Dict[str, Any]:
  val = root.get(key)
  if not isinstance(val, dict):
    val = {}
    root[key] = val
  return val


def _ensure_list_with_object(root: Dict[str, Any], key: str) -> Dict[str, Any]:
  val = root.get(key)
  if not isinstance(val, list) or not val or not isinstance(val[0], dict):
    root[key] = [{}]
    return root[key][0]
  return val[0]


def build_canonical_template(paths: List[str]) -> Dict[str, Any]:
  """Convierte rutas 'a/b/*/c' en un objeto anidado con listas y nulls.

  - Siempre incluye 'radicado' en la raíz.
  - Para 'x/*/y' crea 'x': [{ 'y': null }]. Si hay más subniveles, anida en ese objeto.
  """
  template: Dict[str, Any] = {"radicado": None}
  for raw_path in paths or []:
    if not raw_path:
      continue
    segments = [s for s in raw_path.split('/') if s]
    if not segments:
      continue
    i = 0
    cur: Dict[str, Any] = template
    while i < len(segments):
      name = segments[i]
      is_last = (i == len(segments) - 1)
      # próximo segmento indica lista
      if i + 1 < len(segments) and segments[i + 1] == '*':
        # asegurar lista con un objeto
        cur = _ensure_list_with_object(cur, name)
        i += 2  # saltar el '*'
        # si terminamos después del '*', no hay campo que asignar aquí
        if i >= len(segments):
          break
        # continuar con el objeto dentro de la lista
        continue
      # segmento normal
      if is_last:
        cur[name] = None
      else:
        cur = _ensure_dict_at_path(cur, name)
      i += 1
  return template


def normalize_to_schema(
  data: Dict[str, Any],
  paths: List[str],
  profile_name: str,
) -> Optional[Dict[str, Any]]:
  """
  Usa un LLM para mapear un JSON extraído a un esquema canónico definido por rutas.

  Este proceso asegura que, independientemente de cómo el extractor haya devuelto los datos,
  el resultado final tenga la estructura jerárquica esperada (ej. 'cabecera/fecha', 'items/*/descripcion').

  Args:
      data: El diccionario con los datos extraídos del documento.
      paths: Lista de rutas canónicas (ej. ["cliente/nombre", "totales/iva"]).
      profile_name: Nombre del perfil de extracción (ej. 'invoice') para cargar instrucciones.

  Returns:
      Un diccionario con la clave "normalized" conteniendo el objeto transformado 
      y "_usage" con las estadísticas de tokens. Retorna None si falla.
  """
  try:
    t0 = time.perf_counter()
    template = build_canonical_template(paths)
    # Construcción centralizada de mensajes
    messages = PromptBuilder.build_normalization_messages(profile_name, template, data)
    # Forzar modelo del normalizador a gpt-5-mini
    eff_model = "gpt-4.1-mini"
    eff_base_url = settings.api.base_url or "https://api.openai.com/v1"
    eff_timeout = getattr(settings.api, "request_timeout_s", 30)
    # log inicio
    try:
      logger.info(
        "llm_normalize_start",
        extra={
          "model": eff_model,
          "base_url": eff_base_url,
          "timeout_s": eff_timeout,
          "paths_count": len(paths or []),
          "template_fields": list(template.keys()),
          "radicado": (data.get("radicado") if isinstance(data, dict) else None),
        },
      )
    except Exception:
      pass

    # Llamada centralizada vía LLMClient (JSON estricto)
    try:
      normalized_obj, usage = _llm_client.chat_json(
        messages=messages, model=eff_model,
        response_format="json_object",
        timeout_s=float(eff_timeout),
        component="normalizer",
      )
    except Exception:
      normalized_obj = None
      usage = None

    if not isinstance(normalized_obj, dict):
      try:
        dt_ms = int((time.perf_counter() - t0) * 1000)
        logger.info(
          "llm_normalize_end",
          extra={
            "normalized": False,
            "duration_ms": dt_ms,
            "model": eff_model,
            "base_url": eff_base_url,
            "radicado": (data.get("radicado") if isinstance(data, dict) else None),
          },
        )
      except Exception:
        pass
      return None
    # asegurar radicado en raíz si existía en data
    rad = None
    try:
      rad = data.get("radicado")
    except Exception:
      rad = None
    if rad is not None and isinstance(normalized_obj, dict):
      normalized_obj["radicado"] = rad
    try:
      dt_ms = int((time.perf_counter() - t0) * 1000)
      logger.info(
        "llm_normalize_end",
        extra={
          "normalized": True,
          "duration_ms": dt_ms,
          "model": eff_model,
          "base_url": eff_base_url,
          "radicado": rad,
        },
      )
    except Exception:
      pass
    # Devolver solo el objeto normalizado (compatibilidad) cuando lo llame alguien antiguo
    return {"normalized": normalized_obj, "_usage": usage}
  except Exception as ex:
    try:
      dt_ms = int((time.perf_counter() - t0) * 1000) if 't0' in locals() else None
      logger.info(
        "llm_normalize_end",
        extra={
          "normalized": False,
          "error": str(ex),
          "duration_ms": dt_ms,
          "model": (eff_model if 'eff_model' in locals() else settings.api.model or "gpt-4.1-nano"),
          "base_url": (eff_base_url if 'eff_base_url' in locals() else (settings.api.base_url or "https://api.openai.com/v1")),
          "radicado": (data.get("radicado") if isinstance(data, dict) else None),
        },
      )
    except Exception:
      pass
    return None