Update app.py

app.py

@@ -1614,6 +1614,11 @@ def main():
         st.session_state.custom_model = "gpt-4o"  # Default model
         st.session_state.first_load_complete = False  # Prevent refreshes on first load
         st.session_state.pending_tab_switch = None  # Track pending tab switches
+        
+        # New chat-related session state variables
+        st.session_state.chat_history = []  # Store chat messages
+        st.session_state.user_message = ""  # Current user message
+        st.session_state.chat_export_data = None  # For exporting chat
 
     # Page configuration with improved layout
     st.set_page_config(
@@ -1623,7 +1628,7 @@ def main():
         initial_sidebar_state="expanded"
     )
 
-    # Custom CSS for improved UI
+    # Custom CSS for improved UI (including chat UI styles)
     st.markdown("""
     <style>
     .main-header {
@@ -1721,6 +1726,67 @@ def main():
         margin-top: 1rem;
         border-left: 4px solid #ef4444;
     }
+    
+    /* Chat UI styles */
+    .chat-container {
+        display: flex;
+        flex-direction: column;
+        height: 500px;
+        overflow-y: auto;
+        border: 1px solid #e0e0e0;
+        border-radius: 10px;
+        padding: 1rem;
+        margin-bottom: 1rem;
+        background-color: #f8f9fa;
+    }
+    
+    .user-message {
+        align-self: flex-end;
+        background-color: #4F46E5;
+        color: white;
+        border-radius: 18px 18px 0 18px;
+        padding: 0.8rem 1rem;
+        margin: 0.5rem 0;
+        max-width: 80%;
+    }
+    
+    .ai-message {
+        align-self: flex-start;
+        background-color: #e0e0e0;
+        color: #333;
+        border-radius: 18px 18px 18px 0;
+        padding: 0.8rem 1rem;
+        margin: 0.5rem 0;
+        max-width: 80%;
+    }
+    
+    .message-time {
+        font-size: 0.7rem;
+        color: #777;
+        margin-top: 0.2rem;
+    }
+    
+    .chat-input {
+        display: flex;
+        margin-top: 1rem;
+    }
+    
+    .chat-input input {
+        flex-grow: 1;
+        padding: 0.8rem;
+        border: 1px solid #ccc;
+        border-radius: 5px;
+        margin-right: 0.5rem;
+    }
+    
+    .chat-input button {
+        padding: 0.8rem 1.5rem;
+        background-color: #4F46E5;
+        color: white;
+        border: none;
+        border-radius: 5px;
+        cursor: pointer;
+    }
     </style>
     """, unsafe_allow_html=True)
 
@@ -1977,7 +2043,7 @@ class MyScene(Scene):
                 else:
                     st.success(st.session_state.status)
 
-    # AI ASSISTANT TAB
+    # AI ASSISTANT TAB - ENHANCED WITH CHAT FUNCTIONALITY
     with tabs[1]:
         st.markdown("### 🤖 AI Animation Assistant")
         
@@ -2052,10 +2118,12 @@ class MyScene(Scene):
                 "DeepSeek-R1",
                 "Meta-Llama-3.1-405B-Instruct",
                 "Llama-3.2-90B-Vision-Instruct",
-                "Llama-3.3-70B-Instruct"
+                "Llama-3.3-70B-Instruct",
                 "Llama-4-Scout-17B-16E-Instruct",
                 "Llama-4-Maverick-17B-128E-Instruct-FP8",
+                "gpt-4o-mini",
                 "gpt-4o",
+                "gpt-4.1",
                 "o3-mini",
                 "o1",
                 "o1-mini",
@@ -2080,46 +2148,287 @@ class MyScene(Scene):
                     st.session_state.ai_models['model_name'] = st.session_state.custom_model
                     st.success(f"Model updated to {st.session_state.custom_model}")
             
-            # AI code generation
+            # Enhanced chat interface
             if st.session_state.ai_models and "client" in st.session_state.ai_models:
-                st.markdown("<div class='card'>", unsafe_allow_html=True)
-                st.markdown("#### Generate Animation from Description")
-                st.write("Describe the animation you want to create, or provide partial code to complete.")
+                st.markdown("#### 💬 AI Chat Interface")
+                st.markdown("Chat with the AI about animation concepts, get code suggestions, and troubleshoot issues.")
                 
