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

Descripción general

Cuando se trabaja con grandes conjuntos de datos, 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 toda 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 sobre la tabla historical_travel_time del RMI puede devolver miles de filas de tramos de ruta. Volver a colocar el conjunto de resultados completo en el historial de conversación ralentiza cada turno posterior, aumenta el costo y distrae al modelo de la síntesis de la respuesta final.

Las instrucciones de la instrucción 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 de 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 la herramienta.
  • 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 a la sesión state. Esto te permite almacenar grandes conjuntos de datos en la memoria en segundo plano durante los turnos de conversación sin incluirlos en la instrucción del modelo.

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

Truncamiento y almacenamiento posterior a la ejecución

Objetivo: Mostrarle al modelo una pequeña muestra de los resultados de la búsqueda para la toma de decisiones y guardar el conjunto de datos completo para el resultado final.

En el caso del 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 devuelve una muestra breve 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

Devolver una muestra reduce 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 instrucción del agente explique explícitamente que las herramientas devuelven una muestra por diseño. Sin este contexto, el modelo puede confundirse con los datos parciales, desperdiciando tokens de procesamiento y razonamiento al intentar conciliar los registros faltantes o lanzar llamadas a herramientas redundantes para recuperar el resto.

Ejemplo: Cómo controlar un conjunto de resultados grande

  1. El usuario pregunta "¿Qué rutas tienen índices de retraso superiores a 2 en este momento?"
  2. La consulta del agente en historical_travel_time devuelve miles de filas.
  3. after_tool_callback guarda el resultado completo en el estado de la sesión y solo devuelve 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 clave

  • Evita incluir cargas útiles grandes en las instrucciones: 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 instrucciones.
  • Trunca las salidas grandes de la herramienta: Usa after_tool_callback para devolver muestras pequeñas al modelo y guardar los resultados completos en el estado de la sesión.
  • Informa al modelo sobre las muestras: Agrega una regla de instrucción que explique que los resultados se muestrean para evitar la confusión del modelo, el desperdicio de tokens de razonamiento y las llamadas a herramientas redundantes.

Próximos pasos

Colaboradores

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