rmser/ocr-service/yandex_ocr.py

import os
import time
import json
import base64
import logging
import requests
from typing import Optional

logger = logging.getLogger(__name__)

IAM_TOKEN_URL = "https://iam.api.cloud.yandex.net/iam/v1/tokens"
VISION_URL = "https://ocr.api.cloud.yandex.net/ocr/v1/recognizeText"

class YandexOCREngine:
    def __init__(self):
        self.oauth_token = os.getenv("YANDEX_OAUTH_TOKEN")
        self.folder_id = os.getenv("YANDEX_FOLDER_ID")
        
        # Кэширование IAM токена
        self._iam_token = None
        self._token_expire_time = 0

        if not self.oauth_token or not self.folder_id:
            logger.warning("Yandex OCR credentials (YANDEX_OAUTH_TOKEN, YANDEX_FOLDER_ID) not set. Yandex OCR will be unavailable.")

    def _get_iam_token(self) -> Optional[str]:
        """
        Получает IAM-токен. Если есть живой кэшированный — возвращает его.
        Если нет — обменивает OAuth на IAM.
        """
        current_time = time.time()
        
        # Если токен есть и он "свежий" (с запасом в 5 минут)
        if self._iam_token and current_time < self._token_expire_time - 300:
            return self._iam_token

        logger.info("Obtaining new IAM token from Yandex...")
        try:
            response = requests.post(
                IAM_TOKEN_URL,
                json={"yandexPassportOauthToken": self.oauth_token},
                timeout=10
            )
            response.raise_for_status()
            data = response.json()
            
            self._iam_token = data["iamToken"]
            
            # Токен обычно живет 12 часов, но мы будем ориентироваться на поле expiresAt если нужно,
            # или просто поставим таймер. Для простоты берем 1 час жизни кэша.
            self._token_expire_time = current_time + 3600 
            
            logger.info("IAM token received successfully.")
            return self._iam_token
        except Exception as e:
            logger.error(f"Failed to get IAM token: {e}")
            return None

    def recognize(self, image_bytes: bytes) -> str:
        """
        Отправляет изображение в Yandex Vision и возвращает полный текст.
        """
        if not self.oauth_token or not self.folder_id:
            logger.error("Yandex credentials missing.")
            return ""

        iam_token = self._get_iam_token()
        if not iam_token:
            return ""

        # 1. Кодируем в Base64
        b64_image = base64.b64encode(image_bytes).decode("utf-8")

        # 2. Формируем тело запроса
        # Используем модель 'page' (для документов) и '*' для автоопределения языка
        payload = {
            "mimeType": "JPEG", # Yandex переваривает и PNG под видом JPEG часто, но лучше быть аккуратным.
                                # В идеале определять mime-type из файла, но JPEG - безопасный дефолт для фото.
            "languageCodes": ["*"],
            "model": "page",
            "content": b64_image
        }

        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {iam_token}",
            "x-folder-id": self.folder_id,
            "x-data-logging-enabled": "true"
        }

        # 3. Отправляем запрос
        try:
            logger.info("Sending request to Yandex Vision OCR...")
            response = requests.post(VISION_URL, headers=headers, json=payload, timeout=20)
            
            # Если 401 Unauthorized, возможно токен протух раньше времени (редко, но бывает)
            if response.status_code == 401:
                logger.warning("Got 401 from Yandex. Retrying with fresh token...")
                self._iam_token = None # сброс кэша
                iam_token = self._get_iam_token()
                if iam_token:
                    headers["Authorization"] = f"Bearer {iam_token}"
                    response = requests.post(VISION_URL, headers=headers, json=payload, timeout=20)
            
            response.raise_for_status()
            result_json = response.json()
            
            # 4. Парсим ответ
            # Структура: result -> textAnnotation -> fullText
            # Или (если fullText нет) blocks -> lines -> text
            
            text_annotation = result_json.get("result", {}).get("textAnnotation", {})
            
            if not text_annotation:
                logger.warning("Yandex returned success but no textAnnotation found.")
                return ""
            
            # Самый простой способ - взять fullText, он обычно склеен с \n
            full_text = text_annotation.get("fullText", "")
            
            if not full_text:
                # Фолбэк: если fullText пуст, собираем вручную по блокам
                logger.info("fullText empty, assembling from blocks...")
                lines_text = []
                for block in text_annotation.get("blocks", []):
                    for line in block.get("lines", []):
                        lines_text.append(line.get("text", ""))
                full_text = "\n".join(lines_text)

            return full_text

        except Exception as e:
            logger.error(f"Error during Yandex Vision request: {e}")
            return ""

# Глобальный инстанс
yandex_engine = YandexOCREngine()