@@ -356,13 +356,6 @@ YANDEX_CLOUD_ID = env.str('YANDEX_CLOUD_ID', 'defaultapikey') OPENAI_PROXY_HOST = env.str('OPENAI_PROXY_HOST', 'neuron-proxy:8080') UPSCALE_MULTIPLIER_HOST = env.str('UPSCALE_MULTIPLIER_HOST', 'packet:8080') - -MAX_UPLOAD_SIZE_PER_MODEL = { - 'raifgpt': 50, - 'default': 8, -} - - # Payments YOOKASSA_ACCOUNT_ID = env.str('YOOKASSA_ACCOUNT_ID', default='defaultapikey') YOOKASSA_SECRET_KEY = env.str('YOOKASSA_SECRET_KEY', default='defaultapikey') @@ -40,11 +40,7 @@ class MessageSerializer(serializers.ModelSerializer): def validate(self, data: Dict[str, Any]) -> Dict[str, Any]: file = data.get('file') - version = data.get('info', {}).get('inference', 'default') - max_mb_size = settings.MAX_UPLOAD_SIZE_PER_MODEL.get( - version, - settings.MAX_UPLOAD_SIZE_PER_MODEL['default'] - ) + max_mb_size = 52 if file and file.size > (max_mb_size << 10 << 10): raise ValidationError( _('The file size cannot exceed %(max_mb_size)d MB') % {'max_mb_size': max_mb_size} @@ -41,6 +41,7 @@ class OpenAICompatibleRunner(BaseRunner, ABC): @classmethod def generate(cls, content=None, file=None, parameters={}, history=[], scrape_results=[]): proxies = Proxy.objects.all() + system_prompt = parameters.pop('system_prompt', None) payload = { 'messages': [ *[ @@ -55,6 +56,10 @@ class OpenAICompatibleRunner(BaseRunner, ABC): 'stream': True, **parameters, } + + if system_prompt: + payload['messages'].insert(0, {'role': 'system', 'content': system_prompt}) + for result in scrape_results: if isinstance(result, StringIO): payload['messages'].append( @@ -4,11 +4,13 @@ import re import subprocess import time import zipfile +import fitz +import httpx from datetime import timedelta from decimal import Decimal from io import BytesIO, StringIO from math import ceil -from typing import Any, Callable, Iterable, Literal +from typing import Any, Callable, Literal, Tuple, Union, Optional, Iterable from uuid import UUID import docx2txt @@ -16,7 +18,8 @@ import openpyxl import filetype from django.core.cache import cache from django.core.files.base import ContentFile -from django.db.models import Prefetch +from django.conf import settings +from django.db.models import Prefetch, QuerySet from django.utils.translation import gettext_lazy as _ from PIL import Image as ImageModule from PyPDF2 import PdfReader @@ -110,6 +113,7 @@ class InferenceService: logger.warning('Payment biases are missing; Inference: %s' % (inference.name)) file = None + image_count = None raw_file = input_message.file raw_info = input_message.info.copy() scrape_results = [] @@ -141,9 +145,11 @@ class InferenceService: file.seek(0) elif file_extension in ('pdf',): file = StringIO() - reader = PdfReader(file_buf) - for page in reader.pages: - file.write(page.extract_text()) + pdf_data = self.get_pdf_data(file_buf) + raw_text = pdf_data[0] + image_count = pdf_data[1] + file.write(raw_text) + file.seek(0) elif file_extension in ('doc', 'docx'): extractors: dict[Literal['doc', 'docx'], Callable[[], str]] = { 'doc': lambda: subprocess.Popen( @@ -269,8 +275,7 @@ class InferenceService: file_data ) * payment_rule.cost elif ( - content - and payment_rule.strategy == PaymentRule.StrategyChoices.PER_TEXT_TOKEN + payment_rule.strategy == PaymentRule.StrategyChoices.PER_TEXT_TOKEN and payment_rule.interaction_type == PaymentRule.InteractionTypeChoices.INPUT ): calculated_price += TokenizerTool.token_count(content) * payment_rule.cost @@ -287,8 +292,13 @@ class InferenceService: ) if scrape_results: calculated_price += ( - sum([TokenizerTool.token_count(result.getvalue()) for result in scrape_results]) + sum([TokenizerTool.token_count(result.getvalue()) for result in + scrape_results]) ) * payment_rule.cost + if (system_prompt := parameters.get('system_prompt')): + calculated_price += TokenizerTool.token_count(system_prompt) * payment_rule.cost + if image_count: + calculated_price += image_count * Decimal('0.13') elif ( payment_rule.strategy == PaymentRule.StrategyChoices.PER_TEXT_TOKEN and payment_rule.interaction_type == PaymentRule.InteractionTypeChoices.OUTPUT @@ -397,3 +407,154 @@ class InferenceService: logger.exception(exc) finally: cache.delete(cache_key) + + def get_pdf_data(self, input_data: BytesIO) -> (str, int): + try: + processor = PDFProcessor() + result, image_count = processor.process(file_stream=input_data) + return result, image_count + except Exception as exc: + logger.exception(exc) + return f"Ошибка обработки PDF: {str(exc)}" + finally: + input_data.seek(0) + + +class PDFProcessor: + MAX_BATCH_SIZE = 3.9 * 1024 * 1024 + + def __init__(self): + self.image_count = 0 + + def process(self, file_stream: BytesIO) -> Tuple[str, int]: + try: + pdf_data = file_stream.read() + doc = fitz.open(stream=pdf_data, filetype="pdf") + has_images = any(page.get_images() for page in doc) + if has_images: + text = self._process_ocr(doc) + return text, self.image_count + raw_text = [] + for page in doc: + content = page.get_text("text") + if content: + raw_text.append(content) + return "\n".join(raw_text), 0 + finally: + doc.close() + fitz.TOOLS.store_shrink(100) + + def _process_ocr(self, doc) -> str: + raw_texts, pages_with_image = self._extract_text_and_images(doc) + ocr_texts = self._process_images_with_yandex_vision(doc, pages_with_image) + final_text = self._combine_texts(raw_texts, ocr_texts) + return final_text + + def _extract_text_and_images(self, doc) -> Tuple[dict, list]: + raw_texts = {} + pages_with_image = [] + for page_num, page in enumerate(doc): + text = page.get_text("text") + if text: + raw_texts[page_num] = text + if page.get_images(): + pages_with_image.append(page_num) + return raw_texts, pages_with_image + + def _process_images_with_yandex_vision(self, doc, page_nums) -> dict: + batch_images, page_index_map = self._prepare_image_batches(doc, page_nums) + self.image_count = len(batch_images) + if not batch_images: + return {} + return self._send_to_yandex_vision(batch_images, page_index_map) + + def _prepare_image_batches(self, doc, page_nums) -> Tuple[list, list]: + batch_images = [] + page_index_map = [] + for page_num in page_nums: + try: + page = doc.load_page(page_num) + pix = page.get_pixmap(dpi=150, alpha=False) + img = ImageModule.frombytes("RGB", [pix.width, pix.height], pix.samples) + buffer = BytesIO() + img.save(buffer, format="JPEG", quality=60, optimize=True) + buffer.seek(0) + if buffer.getbuffer().nbytes < self.__class__.MAX_BATCH_SIZE: + batch_images.append(buffer) + page_index_map.append(page_num) + except Exception as e: + logger.error(f"Error processing page {page_num}: {e}") + continue + return batch_images, page_index_map + + def _send_to_yandex_vision(self, batch_images, page_index_map) -> dict: + headers = { + "Authorization": f"Api-Key {settings.YANDEX_CLOUD_API_KEY}", + "Content-Type": "application/json" + } + ocr_results = {} + batches = self._create_batches(batch_images, page_index_map) + for batch, pages in batches: + body = { + "folderId": settings.YANDEX_CLOUD_ID, + "analyze_specs": [{ + "content": base64.b64encode(buf.getvalue()).decode(), + "features": [{ + "type": "TEXT_DETECTION", + "text_detection_config": {"language_codes": ["*"]} + }] + } for buf in batch] + } + try: + resp = httpx.post( + "https://vision.api.cloud.yandex.net/vision/v1/batchAnalyze", + headers=headers, json=body, timeout=60 + ) + if resp.status_code >= 400: + self.image_count = 0 + return {} + response = resp.json() if resp.status_code == 200 else None + if response: + self._parse_vision_response(response, pages, ocr_results) + except Exception as e: + logger.error(f"Yandex Vision API error: {e}") + return {} + return ocr_results + + def _create_batches(self, batch_images, page_index_map) -> list: + batches = [] + current_batch = [] + current_pages = [] + current_size = 0 + for i, buffer in enumerate(batch_images): + size = buffer.getbuffer().nbytes + if current_size + size > self.__class__.MAX_BATCH_SIZE and current_batch: + batches.append((current_batch, current_pages)) + current_batch, current_pages, current_size = [], [], 0 + current_batch.append(buffer) + current_pages.append(page_index_map[i]) + current_size += size + if current_batch: + batches.append((current_batch, current_pages)) + return batches + + def _parse_vision_response(self, response, pages, ocr_results): + for i, spec_result in enumerate(response.get("results", [])): + page_text = [] + for res in spec_result.get("results", []): + for page in res.get("textDetection", {}).get("pages", []): + for block in page.get('blocks', []): + for line in block.get('lines', []): + line_text = " ".join( + word.get('text', '') for word in line.get('words', []) + ) + if line_text: + page_text.append(line_text) + ocr_results[pages[i]] = "\n".join(page_text) + + def _combine_texts(self, raw_texts, ocr_texts) -> str: + all_pages = sorted(set(raw_texts) | set(ocr_texts)) + return "\n\n".join( + f"{raw_texts.get(pn, '')}\n{ocr_texts.get(pn, '')}".strip() + for pn in all_pages + ).strip() or "Не удалось распознать текст" @@ -60,6 +60,10 @@ EMAIL_USE_TLS=1 EMAIL_USE_SSL=0 ERROR_EMAIL_RECIPIENTS=help@root.ru +# YANDEX OCR +YANDEX_CLOUD_API_KEY=AQVNxHnXmp2iZStdOPzMVl6pMuBlCDHPKbWTqqQV +YANDEX_CLOUD_ID=b1gi6c26bkehs7tsa1m9 + # PAYMENTS AND CONFIRMATIONS YOOKASSA_ACCOUNT_ID=322563 YOOKASSA_SECRET_KEY=test_i_Au0KbXnOmdVf1icljT7v4CuDHLG8mXVkyofJQFBns