Snap-Solver/models/mathpix.py

from typing import Generator, Dict, Any
import json
import requests
from .base import BaseModel

class MathpixModel(BaseModel):
    """
    Mathpix OCR model for processing images containing mathematical formulas,
    text, and tables.
    """
    
    def __init__(self, api_key: str, temperature: float = 0.7, system_prompt: str = None):
        """
        Initialize the Mathpix model.
        
        Args:
            api_key: Mathpix API key in format "app_id:app_key"
            temperature: Not used for Mathpix but kept for BaseModel compatibility
            system_prompt: Not used for Mathpix but kept for BaseModel compatibility
            
        Raises:
            ValueError: If the API key format is invalid
        """
        # 只传递必需的参数，不传递language参数
        super().__init__(api_key, temperature, system_prompt)
        try:
            self.app_id, self.app_key = api_key.split(':')
        except ValueError:
            raise ValueError("Mathpix API key must be in format 'app_id:app_key'")
        
        self.api_url = "https://api.mathpix.com/v3/text"
        self.headers = {
            "app_id": self.app_id,
            "app_key": self.app_key,
            "Content-Type": "application/json"
        }
        
        # Content type presets
        self.presets = {
            "math": {
                "formats": ["latex_normal", "latex_styled", "asciimath"],
                "data_options": {
                    "include_asciimath": True,
                    "include_latex": True,
                    "include_mathml": True
                },
                "ocr_options": {
                    "detect_formulas": True,
                    "enable_math_ocr": True,
                    "enable_handwritten": True,
                    "rm_spaces": True
                }
            },
            "text": {
                "formats": ["text"],
                "data_options": {
                    "include_latex": False,
                    "include_asciimath": False
                },
                "ocr_options": {
                    "enable_spell_check": True,
                    "enable_handwritten": True,
                    "rm_spaces": False
                }
            },
            "table": {
                "formats": ["text", "data"],
                "data_options": {
                    "include_latex": True
                },
                "ocr_options": {
                    "detect_tables": True,
                    "enable_spell_check": True,
                    "rm_spaces": True
                }
            },
            "full_text": {
                "formats": ["text"],
                "data_options": {
                    "include_latex": False,
                    "include_asciimath": False
                },
                "ocr_options": {
                    "enable_spell_check": True,
                    "enable_handwritten": True,
                    "rm_spaces": False,
                    "detect_paragraphs": True,
                    "enable_tables": False,
                    "enable_math_ocr": False
                }
            }
        }
        
        # Default to math preset
        self.current_preset = "math"

    def analyze_image(self, image_data: str, proxies: dict = None, content_type: str = None, 
                     confidence_threshold: float = 0.8, max_retries: int = 3) -> Generator[dict, None, None]:
        """
        Analyze an image using Mathpix OCR API.
        
        Args:
            image_data: Base64 encoded image data
            proxies: Optional proxy configuration
            content_type: Type of content to analyze ('math', 'text', or 'table')
            confidence_threshold: Minimum confidence score to accept (0.0 to 1.0)
            max_retries: Maximum number of retry attempts for failed requests
            
        Yields:
            dict: Response chunks with status and content
        """
        if content_type and content_type in self.presets:
            self.current_preset = content_type

        preset = self.presets[self.current_preset]
        
        try:
            # Prepare request payload
            payload = {
                "src": f"data:image/jpeg;base64,{image_data}",
                "formats": preset["formats"],
                "data_options": preset["data_options"],
                "ocr_options": preset["ocr_options"]
            }
            
            # Initialize retry counter
            retry_count = 0
            
            while retry_count < max_retries:
                try:
                    # Send request to Mathpix API with timeout
                    response = requests.post(
                        self.api_url,
                        headers=self.headers,
                        json=payload,
                        proxies=proxies,
                        timeout=25  # 25 second timeout
                    )
                    
                    # Handle specific API error codes
                    if response.status_code == 429:  # Rate limit exceeded
                        if retry_count < max_retries - 1:
                            retry_count += 1
                            continue
                        else:
                            raise requests.exceptions.RequestException("Rate limit exceeded")
                    
                    response.raise_for_status()
                    result = response.json()
                    
                    # Check confidence threshold
                    if 'confidence' in result and result['confidence'] < confidence_threshold:
                        yield {
                            "status": "warning",
                            "content": f"Low confidence score: {result['confidence']:.2%}"
                        }
                    
                    break  # Success, exit retry loop
                    
                except (requests.exceptions.Timeout, requests.exceptions.ConnectionError):
                    if retry_count < max_retries - 1:
                        retry_count += 1
                        continue
                    raise
            
            # Format the response
            formatted_response = self._format_response(result)
            
            # Yield initial status
            yield {
                "status": "started",
                "content": ""
            }
            
            # Yield the formatted response
            yield {
                "status": "completed",
                "content": formatted_response,
                "model": self.get_model_identifier()
            }
            
        except requests.exceptions.RequestException as e:
            yield {
                "status": "error",
                "error": f"Mathpix API error: {str(e)}"
            }
        except Exception as e:
            yield {
                "status": "error",
                "error": f"Error processing image: {str(e)}"
            }

    def analyze_text(self, text: str, proxies: dict = None) -> Generator[dict, None, None]:
        """
        Not implemented for Mathpix model as it only processes images.
        """
        yield {
            "status": "error",
            "error": "Text analysis is not supported by Mathpix model"
        }

    def get_default_system_prompt(self) -> str:
        """
        Not used for Mathpix model.
        """
        return ""

    def get_model_identifier(self) -> str:
        """
        Return the model identifier.
        """
        return "mathpix"

    def _format_response(self, result: Dict[str, Any]) -> str:
        """
        Format the Mathpix API response into a readable string.
        
        Args:
            result: Raw API response from Mathpix
            
        Returns:
            str: Formatted response string with all available formats
        """
        formatted_parts = []
        
        # Add confidence score if available
        if 'confidence' in result:
            formatted_parts.append(f"Confidence: {result['confidence']:.2%}\n")
        
        # Add text content
        if 'text' in result:
            formatted_parts.append("Text Content:")
            formatted_parts.append(result['text'])
            formatted_parts.append("")
        
        # Add LaTeX content
        if 'latex_normal' in result:
            formatted_parts.append("LaTeX (Normal):")
            formatted_parts.append(result['latex_normal'])
            formatted_parts.append("")
            
        if 'latex_styled' in result:
            formatted_parts.append("LaTeX (Styled):")
            formatted_parts.append(result['latex_styled'])
            formatted_parts.append("")
        
        # Add data formats (ASCII math, MathML)
        if 'data' in result and isinstance(result['data'], list):
            for item in result['data']:
                item_type = item.get('type', '')
                if item_type and 'value' in item:
                    formatted_parts.append(f"{item_type.upper()}:")
                    formatted_parts.append(item['value'])
                    formatted_parts.append("")
        
        # Add table data if present
        if 'tables' in result and result['tables']:
            formatted_parts.append("Tables Detected:")
            for i, table in enumerate(result['tables'], 1):
                formatted_parts.append(f"Table {i}:")
                if 'cells' in table:
                    # Format table as a grid
                    cells = table['cells']
                    if cells:
                        max_col = max(cell.get('col', 0) for cell in cells) + 1
                        max_row = max(cell.get('row', 0) for cell in cells) + 1
                        grid = [['' for _ in range(max_col)] for _ in range(max_row)]
                        
                        for cell in cells:
                            row = cell.get('row', 0)
                            col = cell.get('col', 0)
                            text = cell.get('text', '')
                            grid[row][col] = text
                        
                        # Format grid as table
                        col_widths = [max(len(str(grid[r][c])) for r in range(max_row)) for c in range(max_col)]
                        for row in grid:
                            row_str = ' | '.join(f"{str(cell):<{width}}" for cell, width in zip(row, col_widths))
                            formatted_parts.append(f"| {row_str} |")
                formatted_parts.append("")
        
        # Add error message if present
        if 'error' in result:
            error_msg = result['error']
            if isinstance(error_msg, dict):
                error_msg = error_msg.get('message', str(error_msg))
            formatted_parts.append(f"Error: {error_msg}")
        
        return "\n".join(formatted_parts).strip()

    def extract_full_text(self, image_data: str, proxies: dict = None, max_retries: int = 3) -> str:
        """
        专门用于提取图像中的全部文本内容，忽略数学公式和表格等其他元素。
        
        Args:
            image_data: Base64编码的图像数据
            proxies: 可选的代理配置
            max_retries: 请求失败时的最大重试次数
            
        Returns:
            str: 图像中提取的完整文本内容
        """
        try:
            # 准备请求负载，使用专为全文提取配置的参数
            payload = {
                "src": f"data:image/jpeg;base64,{image_data}",
                "formats": ["text"],
                "data_options": {
                    "include_latex": False,
                    "include_asciimath": False
                },
                "ocr_options": {
                    "enable_spell_check": True,
                    "enable_handwritten": True,
                    "rm_spaces": False,
                    "detect_paragraphs": True,
                    "enable_tables": False,
                    "enable_math_ocr": False
                }
            }
            
            # 初始化重试计数器
            retry_count = 0
            
            while retry_count < max_retries:
                try:
                    # 发送请求到Mathpix API
                    response = requests.post(
                        self.api_url,
                        headers=self.headers,
                        json=payload,
                        proxies=proxies,
                        timeout=30  # 30秒超时
                    )
                    
                    # 处理特定API错误代码
                    if response.status_code == 429:  # 超出速率限制
                        if retry_count < max_retries - 1:
                            retry_count += 1
                            continue
                        else:
                            raise requests.exceptions.RequestException("超出API速率限制")
                    
                    response.raise_for_status()
                    result = response.json()
                    
                    # 直接返回文本内容
                    if 'text' in result:
                        return result['text']
                    else:
                        return "未能提取到文本内容"
                    
                except (requests.exceptions.Timeout, requests.exceptions.ConnectionError):
                    if retry_count < max_retries - 1:
                        retry_count += 1
                        continue
                    raise
            
        except requests.exceptions.RequestException as e:
            return f"Mathpix API错误: {str(e)}"
        except Exception as e:
            return f"处理图像时出错: {str(e)}"