src/utils/vector_db_utils.py

"""
벡터 데이터베이스 관련 유틸리티 함수
"""

import os
import json
import logging
import re
from datetime import datetime
from typing import List, Dict, Any, Tuple, Optional

from src import config
from src.utils.exceptions import VectorDBError, DataProcessingError

logger = config.setup_logger("vector_db_utils")

def get_latest_data_files(prefix: str, data_dir: str = None) -> List[str]:
    """
    특정 접두사와 일치하는 최신 데이터 파일 목록 반환
    
    Args:
        prefix: 파일 접두사 (예: 'KDS_details')
        data_dir: 데이터 디렉토리 경로 (기본값: config.DATA_DIR)
        
    Returns:
        최신 데이터 파일 경로 목록
        
    Raises:
        DataProcessingError: 파일을 찾을 수 없는 경우
    """
    data_dir = data_dir or config.DATA_DIR
    files = {}
    
    try:
        for filename in os.listdir(data_dir):
            if not filename.startswith(prefix) or not filename.endswith('.json'):
                continue
                
            parts = filename.split('_')
            if len(parts) < 2:
                continue
                
            doc_type = parts[0]
            timestamp = parts[-1].replace('.json', '')
            
            if doc_type not in files or timestamp > files[doc_type][1]:
                files[doc_type] = (os.path.join(data_dir, filename), timestamp)
        
        if not files:
            logger.warning(f"'{prefix}' 접두사로 시작하는 파일이 '{data_dir}'에 없습니다.")
            return []
            
        return [file_info[0] for file_info in files.values()]
        
    except Exception as e:
        raise DataProcessingError(
            message=f"데이터 파일 검색 중 오류 발생: {str(e)}",
            data_type=prefix
        )

def load_json_data(filepath: str) -> List[Dict[str, Any]]:
    """
    JSON 파일에서 데이터 로드
    
    Args:
        filepath: JSON 파일 경로
        
    Returns:
        로드된 데이터 리스트
        
    Raises:
        DataProcessingError: 파일 로드 중 오류 발생 시
    """
    try:
        logger.info(f"데이터 파일 로드 중: {filepath}")
        with open(filepath, 'r', encoding='utf-8') as f:
            data = json.load(f)
        logger.info(f"데이터 로드 완료: {len(data)}개 항목")
        return data
    except Exception as e:
        raise DataProcessingError(
            message=f"데이터 파일 '{filepath}' 로드 중 오류 발생: {str(e)}",
            data_type="json"
        )

def determine_optimal_chunk_size(text: str, min_size: int = 200, max_size: int = 2000) -> int:
    """
    문서 내용에 따라 최적의 청크 크기 결정
    
    Args:
        text: 문서 내용 텍스트
        min_size: 최소 청크 크기
        max_size: 최대 청크 크기
        
    Returns:
        최적의 청크 크기
    """
    text_length = len(text)
    
    if text_length <= min_size:
        return text_length
    
    if text_length < 1000:
        return min(text_length, max_size)
    elif text_length < 5000:
        return min(800, max_size)
    elif text_length < 10000:
        return min(1000, max_size)
    else:
        return min(1500, max_size)

def split_text_into_chunks(
    text: str, 
    chunk_size: int = None, 
    overlap: int = None,
    respect_paragraphs: bool = True
) -> List[str]:
    """
    텍스트를 청크로 분할
    
    Args:
        text: 분할할 텍스트
        chunk_size: 청크 크기
        overlap: 청크 간 겹치는 문자 수
        respect_paragraphs: 가능한 한 문단 경계를 존중할지 여부
        
    Returns:
        분할된 텍스트 청크 목록
    """
    if not text:
        return []
    
    chunk_size = chunk_size or config.CHUNK_SIZE
    overlap = overlap or config.CHUNK_OVERLAP
    
    if chunk_size > len(text):
        return [text]
    
    chunks = []
    
    if respect_paragraphs:
        paragraphs = text.split('\n\n')
        current_chunk = ""
        
        for para in paragraphs:
            if not para.strip():
                continue
                
            if len(current_chunk) + len(para) > chunk_size:
                if current_chunk:
                    chunks.append(current_chunk)
                
                if len(para) > chunk_size:
                    para_chunks = _split_by_chars(para, chunk_size, overlap)
                    chunks.extend(para_chunks)
                    current_chunk = para_chunks[-1][-overlap:] if overlap > 0 and para_chunks else ""
                else:
                    current_chunk = para
            else:
                if current_chunk:
                    current_chunk += "\n\n" + para
                else:
                    current_chunk = para
        
        if current_chunk:
            chunks.append(current_chunk)
    else:
        chunks = _split_by_chars(text, chunk_size, overlap)
    
    return chunks

