app.py

import os
import json
import gradio as gr
from datetime import datetime

from together import Together
from helper import save_world


client = Together()  # Automatically uses TOGETHER_API_KEY from environment

system_prompt = "You are a creative fantasy world generator. Respond only with structured text in the required format."

def generate_world():
    # Create World
    world_prompt = """
    Generate a creative description for a unique fantasy world with an
    interesting concept around cities built on the backs of massive beasts.

    Output content in the form:
    World Name: <WORLD NAME>
    World Description: <WORLD DESCRIPTION>

    World Name:"""

    output = client.chat.completions.create(
        model="meta-llama/Llama-3-70b-chat-hf",
        temperature=0.0,  # 0.0 = deterministic, 1.0 = creative
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": world_prompt}
        ],
        stream=False,                          # REQUIRED: Together expects stream to be explicitly declared
        # tools=[],                            # REQUIRED: If not using any tools
        # tool_choice={"type": "auto"}         # REQUIRED: If not choosing any tool 
    )

    world_output = output.choices[0].message.content.strip()
    world = {
        "name": world_output.split('\n')[0].replace('World Name: ', '').strip(),
        "description": '\n'.join(world_output.split('\n')[1:]).replace('World Description:', '').strip()
    }

    # Create Kingdoms
    kingdom_prompt = f"""
    Create 3 different kingdoms for a fantasy world.
    For each kingdom generate a description based on the world it's in.
    Describe important leaders, cultures, history of the kingdom.

    Output content in the form:
    Kingdom 1 Name: <KINGDOM NAME>
    Kingdom 1 Description: <KINGDOM DESCRIPTION>
    Kingdom 2 Name: <KINGDOM NAME>
    Kingdom 2 Description: <KINGDOM DESCRIPTION>
    Kingdom 3 Name: <KINGDOM NAME>
    Kingdom 3 Description: <KINGDOM DESCRIPTION>

    World Name: {world['name']}
    World Description: {world['description']}

    Kingdom 1 Name:"""

    kingdoms_schema = {
        "type": "object",
        "properties": {
            "kingdoms": {
                "type": "array",
                "items": {
                    "type": "object",
                    "properties": {
                        "name": {"type": "string"},
                        "description": {"type": "string"}
                    },
                    "required": ["name", "description"]
                },
                "minItems": 3,
                "maxItems": 3
            }
        },
        "required": ["kingdoms"]
}    

    output = client.chat.completions.create(
        model="meta-llama/Llama-3-70b-chat-hf",
        temperature=1.0,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": kingdom_prompt}
        ],
        stream=False,                          # REQUIRED: Together expects stream to be explicitly declared
        response_format={
            "type": "json_schema",
            "schema": kingdoms_schema
        }
    )

    # Parse the JSON string first
    kingdoms_output_str = output.choices[0].message.content.strip()
    
    # Parse the JSON string into a Python dict
    try:
        kingdoms_output = json.loads(kingdoms_output_str)
    except json.JSONDecodeError as e:
        print(f"[ERROR] Failed to parse JSON: {e}")
        return None, None
    
    # Now access the parsed content
    kingdoms = {}
    for k in kingdoms_output["kingdoms"]:
        kingdoms[k["name"]] = {
            "name": k["name"],
            "description": k["description"],
            "world": world["name"]
        }
    
    world["kingdoms"] = kingdoms
    
    # Create Towns
    def get_town_prompt(world, kingdom):
        return f"""
        Create 3 different towns for a fantasy kingdom and world.
        Describe the region it's in, important places of the town,
        and interesting history about it.

        Output content in the form:
        Town 1 Name: <TOWN NAME>
        Town 1 Description: <TOWN DESCRIPTION>
        Town 2 Name: <TOWN NAME>
        Town 2 Description: <TOWN DESCRIPTION>
        Town 3 Name: <TOWN NAME>
        Town 3 Description: <TOWN DESCRIPTION>

        World Name: {world['name']}
        World Description: {world['description']}

        Kingdom Name: {kingdom['name']}
        Kingdom Description: {kingdom['description']}

        Town 1 Name:"""

    def create_towns(world, kingdom):
        output = client.chat.completions.create(
            model="meta-llama/Llama-3-70b-chat-hf",
            temperature=0.0,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": get_town_prompt(world, kingdom)}
            ],
            stream=False,                         # REQUIRED: Together expects stream to be explicitly declared
        )
        towns_output = output.choices[0].message.content.strip()
        towns = {}
        for block in towns_output.split('\n\n'):
            lines = block.strip().split('\n')
            town_name = lines[0].split('Name: ')[1].strip()
            town_description = lines[1].split('Description: ')[1].strip()
            towns[town_name] = {
                "name": town_name,
                "description": town_description,
                "kingdom": kingdom["name"],
                "world": world["name"]
            }
        kingdom["towns"] = towns

    for kingdom in kingdoms.values():
        create_towns(world, kingdom)

    # Create NPCs for one town
    def get_npc_prompt(world, kingdom, town):
        return f"""
        Create 3 different characters based on the world, kingdom
        and town they're in. Describe the character's appearance and
        profession, as well as their deeper pains and desires.

        Output content in the form:
        Character 1 Name: <CHARACTER NAME>
        Character 1 Description: <CHARACTER DESCRIPTION>
        Character 2 Name: <CHARACTER NAME>
        Character 2 Description: <CHARACTER DESCRIPTION>
        Character 3 Name: <CHARACTER NAME>
        Character 3 Description: <CHARACTER DESCRIPTION>

        World Name: {world['name']}
        World Description: {world['description']}

        Kingdom Name: {kingdom['name']}
        Kingdom Description: {kingdom['description']}

        Town Name: {town['name']}
        Town Description: {town['description']}

        Character 1 Name:"""

    def create_npcs(world, kingdom, town):
        output = client.chat.completions.create(
            model="meta-llama/Llama-3-70b-chat-hf",
            temperature=0.0,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": get_npc_prompt(world, kingdom, town)}
            ],
            stream=False,                         # REQUIRED: Together expects stream to be explicitly declared
        )
        npcs_output = output.choices[0].message.content.strip()
        npcs = {}
        for block in npcs_output.split('\n\n'):
            lines = block.strip().split('\n')
            npc_name = lines[0].split('Name: ')[1].strip()
            npc_description = lines[1].split('Description: ')[1].strip()
            npcs[npc_name] = {
                "name": npc_name,
                "description": npc_description,
                "town": town["name"],
                "kingdom": kingdom["name"],
                "world": world["name"]
            }
        town["npcs"] = npcs

    sample_kingdom = next(iter(world["kingdoms"].values()))
    sample_town = next(iter(sample_kingdom["towns"].values()))
    create_npcs(world, sample_kingdom, sample_town)

    # Save and return file path
    filename = f"/tmp/world_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
    save_world(world, filename)

    return json.dumps(world, indent=2), filename

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🐉 Fantasy World Generator")
    gr.Markdown("Click below to generate a fantasy world, its kingdoms, towns, and characters.")

    with gr.Row():
        generate_btn = gr.Button("✨ Generate World")
        json_file = gr.File(label="Download World", interactive=False)

    output_box = gr.Textbox(label="Generated World Summary (JSON)", lines=30)

    generate_btn.click(fn=generate_world, outputs=[output_box, json_file])

demo.launch()