src/contextualizer.py

"""
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