Spaces:

Mistral-AI-Game-Jam
/

ai-interactive-comic-book

Running

ai-interactive-comic-book / server /scripts /test_game.py

tfrere HF Staff

update story

78b81a5 over 1 year ago

7.93 kB

import asyncio
import os
import sys
import time
import argparse
from pathlib import Path
from dotenv import load_dotenv
import uuid
import random

# Add server directory to PYTHONPATH
server_dir = Path(__file__).parent.parent
sys.path.append(str(server_dir))

from core.game_state import GameState
from core.story_generator import StoryGenerator
from core.constants import GameConfig
from core.generators.universe_generator import UniverseGenerator

# Load environment variables
load_dotenv()

def parse_args():
    parser = argparse.ArgumentParser(description="Test the game's story generation")
    parser.add_argument('--show-context', action='store_true', help='Show the full context at each step')
    parser.add_argument('--auto', action='store_true', help='Run in automatic mode with random choices')
    parser.add_argument('--max-turns', type=int, default=15, help='Maximum number of turns before considering test failed (default: 15)')
    return parser.parse_args()

def print_separator(char="=", length=50):
    print(f"\n{char * length}\n")

def print_universe_info(style: str, genre: str, epoch: str, base_story: str):
    print_separator("*")
    print("🌍 UNIVERSE GENERATED")
    print(f"🎨 Style: {style}")
    print(f"📚 Genre: {genre}")
    print(f"⏳ Époque: {epoch}")
    print("\n📖 Base Story:")
    print(base_story)
    print_separator("*")

def print_story_step(step_number, story_text, image_prompts, generation_time: float, story_history: str = None, show_context: bool = False, model_name: str = None, is_death: bool = False, is_victory: bool = False):
    """Print a story step with formatting."""
    print("\n" + "="*80)
    print(f"Step {step_number}")
    print("-"*80)
    
    if show_context and story_history:
        print("\nContext:")
        print("-"*20)
        print(story_history)
        print("-"*20 + "\n")
    
    print(f"Story ({generation_time:.2f}s):")
    print(story_text)
    
    if image_prompts:
        print("\nImage Prompts:")
        for i, prompt in enumerate(image_prompts, 1):
            print(f"{i}. {prompt}")
    
    if is_death:
        print("\n💀 GAME OVER - You died!")
    elif is_victory:
        print("\n🏆 VICTORY - You won!")
    
    print("="*80)

async def play_game(show_context: bool = False, auto_mode: bool = False, max_turns: int = 15):
    # Initialize components
    model_name = "mistral-small"
    story_generator = StoryGenerator(
        api_key=os.getenv("MISTRAL_API_KEY"),
        model_name=model_name
    )
    
    # Create universe generator
    universe_generator = UniverseGenerator(story_generator.mistral_client)
    
    print("\n=== Don't Look Up - Test Mode ===\n")
    print("🎮 Starting adventure...")
    if show_context:
        print("📚 Context display is enabled")
    if auto_mode:
        print("🤖 Running in automatic mode")
        print(f"⏱️  Maximum turns: {max_turns}")
    print_separator()
    
    # Test universe generation
    style, genre, epoch, macguffin, hero = universe_generator._get_random_elements()
    print(f"\nGenerated universe elements:")
    print(f"Style: {style['name']}")
    print(f"Genre: {genre}")
    print(f"Epoch: {epoch}")
    print(f"MacGuffin: {macguffin}")
    print(f"Hero: {hero}")

    base_story = await universe_generator.generate()
    print(f"\nGenerated base story:\n{base_story}")
    
    # Create session and game state
    session_id = str(uuid.uuid4())
    game_state = GameState()
    game_state.set_universe(
        style=style["name"],
        genre=genre,
        epoch=epoch,
        base_story=base_story
    )
    
    # Create story generator
    story_generator.create_segment_generator(
        session_id=session_id,
        style=style,
        genre=genre,
        epoch=epoch,
        base_story=base_story,
        macguffin=macguffin,
        hero=hero
    )
    
    # Display universe information
    print_universe_info(style["name"], genre, epoch, base_story)
    
    last_choice = None
    
    while True:
        # Check for maximum turns in auto mode
        if auto_mode and game_state.story_beat >= max_turns:
            print(f"\n❌ TEST FAILED: Story did not end after {max_turns} turns")
            return False
            
        # Generate story segment
        previous_choice = "none" if game_state.story_beat == 0 else f"Choice {last_choice}"
        
        # Format story history for context
        story_history = ""
        if game_state.story_history:
            segments = []
            for story_response in game_state.story_history:
                segment = story_response.story_text
                time_location = f"[{story_response.time} - {story_response.location}]"
                image_descriptions = "\nVisual panels:\n" + "\n".join(f"- {prompt}" for prompt in story_response.image_prompts)
                segments.append(f"{time_location}\n{segment}{image_descriptions}")
            
            story_history = "\n\n---\n\n".join(segments)
            story_history += "\n\nLast choice made: " + previous_choice
        
        # Mesurer le temps de génération
        start_time = time.time()
        response = await story_generator.generate_story_segment(
            session_id=session_id,
            game_state=game_state,
            previous_choice=previous_choice
        )
        generation_time = time.time() - start_time
        
        # Display current step
        print_story_step(
            game_state.story_beat,
            response.story_text,
            response.image_prompts,
            generation_time,
            story_history,
            show_context,
            model_name,
            response.is_death,
            response.is_victory
        )
        
        if response.is_death:
            print("\n☢️  GAME OVER - Death ☢️")
            print("Sarah has succumbed...")
            return False
            
        # Check for victory
        if response.is_victory:
            print("\n🏆 VICTORY! 🏆")
            print("Sarah has survived and completed her mission!")
            return True
            
        # Display choices
        if len(response.choices) == 2:  # On vérifie qu'on a exactement 2 choix
            print("\n🤔 AVAILABLE CHOICES:")
            for choice in response.choices:
                print(f"{choice.id}. {choice.text}")
                
            # Get player choice
            if auto_mode:
                last_choice = random.randint(1, 2)
                print(f"\n🤖 Auto-choosing: {last_choice}")
                time.sleep(1)  # Small delay for readability
            else:
                while True:
                    try:
                        choice_input = int(input("\n👉 Your choice (1-2): "))
                        if choice_input in [1, 2]:
                            last_choice = choice_input
                            break
                        print("❌ Invalid choice. Please choose 1 or 2.")
                    except ValueError:
                        print("❌ Please enter a number.")
            
            # Update game state
            game_state.story_beat += 1
            # Le StoryResponse est déjà ajouté à l'historique dans generate_story_segment
            
        else:
            print("\n❌ Error: Invalid number of choices received from server")
            return False

