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

Przegląd

W przypadku dużych zbiorów danych przekazywanie surowych danych bezpośrednio do modelu 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 pozyskiwania ich wszystkich do analizy jest kosztowne pod względem tokenów i zwiększa rozmiar okna kontekstu, co obniża wydajność agenta w każdej kolejnej turze.

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

Same instrukcje w prompcie nie wystarczą, aby rozwiązać te problemy. 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 wywołania zwrotne, które umożliwiają przechwytywanie wykonywania narzędzi. Aby obsługiwać duże zbiory 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.

Funkcja wywołania zwrotnego otrzymuje obiekt tool_context z dostępem do sesji state. Umożliwia to przechowywanie dużych zbiorów danych w pamięci w tle podczas kolejnych tur 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: pokazywanie modelowi małej próbki wyników zapytania na potrzeby podejmowania decyzji, przy jednoczesnym zachowaniu pełnego zbioru danych na potrzeby ostatecznych wyników.

W przypadku agenta RMI after_tool_callback przechwytuje execute_sql wyniki, 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

Zwrócenie próbki zmniejsza liczbę tokenów wejściowych we wszystkich przyszłych etapach, co obniża koszty i skraca czas odpowiedzi. Zapobiega to również przepełnianiu okna kontekstu przez duże wyniki zapytań i utrudnianiu modelowi wykonywania instrukcji.

Upewnij się, że prompt agenta wyraźnie wyjaśnia, że narzędzia z założenia zwracają próbkę. Bez tego kontekstu model może zostać wprowadzony w błąd przez niepełne dane, co spowoduje marnowanie tokenów obliczeniowych i tokenów rozumowania na próby uzupełnienia brakujących rekordów lub uruchamianie zbędnych wywołań narzędzi w celu pobrania pozostałych informacji.

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

  1. Użytkownik pyta: „Jakie trasy mają obecnie współczynnik opóźnienia powyżej 2?”
  2. Zapytanie agenta do historical_travel_time zwraca tysiące wierszy.
  3. after_tool_callback zapisuje pełny wynik w stanie sesji i zwraca tylko małą próbkę oraz łączną liczbę wierszy.
  4. Model analizuje próbkę, aby utworzyć podsumowanie.
  5. 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.
Przechwytywanie narzędzia agenta RMI

Najważniejsze punkty

  • Unikaj umieszczania w promptach dużych ładunków: przechowuj duże zbiory danych w stanie sesji zamiast przesyłać je tam i z powrotem w historii promptów.
  • Obcinanie dużych wyników narzędzi: użyj after_tool_callback, aby zwracać do modelu małe próbki, a pełne wyniki zapisywać w stanie sesji.
  • Informowanie modelu o próbkach: dodaj regułę promptu, która wyjaśnia, że wyniki są próbkowane, aby zapobiec dezorientacji modelu, marnowaniu tokenów rozumowania i zbędnym wywołaniom narzędzi.

Dalsze kroki

Współtwórcy

Nathaniel Thomas | Praktykant w zespole inżynierów oprogramowania, Google Maps Platform