feat: add ColPali visual retrieval + contextual retrieval + PDF eval

- src/contextualizer.py: Contextual Retrieval with Gemini Flash, retry on rate limit
- ingest.py: --contextual flag for contextual ingest
- src/colpali_retriever.py: ColPali v1.2 visual retrieval with Qdrant MaxSim
- ingest_colpali.py: PDF pages -> images -> ColPali embeddings pipeline
- compare_retrieval.py: side-by-side comparison ColPali vs text RAG
- evaluate.py: accept golden set path as CLI arg
- data/golden_set_pdf.json: PDF financial report evaluation set
- eval_results.json: baseline results (Faithfulness 0.633 / Relevancy 0.767 / Recall 0.400)

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

.env.example

@@ -2,5 +2,6 @@
 # File .env KHÔNG được commit lên Git
 
 GOOGLE_API_KEY=your_google_api_key_here
-# Hoặc nếu dùng OpenAI
-OPENAI_API_KEY=your_openai_api_key_here
+GROQ_API_KEY=your_groq_api_key_here
+QDRANT_URL=https://xxx.us-east4.gcp.cloud.qdrant.io
+QDRANT_API_KEY=your_qdrant_api_key_here

compare_retrieval.py

@@ -0,0 +1,70 @@
+"""
+So sánh ColPali (visual) vs Text RAG (pypdf + Gemini embed) trên cùng câu hỏi.
+Usage: python compare_retrieval.py
+"""
+from dotenv import load_dotenv
+load_dotenv()
+
+TEST_QUESTIONS = [
+    "Doanh thu thuần của công ty trong quý 1 năm 2026 là bao nhiêu?",
+    "Lợi nhuận sau thuế quý 1 năm 2026?",
+    "Chi phí tài chính của FPT trong nửa đầu năm 2016?",
+    "Tổng tài sản của TIG tại cuối quý 1 năm 2026?",
+    "FPT hoạt động trong những lĩnh vực kinh doanh nào?",
+]
+
+
+def run_text_retrieval(question: str, top_k: int = 3) -> list[dict]:
+    from src.rag import retrieve
+    from src.embedder import embed_query
+    contexts, sources = retrieve(question, top_k=top_k)
+    return sources
+
+
+def run_colpali_retrieval(question: str, top_k: int = 3) -> list[dict]:
+    from src.colpali_retriever import query
+    return query(question, top_k=top_k)
+
+
+def main():
+    print("=" * 60)
+    print("RETRIEVAL COMPARISON: ColPali vs Text RAG")
+    print("=" * 60)
+
+    results = []
+    for q in TEST_QUESTIONS:
+        print(f"\nQ: {q}")
+        print("-" * 50)
+
+        print("[Text RAG]")
+        try:
+            text_hits = run_text_retrieval(q, top_k=3)
+            for h in text_hits:
+                print(f"  {h['source']} p.{h['page']} | rerank={h['rerank_score']:.3f}")
+        except Exception as e:
+            print(f"  ERROR: {e}")
+            text_hits = []
+
+        print("[ColPali]")
+        try:
+            colpali_hits = run_colpali_retrieval(q, top_k=3)
+            for h in colpali_hits:
+                print(f"  {h['source']} p.{h['page']} | score={h['score']:.3f}")
+        except Exception as e:
+            print(f"  ERROR: {e}")
+            colpali_hits = []
+
+        results.append({
+            "question": q,
+            "text_rag": text_hits,
+            "colpali": colpali_hits,
+        })
+
+    import json
+    with open("comparison_results.json", "w", encoding="utf-8") as f:
+        json.dump(results, f, ensure_ascii=False, indent=2)
+    print("\n\nResults saved to comparison_results.json")
+
+
+if __name__ == "__main__":
+    main()

data/golden_set_pdf.json

