llm/tools.py

import uuid
from typing import Dict, Optional

import replicate
from dotenv import load_dotenv
from langchain_core.runnables import RunnableConfig, RunnableLambda
from pydantic import BaseModel

from llm.prompt import generate_image_tool_description
from llm.utils import create_or_update_ip_generation_count, get_ip_generation_count, store_tool_result, upload_generated_image_to_s3

load_dotenv()


# The generate_image tool's input schema
class GenerateImageToolInput(BaseModel):
    prompt: str
    user_id: str
    image_url: str
    title: Optional[str] = "Generated Image"


# The core function that generates an image of the tool
def _generate_image_core(
    prompt: str,
    user_id: str,
    image_url: str,
    title: str,
    client_ip: str,
) -> str:
    """
    Generate an image based on a prompt.
    """
    print(f"[TOOL] generate_image called with prompt: {prompt[:50]}..., user_id: {user_id}, image_url: {image_url[:50]}...")

    # Check if the user has exceeded the generation limit
    if get_ip_generation_count(client_ip) >= 10:
        print("[TOOL] User exceeded the generation limit of 10 this week.")
        return "Failed as user exceeded the max generation limit of 10 this week."

    use_sdxl = False  # True for testing purposes
    if use_sdxl:
        input = {
            "width": 768,
            "height": 768,
            "prompt": prompt,
            "refine": "expert_ensemble_refiner",
            "apply_watermark": False,
            "num_inference_steps": 25,
            "prompt_strength": 0.5,
            "image": image_url,
            "input_image": image_url,
            "output_format": "png",
        }
        version = "stability-ai/sdxl:" "7762fd07cf82c948538e41f63f77d685e02b063e37e496e96eefd46c929f9bdc"
        output = replicate.run(
            version,
            input=input,
        )
        generated_image_url = output[0] if isinstance(output, list) else output
    else:  # Flux Kontext Pro
        # Generate image using Replicate
        input = {
            "prompt": prompt,
            "input_image": image_url,
            "output_format": "png",
        }
        output = replicate.run(
            "black-forest-labs/flux-kontext-pro",
            input=input,
        )
        print(f"[TOOL] Replicate output: {output}")
        print(f"[TOOL] Output type: {type(output)}")
        print(f"[TOOL] Output length: {len(output) if hasattr(output, '__len__') else 'N/A'}")

        # Check if generation was successful
        if not output or (hasattr(output, "__len__") and len(output) == 0):
            print("[TOOL] Replicate generation failed - no output")
            return "Failed to generate image. Please try again."

        # Flux Kontext Pro returns a string URL
        generated_image_url = str(output)
        print(f"[TOOL] Generated image URL: {generated_image_url}")

    # Handle Flux Kontext Pro output format
    image_data: Optional[bytes] = None

    try:
        # Download the image from the URL
        import requests

        response = requests.get(generated_image_url)
        response.raise_for_status()
        image_data = response.content
        print(f"[TOOL] Downloaded image data, size: {len(image_data)} bytes")

    except Exception as e:
        print(f"[TOOL] Error processing output: {e}")
        return f"Failed to process generated image: {str(e)}"

    # Check if we successfully got image data
    if image_data is None:
        return "Failed to get image data from generation output"

    # Update or create a new generation count by + 1 for this ip address
    create_or_update_ip_generation_count(client_ip)

    # Generate unique ID for the image
    image_id = str(uuid.uuid4())

    # Upload to S3
    print(f"[TOOL] Uploading to S3 with image_id: {image_id}")
    print(f"[TOOL] Image data size: {len(image_data)} bytes")
    try:
        s3_result = upload_generated_image_to_s3(
            image_data=image_data,
            image_id=image_id,
            user_id=user_id,
            prompt=prompt,
            title=title,
        )
        print(f"[TOOL] S3 upload result: {s3_result}")
        print(f"[TOOL] S3 upload success: {s3_result.get('success', False)}")

        if s3_result["success"]:
            # Store structured result for the agent to retrieve
            tool_result = {"image_id": image_id, "title": title, "prompt": prompt, "success": True}
            print(f"[TOOL] About to store tool result: {tool_result}")
            store_tool_result(user_id, "generate_image", tool_result)
            print("[TOOL] Tool result stored successfully")

            result_msg = f"Image generated successfully! User can find it his/her gallery. \
                Image ID: {image_id}, Title: {title}"
            print(f"[TOOL] Returning success: {result_msg}")
            return result_msg
        else:
            error_msg = f"Image generated but failed to save: {s3_result.get('error', 'Unknown error')}"
            print(f"[TOOL] Returning error: {error_msg}")
            return error_msg

    except Exception as e:
        error_msg = f"Image generated but failed to save to storage: {str(e)}"
        print(f"[TOOL] Exception during S3 upload: {error_msg}")
        return error_msg

    finally:
        if image_data:
            # Clear image data from memory
            del image_data


def _generate_image_callable(inputs: Dict[str, str], config: RunnableConfig):
    # Normalize inputs whether dict or Pydantic
    if hasattr(inputs, "model_dump"):
        inputs = inputs.model_dump()
    elif hasattr(inputs, "dict"):
        inputs = inputs.dict()

    # Pull the IP from the per-invoke config
    cfg = config.get("configurable") or {}

    ip = cfg.get("client_ip")
    if not isinstance(ip, str) or not ip:
        # Fail fast if it's absent or not a string
        raise ValueError("client_ip is required in config.configurable and must be a non-empty string")

    client_ip: str = ip

    # Call your core with the IP
    return _generate_image_core(
        prompt=inputs["prompt"],
        user_id=inputs["user_id"],
        image_url=inputs["image_url"],
        title=inputs.get("title", "Generated Image"),
        client_ip=client_ip,
    )


def initialize_tools():
    """Initialize the tools for the agent."""
    print("[TOOLS] building generate_image tool")

    # A Runnable that receives (inputs, config) every invoke
    generate_image_runnable = RunnableLambda(_generate_image_callable)

    # As agent creation API expects "tools", convert the runnable to a Tool:
    generate_image_tool = generate_image_runnable.as_tool(
        name="generate_image",
        description=generate_image_tool_description,
        args_schema=GenerateImageToolInput,
    )

    return [generate_image_tool]


if __name__ == "__main__":
    # Test the tool
    generate_image = initialize_tools()[0]
    output = generate_image.invoke(
        {
            "prompt": "A woman in a beautiful sunset over a calm ocean",
            "user_id": "123",
            "image_url": "https://example.com/image.jpg",
            "title": "Test Image",
        },
        config={"configurable": {"client_ip": "127.0.0.1"}},
    )
    print(output)