Update app.py

app.py

@@ -238,51 +238,58 @@ def install_custom_packages(package_list):
     return success, "\n".join(results)
 
 @st.cache_resource(ttl=3600)
-def init_ai_models(model_name=None):
+def init_ai_models_direct():
+    """Direct implementation using the exact pattern from the example code"""
     try:
-        # Get the token directly from secrets
+        # Get token from secrets
         token = get_secret("github_token_api")
         if not token:
-            st.error("GitHub token not found in secrets")
+            st.error("GitHub token not found in secrets. Please add 'github_token_api' to your HuggingFace Spaces secrets.")
             return None
         
-        # Use default model or specified model
-        if not model_name:
-            model_name = "gpt-4o"
+        # Log what we're doing - for debugging
+        logger.info(f"Initializing AI model with token: {token[:5]}...")
         
-        # Import required modules
+        # Use exact imports as in your example
+        import os
         from azure.ai.inference import ChatCompletionsClient
         from azure.ai.inference.models import SystemMessage, UserMessage
         from azure.core.credentials import AzureKeyCredential
         
-        # Define endpoint
+        # Use exact endpoint as in your example
         endpoint = "https://models.inference.ai.azure.com"
         
+        # Use default model
+        model_name = "gpt-4o"
+        
         # Create client exactly as in your example
         client = ChatCompletionsClient(
             endpoint=endpoint,
             credential=AzureKeyCredential(token),
         )
         
-        # Store client and model info
+        # Return the necessary information
         return {
             "client": client,
             "model_name": model_name,
-            "endpoint": endpoint,
-            "token": token,
-            "last_loaded": datetime.now().isoformat()
+            "endpoint": endpoint
         }
+    except ImportError as ie:
+        st.error(f"Import error: {str(ie)}. Please make sure azure-ai-inference is installed.")
+        logger.error(f"Import error: {str(ie)}")
+        return None
     except Exception as e:
-        st.error(f"Failed to initialize AI model: {str(e)}")
+        st.error(f"Error initializing AI model: {str(e)}")
+        logger.error(f"Initialization error: {str(e)}")
         return None
+        
 def suggest_code_completion(code_snippet, models):
+    """Generate code completion using the AI model"""
     if not models or "client" not in models:
-        st.error("AI models not properly initialized")
+        st.error("AI models not properly initialized. Please use the Debug Connection section to test API connectivity.")
         return None
     
     try:
-        from azure.ai.inference.models import UserMessage
-        
         # Create the prompt
         prompt = f"""Write a complete Manim animation scene based on this code or idea:
 {code_snippet}
@@ -296,37 +303,41 @@ The code should be a complete, working Manim animation that includes:
 Here's the complete Manim code:
 """
         
-        # Make API call exactly like your example
-        response = models["client"].complete(
-            messages=[
-                UserMessage(prompt),
-            ],
-            max_tokens=1024,
-            model=models["model_name"]
-        )
-        
-        # Get response content exactly like your example
-        completed_code = response.choices[0].message.content
-        
-        # Clean up the code if it's wrapped in markdown 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]
-        
-        # Add Scene class if missing
-        if "Scene" not in completed_code:
-            completed_code = f"""from manim import *
+        with st.spinner("AI is generating your animation code..."):
+            from azure.ai.inference.models import UserMessage
+            
+            # Make an API call exactly like in your example
+            response = models["client"].complete(
+                messages=[
+                    UserMessage(prompt),
+                ],
+                max_tokens=1000,
+                model=models["model_name"]
+            )
+            
+            # Process the response exactly like in your example
+            completed_code = response.choices[0].message.content
+            
+            # Process the code
+            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}"""
-        
-        return completed_code
+            
+            return completed_code
+            
     except Exception as e:
         st.error(f"Error generating code: {str(e)}")
+        st.code(traceback.format_exc())
         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:
@@ -1972,12 +1983,71 @@ class MyScene(Scene):
         
         # Check password before allowing access
         if check_password():