@@ -0,0 +1,86 @@
+[
+  {
+    "id": "p001",
+    "type": "lookup",
+    "source": "000000016333831_CBTT_BCTC_Q1_2026_signed.pdf",
+    "question": "Doanh thu thuần của công ty trong quý 1 năm 2026 là bao nhiêu?",
+    "answer": "Doanh thu thuần quý 1/2026 được ghi nhận trong báo cáo kết quả hoạt động kinh doanh của công ty."
+  },
+  {
+    "id": "p002",
+    "type": "lookup",
+    "source": "000000016333831_CBTT_BCTC_Q1_2026_signed.pdf",
+    "question": "Lợi nhuận sau thuế của công ty quý 1 năm 2026 là bao nhiêu?",
+    "answer": "Lợi nhuận sau thuế quý 1/2026 được thể hiện trong báo cáo kết quả kinh doanh."
+  },
+  {
+    "id": "p003",
+    "type": "lookup",
+    "source": "FPT_16Q2_BCTC_MSX.pdf",
+    "question": "Doanh thu của FPT trong quý 2 năm 2016 là bao nhiêu?",
+    "answer": "Doanh thu FPT quý 2/2016 được ghi nhận trong báo cáo tài chính giữa niên độ."
+  },
+  {
+    "id": "p004",
+    "type": "lookup",
+    "source": "FPT_16Q2_BCTC_MSX.pdf",
+    "question": "Chi phí tài chính của FPT trong nửa đầu năm 2016 là bao nhiêu?",
+    "answer": "Chi phí tài chính được thể hiện trong báo cáo kết quả hoạt động kinh doanh của FPT."
+  },
+  {
+    "id": "p005",
+    "type": "lookup",
+    "source": "UPCOM_quy1.pdf",
+    "question": "Báo cáo UPCOM quý 1 trình bày thông tin về những công ty nào?",
+    "answer": "Báo cáo UPCOM quý 1 trình bày thông tin giao dịch và tài chính của các công ty niêm yết trên sàn UPCOM."
+  },
+  {
+    "id": "p006",
+    "type": "general",
+    "source": "FPT_16Q2_BCTC_MSX.pdf",
+    "question": "FPT hoạt động trong những lĩnh vực kinh doanh nào?",
+    "answer": "FPT hoạt động trong các lĩnh vực công nghệ thông tin, viễn thông, giáo dục và các dịch vụ liên quan."
+  },
+  {
+    "id": "p007",
+    "type": "general",
+    "source": "000000016402354_EN_Disclosure_Financial_Statements_Q1.2026_TIG.pdf",
+    "question": "Tổng tài sản của TIG tại thời điểm cuối quý 1 năm 2026 là bao nhiêu?",
+    "answer": "Tổng tài sản của TIG được ghi nhận trên bảng cân đối kế toán tại ngày 31/03/2026."
+  },
+  {
+    "id": "p008",
+    "type": "comparison",
+    "source": "FPT_16Q2_BCTC_MSX.pdf",
+    "question": "Doanh thu của FPT quý 2 năm 2016 so với cùng kỳ năm trước tăng hay giảm?",
+    "answer": "Sự biến động doanh thu được thể hiện qua cột so sánh cùng kỳ trong báo cáo tài chính."
+  },
+  {
+    "id": "p009",
+    "type": "lookup",
+    "source": "000000016402354_EN_Disclosure_Financial_Statements_Q1.2026_TIG.pdf",
+    "question": "Vốn chủ sở hữu của TIG tại quý 1 năm 2026 là bao nhiêu?",
+    "answer": "Vốn chủ sở hữu của TIG được trình bày trong phần nguồn vốn của bảng cân đối kế toán."
+  },
+  {
+    "id": "p010",
+    "type": "filter",
+    "source": "FPT_16Q2_BCTC_MSX.pdf",
+    "question": "Trong báo cáo FPT quý 2 năm 2016, khoản mục nào có giá trị lớn nhất trong tài sản ngắn hạn?",
+    "answer": "Khoản mục có giá trị lớn nhất trong tài sản ngắn hạn được thể hiện trên bảng cân đối kế toán của FPT."
+  },
+  {
+    "id": "p011",
+    "type": "out_of_scope",
+    "source": "",
+    "question": "Giá cổ phiếu FPT hôm nay là bao nhiêu?",
+    "answer": ""
+  },
+  {
+    "id": "p012",
+    "type": "out_of_scope",
+    "source": "",
+    "question": "Dự báo tăng trưởng GDP Việt Nam năm 2025 là bao nhiêu?",
+    "answer": ""
+  }
+]

