pdf_rag.py

"""
PDF RAG — chunk, embed, index to Qdrant, hybrid search.

Hybrid search = dense  (OpenAI text-embedding-3-small, cosine similarity)
              + sparse (BM25 via fastembed Qdrant/bm25, dot product)
              merged via Reciprocal Rank Fusion (RRF, k=60).

Sau RRF, mỗi chunk được mở rộng sang các chunk lân cận (N-3 đến N+3) trong cùng
PDF để đưa vào context đầy đủ hơn. Không dùng overlap vì neighbor expansion
đã đảm bảo không mất context tại ranh giới chunk.
"""

import logging
import uuid
from typing import Optional

from fastembed import SparseTextEmbedding
from openai import OpenAI
from qdrant_client import QdrantClient
from qdrant_client.models import (
    Distance,
    FieldCondition,
    Filter,
    MatchAny,
    MatchValue,
    PointStruct,
    SparseVector,
    SparseVectorParams,
    VectorParams,
)

from src.config import OPENAI_API_KEY, QDRANT_API_KEY, QDRANT_URL
from src.pdf_processing import pdf_to_markdown

logger = logging.getLogger(__name__)

# Collection v2: named vectors (dense + sparse). Xóa collection cũ "pdf_chunks" nếu còn.
_PDF_COLLECTION  = "pdf_chunks_v2"
_EMBED_MODEL     = "text-embedding-3-small"
_EMBED_DIMS      = 1536
_BM25_MODEL      = "Qdrant/bm25"

_CHUNK_SIZE      = 1000  # ký tự / chunk
_CHUNK_OVERLAP   = 0     # không cần overlap — neighbor expansion xử lý ranh giới
_EMBED_BATCH     = 32    # số chunk embed song song mỗi lần

_RRF_K           = 60    # hằng số RRF chuẩn
_NEIGHBOR_WINDOW = 3     # fetch N-3 đến N+3 quanh mỗi chunk được retrieve

_qdrant: Optional[QdrantClient]        = None
_openai: Optional[OpenAI]             = None
_bm25:   Optional[SparseTextEmbedding] = None


# ── Client / model helpers ────────────────────────────────────────────────────

def _get_qdrant() -> QdrantClient:
    global _qdrant
    if _qdrant is None:
        if not QDRANT_URL:
            raise RuntimeError("QDRANT_URL chưa được cấu hình.")
        _qdrant = QdrantClient(url=QDRANT_URL, api_key=QDRANT_API_KEY)
        _ensure_collection(_qdrant)
    return _qdrant


def _get_openai() -> OpenAI:
    global _openai
    if _openai is None:
        if not OPENAI_API_KEY:
            raise RuntimeError("OPENAI_API_KEY chưa được cấu hình.")
        _openai = OpenAI(api_key=OPENAI_API_KEY)
    return _openai


def _get_bm25() -> SparseTextEmbedding:
    global _bm25
    if _bm25 is None:
        _bm25 = SparseTextEmbedding(model_name=_BM25_MODEL)
    return _bm25


def _ensure_collection(client: QdrantClient) -> None:
    existing = {c.name for c in client.get_collections().collections}
    if _PDF_COLLECTION not in existing:
        client.create_collection(
            collection_name=_PDF_COLLECTION,
            vectors_config={
                "dense": VectorParams(size=_EMBED_DIMS, distance=Distance.COSINE),
            },
            sparse_vectors_config={
                "sparse": SparseVectorParams(),
            },
        )
        logger.info("Qdrant: collection '%s' created.", _PDF_COLLECTION)

    # Payload indexes — idempotent, an toàn gọi mỗi lần khởi động.
    for field in ("conversation_id", "pdf_name"):
        client.create_payload_index(
            collection_name=_PDF_COLLECTION,
            field_name=field,
            field_schema="keyword",
        )
    client.create_payload_index(
        collection_name=_PDF_COLLECTION,
        field_name="chunk_index",
        field_schema="integer",
    )


# ── Chunking ──────────────────────────────────────────────────────────────────

