app/process_pdf.py - master - mellcorn/vector-search

app/process_pdf.py
73 строки3 KB
Alex
here we go
16 июл 2025, 02:01
16 июл 2025, 02:01654df04
import os
import logging
from app.utils.utils_pdf import PDFProcessor
from app.settings import PDF_DIRECTORY, GIGACHAT_CLIENT_SECRET, BASE64_CLIENT_SECRET
from app.gigachat import GigaChatAuthClient


class ProcessPDF:
    def __init__(self,  pdf_directory: str = PDF_DIRECTORY):
        self.pdf_directory = pdf_directory
        self.logger = logging.getLogger(__name__)
        self.processor = PDFProcessor()
        logging.basicConfig(level=logging.INFO)

        # Инициализация клиента
        self.giga_chat_client = GigaChatAuthClient(
            client_secret=GIGACHAT_CLIENT_SECRET,
            base64_client_secret=BASE64_CLIENT_SECRET
        )
        self.token = self.giga_chat_client.get_token()
        if self.token:
            # Получение заголовков для последующих запросов
            self.auth_headers = self.giga_chat_client.get_auth_headers()



    def process(self):
        """Основной метод для обработки PDF-файлов."""

        # Проверяем наличие токена (в оригинале - Gigachat token
        if not self.token:
            self.logger.error("Нет соединения с Gigachat. Пожалуйста, проверьте ваш токен.")
            return

        # Проверяем существование директории для PDF-файлов
        if not os.path.exists(self.pdf_directory):
            self.logger.info(f"Создание директории: {self.pdf_directory}")
            os.makedirs(self.pdf_directory)
            self.logger.info(f"Пожалуйста, добавьте файлы PDF в директорию {self.pdf_directory} и запустите скрипт снова")
            return

        # Получаем список всех PDF-файлов в директории
        pdf_files = [f for f in os.listdir(self.pdf_directory) if f.lower().endswith('.pdf')]

        # Проверяем наличие PDF-файлов
        if not pdf_files:
            self.logger.info(f"Файлы PDF не найдены в {self.pdf_directory}")
            self.logger.info(f"Пожалуйста, добавьте файлы PDF в директорию {self.pdf_directory} и запустите скрипт снова")
            return

        # Выводим количество найденных PDF-файлов
        self.logger.info(f"Найдено {len(pdf_files)} PDF-файлов для обработки")

        # Обрабатываем все PDF-файлы
        results = self.processor.process_pdf_directory(self.pdf_directory)

        # Выводим результаты обработки
        self.logger.info("Обработка завершена!")
        self.logger.info("Результаты:")
        for pdf_file, saved_count in results.items():
            if saved_count > 0:
                self.logger.info(f"✓ {pdf_file}: Сохранено {saved_count} фрагментов")
            else:
                self.logger.error(f"✗ {pdf_file}: Не удалось обработать")


    def run(self):
        self.process()

if __name__ == "__main__":
    ProcessPDF(
        pdf_directory= PDF_DIRECTORY
    ).run()