def main():
    try:
        args = parse_args()
        success = asyncio.run(play_game(
            show_context=args.show_context,
            auto_mode=args.auto,
            max_turns=args.max_turns
        ))
        if args.auto:
            sys.exit(0 if success else 1)
    except KeyboardInterrupt:
        print("\n\n👋 Game interrupted. See you soon!")
        sys.exit(1)
    except Exception as e:
        print(f"\n❌ An error occurred: {str(e)}")
        sys.exit(1)

if __name__ == "__main__":
    main() 

	import asyncio
	import os
	import sys
	import time
	import argparse
	from pathlib import Path
	from dotenv import load_dotenv
	import uuid
	import random

	# Add server directory to PYTHONPATH
	server_dir = Path(__file__).parent.parent
	sys.path.append(str(server_dir))

	from core.game_state import GameState
	from core.story_generator import StoryGenerator
	from core.constants import GameConfig
	from core.generators.universe_generator import UniverseGenerator

	# Load environment variables
	load_dotenv()

	def parse_args():
	parser = argparse.ArgumentParser(description="Test the game's story generation")
	parser.add_argument('--show-context', action='store_true', help='Show the full context at each step')
	parser.add_argument('--auto', action='store_true', help='Run in automatic mode with random choices')
	parser.add_argument('--max-turns', type=int, default=15, help='Maximum number of turns before considering test failed (default: 15)')
	return parser.parse_args()

	def print_separator(char="=", length=50):
	print(f"\n{char * length}\n")

	def print_universe_info(style: str, genre: str, epoch: str, base_story: str):
	print_separator("*")
	print("🌍 UNIVERSE GENERATED")
	print(f"🎨 Style: {style}")
	print(f"📚 Genre: {genre}")
	print(f"⏳ Époque: {epoch}")
	print("\n📖 Base Story:")
	print(base_story)
	print_separator("*")

	def print_story_step(step_number, story_text, image_prompts, generation_time: float, story_history: str = None, show_context: bool = False, model_name: str = None, is_death: bool = False, is_victory: bool = False):
	"""Print a story step with formatting."""
	print("\n" + "="*80)
	print(f"Step {step_number}")
	print("-"*80)

	if show_context and story_history:
	print("\nContext:")
	print("-"*20)
	print(story_history)
	print("-"*20 + "\n")

	print(f"Story ({generation_time:.2f}s):")
	print(story_text)

	if image_prompts:
	print("\nImage Prompts:")
	for i, prompt in enumerate(image_prompts, 1):
	print(f"{i}. {prompt}")

	if is_death:
	print("\n💀 GAME OVER - You died!")
	elif is_victory:
	print("\n🏆 VICTORY - You won!")

	print("="*80)

	async def play_game(show_context: bool = False, auto_mode: bool = False, max_turns: int = 15):
	# Initialize components
	model_name = "mistral-small"
	story_generator = StoryGenerator(
	api_key=os.getenv("MISTRAL_API_KEY"),
	model_name=model_name
	)

	# Create universe generator
	universe_generator = UniverseGenerator(story_generator.mistral_client)

	print("\n=== Don't Look Up - Test Mode ===\n")
	print("🎮 Starting adventure...")
	if show_context:
	print("📚 Context display is enabled")
	if auto_mode:
	print("🤖 Running in automatic mode")
	print(f"⏱️ Maximum turns: {max_turns}")
	print_separator()

	# Test universe generation
	style, genre, epoch, macguffin, hero = universe_generator._get_random_elements()
	print(f"\nGenerated universe elements:")
	print(f"Style: {style['name']}")
	print(f"Genre: {genre}")
	print(f"Epoch: {epoch}")
	print(f"MacGuffin: {macguffin}")
	print(f"Hero: {hero}")

	base_story = await universe_generator.generate()
	print(f"\nGenerated base story:\n{base_story}")

	# Create session and game state
	session_id = str(uuid.uuid4())
	game_state = GameState()
	game_state.set_universe(
	style=style["name"],
	genre=genre,
	epoch=epoch,
	base_story=base_story
	)

	# Create story generator
	story_generator.create_segment_generator(
	session_id=session_id,
	style=style,
	genre=genre,
	epoch=epoch,
	base_story=base_story,
	macguffin=macguffin,
	hero=hero
	)

	# Display universe information
	print_universe_info(style["name"], genre, epoch, base_story)

	last_choice = None

	while True:
	# Check for maximum turns in auto mode
	if auto_mode and game_state.story_beat >= max_turns:
	print(f"\n❌ TEST FAILED: Story did not end after {max_turns} turns")
	return False

	# Generate story segment
	previous_choice = "none" if game_state.story_beat == 0 else f"Choice {last_choice}"

	# Format story history for context
	story_history = ""
	if game_state.story_history:
	segments = []
	for story_response in game_state.story_history:
	segment = story_response.story_text
	time_location = f"[{story_response.time} - {story_response.location}]"
	image_descriptions = "\nVisual panels:\n" + "\n".join(f"- {prompt}" for prompt in story_response.image_prompts)
	segments.append(f"{time_location}\n{segment}{image_descriptions}")

	story_history = "\n\n---\n\n".join(segments)
	story_history += "\n\nLast choice made: " + previous_choice

	# Mesurer le temps de génération
	start_time = time.time()
	response = await story_generator.generate_story_segment(
	session_id=session_id,
	game_state=game_state,
	previous_choice=previous_choice
	)
	generation_time = time.time() - start_time

	# Display current step
	print_story_step(
	game_state.story_beat,
	response.story_text,
	response.image_prompts,
	generation_time,
	story_history,
	show_context,
	model_name,
	response.is_death,
	response.is_victory
	)

	if response.is_death:
	print("\n☢️ GAME OVER - Death ☢️")
	print("Sarah has succumbed...")
	return False

	# Check for victory
	if response.is_victory:
	print("\n🏆 VICTORY! 🏆")
	print("Sarah has survived and completed her mission!")
	return True

	# Display choices
	if len(response.choices) == 2: # On vérifie qu'on a exactement 2 choix
	print("\n🤔 AVAILABLE CHOICES:")
	for choice in response.choices:
	print(f"{choice.id}. {choice.text}")

	# Get player choice
	if auto_mode:
	last_choice = random.randint(1, 2)
	print(f"\n🤖 Auto-choosing: {last_choice}")
	time.sleep(1) # Small delay for readability
	else:
	while True:
	try:
	choice_input = int(input("\n👉 Your choice (1-2): "))
	if choice_input in [1, 2]:
	last_choice = choice_input
	break
	print("❌ Invalid choice. Please choose 1 or 2.")
	except ValueError:
	print("❌ Please enter a number.")

	# Update game state
	game_state.story_beat += 1
	# Le StoryResponse est déjà ajouté à l'historique dans generate_story_segment

	else:
	print("\n❌ Error: Invalid number of choices received from server")
	return False

	def main():
	try:
	args = parse_args()
	success = asyncio.run(play_game(
	show_context=args.show_context,
	auto_mode=args.auto,
	max_turns=args.max_turns
	))
	if args.auto:
	sys.exit(0 if success else 1)
	except KeyboardInterrupt:
	print("\n\n👋 Game interrupted. See you soon!")
	sys.exit(1)
	except Exception as e:
	print(f"\n❌ An error occurred: {str(e)}")
	sys.exit(1)

	if __name__ == "__main__":
	main()