bot.py - checkDoc - ingvion/TeacherAssistant

bot.py
86 строк4 KB
ingvionio
last commit
22 сен 2024, 06:17
22 сен 2024, 06:17005ae89
import logging
import os
from docx import Document

from telegram import Update
from telegram.ext import ApplicationBuilder, CommandHandler, MessageHandler, filters
import nest_asyncio

from AnalysOllama import analysis_conclusion, analysis_introduction, analysis_introduction_and_conclusion, \
    analysis_literList
from CheckStructure import check_sections_in_docx, write_sections_to_docx

nest_asyncio.apply()

logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', level=logging.INFO)

# Токен, который вы получили от @BotFather
TOKEN = '7172385372:AAETZ_QzAfjJIpZLLfHiLM0ZmDUJhi2_0GA'  # Replace with your actual token
output = "рекомендованные исправления.docx"

section_keywords = ['список исполнителей','реферат',
                    'содержание', 'термины и определения', 'перечень сокращений и обозначений',
                    'введение', 'заключение',
                    'список использованных источников','приложения',
                    ]



# Функция для обработки команды /start
async def start(update: Update, context) -> None:
    await update.message.reply_text('Привет! Я бот который поможет тебе с твоей работой. Отправь мне твою работу и я проверю её на соответствие ГОСТу')

async def handle_message(update: Update, context):
    await update.message.reply_text('Отправь мне твою работу и я проверю её на соответствие ГОСТу')

async def handle_file(update: Update, context) -> None:
    file = update.message.document


    if not file:
        logging.error("No document found in update.message")
        return

    file_path = await file.get_file()

    # Ensure the 'downloads' directory exists
    os.makedirs('downloads', exist_ok=True)

    file_name = os.path.join("downloads", file.file_name)
    await file_path.download_to_drive(file_name)


    await update.message.reply_text(f'Файл {file.file_name} загружен, начинаю анализ.')

    sections = check_sections_in_docx(file_name, section_keywords)
    write_sections_to_docx(output, sections, file_name)
    newDoc = Document(output)


    if sections.get('введение'):
        introduction = analysis_introduction(file_name, 'введение')
        newDoc.add_paragraph(introduction)
    if sections.get('заключение'):
        conclusion = analysis_conclusion(file_name, 'заключение')
        newDoc.add_paragraph(conclusion)
    if sections.get('введение') and sections.get('заключение'):
        introduction_and_conclusion = analysis_introduction_and_conclusion(file_name, 'введение', 'заключение')
        newDoc.add_paragraph(introduction_and_conclusion)
    if sections.get('список использованных источников'):
        literList = analysis_literList(file_name, 'список использованных источников' )
        newDoc.add_paragraph(literList)
    newDoc.save(output)
    await context.bot.send_document(chat_id=update.effective_chat.id, document=open(output, 'rb'))

async def main() -> None:
    application = ApplicationBuilder().token(TOKEN).build()

    application.add_handler(CommandHandler('start', start))
    application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_message))
    application.add_handler(MessageHandler(filters.Document.ALL, handle_file))

    await application.run_polling()

if __name__ == '__main__':
    import asyncio
    asyncio.run(main())