services/chunker.py

"""
텍스트를 적절한 크기로 분할 (Chunking)
"""
from typing import List, Dict
from langchain.text_splitter import RecursiveCharacterTextSplitter
from loguru import logger


class TextChunker:
    """텍스트를 의미 있는 청크로 분할하는 클래스"""

    def __init__(self, chunk_size: int = 1000, chunk_overlap: int = 200):
        """
        Args:
            chunk_size: 각 청크의 최대 문자 수
            chunk_overlap: 청크 간 겹치는 문자 수 (컨텍스트 보존)
        """
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap

        # LangChain의 RecursiveCharacterTextSplitter 사용
        # 문단, 문장, 단어 순서로 분할 시도
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            length_function=len,
            separators=["\n\n", "\n", ". ", " ", ""]
        )

    def chunk_document(self, doc_data: Dict[str, any]) -> List[Dict[str, any]]:
        """
        단일 문서를 여러 청크로 분할

        Args:
            doc_data: PDF processor에서 추출한 문서 데이터

        Returns:
            List of chunks with metadata
        """
        text = doc_data['text']
        chunks = self.text_splitter.split_text(text)

        chunked_docs = []
        for i, chunk in enumerate(chunks):
            chunked_docs.append({
                'text': chunk,
                'chunk_id': i,
                'source_filename': doc_data['filename'],
                'source_filepath': doc_data['filepath'],
                'total_chunks': len(chunks),
                'metadata': doc_data['metadata'],
                'page_count': doc_data['page_count']
            })

        return chunked_docs

    def chunk_all_documents(self, documents: List[Dict[str, any]]) -> List[Dict[str, any]]:
        """
        여러 문서를 모두 청크로 분할

        Args:
            documents: PDF processor에서 추출한 문서 리스트

        Returns:
            List of all chunks from all documents
        """
        all_chunks = []

        logger.info(f"Chunking {len(documents)} documents...")

        for doc in documents:
            chunks = self.chunk_document(doc)
            all_chunks.extend(chunks)

        logger.info(f"Created {len(all_chunks)} chunks from {len(documents)} documents")
        logger.info(f"Average {len(all_chunks) / len(documents):.1f} chunks per document")

        return all_chunks

    def get_chunk_statistics(self, chunks: List[Dict[str, any]]) -> Dict[str, any]:
        """청크 통계 정보"""
        if not chunks:
            return {}

        chunk_lengths = [len(chunk['text']) for chunk in chunks]

        return {
            'total_chunks': len(chunks),
            'avg_chunk_length': sum(chunk_lengths) / len(chunks),
            'min_chunk_length': min(chunk_lengths),
            'max_chunk_length': max(chunk_lengths),
            'total_characters': sum(chunk_lengths),
        }