import os import asyncio import textwrap import re import httpx import gradio as gr from PIL import Image from io import BytesIO from llama_index.core.workflow import ( Workflow, step, StartEvent, StopEvent, Context, Event, ) from llama_index.llms.groq import Groq # --- Secrets --- GROQ_API_KEY = os.environ.get("Groq_Token") HF_TOKEN = os.environ.get("HF_TOKEN2") # --- Event and Workflow Definitions --- class StoryContext(Event): story_part: str inventory: list[str] class SceneReadyEvent(Event): pass class UserChoice(Event): choice: str class StoryEnd(Event): final_message: str # Helper function to generate an image and return its path async def generate_image(prompt: str, hf_token: str) -> str | None: API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell" headers = {"Authorization": f"Bearer {hf_token}"} full_prompt = f"epic fantasy art, digital painting, cinematic lighting, masterpiece, {prompt}" payload = {"inputs": full_prompt} try: async with httpx.AsyncClient() as client: response = await client.post(API_URL, headers=headers, json=payload, timeout=180.0) response.raise_for_status() image = Image.open(BytesIO(response.content)) # Save to a temporary file that Gradio can serve image.save("scene_image.png") return "scene_image.png" except (httpx.TimeoutException, httpx.RequestError, IOError) as e: print(f"Image generation failed: {e}") return None # The full workflow class class StorytellerWorkflow(Workflow): def __init__(self, **kwargs): super().__init__(timeout=300, **kwargs) @step async def generate_story_part(self, ev: StartEvent | UserChoice, ctx: Context) -> StoryContext | StoryEnd: inventory = await ctx.store.get("inventory", []) prompt = "" if isinstance(ev, StartEvent): # The prompt for the first turn prompt = """ You are a creative text adventure game master. Your output is for a console game. Start a new story about a curious explorer entering a recently discovered, glowing cave. Keep the tone mysterious and exciting. After the story part, provide two distinct choices for the player to make. Format your response exactly like this: STORY: [The story text goes here] CHOICES: 1. [First choice] 2. [Second choice] """ elif isinstance(ev, UserChoice): last_story_part = await ctx.store.get("last_story_part") # --- EXPLICIT PROMPT --- prompt = f""" You are a creative text adventure game master. The story so far: "{last_story_part}" The player chose: "{ev.choice}" The player's inventory: {inventory} Continue the story. If a choice results in an item, use `[ADD_ITEM: item name]`. If the story should end, write "[END]". Format your response exactly like this: STORY: [The story text goes here] CHOICES: 1. [First choice on its own line] 2. [Second choice on its own line] """ llm = Groq(model="llama3-8b-8192", api_key=GROQ_API_KEY) response = await llm.acomplete(prompt) response_text = str(response) items_found = re.findall(r"\[ADD_ITEM: (.*?)\]", response_text) if items_found: for item in items_found: if item not in inventory: inventory.append(item) response_text = re.sub(r"\[ADD_ITEM: (.*?)\]", "", response_text).strip() if response_text.strip().startswith("[END]"): final_message = response_text.strip().replace("[END]", "") return StoryEnd(final_message=f"\n--- THE END ---\n{final_message}") try: story_section = response_text.split("STORY:")[1].split("CHOICES:")[0].strip() choices_section = response_text.split("CHOICES:")[1].strip() full_story_part = f"{story_section}\n\nChoices:\n{choices_section}" except IndexError: full_story_part = "The story continues... but the path is blurry." await ctx.store.set("last_story_part", full_story_part) await ctx.store.set("inventory", inventory) # Return the simplified event, without the is_new_scene flag return StoryContext(story_part=full_story_part, inventory=inventory) @step def end_story(self, ev: StoryEnd) -> StopEvent: """This step satisfies the workflow validator by providing a path to a StopEvent.""" return StopEvent(result=ev.final_message) # These two steps are no longer needed, as Gradio's UI will handle the display logic. # @step async def display_scene(...) # @step async def get_user_choice(...) # --- Gradio UI and Application Logic --- async def run_turn(user_input, game_state): # This function no longer needs the component objects passed in. if game_state is None: game_state = {'inventory': [], 'last_story_part': None} event = StartEvent() else: event = UserChoice(choice=user_input) workflow = StorytellerWorkflow() ctx = Context(workflow=workflow) await ctx.store.set("inventory", game_state['inventory']) await ctx.store.set("last_story_part", game_state['last_story_part']) result_event = await workflow.generate_story_part(event, ctx) if isinstance(result_event, StoryEnd): yield (None, result_event.final_message, "", None) return if isinstance(result_event, StoryContext): narrative, choices_text = result_event.story_part.split("Choices:", 1) # --- THIS IS THE NEW DEFENSIVE CODE --- # It finds any choice number (like " 2." or " 3.") that has a space before it # and replaces that space with a newline character. This forces them to be on separate lines. choices_text = re.sub(r" (\d\.)", r"\n\1", choices_text) story_display_text = f"{textwrap.fill(narrative, width=80)}\n\nChoices:{choices_text}" new_game_state = { 'inventory': result_event.inventory, 'last_story_part': result_event.story_part } inventory_text = f"**Inventory:** {', '.join(new_game_state['inventory']) if new_game_state['inventory'] else 'Empty'}" yield (None, story_display_text, inventory_text, new_game_state) # 2. Generate the image. image_path = None if HF_TOKEN: image_path = await generate_image(narrative, HF_TOKEN) # 3. Yield a second, complete tuple with the new image path. yield (image_path, story_display_text, inventory_text, new_game_state) def create_demo(): with gr.Blocks(theme=gr.themes.Soft()) as demo: game_state = gr.State(None) gr.Markdown("# LlamaIndex Workflow: Dynamic Storyteller") gr.Markdown("An AI-powered text adventure game where every scene can be illustrated by AI.") with gr.Row(): with gr.Column(scale=1): image_display = gr.Image(label="Scene", interactive=False) inventory_display = gr.Markdown("**Inventory:** Empty") with gr.Column(scale=2): story_display = gr.Textbox(label="Story", lines=15, interactive=False) user_input = gr.Textbox(label="What do you do?", placeholder="Type your choice and press Enter...") # The inputs and outputs lists should be simple. inputs = [user_input, game_state] outputs = [image_display, story_display, inventory_display, game_state] user_input.submit( fn=run_turn, inputs=inputs, outputs=outputs ) load_inputs = [gr.State(None), game_state] demo.load( fn=run_turn, inputs=load_inputs, outputs=outputs ) return demo if __name__ == "__main__": if not GROQ_API_KEY or not HF_TOKEN: print("ERROR: API keys not found. Make sure to set GROQ_API_KEY and HF_TOKEN in your Hugging Face Space Secrets.") else: app = create_demo() app.launch()