Blame for chunking_methods.py

"""
MIT License
Autor: Fernando Lafuente
Proyecto: Ejercicio opcional NLP - Chunking y Recuperación
Descripción: Implementación de múltiples métodos de chunking para segmentación textual

Este módulo implementa tres estrategias diferentes de chunking (segmentación textual):
1. Chunking fijo: División por tamaño fijo con solapamiento
2. Chunking por oraciones: División en oraciones naturales agrupadas
3. Chunking adaptativo: División basada en similitud semántica entre oraciones

Cada método tiene ventajas específicas según el tipo de documento y la aplicación.
El chunking adaptativo utiliza embeddings para agrupar oraciones semánticamente
relacionadas, lo que puede mejorar la coherencia de los chunks resultantes.
"""

import logging
import re
from typing import List, Tuple
from abc import ABC, abstractmethod

import numpy as np
from sklearn.metrics.pairwise import cosine_similarity

import settings
from utils_logging import setup_logging, log_operation_time, validate_text_length


# Configurar logger específico para este módulo
logger = setup_logging()


class ChunkingStrategy(ABC):
    """
    Clase abstracta base para todas las estrategias de chunking.
    
    Define la interfaz común que deben implementar todos los métodos de chunking.
    Esto permite intercambiar estrategias fácilmente y mantener código limpio
    siguiendo el patrón Strategy de diseño.
    """
    
    @abstractmethod
    def create_chunks(self, text: str) -> List[str]:
        """
        Método abstracto para crear chunks desde un texto.
        
        Args:
            text: Texto a segmentar
            
        Returns:
            List[str]: Lista de chunks generados
        """
        pass
    
    @abstractmethod
    def get_strategy_name(self) -> str:
        """
        Método abstracto para obtener el nombre de la estrategia.
        
        Returns:
            str: Nombre identificativo de la estrategia
        """
        pass


class FixedChunking(ChunkingStrategy):
    """
    Implementa chunking por tamaño fijo con solapamiento.
    
    Este método divide el texto en chunks de tamaño fijo (en caracteres),
    con un solapamiento configurable entre chunks consecutivos. Es el método
    más simple y predecible, útil cuando se requiere uniformidad en el tamaño.
    
    Ventajas:
    - Predecible y rápido
    - Chunks de tamaño uniforme
    - Fácil de implementar y debuggear
    
    Desventajas:
    - Puede cortar oraciones o conceptos a la mitad
    - No considera la estructura semántica del texto
    """
    
    def __init__(self, 
                 chunk_size: int = settings.FIXED_CHUNK_SIZE,
                 overlap: int = settings.FIXED_CHUNK_OVERLAP):
        """
        Inicializa el chunking fijo con parámetros configurables.
        
        Args:
            chunk_size: Tamaño en caracteres de cada chunk
            overlap: Número de caracteres de solapamiento entre chunks
        """
        self.chunk_size = chunk_size
        self.overlap = overlap
        self.logger = logging.getLogger("chunking_system.fixed_chunking")
        
        # Validar parámetros
        if self.overlap >= self.chunk_size:
            raise ValueError("El solapamiento debe ser menor que el tamaño del chunk")
    
    @log_operation_time
    def create_chunks(self, text: str) -> List[str]:
        """
        Crea chunks de tamaño fijo con solapamiento del texto de entrada.
        
        El algoritmo avanza por el texto con pasos de (chunk_size - overlap),
        tomando ventanas de tamaño chunk_size. Esto asegura que la información
        en los límites de chunks se preserve en múltiples chunks.
        
        Args:
            text: Texto a segmentar
            
        Returns:
            List[str]: Lista de chunks de tamaño aproximadamente uniforme
        """
        if not text or not validate_text_length(text):
            self.logger.warning("Texto vacío o demasiado corto para chunking")
            return []
        
        chunks = []
        start = 0
        step = self.chunk_size - self.overlap
        
        while start < len(text):
            end = min(start + self.chunk_size, len(text))
            chunk = text[start:end].strip()
            
            # Solo agregar chunks que tengan contenido válido
            if validate_text_length(chunk):
                chunks.append(chunk)
                self.logger.debug(f"Chunk creado: posición {start}-{end}, "
                                f"longitud {len(chunk)} caracteres")
            
            # Salir si hemos llegado al final
            if end == len(text):
                break
                
            start += step
        
        self.logger.info(f"Chunking fijo completado: {len(chunks)} chunks generados")
        return chunks
    
    def get_strategy_name(self) -> str:
        return f"fixed_chunking_{self.chunk_size}_{self.overlap}"


