Hybrid RAG: BM25+Dense (sqlite-vec/BGE-M3) + cross-encoder reranker (bge-reranker-v2-m3)

src/kpaa/embeddings/index.py

@@ -0,0 +1,226 @@
+"""sqlite-vec 기반 dense embeddings 인덱스 — guide + case 통합.
+
+- 데이터: `data/embeddings.sqlite`
+- 가상 테이블 `chunk_vectors` (vec0): chunk_id, embedding[dim]
+- 메타 테이블 `chunk_meta`: chunk_id, source_type, doc_id, embedded_at, model
+
+`kpaa build-embeddings` CLI 가 `kpaa.guides.builder` / `kpaa.cases.index` 의
+chunks 를 읽어 BM25 인덱스와 별개로 dense 벡터를 빌드. 임베딩 텍스트는 BM25 와
+동일하게 `chunk_context + body` (Anthropic Contextual Retrieval).
+"""
+from __future__ import annotations
+
+import logging
+import sqlite3
+from collections.abc import Iterator
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Literal, NamedTuple
+
+logger = logging.getLogger("kpaa.embeddings.index")
+
+_DEFAULT_DB = "embeddings.sqlite"
+
+
+def default_db_path() -> Path:
+    repo_data = Path(__file__).resolve().parents[3] / "data"
+    if repo_data.exists():
+        return repo_data / _DEFAULT_DB
+    from kpaa.config import get_settings
+    p = get_settings().cache_root / "guides"
+    p.mkdir(parents=True, exist_ok=True)
+    return p / _DEFAULT_DB
+
+
+class EmbedHit(NamedTuple):
+    chunk_id: str
+    source_type: str
+    doc_id: str | None
+    distance: float
+
+
+def _connect(db_path: Path, dim: int) -> sqlite3.Connection:
+    import sqlite_vec
+
+    db_path.parent.mkdir(parents=True, exist_ok=True)
+    conn = sqlite3.connect(db_path)
+    conn.enable_load_extension(True)
+    sqlite_vec.load(conn)
+    conn.enable_load_extension(False)
+    conn.execute(
+        f"CREATE VIRTUAL TABLE IF NOT EXISTS chunk_vectors USING vec0("
+        f"chunk_id TEXT PRIMARY KEY, embedding FLOAT[{dim}])"
+    )
+    conn.execute(
+        """
+        CREATE TABLE IF NOT EXISTS chunk_meta (
+            chunk_id    TEXT PRIMARY KEY,
+            source_type TEXT NOT NULL,
+            doc_id      TEXT,
+            embedded_at TEXT NOT NULL,
+            model       TEXT NOT NULL
+        )
+        """
+    )
+    conn.execute("CREATE INDEX IF NOT EXISTS idx_meta_source ON chunk_meta(source_type)")
+    return conn
+
+
+def _iter_guide_chunks() -> Iterator[tuple[str, str, str, str]]:
+    """yield (chunk_id, doc_id, source_type='guide', text)"""
+    from kpaa.guides.builder import _load_jsonl
+    from kpaa.guides.index import chunks_dir
+
+    for f in sorted(chunks_dir().glob("*.jsonl")):
+        for c in _load_jsonl(f):
+            text = f"{c.chunk_context}\n\n{c.body}" if c.chunk_context else c.body
+            yield c.chunk_id, c.doc_id, "guide", text
+
+
+def _iter_case_chunks() -> Iterator[tuple[str, str, str, str]]:
+    """yield (chunk_id='case_<ntt_id>', doc_id='cases', source_type='case', text)"""
+    from kpaa.cases.index import default_db_path as case_db_path
+
+    p = case_db_path()
+    if not p.exists():
+        logger.warning("cases.sqlite 없음 — 사례 인덱싱 건너뜀")
+        return
+    conn = sqlite3.connect(p)
+    conn.row_factory = sqlite3.Row
+    rows = conn.execute(
+        "SELECT ntt_id, body, summary, chunk_context FROM cases"
+    ).fetchall()
+    conn.close()
+    for r in rows:
+        cid = f"case_{r['ntt_id']}"
+        body = r["body"] or r["summary"] or ""
+        ctx = r["chunk_context"] or ""
+        text = f"{ctx}\n\n{body}" if ctx else body
+        yield cid, "cases", "case", text
+
+
+def build_embed_index(
+    *,
+    source: Literal["guide", "case", "all"] = "all",
+    embedder=None,
+    db_path: Path | None = None,
+    force: bool = False,
+    batch: int = 32,
+) -> int:
+    """모든 chunk → embedding → sqlite-vec. 이미 인덱싱된 chunk_id 는 스킵.
+
+    Returns: 새로 인덱싱된 chunk 수.
+    """
+    from kpaa.embeddings.embedder import Embedder
+
+    embedder = embedder or Embedder.default()
+    path = db_path or default_db_path()
+    if force and path.exists():
+        path.unlink()
+        logger.info("기존 %s 삭제 후 재빌드", path)
+
+    conn = _connect(path, dim=embedder.dim)
+    cur = conn.cursor()
+
+    # 이미 인덱싱된 chunk_id 스킵
+    existing = {r[0] for r in cur.execute("SELECT chunk_id FROM chunk_meta")}
+
+    pending: list[tuple[str, str | None, str, str]] = []
+    iters: list[Iterator[tuple[str, str, str, str]]] = []
+    if source in ("guide", "all"):
+        iters.append(_iter_guide_chunks())
+    if source in ("case", "all"):
+        iters.append(_iter_case_chunks())
+
+    for it in iters:
+        for cid, doc_id, src, text in it:
+            if cid in existing:
+                continue
+            pending.append((cid, doc_id, src, text))
+
+    if not pending:
+        logger.info("인덱싱할 새 청크 없음 (existing=%d)", len(existing))
+        conn.close()
+        return 0
+
+    logger.info(
+        "Embedding %d chunks with %s on %s (batch=%d)...",
+        len(pending), embedder.model_name, embedder.device, batch,
+    )
+    texts = [p[3] for p in pending]
+    vectors = embedder.encode_chunks(texts, batch=batch)
+
+    now = datetime.now(timezone.utc).isoformat(timespec="seconds")
+    for (cid, doc_id, src, _), vec in zip(pending, vectors):
+        cur.execute(
+            "INSERT INTO chunk_vectors(chunk_id, embedding) VALUES (?, ?)",
+            (cid, vec.astype("float32").tobytes()),
+        )
+        cur.execute(
+            "INSERT INTO chunk_meta(chunk_id, source_type, doc_id, embedded_at, model) "
+            "VALUES (?, ?, ?, ?, ?)",
+            (cid, src, doc_id, now, embedder.model_name),
+        )
+    conn.commit()
+    conn.close()
+    logger.info("✓ embeddings.sqlite 빌드 완료: %d 신규 청크 → %s", len(pending), path)
+    return len(pending)
+
+
+def search_embed(
+    query: str,
+    *,
+    source_type: Literal["guide", "case"],
+    k: int = 30,
+    embedder=None,
+    db_path: Path | None = None,
+) -> list[EmbedHit]:
+    """vec0 KNN + source_type 필터.
+
+    vec0 가 KNN top-N 을 반환 → JOIN 으로 source_type 후처리 필터 → 상위 k.
+    KPAA 데이터(~2.2k 청크) 기준 KNN limit = k*3 면 source_type 매칭이 충분히 채워짐.
+    """
+    from kpaa.embeddings.embedder import Embedder
+
+    embedder = embedder or Embedder.default()
+    path = db_path or default_db_path()
+    if not path.exists():
+        return []
+
+    conn = _connect(path, dim=embedder.dim)
+    qvec = embedder.encode_query(query).astype("float32").tobytes()
+    knn_limit = max(k * 3, 60)
+    rows = conn.execute(
+        """
+        SELECT v.chunk_id, m.source_type, m.doc_id, v.distance
+        FROM chunk_vectors v
+        JOIN chunk_meta m USING (chunk_id)
+        WHERE v.embedding MATCH ? AND k = ? AND m.source_type = ?
+        ORDER BY v.distance
+        LIMIT ?
+        """,
+        (qvec, knn_limit, source_type, k),
+    ).fetchall()
+    conn.close()
+    return [EmbedHit(*r) for r in rows]
+
+
+def stats(db_path: Path | None = None) -> dict:
+    """인덱스 상태 — CLI/디버깅용."""
+    path = db_path or default_db_path()
+    if not path.exists():
+        return {"path": str(path), "exists": False}
+    # dim은 메타 없이 모름 — 파일 크기로 추정 가능하나 단순화
+    conn = sqlite3.connect(path)
+    rows = conn.execute(
+        "SELECT source_type, COUNT(*) FROM chunk_meta GROUP BY source_type"
+    ).fetchall()
+    models = [r[0] for r in conn.execute("SELECT DISTINCT model FROM chunk_meta")]
+    conn.close()
+    return {
+        "path": str(path),
+        "exists": True,
+        "size_bytes": path.stat().st_size,
+        "by_source": dict(rows),
+        "models": models,
+    }