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

src/components/ocr_handler.py
500 строк24 KB
alsoalgo
some improvements
16 дек 2025, 13:03
16 дек 2025, 13:0372b7a7e
"""
OCR Handler для работы с OpenRouter API
Поддерживает распознавание документов через OpenRouter OCR движки
"""
import os
import base64
import requests
from typing import Optional, Dict, Any, Tuple, Union
import pandas as pd


class OCRHandler:
    """
    Класс для работы с OpenRouter API для OCR распознавания документов.
    
    Поддерживает:
    - mistral-ocr: для сканированных документов ($2 за 1000 страниц)
    - pdf-text: для структурированных PDF (бесплатно)
    - Vision модели: GPT-4 Vision, Claude Vision и др.
    """
    
    def __init__(self, api_key: Optional[str] = None):
        """
        Инициализация OCR Handler.
        
        :param api_key: API ключ OpenRouter. Если None, берется из переменной окружения OPENROUTER_API_KEY
        """
        self.api_key = api_key or os.getenv("OPENROUTER_API_KEY")
        self.base_url = "https://openrouter.ai/api/v1/chat/completions"
        
        if not self.api_key:
            print("Warning: OPENROUTER_API_KEY not set. OCR functionality will be disabled.")
    
    def _encode_file_to_base64(self, file_path: str) -> Optional[str]:
        """
        Кодирует файл в base64 для отправки в API.
        Согласно документации OpenRouter, для PDF используется формат data:application/pdf;base64,...
        
        :param file_path: Путь к файлу
        :return: base64 строка в формате data URL или None при ошибке
        """
        try:
            with open(file_path, 'rb') as f:
                file_data = f.read()
                base64_data = base64.b64encode(file_data).decode('utf-8')
                
                # Определяем MIME тип
                file_ext = os.path.splitext(file_path)[1].lower()
                mime_types = {
                    '.pdf': 'application/pdf',
                    '.png': 'image/png',
                    '.jpg': 'image/jpeg',
                    '.jpeg': 'image/jpeg',
                }
                mime_type = mime_types.get(file_ext, 'application/octet-stream')
                
                # Формируем data URL согласно документации OpenRouter
                data_url = f"data:{mime_type};base64,{base64_data}"
                return data_url
        except Exception as e:
            print(f"Error encoding file to base64: {e}")
            return None
    
    def extract_text_with_ocr(
        self, 
        file_path: str, 
        engine: str = "mistral-ocr",
        model: str = "anthropic/claude-3.5-sonnet",
        prompt: Optional[str] = None,
        return_annotations: bool = False
    ) -> Union[Optional[str], Tuple[Optional[str], Optional[Dict]]]:
        """
        Извлекает текст из документа используя OpenRouter OCR.
        
        :param file_path: Путь к файлу (PDF, изображение)
        :param engine: Движок OCR ("mistral-ocr", "pdf-text", или "native")
        :param model: Модель для использования (по умолчанию claude-sonnet-4)
        :param prompt: Промпт для модели (опционально)
        :param return_annotations: Если True, возвращает также аннотации с позициями
        :return: Извлеченный текст или (текст, аннотации) если return_annotations=True
        """
        if not self.api_key:
            print("Error: OpenRouter API key not configured.")
            return None
        
        if not os.path.exists(file_path):
            print(f"Error: File not found: {file_path}")
            return None
        
        # Кодируем файл в base64
        file_data_url = self._encode_file_to_base64(file_path)
        if not file_data_url:
            return None
        
        # Получаем имя файла
        filename = os.path.basename(file_path)
        
        # Формируем промпт - гарантируем, что он всегда непустой
        default_prompt = """Extract all text and tables from this document. 
If there are tables, format them as structured data with clear column headers.
Return the extracted information in a clear, organized format."""
        
        if prompt is None:
            prompt = default_prompt
        else:
            prompt = prompt.strip()
            if not prompt:
                prompt = default_prompt
        
        # Финальная проверка - промпт должен быть непустым
        if not prompt or len(prompt) == 0:
            import logging
            logger = logging.getLogger(__name__)
            logger.warning("Prompt is empty after processing, using default prompt")
            prompt = default_prompt
        
        # Убеждаемся, что промпт не пустой перед формированием запроса
        prompt = prompt.strip()
        if not prompt or len(prompt) == 0:
            import logging
            logger = logging.getLogger(__name__)
            logger.warning("Prompt is empty after strip, using default prompt")
            prompt = default_prompt
        
        # Финальная проверка перед отправкой
        if not prompt or len(prompt.strip()) == 0:
            import logging
            logger = logging.getLogger(__name__)
            logger.error("Prompt is still empty, using minimal default")
            prompt = "Extract all text and tables from this document."
        
        # Логируем промпт для отладки (первые 100 символов)
        import logging
        logger = logging.getLogger(__name__)
        logger.debug(f"Using prompt (first 100 chars): {prompt[:100]}...")
        
        # Формируем запрос согласно документации OpenRouter
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
            "HTTP-Referer": "https://github.com/ctf-document-processor",  # Опционально
            "X-Title": "CTF Document Processor",  # Опционально
        }
        
        # Финальная проверка промпта перед формированием запроса
        if not prompt or len(prompt.strip()) == 0:
            logger.error("Prompt is empty before creating API request! Using emergency default.")
            prompt = "Extract all text and tables from this document."
        
        # Убеждаемся, что промпт не пустой и не состоит только из пробелов
        prompt = prompt.strip()
        if not prompt or len(prompt) == 0:
            logger.error("Prompt is empty after final strip! Using emergency default.")
            prompt = "Extract all text and tables from this document."
        
        # Проверяем, что промпт содержит видимые символы (не только пробелы и спецсимволы)
        if not any(c.isalnum() for c in prompt):
            logger.error("Prompt contains no visible alphanumeric characters! Using emergency default.")
            prompt = "Extract all text and tables from this document."
        
        # Логируем финальный промпт перед отправкой
        logger.info(f"Sending prompt to API (length: {len(prompt)}, first 100 chars): {prompt[:100]}...")
        
        # Финальная проверка перед формированием JSON
        if not prompt or len(prompt.strip()) == 0:
            logger.error("CRITICAL: Prompt is still empty before API call! This should not happen.")
            prompt = "Extract all text and tables from this document."
        
        # Формируем сообщения с правильным форматом для PDF
        messages = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": prompt
                    },
                    {
                        "type": "file",
                        "file": {
                            "filename": filename,
                            "file_data": file_data_url
                        }
                    }
                ]
            }
        ]
        
        # Настраиваем плагин для обработки PDF
        plugins = [
            {
                "id": "file-parser",
                "pdf": {
                    "engine": engine
                }
            }
        ]
        
        payload = {
            "model": model,
            "messages": messages,
            "plugins": plugins
        }
        
        try:
            response = requests.post(self.base_url, headers=headers, json=payload, timeout=120)
            response.raise_for_status()
            
            result = response.json()
            if 'choices' in result and len(result['choices']) > 0:
                message = result['choices'][0]['message']
                extracted_text = message.get('content', '')
                
                # Получаем аннотации если они есть (для отслеживания источников)
                annotations = None
                if return_annotations and 'annotations' in message:
                    annotations = message['annotations']
                elif return_annotations:
                    # Если аннотаций нет, создаем базовую структуру
                    annotations = {
                        'file_path': file_path,
                        'engine': engine,
                        'extraction_method': 'ocr'
                    }
                
                if return_annotations:
                    return extracted_text, annotations
                return extracted_text
            else:
                print("Error: Unexpected response format from OpenRouter API")
                return None if not return_annotations else (None, None)
                
        except requests.exceptions.RequestException as e:
            import logging
            logger = logging.getLogger(__name__)
            logger.error(f"Error calling OpenRouter API: {e}")
            if hasattr(e, 'response') and e.response is not None:
                try:
                    error_detail = e.response.json()
                    logger.error(f"API Error details: {error_detail}")
                    # Если ошибка связана с пустым промптом, логируем это
                    if 'text content blocks must be non-empty' in str(error_detail):
                        logger.warning("API returned error about empty text content. This may indicate a prompt issue.")
                except:
                    logger.error(f"API Error status: {e.response.status_code}")
            return None if not return_annotations else (None, None)
    
    def extract_structured_data(
        self, 
        file_path: str, 
        structure_prompt: Optional[str] = None,
        engine: str = "pdf-text"
    ) -> Optional[pd.DataFrame]:
        """
        Извлекает структурированные данные из документа и возвращает DataFrame.
        
        :param file_path: Путь к файлу
        :param structure_prompt: Промпт с описанием структуры данных для извлечения
        :param engine: Движок OCR для использования (по умолчанию pdf-text - бесплатный)
        :return: DataFrame с извлеченными данными или None
        """
        if structure_prompt is None:
            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."""
        
        # Убеждаемся, что промпт не пустой - строгая проверка
        default_structure_prompt = "Extract all structured data from this document and format it as a table with Field | Value format."
        
        if structure_prompt:
            structure_prompt = structure_prompt.strip()
            # Удаляем лишние пробелы и переносы строк, но сохраняем структуру
            cleaned_lines = [line.strip() for line in structure_prompt.split('\n') if line.strip()]
            if cleaned_lines:
                structure_prompt = '\n'.join(cleaned_lines)
            else:
                structure_prompt = None
        
        if not structure_prompt or len(structure_prompt) == 0 or not structure_prompt.strip():
            import logging
            logger = logging.getLogger(__name__)
            logger.warning("structure_prompt is empty or invalid, using default")
            structure_prompt = default_structure_prompt
        
        # Финальная проверка перед передачей - убеждаемся, что есть видимые символы
        structure_prompt = structure_prompt.strip()
        if not structure_prompt or len(structure_prompt) == 0 or not any(c.isalnum() for c in structure_prompt):
            import logging
            logger = logging.getLogger(__name__)
            logger.error("structure_prompt contains no visible characters, using minimal default")
            structure_prompt = default_structure_prompt
        
        # Логируем промпт для отладки
        import logging
        logger = logging.getLogger(__name__)
        logger.info(f"Using structure_prompt (length: {len(structure_prompt)}, first 150 chars): {structure_prompt[:150]}...")
        
        extracted_text = self.extract_text_with_ocr(file_path, engine=engine, prompt=structure_prompt)
        
        if not extracted_text:
            return None
        
        # Пытаемся преобразовать текст в DataFrame
        # Улучшенный парсинг для различных форматов вывода модели
        try:
            lines = [line.strip() for line in extracted_text.split('\n') if line.strip()]
            
            # Вариант 1: Markdown таблица с разделителем "|"
            if "|" in extracted_text:
                table_lines = [line for line in lines if '|' in line and not line.startswith('---')]
                
                if len(table_lines) > 0:
                    # Парсим первую строку как заголовки
                    first_line_parts = [cell.strip() for cell in table_lines[0].split('|') if cell.strip()]
                    
                    # Если первая строка выглядит как заголовок (содержит "Field", "Value", "Column" и т.д.)
                    if len(first_line_parts) >= 2:
                        headers = first_line_parts
                        data_rows = []
                        
                        # Обрабатываем остальные строки
                        for line in table_lines[1:]:
                            row_parts = [cell.strip() for cell in line.split('|') if cell.strip()]
                            # Если количество колонок совпадает или близко
                            if len(row_parts) >= 2:
                                # Нормализуем количество колонок
                                while len(row_parts) < len(headers):
                                    row_parts.append('')
                                if len(row_parts) > len(headers):
                                    row_parts = row_parts[:len(headers)]
                                data_rows.append(row_parts)
                        
                        if data_rows:
                            return pd.DataFrame(data_rows, columns=headers)
            
            # Вариант 2: Таблица с разделителями ":" или "-"
            # Ищем паттерны типа "Field: Value" или "Field - Value"
            key_value_pairs = []
            for line in lines:
                # Пропускаем заголовки и разделители
                if line.startswith('#') or line.startswith('==') or line.startswith('---'):
                    continue
                
                # Пробуем разные разделители
                for sep in ['|', ':', '-', '\t']:
                    if sep in line and len(line.split(sep)) >= 2:
                        parts = [p.strip() for p in line.split(sep, 1)]
                        if len(parts) == 2 and parts[0] and parts[1]:
                            key_value_pairs.append(parts)
                            break
            
            if key_value_pairs:
                # Если нашли пары ключ-значение, создаем DataFrame
                df = pd.DataFrame(key_value_pairs, columns=['Field', 'Value'])
                return df
            
            # Вариант 3: Простой список строк
            # Если ничего не подошло, создаем простой DataFrame
            return pd.DataFrame({'Content': lines})
            
        except Exception as e:
            print(f"Error parsing extracted text to DataFrame: {e}")
            # Возвращаем простой DataFrame с текстом
            lines = [line.strip() for line in extracted_text.split('\n') if line.strip()]
            return pd.DataFrame({'Content': lines})
    
    def extract_with_pdf_text_engine(self, file_path: str) -> Optional[str]:
        """
        Использует бесплатный движок pdf-text для структурированных PDF.
        
        :param file_path: Путь к PDF файлу
        :return: Извлеченный текст
        """
        return self.extract_text_with_ocr(file_path, engine="pdf-text", model="anthropic/claude-3.5-sonnet")
    
    def extract_text_from_image(
        self,
        image_base64: str,
        prompt: Optional[str] = None,
        model: str = "anthropic/claude-3.5-sonnet"
    ) -> Optional[str]:
        """
        Извлекает текст из изображения используя vision-модели OpenRouter.
        Используется для отдельных изображений (не PDF).
        
        :param image_base64: Base64 строка изображения в формате data:image/png;base64,...
        :param prompt: Промпт для модели
        :param model: Vision-модель для использования (по умолчанию claude-3.5-sonnet)
        :return: Извлеченный текст или None
        """
        if not self.api_key:
            import logging
            logger = logging.getLogger(__name__)
            logger.error("Error: OpenRouter API key not configured.")
            return None
        
        # Формируем промпт
        default_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."""
        
        if prompt is None:
            prompt = default_prompt
        else:
            prompt = prompt.strip()
            if not prompt:
                prompt = default_prompt
        
        import logging
        logger = logging.getLogger(__name__)
        
        # Формируем запрос для vision-модели
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
            "HTTP-Referer": "https://github.com/ctf-document-processor",
            "X-Title": "CTF Document Processor",
        }
        
        # Для vision-моделей используем image_url в content
        # Формат для OpenRouter: image_url с url в формате data:image/...;base64,...
        messages = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": prompt
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": image_base64  # data:image/png;base64,... или data:image/jpeg;base64,...
                        }
                    }
                ]
            }
        ]
        
        payload = {
            "model": model,
            "messages": messages
        }
        
        try:
            logger.info(f"Sending image to vision model {model}...")
            logger.debug(f"Payload keys: {payload.keys()}")
            logger.debug(f"Messages structure: {len(messages)} messages")
            
            response = requests.post(self.base_url, headers=headers, json=payload, timeout=120)
            response.raise_for_status()
            
            result = response.json()
            logger.debug(f"API Response keys: {result.keys()}")
            
            if 'choices' in result and len(result['choices']) > 0:
                extracted_text = result['choices'][0]['message'].get('content', '')
                logger.info(f"Successfully extracted text from image ({len(extracted_text)} characters)")
                return extracted_text
            else:
                logger.warning(f"Unexpected response format from OpenRouter API")
                logger.warning(f"Response structure: {list(result.keys())}")
                if 'error' in result:
                    logger.error(f"API Error: {result['error']}")
                # Попробуем найти текст в других возможных местах
                if 'content' in result:
                    logger.info("Found 'content' directly in response")
                    return str(result['content'])
                logger.debug(f"Full response (first 500 chars): {str(result)[:500]}")
                return None
                
        except requests.exceptions.RequestException as e:
            logger.error(f"Error calling OpenRouter API for image: {e}")
            if hasattr(e, 'response') and e.response is not None:
                try:
                    error_detail = e.response.json()
                    logger.error(f"API Error details: {error_detail}")
                except:
                    logger.error(f"API Error status: {e.response.status_code}")
            return None
    
    def extract_with_mistral_ocr(self, file_path: str) -> Optional[str]:
        """
        Использует mistral-ocr для сканированных документов и изображений.
        Работает только с PDF файлами!
        
        :param file_path: Путь к PDF файлу
        :return: Извлеченный текст
        """
        return self.extract_text_with_ocr(file_path, engine="mistral-ocr", model="anthropic/claude-3.5-sonnet")