Intercettazione avanzata degli strumenti per gli agenti RMI ADK

Panoramica

Quando si lavora con set di dati di grandi dimensioni, il passaggio dei dati non elaborati direttamente attraverso il modello causa problemi: rallenta le risposte, aumenta i costi dei token e introduce errori di ragionamento. I dati RMI possono estendersi su milioni di righe, quindi consentire all'agente di eseguire query su una quantità illimitata e di importare tutto per l'analisi è costoso in termini di token e gonfia la finestra contestuale, peggiorando il rendimento dell'agente a ogni turno successivo.

Ad esempio, una query sulla tabella historical_travel_time RMI può restituire migliaia di righe di segmenti stradali. L'inserimento dell'intero set di risultati nella cronologia delle conversazioni rallenta ogni turno successivo, aumenta i costi e distrae il modello dalla sintesi della risposta finale.

Le istruzioni dei prompt da sole non possono risolvere questi problemi. L'unico approccio affidabile è intercettare le chiamate agli strumenti e instradare i payload di dati di grandi dimensioni al di fuori della finestra contestuale del modello.

Callback degli strumenti ADK

ADK fornisce hook di callback per intercettare l'esecuzione degli strumenti. Per gestire set di risultati di grandi dimensioni, utilizza after_tool_callback:

  • before_tool_callback(tool, args, tool_context): viene eseguito dopo che il modello seleziona uno strumento, ma prima che lo strumento venga eseguito.
  • after_tool_callback(tool, args, tool_context, tool_response): viene eseguito dopo che lo strumento ha terminato, ma prima che il risultato venga restituito al modello. Utilizzalo per modificare la risposta dello strumento.

Il callback riceve un oggetto tool_context con accesso allo state della sessione. In questo modo puoi archiviare set di dati di grandi dimensioni nella memoria di sfondo tra i turni di conversazione senza inserirli nel prompt del modello.

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

Troncamento e memorizzazione post-esecuzione

Obiettivo: mostrare al modello un piccolo campione di risultati delle query per il processo decisionale, salvando al contempo il set di dati completo per l'output finale.

Per l'agente RMI, un after_tool_callback intercetta i risultati di execute_sql, salva l'elenco completo delle righe nello stato della sessione e restituisce un breve campione insieme al numero totale di righe:

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

La restituzione di un campione riduce i token di input in tutti i turni futuri, diminuendo i costi e i tempi di risposta. Inoltre, impedisce che i risultati di query di grandi dimensioni affollino la finestra contestuale e compromettano la capacità del modello di seguire le istruzioni.

Assicurati che il prompt dell'agente spieghi esplicitamente che gli strumenti restituiscono un campione per progettazione. Senza questo contesto, il modello può confondersi con i dati parziali, sprecando token di calcolo e ragionamento nel tentativo di riconciliare i record mancanti o avviando chiamate agli strumenti ridondanti per recuperare il resto.

Esempio: gestione di un set di risultati di grandi dimensioni

  1. L'utente chiede "Quali sono le rotte con rapporti di ritardo superiori a 2 al momento?"
  2. La query dell'agente su historical_travel_time restituisce migliaia di righe.
  3. after_tool_callback salva il risultato completo nello stato della sessione e restituisce solo un piccolo campione più il numero totale di righe.
  4. Il modello ragiona su questo campione per comporre il riepilogo.
  5. L'applicazione o gli strumenti downstream leggono il set di dati completo direttamente dallo stato della sessione per eseguire il rendering della tabella completa o esportarla per l'utente.
RMI Agent Tool Interception

Concetti chiave

  • Non inserire payload di grandi dimensioni nei prompt: archivia i set di dati di grandi dimensioni nello stato della sessione anziché passarli avanti e indietro nella cronologia dei prompt.
  • Tronca gli output di strumenti di grandi dimensioni: utilizza after_tool_callback per restituire piccoli campioni al modello salvando al contempo i risultati completi nello stato della sessione.
  • Informa il modello sui campioni: aggiungi una regola di prompt che spieghi che i risultati vengono campionati per evitare confusione del modello, spreco di token di ragionamento e chiamate agli strumenti ridondanti.

Passaggi successivi

Collaboratori

Nathaniel Thomas | Software Engineering Intern, Google Maps Platform