scripts/test_llm.py

# scripts/test_llm.py

"""
Скрипт для тестирования работы LLM-клиента.
Обращается к локальному серверу vLLM и проверяет генерацию JSON для раскладки.
Обновлен для работы с новым API generate_layout (добавлены габариты и каталог оборудования).
"""
import logging
import os
import sys
from typing import Any, Dict, List, Optional

import hydra
from omegaconf import DictConfig

# Добавляем корень проекта в PYTHONPATH, чтобы импорты из src работали корректно
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from src.llm.client import InventoryLLMClient

# Настройка подробного логирования
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S"
)
logger: logging.Logger = logging.getLogger(__name__)

@hydra.main(version_base=None, config_path="../configs", config_name="main_config")
def main(cfg: DictConfig) -> None:
    """
    Основная функция для выполнения тестового запроса к запущенному серверу vLLM.
    
    Args:
        cfg (DictConfig): Конфигурация проекта, загруженная через Hydra.
    """
    logger.info("🚀 Запуск тестирования LLM-клиента...")
    
    # Инициализируем клиент (он сам достанет url и модель из cfg.llm)
    client: InventoryLLMClient = InventoryLLMClient(cfg)
    
    # Задаем тестовый набор категорий
    test_categories: List[str] = ["BEER", "CHIPS", "ENERGY_DRINKS", "WATER", "WINE"]
    logger.info(f"🛒 Доступные категории для теста: {test_categories}")
    
    # Подготавливаем тестовые данные для габаритов и оборудования
    test_size_n: float = 6.0
    test_size_m: float = 6.0
    test_equipment: str = (
        "- showcase_glb (Холодильник): 1.55m x 0.9m\n"
        "- shelf_beer (Стеллаж): 1.55m x 0.55m"
    )
    logger.info(f"📏 Габариты комнаты: {test_size_n}x{test_size_m}м")
    logger.info(f"📦 Каталог оборудования:\n{test_equipment}")
    
    # Формируем тестовый запрос
    test_request: str = "Сделай стенд для вечеринки: нужно много пива, чипсов и немного энергетиков. Поставь два стеллажа и один холодильник."
    logger.info(f"🗣️ Тестовый промпт: '{test_request}'")
    
    # Отправляем запрос на сервер с новыми обязательными параметрами
    result: Optional[Dict[str, Any]] = client.generate_layout(
        user_request=test_request, 
        available_categories=test_categories,
        equipment_catalog=test_equipment,
        size_n=test_size_n,
        size_m=test_size_m
    )
    
    if result:
        logger.info("✅ Успешный ответ от нейросети:")
        # Выводим структуру полученного графа
        nodes: List[Dict[str, Any]] = result.get("nodes", [])
        edges: List[Dict[str, Any]] = result.get("edges", [])
        
        logger.info(f"Сгенерировано узлов: {len(nodes)}")
        for node in nodes:
            logger.info(f"  - Узел: {node.get('id')} ({node.get('type')}) на координатах ({node.get('draft_cx')}, {node.get('draft_cy')})")
            
        logger.info(f"Сгенерировано связей: {len(edges)}")
        for edge in edges:
            logger.info(f"  - Связь: {edge.get('source')} -> {edge.get('target')} ({edge.get('relation')})")
    else:
        logger.error("❌ Нейросеть не вернула результат. Проверь логи сервера vLLM.")

if __name__ == "__main__":
    main()