exAichatbot_agent/exAiChatBot-chatbot2.0-agent/scripts/handlers/search_handler.py

"""
검색 핸들러 모듈
──────────────
벡터 검색 및 재랭킹 처리
"""

import numpy as np
from typing import List, Dict, Any, Optional, Tuple


def _format_rerank_document(candidate: Dict[str, Any]) -> str:
    """리랭커용 FAQ passage — 질문과 답변 전체."""
    q = str(candidate.get("q") or candidate.get("question") or "").strip()
    a = str(candidate.get("a") or candidate.get("answer") or "").strip()
    if q and a:
        return f"질문: {q}\n답변: {a}"
    return q or a


class SearchHandler:
    """벡터 검색 및 재랭킹 핸들러"""
    
    def __init__(self, vector_store, embed_client, rerank_client, config):
        self.vector_store = vector_store
        self.embed_client = embed_client
        self.rerank_client = rerank_client
        self.config = config
        self.last_search_info: Dict[str, Any] = {}
    
    def embed_query(self, query: str, ts: str) -> Optional[np.ndarray]:
        """질문 임베딩
        
        Args:
            query: 검색 질문
            ts: 타임스탬프 (로깅용)
            
        Returns:
            임베딩 벡터 또는 None (실패 시)
        """
        self._last_query_text = query or ""
        try:
            embeddings = self.embed_client.embed([query], normalize=True, is_query=True)
            return np.array(embeddings[0], dtype="float32")
        except Exception as e:
            print(f"[SearchHandler] {ts} ❌ 임베딩 실패: {e}")
            return None
    
    def search(self, query_vec: np.ndarray, threshold: float, ts: str) -> List[Dict[str, Any]]:
        """벡터 검색
        
        Args:
            query_vec: 질문 임베딩 벡터
            threshold: 유사도 임계값
            ts: 타임스탬프 (로깅용)
            
        Returns:
            검색 결과 리스트
        """
        try:
            if (
                getattr(self.config, "hybrid_search_enabled", False)
                and hasattr(self.vector_store, "hybrid_search")
            ):
                results = self.vector_store.hybrid_search(
                    query_vec,
                    getattr(self, "_last_query_text", ""),
                    top_k=self.config.top_k,
                    threshold=threshold,
                    sparse_top_k=getattr(self.config, "sparse_top_k", 30),
                    merge_top_k=getattr(self.config, "hybrid_merge_top_k", 40),
                )
                self.last_search_info = {
                    "searchMode": "hybrid_sparse",
                    "candidateCount": len(results),
                }
            else:
                results = self.vector_store.search(
                    query_vec,
                    top_k=self.config.top_k,
                    threshold=threshold
                )
                self.last_search_info = {
                    "searchMode": "vector",
                    "candidateCount": len(results),
                }
            return results
        except Exception as e:
            print(f"[SearchHandler] {ts} ❌ 벡터 검색 실패: {e}")
            self.last_search_info = {"searchMode": "error", "error": str(e)}
            return []
    
    def rerank(self, query: str, candidates: List[Dict[str, Any]], ts: str) -> Tuple[List[Dict[str, Any]], List[float], bool, Optional[Dict]]:
        """재랭킹
        
        Args:
            query: 검색 질문
            candidates: 후보 문서 리스트
            ts: 타임스탬프 (로깅용)
            
        Returns:
            (상위 N개 결과, 점수 리스트, 재랭킹 사용 여부, 재랭킹 정보)
        """
        if not candidates:
            print(f"[SearchHandler] {ts} 재랭킹 건너뜀: 후보 문서 없음")
            return [], [], False, {"detail": "No candidates to rerank"}

        # 재랭킹 후보 제한
        max_rerank = min(self.config.rerank_candidates, len(candidates))
        rerank_candidates = candidates[:max_rerank]
        
        print(f"[SearchHandler] {ts} 재랭킹 시작 (candidates={len(candidates)} → 상위 {max_rerank}개)")
        
        try:
            # TEI Reranker API 호출 (질문+답변 전체를 passage로 전달)
            documents = [_format_rerank_document(c) for c in rerank_candidates]
            print(f"[SearchHandler] {ts} 리랭커 호출: query={query[:50]}..., documents={len(documents)}개 (Q+A)")
            
            rerank_results = self.rerank_client.rerank(
                query=query,
                documents=documents,
                return_documents=False
            )
            
            print(f"[SearchHandler] {ts} 리랭커 응답: {len(rerank_results) if rerank_results else 0}개 결과")
            
            if not rerank_results:
                print(f"[SearchHandler] {ts} 재랭킹 결과 없음 → FAISS 상위 결과 사용")
                top_n_results = rerank_candidates[:self.config.top_n_for_llm]
                top_scores = [None] * len(top_n_results)
                return top_n_results, top_scores, False, None
            
            # score 기준 정렬 후 상위 N개
            sorted_results = sorted(rerank_results, key=lambda x: x["score"], reverse=True)
            top_n_indices = [r["index"] for r in sorted_results[:self.config.top_n_for_llm]]
            top_n_results = [rerank_candidates[idx] for idx in top_n_indices]
            top_scores = [
                round(sorted_results[i]["score"], 4) if i < len(sorted_results) else None
                for i in range(len(top_n_results))
            ]
            
            # 순위 변화 계산
            rank_changes = [idx - i for i, idx in enumerate(top_n_indices)]
            
            def _question_of(candidate: Dict[str, Any]) -> str:
                return str(candidate.get("q") or candidate.get("question") or "")

            rerank_info = {
                "used": True,
                "original_top_question": _question_of(rerank_candidates[0])[:50] + "...",
                "reranked_top_question": _question_of(top_n_results[0])[:50] + "...",
                "rank_changes": rank_changes
            }
            
            print(f"[SearchHandler] {ts} 재랭킹 완료: 상위 {len(top_n_results)}개, 최고 점수={top_scores[0]}")
            if rank_changes[0] != 0:
                print(f"[SearchHandler] {ts} 순위 변화: FAISS #{top_n_indices[0]+1} → Rerank #1")
            
            return top_n_results, top_scores, True, rerank_info
            
        except Exception as e:
            print(f"[SearchHandler] {ts} ❌ 재랭킹 실패: {e}")
            print(f"[SearchHandler] {ts} FAISS 상위 결과로 대체")
            top_n_results = candidates[:self.config.top_n_for_llm]
            top_scores = [None] * len(top_n_results)
            return top_n_results, top_scores, False, None