class SentenceChunking(ChunkingStrategy):
    """
    Implementa chunking basado en oraciones y párrafos naturales.
    
    Este método respeta la estructura natural del texto, dividiendo en oraciones
    y agrupándolas hasta alcanzar un tamaño objetivo. Preserva la coherencia
    semántica mejor que el chunking fijo al no cortar oraciones.
    
    Ventajas:
    - Respeta límites naturales de oraciones
    - Mantiene coherencia semántica local
    - Más legible para humanos
    
    Desventajas:
    - Chunks de tamaño variable
    - Puede generar chunks muy largos o muy cortos
    - Depende de la calidad de detección de oraciones
    """
    
    def __init__(self,
                 min_length: int = settings.SENTENCE_MIN_LENGTH,
                 max_chunk_size: int = settings.SENTENCE_MAX_CHUNK_SIZE):
        """
        Inicializa el chunking por oraciones con parámetros configurables.
        
        Args:
            min_length: Longitud mínima para considerar una oración válida
            max_chunk_size: Tamaño máximo permitido para un chunk
        """
        self.min_length = min_length
        self.max_chunk_size = max_chunk_size
        self.logger = logging.getLogger("chunking_system.sentence_chunking")
        
        # Patrón regex para detectar finales de oración
        # Busca punto, signo de interrogación o exclamación seguido de espacio o final de línea
        self.sentence_pattern = re.compile(r'[.!?]+\s+|[.!?]+$')
    
    @log_operation_time
    def create_chunks(self, text: str) -> List[str]:
        """
        Crea chunks agrupando oraciones hasta alcanzar el tamaño objetivo.
        
        El algoritmo identifica oraciones usando expresiones regulares,
        las valida y las agrupa secuencialmente hasta alcanzar el tamaño
        máximo permitido. Cada chunk contiene oraciones completas.
        
        Args:
            text: Texto a segmentar
            
        Returns:
            List[str]: Lista de chunks basados en oraciones
        """
        if not text or not validate_text_length(text):
            self.logger.warning("Texto vacío o demasiado corto para chunking")
            return []
        
        # Dividir texto en oraciones
        sentences = self._split_sentences(text)
        
        if not sentences:
            self.logger.warning("No se pudieron extraer oraciones válidas")
            return []
        
        # Agrupar oraciones en chunks
        chunks = self._group_sentences_into_chunks(sentences)
        
        self.logger.info(f"Chunking por oraciones completado: {len(chunks)} chunks "
                        f"generados desde {len(sentences)} oraciones")
        return chunks
    
    def _split_sentences(self, text: str) -> List[str]:
        """
        Divide el texto en oraciones individuales usando regex.
        
        Este método utiliza patrones regex para identificar finales de oración
        y divide el texto en consecuencia. Aplica validaciones para asegurar
        que las oraciones resultantes sean significativas.
        
        Args:
            text: Texto a dividir en oraciones
            
        Returns:
            List[str]: Lista de oraciones válidas
        """
        # Dividir por patrones de final de oración
        potential_sentences = self.sentence_pattern.split(text)
        
        sentences = []
        for sentence in potential_sentences:
            # Limpiar y validar cada oración
            clean_sentence = sentence.strip()
            
            if validate_text_length(clean_sentence, self.min_length):
                sentences.append(clean_sentence)
                self.logger.debug(f"Oración válida: {len(clean_sentence)} caracteres")
        
        return sentences
    
    def _group_sentences_into_chunks(self, sentences: List[str]) -> List[str]:
        """
        Agrupa oraciones secuencialmente en chunks respetando tamaño máximo.
        
        El algoritmo va agregando oraciones a un chunk actual hasta que
        agregar la siguiente oración excedería el tamaño máximo. En ese
        punto, finaliza el chunk actual e inicia uno nuevo.
        
        Args:
            sentences: Lista de oraciones a agrupar
            
        Returns:
            List[str]: Lista de chunks formados por oraciones agrupadas
        """
        chunks = []
        current_chunk = []
        current_length = 0
        
        for sentence in sentences:
            sentence_length = len(sentence)
            
            # Verificar si agregar esta oración excedería el tamaño máximo
            if (current_length + sentence_length > self.max_chunk_size and 
                current_chunk):
                
                # Finalizar chunk actual
                chunk_text = ' '.join(current_chunk)
                chunks.append(chunk_text)
                
                self.logger.debug(f"Chunk completado: {len(current_chunk)} oraciones, "
                                f"{current_length} caracteres")
                
                # Iniciar nuevo chunk
                current_chunk = [sentence]
                current_length = sentence_length
            else:
                # Agregar oración al chunk actual
                current_chunk.append(sentence)
                current_length += sentence_length + 1  # +1 por el espacio separador
        
        # Agregar último chunk si tiene contenido
        if current_chunk:
            chunk_text = ' '.join(current_chunk)
            chunks.append(chunk_text)
            self.logger.debug(f"Chunk final: {len(current_chunk)} oraciones, "
                            f"{current_length} caracteres")
        
        return chunks
    
    def get_strategy_name(self) -> str:
        return f"sentence_chunking_{self.min_length}_{self.max_chunk_size}"


