import gradio as gr import os from huggingface_hub import InferenceClient import random from typing import Generator, Dict, List, Tuple, Optional # Get token from environment variable hf_token = os.environ.get("HF_TOKEN") client = InferenceClient("HuggingFaceH4/zephyr-7b-beta", token=hf_token) # Story genres with genre-specific example prompts GENRE_EXAMPLES = { "fairy_tale": [ "I follow the shimmer of fairy dust into a hidden forest" "I meet a talking rabbit who claims to know a secret about the king’s lost crown" "A tiny dragon appears at my window, asking for help to find its mother" "I step into a clearing where the trees whisper ancient riddles" "A friendly witch invites me into her cozy cottage, offering a warm cup of tea" ], "fantasy": [ "I enter the ancient forest seeking the wizard's tower", "I approach the dragon cautiously with my shield raised", "I examine the mysterious runes carved into the stone altar", "I try to bargain with the elven council for safe passage" ], "sci-fi": [ "I hack into the space station's mainframe", "I investigate the strange signal coming from the abandoned planet", "I negotiate with the alien ambassador about the peace treaty", "I try to repair my damaged spacecraft before oxygen runs out" ], "mystery": [ "I examine the crime scene for overlooked evidence", "I question the nervous butler about the night of the murder", "I follow the suspicious figure through the foggy streets", "I check the victim's diary for hidden clues" ], "horror": [ "I slowly open the creaking door to the basement", "I read the forbidden text while the candles flicker", "I hide under the bed as footsteps approach", "I investigate the strange noises coming from the attic" ], "western": [ "I challenge the outlaw to a duel at high noon", "I track the bandits through the desert canyon", "I enter the saloon looking for information", "I defend the stagecoach from the approaching raiders" ], "cyberpunk": [ "I jack into the corporate mainframe to steal data", "I negotiate with the street gang for cybernetic upgrades", "I hide in the neon-lit alleyway from corporate security", "I meet my mysterious client in the underground bar" ], "historical": [ "I attend the royal ball hoping to meet the mysterious count", "I join the resistance against the occupying forces", "I navigate the dangerous politics of the royal court", "I set sail on a voyage to discover new lands" ], "post-apocalyptic": [ "I scavenge the abandoned shopping mall for supplies", "I approach the fortified settlement seeking shelter", "I navigate through the radioactive zone using my old map", "I hide from the approaching group of raiders" ], "steampunk": [ "I pilot my airship through the lightning storm", "I present my new invention to the Royal Academy", "I investigate the mysterious clockwork automaton", "I sneak aboard the emperor's armored train" ] } # 2. Add constants at the top for magic numbers MAX_HISTORY_LENGTH = 20 MEMORY_WINDOW = 10 MAX_TOKENS = 2048 # Doubled for longer responses TEMPERATURE = 0.8 # Slightly increased for more creative responses TOP_P = 0.95 MIN_RESPONSE_LENGTH = 200 # Minimum characters before yielding response def get_examples_for_genre(genre): """Get example prompts specific to the selected genre""" return GENRE_EXAMPLES.get(genre, GENRE_EXAMPLES["fantasy"]) def get_enhanced_system_prompt(genre=None): """Generate a detailed system prompt with optional genre specification""" selected_genre = genre or "fantasy" system_message = f"""You are an interactive storyteller creating an immersive {selected_genre} choose-your-own-adventure story. For each response you MUST: 1. Write at least 100 words describing the scene, using vivid sensory details 2. Include character dialogue or thoughts that reveal personality and motivations 3. Create a strong sense of atmosphere appropriate for {selected_genre} 4. End EVERY response with exactly three distinct choices, formatted as: - Option 1: [Complete sentence describing a possible action] - Option 2: [Complete sentence describing a possible action] - Option 3: [Complete sentence describing a possible action] Never respond with just one word. Always provide a detailed scene and three choices. If the user's input is unclear, interpret their intent and continue the story naturally. Keep the story cohesive by referencing previous events and choices.""" return system_message def create_story_summary(chat_history): """Create a concise summary of the story so far if the history gets too long""" if len(chat_history) <= 2: return None story_text = "" for user_msg, bot_msg in chat_history: story_text += f"User: {user_msg}\nStory: {bot_msg}\n\n" summary_instruction = { "role": "system", "content": "The conversation history is getting long. Please create a brief summary of the key plot points and character development so far to help maintain context without exceeding token limits." } return summary_instruction def format_history_for_gradio(history_tuples): """Convert (user, bot) tuples into Gradio 'messages' format (role/content dicts).""" messages = [] for user_msg, bot_msg in history_tuples: messages.append({"role": "user", "content": user_msg}) messages.append({"role": "assistant", "content": bot_msg}) return messages # 1. Add type hints for better code maintainability # 4. Add input validation def respond( message: str, chat_history: List[Tuple[str, str]], genre: Optional[str] = None, use_full_memory: bool = True ) -> Generator[List[Dict[str, str]], None, None]: """Generate a response based on the current message and conversation history.""" if not message.strip(): return chat_history if genre and genre not in GENRE_EXAMPLES: genre = "fantasy" # fallback to default system_message = get_enhanced_system_prompt(genre) # Convert your existing (user, bot) history into a format for the API request formatted_history = [] for user_msg, bot_msg in chat_history: formatted_history.append({"role": "user", "content": user_msg}) formatted_history.append({"role": "assistant", "content": bot_msg}) api_messages = [{"role": "system", "content": system_message}] # Use full memory or partial memory if use_full_memory and formatted_history: if len(formatted_history) > MAX_HISTORY_LENGTH: summary_instruction = create_story_summary(chat_history[:len(chat_history)-5]) if summary_instruction: api_messages.append(summary_instruction) for msg in formatted_history[-MEMORY_WINDOW:]: api_messages.append(msg) else: for msg in formatted_history: api_messages.append(msg) else: memory_length = MEMORY_WINDOW if formatted_history: for msg in formatted_history[-memory_length*2:]: api_messages.append(msg) # Add current user message api_messages.append({"role": "user", "content": message}) # Special handling for story initialization if not chat_history or message.lower() in ["start", "begin", "begin my adventure"]: api_messages.append({ "role": "system", "content": f"Begin a new {genre or 'fantasy'} adventure with an intriguing opening scene. Introduce the protagonist without assuming too much about them." }) bot_message = "" try: for response_chunk in client.chat_completion( api_messages, max_tokens=MAX_TOKENS, stream=True, temperature=TEMPERATURE, top_p=TOP_P, ): delta = response_chunk.choices[0].delta.content if delta: bot_message += delta # Only yield complete responses if len(bot_message.strip()) >= MIN_RESPONSE_LENGTH and "Option 3:" in bot_message: new_history = chat_history.copy() new_history.append((message, bot_message)) yield format_history_for_gradio(new_history) except Exception as e: error_message = f"Story magic temporarily interrupted. Please try again. (Error: {str(e)})" broken_history = chat_history + [(message, error_message)] yield format_history_for_gradio(broken_history) def save_story(chat_history): """Convert chat history to markdown for download""" if not chat_history: return "No story to save yet!" story_text = "# My Interactive Adventure\n\n" for user_msg, bot_msg in chat_history: story_text += f"**Player:** {user_msg}\n\n" story_text += f"**Story:** {bot_msg}\n\n---\n\n" return story_text with gr.Blocks(theme=gr.themes.Soft()) as demo: gr.Markdown("# 🔮 Interactive Story Time") gr.Markdown("Create a completely unique literary world, one choice at a time. Dare to explore the unknown.") with gr.Row(): with gr.Column(scale=3): # Chat window + user input chatbot = gr.Chatbot( height=700, # Increased height bubble_full_width=True, # Allow bubbles to use full width show_copy_button=True, avatar_images=(None, "🧙"), type="messages", container=True, scale=1, min_width=800 # Ensure minimum width ) msg = gr.Textbox( placeholder="Describe what you want to do next in the story...", container=False, scale=4, ) with gr.Row(): submit = gr.Button("Continue Story", variant="primary") clear = gr.Button("Start New Adventure") with gr.Column(scale=1): gr.Markdown("## Adventure Settings") genre = gr.Dropdown( choices=list(GENRE_EXAMPLES.keys()), label="Story Genre", info="Choose the theme of your next adventure", value="fantasy" ) full_memory = gr.Checkbox( label="Full Story Memory", value=True, info="When enabled, the AI tries to remember the entire story. If disabled, only the last few exchanges are used." ) gr.Markdown("## Story Starters") # Create four placeholder buttons for story starters starter_btn1 = gr.Button("Starter 1") starter_btn2 = gr.Button("Starter 2") starter_btn3 = gr.Button("Starter 3") starter_btn4 = gr.Button("Starter 4") starter_buttons = [starter_btn1, starter_btn2, starter_btn3, starter_btn4] # 1) We'll return a list of 4 dicts, each dict updating 'value' & 'visible' def update_starter_buttons(selected_genre): """Update starter buttons with examples for the selected genre.""" examples = get_examples_for_genre(selected_genre) results = [] for i in range(4): if i < len(examples): # Return just the string value instead of a dict results.append(examples[i]) else: results.append("") # Empty string for hidden buttons return tuple(results) # Return tuple of strings # 2) Initialize them with "fantasy" so they don't stay "Starter X" on page load # We'll just call the function and store the results in a variable, then apply them in a .load() event initial_button_data = update_starter_buttons("fantasy") # returns 4 dicts # 3) We'll define a "pick_starter" function that sets msg to the chosen text def pick_starter(starter_text, chat_history, selected_genre, memory_flag): # Putting 'starter_text' into the msg return starter_text # 4) Connect each starter button: for starter_button in starter_buttons: starter_button.click( fn=pick_starter, inputs=[starter_button, chatbot, genre, full_memory], outputs=[msg], queue=False ).then( fn=respond, inputs=[msg, chatbot, genre, full_memory], outputs=[chatbot], queue=False ) # 5) Dynamically update the 4 buttons if the user changes the genre genre.change( fn=update_starter_buttons, inputs=[genre], outputs=starter_buttons ) # Handler for user input msg.submit(respond, [msg, chatbot, genre, full_memory], [chatbot]) submit.click(respond, [msg, chatbot, genre, full_memory], [chatbot]) # Clear the chatbot for a new adventure clear.click(lambda: [], None, chatbot, queue=False) clear.click(lambda: "", None, msg, queue=False) # "Download My Story" row with gr.Row(): save_btn = gr.Button("Download My Story", variant="secondary") story_output = gr.Markdown(visible=False) save_btn.click(save_story, inputs=[chatbot], outputs=[story_output]) save_btn.click( fn=lambda: True, inputs=None, outputs=story_output, js="() => {document.getElementById('story_output').scrollIntoView();}", queue=False ) # 6) Finally, run a "load" event to apply initial_button_data to the 4 button outputs on page load def load_initial_buttons(): # Just return our precomputed tuple of 4 dicts return initial_button_data demo.load(fn=load_initial_buttons, outputs=starter_buttons, queue=False) # Run the app if __name__ == "__main__": demo.launch(server_name="0.0.0.0", server_port=7860)