src/components/document_handler.py - master - alsoalgo/CTF

src/components/document_handler.py
906 строк52 KB
alsoalgo
some improvements
16 дек 2025, 13:03
16 дек 2025, 13:0372b7a7e
import pandas as pd
import os
import logging
from typing import Optional, Tuple, Dict, Union, List
import pdfplumber
from docx import Document
import openpyxl
import base64
import io

try:
    import fitz  # PyMuPDF
    PYMUPDF_AVAILABLE = True
except ImportError:
    PYMUPDF_AVAILABLE = False
    fitz = None

try:
    from PIL import Image
    PIL_AVAILABLE = True
except ImportError:
    PIL_AVAILABLE = False
    Image = None

try:
    from src.components.ocr_handler import OCRHandler
    OCR_AVAILABLE = True
except ImportError:
    OCR_AVAILABLE = False
    OCRHandler = None

logger = logging.getLogger(__name__)


class DocumentHandler:
    def __init__(self, use_ocr: bool = True, ocr_api_key: Optional[str] = None):
        """
        Инициализация DocumentHandler.
        
        :param use_ocr: Использовать ли OpenRouter OCR для сложных документов (по умолчанию True)
        :param ocr_api_key: API ключ OpenRouter (если None, берется из переменной окружения)
        """
        self.use_ocr = use_ocr and OCR_AVAILABLE
        self.ocr_handler = None
        if self.use_ocr:
            try:
                self.ocr_handler = OCRHandler(api_key=ocr_api_key)
                logger.info("OCR handler инициализирован")
            except Exception as e:
                logger.warning(f"Не удалось инициализировать OCR handler: {e}")
                self.use_ocr = False
        else:
            logger.info("OCR отключен")
    
    def load_document(self, file_path: str, use_ocr_fallback: bool = False, track_sources: bool = False) -> Union[Optional[pd.DataFrame], Tuple[Optional[pd.DataFrame], Optional[Dict]]]:
        """
        Universal method to load documents of different formats (CSV, XLSX, PDF, DOCX).
        Returns a pandas DataFrame or None if an error occurs.
        """
        if not os.path.exists(file_path):
            print(f"File not found: {file_path}")
            return None

        file_ext = os.path.splitext(file_path)[1].lower()

        if file_ext == '.csv':
            result = self.load_csv_document(file_path)
            if result is not None:
                if track_sources:
                    sources_info = {
                        'file': file_path,
                        'method': 'csv',
                        'rows': len(result),
                        'columns': len(result.columns) if not result.empty else 0
                    }
                    return (result, sources_info)
                return result
            else:
                return (None, None) if track_sources else None
        elif file_ext in ['.xlsx', '.xls']:
            result = self.load_xlsx_document(file_path, track_sources=track_sources)
            if isinstance(result, tuple):
                return result
            elif result is not None:
                return (result, {'method': 'xlsx', 'file': file_path}) if track_sources else result
            else:
                return (None, None) if track_sources else None
        elif file_ext == '.pdf':
            # Всегда используем комплексный анализ PDF (текст + таблицы + OCR)
            result = self.load_pdf_document(file_path, track_sources=track_sources)
            if isinstance(result, tuple):
                return result
            elif result is not None:
                return (result, {'method': 'pdf_comprehensive', 'file': file_path}) if track_sources else result
            else:
                # Если комплексный анализ не дал результатов, пробуем только OCR
                if (use_ocr_fallback or self.use_ocr) and self.ocr_handler:
                    logger.info("Комплексный анализ не дал результатов, пробую только OCR...")
                    ocr_result = self._extract_with_ocr_comprehensive(file_path, track_sources)
                    if ocr_result:
                        return ocr_result
                return (None, None) if track_sources else None
        elif file_ext == '.docx':
            result = self.load_docx_document(file_path, track_sources=track_sources)
            if isinstance(result, tuple):
                return result
            elif result is not None:
                return (result, {'method': 'docx', 'file': file_path}) if track_sources else result
            else:
                return (None, None) if track_sources else None
        else:
            logger.warning(f"Unsupported file format: {file_ext}")
            return (None, None) if track_sources else None

    def load_csv_document(self, file_path: str) -> Optional[pd.DataFrame]:
        """
        Loads a single CSV file into a pandas DataFrame.
        Returns the DataFrame or None if an error occurs.
        """
        try:
            df = pd.read_csv(file_path, index_col=None, dtype=str)
            if df is None or df.empty:
                logger.warning(f"CSV file '{file_path}' is empty or could not be loaded")
                return None
            logger.info(f"CSV file '{file_path}' successfully loaded: {len(df)} rows, {len(df.columns)} columns")
            return df
        except Exception as e:
            logger.error(f"Error loading CSV file {file_path}: {e}")
            return None

    def load_xlsx_document(self, file_path: str, track_sources: bool = False) -> Union[Optional[pd.DataFrame], Tuple[Optional[pd.DataFrame], Optional[Dict]]]:
        """
        Loads an Excel file (XLSX) into a pandas DataFrame.
        Reads the first sheet by default.
        
        :param file_path: Путь к XLSX файлу
        :param track_sources: Если True, возвращает также информацию об источниках
        :return: DataFrame или (DataFrame, sources_info) если track_sources=True
        """
        try:
            df = pd.read_excel(file_path, engine='openpyxl', dtype=str)
            if df is None or df.empty:
                logger.warning(f"XLSX file '{file_path}' is empty or could not be loaded")
                return (None, None) if track_sources else None
            
            logger.info(f"XLSX file '{file_path}' successfully loaded: {len(df)} rows, {len(df.columns)} columns")
            
            if track_sources:
                sources_info = {
                    'file': file_path,
                    'method': 'xlsx',
                    'rows': len(df),
                    'columns': len(df.columns)
                }
                return (df, sources_info)
            return df
        except Exception as e:
            logger.error(f"Error loading XLSX file {file_path}: {e}")
            return (None, None) if track_sources else None

    def _extract_images_from_pdf(self, file_path: str) -> List[Dict]:
        """
        Извлекает все изображения из PDF файла используя PyMuPDF.
        
        :param file_path: Путь к PDF файлу
        :return: Список словарей с информацией об изображениях (page, image_index, base64_data)
        """
        if not PYMUPDF_AVAILABLE or not PIL_AVAILABLE:
            logger.warning("PyMuPDF или PIL не установлены, извлечение изображений недоступно")
            return []
        
        images = []
        try:
            doc = fitz.open(file_path)
            logger.info(f"Извлекаю изображения из PDF через PyMuPDF...")
            
            for page_num in range(len(doc)):
                page = doc.load_page(page_num)
                image_list = page.get_images(full=True)
                
                for img_index, img in enumerate(image_list):
                    try:
                        # Получаем изображение
                        xref = img[0]
                        pix = fitz.Pixmap(doc, xref)
                        
                        # Конвертируем только RGB/GRAY изображения (пропускаем CMYK)
                        if pix.n - pix.alpha < 4:  # GRAY или RGB
                            # Конвертируем в PIL Image
                            img_data = pix.tobytes("png")
                            pil_image = Image.open(io.BytesIO(img_data))
                            
                            # Сжимаем изображение, если оно слишком большое (максимум 5 MB для API)
                            # Сжимаем до максимального размера ~4.5 MB (оставляем запас)
                            max_size_bytes = 4_500_000  # 4.5 MB
                            buffer = io.BytesIO()
                            pil_image.save(buffer, format='PNG')
                            original_size = buffer.tell()
                            
                            # Если изображение слишком большое, сжимаем его
                            if original_size > max_size_bytes:
                                logger.info(f"  Изображение слишком большое ({original_size / 1024 / 1024:.2f} MB), сжимаю...")
                                # Вычисляем коэффициент сжатия
                                scale_factor = (max_size_bytes / original_size) ** 0.5  # Квадратный корень для площади
                                new_width = int(pil_image.width * scale_factor)
                                new_height = int(pil_image.height * scale_factor)
                                
                                # Изменяем размер изображения
                                pil_image = pil_image.resize((new_width, new_height), Image.Resampling.LANCZOS)
                                
                                # Пробуем сохранить с разными уровнями качества
                                buffer = io.BytesIO()
                                pil_image.save(buffer, format='PNG', optimize=True)
                                
                                # Если все еще слишком большое, пробуем JPEG с качеством
                                if buffer.tell() > max_size_bytes:
                                    logger.info(f"  PNG все еще большой, конвертирую в JPEG...")
                                    # Конвертируем в RGB если нужно
                                    if pil_image.mode != 'RGB':
                                        pil_image = pil_image.convert('RGB')
                                    
                                    # Пробуем разные уровни качества
                                    for quality in [85, 75, 65, 55, 45]:
                                        buffer = io.BytesIO()
                                        pil_image.save(buffer, format='JPEG', quality=quality, optimize=True)
                                        if buffer.tell() <= max_size_bytes:
                                            logger.info(f"  Сжато до {buffer.tell() / 1024 / 1024:.2f} MB (JPEG quality {quality})")
                                            img_base64 = base64.b64encode(buffer.getvalue()).decode()
                                            images.append({
                                                'page': page_num + 1,
                                                'image_index': img_index,
                                                'base64': f"data:image/jpeg;base64,{img_base64}",
                                                'width': new_width,
                                                'height': new_height
                                            })
                                            logger.info(f"  Страница {page_num + 1}, изображение {img_index + 1}: {new_width}x{new_height}px (сжато)")
                                            break
                                    else:
                                        logger.warning(f"  Не удалось сжать изображение до приемлемого размера, пропускаю")
                                        pix = None
                                        continue
                                else:
                                    img_base64 = base64.b64encode(buffer.getvalue()).decode()
                                    images.append({
                                        'page': page_num + 1,
                                        'image_index': img_index,
                                        'base64': f"data:image/png;base64,{img_base64}",
                                        'width': new_width,
                                        'height': new_height
                                    })
                                    logger.info(f"  Страница {page_num + 1}, изображение {img_index + 1}: {new_width}x{new_height}px (сжато PNG)")
                            else:
                                # Изображение нормального размера, сохраняем как есть
                                img_base64 = base64.b64encode(buffer.getvalue()).decode()
                                images.append({
                                    'page': page_num + 1,
                                    'image_index': img_index,
                                    'base64': f"data:image/png;base64,{img_base64}",
                                    'width': pix.width,
                                    'height': pix.height
                                })
                                logger.info(f"  Страница {page_num + 1}, изображение {img_index + 1}: {pix.width}x{pix.height}px")
                        else:
                            # CMYK изображения - конвертируем в RGB
                            pix_rgb = fitz.Pixmap(fitz.csRGB, pix)
                            img_data = pix_rgb.tobytes("png")
                            pil_image = Image.open(io.BytesIO(img_data))
                            
                            # Сжимаем изображение, если оно слишком большое
                            max_size_bytes = 4_500_000  # 4.5 MB
                            buffer = io.BytesIO()
                            pil_image.save(buffer, format='PNG')
                            original_size = buffer.tell()
                            
                            if original_size > max_size_bytes:
                                logger.info(f"  Изображение слишком большое ({original_size / 1024 / 1024:.2f} MB), сжимаю...")
                                scale_factor = (max_size_bytes / original_size) ** 0.5
                                new_width = int(pil_image.width * scale_factor)
                                new_height = int(pil_image.height * scale_factor)
                                pil_image = pil_image.resize((new_width, new_height), Image.Resampling.LANCZOS)
                                
                                buffer = io.BytesIO()
                                pil_image.save(buffer, format='PNG', optimize=True)
                                
                                if buffer.tell() > max_size_bytes:
                                    if pil_image.mode != 'RGB':
                                        pil_image = pil_image.convert('RGB')
                                    for quality in [85, 75, 65, 55, 45]:
                                        buffer = io.BytesIO()
                                        pil_image.save(buffer, format='JPEG', quality=quality, optimize=True)
                                        if buffer.tell() <= max_size_bytes:
                                            img_base64 = base64.b64encode(buffer.getvalue()).decode()
                                            images.append({
                                                'page': page_num + 1,
                                                'image_index': img_index,
                                                'base64': f"data:image/jpeg;base64,{img_base64}",
                                                'width': new_width,
                                                'height': new_height
                                            })
                                            logger.info(f"  Страница {page_num + 1}, изображение {img_index + 1}: {new_width}x{new_height}px (CMYK->RGB, сжато JPEG)")
                                            break
                                    else:
                                        logger.warning(f"  Не удалось сжать изображение, пропускаю")
                                        pix_rgb = None
                                        pix = None
                                        continue
                                else:
                                    img_base64 = base64.b64encode(buffer.getvalue()).decode()
                                    images.append({
                                        'page': page_num + 1,
                                        'image_index': img_index,
                                        'base64': f"data:image/png;base64,{img_base64}",
                                        'width': new_width,
                                        'height': new_height
                                    })
                                    logger.info(f"  Страница {page_num + 1}, изображение {img_index + 1}: {new_width}x{new_height}px (CMYK->RGB, сжато PNG)")
                            else:
                                img_base64 = base64.b64encode(buffer.getvalue()).decode()
                                images.append({
                                    'page': page_num + 1,
                                    'image_index': img_index,
                                    'base64': f"data:image/png;base64,{img_base64}",
                                    'width': pix_rgb.width,
                                    'height': pix_rgb.height
                                })
                                logger.info(f"  Страница {page_num + 1}, изображение {img_index + 1}: {pix_rgb.width}x{pix_rgb.height}px (CMYK->RGB)")
                            pix_rgb = None
                        
                        pix = None
                    except Exception as e:
                        logger.warning(f"Ошибка при извлечении изображения {img_index} со страницы {page_num + 1}: {e}")
                        continue
            
            doc.close()
            logger.info(f"Всего извлечено изображений: {len(images)}")
            return images
        except Exception as e:
            logger.error(f"Ошибка при извлечении изображений из PDF: {e}")
            return []
    
    def load_pdf_document(self, file_path: str, track_sources: bool = False) -> Tuple[Optional[pd.DataFrame], Optional[Dict]]:
        """
        Комплексная загрузка PDF файла с полным анализом:
        ПРИОРИТЕТ: Сначала извлекаем и обрабатываем ВСЕ изображения через OCR,
        затем обрабатываем текст и таблицы.
        
        - Извлечение и обработка ВСЕХ изображений через OCR (mistral-ocr)
        - Извлечение текста
        - Извлечение таблиц
        - Агрегация всех источников информации
        
        :param file_path: Путь к PDF файлу
        :param track_sources: Если True, возвращает также информацию об источниках
        :return: DataFrame и опционально словарь с источниками
        """
        logger.info(f"Начинаю комплексный анализ PDF: {file_path}")
        logger.info("=" * 80)
        logger.info("ПРИОРИТЕТ: Сначала извлекаю и обрабатываю ВСЕ изображения")
        logger.info("=" * 80)
        
        aggregated_data = []
        sources_info = {
            'file': file_path,
            'text_sources': [],
            'table_sources': [],
            'ocr_sources': [],
            'image_sources': [],
            'images_found': False,
            'images_processed': 0
        } if track_sources else None
        
        try:
            # ШАГ 1: Извлекаем ВСЕ изображения из PDF и обрабатываем их через OCR
            extracted_images = self._extract_images_from_pdf(file_path)
            has_images = len(extracted_images) > 0
            
            if has_images:
                logger.info(f"Найдено {len(extracted_images)} изображений, обрабатываю каждое через OCR...")
                if track_sources:
                    sources_info['images_found'] = True
                
                # Обрабатываем каждое изображение отдельно через vision-модель
                for img_info in extracted_images:
                    try:
                        logger.info(f"Обрабатываю изображение {img_info['image_index'] + 1} со страницы {img_info['page']}...")
                        
                        # Обрабатываем изображение через vision-модель (не mistral-ocr, т.к. он работает только с PDF)
                        image_prompt = """Extract ALL text and data from this image.
Pay special attention to:
- All visible text (including small text, labels, numbers)
- Tables and structured data
- Key-value pairs (invoice numbers, dates, amounts, names, addresses, etc.)
- Any numbers, codes, or identifiers
- Text in headers, footers, and margins

Return all extracted information in a clear, organized format.
If you see tables, format them with clear headers and rows."""
                        
                        # Сохраняем информацию об изображении в sources_info (даже если OCR не сработал)
                        if track_sources:
                            image_source_info = {
                                'page': img_info['page'],
                                'image_index': img_info['image_index'],
                                'width': img_info['width'],
                                'height': img_info['height'],
                                'text_extracted': 0,
                                'ocr_text': '',
                                'status': 'failed'
                            }
                        
                        if self.ocr_handler:
                            # Используем vision-модель для изображений (mistral-ocr работает только с PDF)
                            extracted_text = self.ocr_handler.extract_text_from_image(
                                img_info['base64'],
                                prompt=image_prompt,
                                model="anthropic/claude-3.5-sonnet"
                            )
                            
                            if extracted_text and extracted_text.strip():
                                # Парсим результат OCR
                                lines = [line.strip() for line in extracted_text.split('\n') if line.strip()]
                                
                                # Добавляем данные из изображения
                                for line in lines:
                                    if line:
                                        aggregated_data.append({
                                            'Content': line,
                                            'Source': 'ocr_image',
                                            'Page': img_info['page'],
                                            'ImageIndex': img_info['image_index'],
                                            'ImageSize': f"{img_info['width']}x{img_info['height']}",
                                            'context': f"Изображение {img_info['image_index'] + 1} со страницы {img_info['page']}"
                                        })
                                
                                if track_sources:
                                    image_source_info['text_extracted'] = len(lines)
                                    image_source_info['ocr_text'] = extracted_text[:500] + "..." if len(extracted_text) > 500 else extracted_text
                                    image_source_info['status'] = 'success'
                                    sources_info['images_processed'] += 1
                                
                                logger.info(f"  ✓ Извлечено {len(lines)} строк текста из изображения {img_info['image_index'] + 1} (страница {img_info['page']})")
                            else:
                                logger.warning(f"  ✗ Vision-модель не извлекла текст из изображения {img_info['image_index'] + 1} (страница {img_info['page']})")
                                if track_sources:
                                    image_source_info['status'] = 'no_text'
                        else:
                            logger.warning("OCR handler недоступен, пропускаю изображение")
                            if track_sources:
                                image_source_info['status'] = 'no_handler'
                        
                        # Сохраняем информацию об изображении (даже если OCR не сработал)
                        if track_sources:
                            sources_info['image_sources'].append(image_source_info)
                    except Exception as e:
                        logger.error(f"Ошибка при обработке изображения {img_info.get('image_index', '?')} со страницы {img_info.get('page', '?')}: {e}")
                        import traceback
                        logger.debug(traceback.format_exc())
                        continue
                
                logger.info(f"✓ Обработано {len(extracted_images)} изображений")
                logger.info("=" * 80)
            
            # ШАГ 2: Извлекаем текст и таблицы из PDF
            logger.info("Теперь извлекаю текст и таблицы из PDF...")
            
            with pdfplumber.open(file_path) as pdf:
                total_pages = len(pdf.pages)
                logger.info(f"PDF содержит {total_pages} страниц")
                
                # 2.1. Извлечение текста из всех страниц
                text_data = []
                for page_num, page in enumerate(pdf.pages, 1):
                    text = page.extract_text()
                    if text and text.strip():
                        lines = [line.strip() for line in text.split('\n') if line.strip()]
                        for line in lines:
                            text_data.append({
                                'Content': line,
                                'Source': 'text',
                                'Page': page_num
                            })
                        if track_sources:
                            sources_info['text_sources'].append({
                                'page': page_num,
                                'lines_count': len(lines)
                            })
                        logger.info(f"  Страница {page_num}: извлечено {len(lines)} строк текста")
                
                # 2.2. Извлечение таблиц из всех страниц
                all_tables = []
                for page_num, page in enumerate(pdf.pages, 1):
                    tables = page.extract_tables()
                    if tables:
                        for table_idx, table in enumerate(tables):
                            all_tables.append({
                                'table': table,
                                'page': page_num,
                                'index': table_idx
                            })
                            if track_sources:
                                sources_info['table_sources'].append({
                                    'page': page_num,
                                    'table_index': table_idx,
                                    'rows': len(table),
                                    'cols': len(table[0]) if table else 0
                                })
                        logger.info(f"  Страница {page_num}: найдено {len(tables)} таблиц")
                
                # 2.3. Если нет текста и таблиц (полностью сканированный документ), используем OCR для всего PDF
                # (Изображения уже обработаны отдельно выше)
                if not text_data and not all_tables and not has_images and self.ocr_handler:
                    logger.info("Текст и таблицы не найдены, изображений тоже нет, запускаю OCR для всего PDF...")
                    try:
                        ocr_result = self._extract_with_ocr_comprehensive(file_path, track_sources, force_mistral_ocr=False)
                        if ocr_result:
                            ocr_df, ocr_sources = ocr_result
                            if ocr_df is not None and not ocr_df.empty:
                                # Добавляем OCR данные
                                for idx, row in ocr_df.iterrows():
                                    for col in ocr_df.columns:
                                        value = str(row[col]).strip()
                                        if value:
                                            aggregated_data.append({
                                                'Content': value,
                                                'Source': 'ocr',
                                                'Page': ocr_sources.get('page', -1) if ocr_sources else -1,
                                                'Column': col,
                                                'context': ocr_sources.get('context', '') if ocr_sources else ''
                                            })
                                if track_sources and ocr_sources:
                                    sources_info['ocr_sources'].append(ocr_sources)
                                logger.info(f"  ✓ OCR извлек {len(ocr_df)} строк данных из всего PDF")
                    except Exception as ocr_error:
                        logger.error(f"  Ошибка при выполнении OCR для всего PDF: {ocr_error}")
                
                # 3. Добавляем текстовые данные
                aggregated_data.extend(text_data)
                
                # 4. Добавляем данные из таблиц
                for table_info in all_tables:
                    table = table_info['table']
                    page_num = table_info['page']
                    
                    if len(table) > 1:
                        # Первая строка как заголовки
                        headers = [str(cell).strip() if cell else f"Column_{i}" 
                                  for i, cell in enumerate(table[0])]
                        for row_idx, row in enumerate(table[1:], 1):
                            row_dict = {'Source': 'table', 'Page': page_num, 'Row': row_idx}
                            for col_idx, cell in enumerate(row):
                                if col_idx < len(headers):
                                    value = str(cell).strip() if cell else ''
                                    if value:
                                        row_dict[headers[col_idx]] = value
                            if any(v for k, v in row_dict.items() if k not in ['Source', 'Page', 'Row']):
                                aggregated_data.append(row_dict)
                    else:
                        # Если только одна строка, добавляем как есть
                        for cell in table[0]:
                            value = str(cell).strip() if cell else ''
                            if value:
                                aggregated_data.append({
                                    'Content': value,
                                    'Source': 'table',
                                    'Page': page_num
                                })
                
                # 5. Агрегируем все данные в единый DataFrame
                if not aggregated_data:
                    logger.warning(f"Не удалось извлечь данные из PDF: {file_path}")
                    return (None, None) if track_sources else None
                
                # Создаем DataFrame из агрегированных данных
                df = pd.DataFrame(aggregated_data)
                
                # Если есть колонки кроме стандартных, объединяем их
                if 'Content' not in df.columns:
                    # Собираем все значения в одну колонку Content
                    content_cols = [col for col in df.columns if col not in ['Source', 'Page', 'Row', 'Column']]
                    if content_cols:
                        df['Content'] = df[content_cols].apply(
                            lambda x: ' | '.join([str(v) for v in x if pd.notna(v) and str(v).strip()]), 
                            axis=1
                        )
                        df = df[['Content', 'Source', 'Page'] + [c for c in df.columns if c not in ['Content', 'Source', 'Page']]]
                
                # Итоговая статистика
                text_count = len([d for d in aggregated_data if d.get('Source') == 'text'])
                table_count = len([d for d in aggregated_data if d.get('Source') == 'table'])
                ocr_image_count = len([d for d in aggregated_data if d.get('Source') == 'ocr_image'])
                ocr_count = len([d for d in aggregated_data if d.get('Source') == 'ocr'])
                
                logger.info("=" * 80)
                logger.info(f"PDF успешно обработан: {file_path}")
                logger.info(f"  Всего извлечено записей: {len(df)}")
                logger.info(f"    - Из изображений (OCR): {ocr_image_count}")
                logger.info(f"    - Из текста: {text_count}")
                logger.info(f"    - Из таблиц: {table_count}")
                logger.info(f"    - Из OCR всего PDF: {ocr_count}")
                logger.info(f"  Найдено таблиц: {len(all_tables)}")
                if has_images:
                    logger.info(f"  ✓ Обработано изображений: {len(extracted_images)}")
                logger.info("=" * 80)
                
                if not track_sources:
                    return df
                return (df, sources_info)

        except Exception as e:
            logger.error(f"Ошибка при загрузке PDF файла {file_path}: {e}")
            return (None, None) if track_sources else None
    
    def _extract_with_ocr_comprehensive(self, file_path: str, track_sources: bool = False, force_mistral_ocr: bool = False) -> Optional[Tuple[pd.DataFrame, Dict]]:
        """
        Комплексное извлечение данных через OCR с детальной информацией об источниках.
        
        :param file_path: Путь к PDF файлу
        :param track_sources: Если True, возвращает информацию об источниках
        :param force_mistral_ocr: Если True, сразу использует mistral-ocr (рекомендуется для PDF с изображениями)
        :return: DataFrame и словарь с источниками или None
        """
        if not self.ocr_handler:
            return None
        
        try:
            # Если force_mistral_ocr=True (для PDF с изображениями), сразу используем mistral-ocr
            # Согласно рекомендациям OpenRouter, mistral-ocr - лучший выбор для PDF с изображениями
            if not force_mistral_ocr:
                # Пробуем сначала бесплатный движок для структурированных PDF (без изображений)
                logger.info("  Пробую извлечь структурированные данные через OCR (pdf-text)...")
                # Убеждаемся, что промпт не пустой
                structure_prompt = """Extract all structured data from this document and format it as a table.
Identify all key-value pairs, invoice fields, table data, and structured information.
Return the data in a clear markdown table format with column headers separated by |.
Format: Field | Value
Include all important information like invoice numbers, dates, amounts, names, addresses, etc.
If multiple tables are present, extract the main one or combine them."""
                df = self.ocr_handler.extract_structured_data(file_path, engine="pdf-text", structure_prompt=structure_prompt)
                if df is not None and not df.empty:
                    logger.info(f"  Успешно извлечено через pdf-text: {df.shape[0]} строк")
                    sources = {'method': 'ocr_pdf_text', 'page': -1} if track_sources else {}
                    return (df, sources) if track_sources else (df, {})
            
            # Используем mistral-ocr (платный, но более мощный, особенно для PDF с изображениями)
            if force_mistral_ocr:
                logger.info("  Использую mistral-ocr (рекомендуется для PDF с изображениями)...")
            else:
                logger.info("  Пробую извлечь через mistral-ocr...")
            
            # Улучшенный промпт для извлечения текста из изображений
            mistral_prompt = """Extract ALL text and data from this document, including text from images and scanned content.
Pay special attention to:
- Text embedded in images
- Tables and structured data
- Key-value pairs (invoice numbers, dates, amounts, names, addresses, etc.)
- Any text that appears in images or graphics

Format tables with clear headers and rows.
Return all extracted information in a clear, organized format.
For each piece of information, try to include the page number if available.
If you see text in images, extract it completely."""
            
            # Запрашиваем аннотации для получения информации о страницах и позициях
            ocr_result = self.ocr_handler.extract_text_with_ocr(
                file_path, 
                engine="mistral-ocr",
                prompt=mistral_prompt,
                return_annotations=True
            )
            
            extracted_text = None
            annotations = None
            if isinstance(ocr_result, tuple):
                extracted_text, annotations = ocr_result
            else:
                extracted_text = ocr_result
            
            if extracted_text:
                # Пробуем распарсить как таблицу
                lines = [line.strip() for line in extracted_text.split('\n') if line.strip()]
                if lines:
                    # Пытаемся найти табличную структуру
                    table_data = []
                    for line in lines:
                        # Проверяем, похоже ли на строку таблицы (с разделителями | или табуляцией)
                        if '|' in line or '\t' in line:
                            parts = [p.strip() for p in (line.split('|') if '|' in line else line.split('\t'))]
                            if len(parts) > 1:
                                table_data.append(parts)
                    
                    if table_data and len(table_data) > 1:
                        # Создаем DataFrame из таблицы
                        headers = table_data[0]
                        df = pd.DataFrame(table_data[1:], columns=headers[:len(table_data[1])])
                    else:
                        # Создаем DataFrame из текста
                        df = pd.DataFrame({'Content': lines})
                    
                    logger.info(f"  ✓ Успешно извлечено через mistral-ocr: {df.shape[0]} строк")
                    
                    # Формируем sources_info с информацией о страницах и контексте
                    if track_sources:
                        sources = {
                            'method': 'ocr_mistral',
                            'page': -1,  # По умолчанию -1, если страница не определена
                            'ocr_raw_text': extracted_text,  # Сохраняем весь текст для контекста
                            'annotations': annotations if annotations else {}
                        }
                        
                        # Пытаемся извлечь информацию о страницах из аннотаций
                        if annotations:
                            # Если в аннотациях есть информация о страницах
                            if isinstance(annotations, dict):
                                if 'page' in annotations:
                                    sources['page'] = annotations.get('page', -1)
                                elif 'pages' in annotations:
                                    # Если несколько страниц, берем первую
                                    pages = annotations.get('pages', [])
                                    if pages:
                                        sources['page'] = pages[0] if isinstance(pages[0], int) else -1
                        
                        # Извлекаем контекст из текста (первые 200 символов для каждого значения)
                        # Это будет использоваться как базовый контекст
                        if extracted_text:
                            # Берем первые 200 символов как общий контекст документа
                            sources['context'] = extracted_text[:200] + "..." if len(extracted_text) > 200 else extracted_text
                        
                        return (df, sources)
                    else:
                        return (df, {})
                else:
                    logger.warning("  mistral-ocr вернул пустой текст")
            else:
                logger.warning("  mistral-ocr не вернул результат")
            
            logger.warning("  OCR не смог извлечь данные (возможно, проблема с подключением к API)")
            return None
        except Exception as e:
            logger.error(f"Ошибка при OCR извлечении: {e}")
            return None

    def load_pdf_with_ocr(self, file_path: str) -> Optional[pd.DataFrame]:
        """
        Загружает PDF используя OpenRouter OCR для распознавания.
        Используется как fallback когда стандартные методы не работают.
        
        :param file_path: Путь к PDF файлу
        :return: DataFrame с извлеченными данными или None
        """
        if not self.ocr_handler:
            print("OCR handler not available")
            return None
        
        try:
            # Пробуем сначала бесплатный движок для структурированных PDF
            df = self.ocr_handler.extract_structured_data(file_path)
            if df is not None and not df.empty:
                print(f"PDF file '{file_path}' successfully loaded using OpenRouter OCR.")
                return df
            
            # Если не получилось, пробуем mistral-ocr (платный, но более мощный)
            print("Trying mistral-ocr engine...")
            extracted_text = self.ocr_handler.extract_with_mistral_ocr(file_path)
            if extracted_text:
                # Преобразуем текст в DataFrame
                lines = [line.strip() for line in extracted_text.split('\n') if line.strip()]
                if lines:
                    df = pd.DataFrame({'Content': lines})
                    print(f"PDF file '{file_path}' successfully loaded using mistral-ocr.")
                    return df
            
            return None
        except Exception as e:
            print(f"Error loading PDF with OCR: {e}")
            return None

    def load_docx_document(self, file_path: str, track_sources: bool = False) -> Union[Optional[pd.DataFrame], Tuple[Optional[pd.DataFrame], Optional[Dict]]]:
        """
        Loads a DOCX file and extracts tables into a pandas DataFrame.
        If multiple tables are found, returns the first one.
        If no tables found, extracts paragraphs as text.
        
        :param file_path: Путь к DOCX файлу
        :param track_sources: Если True, возвращает также информацию об источниках
        :return: DataFrame или (DataFrame, sources_info) если track_sources=True
        """
        logger.info(f"Начинаю загрузку DOCX: {file_path}")
        sources_info = {
            'file': file_path,
            'method': 'docx',
            'tables_found': 0,
            'paragraphs_found': 0,
            'table_sources': [],
            'text_sources': []
        } if track_sources else None
        
        try:
            doc = Document(file_path)
            tables = doc.tables

            if not tables:
                # If no tables found, extract paragraphs as text
                paragraphs = [para.text.strip() for para in doc.paragraphs if para.text.strip()]
                if paragraphs:
                    df = pd.DataFrame({'Content': paragraphs})
                    logger.info(f"DOCX file '{file_path}' loaded as text: {len(paragraphs)} paragraphs found")
                    if track_sources:
                        sources_info['paragraphs_found'] = len(paragraphs)
                        sources_info['text_sources'] = [
                            {'paragraph_index': i, 'content_preview': para[:50] + '...' if len(para) > 50 else para}
                            for i, para in enumerate(paragraphs)
                        ]
                    return (df, sources_info) if track_sources else df
                else:
                    logger.warning(f"No tables or text found in DOCX file '{file_path}'.")
                    return (None, None) if track_sources else None

            # Extract all tables and aggregate them
            all_data = []
            all_headers = None
            
            for table_idx, table in enumerate(tables):
                table_data = []
                headers = None
                
                for i, row in enumerate(table.rows):
                    row_data = [cell.text.strip() for cell in row.cells]
                    if i == 0:
                        headers = row_data
                        if all_headers is None:
                            all_headers = headers
                    else:
                        # Проверяем, что строка не пустая (есть хотя бы одно непустое значение)
                        if any(cell.strip() for cell in row_data):
                            table_data.append(row_data)
                
                if headers:
                    # Если таблица содержит только заголовки, создаем DataFrame с заголовками
                    # и одной пустой строкой, чтобы DataFrame не был пустым
                    if not table_data:
                        logger.warning(f"Таблица {table_idx} содержит только заголовки, создаю пустую строку данных")
                        # Создаем одну пустую строку с количеством колонок равным заголовкам
                        table_data = [[''] * len(headers)]
                    
                    # Ensure all rows have the same length as headers
                    for row in table_data:
                        while len(row) < len(headers):
                            row.append('')
                        if len(row) > len(headers):
                            row = row[:len(headers)]
                    
                    all_data.extend(table_data)
                    
                    if track_sources:
                        sources_info['table_sources'].append({
                            'table_index': table_idx,
                            'rows': len(table_data),
                            'columns': len(headers),
                            'headers': headers
                        })

            if all_headers:
                # Если нет данных, создаем DataFrame с заголовками и одной пустой строкой
                if not all_data:
                    logger.warning("Все таблицы пустые, создаю DataFrame с заголовками и пустой строкой")
                    all_data = [[''] * len(all_headers)]
                
                # Ensure all rows have the same length as all_headers
                for row in all_data:
                    while len(row) < len(all_headers):
                        row.append('')
                    if len(row) > len(all_headers):
                        row = row[:len(all_headers)]
                
                df = pd.DataFrame(all_data, columns=all_headers)
            else:
                # Если нет заголовков, создаем простой DataFrame из данных
                if not all_data:
                    logger.warning("Нет заголовков и данных, создаю пустой DataFrame")
                    df = pd.DataFrame()
                else:
                    df = pd.DataFrame(all_data)

            # Convert all values to string
            df = df.astype(str)
            
            if track_sources:
                sources_info['tables_found'] = len(tables)
            
            logger.info(f"DOCX file '{file_path}' successfully loaded: {len(tables)} table(s) found, {len(df)} rows extracted")
            
            return (df, sources_info) if track_sources else df

        except Exception as e:
            logger.error(f"Error loading DOCX file {file_path}: {e}")
            return (None, None) if track_sources else None

    def save_dataframe_to_csv(self, df: pd.DataFrame, output_path: str):
        """
        Saves a DataFrame to the specified CSV file.
        """
        try:
            df.to_csv(output_path, index=False)
            print(f"Table successfully saved to {output_path}")
        except Exception as e:
            print(f"Error saving file {output_path}: {e}")