-                # Predefined animation ideas dropdown
-                animation_ideas = [
-                    "Select an idea...",
-                    "Create a 3D animation showing a sphere morphing into a torus",
-                    "Show a visual proof of the Pythagorean theorem",
-                    "Visualize a Fourier transform converting a signal from time domain to frequency domain",
-                    "Create an animation explaining neural network forward propagation",
-                    "Illustrate the concept of integration with area under a curve"
-                ]
+                # Chat history display container
+                st.markdown("<div class='chat-container' id='chat-container'>", unsafe_allow_html=True)
                 
-                selected_idea = st.selectbox(
-                    "Try one of these ideas", 
-                    options=animation_ideas
-                )
+                # System message at the start if chat is empty
+                if not st.session_state.chat_history:
+                    st.markdown(
+                        "<div class='ai-message'>"
+                        "Hello! I'm your Manim Animation Assistant. I can help you create and explain mathematical animations. "
+                        "Ask me about creating specific animations, explaining concepts, or troubleshooting your code."
+                        "<div class='message-time'>System</div>"
+                        "</div>",
+                        unsafe_allow_html=True
+                    )
                 
-                prompt_value = selected_idea if selected_idea != "Select an idea..." else ""
+                # Display chat history
+                for message in st.session_state.chat_history:
+                    if message["role"] == "user":
+                        st.markdown(
+                            f"<div class='user-message'>{message['content']}"
+                            f"<div class='message-time'>{message['time']}</div>"
+                            "</div>",
+                            unsafe_allow_html=True
+                        )
+                    else:
+                        st.markdown(
+                            f"<div class='ai-message'>{message['content']}"
+                            f"<div class='message-time'>{message['time']}</div>"
+                            "</div>",
+                            unsafe_allow_html=True
+                        )
                 
-                code_input = st.text_area(
-                    "Your Prompt or Code",
-                    value=prompt_value,
-                    placeholder="Example: Create an animation that shows a circle morphing into a square while changing color from red to blue",
-                    height=150
-                )
+                st.markdown("</div>", unsafe_allow_html=True)
                 
-                if st.button("Generate Animation Code", key="gen_ai_code"):
-                    if code_input:
-                        with st.spinner("AI is generating your animation code..."):
-                            try:
-                                # Direct implementation of code generation
-                                client = st.session_state.ai_models["client"]
-                                model_name = st.session_state.ai_models["model_name"]
+                # Input form for new messages
+                with st.form(key="chat_form", clear_on_submit=True):
+                    col1, col2 = st.columns([5, 1])
+                    
+                    with col1:
+                        user_input = st.text_area(
+                            "Your message", 
+                            value=st.session_state.user_message, 
+                            key="user_input_area",
+                            placeholder="Ask about animation concepts, request code examples, or describe what you want to create...",
+                            height=100
+                        )
+                    
+                    with col2:
+                        submit_button = st.form_submit_button("Send")
+                
+                # Process form submission
+                if submit_button and user_input:
+                    # Store user message in history
+                    current_time = datetime.now().strftime("%H:%M:%S")
+                    st.session_state.chat_history.append({
+                        "role": "user",
+                        "content": user_input,
+                        "time": current_time
+                    })
+                    
+                    # Get AI response
+                    with st.spinner("AI is thinking..."):
+                        try:
+                            client = st.session_state.ai_models["client"]
+                            model_name = st.session_state.ai_models["model_name"]
+                            
+                            # Convert chat history to API format
+                            from azure.ai.inference.models import UserMessage, AssistantMessage
+                            
+                            # Create messages array for API
+                            messages = []
+                            
+                            # Add system message with Manim context
+                            system_message = """
+                            I am a Manim Animation Assistant. I can help with:
+                            1. Creating mathematical animations using Manim
+                            2. Explaining animation concepts and techniques
+                            3. Generating and debugging Manim code
+                            4. Offering creative suggestions for visualizing concepts
+                            
+                            When providing code, I'll make sure it's complete, runnable Manim code with proper imports and structure.
+                            """
+                            
+                            # Add system message as first user message (workaround for some models)
+                            messages.append(UserMessage("You are a Manim animation expert assistant. Respond to my questions with helpful information about creating mathematical animations with Manim."))
+                            
+                            # Add chat history (limited to last 10 messages for context window)
+                            for msg in st.session_state.chat_history[-10:]:
+                                if msg["role"] == "user":
+                                    messages.append(UserMessage(msg["content"]))
+                                else:
+                                    messages.append(AssistantMessage(msg["content"]))
+                            
+                            # Get response from API
+                            response = client.complete(
+                                messages=messages,
+                                max_tokens=1000,
+                                model=model_name
+                            )
+                            
+                            # Process the response
+                            if response and response.choices and len(response.choices) > 0:
+                                ai_response = response.choices[0].message.content
+                                
+                                # Store AI response in history
+                                st.session_state.chat_history.append({
+                                    "role": "assistant",
+                                    "content": ai_response,
+                                    "time": datetime.now().strftime("%H:%M:%S")
+                                })
                                 
