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

src/kpaa/retrieval/retriever.py

@@ -27,12 +27,66 @@ from functools import lru_cache
 from pathlib import Path
 from typing import Any
 
+import os
+
 from kpaa.cases import CasesIndex
 from kpaa.guides import GuidesIndex
 from kpaa.law_api import KoreanLawClient
 from kpaa.retrieval.excerpts import Excerpt
 from kpaa.retrieval.router import RouterPlan
 
+
+# ─── Hybrid retrieval (BM25 + Dense via sqlite-vec) ────────────────────────
+# `kpaa build-embeddings` 로 data/embeddings.sqlite 가 빌드되어 있으면 자동 사용.
+# 빌드 안 된 환경 / 실패 시 BM25 단독으로 fallback.
+#
+# - KPAA_DENSE_RETRIEVAL=off 면 비활성. (default on)
+# - RRF k 상수: 60 (Cormack et al. 2009 권장값).
+_RRF_K = 60
+
+
+def _dense_enabled() -> bool:
+    return os.environ.get("KPAA_DENSE_RETRIEVAL", "on").lower() not in ("off", "0", "false", "no")
+
+
+def _safe_dense_search(query: str, *, source_type: str, k: int) -> list:
+    """Dense 검색 — 인덱스 없거나 모델 로드 실패 시 [] 반환 (BM25 fallback)."""
+    if not _dense_enabled() or not query:
+        return []
+    try:
+        from kpaa.embeddings.index import search_embed
+
+        return search_embed(query, source_type=source_type, k=k)
+    except Exception as e:  # noqa: BLE001
+        logger.warning("Dense retrieval skipped (%s)", e)
+        return []
+
+
+def _disabled_rerank() -> bool:
+    """Reranker 비활성 여부 — 후보 풀 크기 결정에 사용 (활성 시 크게)."""
+    return os.environ.get("KPAA_RERANKER", "").lower() in ("off", "0", "false", "no", "disabled")
+
+
+def _maybe_rerank(query: str, excerpts: list[Excerpt], *, k: int) -> list[Excerpt]:
+    """Cross-encoder reranker 활성 시 top-k 정밀 정렬, 미설치/disabled 면 원순서 유지."""
+    if not query or not excerpts:
+        return excerpts[:k]
+    try:
+        from kpaa.retrieval.reranker import Reranker
+
+        rr = Reranker.default()
+    except Exception as e:  # noqa: BLE001
+        logger.warning("Reranker import failed (%s) — original order", e)
+        return excerpts[:k]
+    if rr is None or len(excerpts) <= k:
+        return excerpts[:k]
+    return rr.rerank(
+        query, excerpts,
+        text_fn=lambda e: f"{e.title}\n{(e.content or '')[:1500]}" if e.title else (e.content or "")[:1500],
+        top_k=k,
+    )
+# ──────────────────────────────────────────────────────────────────────────
+
 # Progress callback signature: async fn(stage: str, payload: dict).
 # 단계별 SSE prelude 표시용. None 이면 silent.
 ProgressCB = Callable[[str, dict[str, Any]], Awaitable[None]] | None
@@ -211,22 +265,50 @@ async def _fetch_cases(
         if on_progress:
             await on_progress("fetch_done", {"source": "case", "count": 0, "keyword": ""})
         return []
-    # 키워드 + 원본 질문 두 쿼리로 검색 후 RRF(Reciprocal Rank Fusion)로 결합 —
-    # LLM 추출 키워드가 핵심 주제어 누락 시 원본 질문이 안전망. 단순 concat은
-    # BM25 토큰 가중치 차이로 한쪽이 독점 가능하므로 rank 기반 결합 필요.
-    _RRF_K = 60
-    queries: list[str] = []
+    # Hybrid retrieval — BM25 + Dense → RRF.
+    #
+    # BM25 입력: search_keywords 결합 query + 원본 질문 (LLM 키워드 추출이 핵심 주제어
+    #   누락 시 원본이 안전망)
+    # Dense 입력: 원본 질문 1회 (semantic 검색은 자연어 길수록 좋음. 임베딩 비용도
+    #   1회만)
+    #
+    # RRF 로 두 신호 통합. dense 인덱스 없으면 BM25 단독으로 fallback.
+    bm25_queries: list[str] = []
     if plan.search_keywords:
-        queries.append(" ".join(plan.search_keywords[:3]))
-    if plan.query and plan.query not in queries:
-        queries.append(plan.query)
-    rrf_scores: dict = {}
+        bm25_queries.append(" ".join(plan.search_keywords[:3]))
+    if plan.query and plan.query not in bm25_queries:
+        bm25_queries.append(plan.query)
+
+    rrf_scores: dict[str, float] = {}
     hit_map: dict = {}
-    for q in queries:
-        for rank, h in enumerate(idx.search(q, k=k)):
+    pool = max(k * 3, 30)
+
+    # BM25 — 두 query 시도해 RRF 누적
+    for q in bm25_queries:
+        for rank, h in enumerate(idx.search(q, k=pool)):
             rrf_scores[h.ntt_id] = rrf_scores.get(h.ntt_id, 0.0) + 1.0 / (_RRF_K + rank)
             hit_map.setdefault(h.ntt_id, h)
-    top_ids = sorted(rrf_scores, key=lambda i: -rrf_scores[i])[:k]
+
+    # Dense — 원본 질문으로 1회. 결과는 EmbedHit(chunk_id='case_<ntt_id>'..)
+    dense_ids: list[str] = []
+    for rank, eh in enumerate(_safe_dense_search(plan.query, source_type="case", k=pool)):
+        ntt_id = eh.chunk_id.removeprefix("case_")
+        rrf_scores[ntt_id] = rrf_scores.get(ntt_id, 0.0) + 1.0 / (_RRF_K + rank)
+        if ntt_id not in hit_map:
+            dense_ids.append(ntt_id)
+
+    # Dense-only id 들의 Case 본문 lookup (BM25 결과에 없는 것만)
+    if dense_ids:
+        from kpaa.cases.index import get_cases
+
+        extra = get_cases(dense_ids)
+        hit_map.update(extra)
+
+    # Reranker 가용 시 더 큰 후보 풀(k*3 ~ 20)을 reranker 에 넘겨 정밀 정렬
+    rerank_pool = max(k * 3, 15) if not _disabled_rerank() else k
+    top_ids = [i for i in sorted(rrf_scores, key=lambda i: -rrf_scores[i]) if i in hit_map][
+        :rerank_pool
+    ]
     hits = [hit_map[i] for i in top_ids]
     out: list[Excerpt] = []
     for h in hits:
@@ -253,6 +335,7 @@ async def _fetch_cases(
                 recency_score=_recency_score(_yyyy_mmdd_to_year(h.case_year or h.reg_dt)),
             )
         )
