chwellofficial/nt360Slides

1import asyncio2from typing import List, Optional3from models.image_prompt import ImagePrompt4from models.sql.image_asset import ImageAsset5from models.sql.slide import SlideModel6from services.icon_finder_service import ICON_FINDER_SERVICE7from services.image_generation_service import ImageGenerationService8from utils.dict_utils import get_dict_at_path, get_dict_paths_with_key, set_dict_at_path9 10 11async def process_slide_and_fetch_assets(12    image_generation_service: ImageGenerationService,13    slide: SlideModel,14    outline_image_urls: Optional[List[str]] = None,15) -> List[ImageAsset]:16 17    async_tasks = []18    async_task_meta = []19 20    image_paths = get_dict_paths_with_key(slide.content, "__image_prompt__")21    icon_paths = get_dict_paths_with_key(slide.content, "__icon_query__")22 23    for image_index, image_path in enumerate(image_paths):24        __image_prompt__parent = get_dict_at_path(slide.content, image_path)25 26        if (27            outline_image_urls28            and image_index < len(outline_image_urls)29            and outline_image_urls[image_index]30        ):31            __image_prompt__parent["__image_url__"] = outline_image_urls[image_index]32            set_dict_at_path(slide.content, image_path, __image_prompt__parent)33            continue34 35        async_tasks.append(36            image_generation_service.generate_image(37                ImagePrompt(38                    prompt=__image_prompt__parent["__image_prompt__"],39                )40            )41        )42        async_task_meta.append(("image", image_path))43 44    for icon_path in icon_paths:45        __icon_query__parent = get_dict_at_path(slide.content, icon_path)46        async_tasks.append(47            ICON_FINDER_SERVICE.search_icons(__icon_query__parent["__icon_query__"])48        )49        async_task_meta.append(("icon", icon_path))50 51    results = await asyncio.gather(*async_tasks) if async_tasks else []52 53    return_assets = []54    for (task_type, asset_path), result in zip(async_task_meta, results):55        if task_type == "image":56            image_dict = get_dict_at_path(slide.content, asset_path)57            if isinstance(result, ImageAsset):58                return_assets.append(result)59                image_dict["__image_url__"] = result.path60            else:61                image_dict["__image_url__"] = result62            set_dict_at_path(slide.content, asset_path, image_dict)63            continue64 65        icon_dict = get_dict_at_path(slide.content, asset_path)66        # ICON_FINDER_SERVICE.search_icons returns a list of URLs67        if isinstance(result, list) and result:68            icon_dict["__icon_url__"] = result[0]69        else:70            # Fallback to FastAPI static placeholder if no icon found71            icon_dict["__icon_url__"] = "/static/icons/placeholder.svg"72        set_dict_at_path(slide.content, asset_path, icon_dict)73 74    return return_assets75 76 77async def process_old_and_new_slides_and_fetch_assets(78    image_generation_service: ImageGenerationService,79    old_slide_content: dict,80    new_slide_content: dict,81) -> List[ImageAsset]:82    # Finds all old images83    old_image_dict_paths = get_dict_paths_with_key(84        old_slide_content, "__image_prompt__"85    )86    old_image_dicts = [87        get_dict_at_path(old_slide_content, path) for path in old_image_dict_paths88    ]89    old_image_prompts = [90        old_image_dict["__image_prompt__"] for old_image_dict in old_image_dicts91    ]92 93    # Finds all old icons94    old_icon_dict_paths = get_dict_paths_with_key(old_slide_content, "__icon_query__")95    old_icon_dicts = [96        get_dict_at_path(old_slide_content, path) for path in old_icon_dict_paths97    ]98    old_icon_queries = [99        old_icon_dict["__icon_query__"] for old_icon_dict in old_icon_dicts100    ]101 102    # Finds all new images103    new_image_dict_paths = get_dict_paths_with_key(104        new_slide_content, "__image_prompt__"105    )106    new_image_dicts = [107        get_dict_at_path(new_slide_content, path) for path in new_image_dict_paths108    ]109 110    # Finds all new icons111    new_icon_dict_paths = get_dict_paths_with_key(new_slide_content, "__icon_query__")112    new_icon_dicts = [113        get_dict_at_path(new_slide_content, path) for path in new_icon_dict_paths114    ]115 116    # Creates async tasks for fetching new images117    async_image_fetch_tasks = []118    new_images_fetch_status = []119 120    # Creates async tasks for fetching new icons121    async_icon_fetch_tasks = []122    new_icons_fetch_status = []123 124    # Creates async tasks for fetching new images125    # Use old image url if prompt is same126    for new_image in new_image_dicts:127        if new_image["__image_prompt__"] in old_image_prompts:128            old_image_url = old_image_dicts[129                old_image_prompts.index(new_image["__image_prompt__"])130            ]["__image_url__"]131            new_image["__image_url__"] = old_image_url132            new_images_fetch_status.append(False)133            continue134 135        async_image_fetch_tasks.append(136            image_generation_service.generate_image(137                ImagePrompt(138                    prompt=new_image["__image_prompt__"],139                )140            )141        )142        new_images_fetch_status.append(True)143 144    # Creates async tasks for fetching new icons145    # Use old icon url if query is same146    for new_icon in new_icon_dicts:147        if new_icon["__icon_query__"] in old_icon_queries:148            old_icon_url = old_icon_dicts[149                old_icon_queries.index(new_icon["__icon_query__"])150            ]["__icon_url__"]151            new_icon["__icon_url__"] = old_icon_url152            new_icons_fetch_status.append(False)153            continue154 155        async_icon_fetch_tasks.append(156            ICON_FINDER_SERVICE.search_icons(new_icon["__icon_query__"])157        )158        new_icons_fetch_status.append(True)159 160    new_images = await asyncio.gather(*async_image_fetch_tasks)161    new_icons = await asyncio.gather(*async_icon_fetch_tasks)162 163    # list of new assets164    new_assets = []165 166    # Sets new image and icon urls for assets that were fetched167    for i, _ in enumerate(new_images):168        if new_images_fetch_status[i]:169            fetched_image = new_images[i]170            if isinstance(fetched_image, ImageAsset):171                new_assets.append(fetched_image)172                image_url = fetched_image.path173            else:174                image_url = fetched_image175            new_image_dicts[i]["__image_url__"] = image_url176 177    for i, _ in enumerate(new_icons):178        if new_icons_fetch_status[i]:179            icon_result = new_icons[i]180            if icon_result and len(icon_result) > 0:181                new_icon_dicts[i]["__icon_url__"] = icon_result[0]182            else:183                # Fallback to placeholder if no icon found184                new_icon_dicts[i]["__icon_url__"] = "/static/icons/placeholder.svg"185 186    for i, new_image_dict in enumerate(new_image_dicts):187        set_dict_at_path(new_slide_content, new_image_dict_paths[i], new_image_dict)188 189    for i, new_icon_dict in enumerate(new_icon_dicts):190        set_dict_at_path(new_slide_content, new_icon_dict_paths[i], new_icon_dict)191 192    return new_assets193 194 195def process_slide_add_placeholder_assets(slide: SlideModel):196 197    image_paths = get_dict_paths_with_key(slide.content, "__image_prompt__")198    icon_paths = get_dict_paths_with_key(slide.content, "__icon_query__")199 200    for image_path in image_paths:201        image_dict = get_dict_at_path(slide.content, image_path)202        # Use FastAPI static path for placeholder image203        image_dict["__image_url__"] = "/static/images/placeholder.jpg"204        set_dict_at_path(slide.content, image_path, image_dict)205 206    for icon_path in icon_paths:207        icon_dict = get_dict_at_path(slide.content, icon_path)208        # Use FastAPI static path for placeholder icon209        icon_dict["__icon_url__"] = "/static/icons/placeholder.svg"210        set_dict_at_path(slide.content, icon_path, icon_dict)211