Update app.py

app.py

@@ -1614,11 +1614,6 @@ 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(
@@ -1628,7 +1623,7 @@ def main():
         initial_sidebar_state="expanded"
     )
 
-    # Custom CSS for improved UI (including chat UI styles)
+    # Custom CSS for improved UI
     st.markdown("""
     <style>
     .main-header {
@@ -1726,67 +1721,6 @@ 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)
 
@@ -2043,7 +1977,7 @@ class MyScene(Scene):
                 else:
                     st.success(st.session_state.status)
 
-    # AI ASSISTANT TAB - ENHANCED WITH CHAT FUNCTIONALITY
+    # AI ASSISTANT TAB
     with tabs[1]:
         st.markdown("### 🤖 AI Animation Assistant")
         
@@ -2083,7 +2017,7 @@ class MyScene(Scene):
                                 messages=[
                                     UserMessage("Hello, this is a connection test."),
                                 ],
-                                max_tokens=10000,
+                                max_tokens=1000000,
                                 model=model_name
                             )
                             
@@ -2118,7 +2052,7 @@ 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",
@@ -2148,287 +2082,46 @@ 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}")
             
-            # Enhanced chat interface
+            # AI code generation
             if st.session_state.ai_models and "client" in st.session_state.ai_models:
-                st.markdown("#### 💬 AI Chat Interface")
-                st.markdown("Chat with the AI about animation concepts, get code suggestions, and troubleshoot issues.")
-                
-                # Chat history display container
-                st.markdown("<div class='chat-container' id='chat-container'>", unsafe_allow_html=True)
+                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.")
                 
-                # 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
-                    )
+                # 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"
+                ]
                 
-                # 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
-                        )
+                selected_idea = st.selectbox(
+                    "Try one of these ideas", 
+                    options=animation_ideas
+                )
                 
-                st.markdown("</div>", unsafe_allow_html=True)
+                prompt_value = selected_idea if selected_idea != "Select an idea..." else ""
                 
-                # 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")
+                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
+                )
                 
-                # 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")
-                                })
+                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"]
                                 
-                                # 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:
+                                # 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:
@@ -2439,93 +2132,84 @@ 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
                                     
-                                    # Make API call directly
-                                    from azure.ai.inference.models import UserMessage
-                                    response = client.complete(
-                                        messages=[
-                                            UserMessage(prompt),
-                                        ],
-                                        max_tokens=1000,
-                                        model=model_name
-                                    )
+                                    # 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]
                                     
-                                    # 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 *
+                                    # 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
-                                        
-                                        # 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")
+                                    
+                                    # 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")
                     
-                    # 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"]
+                    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"]
+                                )
+                                
+                                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"
                                     )
-                                    
-                                    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.error(f"Failed to generate preview: {status}")
+                    st.markdown("</div>", unsafe_allow_html=True)
             else:
                 st.warning("AI models not initialized. Please use the Debug Connection section to test API connectivity.")
         else:
@@ -3363,6 +3047,4 @@ class PlotScene(Scene):
         st.session_state.first_load_complete = True
 
 if __name__ == "__main__":
-    main()
-
-
+    main()