Update app.py

app.py

@@ -101,6 +101,8 @@ def ensure_packages():
         'seaborn': '0.11.2',  # For enhanced visualizations
         'scipy': '1.7.3',    # For scientific computations
         'huggingface_hub': '0.16.0',  # For Hugging Face API
+        'azure-ai-inference': '1.0.0',
+        'azure-core': '1.29.5',
     }
     
     with st.spinner("Checking required packages..."):
@@ -239,7 +241,7 @@ def install_custom_packages(package_list):
 
 @st.cache_resource(ttl=3600)
 def init_ai_models(model_name=None):
-    try:        
+    try:
         # Get the GitHub token from secrets
         github_token = get_secret("github_token_api")
         if not github_token:
@@ -248,27 +250,23 @@ def init_ai_models(model_name=None):
         
         # Use default model if none specified
         if not model_name:
-            model_name = "deepseek-ai/deepseek-coder-1.3b-base"
-            
-        with st.spinner(f"Loading AI model: {model_name} - this may take a moment..."):
-            # Configure the Hugging Face API with the token
-            from huggingface_hub import login
-            login(token=github_token)
-            
-            # Now initialize the model with the token
-            code_model = pipeline(
-                "text-generation",
-                model=model_name,
-                token=github_token,  # Pass the token for authentication
-                trust_remote_code=True
+            model_name = "gpt-4o"
+            
+        with st.spinner(f"Initializing AI model: {model_name}..."):
+            # Import Azure modules
+            from azure.ai.inference import ChatCompletionsClient
+            from azure.ai.inference.models import SystemMessage, UserMessage
+            from azure.core.credentials import AzureKeyCredential
+            
+            # Setup Azure Inference client
+            endpoint = "https://models.inference.ai.azure.com"
+            client = ChatCompletionsClient(
+                endpoint=endpoint,
+                credential=AzureKeyCredential(github_token),
             )
             
-            if code_model is None:
-                st.error("Failed to initialize code model")
-                return None
-                
             return {
-                "code_model": code_model,
+                "client": client,
                 "model_name": model_name,
                 "last_loaded": datetime.now().isoformat()
             }
@@ -277,6 +275,61 @@ def init_ai_models(model_name=None):
         logger.error(f"AI model initialization error: {str(e)}")
         return None
 
+def suggest_code_completion(code_snippet, models):
+    if not models or "client" not in models:
+        st.error("AI models not properly initialized")
+        return None
+        
+    try:
+        prompt = f"""Write a complete Manim animation scene based on this code or idea:
+{code_snippet}
+
+The code should be a complete, working Manim animation that includes:
+- Proper Scene class definition
+- Constructor with animations
+- Proper use of self.play() for animations
+- Proper wait times between animations
+
+Here's the complete Manim code:
+"""
+        with st.spinner("AI is generating your animation code..."):
+            from azure.ai.inference.models import UserMessage
+            
+            # Generate the response using Azure Inference API
+            response = models["client"].complete(
+                messages=[
+                    UserMessage(prompt),
+                ],
+                max_tokens=1024,
+                model=models["model_name"]
+            )
+            
+            # Extract the generated code from the response
+            if response and response.choices and len(response.choices) > 0:
+                completed_code = response.choices[0].message.content
+            else:
+                st.error("No valid completion generated")
+                return None
+            
+            # Extract code if wrapped in markdown code blocks
+            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]
+        
+        if "Scene" not in completed_code:
+            completed_code = f"""from manim import *
+
+class MyScene(Scene):
+    def construct(self):
+        {completed_code}"""
+        
+        return completed_code
+    except Exception as e:
+        st.error(f"Error suggesting code: {str(e)}")
+        logger.error(f"Code suggestion error: {str(e)}")
+        return None
+
 def check_model_freshness():
     """Check if models need to be reloaded based on TTL"""
     if 'ai_models' not in st.session_state or st.session_state.ai_models is None:
@@ -1550,7 +1603,9 @@ def main():
             "animation_speed": "Normal"
         }
         st.session_state.password_entered = False  # Track password authentication
-        st.session_state.custom_model = "deepseek-ai/deepseek-coder-1.3b-base"  # Default model
+        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
 
     # Page configuration with improved layout
     st.set_page_config(
@@ -1773,7 +1828,7 @@ def main():
                     language="python",
                     theme="monokai",
                     min_lines=20,
-                    key="ace_editor"
+                    key=f"ace_editor_{st.session_state.editor_key}"
                 )
             else:
                 current_code = st.session_state.code if hasattr(st.session_state, 'code') and st.session_state.code else ""
@@ -1781,7 +1836,7 @@ def main():
                     "Manim Python Code",
                     value=current_code,
                     height=400,
-                    key="code_textarea"
+                    key=f"code_textarea_{st.session_state.editor_key}"
                 )
             
             # Update code in session state if it changed
@@ -1923,14 +1978,14 @@ class MyScene(Scene):
             # Model selection - only visible after authentication
             st.markdown("#### Model Selection")
             
