ai-benchmark/scripts/generate_summarization_from_mongo.py

#!/usr/bin/env python3
"""
Скрипт для генерации тестов пересказов из MongoDB.

Извлекает текст статьи из коллекции rssNotification (поле .meta.topicContent)
и генерирует тестовые данные в формате TXT для бенчмарка AI.
"""

import argparse
import json
import sys
from pathlib import Path
from typing import Dict, Optional

import pymongo
from pymongo import MongoClient

# Добавляем путь к исходникам, чтобы импортировать константы
sys.path.insert(0, str(Path(__file__).parent.parent))

from src.constants import TEST_SEPARATOR

def sanitize_filename(filename: str) -> str:
    """
    Очищает строку от недопустимых символов для имени файла.

    Args:
        filename: Исходное имя файла

    Returns:
        Очищенное имя файла или пустая строка, если очистка невозможна
    """
    import re
    # Заменяем недопустимые символы на подчеркивание
    # Допустимые символы: буквы, цифры, подчеркивание, тире, точка
    sanitized = re.sub(r'[^\w\-\.]', '_', filename)
    return sanitized if sanitized else filename

def connect_to_mongo() -> MongoClient:
    """Подключается к MongoDB кластеру."""
    client = MongoClient(
        "mongodb://10.0.0.3:27017,10.0.0.4:27017,10.0.0.5:27017/",
        connectTimeoutMS=30000,
        socketTimeoutMS=30000,
        serverSelectionTimeoutMS=30000,
        retryWrites=True,
        retryReads=True
    )
    return client

def extract_text_from_topic_content(topic_content: Dict) -> Optional[str]:
    """
    Извлекает текст статьи из .meta.topicContent.

    Args:
        topic_content: Содержимое поля .meta.topicContent из MongoDB

    Returns:
        Текст статьи или None, если не удалось извлечь
    """
    if not topic_content:
        return None

    # Преобразуем в строку, если это не строка
    content_str = str(topic_content)

    return content_str

def generate_test_from_mongo_record(record_id: str) -> bool:
    """
    Генерирует тест пересказа из записи MongoDB.

    Args:
        record_id: ID записи в MongoDB

    Returns:
        True, если тест успешно generated, False в случае ошибки
    """
    try:
        client = connect_to_mongo()
        db = client['tracker_conbot']
        collection = db['rssNotification']

        # Извлекаем запись по ID
        record = collection.find_one({"_id": record_id})
        if not record:
            print(f"❌ Запись с ID {record_id} не найдена в коллекции")
            return False

        # Отладочная информация
        print(f"🔍 Найдена запись: {record_id}")
        print(f"📋 Полная структура записи:")
        print(json.dumps(record, ensure_ascii=False, indent=2, default=str))

        # Извлекаем текст из meta.topicContent
        meta_data = record.get('meta', {})
        topic_content = meta_data.get('topicContent')
        if not topic_content:
            print(f"❌ В записи {record_id} отсутствует поле meta.topicContent")
            return False

        print(f"📝 Тип поля meta.topicContent: {type(topic_content)}")
        print(f"📝 Содержимое meta.topicContent (первые 500 символов):")
        print(str(topic_content)[:500])

        # Извлекаем текст
        article_text = extract_text_from_topic_content(topic_content)
        if not article_text:
            print(f"❌ Не удалось извлечь текст из meta.topicContent записи {record_id}")
            return False

        print(f"📝 Итоговый текст (первые 500 символов): {article_text[:500]}")

        # Генерируем пересказ через LLM (если доступно)
        expected_summary = ""
        try:
            # Пытаемся сгенерировать пересказ
            summary_prompt = f"""Summarize the following text in 1-2 sentences:
"{article_text}"
Provide only the summary, no additional text."""
            # Если Ollama доступен, можно было бы использовать его здесь
            # Для простоты оставляем пустым или можно добавить логику позже
            expected_summary = ""
        except:
            # Если генерация не удалась, оставляем пустым
            expected_summary = ""

        # Создаем директорию для сохранения теста (всегда в tests/summarization)
        output_path = Path("tests") / "summarization"
        output_path.mkdir(parents=True, exist_ok=True)

        # Очищаем ID от недопустимых символов для имени файла
        filename = sanitize_filename(record_id)
        if not filename:
            print(f"❌ Не удалось создать допустимое имя файла из ID записи {record_id}")
            return False

        # Используем очищенный ID записи как имя файла
        test_file = output_path / f"{filename}.txt"

        # Сохраняем текст статьи и ожидаемый пересказ с разделителем
        with open(test_file, "w", encoding="utf-8") as f:
            f.write(f"{article_text}{TEST_SEPARATOR}{expected_summary}")

        print(f"✅ Создан тест tests/summarization/{filename}.txt")
        print(f"   Источник: MongoDB запись {record_id}")
        print(f"   Текст статьи (первые 100 символов): {article_text[:100]}...")

        return True

    except Exception as e:
        print(f"❌ Ошибка при генерации теста: {e}")
        return False
    finally:
        if 'client' in locals():
            client.close()

def validate_test(test_data: Dict[str, str]) -> bool:
    """Валидирует тестовые данные."""
    if not isinstance(test_data, dict):
        print("❌ Тест должен быть словарём (JSON объект)")
        return False

    if "prompt" not in test_data:
        print("❌ Отсутствует поле 'prompt'")
        return False

    if "expected" not in test_data:
        print("❌ Отсутствует поле 'expected'")
        return False

    if not isinstance(test_data["prompt"], str):
        print("❌ Поле 'prompt' должно быть строкой")
        return False

    if not isinstance(test_data["expected"], str):
        print("❌ Поле 'expected' должно быть строкой")
        return False

    if not test_data["prompt"].strip():
        print("❌ Поле 'prompt' не может быть пустым")
        return False

    if not test_data["expected"].strip():
        print("❌ Поле 'expected' не может быть пустым")
        return False

    return True

def main():
    """Основная функция скрипта."""
    parser = argparse.ArgumentParser(
        description="Генератор тестов пересказов из MongoDB",
        epilog="Примеры использования:\n"
               "  python scripts/generate_summarization_from_mongo.py --record-id 507f1f77bcf86cd799439011"
    )
    parser.add_argument(
        "--record-id",
        type=str,
        required=True,
        help="ID записи в MongoDB (обязательный параметр)"
    )

    args = parser.parse_args()

    print(f"🔍 Подключаюсь к MongoDB кластеру...")
    print(f"📄 Извлекаю запись с ID: {args.record_id}")
    print(f"💾 Сохраняю тест в: tests/summarization/")

    success = generate_test_from_mongo_record(args.record_id)

    if success:
        print("\n✨ Готово! Тест успешно generated.")
    else:
        print("\n❌ Не удалось generated тест.")
        sys.exit(1)

if __name__ == "__main__":
    main()