Добавил черновики накладных и OCR через Яндекс. LLM для расшифровки универсальный

ocr-service/yandex_ocr.py

@@ -0,0 +1,137 @@
+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()