Interception avancée des outils pour les agents RMI ADK

Présentation

Lorsque vous traitez de grands ensembles de données, le fait de transmettre directement des données brutes via le modèle pose des problèmes : cela ralentit les réponses, augmente les coûts des jetons et introduit des erreurs de raisonnement. Les données RMI peuvent s'étendre sur des millions de lignes. Par conséquent, laisser l'agent interroger une quantité illimitée et tout ingérer pour l'analyse est coûteux en jetons et gonfle la fenêtre de contexte, ce qui dégrade les performances de l'agent à chaque tour suivant.

Par exemple, une requête sur la table RMI historical_travel_time peut renvoyer des milliers de lignes de segments de route. Le fait de remettre l'ensemble complet des résultats dans l'historique des conversations ralentit chaque tour ultérieur, augmente les coûts et empêche le modèle de synthétiser la réponse finale.

Les instructions d'invite seules ne peuvent pas résoudre ces problèmes. La seule approche fiable consiste à intercepter les appels d'outils et à acheminer les charges utiles de données volumineuses en dehors de la fenêtre de contexte du modèle.

Rappels d'outils ADK

ADK fournit des hooks de rappel pour intercepter l'exécution des outils. Pour gérer de grands ensembles de résultats, utilisez after_tool_callback :

  • before_tool_callback(tool, args, tool_context): s'exécute après que le modèle a sélectionné un outil, mais avant que l'outil ne s'exécute.
  • after_tool_callback(tool, args, tool_context, tool_response): s'exécute une fois l'outil terminé, mais avant que le résultat ne soit renvoyé au modèle. Utilisez-le pour modifier la réponse de l'outil.

Le rappel reçoit un objet tool_context avec accès à state de la session. Cela vous permet de stocker de grands ensembles de données dans la mémoire d'arrière-plan sur plusieurs tours de conversation sans les placer dans l'invite du modèle.

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

Troncature et mise en cache après l'exécution

Objectif : Afficher au modèle un petit échantillon de résultats de requête pour la prise de décision, tout en enregistrant l'ensemble de données complet pour la sortie finale.

Pour l'agent RMI, un after_tool_callback intercepte les résultats execute_sql, enregistre la liste complète des lignes dans l'état de la session et renvoie un petit échantillon ainsi que le nombre total de lignes :

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

Le fait de renvoyer un échantillon réduit les jetons d'entrée sur tous les tours futurs, ce qui diminue les coûts et les temps de réponse. Cela empêche également les résultats de requête volumineux d'encombrer la fenêtre de contexte et de nuire à la capacité du modèle à suivre les instructions.

Assurez-vous que l'invite de votre agent explique explicitement que les outils renvoient un échantillon par conception. Sans ce contexte, le modèle peut être dérouté par des données partielles, ce qui gaspille des jetons de calcul et de raisonnement en essayant de rapprocher les enregistrements manquants ou en lançant des appels d'outils redondants pour récupérer le reste.

Exemple : Gérer un grand ensemble de résultats

  1. L'utilisateur demande "Quelles sont les routes qui ont des taux de retard supérieurs à 2 en ce moment ?"
  2. La requête de l'agent par rapport à historical_travel_time renvoie des milliers de lignes.
  3. Le after_tool_callback enregistre le résultat complet dans l'état de la session et ne renvoie qu'un petit échantillon ainsi que le nombre total de lignes.
  4. Le modèle raisonne sur cet échantillon pour composer son résumé.
  5. L'application ou les outils en aval lisent l'ensemble de données complet directement à partir de l'état de la session pour afficher le tableau complet ou l'exporter pour l'utilisateur.
Interception d'outil d'agent RMI

Points à retenir

  • Évitez les charges utiles volumineuses dans les invites : stockez de grands ensembles de données dans l’état de la session plutôt que de les transmettre dans l’historique des invites.
  • Tronquez les sorties d'outils volumineuses : utilisez after_tool_callback pour renvoyer de petits échantillons au modèle tout en enregistrant les résultats complets dans l'état de la session.
  • Informez le modèle sur les échantillons : ajoutez une règle d'invite expliquant que les résultats sont échantillonnés pour éviter toute confusion du modèle, tout gaspillage de jetons de raisonnement et tout appel d'outils redondant.

Étapes suivantes

  • Référence des rappels ADK : consultez le guide officiel des rappels ADK pour en savoir plus sur l'interception au niveau du modèle, de l'agent et de l'outil.

Contributeurs

Nathaniel Thomas | Stagiaire en génie logiciel, Google Maps Platform