Erweiterte Tool-Abfangfunktion für RMI-ADK-Agents

Hinweis: Dieser Leitfaden basiert auf der Referenzarchitektur des KI-Agenten für Roads Management Insights und dem zugehörigen Beispielquellcode auf GitHub. Er nutzt die Erkenntnisse aus diesem Agenten, aber diese Abfangtechniken gelten für alle Agent Development Kit-Agenten (ADK), deren Tools große Nutzlasten erzeugen oder verarbeiten.

Übersicht

Wenn Sie mit großen Datasets arbeiten, verursacht die direkte Übergabe von Rohdaten an das Modell Probleme: Die Antworten werden langsamer, die Tokenkosten steigen und es treten Fehler bei der Argumentation auf. RMI-Daten können Millionen von Zeilen umfassen. Wenn der Agent eine unbegrenzte Menge abfragen und alle Daten zur Analyse aufnehmen kann, ist das teuer in Bezug auf Token und bläht das Kontextfenster auf, was die Leistung des Agenten bei jeder nachfolgenden Anfrage beeinträchtigt.

Eine Abfrage der RMI-Tabelle historical_travel_time kann beispielsweise Tausende von Zeilen mit Straßenabschnitten zurückgeben. Wenn Sie den vollständigen Ergebnissatz in den Unterhaltungsverlauf einfügen, wird jede spätere Anfrage verlangsamt, die Kosten steigen und das Modell wird davon abgelenkt, die endgültige Antwort zu synthetisieren.

Prompt-Anweisungen allein können diese Probleme nicht lösen. Der einzige zuverlässige Ansatz besteht darin, Toolaufrufe abzufangen und große Daten-Nutzlasten außerhalb des Kontextfensters des Modells weiterzuleiten.

ADK-Tool-Callbacks

ADK bietet Callback-Hooks, um die Toolausführung abzufangen. Verwenden Sie after_tool_callback, um große Ergebnissätze zu verarbeiten:

• before_tool_callback(tool, args, tool_context): Wird ausgeführt, nachdem das Modell ein Tool ausgewählt hat, aber bevor das Tool ausgeführt wird.
• after_tool_callback(tool, args, tool_context, tool_response): Wird ausgeführt, nachdem das Tool beendet wurde, aber bevor das Ergebnis an das Modell zurückgegeben wird. Verwenden Sie es, um die Toolantwort zu ändern.

Der Callback empfängt ein tool_context-Objekt mit Zugriff auf den Sitzungsstatus state. So können Sie große Datasets über mehrere Unterhaltungsrunden hinweg im Hintergrundspeicher speichern, ohne sie in den Prompt des Modells einzufügen.

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

Nach der Ausführung kürzen und speichern

Ziel: Dem Modell eine kleine Stichprobe der Abfrageergebnisse zur Entscheidungsfindung zeigen und gleichzeitig das vollständige Dataset für die endgültige Ausgabe speichern.

Für den RMI-Agenten fängt ein after_tool_callback die execute_sql-Ergebnisse ab, speichert die vollständige Liste der Zeilen im Sitzungsstatus und gibt eine kurze Stichprobe zusammen mit der Gesamtzahl der Zeilen zurück:

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

Durch die Rückgabe einer Stichprobe werden die Eingabetoken bei allen zukünftigen Anfragen reduziert, was Kosten und Antwortzeiten senkt. Außerdem wird verhindert, dass große Abfrageergebnisse das Kontextfenster überlasten und die Fähigkeit des Modells beeinträchtigen, Anweisungen zu befolgen.

Achten Sie darauf, dass in Ihrem Agenten-Prompt ausdrücklich erklärt wird, dass Tools standardmäßig eine Stichprobe zurückgeben. Ohne diesen Kontext kann das Modell durch unvollständige Daten verwirrt werden und Rechen- und Argumentationstoken verschwenden, um fehlende Datensätze abzugleichen oder redundante Toolaufrufe zu starten, um den Rest abzurufen.

Beispiel: Umgang mit einem großen Ergebnissatz

1. Der Nutzer fragt „Welche Routen haben derzeit Verzögerungsraten über 2?“
2. Die Abfrage des Agenten für historical_travel_time gibt Tausende von Zeilen zurück.
3. Der after_tool_callback speichert das vollständige Ergebnis im Sitzungsstatus und gibt nur eine kleine Stichprobe sowie die Gesamtzahl der Zeilen zurück.
4. Das Modell argumentiert anhand dieser Stichprobe, um die Zusammenfassung zu erstellen.
5. Die Anwendung oder nachgelagerte Tools lesen das vollständige Dataset direkt aus dem Sitzungsstatus, um die vollständige Tabelle zu rendern oder für den Nutzer zu exportieren.

Wichtigste Ergebnisse

• Große Nutzlasten aus Prompts entfernen: Speichern Sie große Datasets im Sitzungs status, anstatt sie über den Promptverlauf hin und her zu senden.
• Große Toolausgaben kürzen: Verwenden Sie after_tool_callback, um kleine Stichproben an das Modell zurückzugeben und gleichzeitig vollständige Ergebnisse im Sitzungsstatus zu speichern.
• Das Modell über Stichproben informieren: Fügen Sie eine Prompt-Regel hinzu, in der erklärt wird, dass Ergebnisse als Stichprobe zurückgegeben werden, um Verwirrung des Modells, verschwendete Argumentationstoken und redundante Toolaufrufe zu vermeiden.

Nächste Schritte

• ADK-Callbacks-Referenz: Im offiziellen ADK-Leitfaden zu Callbacks finden Sie Details zum Abfangen auf Modell-, Agenten- und Tool-Ebene.

Beitragende

Nathaniel Thomas | Software Engineering Intern, Google Maps Platform