def _chunk_text(text: str) -> list[str]:
    if len(text) <= _CHUNK_SIZE:
        return [text.strip()] if text.strip() else []

    chunks: list[str] = []
    start = 0
    while start < len(text):
        end = min(start + _CHUNK_SIZE, len(text))

        if end < len(text):
            for boundary in ('\n\n', '\n', '.', '!', '?'):
                pos = text.rfind(boundary, start + _CHUNK_SIZE // 2, end)
                if pos != -1:
                    end = pos + len(boundary)
                    break

        chunk = text[start:end].strip()
        if chunk:
            chunks.append(chunk)

        # _CHUNK_OVERLAP = 0, nhưng giữ công thức chung để dễ điều chỉnh sau
        next_start = end - _CHUNK_OVERLAP
        if next_start <= start:
            next_start = end
        start = next_start

    return chunks


# ── Embedding ─────────────────────────────────────────────────────────────────

def _embed_batch(texts: list[str]) -> list[list[float]]:
    response = _get_openai().embeddings.create(model=_EMBED_MODEL, input=texts)
    return [item.embedding for item in response.data]


def _embed_one(text: str) -> list[float]:
    return _embed_batch([text])[0]


def _bm25_batch(texts: list[str]) -> list[SparseVector]:
    embeddings = list(_get_bm25().embed(texts))
    return [
        SparseVector(indices=e.indices.tolist(), values=e.values.tolist())
        for e in embeddings
    ]


def _bm25_one(text: str) -> SparseVector:
    return _bm25_batch([text])[0]


# ── Neighbor expansion ────────────────────────────────────────────────────────

def _expand_chunks(
    client: QdrantClient,
    conversation_id: str,
    hits: list[tuple[str, int]],  # (pdf_name, chunk_index)
    window: int = _NEIGHBOR_WINDOW,
) -> list[str]:
    """
    Với mỗi (pdf_name, chunk_index) được retrieve, fetch thêm chunk N-window đến N+window
    từ cùng PDF. Các cửa sổ chồng lấp được merge thành một đoạn liên tục để
    tránh đưa nội dung trùng lặp vào context.

    Returns:
        Danh sách đoạn văn bản, mỗi đoạn là một cửa sổ liên tục (đã merge nếu chồng lấp).
    """
    # Gom tất cả chunk_index cần fetch theo từng pdf_name
    pdf_needed: dict[str, set[int]] = {}
    for pdf_name, chunk_index in hits:
        indices = set(range(max(0, chunk_index - window), chunk_index + window + 1))
        pdf_needed.setdefault(pdf_name, set()).update(indices)

    results: list[str] = []

    for pdf_name, needed in pdf_needed.items():
        fetch_filter = Filter(must=[
            FieldCondition(key="conversation_id", match=MatchValue(value=conversation_id)),
            FieldCondition(key="pdf_name",        match=MatchValue(value=pdf_name)),
            FieldCondition(key="chunk_index",     match=MatchAny(any=sorted(needed))),
        ])
        fetched, _ = client.scroll(
            collection_name=_PDF_COLLECTION,
            scroll_filter=fetch_filter,
            limit=len(needed) + 5,
            with_payload=True,
            with_vectors=False,
        )

        # Map chunk_index → text, rồi sort
        chunk_map = {
            p.payload["chunk_index"]: p.payload.get("chunk_text", "")
            for p in fetched
            if "chunk_index" in p.payload
        }
        sorted_indices = sorted(chunk_map)
        if not sorted_indices:
            continue

        # Gom các chunk_index liên tiếp thành từng run (merge overlapping windows)
        runs: list[list[int]] = []
        current: list[int] = [sorted_indices[0]]
        for idx in sorted_indices[1:]:
            if idx == current[-1] + 1:
                current.append(idx)
            else:
                runs.append(current)
                current = [idx]
        runs.append(current)

        for run in runs:
            text = "\n\n".join(chunk_map[i] for i in run)
            if text.strip():
                results.append(text)

    return results


# ── Public API ────────────────────────────────────────────────────────────────

def index_pdf(pdf_path: str, pdf_name: str, conversation_id: str) -> int:
    """
    Đọc PDF, chunk, embed (dense + sparse) và upsert vào Qdrant.
    UUID v5 làm point ID đảm bảo idempotent — gửi lại cùng file không tạo duplicate.

    Returns:
        Số chunk đã index.
    """
    text   = pdf_to_markdown(pdf_path)
    chunks = _chunk_text(text)
    if not chunks:
        logger.warning("PDF '%s' không có nội dung để index.", pdf_name)
        return 0

    client  = _get_qdrant()
    indexed = 0

    for batch_start in range(0, len(chunks), _EMBED_BATCH):
        batch       = chunks[batch_start : batch_start + _EMBED_BATCH]
        dense_vecs  = _embed_batch(batch)
        sparse_vecs = _bm25_batch(batch)

        points = [
            PointStruct(
                id=str(uuid.uuid5(
                    uuid.NAMESPACE_DNS,
                    f"{conversation_id}::{pdf_name}::{batch_start + i}",
                )),
                vector={
                    "dense":  dense_vecs[i],
                    "sparse": sparse_vecs[i],
                },
                payload={
                    "conversation_id": conversation_id,
                    "pdf_name":        pdf_name,
                    "chunk_index":     batch_start + i,
                    "chunk_text":      batch[i],
                },
            )
            for i in range(len(batch))
        ]

        client.upsert(collection_name=_PDF_COLLECTION, points=points)
        indexed += len(points)

    logger.info(
        "Đã index %d chunks từ '%s' cho conversation '%s'.",
        indexed, pdf_name, conversation_id,
    )
    return indexed


def hybrid_search(query: str, conversation_id: str, top_k: int = 5) -> list[str]:
    """
    Hybrid search:
      Dense  — OpenAI cosine similarity, Qdrant trả cosine score.
      Sparse — BM25 dot product, Qdrant trả BM25 score.
      Merge  — RRF: score = 1/(k + rank_dense) + 1/(k + rank_sparse).

    Sau RRF, mỗi chunk được mở rộng sang N-3 đến N+3 trong cùng PDF.
    Các cửa sổ chồng lấp tự động được merge thành đoạn liên tục.
    """
    client      = _get_qdrant()
    conv_filter = Filter(must=[
        FieldCondition(key="conversation_id", match=MatchValue(value=conversation_id)),
    ])

    # ── Dense search ──────────────────────────────────────────────────────────
    dense_hits = client.query_points(
        collection_name=_PDF_COLLECTION,
        query=_embed_one(query),
        using="dense",
        query_filter=conv_filter,
        limit=top_k * 3,
        with_payload=True,
    ).points

    # ── Sparse (BM25) search ──────────────────────────────────────────────────
    bm25_vec    = _bm25_one(query)
    sparse_hits = client.query_points(
        collection_name=_PDF_COLLECTION,
        query=SparseVector(indices=bm25_vec.indices, values=bm25_vec.values),
        using="sparse",
        query_filter=conv_filter,
        limit=top_k * 3,
        with_payload=True,
    ).points

    # ── RRF merge ─────────────────────────────────────────────────────────────
    scores:   dict[str, float] = {}
    payloads: dict[str, dict]  = {}

    for rank, hit in enumerate(dense_hits):
        sid = str(hit.id)
        scores[sid]   = scores.get(sid, 0.0) + 1.0 / (rank + _RRF_K)
        payloads[sid] = hit.payload

    for rank, hit in enumerate(sparse_hits):
        sid = str(hit.id)
        scores[sid] = scores.get(sid, 0.0) + 1.0 / (rank + _RRF_K)
        if sid not in payloads:
            payloads[sid] = hit.payload

    top_ids = sorted(scores, key=scores.__getitem__, reverse=True)[:top_k]

    # ── Neighbor expansion ────────────────────────────────────────────────────
    hits_meta = [
        (payloads[sid].get("pdf_name", ""), payloads[sid].get("chunk_index", 0))
        for sid in top_ids
        if sid in payloads
    ]
    return _expand_chunks(client, conversation_id, hits_meta)