Interceptação avançada de ferramentas para agentes do ADK de RMI

Visão geral

Ao lidar com grandes conjuntos de dados, a transmissão de dados brutos diretamente pelo modelo causa problemas: ela deixa as respostas mais lentas, aumenta os custos de token e introduz erros de raciocínio. Os dados do RMI podem abranger milhões de linhas. Portanto, permitir que o agente consulte uma quantidade ilimitada e ingira tudo para análise é caro em termos de tokens e incha a janela de contexto, degradando a performance do agente em cada turno subsequente.

Por exemplo, uma consulta na tabela historical_travel_time do RMI pode retornar milhares de linhas de segmentos de estradas. Colocar o conjunto de resultados completo de volta no histórico de conversas deixa cada turno posterior mais lento, aumenta o custo e distrai o modelo da síntese da resposta final.

As instruções de comando sozinhas não podem resolver esses problemas. A única abordagem confiável é interceptar chamadas de ferramentas e encaminhar payloads de dados grandes para fora da janela de contexto do modelo.

Callbacks de ferramentas do ADK

O ADK fornece hooks de callback para interceptar a execução de ferramentas. Para lidar com grandes conjuntos de resultados, use after_tool_callback:

  • before_tool_callback(tool, args, tool_context): é executado depois que o modelo seleciona uma ferramenta, mas antes da execução dela.
  • after_tool_callback(tool, args, tool_context, tool_response): é executado depois que a ferramenta termina, mas antes que o resultado retorne ao modelo. Use-o para modificar a resposta da ferramenta.

O callback recebe um objeto tool_context com acesso ao state da sessão. Isso permite armazenar grandes conjuntos de dados na memória em segundo plano em turnos de conversa sem colocá-los no comando do modelo.

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

Truncamento e armazenamento pós-execução

Objetivo: mostrar ao modelo uma pequena amostra dos resultados da consulta para tomada de decisão, salvando o conjunto de dados completo para a saída final.

Para o agente do RMI, um after_tool_callback intercepta os resultados de execute_sql, salva a lista completa de linhas no estado da sessão e retorna uma pequena amostra com a contagem total de linhas:

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

Retornar uma amostra reduz os tokens de entrada em todos os turnos futuros, diminuindo os custos e os tempos de resposta. Ele também impede que grandes resultados de consultas lotem a janela de contexto e prejudiquem a capacidade do modelo de seguir instruções.

Verifique se o comando do agente explica explicitamente que as ferramentas retornam uma amostra por design. Sem esse contexto, o modelo pode ficar confuso com dados parciais, desperdiçando tokens de computação e raciocínio ao tentar reconciliar registros ausentes ou iniciar chamadas de ferramentas redundantes para buscar o restante.

Exemplo: como lidar com um grande conjunto de resultados

  1. O usuário pergunta "Quais são as rotas que têm taxas de atraso acima de 2 agora?"
  2. A consulta do agente em historical_travel_time retorna milhares de linhas.
  3. O after_tool_callback salva o resultado completo no estado da sessão e retorna apenas uma pequena amostra mais a contagem total de linhas.
  4. O modelo raciocina sobre essa amostra para compor o resumo.
  5. O aplicativo ou as ferramentas downstream leem o conjunto de dados completo diretamente do estado da sessão para renderizar a tabela completa ou exportar para o usuário.
Interceptação de ferramentas do agente RMI

Pontos principais

  • Mantenha payloads grandes fora dos comandos: armazene grandes conjuntos de dados no estado da sessão em vez de transmiti-los de um lado para o outro no histórico de comandos.
  • Trunque grandes saídas de ferramentas: use after_tool_callback para retornar pequenas amostras ao modelo enquanto salva os resultados completos no estado da sessão.
  • Informe o modelo sobre as amostras: adicione uma regra de comando explicando que os resultados são amostrados para evitar confusão do modelo, tokens de raciocínio desperdiçados e chamadas de ferramentas redundantes.

Próximas etapas

Colaboradores

Nathaniel Thomas | Estagiário de engenharia de software, Plataforma Google Maps