IPSExtractor/application/postprocess_agent.py

"""Agente inteligente de post-procesamiento de CSV.

Genera código Python con un LLM para transformar DataFrames según
instrucciones del usuario en lenguaje natural.
Usa la API de Responses de OpenAI (requerida por modelos codex).
"""
from __future__ import annotations

import logging
import re
import threading
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import pandas as pd
from openai import OpenAI

logger = logging.getLogger(__name__)

POSTPROCESS_MODEL = "gpt-5.1-codex-mini"

SYSTEM_PROMPT = """\
Eres un agente experto en transformación de datos con pandas.

REGLAS ESTRICTAS:
1. Recibes una MUESTRA representativa de un CSV y una instrucción del usuario.
2. Debes devolver ÚNICAMENTE un bloque de código Python válido dentro de ```python ... ```.
3. El código opera sobre una variable `df` que es un `pd.DataFrame` ya cargado.
4. El código debe modificar `df` in-place o reasignar a `df`. Al final, `df` debe contener el resultado.
5. Solo puedes usar las siguientes librerías (ya importadas): `pd` (pandas), `np` (numpy), `re`, `datetime`.
6. NO puedes usar: open, exec, eval, __import__, compile, subprocess, os, sys, requests, ni ninguna forma de I/O o acceso al sistema.
7. El código debe ser conciso y eficiente.
8. NO incluyas explicaciones fuera del bloque de código.
"""

RETRY_TEMPLATE = """\
El código que generaste produjo el siguiente error al ejecutarse:

CÓDIGO:
```python
{code}
```

ERROR:
{error}

INSTRUCCIÓN ORIGINAL DEL USUARIO:
{instructions}

MUESTRA DEL CSV:
{sample}

Por favor genera un código corregido que solucione el error. Devuelve ÚNICAMENTE el bloque ```python ... ```.
"""