-            # Model selection - only visible after authentication
+            # Debug section
+            with st.expander("🔧 Debug Connection"):
+                st.markdown("Test the AI model connection directly")
+                
+                if st.button("Test API Connection", key="test_api_btn"):
+                    with st.spinner("Testing API connection..."):
+                        try:
+                            # Get token from secrets
+                            token = get_secret("github_token_api")
+                            if not token:
+                                st.error("GitHub token not found in secrets")
+                                st.stop()
+                            
+                            # Import required modules
+                            import os
+                            from azure.ai.inference import ChatCompletionsClient
+                            from azure.ai.inference.models import SystemMessage, UserMessage
+                            from azure.core.credentials import AzureKeyCredential
+                            
+                            # Define endpoint
+                            endpoint = "https://models.inference.ai.azure.com"
+                            model_name = "gpt-4o"
+                            
+                            # Create client directly following example
+                            client = ChatCompletionsClient(
+                                endpoint=endpoint,
+                                credential=AzureKeyCredential(token),
+                            )
+                            
+                            # Test with a simple prompt
+                            response = client.complete(
+                                messages=[
+                                    UserMessage("Hello, this is a connection test."),
+                                ],
+                                max_tokens=100,
+                                model=model_name
+                            )
+                            
+                            # Check if response is valid
+                            if response and response.choices and len(response.choices) > 0:
+                                test_response = response.choices[0].message.content
+                                st.success(f"✅ Connection successful! Response: {test_response[:50]}...")
+                                
+                                # Save working connection to session state
+                                st.session_state.ai_models = {
+                                    "client": client,
+                                    "model_name": model_name,
+                                    "endpoint": endpoint,
+                                    "last_loaded": datetime.now().isoformat()
+                                }
+                            else:
+                                st.error("❌ API returned an empty response")
+                        except ImportError as ie:
+                            st.error(f"Module import error: {str(ie)}")
+                            st.info("Try installing required packages: azure-ai-inference and azure-core")
+                        except Exception as e:
+                            st.error(f"❌ API test failed: {str(e)}")
+                            import traceback
+                            st.code(traceback.format_exc())
+            
+            # Model selection
             st.markdown("#### Model Selection")
             
             # Predefined Azure models
             popular_models = [
-                "Select a predefined model...",
                 "gpt-4o",
                 "gpt-4",
                 "gpt-35-turbo",
@@ -1985,50 +2055,23 @@ class MyScene(Scene):
                 "claude-3-sonnet-20240229",
             ]
             
-            custom_model_option = st.radio(
-                "Choose model option:",
-                ["Use predefined model", "Enter custom model name"]
+            selected_model = st.selectbox(
+                "Select a model:", 
+                options=popular_models,
+                index=0
             )
             
-            if custom_model_option == "Use predefined model":
-                selected_model = st.selectbox(
-                    "Select a model:", 
-                    options=popular_models
-                )
-                
-                if selected_model != "Select a predefined model...":
-                    st.session_state.custom_model = selected_model
-            else:
-                custom_model_name = st.text_input(
-                    "Enter custom model name:", 
-                    value=st.session_state.custom_model,
-                    placeholder="e.g., gpt-4o"
-                )
-                
-                if custom_model_name:
-                    st.session_state.custom_model = custom_model_name
-            
-            # Display currently selected model
+            st.session_state.custom_model = selected_model
             st.info(f"Currently selected model: {st.session_state.custom_model}")
             
-            # Only initialize models if password check passes and we have a model name
-            if st.session_state.custom_model:
-                model_changed = False
-                
-                # Check if we need to reinitialize the model (if model changed)
-                if (st.session_state.ai_models is not None and 
-                    'model_name' in st.session_state.ai_models and
-                    st.session_state.ai_models['model_name'] != st.session_state.custom_model):
-                    model_changed = True
-                
-                # Initialize model if needed
-                if (st.session_state.ai_models is None or 
-                    not check_model_freshness() or
-                    model_changed):
-                    st.session_state.ai_models = init_ai_models(st.session_state.custom_model)
-            
-            if st.session_state.ai_models:
-                # AI code generation
+            # Update model if it changed
+            if st.session_state.ai_models and 'model_name' in st.session_state.ai_models:
+                if st.session_state.ai_models['model_name'] != st.session_state.custom_model:
+                    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
+            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.")
@@ -2060,16 +2103,67 @@ class MyScene(Scene):
                 if st.button("Generate Animation Code", key="gen_ai_code"):
                     if code_input:
                         with st.spinner("AI is generating your animation code..."):
-                            response = suggest_code_completion(code_input, st.session_state.ai_models)
-                            if response:
-                                st.session_state.generated_code = response
+                            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:
+- 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:
+"""
+                                
+                                # 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]
+                                    
+                                    # 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 st.session_state.generated_code:
+                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")
@@ -2105,7 +2199,7 @@ class MyScene(Scene):
                                     st.error(f"Failed to generate preview: {status}")
                     st.markdown("</div>", unsafe_allow_html=True)
             else:
-                st.warning("AI models failed to load. Generation features are unavailable.")
+                st.warning("AI models not initialized. Please use the Debug Connection section to test API connectivity.")
         else:
             st.info("Please enter the correct password to access AI features")