Update app.py

app.py

@@ -239,7 +239,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:
@@ -251,23 +251,46 @@ def init_ai_models(model_name=None):
             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),
-            )
-            
-            return {
-                "client": client,
-                "model_name": model_name,
-                "last_loaded": datetime.now().isoformat()
-            }
+            try:
+                # Import Azure modules here to prevent import errors
+                from azure.ai.inference import ChatCompletionsClient
+                from azure.ai.inference.models import UserMessage
+                from azure.core.credentials import AzureKeyCredential
+                
+                # Setup Azure Inference client with proper endpoint
+                endpoint = "https://models.inference.ai.azure.com"
+                
+                # Create the client with proper credential
+                client = ChatCompletionsClient(
+                    endpoint=endpoint,
+                    credential=AzureKeyCredential(github_token),
+                )
+                
+                # Test the client with a simple request to verify it works
+                try:
+                    test_response = client.complete(
+                        messages=[UserMessage("Hello")],
+                        max_tokens=10,
+                        model=model_name
+                    )
+                    if not test_response or not test_response.choices:
+                        st.error(f"API connection test failed for model {model_name}")
+                        return None
+                except Exception as api_err:
+                    st.error(f"API test failed: {str(api_err)}")
+                    logger.error(f"API test error: {str(api_err)}")
+                    return None
+                
+                logger.info(f"Successfully initialized AI model: {model_name}")
+                return {
+                    "client": client,
+                    "model_name": model_name,
+                    "last_loaded": datetime.now().isoformat()
+                }
+            except ImportError as ie:
+                st.error(f"Azure AI modules not available: {str(ie)}")
+                logger.error(f"Import error: {str(ie)}")
+                return None
     except Exception as e:
         st.error(f"Failed to load AI models: {str(e)}")
         logger.error(f"AI model initialization error: {str(e)}")
@@ -279,6 +302,8 @@ def suggest_code_completion(code_snippet, models):
         return None
         
     try:
+        from azure.ai.inference.models import UserMessage
+        
         prompt = f"""Write a complete Manim animation scene based on this code or idea:
 {code_snippet}
 
@@ -291,8 +316,6 @@ The code should be a complete, working Manim animation that includes:
 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=[