fckbot/rag/functions/generate_document.py

# -*- coding: utf-8 -*-
"""
Модуль генерации документа по шаблону.
Заполняет плейсхолдеры {{...}} на основе данных из базы знаний.

АДАПТАЦИЯ: теперь функция принимает промпт и конфиг как аргументы.
"""

import re
import asyncio
import os
import logging
import shutil
from docx import Document
from core.functions.file_processor import send_file_to_user

logger = logging.getLogger(__name__)


async def generate_document_from_template(
    bot,
    msg,
    room_jid,
    template_path,
    template_name,
    user_jid,
    giga,
    kb,
    config,
    prompt_text,
    temperature=0.1,
    allow_public=False
):
    """
    Фоновая задача: парсит шаблон, ищет плейсхолдеры {{...}},
    для каждого делает контекстный поиск в БЗ комнаты, генерирует значение через LLM,
    заменяет в копии шаблона и отправляет готовый файл в комнату.

    Аргументы:
        bot: экземпляр бота
        msg: объект XMPP-сообщения
        room_jid: JID комнаты
        template_path: путь к шаблону
        template_name: имя шаблона
        user_jid: JID пользователя
        giga: клиент GigaChat
        kb: сервис базы знаний
        config: объект конфигурации
        prompt_text: содержимое промпта для генерации
        temperature: температура генерации
        allow_public: разрешено ли использовать общие знания
    """
    logger.info(f"🚀 [ГЕНЕРАЦИЯ] Запущена для шаблона '{template_name}' в комнате {room_jid}")

    try:
        doc = Document(template_path)

        # Собираем все плейсхолдеры и их контекст
        placeholders_info = {}
        for paragraph in doc.paragraphs:
            for match in re.finditer(r'\{\{([^}]+)\}\}', paragraph.text):
                ph = match.group(1)
                if ph not in placeholders_info:
                    placeholders_info[ph] = paragraph.text

        for table in doc.tables:
            for row in table.rows:
                for cell in row.cells:
                    for match in re.finditer(r'\{\{([^}]+)\}\}', cell.text):
                        ph = match.group(1)
                        if ph not in placeholders_info:
                            placeholders_info[ph] = cell.text

        placeholders = set(placeholders_info.keys())
        logger.info(f"🔍 [ГЕНЕРАЦИЯ] Найдено плейсхолдеров: {len(placeholders)}")

        if not placeholders:
            reply = msg.reply(f"❌ В шаблоне '{template_name}' нет плейсхолдеров {{...}}.")
            if reply:
                reply.send()
            return

        generated = {}
        for ph, surrounding in placeholders_info.items():
            logger.info(f"📖 [ГЕНЕРАЦИЯ] Обработка плейсхолдера '{ph}'")

            # Формируем поисковый запрос
            search_query = f"{ph} {surrounding[:100]}"
            context = await kb.find_relevant_info(search_query, user_jid, room_jid, top_k=20)

            if not context and allow_public:
                logger.info(f"⚠️ [ГЕНЕРАЦИЯ] Контекст для '{ph}' не найден, используем общие знания")
                prompt = f"Ты — эксперт. Ответь на вопрос: {ph}. Дай краткое значение, без пояснений."
                try:
                    response = await giga.chat(
                        history=[],
                        query=prompt,
                        system_prompt=None,
                        file_id=None,
                        temperature=temperature
                    )
                    generated[ph] = response.strip()
                    continue
                except Exception as e:
                    logger.error(f"❌ [ГЕНЕРАЦИЯ] Ошибка LLM для '{ph}': {e}")
                    generated[ph] = "[ошибка]"
                    continue

            if not context:
                logger.warning(f"⚠️ [ГЕНЕРАЦИЯ] Нет контекста для '{ph}', и общие знания запрещены")
                generated[ph] = "не определено"
                continue

            # Генерация значения на основе контекста
            prompt = prompt_text.format(ph=ph, context=context)
            try:
                response = await giga.chat(
                    history=[],
                    query=prompt,
                    system_prompt=None,
                    file_id=None,
                    temperature=temperature
                )
                cleaned = re.sub(r'^```\w*\n?|```$', '', response.strip())
                generated[ph] = cleaned
                logger.info(f"✅ [ГЕНЕРАЦИЯ] Значение для '{ph}': {generated[ph][:100]}")
            except Exception as e:
                logger.error(f"❌ [ГЕНЕРАЦИЯ] Ошибка LLM для '{ph}': {e}", exc_info=True)
                generated[ph] = "[ошибка]"

        # Замена плейсхолдеров в копии документа
        temp_dir = config.temp_dir
        new_path = temp_dir / f"generated_{template_name}_{int(asyncio.get_event_loop().time())}.docx"
        shutil.copy2(template_path, new_path)

        new_doc = Document(new_path)

        def replace_in_paragraph(paragraph, replacements):
            full_text = ''.join(run.text for run in paragraph.runs)
            modified = False
            for key, value in replacements.items():
                placeholder = f"{{{{{key}}}}}"
                if placeholder in full_text:
                    for run in paragraph.runs:
                        if placeholder in run.text:
                            run.text = run.text.replace(placeholder, value)
                            modified = True
            if not modified:
                for key, value in replacements.items():
                    placeholder = f"{{{{{key}}}}}"
                    if placeholder in paragraph.text:
                        paragraph.text = paragraph.text.replace(placeholder, value)

        def walk_document(doc_obj, replacements):
            for paragraph in doc_obj.paragraphs:
                replace_in_paragraph(paragraph, replacements)
            for table in doc_obj.tables:
                for row in table.rows:
                    for cell in row.cells:
                        for paragraph in cell.paragraphs:
                            replace_in_paragraph(paragraph, replacements)

        walk_document(new_doc, generated)
        new_doc.save(new_path)

        await send_file_to_user(bot, msg, str(new_path), prefix=f"📄 Сгенерированный документ по шаблону '{template_name}':", suffix="generated")

        try:
            os.unlink(new_path)
        except Exception as e:
            logger.warning(f"Не удалось удалить временный файл {new_path}: {e}")

        logger.info(f"🏁 [ГЕНЕРАЦИЯ] Завершено для шаблона '{template_name}'")

    except Exception as e:
        logger.exception(f"💥 [ГЕНЕРАЦИЯ] Критическая ошибка: {e}")
        if msg:
            reply = msg.reply(f"❌ Ошибка при генерации документа: {e}")
            if reply:
                reply.send()