+    out = _maybe_rerank(plan.query, out, k=k)
     if on_progress:
         await on_progress(
             "fetch_done", {"source": "case", "count": len(out), "keyword": plan.top_keyword}
@@ -274,23 +357,46 @@ async def _fetch_guides(
         if on_progress:
             await on_progress("fetch_done", {"source": "guide", "count": 0, "keyword": ""})
         return []
-    # 키워드 + 원본 질문 두 쿼리로 검색 후 RRF로 결합 — LLM 추출 키워드가 핵심
-    # 주제어를 누락해도 원본 질문이 안전망 (e.g. "처방전 보관기간" → ["보관기간"]만
-    # 추출돼도 원본 query에서 "처방전" 토큰 hit 가능). 단순 concat은 BM25 가중치
-    # 차이로 한쪽이 독점하므로 rank 기반 union 필요.
-    _RRF_K = 60
-    queries: list[str] = []
+    # Hybrid retrieval — BM25 + Dense → RRF.
+    #
+    # BM25 입력: search_keywords 결합 query + 원본 질문 (LLM 키워드 추출이 핵심 주제어
+    #   누락 시 원본이 안전망. e.g. "처방전 보관기간" → ["보관기간"]만 추출돼도 원본
+    #   query 에서 "처방전" 토큰 hit 가능)
+    # Dense 입력: 원본 질문 1회 (semantic 검색은 자연어 길수록 좋음)
+    bm25_queries: list[str] = []
     if plan.search_keywords:
-        queries.append(" ".join(plan.search_keywords[:3]))
-    if plan.query and plan.query not in queries:
-        queries.append(plan.query)
-    rrf_scores: dict = {}
+        bm25_queries.append(" ".join(plan.search_keywords[:3]))
+    if plan.query and plan.query not in bm25_queries:
+        bm25_queries.append(plan.query)
+
+    rrf_scores: dict[str, float] = {}
     hit_map: dict = {}
-    for q in queries:
-        for rank, h in enumerate(idx.search(q, k=k)):
+    pool = max(k * 3, 30)
+
+    # BM25
+    for q in bm25_queries:
+        for rank, h in enumerate(idx.search(q, k=pool)):
             rrf_scores[h.chunk_id] = rrf_scores.get(h.chunk_id, 0.0) + 1.0 / (_RRF_K + rank)
             hit_map.setdefault(h.chunk_id, h)
-    top_ids = sorted(rrf_scores, key=lambda i: -rrf_scores[i])[:k]
+
+    # Dense — 원본 질문 1회 (chunk_id 그대로)
+    dense_ids: list[str] = []
+    for rank, eh in enumerate(_safe_dense_search(plan.query, source_type="guide", k=pool)):
+        rrf_scores[eh.chunk_id] = rrf_scores.get(eh.chunk_id, 0.0) + 1.0 / (_RRF_K + rank)
+        if eh.chunk_id not in hit_map:
+            dense_ids.append(eh.chunk_id)
+
+    # Dense-only chunk_id 들의 GuideChunk 본문 lookup
+    if dense_ids:
+        from kpaa.guides.index import get_chunks
+
+        extra = get_chunks(dense_ids)
+        hit_map.update(extra)
+
+    rerank_pool = max(k * 3, 15) if not _disabled_rerank() else k
+    top_ids = [i for i in sorted(rrf_scores, key=lambda i: -rrf_scores[i]) if i in hit_map][
+        :rerank_pool
+    ]
     hits = [hit_map[i] for i in top_ids]
     out: list[Excerpt] = []
     for h in hits:
@@ -318,6 +424,7 @@ async def _fetch_guides(
                 recency_score=_recency_score(year),
             )
         )
+    out = _maybe_rerank(plan.query, out, k=k)
     if on_progress:
         await on_progress(
             "fetch_done",