Расширенный инструмент перехвата данных для агентов RMI ADK

Примечание: Данное руководство основано на эталонной архитектуре ИИ-агента Roads Management Insights и его примере исходного кода на GitHub . Оно опирается на опыт, полученный при работе с этим агентом, но описанные методы перехвата применимы к любому агенту из комплекта разработки агентов (ADK), инструменты которого создают или обрабатывают большие объемы данных.

Обзор

При работе с большими наборами данных прямая передача необработанных данных через модель вызывает проблемы: замедляет ответы, увеличивает стоимость токенов и приводит к ошибкам рассуждений. Данные RMI могут содержать миллионы строк, поэтому предоставление агенту возможности запрашивать неограниченное количество данных и обрабатывать их все для анализа обходится дорого с точки зрения токенов и раздувает контекстное окно, ухудшая производительность агента на каждом последующем ходу.

Например, запрос к таблице RMI historical_travel_time может вернуть тысячи строк с данными о дорожных участках. Возвращение всего набора результатов в историю переписки замедляет каждый последующий поворот, увеличивает стоимость и отвлекает модель от синтеза окончательного ответа.

Одних лишь подсказок недостаточно для решения этих проблем. Единственный надежный подход — перехватывать вызовы инструментов и направлять большие объемы данных за пределы контекстного окна модели.

Обратные вызовы инструментов ADK

ADK предоставляет обработчики обратного вызова для перехвата выполнения инструментов. Для обработки больших наборов результатов используйте after_tool_callback :

• before_tool_callback(tool, args, tool_context) : Выполняется после того, как модель выбирает инструмент, но до выполнения самого инструмента.
• after_tool_callback(tool, args, tool_context, tool_response) : Выполняется после завершения работы инструмента, но до того, как результат будет отправлен обратно в модель. Используйте его для изменения ответа инструмента .

Функция обратного вызова получает объект tool_context с доступом к state сессии. Это позволяет хранить большие наборы данных в фоновом режиме в памяти между ходами разговора, не помещая их в подсказку модели.

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

Усечение и сохранение данных после выполнения.

Цель : показать модели небольшой пример результатов запроса для принятия решения, сохранив при этом полный набор данных для окончательного вывода.

Для агента RMI функция after_tool_callback перехватывает результаты execute_sql , сохраняет полный список строк в состоянии сессии и возвращает краткий пример вместе с общим количеством строк:

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

Возврат образца сокращает количество входных токенов на всех последующих этапах, снижая затраты и время ответа. Это также предотвращает переполнение контекстного окна большими результатами запроса и ухудшение способности модели следовать инструкциям.

Убедитесь, что в подсказке вашего агента четко указано, что инструменты по своей сути возвращают образец. Без этого контекста модель может запутаться из-за неполных данных, тратя вычислительные ресурсы и токены рассуждений на попытки согласовать отсутствующие записи или запуская избыточные вызовы инструментов для получения оставшейся информации.

Пример: Обработка большого набора результатов.

1. Пользователь спрашивает : «Какие маршруты сейчас имеют коэффициент задержки более 2?»
2. Запрос агента к базе данных historical_travel_time возвращает тысячи строк.
3. Функция after_tool_callback сохраняет полный результат в состоянии сессии и возвращает только небольшую выборку плюс общее количество строк.
4. Модель анализирует эти данные для составления своего резюме.
5. Приложение или используемые инструменты считывают весь набор данных непосредственно из состояния сессии для отображения полной таблицы или экспорта для пользователя.

Основные выводы

• Не передавайте большие объемы данных в запросы : храните большие наборы данных в состоянии сессии, а не передавайте их туда-обратно через историю запросов.
• Усечение больших выходных данных инструмента : используйте after_tool_callback для возврата небольших выборок в модель, сохраняя при этом полные результаты в состоянии сессии.
• Сообщите модели о выборках : добавьте правило-подсказку, объясняющее, что результаты выбираются путем выборки, чтобы предотвратить путаницу в модели, нерациональное использование токенов рассуждений и избыточные вызовы инструментов.

Следующие шаги

• Справочник по обратным вызовам ADK : Подробную информацию о перехвате на уровне модели, агента и инструмента см. в официальном руководстве по обратным вызовам ADK .

Авторы

Натаниэль Томас | Стажер-разработчик программного обеспечения, платформа Google Maps