Google używa technologii AI do tłumaczenia treści na Twój preferowany język. Tłumaczenia wygenerowane przez AI mogą zawierać błędy.

Zaawansowane przechwytywanie narzędzi w przypadku agentów pakietu ADK RMI

Przegląd

Podczas pracy z dużymi zbiorami danych przekazywanie surowych danych bezpośrednio przez model powoduje problemy: spowalnia odpowiedzi, zwiększa koszty tokenów i wprowadza błędy w rozumowaniu. Dane RMI mogą obejmować miliony wierszy, więc umożliwienie agentowi wykonywania zapytań o nieograniczoną ilość danych i wczytywania ich wszystkich do analizy jest kosztowne pod względem tokenów i zwiększa okno kontekstu, co pogarsza wydajność agenta w każdym kolejnym kroku.

Na przykład zapytanie w tabeli RMI historical_travel_time może zwrócić tysiące wierszy segmentów dróg. Umieszczenie pełnego zestawu wyników w historii rozmowy spowalnia każdy kolejny krok, zwiększa koszty i odwraca uwagę modelu od syntezy ostatecznej odpowiedzi.

Same instrukcje w prompcie nie rozwiążą tych problemów. Jedynym niezawodnym sposobem jest przechwytywanie wywołań narzędzi i kierowanie dużych ładunków danych poza okno kontekstu modelu.

Wywołania zwrotne narzędzia ADK

Pakiet ADK udostępnia haki wywołań zwrotnych do przechwytywania wykonywania narzędzi. Aby obsługiwać duże zestawy wyników, użyj after_tool_callback:

before_tool_callback(tool, args, tool_context): uruchamia się po wybraniu narzędzia przez model, ale przed jego wykonaniem.
after_tool_callback(tool, args, tool_context, tool_response): uruchamia się po zakończeniu działania narzędzia, ale przed przekazaniem wyniku do modelu. Użyj go, aby zmodyfikować odpowiedź narzędzia.

Wywołanie zwrotne otrzymuje obiekt tool_context z dostępem do state sesji. Pozwala to przechowywać duże zbiory danych w pamięci w tle między kolejnymi krokami rozmowy bez umieszczania ich w prompcie modelu.

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

Obcinanie i przechowywanie po wykonaniu

Cel: pokazanie modelowi małej próbki wyników zapytania w celu podjęcia decyzji, a jednocześnie zapisanie pełnego zbioru danych na potrzeby ostatecznego wyniku.

W przypadku agenta RMI after_tool_callback przechwytuje wyniki execute_sql, zapisuje pełną listę wierszy w stanie sesji i zwraca krótką próbkę wraz z łączną liczbą wierszy:

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

Zwracanie próbki zmniejsza liczbę tokenów wejściowych we wszystkich przyszłych krokach, co obniża koszty i skraca czas odpowiedzi. Zapobiega też przepełnianiu okna kontekstu przez duże wyniki zapytań i pogarszaniu zdolności modelu do wykonywania instrukcji.

Upewnij się, że prompt agenta wyraźnie wyjaśnia, że narzędzia zwracają próbkę. Bez tego kontekstu model może się pomylić z powodu niepełnych danych, marnując tokeny obliczeniowe i tokeny rozumowania na próby pogodzenia brakujących rekordów lub uruchamiając nadmiarowe wywołania narzędzi w celu pobrania reszty.

Przykład: obsługa dużego zestawu wyników

Użytkownik pyta: „Jakie trasy mają obecnie współczynniki opóźnienia powyżej 2?”
Zapytanie agenta w tabeli historical_travel_time zwraca tysiące wierszy.
after_tool_callback zapisuje pełny wynik w stanie sesji i zwraca tylko małą próbkę oraz łączną liczbę wierszy.
Model analizuje tę próbkę, aby utworzyć podsumowanie.
Aplikacja lub narzędzia podrzędne odczytują pełny zbiór danych bezpośrednio ze stanu sesji, aby wyrenderować pełną tabelę lub wyeksportować ją dla użytkownika.

Najważniejsze punkty

Nie umieszczaj dużych ładunków w promptach: przechowuj duże zbiory danych w stanie sesji zamiast przekazywać je tam i z powrotem w historii promptów.
Obcinaj duże dane wyjściowe narzędzi: użyj after_tool_callback, aby zwracać modelowi małe próbki, a jednocześnie zapisywać pełne wyniki w stanie sesji.
Informuj model o próbkach: dodaj regułę promptu wyjaśniającą, że wyniki są próbkowane, aby zapobiec pomyłkom modelu, marnowaniu tokenów rozumowania i nadmiarowym wywołaniom narzędzi.

Dalsze kroki

Dokumentacja wywołań zwrotnych ADK: przeczytaj oficjalny przewodnik po wywołaniach zwrotnych ADK, aby dowiedzieć się więcej o przechwytywaniu na poziomie modelu, agenta i narzędzia.

Współtwórcy

Nathaniel Thomas | Stażysta w dziale inżynierii oprogramowania, Google Maps Platform