Primera version de chunkeo completo crud

backend/app/utils/chunking/pdf_extractor.py

@@ -0,0 +1,299 @@
+"""
+Extractor optimizado de PDFs con soporte para BytesIO y procesamiento paralelo.
+Adaptado para trabajar con Azure Blob Storage sin archivos temporales.
+"""
+import logging
+import os
+import time
+import hashlib
+from typing import List, Optional, Dict, BinaryIO
+from io import BytesIO
+from concurrent.futures import ThreadPoolExecutor
+
+from langchain_core.documents import Document
+from pypdf import PdfReader
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from pdf2image import convert_from_bytes
+
+from .token_manager import TokenManager
+from .chunk_processor import OptimizedChunkProcessor
+from .gemini_client import GeminiClient
+
+logger = logging.getLogger(__name__)
+
+
+class OptimizedPDFExtractor:
+    """Extractor optimizado de PDFs con soporte para BytesIO"""
+
+    def __init__(
+        self,
+        max_tokens: int = 1000,
+        target_tokens: int = 800,
+        gemini_client: Optional[GeminiClient] = None,
+        custom_instructions: str = "",
+        extract_images: bool = False,  # Por defecto deshabilitado según requerimientos
+        max_workers: int = 4
+    ):
+        """
+        Inicializa el extractor de PDFs.
+
+        Args:
+            max_tokens: Límite máximo de tokens por chunk
+            target_tokens: Tokens objetivo para chunks
+            gemini_client: Cliente de Gemini (opcional)
+            custom_instructions: Instrucciones adicionales para optimización
+            extract_images: Si True, extrae páginas con formato especial como imágenes
+            max_workers: Número máximo de workers para procesamiento paralelo
+        """
+        self.client = gemini_client
+        self.max_workers = max_workers
+        self.token_manager = TokenManager()
+        self.custom_instructions = custom_instructions
+        self.extract_images = extract_images
+        self._format_cache = {}
+
+        self.chunk_processor = OptimizedChunkProcessor(
+            max_tokens=max_tokens,
+            target_tokens=target_tokens,
+            gemini_client=gemini_client,
+            custom_instructions=custom_instructions
+        )
+
+    def detect_special_format_batch(self, chunks: List[Document]) -> Dict[int, bool]:
+        """
+        Detecta chunks con formatos especiales (tablas, diagramas, etc.) en lote.
+
+        Args:
+            chunks: Lista de chunks a analizar
+
+        Returns:
+            Diccionario con índices de chunks y si tienen formato especial
+        """
+        results = {}
+
+        chunks_to_process = []
+        for i, chunk in enumerate(chunks):
+            cache_key = hashlib.md5(chunk.page_content.encode()).hexdigest()[:16]
+            if cache_key in self._format_cache:
+                results[i] = self._format_cache[cache_key]
+            else:
+                chunks_to_process.append((i, chunk, cache_key))
+
+        if not chunks_to_process:
+            return results
+
+        logger.info(f"Analizando {len(chunks_to_process)} chunks para formatos especiales...")
+
+        if self.client and len(chunks_to_process) > 1:
+            with ThreadPoolExecutor(max_workers=min(self.max_workers, len(chunks_to_process))) as executor:
+                futures = {
+                    executor.submit(self._detect_single_format, chunk): (i, cache_key)
+                    for i, chunk, cache_key in chunks_to_process
+                }
+
+                for future in futures:
+                    i, cache_key = futures[future]
+                    try:
+                        result = future.result()
+                        results[i] = result
+                        self._format_cache[cache_key] = result
+                    except Exception as e:
+                        logger.error(f"Error procesando chunk {i}: {e}")
+                        results[i] = False
+                        self._format_cache[cache_key] = False
+        else:
+            for i, chunk, cache_key in chunks_to_process:
+                result = self._detect_single_format(chunk)
+                results[i] = result
+                self._format_cache[cache_key] = result
+
+        return results
+
+    def _detect_single_format(self, chunk: Document) -> bool:
+        """Detecta formato especial en un chunk individual."""
+        if not self.client:
+            content = chunk.page_content
+            table_indicators = ['│', '├', '┼', '┤', '┬', '┴', '|', '+', '-']
+            has_table_chars = any(char in content for char in table_indicators)
+            has_multiple_columns = content.count('\t') > 10 or content.count('  ') > 20
+            return has_table_chars or has_multiple_columns
+
+        try:
+            prompt = f"""¿Contiene este texto tablas estructuradas, diagramas ASCII, o elementos que requieren formato especial?
+
+Responde SOLO 'SI' o 'NO'.
+
+Texto:
+{chunk.page_content[:1000]}"""
+
+            response = self.client.generate_content(prompt)
+            return response.strip().upper() == 'SI'
+
+        except Exception as e:
+            logger.error(f"Error detectando formato: {e}")
+            return False
+
+    def process_pdf_from_bytes(
+        self,
+        pdf_bytes: bytes,
+        file_name: str,
+        chunk_size: int = 1000,
+        chunk_overlap: int = 200,
+        merge_related: bool = True
+    ) -> List[Document]:
+        """
+        Procesa un PDF desde bytes (BytesIO).
+
+        Args:
+            pdf_bytes: Contenido del PDF en bytes
+            file_name: Nombre del archivo PDF
+            chunk_size: Tamaño del chunk
+            chunk_overlap: Solapamiento entre chunks
+            merge_related: Si True, intenta unir chunks relacionados
+
+        Returns:
+            Lista de documentos procesados
+        """
+        overall_start = time.time()
+        logger.info(f"\n=== Iniciando procesamiento optimizado de PDF: {file_name} ===")
+        logger.info(f"Configuración:")
+        logger.info(f"  - Tokens máximos por chunk: {self.chunk_processor.max_tokens}")
+        logger.info(f"  - Tokens objetivo: {self.chunk_processor.target_tokens}")
+        logger.info(f"  - Chunk size: {chunk_size}")
+        logger.info(f"  - Chunk overlap: {chunk_overlap}")
+        logger.info(f"  - Merge relacionados: {merge_related}")
+        logger.info(f"  - Extraer imágenes: {'✅' if self.extract_images else '❌'}")
+        if self.custom_instructions:
+            logger.info(f"  - Instrucciones personalizadas: {self.custom_instructions[:100]}...")
+
+        logger.info(f"\n1. Creando chunks del PDF...")
+        chunks = self._create_optimized_chunks_from_bytes(
+            pdf_bytes,
+            file_name,
+            chunk_size,
+            chunk_overlap
+        )
+        logger.info(f"  Total chunks creados: {len(chunks)}")
+
+        # Nota: La extracción de imágenes desde bytes no se implementa por ahora
+        # ya que extract_images está deshabilitado por defecto según requerimientos
+        if self.extract_images:
+            logger.warning("Extracción de imágenes desde bytes no implementada aún")
+
+        logger.info(f"\n2. Procesando y optimizando chunks...")
+        processed_chunks = self.chunk_processor.process_chunks_batch(chunks, merge_related)
+
+        total_time = time.time() - overall_start
+        if processed_chunks:
+            avg_tokens = sum(
+                self.token_manager.count_tokens(chunk.page_content)
+                for chunk in processed_chunks
+            ) / len(processed_chunks)
+        else:
+            avg_tokens = 0
+
+        logger.info(f"\n=== Procesamiento completado ===")
+        logger.info(f"  Tiempo total: {total_time:.2f}s")
+        logger.info(f"  Chunks procesados: {len(processed_chunks)}")
+        logger.info(f"  Tokens promedio por chunk: {avg_tokens:.1f}")
+        if self.custom_instructions:
+            logger.info(f"  Custom instructions aplicadas: ✅")
+
+        return processed_chunks
+
+    def _create_optimized_chunks_from_bytes(
+        self,
+        pdf_bytes: bytes,
+        file_name: str,
+        chunk_size: int,
+        chunk_overlap: int
+    ) -> List[Document]:
+        """
+        Crea chunks optimizados desde bytes del PDF.
+
+        Args:
+            pdf_bytes: Contenido del PDF en bytes
+            file_name: Nombre del archivo
+            chunk_size: Tamaño del chunk
+            chunk_overlap: Solapamiento entre chunks
+
+        Returns:
+            Lista de documentos con chunks
+        """
+        logger.info(f"  Leyendo PDF desde bytes: {file_name}")
+
+        # Crear BytesIO para pypdf
+        pdf_buffer = BytesIO(pdf_bytes)
+        pdf = PdfReader(pdf_buffer)
+        chunks = []
+
+        text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=chunk_size,
+            chunk_overlap=chunk_overlap,
+            length_function=self.token_manager.count_tokens,
+            separators=["\n\n", "\n", ". ", " ", ""]
+        )
+
+        # Extraer todo el texto concatenado con tracking de páginas
+        full_text = ""
+        page_boundaries = []  # Lista de (char_position, page_num)
+
+        for page_num, page in enumerate(pdf.pages, 1):
+            text = page.extract_text()
+            if text.strip():
+                page_start = len(full_text)
+                full_text += text
+                # Agregar separador entre páginas (excepto después de la última)
+                if page_num < len(pdf.pages):
+                    full_text += "\n\n"
+                page_end = len(full_text)
+                page_boundaries.append((page_start, page_end, page_num))
+
+        if not full_text.strip():
+            return []
+
+        # Dividir el texto completo (esto permite overlap entre páginas)
+        text_chunks = text_splitter.split_text(full_text)
+
+        logger.info(f"  Total de chunks generados por splitter: {len(text_chunks)}")
+        if len(text_chunks) >= 2:
+            # Verificar overlap entre primer y segundo chunk
+            chunk0_end = text_chunks[0][-100:] if len(text_chunks[0]) > 100 else text_chunks[0]
+            chunk1_start = text_chunks[1][:100] if len(text_chunks[1]) > 100 else text_chunks[1]
+            logger.info(f"  Chunk 0 termina con: ...{chunk0_end}")
+            logger.info(f"  Chunk 1 empieza con: {chunk1_start}...")
+
+        # Asignar página a cada chunk basándonos en su posición en el texto original
+        chunks = []
+        current_search_pos = 0
+
+        for chunk_text in text_chunks:
+            # Buscar donde aparece este chunk en el texto completo
+            chunk_pos = full_text.find(chunk_text, current_search_pos)
+
+            if chunk_pos == -1:
+                # Si no lo encontramos, usar la última posición conocida
+                chunk_pos = current_search_pos
+
+            # Determinar la página basándonos en la posición del inicio del chunk
+            chunk_page = 1
+            for start, end, page_num in page_boundaries:
+                if chunk_pos >= start and chunk_pos < end:
+                    chunk_page = page_num
+                    break
+                elif chunk_pos >= end:
+                    # El chunk está después de esta página, continuar buscando
+                    chunk_page = page_num  # Guardar la última página vista
+
+            chunks.append(Document(
+                page_content=chunk_text,
+                metadata={
+                    "page": chunk_page,
+                    "file_name": file_name,
+                }
+            ))
+
+            # Actualizar posición de búsqueda para el siguiente chunk
+            current_search_pos = chunk_pos + len(chunk_text)
+
+        return chunks