app.py

import os
import json
import glob
import gradio as gr
from datetime import datetime
from together import Together
from helper import load_world, save_world

# Global variable to allow demo restart
demo = None

client = Together()

# Moderate Content Policy
# Set the content policy
mature_content_policy = {
    "policy": """O2: Sexual Content.
Should not
- Write about graphic or illegal content.
Can
- Write about mature content between adults"""
}

# Safe Content Policy
safe_content_policy = {
    "policy": """O1: Violence and Hate.
Should not
- Express hateful, derogatory, or demeaning sentiments against people based on sensitive personal characteristics.
- Use slurs based on personal characteristics
- Write gruesome / graphic violence.
Can
- Write tame fantasy violence
O2: Sexual Content.
Should not
- Write any somewhat sexual content that mentions sexual functions or sexual body parts.
- Write any content that would be inappropriate for kids.
Can
- Write romantic content (including kissing, flirting etc...)
O3: Self-Harm.
Should not
- Encourage people to harm themselves.
- Romanticize or condone self-harm or suicide.
- Write story depicting suicide.
O4: Profanity.
Should not
- Write with any profane language that wouldn't be appropriate for kids.
Can
- Write with kid appropriate language
O5: Regulated or Controlled Substances.
Should not
- Write stories that depicts, glorifies or encourages drug use."""
}

everyone_content_policy = {
    "policy": """O1: Violence and Hate.
Should not
- Express hateful, derogatory, or demeaning sentiments against people based on sensitive personal characteristics.
- Use slurs based on personal characteristics
- Write gruesome / graphic violence.
Can
- Write tame fantasy violence
O2: Sexual Content.
Should not
- Write any somewhat sexual content that mentions sexual functions or sexual body parts.
- Write any content that would be inappropriate for kids.
Can
- Write romantic content (including kissing, flirting etc...)
O3: Self-Harm.
Should not
- Encourage people to harm themselves.
- Romanticize or condone self-harm or suicide.
- Write story depicting suicide.
O4: Profanity.
Should not
- Write with any profane language that wouldn't be appropriate for kids.
Can
- Write with kid appropriate language
O5: Regulated or Controlled Substances.
Should not
- Write stories that depicts, glorifies or encourages drug use."""
}

content_policies = {
    "Safe": safe_content_policy["policy"],
    "Everyone": everyone_content_policy["policy"],
    "Mature": mature_content_policy["policy"]
}

# Define Inventory Detector
inventory_detector_system_prompt = """You are an AI Game Assistant. \
Your job is to detect changes to a player's \
inventory based on the most recent story and game state.
If a player picks up, or gains an item add it to the inventory \
with a positive change_amount.
If a player loses an item remove it from their inventory \
with a negative change_amount.
Given a player name, inventory and story, return a list of json update
of the player's inventory in the following form.
Only take items that it's clear the player (you) lost.
Only give items that it's clear the player gained. 
Don't make any other item updates.
If no items were changed return {"itemUpdates": []}
and nothing else.

Response must be in Valid JSON
Don't add items that were already added in the inventory

Inventory Updates:
{
    "itemUpdates": [
        {"name": <ITEM NAME>, 
        "change_amount": <CHANGE AMOUNT>}...
    ]
}
"""

# World generation logic
system_prompt_world = "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,
        messages=[
            {"role": "system", "content": system_prompt_world},
            {"role": "user", "content": world_prompt}
        ],
        stream=False
    )

    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_world},
            {"role": "user", "content": kingdom_prompt}
        ],
        stream=False,
        response_format={"type": "json_schema", "schema": kingdoms_schema}
    )

    kingdoms_output = json.loads(output.choices[0].message.content.strip())
    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_world},
                {"role": "user", "content": get_town_prompt(world, kingdom)}
            ],
            stream=False
        )
        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
    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_world},
                {"role": "user", "content": get_npc_prompt(world, kingdom, town)}
            ],
            stream=False
        )
        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 world
    filename = f"/tmp/world_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
    save_world(world, filename)
    return world, filename

# Helper function to find the latest world file
def get_latest_world_file():
    world_files = glob.glob("/tmp/world_*.json")
    if not world_files:
        return None
    return max(world_files, key=os.path.getctime)

# Initialize game state
def initialize_game_state():
    # Check for existing world file
    world_file = get_latest_world_file()
    if not world_file:
        # Generate a new world if none exists
        world, world_file = generate_world()
    else:
        world = load_world(world_file)

    # Select kingdom, town, and character
    kingdom = next(iter(world['kingdoms'].values()))
    town = next(iter(kingdom['towns'].values()))
    character = next(iter(town['npcs'].values()))

    # Generate initial story
    system_prompt = """You are an AI Game master. Your job is to create a 
start to an adventure based on the world, kingdom, town and character 
a player is playing as. 
Instructions:
You must only use 2-4 sentences
Write in second person. For example: "You are Jack"
Write in present tense. For example "You stand at..."
First describe the character and their backstory.
Then describe where they start and what they see around them."""
    world_info = f"""
World: {world['name']}
Kingdom: {kingdom['name']}
Town: {town['name']}
Your Character: {character['name']}
"""
    model_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": world_info + '\nYour Start:'}
        ],
    )
    start = model_output.choices[0].message.content
    world['start'] = start
    save_world(world, world_file)

    # Return game state
    return {
        "world": world['description'],
        "kingdom": kingdom['description'],
        "town": town['description'],
        "character": character['description'],
        "start": start,
        "inventory": {
            "cloth pants": 1,
            "cloth shirt": 1,
            "goggles": 1,
            "leather bound journal": 1,
            "gold": 5
        }
    }

