Thiết Kế Stateful Agent Memory Engine 3 Tầng Trong Node.js

flowchart TD
    A["Tương Tác Mới Của Người Dùng"] --> B["TẦNG 1: Short-term Memory (Sliding Buffer)"]
    B -->|Khi vượt quá 2.000 Tokens| C["TẦNG 2: Mid-term Memory (Semantic Compactor)"]
    C -->|Nén thành Bản tóm tắt ngữ cảnh| D["Cập nhật System Prompt / Rolling Summary"]
    
    A --> E["Trích xuất thực thể & Facts quan trọng"]
    E --> F["TẦNG 3: Long-term Memory (Vector & KV Store)"]
    
    D & F --> G["Prompt Hoàn Chỉnh Gửi Đến LLM Core"]

const { get_encoding } = require("js-tiktoken");
const OpenAI = require("openai");

const tokenizer = get_encoding("cl100k_base");
const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY });

// 1. TẦNG 1: Bộ nhớ đệm trượt (Sliding Buffer Memory)
class SlidingBufferMemory {
  constructor(maxTokenBudget = 2000) {
    this.messages = [];
    this.maxTokenBudget = maxTokenBudget;
    this.rollingSummary = "";
  }

  addMessage(role, content) {
    this.messages.push({ role, content, timestamp: Date.now() });
  }

  calculateTotalTokens() {
    let total = 0;
    for (const msg of this.messages) {
      total += tokenizer.encode(msg.content).length;
    }
    return total;
  }

  // 2. TẦNG 2: Kích hoạt nén ngữ cảnh ngữ nghĩa (Semantic Compaction)
  async compactMemoryIfNeeded() {
    const currentTokens = this.calculateTotalTokens();

    // Nếu dung lượng vượt quá ngân sách cho phép
    if (currentTokens > this.maxTokenBudget && this.messages.length > 4) {
      console.log(`[MEMORY COMPACTION]: Đang nén bộ nhớ (${currentTokens} tokens vượt ngưỡng ${this.maxTokenBudget})...`);
      
      // Lấy 50% tin nhắn cũ nhất để nén
      const halfIndex = Math.floor(this.messages.length / 2);
      const messagesToCompress = this.messages.splice(0, halfIndex);

      const conversationText = messagesToCompress
        .map(m => `${m.role.toUpperCase()}: ${m.content}`)
        .join("
");

      const prompt = `Bạn là Memory Controller. Hãy cập nhật bản tóm tắt ngắn gọn các sự thật và quyết định quan trọng nhất từ đoạn hội thoại sau. Bỏ qua các câu xã giao.
Bản tóm tắt trước đó: "${this.rollingSummary}"
Đoạn hội thoại cần nén:
${conversationText}`;

      const response = await openai.chat.completions.create({
        model: "gpt-4o-mini", // Dùng model nhỏ gọn, tiết kiệm chi phí
        messages: [{ role: "user", content: prompt }]
      });

      this.rollingSummary = response.choices[0].message.content;
      console.log("[MEMORY COMPACTION]: Bản tóm tắt mới:", this.rollingSummary);
    }
  }

  // Tạo mảng tin nhắn hoàn chỉnh để gửi cho LLM
  getPromptMessages() {
    const output = [];
    if (this.rollingSummary) {
      output.push({
        role: "system",
        content: `[NGỮ CẢNH CÁC PHIÊN TRƯỚC]: ${this.rollingSummary}`
      });
    }
    return output.concat(this.messages.map(m => ({ role: m.role, content: m.content })));
  }
}

async function injectLongTermMemory(userQuery, userId, vectorStoreClient) {
  // 1. Tạo vector embedding cho câu hỏi hiện tại
  const embeddingResponse = await openai.embeddings.create({
    model: "text-embedding-3-small",
    input: userQuery
  });
  const queryVector = embeddingResponse.data[0].embedding;

  // 2. Tìm kiếm các sự thật liên quan nhất trong Vector DB
  const searchResults = await vectorStoreClient.query({
    vector: queryVector,
    filter: { userId },
    topK: 3
  });

  if (searchResults.length === 0) return "";

  const retrievedFacts = searchResults
    .map(item => `- ${item.metadata.fact}`)
    .join("
");

  return `[THÔNG TIN DÀI HẠN VỀ NGƯỜI DÙNG]:
${retrievedFacts}
`;
}