def sample_dataframe(df: pd.DataFrame, n: int = 5) -> Tuple[str, Dict[str, Any]]:
  """Genera una muestra inteligente del DataFrame.

  Retorna (sample_text, column_info) donde sample_text es una
  representación legible y column_info contiene metadatos.
  """
  total_rows = len(df)
  total_cols = len(df.columns)

  col_info = {
    "total_rows": total_rows,
    "total_cols": total_cols,
    "columns": {col: str(df[col].dtype) for col in df.columns},
  }

  used_indices: set[int] = set()
  sections: List[str] = []

  sections.append(
    f"Total filas: {total_rows:,} | Total columnas: {total_cols}\n"
    f"Columnas y tipos: {', '.join(f'{c} ({t})' for c, t in col_info['columns'].items())}"
  )

  head_n = min(n, total_rows)
  head_idx = list(range(head_n))
  used_indices.update(head_idx)
  sections.append(f"--- Primeras {head_n} filas ---\n{df.iloc[head_idx].to_csv(index=False)}")

  if total_rows > n:
    tail_start = max(total_rows - n, head_n)
    tail_idx = list(range(tail_start, total_rows))
    new_tail = [i for i in tail_idx if i not in used_indices]
    if new_tail:
      used_indices.update(new_tail)
      sections.append(f"--- Últimas {len(new_tail)} filas ---\n{df.iloc[new_tail].to_csv(index=False)}")

  if total_rows > 2 * n:
    mid_center = total_rows // 2
    mid_start = max(mid_center - n // 2, 0)
    mid_idx = list(range(mid_start, min(mid_start + n, total_rows)))
    new_mid = [i for i in mid_idx if i not in used_indices]
    if new_mid:
      used_indices.update(new_mid)
      sections.append(f"--- Filas del medio ---\n{df.iloc[new_mid].to_csv(index=False)}")

  remaining = [i for i in range(total_rows) if i not in used_indices]
  if remaining:
    rand_n = min(n, len(remaining))
    rng = np.random.default_rng(42)
    rand_idx = sorted(rng.choice(remaining, size=rand_n, replace=False).tolist())
    used_indices.update(rand_idx)
    sections.append(f"--- {rand_n} filas aleatorias ---\n{df.iloc[rand_idx].to_csv(index=False)}")

  sample_text = "\n\n".join(sections)
  return sample_text, col_info


def _extract_python_code(text: str) -> str:
  """Extrae el bloque de código Python de la respuesta del LLM."""
  pattern = r"```python\s*\n(.*?)```"
  match = re.search(pattern, text, re.DOTALL)
  if match:
    return match.group(1).strip()
  pattern2 = r"```\s*\n(.*?)```"
  match2 = re.search(pattern2, text, re.DOTALL)
  if match2:
    return match2.group(1).strip()
  return text.strip()


def _call_responses_api(
  client: OpenAI,
  instructions: str,
  user_input: str,
) -> Tuple[str, Optional[Dict[str, Any]]]:
  """Llama a la API de Responses (requerida por modelos codex)."""
  response = client.responses.create(
    model=POSTPROCESS_MODEL,
    instructions=instructions,
    input=user_input,
  )
  usage = None
  raw_usage = getattr(response, "usage", None)
  if raw_usage is not None:
    usage = {
      "input_tokens": getattr(raw_usage, "input_tokens", None),
      "output_tokens": getattr(raw_usage, "output_tokens", None),
      "total_tokens": getattr(raw_usage, "total_tokens", None),
    }
  return response.output_text, usage


def generate_code(
  client: OpenAI,
  csv_sample: str,
  user_instructions: str,
) -> Tuple[str, Optional[Dict[str, Any]]]:
  """Genera código Python invocando al LLM via Responses API.

  Retorna (code_string, usage_dict).
  """
  user_input = (
    f"MUESTRA DEL CSV:\n{csv_sample}\n\n"
    f"INSTRUCCIONES:\n{user_instructions}"
  )
  raw_text, usage = _call_responses_api(client, SYSTEM_PROMPT, user_input)
  code = _extract_python_code(raw_text)
  return code, usage


def _build_safe_builtins() -> dict:
  import builtins as _b
  blocked = {"open", "exec", "eval", "__import__", "compile",
             "exit", "quit", "breakpoint", "input"}
  return {k: v for k, v in vars(_b).items() if k not in blocked}

_SAFE_BUILTINS = _build_safe_builtins()


def execute_code(
  code: str,
  df: pd.DataFrame,
  timeout_s: float = 60,
) -> Tuple[Optional[pd.DataFrame], Optional[str]]:
  """Ejecuta código Python generado en un entorno restringido.

  Retorna (df_result, error_message). Si error_message es None, la
  ejecución fue exitosa.
  """
  import datetime as _dt

  df_copy = df.copy()
  namespace: Dict[str, Any] = {
    "__builtins__": _SAFE_BUILTINS,
    "pd": pd,
    "np": np,
    "re": re,
    "datetime": _dt,
    "df": df_copy,
  }

  result_holder: Dict[str, Any] = {"df": None, "error": None}

  def _run():
    try:
      exec(code, namespace)  # noqa: S102
      result_holder["df"] = namespace.get("df", df_copy)
    except Exception as exc:
      result_holder["error"] = f"{type(exc).__name__}: {exc}"

  thread = threading.Thread(target=_run, daemon=True)
  thread.start()
  thread.join(timeout=timeout_s)

  if thread.is_alive():
    return None, "Timeout: la ejecución excedió el límite de tiempo."

  if result_holder["error"]:
    return None, result_holder["error"]

  result_df = result_holder["df"]
  if not isinstance(result_df, pd.DataFrame):
    return None, "El código no produjo un DataFrame válido en la variable `df`."

  return result_df, None


def retry_with_error(
  client: OpenAI,
  csv_sample: str,
  user_instructions: str,
  failed_code: str,
  error_msg: str,
) -> Tuple[str, Optional[Dict[str, Any]]]:
  """Re-envía al LLM el código fallido y el error para corrección."""
  retry_content = RETRY_TEMPLATE.format(
    code=failed_code,
    error=error_msg,
    instructions=user_instructions,
    sample=csv_sample,
  )
  raw_text, usage = _call_responses_api(client, SYSTEM_PROMPT, retry_content)
  code = _extract_python_code(raw_text)
  return code, usage


def detect_delimiter(filepath: str) -> str:
  """Detecta el delimitador de un archivo CSV leyendo las primeras líneas."""
  import csv as _csv

  with open(filepath, "r", encoding="utf-8", errors="replace") as f:
    sample = f.read(8192)

  try:
    dialect = _csv.Sniffer().sniff(sample, delimiters=",;\t|")
    return dialect.delimiter
  except _csv.Error:
    return ","