-                                # Create the prompt
-                                prompt = f"""Write a complete Manim animation scene based on this code or idea:
+                                # Look for code in the response
+                                if "```python" in ai_response and "```" in ai_response:
+                                    # Extract code blocks
+                                    code_blocks = []
+                                    parts = ai_response.split("```python")
+                                    for part in parts[1:]:  # Skip the first part (before first code block)
+                                        if "```" in part:
+                                            code = part.split("```")[0].strip()
+                                            code_blocks.append(code)
+                                    
+                                    # If we have code blocks, offer to use the first one
+                                    if code_blocks and "class" in code_blocks[0] and "Scene" in code_blocks[0]:
+                                        st.info("Manim code detected in the response. Would you like to use it?")
+                                        if st.button("Use This Code in Editor", key="use_chat_code"):
+                                            st.session_state.code = code_blocks[0]
+                                            st.session_state.temp_code = code_blocks[0]
+                                            st.session_state.pending_tab_switch = 0  # Switch to editor tab
+                                            st.rerun()
+                            else:
+                                # Handle empty response
+                                st.error("The AI model returned an empty response. Please try again.")
+                        except Exception as e:
+                            st.error(f"Error getting AI response: {str(e)}")
+                            import traceback
+                            st.code(traceback.format_exc())
+                            
+                            # Add error message to chat
+                            st.session_state.chat_history.append({
+                                "role": "assistant",
+                                "content": f"Sorry, I encountered an error: {str(e)}. Please try again or check the connection.",
+                                "time": datetime.now().strftime("%H:%M:%S")
+                            })
+                    
+                    # Clear the user input
+                    st.session_state.user_message = ""
+                    # Force a rerun to update the chat display
+                    st.rerun()
+                
+                # Chat export options
+                if st.session_state.chat_history:
+                    st.markdown("#### Chat Options")
+                    col1, col2 = st.columns(2)
+                    
+                    with col1:
+                        if st.button("Export Chat", key="export_chat_btn"):
+                            # Format chat history as text
+                            chat_text = "# Manim Animation Chat Export\n\n"
+                            for msg in st.session_state.chat_history:
+                                role = "You" if msg["role"] == "user" else "AI Assistant"
+                                chat_text += f"## {role} ({msg['time']})\n\n{msg['content']}\n\n"
+                            
+                            # Store in session state for download
+                            st.session_state.chat_export_data = chat_text
+                            
+                            # Show download button
+                            st.download_button(
+                                "📥 Download Chat Text",
+                                data=st.session_state.chat_export_data,
+                                file_name=f"manim_chat_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md",
+                                mime="text/markdown"
+                            )
+                    
+                    with col2:
+                        if st.button("Clear Chat", key="clear_chat_btn"):
+                            # Clear chat history
+                            st.session_state.chat_history = []
+                            st.success("Chat history cleared!")
+                            st.rerun()
+                
+                # Additional AI features (code generation, etc.)
+                with st.expander("✨ Animation Idea Generator"):
+                    st.markdown("Let the AI help you generate animation ideas based on concepts.")
+                    
+                    concept_input = st.text_input(
+                        "Enter a mathematical or scientific concept:",
+                        placeholder="e.g., Pythagorean theorem, Wave interference, Neural networks"
+                    )
+                    
+                    if st.button("Generate Animation Ideas", key="gen_ideas_btn"):
+                        if concept_input:
+                            with st.spinner("Generating ideas..."):
+                                try:
+                                    # Get ideas from AI
+                                    client = st.session_state.ai_models["client"]
+                                    model_name = st.session_state.ai_models["model_name"]
+                                    
+                                    idea_prompt = f"""Generate 3 creative animation ideas for visualizing the concept of "{concept_input}" using Manim.
+                                    For each idea, provide:
+                                    1. A title
+                                    2. A brief description of what the animation would show
+                                    3. Key visual elements to include
+                                    
+                                    Format your response with clear sections for each idea."""
+                                    
+                                    from azure.ai.inference.models import UserMessage
+                                    response = client.complete(
+                                        messages=[
+                                            UserMessage(idea_prompt),
+                                        ],
+                                        max_tokens=800,
+                                        model=model_name
+                                    )
+                                    
+                                    if response and response.choices and len(response.choices) > 0:
+                                        ideas = response.choices[0].message.content
+                                        st.markdown("### Animation Ideas")
+                                        st.markdown(ideas)
+                                        
+                                        # Add to chat history
+                                        st.session_state.chat_history.append({
+                                            "role": "user",
+                                            "content": f"Generate animation ideas for: {concept_input}",
+                                            "time": datetime.now().strftime("%H:%M:%S")
+                                        })
+                                        
+                                        st.session_state.chat_history.append({
+                                            "role": "assistant",
+                                            "content": ideas,
+                                            "time": datetime.now().strftime("%H:%M:%S")
+                                        })
+                                    else:
+                                        st.error("Failed to generate ideas. The AI returned an empty response.")
+                                except Exception as e:
+                                    st.error(f"Error generating ideas: {str(e)}")
+                        else:
+                            st.warning("Please enter a concept first")
+                
+                # AI code generation
+                with st.expander("🧩 Quick Code Generator"):
+                    st.markdown("#### Generate Animation from Description")
+                    st.write("Describe the animation you want to create, or provide partial code to complete.")
+                    
+                    # Predefined animation ideas dropdown
+                    animation_ideas = [
+                        "Select an idea...",
+                        "Create a 3D animation showing a sphere morphing into a torus",
+                        "Show a visual proof of the Pythagorean theorem",
+                        "Visualize a Fourier transform converting a signal from time domain to frequency domain",
+                        "Create an animation explaining neural network forward propagation",
+                        "Illustrate the concept of integration with area under a curve"
+                    ]
+                    
+                    selected_idea = st.selectbox(
+                        "Try one of these ideas", 
+                        options=animation_ideas
+                    )
+                    
+                    prompt_value = selected_idea if selected_idea != "Select an idea..." else ""
+                    
+                    code_input = st.text_area(
+                        "Your Prompt or Code",
+                        value=prompt_value,
+                        placeholder="Example: Create an animation that shows a circle morphing into a square while changing color from red to blue",
+                        height=120
+                    )
+                    
+                    if st.button("Generate Animation Code", key="gen_ai_code"):
+                        if code_input:
+                            with st.spinner("AI is generating your animation code..."):
+                                try:
+                                    # Direct implementation of code generation
+                                    client = st.session_state.ai_models["client"]
+                                    model_name = st.session_state.ai_models["model_name"]
+                                    
+                                    # Create the prompt
+                                    prompt = f"""Write a complete Manim animation scene based on this code or idea:
 {code_input}
 
 The code should be a complete, working Manim animation that includes:
