Agentes de IA con LangGraph para Análisis de Datos

pip install langgraph langchain langchain-openai pandas numpy

export OPENAI_API_KEY="sk-..."

from typing import TypedDict, Annotated, Sequence
from langchain_core.messages import BaseMessage
import operator

class EstadoAgente(TypedDict):
    messages: Annotated[Sequence[BaseMessage], operator.add]
    datos_cargados: bool
    resumen_dataset: str

import pandas as pd
import numpy as np
from langchain_core.tools import tool

# "Base de datos" en memoria para el demo
_df: pd.DataFrame | None = None

@tool
def cargar_dataset(path: str) -> str:
    """Carga un dataset CSV en memoria. Devuelve un resumen de columnas y forma."""
    global _df
    _df = pd.read_csv(path)
    resumen = (
        f"Dataset cargado: {_df.shape[0]} filas × {_df.shape[1]} columnas.\n"
        f"Columnas: {', '.join(_df.columns)}\n"
        f"Tipos:\n{_df.dtypes.to_string()}"
    )
    return resumen

@tool
def ejecutar_consulta(consulta: str) -> str:
    """
    Ejecuta una consulta pandas sobre el dataset cargado.
    Usa expresiones pandas estándar, por ejemplo:
      - df[df['churn'] == 1]['cargo_mensual'].mean()
      - df.groupby('plan')['churn'].value_counts()
    La variable 'df' se refiere al dataset cargado.
    """
    global _df
    if _df is None:
        return "No hay dataset cargado. Llama a cargar_dataset primero."
    try:
        resultado = eval(consulta, {'df': _df, 'pd': pd, 'np': np})
        return str(resultado)
    except Exception as e:
        return f"Error en consulta: {e}"

@tool
def describir_columna(nombre_columna: str) -> str:
    """Devuelve estadísticas descriptivas para una columna específica del dataset cargado."""
    global _df
    if _df is None:
        return "No hay dataset cargado."
    if nombre_columna not in _df.columns:
        return f"Columna '{nombre_columna}' no encontrada. Disponibles: {list(_df.columns)}"
    stats = _df[nombre_columna].describe().to_string()
    pct_nulos = _df[nombre_columna].isnull().mean() * 100
    return f"{stats}\nValores nulos: {pct_nulos:.1f}%"

herramientas = [cargar_dataset, ejecutar_consulta, describir_columna]

from langchain_openai import ChatOpenAI
from langchain_core.messages import SystemMessage
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode

# LLM con herramientas vinculadas
llm = ChatOpenAI(model='gpt-4o-mini', temperature=0)
llm_con_herramientas = llm.bind_tools(herramientas)

PROMPT_SISTEMA = """Eres un asistente de análisis de datos. Tienes acceso a herramientas 
que te permiten cargar datasets CSV, ejecutar consultas pandas y describir columnas.

Cuando el usuario haga una pregunta:
1. Carga el dataset si aún no está cargado.
2. Explora los datos para entender su estructura.
3. Ejecuta las consultas específicas necesarias para responder la pregunta.
4. Resume tus hallazgos claramente, con números.

Siempre verifica que tus consultas funcionaron antes de sacar conclusiones."""

def nodo_agente(state: EstadoAgente) -> EstadoAgente:
    """El nodo de razonamiento: llama al LLM para decidir la siguiente acción."""
    messages = [SystemMessage(content=PROMPT_SISTEMA)] + list(state['messages'])
    respuesta = llm_con_herramientas.invoke(messages)
    return {'messages': [respuesta]}

def debe_continuar(state: EstadoAgente) -> str:
    """Función de enrutamiento: ¿continuar el ciclo o terminar?"""
    ultimo_mensaje = state['messages'][-1]
    if hasattr(ultimo_mensaje, 'tool_calls') and ultimo_mensaje.tool_calls:
        return 'herramientas'
    return END

# Construir el grafo
nodo_herramientas = ToolNode(herramientas)

grafo = StateGraph(EstadoAgente)
grafo.add_node('agente', nodo_agente)
grafo.add_node('herramientas', nodo_herramientas)

grafo.set_entry_point('agente')
grafo.add_conditional_edges('agente', debe_continuar)
grafo.add_edge('herramientas', 'agente')  # Después de la llamada, vuelve al agente

agente = grafo.compile()

from langchain_core.messages import HumanMessage

def ejecutar_agente(pregunta: str, data_path: str | None = None):
    mensaje_inicial = HumanMessage(
        content=f"Ruta del dataset: {data_path}\n\nPregunta: {pregunta}"
        if data_path else pregunta
    )

    resultado = agente.invoke(
        {'messages': [mensaje_inicial], 'datos_cargados': False, 'resumen_dataset': ''},
        config={'recursion_limit': 15},
    )

    mensaje_final = resultado['messages'][-1]
    return mensaje_final.content

# Ejemplo: analizar un dataset de muestra
respuesta = ejecutar_agente(
    pregunta="¿Cuál es la tasa de churn por número de productos? ¿Qué grupo tiene el mayor churn?",
    data_path="datos/churn.csv",
)
print(respuesta)

Cargué el dataset (2000 filas × 6 columnas) y ejecuté el análisis.

Tasa de churn por número de productos:
  1 producto:   42.1%  ← mayor churn
  2 productos:  31.7%
  3 productos:  22.4%
  4 productos:  15.8%
  5 productos:   9.3%  ← menor churn

Los clientes con solo 1 producto tienen la mayor tasa de churn con 42.1%, 
que es 4.5× más alta que los clientes con 5 productos. Esto sugiere que 
la venta cruzada de productos adicionales mejora significativamente la retención.

from langchain_core.messages import AIMessage, ToolMessage

resultado = agente.invoke({
    'messages': [HumanMessage(content="Carga datos/churn.csv y dime el cargo mensual promedio para clientes que se fueron vs los que se quedaron.")],
    'datos_cargados': False,
    'resumen_dataset': '',
})

for msg in resultado['messages']:
    if isinstance(msg, AIMessage) and msg.tool_calls:
        for tc in msg.tool_calls:
            print(f"→ Llamada: {tc['name']}({tc['args']})")
    elif isinstance(msg, ToolMessage):
        print(f"← Resultado: {msg.content[:200]}")
    elif isinstance(msg, AIMessage):
        print(f"✓ Respuesta final: {msg.content}")

→ Llamada: cargar_dataset({'path': 'datos/churn.csv'})
← Resultado: Dataset cargado: 2000 filas × 6 columnas...
→ Llamada: ejecutar_consulta({'consulta': "df.groupby('churn')['cargo_mensual'].mean()"})
← Resultado: churn\n0    58.3\n1    79.2
✓ Respuesta final: Los clientes que se fueron pagan significativamente más en promedio ($79.2/mes)...