11. Middleware Hooks (beforeModel & afterModel)

Chào mừng bạn đến với Phần 2: Nâng cao. Trong phần này, chúng ta sẽ không chỉ sử dụng các công cụ có sẵn mà sẽ học cách can thiệp sâu vào "hệ thần kinh" của Agent. Bài học này tập trung vào hai "điểm chạm" quan trọng nhất trong vòng lặp ReAct: Trước khi gọi AI và Sau khi AI phản hồi, sử dụng hệ thống tiện ích chuyên nghiệp đã thiết lập từ Bài 1.

I. Node-style Hooks là gì?

Trong LangChain, Node-style hooks là các hàm chạy tuần tự tại các thời điểm cố định trong vòng đời của Agent. Có 2 hook phổ biến nhất:

1. beforeModel: Chạy ngay trước khi Agent gửi dữ liệu ngữ cảnh (messages) cho LLM. Thường dùng để: Ghi log câu hỏi thực tế, kiểm tra quyền truy cập, hoặc đếm số tin nhắn gửi đi.
2. afterModel: Chạy ngay sau khi LLM trả về phản hồi nhưng trước khi Agent thực hiện các bước tiếp theo (như gọi Tool). Thường dùng để: Kiểm tra nội dung đầu ra, đếm token, hoặc ghi log suy luận của AI.

II. Cách tạo Middleware Custom

Bạn sử dụng hàm createMiddleware và định nghĩa các hàm hook bên trong. Mỗi hook nhận vào tham số state, giúp bạn truy cập toàn bộ lịch sử tin nhắn và trạng thái hiện tại của Agent.

III. Mã nguồn: 11-custom-middleware-hooks.ts

Chúng ta sẽ tạo một Middleware tự chế để in ra terminal chính xác những gì đang diễn ra bên trong "não" của AI tại mỗi bước.

// apps/langchain/scripts/11-custom-middleware-hooks.ts
// pnpm --filter=ai-notes-langchain run tsx scripts/11-custom-middleware-hooks.ts

import './env';
import { createGeminiModel, generateImage } from '@workspace/util-langchain';
import { createAgent, createMiddleware } from 'langchain';

// 1. Define a Custom Middleware with Activity Tracking logic
const activityTracker = createMiddleware({
  name: 'ActivityTracker',

  // Runs BEFORE calling the LLM
  beforeModel: (state) => {
    console.log(`\n[BEFORE_MODEL]: Preparing to call AI...`);
    console.log(`> Current context size: ${state.messages.length} messages.`);

    // You can return an object to update state, or void to just observe
    return;
  },

  // Runs AFTER the LLM responds
  afterModel: (state) => {
    const lastMessage = state.messages[state.messages.length - 1];
    console.log(`[AFTER_MODEL]: AI has responded.`);

    // Check if AI wants to use tools or just reply with text
    if (lastMessage.tool_calls?.length) {
      console.log(
        `> AI decision: Calling ${lastMessage.tool_calls.length} tool(s).`,
      );
    } else {
      const preview = lastMessage.content?.toString().slice(0, 50);
      console.log(`> AI Text Response: "${preview}..."`);
    }
    return;
  },
});

async function main() {
  const model = createGeminiModel();

  // 2. Initialize Agent and attach our custom middleware
  const agent = createAgent({
    model: model,
    tools: [], // Simple test without tools
    middleware: [activityTracker],
  });

  // 3. Generate visual diagram
  await generateImage(
    agent.graph,
    'images/courses/langchain-course/scripts-11-custom-hooks.jpg',
  );

  console.log('--- Starting Activity Tracker Test ---');

  // 4. Execute Agent
  await agent.invoke({
    messages: [
      {
        role: 'user',
        content: 'Tell me a very short joke about a programmer.',
      },
    ],
  });

  console.log('\n--- Test Completed ---');
}

main().catch(console.error);

IV. Phân tích kỹ thuật

• Tính tuần tự: Nếu bạn gắn nhiều middleware, các hàm beforeModel sẽ chạy theo thứ tự từ trên xuống dưới trong mảng middleware. Tuy nhiên, afterModel sẽ chạy ngược lại từ dưới lên trên.
• Tham số state: Đây là cửa sổ nhìn vào bộ nhớ của Agent. Bạn có thể đọc state.messages để biết AI đã học được những gì từ các câu hỏi trước (nếu có dùng Checkpointer).
• Ứng dụng thực tế: Bạn có thể dùng beforeModel để tự động chèn thêm câu: "Hãy trả lời bằng tiếng Việt chuyên nghiệp" vào mọi yêu cầu của người dùng mà không cần sửa prompt gốc.

V. Hình ảnh sơ đồ Graph

Middleware được nhúng trực tiếp vào các node xử lý của Agent, bạn có thể quan sát cấu trúc này qua sơ đồ:

VI. Tổng kết

• Hooks giúp bạn xây dựng khả năng quan sát (Observability) cho Agent mà không làm thay đổi logic nghiệp vụ.
• beforeModel dùng cho các tác vụ chuẩn bị và bảo mật đầu vào.
• afterModel dùng cho các tác vụ kiểm tra và thống kê kết quả đầu ra.
• Việc nạp env.ts và dùng createGeminiModel đảm bảo Middleware của bạn luôn chạy trên cấu hình model ổn định nhất.

Trong bài học tiếp theo, chúng ta sẽ học về Wrap Hooks — một kỹ thuật cao cấp hơn để triển khai cơ chế tự động thử lại (Retry) khi AI gặp sự cố API!

👉 Bài tiếp theo: 12. Xử lý lỗi với Wrap Hooks (Retry & Fallback)