Update reg_embedding_system.py

reg_embedding_system.py

@@ -255,7 +255,7 @@ def search_vectorstore(retriever, query, k=5):
     results = retriever.invoke(query)
     return results[:k]
 
-# === search_with_metadata_filter ===
+# === search_with_metadata_filter (사전 필터링 버전) ===
 def search_with_metadata_filter(
     ensemble_retriever: EnsembleRetriever,
     vectorstore: FAISS,
@@ -263,11 +263,16 @@ def search_with_metadata_filter(
     k: int = 5,
     metadata_filter: Optional[Dict[str, Any]] = None,
     sqlite_conn: Optional[sqlite3.Connection] = None,
-    exact_match: bool = True
+    failsafe_search: bool = True
 ) -> List[Document]:
-    """SQLite로 사전 필터링 후 FAISS 검색"""
+    """
+    SQLite로 사전 필터링 → FAISS ID 추출 → IDSelector로 FAISS 검색 제한
+    → BM25는 post-filtering (BM25는 IDSelector 미지원)
+    """
     vector_ret, bm25_ret = ensemble_retriever.retrievers
 
+    vector_docs = []
+
     # === 1. SQLite에서 필터링된 FAISS ID 추출 ===
     filtered_ids = None
     if metadata_filter and sqlite_conn:
@@ -276,18 +281,21 @@ def search_with_metadata_filter(
         params = []
 
         for key, value in metadata_filter.items():
-            #logger.info(f"[key] {key}")
-            #logger.info(f"[value] {value}")
+            print(f"[key] {key}")
+            print(f"[value] {value}")
             if isinstance(value, list): 
+                # IN 쿼리: 리스트 값 지원 
                 if not value: 
-                    continue
+                    continue # 빈 리스트면 무시 
                 placeholders = ', '.join(['?'] * len(value)) 
                 where_clauses.append(f"{key} IN ({placeholders})") 
                 params.extend(value) 
             else: 
+                # 단일 값 
                 where_clauses.append(f"{key} = ?") 
                 params.append(value)
 
+
         if where_clauses:
             where_sql = " OR ".join(where_clauses)
             sql_query = f"SELECT faiss_id FROM documents WHERE {where_sql}"
@@ -295,34 +303,39 @@ def search_with_metadata_filter(
             try:
                 cursor.execute(sql_query, params)
                 filtered_ids = {row[0] for row in cursor.fetchall()}
-                #logger.info(f"[사전 필터링] {len(filtered_ids)}개 ID 획득 → FAISS 검색 제한")
             except Exception as e:
                 logger.info(f"[경고] SQLite 필터링 실패: {e}")
                 filtered_ids = None
-        #else:
-            #logger.info("[안내] 필터 조건 없음 → 전체 검색")
-    #else:
-        #logger.info("[안내] 필터 또는 DB 없음 → 전체 검색")
+        else:
+            logger.info("[안내] 필터 조건 없음 → 전체 검색")
+    else:
+        logger.info("[안내] 필터 또는 DB 없음 → 전체 검색")
 
-    # === 2. FAISS 벡터 검색 ===
+    # === 2. FAISS 벡터 검색 (IDSelector 기반 사전 필터링) ===
     if filtered_ids and len(filtered_ids) > 0:
+        # IDSelector 생성
         selector = MetadataIDSelector(filtered_ids)
+
+        # FAISS 인덱스 추출
         index: faiss.Index = vectorstore.index
-        
         if not hasattr(index, "search"):
             raise ValueError("FAISS 인덱스가 검색을 지원하지 않습니다.")
 
+        # 쿼리 임베딩
         query_embedding = np.array(vectorstore.embeddings.embed_query(query)).astype('float32')
         query_embedding = query_embedding.reshape(1, -1)
 
+        # 검색 파라미터 설정
         search_params = faiss.SearchParametersIVF(
             sel=selector,
-            nprobe=50
+            nprobe=50  # 필요시 조정 (성능 vs 재현율)
         )
 
+        # 여유 있게 k * 10개 후보 요청 (필터 후 부족 방지)
         _k = max(k * 10, 100)
         D, I = index.search(query_embedding, _k, params=search_params)
 
+        # 유효한 결과만 추출
         valid_indices = [i for i in I[0] if i != -1]
         vector_docs = []
         for idx in valid_indices[:k]:
@@ -330,29 +343,27 @@ def search_with_metadata_filter(
             doc = vectorstore.docstore.search(doc_id)
             if isinstance(doc, Document):
                 vector_docs.append(doc)
-
-        #logger.info(f"[벡터 검색] {len(valid_indices)}개 후보 → {len(vector_docs)}개 유효")
     else:
-        search_k = k * 5
-        vector_docs = vector_ret.invoke(query, config={"search_kwargs": {"k": search_k}})
-        #logger.info(f"[벡터 검색] 전체 검색 → {len(vector_docs)}개 후보")
+        if failsafe_search:
+            # 필터 없거나 실패 → 일반 검색 (기존 방식)
+            search_k = k * 5
+            vector_docs = vector_ret.invoke(query, config={"search_kwargs": {"k": search_k}})
 
-    # === 3. BM25 검색 ===
+    # === 3. BM25 검색 (post-filtering, BM25는 IDSelector 미지원) ===
     bm25_docs = []
-    if hasattr(bm25_ret, "invoke"):
-        search_k = k * 5
-        candidates = bm25_ret.invoke(query, config={"search_kwargs": {"k": search_k}})
-        if filtered_ids:
-            bm25_docs = [d for d in candidates if d.metadata.get('faiss_id') in filtered_ids]
-        else:
-            bm25_docs = candidates[:k]
-        #logger.info(f"[BM25 검색] {len(candidates)}개 후보 → {len(bm25_docs)}개 필터링 후")
+    if failsafe_search:
+        if hasattr(bm25_ret, "invoke"):
+            search_k = k * 5
+            candidates = bm25_ret.invoke(query, config={"search_kwargs": {"k": search_k}})
+            if filtered_ids:
+                bm25_docs = [d for d in candidates if d.metadata.get('faiss_id') in filtered_ids]
+            else:
+                bm25_docs = candidates[:k]
 
     # === 4. 병합 및 최종 k개 반환 ===
     combined = {id(d): d for d in (vector_docs + bm25_docs)}.values()
     final_results = list(combined)[:k]
 
-    #logger.info(f"[최종 결과] {len(final_results)}개 문서 반환")
     return final_results
 
 def get_unique_metadata_values(