def _split_by_chars(text: str, chunk_size: int, overlap: int) -> List[str]:
    """
    텍스트를 문자 단위로 청크로 분할
    """
    chunks = []
    start = 0
    
    while start < len(text):
        end = min(start + chunk_size, len(text))
        
        if end < len(text):
            sentence_end = max(
                text.rfind('. ', start, end),
                text.rfind('.\n', start, end),
                text.rfind('? ', start, end),
                text.rfind('?\n', start, end),
                text.rfind('! ', start, end),
                text.rfind('!\n', start, end)
            )
            
            if sentence_end == -1 or (sentence_end - start) < (chunk_size // 2):
                sentence_end = text.rfind(' ', start, end)
            
            if sentence_end != -1 and (sentence_end - start) >= (chunk_size // 3):
                end = sentence_end + 1
        
        chunk = text[start:end].strip()
        if chunk:
            chunks.append(chunk)
        
        start = max(start + 1, end - overlap)
        
        if start >= end:
            start = end
    
    return chunks

def check_index_status(collection_name: str, num_docs: int) -> Tuple[bool, str]:
    """
    인덱스 상태 확인
    
    Args:
        collection_name: 컬렉션 이름
        num_docs: 문서 수
        
    Returns:
        (성공 여부, 상태 메시지) 튜플
    """
    if num_docs == 0:
        return False, f"'{collection_name}' 컬렉션에 문서가 없습니다."
    return True, f"'{collection_name}' 컬렉션에 {num_docs}개 문서가 있습니다."


def structured_chunking(text: str, base_metadata: Dict[str, Any] = None, chunk_size: int = 2000) -> List[Dict[str, Any]]:
    """
    KDS/KCS 표준 문서를 섹션 단위로 청킹합니다.
    
    전략:
    - 최상위 조항(1., 2., 3. 등 한~두 자리 숫자 + 점 + 공백)을 섹션 경계로 인식
    - 섹션이 chunk_size를 초과할 경우에만 추가 분할
    - 문서당 보통 5~15개 청크 생성 (목표: 전체 ~5~10만 청크)
    
    Args:
        text: 분할할 표준 문서 전체 텍스트
        base_metadata: 기본 메타데이터 (문서명, 코드 등)
        chunk_size: 최대 청크 크기 (기본 2000자)
        
    Returns:
        메타데이터가 포함된 청크 딕셔너리 목록
    """
    if not text:
        return []
        
    base_metadata = base_metadata or {}
    chunks = []
    
    # 최상위 조항 경계: 줄 시작에 "1. " "10. " 형식의 한두 자리 숫자
    top_section_pattern = re.compile(r"^(\d{1,2})\.\s+\S")
    
    lines = text.split('\n')
    
    # 최상위 섹션 경계 찾기
    section_starts: List[Tuple[int, str]] = []
    for i, line in enumerate(lines):
        stripped = line.strip()
        if stripped and top_section_pattern.match(stripped):
            section_starts.append((i, stripped))
    
    def make_chunk(content: str, section_heading: str, chunk_idx: int = 0) -> Dict[str, Any]:
        prefix = f"[{section_heading}]\n" if section_heading else ""
        return {
            "page_content": (prefix + content).strip(),
            "metadata": {
                **base_metadata,
                "element_type": "section",
                "heading_context": section_heading,
                "chunk_index": chunk_idx
            }
        }
    
    def add_section_as_chunks(section_text: str, heading: str):
        """섹션이 chunk_size 이하면 1청크, 초과면 단락 단위 분할"""
        if not section_text.strip():
            return
        if len(section_text) <= chunk_size:
            chunks.append(make_chunk(section_text, heading))
        else:
            sub = split_text_into_chunks(section_text, chunk_size=chunk_size, overlap=200, respect_paragraphs=True)
            for i, sc in enumerate(sub):
                chunks.append(make_chunk(sc, heading, chunk_idx=i))
    
    if not section_starts:
        # 섹션 구분이 없으면 전체를 chunk_size 단위로 분할
        sub = split_text_into_chunks(text, chunk_size=chunk_size, overlap=200, respect_paragraphs=True)
        for i, sc in enumerate(sub):
            chunks.append(make_chunk(sc, "", chunk_idx=i))
        return chunks
    
    # 섹션 이전 머리말 처리
    if section_starts[0][0] > 0:
        preamble = '\n'.join(lines[:section_starts[0][0]]).strip()
        if preamble:
            add_section_as_chunks(preamble, "머리말")
    
    # 각 섹션 처리
    for idx, (start_line, heading) in enumerate(section_starts):
        end_line = section_starts[idx + 1][0] if idx + 1 < len(section_starts) else len(lines)
        section_text = '\n'.join(lines[start_line:end_line]).strip()
        if section_text:
            add_section_as_chunks(section_text, heading)
    
    return chunks