class AdaptiveChunking(ChunkingStrategy):
    """
    Implementa chunking adaptativo basado en similitud semántica.
    
    Este método utiliza embeddings para calcular la similitud entre oraciones
    consecutivas y agrupa aquellas que son semánticamente similares. Requiere
    un modelo de embeddings externo pero puede producir chunks más coherentes
    temáticamente.
    
    Ventajas:
    - Chunks semánticamente coherentes
    - Adaptativo al contenido del texto
    - Preserva temas y conceptos relacionados
    
    Desventajas:
    - Computacionalmente más costoso
    - Requiere modelo de embeddings
    - Tamaños de chunk muy variables
    """
    
    def __init__(self,
                 similarity_threshold: float = settings.ADAPTIVE_SIMILARITY_THRESHOLD,
                 min_chunk_size: int = settings.ADAPTIVE_MIN_CHUNK_SIZE,
                 max_chunk_size: int = settings.ADAPTIVE_MAX_CHUNK_SIZE):
        """
        Inicializa el chunking adaptativo con parámetros configurables.
        
        Args:
            similarity_threshold: Umbral de similitud coseno para agrupar oraciones
            min_chunk_size: Tamaño mínimo de chunk en caracteres
            max_chunk_size: Tamaño máximo de chunk en caracteres
        """
        self.similarity_threshold = similarity_threshold
        self.min_chunk_size = min_chunk_size
        self.max_chunk_size = max_chunk_size
        self.logger = logging.getLogger("chunking_system.adaptive_chunking")
        
        # Validar umbral de similitud
        if not 0 <= similarity_threshold <= 1:
            raise ValueError("similarity_threshold debe estar entre 0 y 1")
    
    @log_operation_time
    def create_chunks(self, text: str, embeddings_generator=None) -> List[str]:
        """
        Crea chunks basados en similitud semántica entre oraciones.
        
        Este algoritmo:
        1. Divide el texto en oraciones
        2. Genera embeddings para cada oración
        3. Calcula similitud coseno entre oraciones consecutivas
        4. Agrupa oraciones con similitud superior al umbral
        5. Respeta límites de tamaño mínimo y máximo
        
        Args:
            text: Texto a segmentar
            embeddings_generator: Función para generar embeddings de oraciones
            
        Returns:
            List[str]: Lista de chunks semánticamente coherentes
        """
        if not text or not validate_text_length(text):
            self.logger.warning("Texto vacío o demasiado corto para chunking")
            return []
        
        if embeddings_generator is None:
            self.logger.warning("No se proporcionó generador de embeddings. "
                              "Fallback a chunking por oraciones")
            fallback = SentenceChunking()
            return fallback.create_chunks(text)
        
        # Dividir en oraciones usando la misma lógica que SentenceChunking
        sentence_chunker = SentenceChunking()
        sentences = sentence_chunker._split_sentences(text)
        
        if len(sentences) < 2:
            self.logger.info("Pocas oraciones para chunking adaptativo")
            return sentences
        
        # Generar embeddings para cada oración
        sentence_embeddings = self._generate_sentence_embeddings(
            sentences, embeddings_generator
        )
        
        # Calcular similitudes y agrupar
        chunks = self._group_by_similarity(sentences, sentence_embeddings)
        
        self.logger.info(f"Chunking adaptativo completado: {len(chunks)} chunks "
                        f"generados desde {len(sentences)} oraciones")
        return chunks
    
    def _generate_sentence_embeddings(self, 
                                    sentences: List[str], 
                                    embeddings_generator) -> np.ndarray:
        """
        Genera embeddings para una lista de oraciones.
        
        Args:
            sentences: Lista de oraciones
            embeddings_generator: Función que genera embeddings
            
        Returns:
            np.ndarray: Array de embeddings de forma (n_sentences, embedding_dim)
        """
        embeddings = []
        
        for i, sentence in enumerate(sentences):
            try:
                embedding = embeddings_generator(sentence)
                embeddings.append(embedding)
                self.logger.debug(f"Embedding generado para oración {i+1}")
            except Exception as e:
                self.logger.error(f"Error generando embedding para oración {i+1}: {e}")
                # Usar vector cero como fallback
                embeddings.append(np.zeros(settings.EMBEDDING_DIMENSION))
        
        return np.array(embeddings)
    
    def _group_by_similarity(self, 
                           sentences: List[str], 
                           embeddings: np.ndarray) -> List[str]:
        """
        Agrupa oraciones basándose en similitud de sus embeddings.
        
        El algoritmo utiliza similitud coseno para determinar si dos oraciones
        consecutivas deben ser agrupadas en el mismo chunk. La fórmula de
        similitud coseno es: sim(A,B) = (A·B) / (||A|| × ||B||)
        
        Args:
            sentences: Lista de oraciones
            embeddings: Embeddings correspondientes a cada oración
            
        Returns:
            List[str]: Lista de chunks agrupados por similitud
        """
        if len(sentences) != len(embeddings):
            raise ValueError("Número de oraciones y embeddings no coincide")
        
        chunks = []
        current_chunk = [sentences[0]]
        current_length = len(sentences[0])
        
        for i in range(1, len(sentences)):
            # Calcular similitud coseno entre oración actual y anterior
            similarity = cosine_similarity(
                embeddings[i-1].reshape(1, -1),
                embeddings[i].reshape(1, -1)
            )[0, 0]
            
            sentence_length = len(sentences[i])
            new_length = current_length + sentence_length + 1
            
            # Decidir si agregar a chunk actual o crear nuevo chunk
            should_group = (
                similarity >= self.similarity_threshold and
                new_length <= self.max_chunk_size
            )
            
            if should_group:
                # Agregar a chunk actual
                current_chunk.append(sentences[i])
                current_length = new_length
                self.logger.debug(f"Oraciones {i} y {i+1} agrupadas "
                                f"(similitud: {similarity:.4f})")
            else:
                # Finalizar chunk actual. Se añade siempre para asegurar
                # que el documento quede completamente cubierto, incluso
                # si el chunk es menor que el tamaño mínimo configurado.
                chunk_text = ' '.join(current_chunk)
                chunks.append(chunk_text)
                self.logger.debug(f"Chunk completado: {len(current_chunk)} oraciones, "
                                f"{current_length} caracteres")
                
                # Iniciar nuevo chunk
                current_chunk = [sentences[i]]
                current_length = sentence_length
                
                if similarity < self.similarity_threshold:
                    self.logger.debug(f"Nueva sección detectada en oración {i+1} "
                                    f"(similitud: {similarity:.4f})")
        
        # Agregar último chunk (siempre que tenga contenido)
        if current_chunk:
            chunk_text = ' '.join(current_chunk)
            chunks.append(chunk_text)
            self.logger.debug(f"Chunk final: {len(current_chunk)} oraciones, "
                            f"{current_length} caracteres")
        
        return chunks
    
    def get_strategy_name(self) -> str:
        return f"adaptive_chunking_{self.similarity_threshold}_{self.min_chunk_size}_{self.max_chunk_size}"