eval_results.json

@@ -0,0 +1,57 @@
+{
+  "summary": {
+    "faithfulness": 0.6333333333333334,
+    "answer_relevancy": 0.7666666666666667,
+    "context_recall": 0.4000000000000001
+  },
+  "details": [
+    {
+      "question": "Doanh thu thuần của công ty trong quý 1 năm 2026 là bao nhiêu?",
+      "ground_truth": "Doanh thu thuần quý 1/2026 được ghi nhận trong báo cáo kết quả hoạt động kinh doanh của công ty.",
+      "answer": "Doanh thu thuần của công ty trong quý 1 năm 2026 là 63,904,435,716 VND [1].",
+      "faithfulness": 0.9,
+      "answer_relevancy": 1.0,
+      "context_recall": 0.8
+    },
+    {
+      "question": "Lợi nhuận sau thuế của công ty quý 1 năm 2026 là bao nhiêu?",
+      "ground_truth": "Lợi nhuận sau thuế quý 1/2026 được thể hiện trong báo cáo kết quả kinh doanh.",
+      "answer": "Không tìm thấy thông tin về lợi nhuận sau thuế của công ty quý 1 năm 2026 trong các tài liệu cung cấp. Tuy nhiên, có thông tin về lợi nhuận sau thuế thu nhập doanh nghiệp tại báo cáo kết quả kinh doanh trong một số tài liệu, nhưng không rõ ràng về quý 1 năm 2026 [1], [2], [4].",
+      "faithfulness": 0.2,
+      "answer_relevancy": 0.8,
+      "context_recall": 0.0
+    },
+    {
+      "question": "Doanh thu của FPT trong quý 2 năm 2016 là bao nhiêu?",
+      "ground_truth": "Doanh thu FPT quý 2/2016 được ghi nhận trong báo cáo tài chính giữa niên độ.",
+      "answer": "Không tìm thấy thông tin về doanh thu của FPT trong quý 2 năm 2016 trong các tài liệu cung cấp. Tuy nhiên, trong tài liệu [4], có thông tin về doanh thu bán hàng và cung cấp dịch vụ của FPT trong kỳ hoạt động từ ngày 01/01/2016 đến ngày 30/6/2016 là 121.602.490.344 VND, nhưng không có thông tin cụ thể về doanh thu trong quý 2 năm 2016.",
+      "faithfulness": 0.5,
+      "answer_relevancy": 0.8,
+      "context_recall": 0.8
+    },
+    {
+      "question": "Chi phí tài chính của FPT trong nửa đầu năm 2016 là bao nhiêu?",
+      "ground_truth": "Chi phí tài chính được thể hiện trong báo cáo kết quả hoạt động kinh doanh của FPT.",
+      "answer": "Không tìm thấy thông tin cụ thể về chi phí tài chính của FPT trong nửa đầu năm 2016 trong các đoạn tài liệu cung cấp. Tuy nhiên, các thông tin về doanh thu, lợi nhuận và các chỉ tiêu tài chính khác của FPT trong năm 2016 và сравнение với năm 2015 được đề cập trong các tài liệu [1], [2], [3], [4] và [5]. Nếu bạn cần thông tin cụ thể về chi phí tài chính, có thể cần tham khảo thêm các tài liệu khác hoặc phần cụ thể hơn trong báo cáo tài chính của FPT.",
+      "faithfulness": 0.5,
+      "answer_relevancy": 0.2,
+      "context_recall": 0.0
+    },
+    {
+      "question": "Báo cáo UPCOM quý 1 trình bày thông tin về những công ty nào?",
+      "ground_truth": "Báo cáo UPCOM quý 1 trình bày thông tin giao dịch và tài chính của các công ty niêm yết trên sàn UPCOM.",
+      "answer": "Báo cáo UPCOM quý 1 trình bày thông tin về các công ty niêm yết trên sàn UPCOM. [Không có thông tin cụ thể trong phần NGỮ CẢNH về các công ty cụ thể, tuy nhiên có thể suy luận thông tin chung về sàn UPCOM]",
+      "faithfulness": 0.8,
+      "answer_relevancy": 0.8,
+      "context_recall": 0.0
+    },
+    {
+      "question": "FPT hoạt động trong những lĩnh vực kinh doanh nào?",
+      "ground_truth": "FPT hoạt động trong các lĩnh vực công nghệ thông tin, viễn thông, giáo dục và các dịch vụ liên quan.",
+      "answer": "FPT hoạt động trong lĩnh vực cung cấp dịch vụ và giải pháp trong lĩnh vực công nghệ thông tin, bao gồm tích hợp hệ thống, cung cấp dịch vụ phần mềm, và các dịch vụ công nghệ khác [1].",
+      "faithfulness": 0.9,
+      "answer_relevancy": 1.0,
+      "context_recall": 0.8
+    }
+  ]
+}

