Intercepción avanzada de herramientas para agentes del ADK de RMI

Descripción general

Cuando se trabaja con conjuntos de datos grandes, pasar datos sin procesar directamente a través del modelo causa problemas: ralentiza las respuestas, aumenta los costos de tokens y genera errores de razonamiento. Los datos de RMI pueden abarcar millones de filas, por lo que permitir que el agente consulte una cantidad ilimitada y la ingiera para el análisis es costoso en tokens y aumenta la ventana de contexto, lo que degrada el rendimiento del agente en cada turno posterior.

Por ejemplo, una consulta en la tabla historical_travel_time de RMI puede mostrar miles de filas de segmentos de ruta. Si vuelves a colocar el conjunto de resultados completo en el historial de conversaciones, se ralentizará cada turno posterior, aumentará el costo y distraerá al modelo de la síntesis de la respuesta final.

Las instrucciones de la solicitud por sí solas no pueden resolver estos problemas. El único enfoque confiable es interceptar las llamadas a herramientas y enrutar las cargas útiles de datos grandes fuera de la ventana de contexto del modelo.

Devoluciones de llamada a herramientas del ADK

El ADK proporciona hooks de devolución de llamada para interceptar la ejecución de herramientas. Para controlar conjuntos de resultados grandes, usa after_tool_callback:

  • before_tool_callback(tool, args, tool_context): Se ejecuta después de que el modelo selecciona una herramienta, pero antes de que se ejecute.
  • after_tool_callback(tool, args, tool_context, tool_response): Se ejecuta después de que finaliza la herramienta, pero antes de que el resultado vuelva al modelo. Úsala para modificar la respuesta de la herramienta.

La devolución de llamada recibe un objeto tool_context con acceso al state de la sesión. Esto te permite almacenar conjuntos de datos grandes en la memoria en segundo plano en los turnos de conversación sin colocarlos en la solicitud del modelo.

# root_agent = llm_agent.Agent(
#     ...,
#     after_tool_callback=stash_and_truncate_results,
# )

Truncamiento y almacenamiento posteriores a la ejecución

Objetivo: Mostrarle al modelo una pequeña muestra de los resultados de la consulta para la toma de decisiones, mientras se guarda el conjunto de datos completo para el resultado final.

Para el agente de RMI, un after_tool_callback intercepta los resultados de execute_sql, guarda la lista completa de filas en el estado de la sesión y muestra una muestra corta junto con el recuento total de filas:

def stash_and_truncate_results(tool, args, tool_context, tool_response):
  if tool.name == "execute_sql" and "rows" in tool_response:
    rows = tool_response["rows"]
    tool_context.state["last_sql_result"] = rows      # keep everything
    tool_response["total_rows_fetched"] = len(rows)
    tool_response["rows"] = rows[:SAMPLE_LIMIT]        # show a sample
    tool_response["message"] = (
        f"Showing {SAMPLE_LIMIT} of {len(rows)} rows. "
        "Full result stashed in background state."
    )
  return tool_response

Si se muestra una muestra, se reducen los tokens de entrada en todos los turnos futuros, lo que disminuye los costos y los tiempos de respuesta. También evita que los resultados de consultas grandes saturen la ventana de contexto y perjudiquen la capacidad del modelo para seguir instrucciones.

Asegúrate de que la solicitud de tu agente explique explícitamente que las herramientas muestran una muestra por diseño. Sin este contexto, el modelo puede confundirse con datos parciales, desperdiciar tokens de razonamiento y procesamiento tratando de conciliar los registros faltantes o iniciar llamadas a herramientas redundantes para recuperar el resto.

Ejemplo: Cómo controlar un conjunto de resultados grande

  1. El usuario pregunta "¿Cuáles son las rutas que tienen índices de retraso superiores a 2 en este momento?"
  2. La consulta del agente en historical_travel_time muestra miles de filas.
  3. El after_tool_callback guarda el resultado completo en el estado de la sesión y muestra solo una pequeña muestra más el recuento total de filas.
  4. El modelo razona sobre esa muestra para componer su resumen.
  5. La aplicación o las herramientas posteriores leen el conjunto de datos completo directamente del estado de la sesión para renderizar la tabla completa o exportarla para el usuario.
Interceptación de herramientas del agente de RMI

Conclusiones principales

  • Mantén las cargas útiles grandes fuera de las solicitudes: Almacena conjuntos de datos grandes en el estado de la sesión en lugar de pasarlos de un lado a otro a través del historial de solicitudes.
  • Trunca los resultados de herramientas grandes: Usa after_tool_callback para mostrar muestras pequeñas al modelo mientras guardas los resultados completos en el estado de la sesión.
  • Informa al modelo sobre las muestras: Agrega una regla de solicitud que explique que los resultados se muestrean para evitar la confusión del modelo, los tokens de razonamiento desperdiciados y las llamadas a herramientas redundantes.

Próximos pasos

Colaboradores

Nathaniel Thomas | Pasante de ingeniería de software, Google Maps Platform