src/vector_db.py

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]}...")