evaluate.py

@@ -138,4 +138,7 @@ def evaluate(golden_set_path: str = "data/golden_set.json", sample: int = 20):
 
 
 if __name__ == "__main__":
-    evaluate(sample=20)
+    import sys
+    golden_path = sys.argv[1] if len(sys.argv) > 1 else "data/golden_set.json"
+    sample = int(sys.argv[2]) if len(sys.argv) > 2 else 20
+    evaluate(golden_set_path=golden_path, sample=sample)

ingest.py

@@ -4,23 +4,34 @@ PDF_FOLDER = "data/pdf_text"
 IMAGE_FOLDER = "data/images"
 
 if __name__ == "__main__":
-    mode = sys.argv[1] if len(sys.argv) > 1 else "all"
+    args = sys.argv[1:]
+    mode = args[0] if args else "all"
+    contextual = "--contextual" in args
 
     if mode == "all":
         import subprocess
-        subprocess.run([sys.executable, __file__, "pdf"], check=True)
-        subprocess.run([sys.executable, __file__, "images"], check=True)
+        extra = ["--contextual"] if contextual else []
+        subprocess.run([sys.executable, __file__, "pdf"] + extra, check=True)
+        subprocess.run([sys.executable, __file__, "images"] + extra, check=True)
 
     elif mode == "pdf":
         from src.parsers.router import parse_folder as parse_folder_pdf
         from src.chunker import chunk_pages
         from src.vector_store import add_chunks
 
-        print("=== INGEST PDF (pypdf - CPU) ===")
+        label = "CONTEXTUAL" if contextual else "STANDARD"
+        print(f"=== INGEST PDF ({label}) ===")
         pdf_pages = parse_folder_pdf(PDF_FOLDER, extensions={".pdf"})
         print(f"Tổng: {len(pdf_pages)} đoạn từ PDF")
         chunks = chunk_pages(pdf_pages)
         print(f"Tổng: {len(chunks)} chunks")
+
+        if contextual and chunks:
+            from src.contextualizer import add_context_to_chunks
+            print("Đang sinh context cho từng chunk...")
+            chunks = add_context_to_chunks(chunks, pdf_pages)
+            print("Context xong!")
+
         if chunks:
             add_chunks(chunks)
         print("PDF xong!\n")
@@ -30,11 +41,19 @@ if __name__ == "__main__":
         from src.chunker import chunk_pages
         from src.vector_store import add_chunks
 
-        print("=== INGEST ẢNH (EasyOCR - GPU) ===")
+        label = "CONTEXTUAL" if contextual else "STANDARD"
+        print(f"=== INGEST ẢNH ({label}) ===")
         image_pages = parse_folder_img(IMAGE_FOLDER, extensions={".jpg", ".jpeg", ".png"})
         print(f"Tổng: {len(image_pages)} đoạn từ ảnh")
         chunks = chunk_pages(image_pages)
         print(f"Tổng: {len(chunks)} chunks")
