app.py

import os
import json
import random
import logging
import warnings

import gradio as gr

from together import Together
from mistralai import Mistral
from dotenv import load_dotenv

logging.basicConfig(level=logging.INFO)
warnings.filterwarnings("ignore", category=UserWarning)

# Load the environment variables
load_dotenv()
TOGETHER_API_TOKEN = os.environ["TOGETHER_API_KEY"]
MISTRAL_API_TOKEN = os.environ["MISTRAL_API_KEY"]

demo = None # Added to allow restart

def start_game(main_loop, share=False):
    # Added code to support restart
    global demo
    
    # If demo is already running, close it first
    if demo is not None:
        demo.close()

    demo = gr.ChatInterface(
        main_loop,
        chatbot=gr.Chatbot(placeholder="Type 'Start' to begin", type='messages', height="75vh"),
        textbox=gr.Textbox(placeholder="What do you do next?", submit_btn="⬆", autofocus=True),
        title="Joe's AI Adventure Game ✨",
        description="This is a simple text-based adventure game powered by LLMs 🎮",
        theme="glass",
        examples=["What am I doing here?", "Tell me the background story", "Look around", "Continue the story"],
        cache_examples=True,
    )
    demo.launch(share=share)

def get_mistral_response(client=None, system_prompt="", user_prompt="",
                         model="mistral-small-latest", messages=[]):
    """
    Get response from Mistral model.
    """

    assert client is not None
    
    if len(messages) == 0:
        assert len(user_prompt) > 0, "User prompt cannot be empty!"
        
        chat_response = client.chat.complete(
            model=model,
            messages=[
                {
                    "role": "system",
                    "content": system_prompt,
                    "role": "user",
                    "content": user_prompt,
                },
            ]
        )
    else:
        chat_response = client.chat.complete(model=model, 
                                             messages=messages)
    
    return chat_response.choices[0].message.content

def run_action(client, message, history, game_state, use_mistral=True):
    """
    Run the action and get the next part of the story.
    """

    if(message.lower() == 'start'):
        # Include character image with the start message
        char_image_path = game_state.get('character_image', None)
        if char_image_path and os.path.exists(char_image_path):
            return {"text": game_state['start'], "files": [char_image_path]}
        return game_state['start']

    system_prompt = """You are an AI same 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 and present tense. \
    Example: You look north and see..."""

    world_info = f"""
    World: {game_state['world']}
    Kingdom: {game_state['kingdom']}
    Town: {game_state['town']}
    Your character:  {game_state['character']}"""

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": world_info}
    ]

    # Handle both message format (dict with role/content) and tuple format
    for action in history:
        if isinstance(action, dict):
            # New format: messages are dicts with 'role' and 'content'
            if 'content' in action:
                content = action['content']

                # Skip if content is a tuple (file paths from multimodal messages)
                if isinstance(content, tuple):
                    continue

                # Handle multimodal messages (text + files)
                if isinstance(content, dict) and 'text' in content:
                    messages.append({"role": action['role'], "content": content['text']})
                elif isinstance(content, str):
                    messages.append({"role": action['role'], "content": content})
        else:
            # Old format: messages are tuples [user_message, assistant_message]
            # Extract text if the message is a dict with files
            user_msg = action[0]
            assistant_msg = action[1]

            if isinstance(assistant_msg, dict) and 'text' in assistant_msg:
                assistant_msg = assistant_msg['text']

            messages.append({"role": "user", "content": user_msg})
            messages.append({"role": "assistant", "content": assistant_msg})

    messages.append({"role": "user", "content": message})
    
    if use_mistral:
        model_output = get_mistral_response(client=client, messages=messages)
        result = model_output
    else:
        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

def main_loop(client, message, history, game_state):
    """
    Run the game
    """
    
    return run_action(client, message, history, game_state)

def main():
    # Initialize the clients
    assert TOGETHER_API_TOKEN is not None, "Please set TOGETHER_API_TOKEN environment variable."
    assert MISTRAL_API_TOKEN is not None, "Please set MISTRAL_API_TOKEN environment variable."
    together_client = Together(api_key=TOGETHER_API_TOKEN)
    mistral_client = Mistral(api_key=MISTRAL_API_TOKEN)

    # Read the system prompt
    sys_prompt_file = "./assets/sys_prompt.txt"
    sys_prompt = None
    with open(sys_prompt_file, 'r', encoding='utf-8') as file:
        sys_prompt = file.read()
    assert sys_prompt is not None, "System prompt cannot be read!"

    # Load the world file
    world_file = "./assets/MyWorld.json"
    my_world = None
    with open(world_file, 'r') as file:
        my_world = json.load(file)
    assert my_world is not None, "Failed to load world data!"

    # Retrieve a random kingdom
    kingdoms = list(my_world["kingdoms"].values())
    selected_kingdom = random.choice(kingdoms)
    logging.info(f"Selected Kingdom: {selected_kingdom['name']}")

    # Retrieve a random town from the selected kingdom
    towns = list(selected_kingdom["towns"].values())
    selected_town = random.choice(towns)
    logging.info(f"Selected Town: {selected_town['name']}")

    # Retrieve a random character from the selected town
    characters = list(selected_town["npcs"].values())
    selected_character = random.choice(characters)
    char_name = selected_character['name']
    logging.info(f"Selected Character: {char_name}")

    # Construct character image path
    char_image_filename = char_name.replace(" ", "_") + ".png"
    char_image_path = os.path.join("./assets/images", char_image_filename)
    logging.info(f"Character Image Path: {char_image_path}")

    world_info = f"""
    World: {my_world}
    Kingdom: {selected_kingdom}
    Town: {selected_town}
    Your character: {selected_character}
    """

    # Start the game
    user_prompt = world_info + '\nYour start:'
    model_output = get_mistral_response(client=mistral_client,
                                        system_prompt=sys_prompt, user_prompt=user_prompt)
    start = model_output.strip()
    my_world['start'] = start

    # Define the game state
    game_state = {
        "world": my_world['description'],
        "kingdom": selected_kingdom['description'],
        "town": selected_town['description'],
        "character": selected_character['description'],
        "character_image": char_image_path,
        "start": start,
    }
    start_game(lambda msg, hist: main_loop(mistral_client, msg, hist, game_state))

if __name__ == "__main__":
    main()