@@ -2130,84 +2439,93 @@ The code should be a complete, working Manim animation that includes:
 
 Here's the complete Manim code:
 """
-                                
-                                # Make API call directly
-                                from azure.ai.inference.models import UserMessage
-                                response = client.complete(
-                                    messages=[
-                                        UserMessage(prompt),
-                                    ],
-                                    max_tokens=1000,
-                                    model=model_name
-                                )
-                                
-                                # Process the response
-                                if response and response.choices and len(response.choices) > 0:
-                                    completed_code = response.choices[0].message.content
                                     
-                                    # Extract code from markdown if present
-                                    if "```python" in completed_code:
-                                        completed_code = completed_code.split("```python")[1].split("```")[0]
-                                    elif "```" in completed_code:
-                                        completed_code = completed_code.split("```")[1].split("```")[0]
+                                    # Make API call directly
+                                    from azure.ai.inference.models import UserMessage
+                                    response = client.complete(
+                                        messages=[
+                                            UserMessage(prompt),
+                                        ],
+                                        max_tokens=1000,
+                                        model=model_name
+                                    )
                                     
-                                    # Add Scene class if missing
-                                    if "Scene" not in completed_code:
-                                        completed_code = f"""from manim import *
+                                    # Process the response
+                                    if response and response.choices and len(response.choices) > 0:
+                                        completed_code = response.choices[0].message.content
+                                        
+                                        # Extract code from markdown if present
+                                        if "```python" in completed_code:
+                                            completed_code = completed_code.split("```python")[1].split("```")[0]
+                                        elif "```" in completed_code:
+                                            completed_code = completed_code.split("```")[1].split("```")[0]
+                                        
+                                        # Add Scene class if missing
+                                        if "Scene" not in completed_code:
+                                            completed_code = f"""from manim import *
 
 class MyScene(Scene):
     def construct(self):
         {completed_code}"""
