Create app.py

app.py

@@ -0,0 +1,347 @@
+import os
+import gradio as gr
+from langchain.chat_models import ChatOpenAI
+from langchain.prompts import ChatPromptTemplate
+from langchain.chains import LLMChain, SequentialChain
+from langchain.memory import ConversationBufferMemory
+import random
+
+# Global variables
+current_story = ""
+story_branches = {}
+story_history = []
+
+def generate_story_idea(llm, genre, theme, length):
+    """Generate a story idea based on user inputs"""
+    idea_prompt = ChatPromptTemplate.from_template(
+        """You are a creative writing assistant.
+        Generate a compelling story idea for a {length} story in the {genre} genre 
+        that explores the theme of {theme}.
+        Include a brief plot outline, main character, and setting.
+        Respond with just the idea, no additional commentary.
+        """
+    )
+    
+    idea_chain = LLMChain(llm=llm, prompt=idea_prompt, output_key="story_idea")
+    return idea_chain.run(genre=genre, theme=theme, length=length)
+
+def create_story_beginning(llm, story_idea):
+    """Create the beginning of a story based on the idea"""
+    beginning_prompt = ChatPromptTemplate.from_template(
+        """You are a creative writing assistant.
+        Based on this story idea: {story_idea}
+        
+        Write an engaging opening for this story that introduces the main character and setting.
+        Make it compelling and hook the reader immediately.
+        Write approximately 250-350 words.
+        """
+    )
+    
+    beginning_chain = LLMChain(llm=llm, prompt=beginning_prompt, output_key="story_beginning")
+    return beginning_chain.run(story_idea=story_idea)
+
+def continue_linear_story(llm, story_so_far):
+    """Continue a linear story based on the story so far"""
+    continue_prompt = ChatPromptTemplate.from_template(
+        """You are a creative writing assistant.
+        Continue this story:
+        
+        {story_so_far}
+        
+        Write the next part of the story (approximately 250-350 words), advancing the plot in an interesting way.
+        End at a point that feels satisfying but leaves room for more story.
+        """
+    )
+    
+    continue_chain = LLMChain(llm=llm, prompt=continue_prompt, output_key="story_continuation")
+    return continue_chain.run(story_so_far=story_so_far)
+
+def generate_story_branches(llm, story_so_far):
+    """Generate three possible story continuations as branches"""
+    branch_prompt = ChatPromptTemplate.from_template(
+        """You are a creative writing assistant.
+        Based on this story so far:
+        
+        {story_so_far}
+        
+        Generate THREE possible directions the story could take next. For each:
+        1. Provide a brief title (10 words or less)
+        2. Write a short description (1-2 sentences)
+        
+        Format as:
+        Option 1: [Title]
+        [Description]
+        
+        Option 2: [Title]
+        [Description]
+        
+        Option 3: [Title]
+        [Description]
+        """
+    )
+    
+    branch_chain = LLMChain(llm=llm, prompt=branch_prompt, output_key="story_branches")
+    branches_text = branch_chain.run(story_so_far=story_so_far)
+    
+    # Parse the branches
+    branches = []
+    lines = branches_text.strip().split('\n')
+    current_option = None
+    current_description = ""
+    
+    for line in lines:
+        if line.startswith("Option "):
+            if current_option:
+                branches.append((current_option, current_description.strip()))
+            current_option = line.split(": ", 1)[1] if ": " in line else line
+            current_description = ""
+        elif current_option is not None:
+            current_description += line + " "
+    
+    if current_option:
+        branches.append((current_option, current_description.strip()))
+    
+    return branches
+
+def continue_branch(llm, story_so_far, branch_title, branch_description):
+    """Continue the story based on the selected branch"""
+    branch_continue_prompt = ChatPromptTemplate.from_template(
+        """You are a creative writing assistant.
+        Continue this story:
+        
+        {story_so_far}
+        
+        The story should now follow this direction:
+        {branch_title} - {branch_description}
+        
+        Write the next part of the story (approximately 250-350 words) following this direction.
+        Make it engaging and consistent with what came before.
+        """
+    )
+    
+    branch_chain = LLMChain(llm=llm, prompt=branch_continue_prompt, output_key="branch_continuation")
+    return branch_chain.run(
+        story_so_far=story_so_far, 
+        branch_title=branch_title, 
+        branch_description=branch_description
+    )
+
+def initialize_story(api_key, genre, theme, length):
+    """Initialize a new story with the given parameters"""
+    global current_story, story_branches, story_history
+    
+    try:
+        # Initialize LLM
+        llm = ChatOpenAI(
+            openai_api_key=api_key,
+            model="gpt-4-turbo",
+            temperature=0.7
+        )
+        
+        # Reset story state
+        current_story = ""
+        story_branches = {}
+        story_history = []
+        
+        # Generate story idea and beginning
+        story_idea = generate_story_idea(llm, genre, theme, length)
+        story_beginning = create_story_beginning(llm, story_idea)
+        
+        current_story = story_beginning
+        story_history.append(current_story)
+        
+        # Return the story beginning and available actions
+        return (
+            f"Story Idea:\n{story_idea}\n\n" + 
+            f"Story Beginning:\n{story_beginning}", 
+            gr.update(visible=True), 
+            gr.update(visible=True),
+            gr.update(visible=True)
+        )
+    except Exception as e:
+        return f"Error initializing story: {str(e)}", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+
+def continue_story(api_key):
+    """Continue the linear story"""
+    global current_story, story_history
+    
+    try:
+        # Initialize LLM
+        llm = ChatOpenAI(
+            openai_api_key=api_key,
+            model="gpt-4-turbo",
+            temperature=0.7
+        )
+        
+        # Continue the story
+        continuation = continue_linear_story(llm, current_story)
+        current_story += "\n\n" + continuation
+        story_history.append(continuation)
+        
+        return current_story, gr.update(value="Story continued successfully!")
+    except Exception as e:
+        return current_story, gr.update(value=f"Error continuing story: {str(e)}")
+
+def generate_branches(api_key):
+    """Generate branching options for the story"""
+    global current_story, story_branches
+    
+    try:
+        # Initialize LLM
+        llm = ChatOpenAI(
+            openai_api_key=api_key,
+            model="gpt-4-turbo",
+            temperature=0.8  # Higher temperature for more creative branches
+        )
+        
+        # Generate branches
+        branches = generate_story_branches(llm, current_story)
+        story_branches = {f"Option {i+1}: {title}": (title, desc) for i, (title, desc) in enumerate(branches)}
+        
+        # Create formatted output for display
+        branches_output = "\n\n".join([f"{option}\n{desc}" for option, (_, desc) in story_branches.items()])
+        
+        # Create radio options
+        radio_options = list(story_branches.keys())
+        
+        return branches_output, gr.update(choices=radio_options, value=radio_options[0] if radio_options else None, visible=True), gr.update(visible=True)
+    except Exception as e:
+        return f"Error generating branches: {str(e)}", gr.update(visible=False), gr.update(visible=False)
+
+def select_branch(api_key, selected_branch):
+    """Continue the story based on the selected branch"""
+    global current_story, story_branches, story_history
+    
+    try:
+        if not selected_branch or selected_branch not in story_branches:
+            return current_story, "Please select a valid branch."
+        
+        # Initialize LLM
+        llm = ChatOpenAI(
+            openai_api_key=api_key,
+            model="gpt-4-turbo",
+            temperature=0.7
+        )
+        
+        # Get branch details
+        branch_title, branch_description = story_branches[selected_branch]
+        
+        # Continue along the selected branch
+        continuation = continue_branch(llm, current_story, branch_title, branch_description)
+        
+        # Update story state
+        current_story += f"\n\n[{selected_branch}]\n\n" + continuation
+        story_history.append(f"[Branch: {branch_title}] {continuation}")
+        
+        return current_story, gr.update(value="Branch selected and story continued!")
+    except Exception as e:
+        return current_story, gr.update(value=f"Error selecting branch: {str(e)}")
+
+def create_app():
+    """Create the Gradio interface"""
+    with gr.Blocks(title="AI Story Generator", theme=gr.themes.Soft()) as app:
+        gr.Markdown("# 📚 AI-Powered Story Generator")
+        gr.Markdown("Create interactive stories with branching narratives using OpenAI and LangChain")
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                api_key = gr.Textbox(
+                    label="OpenAI API Key", 
+                    placeholder="Enter your OpenAI API key here", 
+                    type="password"
+                )
+                
+                with gr.Group():
+                    gr.Markdown("### Story Parameters")
+                    genre = gr.Dropdown(
+                        label="Genre",
+                        choices=[
+                            "Fantasy", "Science Fiction", "Mystery", "Romance", 
+                            "Horror", "Adventure", "Historical Fiction", "Comedy"
+                        ],
+                        value="Fantasy"
+                    )
+                    theme = gr.Textbox(
+                        label="Theme", 
+                        placeholder="e.g., Redemption, Loss, Discovery", 
+                        value="Adventure"
+                    )
+                    length = gr.Radio(
+                        label="Story Length", 
+                        choices=["Short Story", "Novella", "Novel Chapter"],
+                        value="Short Story"
+                    )
+                
+                with gr.Row():
+                    start_btn = gr.Button("Start New Story", variant="primary")
+                
+                with gr.Row():
+                    continue_btn = gr.Button("Continue Story", visible=False)
+                    branch_btn = gr.Button("Generate Branch Options", visible=False)
+                
+                status = gr.Textbox(label="Status", value="", visible=True)
+                
+                with gr.Group(visible=False) as branch_group:
+                    branch_output = gr.Textbox(label="Story Branches")
+                    branch_selection = gr.Radio(label="Select a Branch", choices=[])
+                    select_branch_btn = gr.Button("Continue with Selected Branch")
+            
+            with gr.Column(scale=2):
+                story_output = gr.Textbox(
+                    label="Generated Story", 
+                    placeholder="Your story will appear here...",
+                    lines=20
+                )
+        
+        # Define button click events
+        start_btn.click(
+            initialize_story,
+            inputs=[api_key, genre, theme, length],
+            outputs=[story_output, continue_btn, branch_btn, branch_group]
+        )
+        
+        continue_btn.click(
+            continue_story,
+            inputs=[api_key],
+            outputs=[story_output, status]
+        )
+        
+        branch_btn.click(
+            generate_branches,
+            inputs=[api_key],
+            outputs=[branch_output, branch_selection, select_branch_btn]
+        )
+        
+        select_branch_btn.click(
+            select_branch,
+            inputs=[api_key, branch_selection],
+            outputs=[story_output, status]
+        )
+        
+        # Example usage
+        gr.Examples(
+            examples=[
+                ["Fantasy", "Coming of age", "Short Story"],
+                ["Science Fiction", "Artificial intelligence", "Novella"],
+                ["Mystery", "Betrayal", "Novel Chapter"]
+            ],
+            inputs=[genre, theme, length],
+            outputs=[]
+        )
+        
+        gr.Markdown("""
+        ## How to Use
+        1. Enter your OpenAI API key
+        2. Select a genre, theme, and length
+        3. Click "Start New Story" to begin
+        4. Continue the linear story or generate branching options
+        5. If you chose branching, select a branch to follow
+        
+        This app uses LangChain to orchestrate the story generation process and OpenAI's models to create the content.
+        """)
+    
+    return app
+
+if __name__ == "__main__":
+    app = create_app()
+    app.launch()