class ChunkingManager:
    """
    Gestor principal para seleccionar y ejecutar estrategias de chunking.
    
    Esta clase actúa como fachada para todos los métodos de chunking,
    permitiendo seleccionar la estrategia desde configuración y proporcionando
    una interfaz unificada para el sistema principal.
    """
    
    def __init__(self):
        """
        Inicializa el gestor con la estrategia configurada en settings.
        """
        self.logger = logging.getLogger("chunking_system.manager")
        self.strategy = self._create_strategy()
    
    def _create_strategy(self) -> ChunkingStrategy:
        """
        Crea la estrategia de chunking según la configuración.
        
        Returns:
            ChunkingStrategy: Instancia de la estrategia configurada
            
        Raises:
            ValueError: Si el método configurado no es válido
        """
        method = settings.CHUNKING_METHOD.lower()
        
        if method == "fixed":
            return FixedChunking()
        elif method == "sentence":
            return SentenceChunking()
        elif method == "adaptive":
            return AdaptiveChunking()
        else:
            raise ValueError(f"Método de chunking no válido: {method}")
    
    @log_operation_time
    def create_chunks(self, text: str, embeddings_generator=None) -> List[str]:
        """
        Crea chunks usando la estrategia configurada.
        
        Args:
            text: Texto a segmentar
            embeddings_generator: Generador de embeddings (requerido para adaptativo)
            
        Returns:
            List[str]: Lista de chunks generados
        """
        self.logger.info(f"Iniciando chunking con estrategia: {self.strategy.get_strategy_name()}")
        
        # Para chunking adaptativo, pasar el generador de embeddings
        if isinstance(self.strategy, AdaptiveChunking):
            chunks = self.strategy.create_chunks(text, embeddings_generator)
        else:
            chunks = self.strategy.create_chunks(text)
        
        self.logger.info(f"Chunking completado: {len(chunks)} chunks generados")
        
        # Log estadísticas de chunks
        if chunks:
            lengths = [len(chunk) for chunk in chunks]
            self.logger.info(f"Estadísticas de chunks - "
                           f"Min: {min(lengths)}, Max: {max(lengths)}, "
                           f"Promedio: {sum(lengths)/len(lengths):.1f} caracteres")
        
        return chunks
    
    def get_current_strategy_name(self) -> str:
        """
        Obtiene el nombre de la estrategia actualmente configurada.
        
        Returns:
            str: Nombre de la estrategia actual
        """
        return self.strategy.get_strategy_name()