IPSExtractor/services/extraction/domain/normalizer.py

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

Usa un backend compatible con 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 normalizer_llm_config, settings
from .prompts import PromptBuilder
from .llm_client import LLMClient

logger = logging.getLogger('domain.normalizer')


def _build_normalizer_client() -> tuple[LLMClient, dict[str, Any]]:
  config = normalizer_llm_config()
  client = LLMClient(
    api_key=str(config["api_key"]),
    base_url=str(config["base_url"]),
    default_model=str(config["model"]),
    request_timeout_s=float(config["timeout_s"]),
  )
  return client, config


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)
    llm_client, normalizer_config = _build_normalizer_client()
    eff_model = str(normalizer_config["model"])
    eff_base_url = str(normalizer_config["base_url"])
    eff_timeout = float(normalizer_config["timeout_s"])
    # 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",
        max_completion_tokens=16384,
      )
    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.normalizer_model),
          "base_url": (eff_base_url if 'eff_base_url' in locals() else settings.api.effective_normalizer_base_url),
          "radicado": (data.get("radicado") if isinstance(data, dict) else None),
        },
      )
    except Exception:
      pass
    return None