src/layout_generator/pipeline.py

# src/layout_generator/pipeline.py

"""
Оркестратор генерации 3D-планировок.
Собирает этапы в единый объектно-ориентированный пайплайн и прогоняет контекст.
Включает поддержку потоковой передачи статусов (streaming) для UI.
"""
import logging
import time
from pathlib import Path
from typing import List, Generator, Tuple, Optional, Dict, Any

from hydra import compose, initialize_config_dir
from hydra.core.config_store import ConfigStore
from omegaconf import DictConfig, OmegaConf

from dsynth.scene_gen.hydra_configs import DsConfig, ShelfConfig
from utils.logger import CollisionTracker

from .base import BaseStep, LayoutContext
from .assets import LoadAssetsStep
from .planner import CalculatePlanogramStep
from .topology import GenerateTopologyStep
from .assembler import AssembleSceneStep
from .exporter import ExportAndCompressStep

logger: logging.Logger = logging.getLogger(__name__)

def create_layout_pipeline() -> List[BaseStep]:
    """Фабрика для сборки стандартного пайплайна генерации."""
    return [
        LoadAssetsStep(),
        CalculatePlanogramStep(),
        GenerateTopologyStep(),
        AssembleSceneStep(),
        ExportAndCompressStep()
    ]

def generate_layout_stream(
    project_dir: Path,
    tracker: CollisionTracker,
    size_n: int = 10,
    size_m: int = 10,
    layout_graph: Optional[Dict[str, Any]] = None
) -> Generator[Tuple[Optional[str], Dict[str, Any]], None, None]:
    """
    Потоковая версия оркестратора.
    Возвращает кортежи: (Путь_к_модели_или_None, Словарь_статуса_для_UI).
    Работает исключительно на базе ИИ-графа (layout_graph).
    """
    # Шаг 1: Инициализация (Шаг 0 зарезервирован для LLM в интерфейсе)
    yield None, {"step_idx": 1, "status": "running"}
    start_init = time.time()
    
    conf_dir: str = str(project_dir / 'configs')
    assets_dir: str = str(project_dir / 'assets')

    cs: ConfigStore = ConfigStore.instance()
    try:
        cs.store(group="shelves", name="base_shelf_config", node=ShelfConfig)
        cs.store(group="ds", name="main_darkstore_config_base", node=DsConfig)
    except Exception:
        pass

    logger.info(f"🏗 ЗАДАЧА ПЛАНОГРАММИРОВАНИЯ: Помещение {size_n}x{size_m}")

    with initialize_config_dir(version_base=None, config_dir=conf_dir):
        cfg: DictConfig = compose(config_name="main_config", overrides=[f"ds.size_n={size_n}", f"ds.size_m={size_m}"])

    OmegaConf.set_struct(cfg, False)
    cfg.assets.assets_dir_path = assets_dir
    OmegaConf.resolve(cfg)

    context = LayoutContext(
        project_dir=str(project_dir),
        size_n=size_n, size_m=size_m,
        layout_graph=layout_graph,
        cfg=cfg, tracker=tracker
    )
    
    elapsed_init = time.time() - start_init
    yield None, {"step_idx": 1, "status": "done", "time": elapsed_init}

    pipeline = create_layout_pipeline()

    # Шаги 2-6: Прогоняем пайплайн с замером времени
    for idx, step in enumerate(pipeline):
        step_idx = idx + 2
        yield None, {"step_idx": step_idx, "status": "running"}
        
        start_t = time.time()
        context = step.process(context)
        elapsed_t = time.time() - start_t
        
        # Обработка критических ошибок на лету
        if not context.product_assets_lib:
            yield None, {"step_idx": step_idx, "status": "error", "message": "В кэше отсутствуют модели"}
            return
            
        if isinstance(step, GenerateTopologyStep) and not context.is_gen:
            yield None, {"step_idx": step_idx, "status": "error", "message": "Помещение слишком мало"}
            return

        yield None, {"step_idx": step_idx, "status": "done", "time": elapsed_t}

    yield context.output_filename, {"status": "finished"}

def generate_inventory_driven_layout(
    project_dir: Path,
    tracker: CollisionTracker,
    size_n: int = 10,
    size_m: int = 10
) -> str:
    """Обычная функция для обратной совместимости."""
    final_path = ""
    for path, _ in generate_layout_stream(
        project_dir, tracker, size_n, size_m
    ):
        if path:
            final_path = path
    return final_path