Lắp Ráp Ngữ Cảnh RAG và Kiểm Soát Giới Hạn Token (Context Window Budget)

Mục tiêu kiến thức

• Hiểu kiến trúc tạo Prompt cho LLM trong RAG:

Một Prompt hoàn chỉnh gồm:

• Chỉ thị hệ thống (System Instruction).
• Ngữ cảnh tài liệu được nhồi vào (Injected Context), kèm chỉ mục trích dẫn [Doc 1], [Doc 2].
• Câu hỏi của người dùng (User Question).
• Kiểm soát Ngân sách Token (Token Budget):

Context window của LLM có giới hạn (và chi phí tính theo token). Cần thuật toán đóng gói tham lam (Greedy Context Packing) để tối đa hóa số tài liệu hữu ích đưa vào prompt mà tuyệt đối không vượt quá max_tokens.

Yêu cầu

Viết hàm assemble_rag_prompt(system_prompt: str, user_query: str, retrieved_docs: list[dict[str, str]], max_total_tokens: int, token_counter_fn: Any) -> dict[str, Any]:

• retrieved_docs: Danh sách tài liệu đã sắp xếp theo độ liên quan giảm dần, mỗi phần tử có key "id", "text".
• token_counter_fn(text: str) -> int: Hàm đếm số token của một chuỗi văn bản.
• Quy trình lắp ráp:
• Tính số token tối thiểu bắt buộc: base_tokens = token_counter_fn(system_prompt) + token_counter_fn(user_query) + overhead. Nếu base_tokens > max_total_tokens, ném ValueError.
• Ngân sách dành cho Context: budget = max_total_tokens - base_tokens.
• Duyệt qua từng tài liệu trong retrieved_docs:
• Định dạng: f"[Source {i+1}: {doc['id']}]\n{doc['text']}".
• Nếu thêm tài liệu này vào mà tổng số token của toàn bộ prompt không vượt quá max_total_tokens, nạp vào danh sách được chọn.
• Nếu vượt quá, dừng hoặc bỏ qua.
• Lắp ráp Prompt hoàn chỉnh theo mẫu:

     {system_prompt}
     
     Context:
     {formatted_sources_joined_by_double_newline}
     
     Question: {user_query}

• Trả về dictionary:

{"prompt": str, "used_sources": list[str], "total_tokens": int}.

Ràng buộc

• Thời gian chạy tối đa: 2000ms.
• Giới hạn bộ nhớ: 256MB.
• Dữ liệu đầu vào hợp lệ theo đúng kiểu dữ liệu và miền giá trị được mô tả.