#!/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()