การสกัดกั้นเครื่องมือขั้นสูงสำหรับเอเจนต์ RMI ADK

ภาพรวม

เมื่อจัดการกับชุดข้อมูลขนาดใหญ่ การส่งข้อมูลดิบผ่านโมเดลโดยตรงจะทำให้เกิดปัญหาต่างๆ เช่น การตอบสนองช้าลง ค่าใช้จ่ายโทเค็นสูงขึ้น และเกิดข้อผิดพลาดในการให้เหตุผล ข้อมูล RMI อาจมีหลายล้านแถว ดังนั้นการอนุญาตให้ Agent ค้นหาข้อมูลจำนวนไม่จำกัดและนำเข้าข้อมูลทั้งหมดเพื่อวิเคราะห์จึงมีค่าใช้จ่ายโทเค็นสูงและทำให้หน้าต่างบริบทใหญ่ขึ้น ซึ่งจะทำให้ประสิทธิภาพของ Agent ลดลงในทุกๆ รอบถัดไป

ตัวอย่างเช่น การค้นหาในตาราง historical_travel_time ของ RMI อาจส่งกลับมาเป็นแถวของส่วนถนนหลายพันแถว การใส่ชุดผลลัพธ์ทั้งหมดกลับไปในประวัติการสนทนาจะทำให้ทุกๆ รอบถัดไปช้าลง เพิ่มค่าใช้จ่าย และทำให้โมเดลไม่สามารถสังเคราะห์คำตอบสุดท้ายได้

คำแนะนำในพรอมต์เพียงอย่างเดียวไม่สามารถแก้ไขปัญหาเหล่านี้ได้ แนวทางที่เชื่อถือได้เพียงอย่างเดียวคือการดักจับการเรียกใช้เครื่องมือ และกำหนดเส้นทางเพย์โหลดข้อมูลขนาดใหญ่ออกไปนอกหน้าต่างบริบทของโมเดล

การเรียกกลับของเครื่องมือ 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 Agent นั้น 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

การส่งกลับตัวอย่างจะลดโทเค็นอินพุตในทุกๆ รอบถัดไป ซึ่งจะช่วยลดค่าใช้จ่ายและเวลาในการตอบสนอง นอกจากนี้ยังป้องกันไม่ให้ผลการค้นหาขนาดใหญ่ทำให้หน้าต่างบริบทเต็มและส่งผลเสียต่อความสามารถของโมเดลในการทำตามคำแนะนำ

ตรวจสอบว่าพรอมต์ของ Agent อธิบายอย่างชัดเจนว่าเครื่องมือจะส่งกลับตัวอย่างตามที่ออกแบบไว้ หากไม่มีบริบทนี้ โมเดลอาจสับสนกับข้อมูลบางส่วน ทำให้เสียโทเค็นการคำนวณและการให้เหตุผลไปกับการพยายามรวมบันทึกที่ขาดหายไปหรือเปิดใช้การเรียกใช้เครื่องมือที่ซ้ำซ้อนเพื่อดึงข้อมูลที่เหลือ

ตัวอย่าง: การจัดการชุดผลลัพธ์ขนาดใหญ่

  1. ผู้ใช้ถามว่า "ตอนนี้มีเส้นทางใดบ้างที่มีอัตราส่วนความล่าช้ามากกว่า 2"
  2. การค้นหาของ Agent ใน historical_travel_time จะส่งกลับมาเป็นแถวหลายพันแถว
  3. after_tool_callback จะบันทึกผลลัพธ์ทั้งหมดไว้ในสถานะเซสชันและส่งกลับมาเพียงตัวอย่างเล็กๆ พร้อมจำนวนแถวทั้งหมด
  4. โมเดลจะให้เหตุผลจากตัวอย่างนั้นเพื่อสร้างสรุป
  5. แอปพลิเคชันหรือเครื่องมือปลายทางจะอ่านชุดข้อมูลทั้งหมดจากสถานะเซสชันโดยตรงเพื่อแสดงตารางที่สมบูรณ์หรือส่งออกให้ผู้ใช้
การสกัดกั้นเครื่องมือ Agent ของ RMI

สรุปประเด็นสำคัญ

  • เก็บเพย์โหลดขนาดใหญ่ออกไปจากพรอมต์: จัดเก็บชุดข้อมูลขนาดใหญ่ไว้ในสถานะเซสชัน แทนที่จะส่งข้อมูลไปมาผ่านประวัติพรอมต์
  • ตัดทอนเอาต์พุตขนาดใหญ่ของเครื่องมือ: ใช้ after_tool_callback เพื่อส่งกลับตัวอย่างเล็กๆ ไปยังโมเดล ขณะเดียวกันก็บันทึกผลลัพธ์ทั้งหมดไว้ในสถานะเซสชัน
  • แจ้งให้โมเดลทราบเกี่ยวกับตัวอย่าง: เพิ่มกฎพรอมต์ที่อธิบายว่า ผลลัพธ์เป็นตัวอย่างเพื่อป้องกันไม่ให้โมเดลสับสน เสียโทเค็นการให้เหตุผล และมีการเรียกใช้เครื่องมือที่ซ้ำซ้อน

ขั้นตอนถัดไป

  • ข้อมูลอ้างอิงการเรียกกลับของ ADK: อ่าน คู่มือการเรียกกลับของ ADKอย่างเป็นทางการเพื่อดูรายละเอียดเกี่ยวกับการดักจับระดับโมเดล, Agent และเครื่องมือ

ผู้ร่วมให้ข้อมูล

Nathaniel Thomas | นักศึกษาฝึกงานด้านวิศวกรรมซอฟต์แวร์, Google Maps Platform