+
+        if contextual and chunks:
+            from src.contextualizer import add_context_to_chunks
+            print("Đang sinh context cho từng chunk...")
+            chunks = add_context_to_chunks(chunks, image_pages)
+            print("Context xong!")
+
         if chunks:
             add_chunks(chunks)
-        print("Ảnh xong!")
+        print("Ảnh xong!")

ingest_colpali.py

@@ -0,0 +1,74 @@
+"""
+Ingest PDF pages as images into ColPali collection.
+Usage: python ingest_colpali.py [pdf_folder]
+"""
+import sys
+import os
+from pathlib import Path
+from PIL import Image
+
+PDF_FOLDER = sys.argv[1] if len(sys.argv) > 1 else "data/pdf_text"
+IMAGE_OUT = "data/colpali_pages"
+
+
+def pdf_to_images(pdf_path: Path, dpi: int = 150) -> list[Image.Image]:
+    """Convert PDF pages to PIL Images."""
+    try:
+        from pdf2image import convert_from_path
+        return convert_from_path(str(pdf_path), dpi=dpi)
+    except Exception as e:
+        print(f"  pdf2image failed ({e}), trying pypdfium2...")
+        import pypdfium2 as pdfium
+        doc = pdfium.PdfDocument(str(pdf_path))
+        images = []
+        for i in range(len(doc)):
+            page = doc[i]
+            bitmap = page.render(scale=dpi / 72)
+            images.append(bitmap.to_pil())
+        return images
+
+
+def main():
+    from src.colpali_retriever import ingest_images
+
+    os.makedirs(IMAGE_OUT, exist_ok=True)
+    pdf_files = list(Path(PDF_FOLDER).glob("*.pdf"))
+
+    if not pdf_files:
+        print(f"No PDFs found in {PDF_FOLDER}")
+        return
+
+    print(f"Found {len(pdf_files)} PDFs")
+
+    all_images, all_meta = [], []
+    for pdf_path in pdf_files:
+        print(f"\nConverting {pdf_path.name}...")
+        try:
+            pages = pdf_to_images(pdf_path)
+            print(f"  {len(pages)} pages")
+            for page_num, img in enumerate(pages):
+                # Save image for reference/display
+                img_filename = f"{pdf_path.stem}_page{page_num+1}.jpg"
+                img_path = os.path.join(IMAGE_OUT, img_filename)
+                img.save(img_path, "JPEG", quality=85)
+
+                all_images.append(img)
+                all_meta.append({
+                    "source": pdf_path.name,
+                    "page": page_num + 1,
+                    "image_path": img_path,
+                })
+        except Exception as e:
+            print(f"  ERROR: {e}")
+
+    if not all_images:
+        print("No images to ingest.")
+        return
+
+    print(f"\nTotal: {len(all_images)} pages to ingest")
+    ingest_images(all_images, all_meta)
+    print("\nColPali ingest complete!")
+
+
+if __name__ == "__main__":
+    main()

src/colpali_retriever.py