def is_safe(message, policy_name="Safe"):
    """Check if message is safe according to the selected policy"""
    policy = content_policies[policy_name]

    prompt = f"""[INST] Task: Check if there is unsafe content in 
'user' messages in conversations according to our safety policy 
with the below categories.

<BEGIN UNSAFE CONTENT CATEGORIES>
{policy}
<END UNSAFE CONTENT CATEGORIES>

<BEGIN CONVERSATION>
user: {message}

Provide your safety assessment for user in the above conversation:
- First line must read 'safe' or 'unsafe'.
- If unsafe, a second line must include a comma-separated list of 
violated categories. [/INST]"""

    # Build the prompt with embedded values
    response = client.completions.create(
        model="Meta-Llama/LlamaGuard-2-8b",
        prompt=prompt,
    )
    result = response.choices[0].text
    return result.strip().lower().startswith("safe")

def detect_inventory_changes(game_state, output):
    messages = [
        {"role": "system", "content": inventory_detector_system_prompt},
        {"role": "user", "content": f"Current Inventory: {json.dumps(game_state['inventory'])}"},
        {"role": "user", "content": f"Recent Story: {output}"},
        {"role": "user", "content": "Inventory Updates"}
    ]

    chat_completion = client.chat.completions.create(
        # response_format={"type": "json_object", "schema": InventoryUpdate.model_json_schema()},
        model="meta-llama/Llama-3-70b-chat-hf",
        temperature=0.0,
        messages=messages
    )
    response = chat_completion.choices[0].message.content
    return json.loads(response)['itemUpdates'] 

# Update inventory
def update_inventory(inventory, item_updates):
    update_msg = ''
    
    for update in item_updates:
        name = update['name']
        change_amount = update['change_amount']
        
        if change_amount > 0:
            if name not in inventory:
                inventory[name] = change_amount
            else:
                inventory[name] += change_amount
            update_msg += f'\nInventory: {name} +{change_amount}'
        elif name in inventory and change_amount < 0:
            inventory[name] += change_amount
            update_msg += f'\nInventory: {name} {change_amount}'
            
        if name in inventory and inventory[name] < 0:
            del inventory[name]
            
    return update_msg

# Main action loop (include inventory in the story)
def run_action(message, history, game_state):
    
    if(message == 'start game'):
        return game_state['start']

    system_prompt = """You are an AI Game master. Your job is to write what \
happens next in a player's adventure game.\
Instructions: \
You must on only write 1-3 sentences in response. \
Always write in second person present tense. \
Ex. (You look north and see...) \
Don't let the player use items they don't have in their inventory.
"""
    world_info = f"""
World: {game_state['world']}
Kingdom: {game_state['kingdom']}
Town: {game_state['town']}
Your Character: {game_state['character']}
Inventory: {json.dumps(game_state['inventory'])}"""
        
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": world_info}
    ]
    # Process history as a list of message dictionaries
    for msg in history:
        messages.append({"role": msg["role"], "content": msg["content"]})

    messages.append({"role": "user", "content": message})

    model_output = client.chat.completions.create(
        model="meta-llama/Llama-3-70b-chat-hf",
        messages=messages
    )
    
    result = model_output.choices[0].message.content
    return result

# Gradio main loop
def main_loop(message, history, policy_name="Safe"):
    """Main game loop that processes player actions"""
    game_state = initialize_game_state()

    # Check if input message is safe
    if not is_safe(message, policy_name):
        return 'Invalid action.'

    result = run_action(message, history, game_state)

    # Check if output is safe
    if not is_safe(result, policy_name):
        return 'Invalid output.'

    # Detect and update inventory changes
    try:
        item_updates = detect_inventory_changes(game_state, result)
        update_msg = update_inventory(game_state['inventory'], item_updates)
        result += update_msg
    except Exception as e:
        # If inventory detection fails, continue without inventory updates
        print(f"Inventory detection failed: {e}")
        
    return result

def gradio_game_interface(message, history, policy_name):
    return main_loop(message, history, policy_name)

# Start the Gradio UI Game
def start_game(main_loop_func, share=False):
    """Create and launch the Gradio interface for the RPG game"""
    # Added code to support restart
    global demo
    # If demo is already running, close it first
    if demo is not None:
        demo.close()

    with gr.Blocks() as demo:
        gr.Markdown("## 🐉 AI RPG Game")
        gr.Markdown("Generate a fantasy world or use an existing one, then start your adventure!")
        
        with gr.Row():
            generate_btn = gr.Button("✨ Generate New World")
            json_file = gr.File(label="Download World", interactive=False)
        
        output_box = gr.Textbox(label="Generated World Summary (JSON)", lines=10, visible=False)

        policy_dropdown = gr.Dropdown(
            choices=["Safe", "Everyone", "Mature"],
            value="Safe",
            label="Select Safety Policy"
        )
            
        chatbot=gr.Chatbot(height=250, placeholder="Type 'start game' to begin", type="messages") # OpenAI-style format
        textbox = gr.Textbox(placeholder="What do you do next?", container=False, scale=7)
        state = gr.State([])

        def wrapped_main(message, history):
            return main_loop_func(message, history, policy_dropdown.value)

        def generate_and_display_world():
            world, filename = generate_world()
            return json.dumps(world, indent=2), filename, gr.update(visible=True)

        chat_interface = gr.ChatInterface(
            wrapped_main,
            chatbot=chatbot,
            textbox=textbox,
            title="AI RPG",
            theme="soft",
            examples=["Look around", "Continue the story"],
            cache_examples=False
        )

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

        policy_dropdown.change(fn=lambda x: None, inputs=policy_dropdown, outputs=[])        

    demo.launch(share=share, server_name="0.0.0.0")


# Launch the app if run directly
if __name__ == "__main__":
    # Launch the game
    start_game(main_loop)