""" 벡터 데이터베이스 관리 모듈 """ import os import json import logging import re import numpy as np from concurrent.futures import ThreadPoolExecutor, as_completed from typing import Dict, List, Any, Optional, Union, Tuple from datetime import datetime from tqdm import tqdm import chromadb from chromadb.config import Settings from sentence_transformers import SentenceTransformer from src import config from src.utils import ( VectorDBError, DataProcessingError, get_latest_data_files, load_json_data, determine_optimal_chunk_size, split_text_into_chunks, check_index_status, structured_chunking ) from src.bm25_client import BM25Client from src.standard_codes import ( STANDARD_CODE_PATTERN, canonical_doc_type, collection_for_doc_type, collection_names_for_doc_types, doc_type_matches, normalize_standard_code, standard_family_name, ) from src.ranking_rules import score_ranking_rules logger = config.setup_logger("vector_db") STANDARD_CODE_LABEL_PATTERN = re.compile(r"(?:표준코드|기준코드)\W*[:：]?\s*", re.IGNORECASE) def _normalize_standard_code(value: Any) -> str: return normalize_standard_code(value) def _extract_standard_codes(value: Any) -> List[str]: codes: List[str] = [] seen = set() for match in STANDARD_CODE_PATTERN.finditer(str(value or "")): code = f"{match.group(1).upper()} {match.group(2)} {match.group(3)} {match.group(4)}" if code not in seen: codes.append(code) seen.add(code) return codes def _extract_labeled_standard_code(value: Any) -> str: text = str(value or "") label_match = STANDARD_CODE_LABEL_PATTERN.search(text) if not label_match: return "" codes = _extract_standard_codes(text[label_match.end(): label_match.end() + 100]) return codes[0] if codes else "" class KCSCVectorDB: """ ChromaDB를 활용한 벡터 데이터베이스 관리 한국건설기준센터(KCSC) 데이터를 임베딩하고 검색하는 기능 제공 """ _instance = None _model = None def __new__(cls, *args, **kwargs): """싱글턴 패턴 구현""" if cls._instance is None: cls._instance = super(KCSCVectorDB, cls).__new__(cls) return cls._instance def __init__(self, data_dir=None, db_dir=None, embedding_model=None): """ 벡터 데이터베이스 초기화 Args: data_dir: 데이터 디렉토리 경로 db_dir: 벡터 DB 저장 경로 embedding_model: 임베딩 모델 이름 또는 경로 """ # 이미 초기화된 경우 중복 초기화 방지 if hasattr(self, 'initialized') and self.initialized: return self.data_dir = data_dir or config.DATA_DIR self.db_dir = db_dir or config.VECTOR_DB_DIR # 디렉토리 생성 os.makedirs(self.data_dir, exist_ok=True) os.makedirs(self.db_dir, exist_ok=True) # 임베딩 모델 초기화 self._init_embedding_model(embedding_model or config.EMBEDDING_MODEL) # BM25 클라이언트 초기화 self.bm25_client = BM25Client() self.bm25_client.load_index("kcsc_bm25") # ChromaDB 클라이언트 초기화 logger.info(f"ChromaDB 초기화 중... (경로: {self.db_dir})") try: self.client = chromadb.PersistentClient( path=self.db_dir, settings=Settings( anonymized_telemetry=False, allow_reset=True ) ) # 컬렉션 초기화 self.collections = {} self._init_collections() self.initialized = True logger.info("벡터 데이터베이스 초기화 완료") except Exception as e: logger.error(f"ChromaDB 초기화 실패: {str(e)}") raise VectorDBError(f"벡터 데이터베이스 초기화 실패: {str(e)}", "init") @classmethod def _init_embedding_model(cls, model_name): """ 임베딩 모델 초기화 (클래스 메서드로 싱글턴 구현) Args: model_name: 임베딩 모델 이름 또는 경로 """ if cls._model is None: try: logger.info(f"임베딩 모델 로드 중... (모델: {model_name})") cls._model = SentenceTransformer(model_name) logger.info("임베딩 모델 로드 완료") except Exception as e: logger.error(f"임베딩 모델 로드 실패: {str(e)}") raise VectorDBError(f"임베딩 모델 '{model_name}' 로드 실패: {str(e)}", "model_load") def _init_collections(self): """컬렉션 초기화 또는 로드""" for doc_type in ["KDS", "KCS"]: try: self.collections[doc_type] = self.client.get_or_create_collection( name=doc_type, metadata={"description": f"{doc_type} 설계기준 문서"} ) count = self.collections[doc_type].count() logger.info(f"{doc_type} 컬렉션 초기화 완료 (항목 수: {count})") except Exception as e: logger.error(f"{doc_type} 컬렉션 초기화 실패: {str(e)}") raise VectorDBError(f"{doc_type} 컬렉션 초기화 실패: {str(e)}", "collection_init") def process_document(self, doc: Dict[str, Any], doc_type: str) -> Optional[Dict[str, Any]]: """ 문서를 처리하여 인덱싱 가능한 형태로 변환 Args: doc: 문서 데이터 doc_type: 문서 유형 Returns: 처리된 문서 데이터 또는 None (내용이 없는 경우) """ # 기본 필드 설정 processed = { "id": f"{doc_type}-{doc.get('Code', '')}-{doc.get('No', '0')}", "code": doc.get('Code', ''), "full_code": doc.get('FullCode', ''), "name": doc.get('Name', ''), "title": doc.get('Title', ''), "content": doc.get('Contents', ''), "version": doc.get('Version', ''), "update_date": doc.get('UpdateDate', ''), "doc_type": doc_type } # 내용이 있는 경우에만 반환 if processed["content"]: return processed logger.warning(f"문서 {processed['id']}에 내용이 없습니다.") return None def create_document_chunks( self, doc: Dict[str, Any], chunk_size: int = None, overlap: int = None ) -> List[Dict[str, Any]]: """ 문서를 청크로 나누기 Args: doc: 처리된 문서 데이터 chunk_size: 청크 크기 (None이면 자동 결정) overlap: 청크 간 겹치는 문자 수 Returns: 청크 목록 """ content = doc.get("content", "") if not content: return [doc] # 청크 크기 결정 if chunk_size is None: chunk_size = determine_optimal_chunk_size(content) overlap = overlap or config.CHUNK_OVERLAP # 청크로 나누기 if len(content) <= chunk_size: # 작은 문서는 그대로 사용 doc["chunk_id"] = f"{doc['id']}-0" return [doc] # 청크 생성 text_chunks = split_text_into_chunks(content, chunk_size, overlap) chunks = [] for i, chunk_content in enumerate(text_chunks): chunk = doc.copy() chunk["content"] = chunk_content chunk["chunk_id"] = f"{doc['id']}-{i}//{len(text_chunks)}" chunks.append(chunk) return chunks def index_documents( self, doc_type: str, documents: List[Dict[str, Any]], chunk_size: int = None, overlap: int = None, batch_size: int = 100, reset: bool = True ) -> Tuple[int, int]: """ 문서를 인덱싱 Args: doc_type: 문서 유형 documents: 문서 목록 chunk_size: 청크 크기 overlap: 청크 간 겹치는 문자 수 batch_size: 배치 크기 reset: 기존 데이터 리셋 여부 Returns: (인덱싱된 문서 수, 청크 수) 튜플 """ if doc_type not in self.collections: raise VectorDBError(f"{doc_type} 컬렉션이 없습니다.", "collection_missing") collection = self.collections[doc_type] # 청크 생성 chunks = [] processed_docs = 0 for doc in documents: processed = self.process_document(doc, doc_type) if not processed: continue # 문서를 청크로 나누기 doc_chunks = self.create_document_chunks(processed, chunk_size, overlap) chunks.extend(doc_chunks) processed_docs += 1 logger.info(f"{doc_type}: {processed_docs}개 문서에서 {len(chunks)}개 청크 생성됨") # 컬렉션 리셋 (필요시) if reset and collection.count() > 0: logger.info(f"{doc_type} 컬렉션 초기화 중... (기존 {collection.count()}개 항목)") collection.delete(where={}) # 청크 데이터 준비 if not chunks: logger.warning(f"{doc_type}: 인덱싱할 청크가 없습니다.") return 0, 0 # 배치 처리 total_batches = (len(chunks) + batch_size - 1) // batch_size for batch_idx in range(total_batches): start_idx = batch_idx * batch_size end_idx = min((batch_idx + 1) * batch_size, len(chunks)) batch_chunks = chunks[start_idx:end_idx] try: ids = [] documents = [] metadatas = [] embeddings = [] for chunk in batch_chunks: # 임베딩할 텍스트 준비 text = f"{chunk['name']} {chunk['title'] or ''}: {chunk['content']}" # 메타데이터 준비 metadata = { "id": chunk["id"], "code": chunk["code"], "full_code": chunk["full_code"], "name": chunk["name"], "title": chunk["title"], "doc_type": chunk["doc_type"], "version": chunk["version"] } # 임베딩 생성 if self._model is None: raise VectorDBError("임베딩 모델이 초기화되지 않았습니다.", "model_missing") embedding = self._model.encode(text) ids.append(chunk["chunk_id"]) documents.append(text) metadatas.append(metadata) embeddings.append(embedding.tolist()) # 데이터 추가 collection.add( ids=ids, documents=documents, metadatas=metadatas, embeddings=embeddings ) logger.info(f"{doc_type} 배치 {batch_idx+1}/{total_batches} 인덱싱 완료 ({len(ids)}개 항목)") except Exception as e: logger.error(f"{doc_type} 배치 {batch_idx+1} 인덱싱 실패: {str(e)}") raise VectorDBError(f"{doc_type} 배치 인덱싱 실패: {str(e)}", "batch_indexing") return processed_docs, len(chunks) def build_hybrid_index_from_standards(self, file_path: str = None, chunk_size: int = 4000, limit: int = None) -> Dict[str, Any]: """ standards.json을 로드하여 필드 기반 청킹 후 하이브리드 인덱스 구축. 전략: scope / materials / construction / quality / safety 5개 핵심 필드의 발췌문을 합쳐 문서당 1개 대표 청크를 만든다. """ file_path = file_path or os.path.join(config.DATA_DIR, "raw", "standards.json") if not os.path.exists(file_path): logger.warning(f"표준 데이터 파일을 찾을 수 없습니다: {file_path}") return {"status": "error", "message": "표준 데이터 파일 없음"} logger.info(f"표준 데이터 로드 중... ({file_path})") try: with open(file_path, 'r', encoding='utf-8') as f: standards_data = json.load(f) except Exception as e: return {"status": "error", "message": f"로드 실패: {e}"} items = list(standards_data.values()) if isinstance(standards_data, dict) else standards_data if limit: items = items[:limit] # 문서당 1개 대표 청크 전략: # standard_code + title + 각 핵심 필드 앞 N자를 합산 → 하나의 검색 가능한 단위 생성 # 목표: 5,233 문서 = 5,233 청크 → 임베딩 배치 11개 → ~2분 완료 FIELD_EXCERPT = [ ("scope", "적용범위", 600), ("materials", "자재", 600), ("construction", "시공", 600), ("quality", "품질관리", 400), ("safety", "안전", 400), ] SKIP_VALUES = {"내용 없음", "", None} all_chunks: List[Dict[str, Any]] = [] processed_docs = 0 fallback_code_docs = [] skipped_docs = [] failed_docs = [] def resolve_standard_code(doc: Dict[str, Any], doc_idx: int) -> Tuple[str, str]: raw_code = doc.get("standard_code") or doc.get("code") or "" code = _normalize_standard_code(raw_code) if code: return code, "standard_code" title_codes = _extract_standard_codes(doc.get("title", "")) if len(title_codes) == 1: return title_codes[0], "title" if len(title_codes) > 1: failed_docs.append({ "doc_idx": doc_idx, "title": str(doc.get("title", ""))[:160], "reason": "multiple_codes_in_title", "candidates": title_codes, }) return "", "" for field_name in ("full_content", "summary", "scope"): labeled_code = _extract_labeled_standard_code(doc.get(field_name, "")) if labeled_code: return labeled_code, f"{field_name}_label" for field_name in ("summary", "scope", "full_content"): field_codes = _extract_standard_codes(doc.get(field_name, "")) if len(field_codes) == 1: return field_codes[0], field_name if len(field_codes) > 1: failed_docs.append({ "doc_idx": doc_idx, "title": str(doc.get("title", ""))[:160], "reason": f"multiple_codes_in_{field_name}", "candidates": field_codes[:10], }) return "", "" failed_docs.append({ "doc_idx": doc_idx, "title": str(doc.get("title", ""))[:160], "reason": "code_not_found", "candidates": [], }) return "", "" for doc_idx, doc in enumerate(tqdm(items, desc="문서 대표 청크 생성 중")): code, code_source = resolve_standard_code(doc, doc_idx) title = doc.get("title", "") if not code: continue if code_source != "standard_code": fallback_code_docs.append({ "doc_idx": doc_idx, "code": code, "source": code_source, "title": str(title)[:160], }) doc_type = canonical_doc_type(code) collection_type = collection_for_doc_type(doc_type) standard_family = standard_family_name(doc_type) # 짧은 단어형 제목은 제목만으로 의미가 모호하므로 기준코드를 청크 앞에 명시한다. parts = [ f"[기준코드] {code}", f"[기준종류] {standard_family}", f"[제목] {code} {title}".strip(), ] for field_key, field_label, max_len in FIELD_EXCERPT: val = str(doc.get(field_key, "") or "").strip() if val and val not in SKIP_VALUES: effective_max_len = 1000 if field_key == "scope" and "(전문)" in title else max_len excerpt = val[:effective_max_len] parts.append(f"[{field_label}] {excerpt}") page_content = "\n".join(parts) if len(page_content) < 50: # 내용이 거의 없는 문서 건너뜀 skipped_docs.append({ "doc_idx": doc_idx, "code": code, "title": str(title)[:160], "reason": "representative_chunk_too_short", }) continue chunk_id = f"{code}_main_{doc_idx}" all_chunks.append({ "id": chunk_id, "page_content": page_content, "metadata": { "code": code, "full_code": code, "name": title, "title": title, "doc_type": doc_type, "collection_type": collection_type, "standard_family": standard_family, "standard_code_source": code_source, "field": "combined", "chunk_index": 0, } }) processed_docs += 1 logger.info(f"총 {processed_docs}개 문서에서 {len(all_chunks)}개 청크 생성됨.") if fallback_code_docs: logger.info(f"표준코드 fallback 적용: {len(fallback_code_docs)}개 문서") if failed_docs: logger.error(f"표준코드 추출 실패: {len(failed_docs)}개 문서") return { "status": "error", "message": "표준코드 추출 실패 문서가 있어 인덱스 빌드를 중단했습니다.", "processed_docs": processed_docs, "total_input_docs": len(items), "fallback_code_docs": len(fallback_code_docs), "skipped_docs": len(skipped_docs), "failed_docs": len(failed_docs), "failed_doc_samples": failed_docs[:20], } # ── BM25 인덱스 빌드 ───────────────────────────────────────── self.bm25_client.build_index(all_chunks) self.bm25_client.save_index("kcsc_bm25") # ── ChromaDB 인덱스 빌드 ────────────────────────────────────── for doc_type in ["KDS", "KCS"]: try: self.client.delete_collection(name=doc_type) except Exception: pass collection = self.client.create_collection( name=doc_type, metadata={"description": f"{doc_type} 설계기준 문서"} ) self.collections[doc_type] = collection type_chunks = [c for c in all_chunks if c["metadata"].get("collection_type") == doc_type] if not type_chunks: logger.info(f"{doc_type}: 인덱싱할 청크 없음 (건너뜀)") continue batch_size = 512 # 임베딩 배치 크기 (속도 최적화) total_batches = (len(type_chunks) + batch_size - 1) // batch_size logger.info(f"{doc_type} ChromaDB 인덱싱 시작: {len(type_chunks)}개 청크, {total_batches}개 배치") for batch_idx in tqdm(range(total_batches), desc=f"{doc_type} 임베딩"): s = batch_idx * batch_size e = min(s + batch_size, len(type_chunks)) batch = type_chunks[s:e] ids = [c["id"] for c in batch] documents = [c["page_content"] for c in batch] metadatas = [c["metadata"] for c in batch] embeddings = self._model.encode(documents, show_progress_bar=False).tolist() collection.add(ids=ids, documents=documents, metadatas=metadatas, embeddings=embeddings) logger.info(f"ChromaDB {doc_type} 인덱싱 완료 ({len(type_chunks)}개 청크)") return { "status": "success", "processed_docs": processed_docs, "total_input_docs": len(items), "fallback_code_docs": len(fallback_code_docs), "skipped_docs": len(skipped_docs), "failed_docs": len(failed_docs), "total_chunks": len(all_chunks), "kds_chunks": sum(1 for c in all_chunks if c["metadata"].get("doc_type") == "KDS"), "kcs_chunks": sum(1 for c in all_chunks if c["metadata"].get("doc_type") == "KCS"), "excs_chunks": sum(1 for c in all_chunks if c["metadata"].get("doc_type") == "EXCS"), "kcs_collection_chunks": sum(1 for c in all_chunks if c["metadata"].get("collection_type") == "KCS"), } def build_index_from_files( self, chunk_size: int = None, overlap: int = None, reset: bool = True ) -> Dict[str, Dict[str, int]]: """ 모든 문서 파일을 인덱싱 Args: chunk_size: 청크 크기 overlap: 청크 간 겹치는 문자 수 reset: 기존 데이터 리셋 여부 Returns: 인덱싱 통계 정보 """ stats = {} # 세부 정보 파일 가져오기 detail_files = get_latest_data_files("KDS_details", self.data_dir) + \ get_latest_data_files("KCS_details", self.data_dir) if not detail_files: logger.warning("인덱싱할 데이터 파일이 없습니다.") return {"status": "warning", "message": "인덱싱할 데이터 파일이 없습니다."} for file_path in detail_files: try: # 파일명에서 문서 유형 추출 file_name = os.path.basename(file_path) doc_type = file_name.split('_')[0] logger.info(f"{file_path} 파일 인덱싱 시작") # 데이터 로드 documents = load_json_data(file_path) if documents: # 인덱싱 doc_count, chunk_count = self.index_documents( doc_type, documents, chunk_size=chunk_size, overlap=overlap, reset=reset ) stats[doc_type] = { "documents": doc_count, "chunks": chunk_count, "file": file_path } logger.info(f"{doc_type} 인덱싱 완료: {doc_count}개 문서, {chunk_count}개 청크") else: logger.warning(f"{file_path} 파일에 인덱싱할 문서가 없습니다.") stats[doc_type] = {"documents": 0, "chunks": 0, "file": file_path} except Exception as e: logger.error(f"{file_path} 파일 인덱싱 실패: {str(e)}") stats[os.path.basename(file_path)] = {"error": str(e)} return { "status": "success" if stats else "error", "stats": stats } def search( self, query: str, doc_types: List[str] = None, limit: int = None, min_relevance: float = 0.3 ) -> List[Dict[str, Any]]: """ 자연어 쿼리로 관련 문서 검색 Args: query: 검색 쿼리 doc_types: 검색할 문서 유형 목록 limit: 결과 제한 수 min_relevance: 최소 관련성 점수 Returns: 검색 결과 목록 """ if not query: logger.warning("검색 쿼리가 비어 있습니다.") return [] limit = limit or config.SEARCH_LIMIT allowed_doc_types = {str(dt).upper() for dt in doc_types} if doc_types else None collection_doc_types = ( sorted(collection_names_for_doc_types(doc_types)) if doc_types else list(self.collections.keys()) ) logger.debug(f"검색 쿼리: '{query}', 대상 문서 유형: {doc_types or collection_doc_types}") try: standard_code_pattern = STANDARD_CODE_PATTERN def clean_title_text(value: Any) -> str: title = standard_code_pattern.sub(" ", str(value or "")) title = re.sub(r"[_\-]+", " ", title) title = re.sub(r"\([^)]*(?:\d{2,4}(?:[.\-]\d{1,2})?|전문|수정|제정)[^)]*\)", " ", title) title = re.sub(r"(?:수정|제정|전문)", " ", title) return re.sub(r"\s+", " ", title).strip() def normalize_title_match_text(value: Any) -> str: return "".join(ch for ch in str(value or "").lower() if ch.isalnum()) normalized_query = normalize_title_match_text(query) clean_normalized_query = normalize_title_match_text(clean_title_text(query)) query_tokens = set(self.bm25_client._tokenize(query)) if self.bm25_client.documents else set() wants_excs = ( "고속도로공사전문시방서" in normalized_query or "전문시방서" in normalized_query or "고속도로공사" in normalized_query ) wants_kds = "설계기준" in normalized_query wants_kcs = "표준시방서" in normalized_query generic_signal_tokens = { "기준", "설계", "설계기준", "공사", "구조", "일반", "수정", "제정", "검토", "검토할", "기준으로", "찾아줘", "알려주세요", "콘크", "크리", "리트", "콘크리", "크리트", "트구", "리트구", "트구조", "크리트구", "리트구조", "계기", "계기준", "길이", "조건", "확인", "결과", "어떤", "하는", "하는데", "해야", "하나요", "검토해", "검토해야", "토해", "토해야", "적합", "적합한", "충분", "충분한", } important_short_terms = { "내진", "지진", "정착", "이음", "철근", "수압", "시험", "검사", "안전", "보건", "관리", "계획", "합성", "강도", "피복", "전단", "압축", "볼트", "접합", "용접", "방수", "지반", "기초", "품질", "앵커", "가설", "슬래브", "말뚝", "교량", "터널", "배관", "네일", "옹벽", "한중", "조임", "체결", } query_signal_tokens = { token for token in query_tokens if token not in generic_signal_tokens and (len(token) >= 3 or token in important_short_terms) } token_idf = getattr(self.bm25_client.bm25, "idf", {}) if self.bm25_client.bm25 else {} def token_importance(token: str) -> float: try: return max(float(token_idf.get(token, 0.0)), 0.1) except Exception: return 0.1 vector_candidate_count = max(limit, 5) lexical_candidate_count = max(limit * 30, 150) # 쿼리 임베딩. 모델/Chroma 쪽이 메모리 부족으로 실패해도 # 아래 BM25 경로는 계속 사용할 수 있어야 한다. query_embedding = None if self._model is None: logger.warning("임베딩 모델이 초기화되지 않아 벡터 검색을 건너뜁니다.") else: try: query_embedding = self._model.encode(query) except Exception as e: logger.warning(f"쿼리 임베딩 실패, BM25 검색만 수행합니다: {e}") all_results = [] # 각 컬렉션에서 검색 if query_embedding is not None: for doc_type in collection_doc_types: if doc_type not in self.collections: logger.warning(f"{doc_type} 컬렉션이 없습니다.") continue collection = self.collections[doc_type] try: collection_count = collection.count() except Exception as e: logger.warning(f"{doc_type} 컬렉션 count 실패, 벡터 검색 건너뜀: {e}") continue # 컬렉션에 데이터가 없으면 건너뛰기 if collection_count == 0: logger.warning(f"{doc_type} 컬렉션에 데이터가 없습니다.") continue try: # 검색 실행 results = collection.query( query_embeddings=[query_embedding.tolist()], n_results=min(vector_candidate_count, collection_count) # 필터링 전에 더 많이 가져옴 ) except Exception as e: logger.warning(f"{doc_type} Chroma 벡터 검색 실패, BM25 검색으로 계속 진행: {e}") continue # 결과 형식 변환 for i in range(len(results['ids'][0])): distance = float(results['distances'][0][i]) relevance = float(1.0 / (1.0 + max(distance, 0.0))) # 거리값을 0~1 유사도로 정규화 result = { "id": results['ids'][0][i], "text": results['documents'][0][i], "metadata": results['metadatas'][0][i], "distance": distance, "relevance": relevance } all_results.append(result) # BM25 검색 (Lexical) bm25_results = self.bm25_client.search(query, top_k=lexical_candidate_count) if self.bm25_client.bm25 else [] title_results = [] if self.bm25_client.documents and normalized_query: for doc in self.bm25_client.documents: metadata = doc.get("metadata", {}) title = metadata.get("title") or metadata.get("name") clean_title = clean_title_text(title) normalized_title = normalize_title_match_text(clean_title) if not normalized_title: continue code = metadata.get("code", "") raw_type = canonical_doc_type(code, metadata.get("doc_type", "")) title_score = 0.0 if normalized_title in clean_normalized_query: title_score += 80.0 + min(len(normalized_title), 30) else: title_tokens = set(self.bm25_client._tokenize(clean_title)) overlap = title_tokens & query_tokens if overlap: title_score += float(len(overlap)) title_score += sum( min(token_importance(token), 5.0) for token in overlap & query_signal_tokens ) if len(overlap) / max(len(title_tokens), 1) >= 0.5: title_score += 3.0 if title_score > 0: if wants_excs and raw_type == "EXCS": title_score += 10.0 elif wants_kds and raw_type == "KDS": title_score += 8.0 elif wants_kcs and raw_type == "KCS": title_score += 8.0 title_results.append((doc, title_score)) title_results.sort(key=lambda item: item[1], reverse=True) title_results = title_results[:lexical_candidate_count] substring_results = [] if self.bm25_client.documents: query_terms = [normalize_title_match_text(term) for term in query.split()] query_terms = [term for term in query_terms if len(term) >= 2] for doc in self.bm25_client.documents: metadata = doc.get("metadata", {}) title_text = normalize_title_match_text(clean_title_text(metadata.get("title") or metadata.get("name"))) content_text = normalize_title_match_text(doc.get("page_content", "")) score = 0 for term in query_terms: if term in title_text: score += 3 elif term in content_text: score += 1 if score > 0: substring_results.append((doc, score)) substring_results.sort(key=lambda item: item[1], reverse=True) substring_results = substring_results[:lexical_candidate_count] # RRF (Reciprocal Rank Fusion) 결합 rrf_k = 60 fused_scores = {} fused_results = {} # ChromaDB 점수 반영 for rank, res in enumerate(sorted(all_results, key=lambda x: x['distance'])): doc_id = res["id"] if doc_id not in fused_scores: fused_scores[doc_id] = 0 fused_results[doc_id] = res fused_scores[doc_id] += 1.0 / (rrf_k + rank + 1) # BM25 점수 반영 for rank, (bm25_doc, score) in enumerate(bm25_results): # BM25 인덱스는 같은 청크 id를 보유하므로 문자열 비교 없이 id로 결합한다. matched_id = bm25_doc.get("id") or f"bm25_only_{rank}" if matched_id not in fused_results: fused_results[matched_id] = { "id": matched_id, "text": bm25_doc["page_content"], "metadata": bm25_doc["metadata"], "distance": 0, "relevance": min(score / 10.0, 1.0) # 임의 정규화 } fused_scores[matched_id] = 0 fused_scores[matched_id] += 1.0 / (rrf_k + rank + 1) # 자연어 질의에 포함된 제목구 후보는 BM25 상위권에 없어도 id 기준으로 결합한다. for rank, (title_doc, score) in enumerate(title_results): matched_id = title_doc.get("id") or f"title_only_{rank}" if matched_id not in fused_results: fused_results[matched_id] = { "id": matched_id, "text": title_doc["page_content"], "metadata": title_doc["metadata"], "distance": 0, "relevance": min(score / 50.0, 1.0) } fused_scores[matched_id] = 0 fused_scores[matched_id] += 4.0 / (rrf_k + rank + 1) # BM25가 공백 기반 토큰화로 복합어를 놓칠 때만 부분문자열 후보를 보강한다. for rank, (substring_doc, score) in enumerate(substring_results): doc_text = substring_doc["page_content"] matched_id = substring_doc.get("id") if matched_id not in fused_results: fused_results[matched_id] = { "id": matched_id, "text": doc_text, "metadata": substring_doc["metadata"], "distance": 0, "relevance": min(score / 10.0, 1.0) } fused_scores[matched_id] = 0 fused_scores[matched_id] += 2.0 / (rrf_k + rank + 1) # 제목 그대로 검색하는 질의는 같은 제목의 기준을 우선한다. if normalized_query: for doc_id, res in fused_results.items(): metadata = res.get("metadata", {}) clean_title = clean_title_text(metadata.get("title") or metadata.get("name")) normalized_title = normalize_title_match_text(clean_title) if not normalized_title: continue raw_type = canonical_doc_type(metadata.get("code", ""), metadata.get("doc_type", "")) if clean_normalized_query == normalized_title: fused_scores[doc_id] += 1.0 elif clean_normalized_query in normalized_title or normalized_title in clean_normalized_query: fused_scores[doc_id] += 0.75 title_tokens = set(self.bm25_client._tokenize(clean_title)) query_tokens = set(self.bm25_client._tokenize(query)) meaningful_overlap = { token for token in title_tokens & query_tokens if token in query_signal_tokens } if meaningful_overlap: fused_scores[doc_id] += min( 0.08 * sum(token_importance(token) for token in meaningful_overlap), 0.65, ) content_text = res.get("text", "") content_tokens = set(self.bm25_client._tokenize(content_text[:2500])) content_overlap = { token for token in query_tokens & content_tokens if token in query_signal_tokens } if content_overlap: fused_scores[doc_id] += min( 0.05 * sum(token_importance(token) for token in content_overlap), 0.75, ) rare_title_hits = { token for token in meaningful_overlap if token_importance(token) >= 2.0 } rare_content_hits = { token for token in content_overlap if token_importance(token) >= 2.0 } if rare_title_hits: fused_scores[doc_id] += min( 0.18 * sum(token_importance(token) for token in rare_title_hits), 0.85, ) if rare_content_hits: fused_scores[doc_id] += min( 0.12 * sum(token_importance(token) for token in rare_content_hits), 0.75, ) code = metadata.get("code", "") fused_scores[doc_id] += score_ranking_rules(normalized_query, code, "vector") if wants_excs: if raw_type == "EXCS": fused_scores[doc_id] += 0.75 elif raw_type in {"KCS", "KDS"}: fused_scores[doc_id] -= 0.35 # 최종 RRF 스코어로 정렬 final_results = sorted(list(fused_results.values()), key=lambda x: fused_scores[x["id"]], reverse=True) deduped_results = [] seen_codes = set() for res in final_results: res = dict(res) raw_fusion_score = float(fused_scores.get(res["id"], 0.0)) res["fusion_score"] = raw_fusion_score res["relevance"] = min(raw_fusion_score / 0.08, 1.0) if res["relevance"] < min_relevance: continue metadata = dict(res.get("metadata", {}) or {}) if not doc_type_matches(metadata, allowed_doc_types): continue metadata["doc_type"] = canonical_doc_type(metadata.get("code", ""), metadata.get("doc_type", "")) metadata.setdefault("collection_type", collection_for_doc_type(metadata["doc_type"])) res["metadata"] = metadata code = metadata.get("code") or res.get("id") if code in seen_codes: continue seen_codes.add(code) deduped_results.append(res) # 결과 수 제한 limited_results = deduped_results[:limit] logger.debug(f"하이브리드 RRF 검색 결과: {len(limited_results)}개 항목") return limited_results except Exception as e: logger.error(f"검색 중 오류 발생: {str(e)}") raise VectorDBError(f"검색 중 오류 발생: {str(e)}", "search") def get_document_by_id(self, doc_type: str, doc_id: str) -> Optional[Dict[str, Any]]: """ ID로 문서 가져오기 Args: doc_type: 문서 유형 doc_id: 문서 ID Returns: 문서 정보 또는 None (찾지 못한 경우) """ if doc_type not in self.collections: logger.warning(f"{doc_type} 컬렉션이 없습니다.") return None collection = self.collections[doc_type] try: results = collection.get(ids=[doc_id]) if not results or not results["ids"]: return None return { "id": results["ids"][0], "text": results["documents"][0], "metadata": results["metadatas"][0] } except Exception as e: logger.error(f"ID {doc_id}로 문서 조회 중 오류 발생: {str(e)}") return None def get_collection_stats(self) -> Dict[str, int]: """ 컬렉션 통계 정보 조회 Returns: 컬렉션별 문서 수 """ stats = {} for doc_type, collection in self.collections.items(): try: stats[doc_type] = collection.count() except Exception as e: logger.error(f"{doc_type} 컬렉션 통계 조회 실패: {str(e)}") stats[doc_type] = -1 return stats def similarity_search( self, text: str, doc_types: List[str] = None, limit: int = None ) -> List[Dict[str, Any]]: """ 텍스트 유사성 검색 (자연어 검색 대신 직접 텍스트 비교) Args: text: 비교할 텍스트 doc_types: 검색할 문서 유형 목록 limit: 결과 제한 수 Returns: 유사한 문서 목록 """ # 일반 검색 메서드 활용 return self.search(text, doc_types, limit) def reset_collection(self, doc_type: str = None) -> Dict[str, Any]: """ 컬렉션 데이터 초기화 Args: doc_type: 초기화할 문서 유형 (None이면 모든 컬렉션) Returns: 초기화 결과 통계 """ results = {} if doc_type: if doc_type not in self.collections: raise VectorDBError(f"{doc_type} 컬렉션이 없습니다.", "collection_missing") # 단일 컬렉션 초기화 collection = self.collections[doc_type] count_before = collection.count() collection.delete(where={}) results[doc_type] = {"before": count_before, "after": collection.count()} else: # 모든 컬렉션 초기화 for type_name, collection in self.collections.items(): count_before = collection.count() collection.delete(where={}) results[type_name] = {"before": count_before, "after": collection.count()} return { "status": "success", "reset_stats": results } # 실행 예시 if __name__ == "__main__": # 벡터 DB 초기화 vector_db = KCSCVectorDB() # 기존 데이터 통계 확인 print("벡터 DB 초기 상태:", vector_db.get_collection_stats()) # 테스트 검색 (데이터가 있는 경우) for collection_name, count in vector_db.get_collection_stats().items(): if count > 0: print(f"\n{collection_name} 컬렉션에서 테스트 검색:") results = vector_db.search("철근콘크리트 기둥 설계", [collection_name], 3) for i, result in enumerate(results): print(f"\n--- 결과 {i+1} ---") print(f"문서: {result['metadata']['name']} ({result['metadata']['code']})") print(f"유사도: {result['relevance']:.4f}") print(f"내용 미리보기: {result['text'][:150]}...")