@@ -0,0 +1,139 @@
+"""
+ColPali Visual Retrieval — truy vấn tài liệu bằng hình ảnh thay vì text.
+Model: vidore/colpali-v1.2 (multi-vector, MaxSim scoring)
+So sánh với pipeline text hiện tại để thấy sự khác biệt trên tài liệu có bảng biểu.
+"""
+import os
+import torch
+from pathlib import Path
+from PIL import Image
+from colpali_engine.models import ColPali, ColPaliProcessor
+from qdrant_client import QdrantClient
+from qdrant_client.models import (
+    Distance, VectorParams, MultiVectorConfig,
+    MultiVectorComparator, PointStruct, NamedSparseVector,
+)
+from dotenv import load_dotenv
+
+load_dotenv()
+
+COLPALI_MODEL = "vidore/colpali-v1.2"
+COLLECTION_NAME = "colpali_docs"
+VECTOR_DIM = 128  # ColPali output dim per patch
+
+_model = None
+_processor = None
+
+
+def _get_model():
+    global _model, _processor
+    if _model is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"Loading ColPali on {device}...")
+        _model = ColPali.from_pretrained(
+            COLPALI_MODEL,
+            torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
+            device_map=device,
+        ).eval()
+        _processor = ColPaliProcessor.from_pretrained(COLPALI_MODEL)
+        print("ColPali loaded.")
+    return _model, _processor
+
+
+def _get_client() -> QdrantClient:
+    url = os.getenv("QDRANT_URL")
+    api_key = os.getenv("QDRANT_API_KEY")
+    if url and api_key:
+        return QdrantClient(url=url, api_key=api_key, timeout=60)
+    return QdrantClient(path="qdrant_storage")
+
+
+def _get_or_create_collection(client: QdrantClient):
+    existing = [c.name for c in client.get_collections().collections]
+    if COLLECTION_NAME not in existing:
+        client.create_collection(
+            collection_name=COLLECTION_NAME,
+            vectors_config={
+                "colpali": VectorParams(
+                    size=VECTOR_DIM,
+                    distance=Distance.COSINE,
+                    multivector_config=MultiVectorConfig(
+                        comparator=MultiVectorComparator.MAX_SIM,
+                    ),
+                )
+            },
+        )
+    return client.get_collection(COLLECTION_NAME)
+
+
+def encode_images(images: list[Image.Image]) -> list[list[list[float]]]:
+    """Encode list ảnh → list multi-vector embeddings."""
+    model, processor = _get_model()
+    device = next(model.parameters()).device
+
+    batch = processor.process_images(images).to(device)
+    with torch.no_grad():
+        embeddings = model(**batch)  # (B, num_patches, dim)
+
+    return embeddings.cpu().float().tolist()
+
+
+def encode_query(query_text: str) -> list[list[float]]:
+    """Encode text query → multi-vector embedding."""
+    model, processor = _get_model()
+    device = next(model.parameters()).device
+
+    batch = processor.process_queries([query_text]).to(device)
+    with torch.no_grad():
+        embedding = model(**batch)  # (1, num_tokens, dim)
+
+    return embedding[0].cpu().float().tolist()
+
+
+def ingest_images(
+    images: list[Image.Image],
+    metadatas: list[dict],
+    start_id: int = 0,
+):
+    """Ingest ảnh PDF pages vào Qdrant ColPali collection."""
+    client = _get_client()
+    _get_or_create_collection(client)
+
+    print(f"Encoding {len(images)} pages with ColPali...")
+    embeddings = encode_images(images)
+
+    points = [
+        PointStruct(
+            id=start_id + i,
+            vector={"colpali": emb},
+            payload=meta,
+        )
+        for i, (emb, meta) in enumerate(zip(embeddings, metadatas))
+    ]
+
+    client.upsert(collection_name=COLLECTION_NAME, points=points)
+    print(f"Ingested {len(points)} pages into '{COLLECTION_NAME}'")
+
+
+def query(query_text: str, top_k: int = 3) -> list[dict]:
+    """Tìm kiếm ảnh tương tự với query text qua MaxSim."""
+    client = _get_client()
+    query_emb = encode_query(query_text)
+
+    results = client.query_points(
+        collection_name=COLLECTION_NAME,
+        query=query_emb,
+        using="colpali",
+        limit=top_k,
+        with_payload=True,
+    ).points
+
+    return [
+        {
+            "score": round(r.score, 4),
+            "source": r.payload.get("source", ""),
+            "page": r.payload.get("page", 0),
+            "image_path": r.payload.get("image_path", ""),
+        }
+        for r in results
+    ]

src/contextualizer.py

