Middlewares personalizados

Creación de Middlewares personalizados

Los middlewares personalizados permiten extender la funcionalidad de los agentes inyectando lógica propia en puntos estratégicos de su ciclo de vida. A diferencia de los middlewares nativos, que cubren casos de uso comunes, los personalizados te ofrecen control total para implementar requisitos de negocio específicos, como reglas de cumplimiento, auditoría avanzada o integración con sistemas propietarios.

La arquitectura de middlewares en LangChain se basa en un sistema de hooks (ganchos) que interceptan la ejecución. Puedes implementar estos hooks mediante funciones decoradas para lógica simple o clases para comportamientos más complejos y configurables.

Tipos de Hooks

Existen dos categorías principales de hooks que determinan cómo y cuándo se ejecuta tu código personalizado:

• Node-style Hooks: Se ejecutan en puntos discretos del flujo, como antes o después de una acción. Son ideales para validaciones, logging o modificación del estado.

  • before_agent: Al inicio de la ejecución del agente (una vez).
  • before_model: Justo antes de invocar al LLM.
  • after_model: Inmediatamente después de recibir la respuesta del LLM.
  • after_agent: Al finalizar la ejecución del agente (una vez).

• Wrap-style Hooks: Envuelven la ejecución de una llamada, permitiéndote ejecutar código antes y después, o incluso modificar la llamada misma o su resultado.

  • wrap_model_call: Envuelve la llamada al modelo de lenguaje.
  • wrap_tool_call: Envuelve la ejecución de una herramienta.

Implementación basada en funciones

Para lógica sencilla y sin estado interno, el enfoque más directo es utilizar decoradores. Cada decorador corresponde a uno de los hooks mencionados y recibe como argumentos el estado actual (AgentState) y el contexto de ejecución (Runtime).

1. Logger simple con decorators:

from langchain.agents.middleware import before_model, after_model, AgentState
from langgraph.runtime import Runtime
from typing import Any

@before_model
def log_model_input(state: AgentState, runtime: Runtime) -> None:
    print(f"Llamando al modelo con {len(state['messages'])} mensajes en el historial")

@after_model
def log_model_output(state: AgentState, runtime: Runtime) -> None:
    last_message = state["messages"][-1]
    print(f"Respuesta recibida. Tokens: {last_message.usage_metadata['total_tokens']}")

En este ejemplo, las funciones simplemente leen el estado para imprimir información. Al retornar None, indican al agente que continúe con su flujo normal sin modificaciones.

Control de flujo y modificación de estado

Los middlewares no son solo observadores pasivos; pueden alterar el curso de la ejecución. Retornando un diccionario, puedes actualizar el estado del agente o forzar saltos en el flujo mediante la clave especial jump_to.

2. Middleware de límite de mensajes (Guardrail):

from langchain.agents.middleware import before_model, AgentState
from langchain.messages import AIMessage
from langgraph.runtime import Runtime
from typing import Any

@before_model(can_jump_to=["end"])
def check_message_limit(state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
    # Si superamos los 50 mensajes, cortamos la ejecución
    if len(state["messages"]) >= 50:
        return {
            "messages": [AIMessage(content="Límite de conversación alcanzado.")],
            "jump_to": "end"  # Fuerza la finalización inmediata del agente
        }
    return None

Observa que el decorador requiere declarar explícitamente los destinos permitidos con can_jump_to. Esto garantiza que el flujo no salte a estados no válidos. Si la condición se cumple, inyectamos un mensaje final y terminamos el agente.

Nota sobre persistencia: Este middleware verifica la longitud del historial acumulado. Para que funcione correctamente a través de múltiples interacciones con el usuario, es necesario que el agente tenga configurada una memoria (checkpointer) y se invoque utilizando un thread_id persistente. De lo contrario, cada ejecución comenzará con un historial vacío y el límite nunca se alcanzará.

Interceptación avanzada (Wrap-style)

Los hooks de tipo "wrap" son más potentes, ya que te dan el control de la llamada misma. Reciben un argumento adicional handler, que es una función que ejecuta la lógica original. Tú decides cuándo llamar a handler() y qué hacer con su resultado.

3. Middleware para inyección de headers:

from langchain.agents.middleware import wrap_model_call

@wrap_model_call
async def add_custom_headers(state, runtime, handler):
    # Modificamos la configuración del runtime antes de llamar
    runtime.config["configurable"]["headers"] = {"X-Custom-Trace": "agent-123"}
    
    # Ejecutamos la llamada original
    response = await handler(state, runtime)
    
    # Podríamos modificar la respuesta aquí si fuera necesario
    return response

Este patrón es fundamental para integraciones de infraestructura, como pasar trazas distribuidas, gestionar tokens de autenticación dinámicos o realizar reintentos personalizados a nivel de red.

Registro y uso en el agente

Una vez definidos tus middlewares, se integran en el agente pasándolos en una lista al momento de la creación.

4. Configuración final del agente:

from langchain.agents import create_agent

agent = create_agent(
    model="gpt-4-turbo",
    tools=[],
    middleware=[
        log_model_input,      # Función decorada
        check_message_limit,  # Función decorada con lógica de control
        # add_custom_headers  # Función wrap (si fuera necesaria)
    ]
)

El orden es importante: los hooks before se ejecutan en el orden de la lista, mientras que los after y wrap se resuelven de forma anidada (como una cebolla), donde el primer middleware de la lista es el más externo.

AgentState y Runtime

En todos los ejemplos anteriores, hemos visto que los middlewares reciben dos argumentos fundamentales: state y runtime. Comprender qué contienen estos objetos es clave para escribir lógica efectiva.

AgentState: La memoria del agente

El AgentState es un diccionario tipado que representa la instantánea actual de la conversación. Es el objeto que viaja a través de todos los nodos del grafo (modelo, herramientas, middlewares) acumulando información.

Sus claves principales son:

• messages: Una lista de objetos BaseMessage (HumanMessage, AIMessage, ToolMessage) que forman el historial de la conversación.

Cuando tu middleware retorna un diccionario, LangGraph realiza una operación de merge (fusión) con este estado. Por eso, al retornar {"messages": [nuevo_mensaje]}, el mensaje se añade a la lista existente en lugar de sobrescribirla.

Runtime: El contexto de ejecución

El Runtime es el objeto que proporciona acceso a la infraestructura y dependencias externas durante la ejecución del agente. Según la documentación moderna, expone principalmente:

• context: Información estática inyectada al invocar el agente (como IDs de usuario, conexiones a BD, etc.). Funciona como inyección de dependencias.
• store: Acceso al almacenamiento clave-valor global (si está configurado) para memoria a largo plazo.
• stream_writer: Utilidad para emitir eventos personalizados al stream de salida.

Para utilizar el context, primero debes definir su esquema en create_agent y luego pasar los valores al invocar.

Ejemplo de inspección con Contexto:

from dataclasses import dataclass

# 1. Definimos el esquema del contexto
@dataclass
class UserContext:
    user_id: str
    plan: str

# 2. Middleware que accede al contexto
@before_model
def check_user_tier(state: AgentState, runtime: Runtime[UserContext]):
    # Acceso seguro y tipado a las dependencias
    print(f"Usuario: {runtime.context.user_id}")
    
    if runtime.context.plan == "free":
        print("⚠️ Usuario en plan gratuito")

# 3. Configuración del agente
# agent = create_agent(..., middleware=[check_user_tier], context_schema=UserContext)

Dominar el Runtime te permite separar la configuración (IDs, secretos, preferencias) de la lógica del agente, haciendo tu código más limpio y testable.