Tính năng chặn công cụ nâng cao cho các tác nhân RMI ADK

Tổng quan

Khi xử lý các tập dữ liệu lớn, việc truyền trực tiếp dữ liệu thô qua mô hình sẽ gây ra các vấn đề: làm chậm phản hồi, tăng chi phí mã thông báo và gây ra lỗi suy luận. Dữ liệu RMI có thể trải rộng trên hàng triệu hàng, vì vậy, việc cho phép tác nhân truy vấn một lượng không giới hạn và tiếp nhận tất cả dữ liệu đó để phân tích sẽ tốn nhiều mã thông báo và làm phình to cửa sổ ngữ cảnh, làm giảm hiệu suất của tác nhân ở mọi lượt tiếp theo.

Ví dụ: một truy vấn trên bảng historical_travel_time RMI có thể trả về hàng nghìn hàng đoạn đường. Việc đưa toàn bộ tập kết quả trở lại nhật ký trò chuyện sẽ làm chậm mọi lượt sau đó, tăng chi phí và khiến mô hình không thể tổng hợp câu trả lời cuối cùng.

Chỉ có hướng dẫn về câu lệnh không thể giải quyết những vấn đề này. Cách duy nhất đáng tin cậy là chặn các lệnh gọi công cụ và định tuyến các tải trọng dữ liệu lớn ra ngoài cửa sổ ngữ cảnh của mô hình.

Lệnh gọi lại công cụ ADK

ADK cung cấp các giao diện kết nối lệnh gọi lại để chặn quá trình thực thi công cụ. Để xử lý các tập hợp kết quả lớn, hãy sử dụng after_tool_callback:

  • before_tool_callback(tool, args, tool_context): Chạy sau khi mô hình chọn một công cụ, nhưng trước khi công cụ thực thi.
  • after_tool_callback(tool, args, tool_context, tool_response): Chạy sau khi công cụ hoàn tất, nhưng trước khi kết quả quay lại mô hình. Sử dụng tính năng này để sửa đổi câu trả lời của công cụ.

Lệnh gọi lại sẽ nhận được một đối tượng tool_context có quyền truy cập vào phiên state. Nhờ đó, bạn có thể lưu trữ các tập dữ liệu lớn trong bộ nhớ nền qua các lượt trò chuyện mà không cần đưa chúng vào câu lệnh của mô hình.

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

Việc cắt bớt và lưu trữ sau khi thực thi

Mục tiêu: Cho mô hình xem một mẫu nhỏ kết quả truy vấn để đưa ra quyết định, đồng thời lưu tập dữ liệu hoàn chỉnh cho đầu ra cuối cùng.

Đối với tác nhân RMI, after_tool_callback sẽ chặn kết quả execute_sql, lưu toàn bộ danh sách hàng ở trạng thái phiên và trả về một mẫu ngắn cùng với tổng số hàng:

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

Việc trả về một mẫu sẽ giảm số lượng mã thông báo đầu vào trong tất cả các lượt tiếp theo, giúp giảm chi phí và thời gian phản hồi. Điều này cũng ngăn kết quả của các truy vấn lớn làm đầy cửa sổ ngữ cảnh và làm giảm khả năng làm theo hướng dẫn của mô hình.

Đảm bảo câu lệnh cho trợ lý ảo giải thích rõ ràng rằng các công cụ trả về mẫu theo thiết kế. Nếu không có ngữ cảnh này, mô hình có thể bị nhầm lẫn bởi dữ liệu một phần, lãng phí các mã thông báo tính toán và suy luận khi cố gắng đối chiếu các bản ghi bị thiếu hoặc khởi chạy các lệnh gọi công cụ dư thừa để tìm nạp phần còn lại.

Ví dụ: Xử lý một tập hợp kết quả lớn

  1. Người dùng hỏi "What are routes that have delay ratios over 2 right now?" (Những tuyến đường nào hiện có tỷ lệ trễ trên 2?)
  2. Truy vấn của tác nhân đối với historical_travel_time trả về hàng nghìn hàng.
  3. after_tool_callback lưu kết quả hoàn chỉnh ở trạng thái phiên và chỉ trả về một mẫu nhỏ cùng với tổng số hàng.
  4. Mô hình sẽ suy luận dựa trên mẫu đó để tạo bản tóm tắt.
  5. Ứng dụng hoặc các công cụ hạ nguồn đọc toàn bộ tập dữ liệu trực tiếp từ trạng thái phiên để hiển thị toàn bộ bảng hoặc xuất cho người dùng.
Chặn công cụ tác nhân RMI

Những điểm chính cần ghi nhớ

  • Không đưa tải trọng lớn vào câu lệnh: Lưu trữ các tập dữ liệu lớn trong trạng thái phiên thay vì truyền qua lại chúng thông qua nhật ký câu lệnh.
  • Cắt bớt đầu ra của công cụ có kích thước lớn: Sử dụng after_tool_callback để trả về các mẫu nhỏ cho mô hình trong khi lưu kết quả đầy đủ ở trạng thái phiên.
  • Thông báo cho mô hình về các mẫu: Thêm một quy tắc câu lệnh giải thích rằng kết quả được lấy mẫu để ngăn mô hình nhầm lẫn, lãng phí mã thông báo suy luận và các lệnh gọi công cụ dư thừa.

Các bước tiếp theo

Người đóng góp

Nathaniel Thomas | Thực tập sinh kỹ thuật phần mềm, Google Maps Platform