@@ -0,0 +1,110 @@
+"""
+Contextual Retrieval — sinh 1-2 câu context cho mỗi chunk dựa trên toàn văn tài liệu.
+Context được prepend vào chunk trước khi embed + BM25, giúp retrieval chính xác hơn.
+Tham khảo: Anthropic Contextual Retrieval paper (~35% giảm retrieval failure).
+"""
+import os
+import google.generativeai as genai
+from dotenv import load_dotenv
+
+load_dotenv()
+genai.configure(api_key=os.environ["GOOGLE_API_KEY"].strip())
+_model = genai.GenerativeModel("gemini-2.0-flash")
+
+CONTEXT_MODEL = "gemini-2.0-flash"
+
+CONTEXT_PROMPT = """Dưới đây là toàn bộ tài liệu:
+<document>
+{doc_text}
+</document>
+
+Đây là đoạn văn cần đặt vào ngữ cảnh:
+<chunk>
+{chunk_text}
+</chunk>
+
+Hãy viết 1-2 câu ngắn gọn mô tả vị trí và vai trò của đoạn này trong tài liệu, \
+giúp người đọc hiểu đoạn này nằm ở đâu và nói về điều gì.
+Chỉ trả về 1-2 câu đó, không giải thích thêm."""
+
+
+def generate_context(doc_text: str, chunk_text: str) -> str:
+    """Sinh context cho 1 chunk. Trả về context string."""
+    # Giới hạn doc_text để tránh vượt context window
+    max_doc_chars = 8000
+    if len(doc_text) > max_doc_chars:
+        doc_text = doc_text[:max_doc_chars] + "\n...[nội dung tiếp theo]..."
+
+    import time
+    prompt = CONTEXT_PROMPT.format(doc_text=doc_text, chunk_text=chunk_text)
+    for attempt in range(5):
+        try:
+            response = _model.generate_content(
+                prompt,
+                generation_config={"max_output_tokens": 150, "temperature": 0.1},
+            )
+            return response.text.strip()
+        except Exception as e:
+            msg = str(e)
+            if "429" in msg:
+                # TPD (per day) exhausted — no point retrying
+                if "PerDay" in msg or "limit: 0" in msg:
+                    print(f"  [contextualizer] hết quota ngày, dừng sinh context")
+                    raise RuntimeError("QUOTA_DAY_EXHAUSTED") from e
+                # TPM (per minute) — wait and retry
+                wait = 60 * (attempt + 1)
+                print(f"  [contextualizer] rate limit, chờ {wait}s rồi thử lại...")
+                time.sleep(wait)
+            else:
+                print(f"  [contextualizer] lỗi: {e}, bỏ qua context")
+                return ""
+    return ""
+
+
+def add_context_to_chunks(chunks: list[dict], source_pages: list[dict]) -> list[dict]:
+    """
+    Prepend context vào text của mỗi chunk.
+    source_pages: list các page dict có 'text' và 'metadata.source' để build doc_text.
+    """
+    # Group pages theo source file
+    from collections import defaultdict
+    doc_texts = defaultdict(list)
+    for page in source_pages:
+        src = page["metadata"].get("source", "unknown")
+        doc_texts[src].append(page["text"])
+
+    # Full text per document
+    full_docs = {src: "\n\n".join(pages) for src, pages in doc_texts.items()}
+
+    import time
+    enriched = []
+    total = len(chunks)
+    for i, chunk in enumerate(chunks):
+        src = chunk["metadata"].get("source", "unknown")
+        doc_text = full_docs.get(src, "")
+
+        print(f"  Generating context {i+1}/{total}: {src}...")
+        try:
+            context = generate_context(doc_text, chunk["text"])
+        except RuntimeError as e:
+            if "QUOTA_DAY_EXHAUSTED" in str(e):
+                print(f"  [contextualizer] quota ngày hết tại chunk {i+1}/{total}, dừng lại")
+                # Append remaining chunks without context
+                for remaining in chunks[i:]:
+                    r = dict(remaining)
+                    r["metadata"] = {**remaining["metadata"], "has_context": False}
+                    enriched.append(r)
+                return enriched
+            context = ""
+        time.sleep(4)  # Gemini Flash free tier: 15 RPM
+
+        new_chunk = dict(chunk)
+        if context:
+            new_chunk["text"] = f"{context}\n\n{chunk['text']}"
+            new_chunk["metadata"] = {**chunk["metadata"], "has_context": True}
+        else:
+            new_chunk["metadata"] = {**chunk["metadata"], "has_context": False}
+
+        enriched.append(new_chunk)
+
+    return enriched