-                                    
-                                    # Store the generated code
-                                    st.session_state.generated_code = completed_code
-                                else:
-                                    st.error("Failed to generate code. API returned an empty response.")
-                            except Exception as e:
-                                st.error(f"Error generating code: {str(e)}")
-                                import traceback
-                                st.code(traceback.format_exc())
-                    else:
-                        st.warning("Please enter a description or prompt first")
-                
-                st.markdown("</div>", unsafe_allow_html=True)
-                
-                # AI generated code display and actions
-                if "generated_code" in st.session_state and st.session_state.generated_code:
-                    st.markdown("<div class='card'>", unsafe_allow_html=True)
-                    st.markdown("#### Generated Animation Code")
-                    st.code(st.session_state.generated_code, language="python")
-                    
-                    col_ai1, col_ai2 = st.columns(2)
-                    with col_ai1:
-                        if st.button("Use This Code", key="use_gen_code"):
-                            st.session_state.code = st.session_state.generated_code
-                            st.session_state.temp_code = st.session_state.generated_code
-                            # Set pending tab switch to editor tab
-                            st.session_state.pending_tab_switch = 0
-                            st.rerun()
+                                        
+                                        # Store the generated code
+                                        st.session_state.generated_code = completed_code
+                                        
+                                        # Add to chat history
+                                        st.session_state.chat_history.append({
+                                            "role": "user",
+                                            "content": f"Generate code for: {code_input}",
+                                            "time": datetime.now().strftime("%H:%M:%S")
+                                        })
+                                        
+                                        st.session_state.chat_history.append({
+                                            "role": "assistant",
+                                            "content": f"Here's the Manim code for your request:\n\n```python\n{completed_code}\n```",
+                                            "time": datetime.now().strftime("%H:%M:%S")
+                                        })
+                                    else:
+                                        st.error("Failed to generate code. API returned an empty response.")
+                                except Exception as e:
+                                    st.error(f"Error generating code: {str(e)}")
+                                    import traceback
+                                    st.code(traceback.format_exc())
+                        else:
+                            st.warning("Please enter a description or prompt first")
                     
-                    with col_ai2:
-                        if st.button("Render Preview", key="render_preview"):
-                            with st.spinner("Rendering preview..."):
-                                video_data, status = generate_manim_video(
-                                    st.session_state.generated_code, 
-                                    "mp4", 
-                                    "480p",  # Use lowest quality for preview
-                                    ANIMATION_SPEEDS["Normal"]
-                                )
-                                
-                                if video_data:
-                                    st.video(video_data)
-                                    st.download_button(
-                                        label="Download Preview",
-                                        data=video_data,
-                                        file_name=f"manim_preview_{int(time.time())}.mp4",
-                                        mime="video/mp4"
+                    # AI generated code display and actions
+                    if "generated_code" in st.session_state and st.session_state.generated_code:
+                        st.markdown("#### Generated Animation Code")
+                        st.code(st.session_state.generated_code, language="python")
+                        
+                        col_ai1, col_ai2 = st.columns(2)
+                        with col_ai1:
+                            if st.button("Use This Code", key="use_gen_code"):
+                                st.session_state.code = st.session_state.generated_code
+                                st.session_state.temp_code = st.session_state.generated_code
+                                # Set pending tab switch to editor tab
+                                st.session_state.pending_tab_switch = 0
+                                st.rerun()
+                        
+                        with col_ai2:
+                            if st.button("Render Preview", key="render_preview"):
+                                with st.spinner("Rendering preview..."):
+                                    video_data, status = generate_manim_video(
+                                        st.session_state.generated_code, 
+                                        "mp4", 
+                                        "480p",  # Use lowest quality for preview
+                                        ANIMATION_SPEEDS["Normal"]
                                     )
-                                else:
-                                    st.error(f"Failed to generate preview: {status}")
-                    st.markdown("</div>", unsafe_allow_html=True)
+                                    
+                                    if video_data:
+                                        st.video(video_data)
+                                        st.download_button(
+                                            label="Download Preview",
+                                            data=video_data,
+                                            file_name=f"manim_preview_{int(time.time())}.mp4",
+                                            mime="video/mp4"
+                                        )
+                                    else:
+                                        st.error(f"Failed to generate preview: {status}")
             else:
                 st.warning("AI models not initialized. Please use the Debug Connection section to test API connectivity.")
         else:
@@ -3046,3 +3364,6 @@ class PlotScene(Scene):
 
 if __name__ == "__main__":
     main()
+
+if __name__ == "__main__":
+    main()