Update app.py

app.py

@@ -239,71 +239,51 @@ def install_custom_packages(package_list):
 
 @st.cache_resource(ttl=3600)
 def init_ai_models(model_name=None):
-    try:        
-        # Get the GitHub token from secrets
-        github_token = get_secret("github_token_api")
-        if not github_token:
-            st.error("GitHub API token not configured in HuggingFace Spaces secrets")
+    try:
+        # Get the token directly from secrets
+        token = get_secret("github_token_api")
+        if not token:
+            st.error("GitHub token not found in secrets")
             return None
         
-        # Use default model if none specified
+        # Use default model or specified model
         if not model_name:
             model_name = "gpt-4o"
-            
-        with st.spinner(f"Initializing AI model: {model_name}..."):
-            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
+        
+        # Import required modules
+        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"
+        
+        # Create client exactly as in your example
+        client = ChatCompletionsClient(
+            endpoint=endpoint,
+            credential=AzureKeyCredential(token),
+        )
+        
+        # Store client and model info
+        return {
+            "client": client,
+            "model_name": model_name,
+            "endpoint": endpoint,
+            "token": token,
+            "last_loaded": datetime.now().isoformat()
+        }
     except Exception as e:
-        st.error(f"Failed to load AI models: {str(e)}")
-        logger.error(f"AI model initialization error: {str(e)}")
+        st.error(f"Failed to initialize AI model: {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:
         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}
 
@@ -315,29 +295,26 @@ 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..."):
-            # 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]
         
+        # 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 *
 
@@ -347,8 +324,7 @@ class MyScene(Scene):
         
         return completed_code
     except Exception as e:
-        st.error(f"Error suggesting code: {str(e)}")
-        logger.error(f"Code suggestion error: {str(e)}")
+        st.error(f"Error generating code: {str(e)}")
         return None
 
 def check_model_freshness():