-            # Predefined popular models
+            # Predefined Azure models
             popular_models = [
                 "Select a predefined model...",
-                "deepseek-ai/deepseek-coder-1.3b-base",
-                "mistralai/Mistral-7B-Instruct-v0.2",
-                "Phind/Phind-CodeLlama-34B-v2",
-                "facebook/incoder-1B", 
-                "Salesforce/codegen2-1B",
+                "gpt-4o",
+                "gpt-4",
+                "gpt-35-turbo",
+                "claude-3-opus-20240229",
+                "claude-3-sonnet-20240229",
             ]
             
             custom_model_option = st.radio(
@@ -1948,9 +2003,9 @@ class MyScene(Scene):
                     st.session_state.custom_model = selected_model
             else:
                 custom_model_name = st.text_input(
-                    "Enter custom model name/path:", 
+                    "Enter custom model name:", 
                     value=st.session_state.custom_model,
-                    placeholder="e.g., organization/model-name"
+                    placeholder="e.g., gpt-4o"
                 )
                 
                 if custom_model_name:
@@ -2027,8 +2082,8 @@ class MyScene(Scene):
                         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
-                            # Switch to Editor tab
-                            st.session_state.active_tab = 0
+                            # Set pending tab switch to editor tab
+                            st.session_state.pending_tab_switch = 0
                             st.rerun()
                     
                     with col_ai2:
@@ -2125,7 +2180,7 @@ class MyScene(Scene):
                         if st.button(formula["name"], key=f"latex_{formula['name']}"):
                             # Insert formula into the text area
                             st.session_state.latex_formula = formula["latex"]
-                            # Rerun to update the text area and preview
+                            # Refresh without full page rerun
                             st.rerun()
             
             # LaTeX code snippet
@@ -2168,8 +2223,8 @@ self.wait(2)
                                 st.session_state.temp_code = st.session_state.code
                                 st.success("LaTeX formula inserted into the editor!")
                                 
-                                # Switch to Editor tab
-                                st.session_state.active_tab = 0
+                                # Set pending tab switch to editor tab
+                                st.session_state.pending_tab_switch = 0
                                 st.rerun()
                         else:
                             st.warning("Could not find 'construct' method in your code. Please add a scene class first.")
@@ -2188,8 +2243,8 @@ class LatexScene(Scene):
                         st.session_state.temp_code = basic_scene
                         st.success("Created new scene with your LaTeX formula!")
                         
-                        # Switch to Editor tab
-                        st.session_state.active_tab = 0
+                        # Set pending tab switch to editor tab
+                        st.session_state.pending_tab_switch = 0
                         st.rerun()
         
         with col_latex2:
@@ -2259,7 +2314,6 @@ class LatexScene(Scene):
                     if "image_paths" not in st.session_state:
                         st.session_state.image_paths = []
                     
-                    # Check if this image was already added
                     # Check if this image was already added
                     image_already_added = False
                     for img in st.session_state.image_paths:
@@ -2305,8 +2359,8 @@ class ImageScene(Scene):
                                 st.session_state.temp_code = st.session_state.code
                                 st.success(f"Added {img_info['name']} to your code!")
                                 
-                                # Switch to Editor tab
-                                st.session_state.active_tab = 0
+                                # Set pending tab switch to editor tab
+                                st.session_state.pending_tab_switch = 0
                                 st.rerun()
                         except Exception as e:
                             st.error(f"Error loading image {img_info['name']}: {e}")
@@ -2708,14 +2762,14 @@ plt.legend()
                 language="python",
                 theme="monokai",
                 min_lines=15,
-                key="python_editor"
+                key=f"python_editor_{st.session_state.editor_key}"
             )
         else:
             python_code = st.text_area(
                 "Python Code",
                 value=python_code,
                 height=400,
-                key="python_textarea"
+                key=f"python_textarea_{st.session_state.editor_key}"
             )
         
         # Store script in session state (without clearing existing code)
@@ -2784,8 +2838,8 @@ self.wait(1)
                                     st.session_state.code += "\n" + plot_code
                                     st.session_state.temp_code = st.session_state.code
                                     st.success(f"Plot {i+1} added to your animation code!")
-                                    # Switch to Editor tab
-                                    st.session_state.active_tab = 0
+                                    # Set pending tab switch to editor tab
+                                    st.session_state.pending_tab_switch = 0
                                     st.rerun()
                                 else:
                                     basic_scene = f"""from manim import *
@@ -2797,8 +2851,8 @@ class PlotScene(Scene):
                                     st.session_state.code = basic_scene
                                     st.session_state.temp_code = basic_scene
                                     st.success(f"Created new scene with Plot {i+1}!")
-                                    # Switch to Editor tab
-                                    st.session_state.active_tab = 0
+                                    # Set pending tab switch to editor tab
+                                    st.session_state.pending_tab_switch = 0
                                     st.rerun()
             
             # Provide option to save the script
@@ -2875,12 +2929,19 @@ class PlotScene(Scene):
         ```
         """)
     
-    # Restore active tab if specified
-    if st.session_state.active_tab is not None and st.session_state.active_tab < len(tabs):
+    # Handle tab switching with session state to prevent refresh loop
+    if st.session_state.pending_tab_switch is not None:
+        st.session_state.active_tab = st.session_state.pending_tab_switch
+        st.session_state.pending_tab_switch = None
+        
+        # Set tabs active state
         for i, tab in enumerate(tabs):
             if i == st.session_state.active_tab:
                 tab.active = True
-                break
+    
+    # Mark first load as complete to prevent unnecessary refreshes
+    if not st.session_state.first_load_complete:
+        st.session_state.first_load_complete = True
 
 if __name__ == "__main__":
-    main()
+    main()