src/agent/tools.py

from langchain_core.tools import tool
from typing import Annotated, Any, Dict, List
from images.image_generator import generate_image
from langgraph.prebuilt import InjectedState
import logging

logger = logging.getLogger(__name__)

def _err(msg: str) -> str:
    logger.error(msg)
    return f"{{ 'error': '{msg}' }}"

def _success(msg: str) -> str:
    logger.info(msg)
    return f"{{ 'success': '{msg}' }}"

@tool
async def generate_scene_image(
    prompt: Annotated[
        str,
        "The prompt to generate an image from"
    ],
    state: InjectedState,
) -> Annotated[
    str,
    "The path to the generated image"
]:
    """
    Generate an image from a prompt and set current scene image.
    """
    try:
        image_path, img_description = generate_image(prompt)
        state["current_scene"]["image"] = image_path
        state["current_scene"]["image_description"] = img_description
        return _success(f"Image generated and set as current scene image: {img_description}")
    except Exception as e:
        return _err(str(e))
    
    
available_tools = [generate_scene_image]