import os import json import numpy as np from sentence_transformers import SentenceTransformer import chromadb from chromadb.config import Settings import logging os.makedirs("logs", exist_ok=True) logging.basicConfig(filename='logs/vector_db.log', level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') class KCSCVectorDB: def __init__(self, data_dir="data", db_dir="vector_db"): self.data_dir = data_dir self.db_dir = db_dir os.makedirs(self.db_dir, exist_ok=True) # 임베딩 모델 로드 (다국어 모델 사용) self.model = SentenceTransformer('jhgan/ko-sroberta-multitask') # Chroma DB 클라이언트 초기화 self.client = chromadb.PersistentClient(path=self.db_dir, settings=Settings(allow_reset=True)) # 각 문서 타입별 컬렉션 생성 self.collections = {} 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} 설계기준 문서"} ) except Exception as e: logging.error(f"{doc_type} 컬렉션 생성 중 오류: {str(e)}") def load_processed_docs(self): """처리된 문서 데이터 로드""" try: file_path = os.path.join(self.data_dir, "processed_docs.json") with open(file_path, 'r', encoding='utf-8') as f: return json.load(f) except Exception as e: logging.error(f"처리된 문서 로드 중 오류: {str(e)}") return [] def create_document_chunks(self, docs, chunk_size=1000, overlap=200): """문서를 청크로 분할""" chunks = [] for doc in docs: content = doc.get('content', '') if not content or len(content) <= chunk_size: # 내용이 없거나 청크 크기보다 작은 경우 그대로 사용 chunk_doc = doc.copy() chunk_doc['chunk_id'] = f"{doc['id']}-0" chunks.append(chunk_doc) continue # 긴 문서를 청크로 분할 for i in range(0, len(content), chunk_size - overlap): chunk_content = content[i:i + chunk_size] if len(chunk_content) < 100: # 너무 작은 청크는 건너뛰기 continue chunk_doc = doc.copy() chunk_doc['content'] = chunk_content chunk_doc['chunk_id'] = f"{doc['id']}-{i//(chunk_size-overlap)}" chunks.append(chunk_doc) logging.info(f"총 {len(docs)}개 문서에서 {len(chunks)}개의 청크 생성됨") return chunks def build_index(self): """벡터 데이터베이스 구축""" docs = self.load_processed_docs() if not docs: logging.error("문서를 로드할 수 없습니다.") return False # 문서를 청크로 분할 chunks = self.create_document_chunks(docs) # 문서 타입별로 그룹화 doc_type_groups = {} for chunk in chunks: doc_type = chunk.get('doc_type', 'unknown') if doc_type not in doc_type_groups: doc_type_groups[doc_type] = [] doc_type_groups[doc_type].append(chunk) # 각 문서 타입별로 인덱싱 for doc_type, type_chunks in doc_type_groups.items(): if doc_type not in self.collections: logging.warning(f"{doc_type} 컬렉션이 없습니다. 건너뜁니다.") continue collection = self.collections[doc_type] # 기존 데이터 확인 및 필요시 초기화 if collection.count() > 0: logging.info(f"{doc_type} 컬렉션 초기화 (기존 {collection.count()}개 항목)") collection.delete(where={}) # 청크 데이터 준비 ids = [] documents = [] metadatas = [] embeddings = [] for chunk in type_chunks: # 텍스트 준비 (제목 + 내용) text_for_embedding = f"{chunk.get('name', '')} {chunk.get('title', '')}: {chunk.get('content', '')}" # 텍스트가 비어있으면 건너뜀 if not text_for_embedding.strip(): continue # 임베딩 생성 embedding = self.model.encode(text_for_embedding) ids.append(chunk['chunk_id']) documents.append(text_for_embedding) # 메타데이터 준비 (검색 결과에서 필요한 정보) metadata = { "id": chunk['id'], "code": chunk.get('code', ''), "full_code": chunk.get('full_code', ''), "name": chunk.get('name', ''), "title": chunk.get('title', ''), "doc_type": doc_type, "version": chunk.get('version', '') } metadatas.append(metadata) embeddings.append(embedding.tolist()) # 데이터 추가 if ids: collection.add( ids=ids, documents=documents, metadatas=metadatas, embeddings=embeddings ) logging.info(f"{doc_type} 컬렉션에 {len(ids)}개 항목 추가됨") return True def search(self, query, doc_types=None, limit=5): """자연어 쿼리로 관련 문서 검색""" if not doc_types: doc_types = list(self.collections.keys()) # 쿼리 임베딩 query_embedding = self.model.encode(query) all_results = [] # 각 컬렉션에서 검색 for doc_type in doc_types: if doc_type not in self.collections: continue collection = self.collections[doc_type] results = collection.query( query_embeddings=[query_embedding.tolist()], n_results=limit ) # 결과 형식 변환 for i in range(len(results['ids'][0])): result = { "id": results['ids'][0][i], "text": results['documents'][0][i], "metadata": results['metadatas'][0][i], "distance": float(results['distances'][0][i]) } all_results.append(result) # 전체 결과를 거리(유사도)를 기준으로 정렬 all_results = sorted(all_results, key=lambda x: x['distance']) return all_results[:limit] # 사용 예시 if __name__ == "__main__": vector_db = KCSCVectorDB() # 벡터 DB 구축 vector_db.build_index() # 검색 테스트 results = vector_db.search("철근콘크리트 기둥의 설계 방법") for i, result in enumerate(results): print(f"\n--- 결과 {i+1} ---") print(f"문서: {result['metadata']['name']} ({result['metadata']['code']})") print(f"유사도: {1 - result['distance']:.4f}") print(f"내용 미리보기: {result['text'][:200]}...")