Update app.py

app.py

@@ -71,6 +71,25 @@ except ImportError:
     ACE_EDITOR_AVAILABLE = False
     logger.warning("streamlit-ace not available, falling back to standard text editor")
 
+def prepare_api_params(messages, model_name):
+    """Create appropriate API parameters based on model configuration"""
+    # Get model configuration
+    config = MODEL_CONFIGS.get(model_name, MODEL_CONFIGS["default"])
+    
+    # Base parameters common to all models
+    api_params = {
+        "messages": messages,
+        "model": model_name
+    }
+    
+    # Add the appropriate token parameter based on model
+    if config["param_name"] == "max_completion_tokens":
+        api_params["max_completion_tokens"] = config["max_tokens"]
+    else:
+        api_params["max_tokens"] = config["max_tokens"]
+    
+    return api_params, config
+
 # New functions for accessing secrets and password verification
 def get_secret(github_token_api):
     """Retrieve a secret from HuggingFace Spaces environment variables"""
@@ -331,12 +350,9 @@ Here's the complete Manim code:
             # Get the current model name
             model_name = models["model_name"]
             
-            # Get configuration for this model (or use default)
-            config = MODEL_CONFIGS.get(model_name, MODEL_CONFIGS["default"])
-            
-            # Check if this is a model that requires max_completion_tokens
-            # Prepare common API parameters
+            # Prepare API parameters based on model
             messages = [UserMessage(prompt)]
+            api_params, config = prepare_api_params(messages, model_name)
             
             # Check if we need to specify API version
             if config["api_version"]:
@@ -360,21 +376,11 @@ Here's the complete Manim code:
                     api_version=config["api_version"]
                 )
                 
-                # Make the API call with the version-specific client and normal max_tokens
-                response = version_specific_client.complete(
-                    messages=messages,
-                    model=model_name,
-                    max_tokens=config["max_tokens"],
-                    max_completion_tokens=config["max_completion_tokens"]
-                )
+                # Make the API call with the version-specific client
+                response = version_specific_client.complete(**api_params)
             else:
-                # Use the existing client with normal max_tokens
-                response = models["client"].complete(
-                    messages=messages,
-                    model=model_name,
-                    max_tokens=config["max_tokens"],
-                    max_completion_tokens=config["max_completion_tokens"]
-                )
+                # Use the existing client
+                response = models["client"].complete(**api_params)
             
             # Process the response
             completed_code = response.choices[0].message.content
@@ -2133,14 +2139,9 @@ class MyScene(Scene):
                             endpoint = "https://models.inference.ai.azure.com"
                             model_name = st.session_state.custom_model
                             
-                            # Get model configuration
-                            config = MODEL_CONFIGS.get(model_name, MODEL_CONFIGS["default"])
-                            
-                            # Check if this is a model that requires max_completion_tokens
-                            if config["param_name"] == "max_com343434pletion_tokens":
-                                st.warning(f"Model {model_name} requires special handling that isn't fully supported. Testing with gpt-4o instead.")
-                                model_name = "gpt-4o"
-                                config = MODEL_CONFIGS["gpt-4o"]
+                            # Prepare API parameters 
+                            messages = [UserMessage("Hello, this is a connection test.")]
+                            api_params, config = prepare_api_params(messages, model_name)
                             
                             # Create client with appropriate API version
                             api_version = config.get("api_version")
@@ -2156,13 +2157,8 @@ class MyScene(Scene):
                                     credential=AzureKeyCredential(token),
                                 )
                             
-                            # Test with a simple prompt
-                            response = client.complete(
-                                messages=[UserMessage("Hello, this is a connection test.")],
-                                model=model_name,
-                                max_tokens=50 , # Small value for quick testing
-                                max_completion_tokens=50
-                            )
+                            # Test with the prepared parameters
+                            response = client.complete(**api_params)
                             
                             # Check if response is valid
                             if response and response.choices and len(response.choices) > 0:
@@ -2185,7 +2181,7 @@ class MyScene(Scene):
                             st.error(f"❌ API test failed: {str(e)}")
                             import traceback
                             st.code(traceback.format_exc())
-            
+                            
             # Model selection with enhanced UI
             st.markdown("### 🤖 Model Selection")
             st.markdown("Select an AI model for generating animation code:")
@@ -2216,7 +2212,7 @@ class MyScene(Scene):
                         <div class="model-card {'selected-model' if is_selected else ''}">
                             <h4>{model_name}</h4>
                             <div class="model-details">
-                                <p>Max Tokens: {config['max_tokens']:,}</p>
+                                <p>Max Tokens: {config['max_tokens']:,config[max]}</p>
                                 <p>API Version: {config['api_version'] if config['api_version'] else 'Default'}</p>
                                 {warning_html}
                             </div>
@@ -2241,21 +2237,31 @@ class MyScene(Scene):
                         # Test connection with minimal prompt
                         from azure.ai.inference.models import UserMessage
                         model_name = st.session_state.custom_model
-                        config = MODEL_CONFIGS.get(model_name, MODEL_CONFIGS["default"])
                         
-                        # Check if this is a model that requires max_completion_tokens
-                        if config["param_name"] == "max_completion_tokens":
-                            st.warning(f"Model {model_name} requires special handling that isn't fully supported. Using gpt-4o instead.")
-                            model_name = "gpt-4o"
-                            config = MODEL_CONFIGS["gpt-4o"]
+                        # Prepare parameters
+                        messages = [UserMessage("Hello")]
+                        api_params, config = prepare_api_params(messages, model_name)
                         
-                        # Use appropriate parameters based on model configuration
-                        response = st.session_state.ai_models["client"].complete(
-                            messages=[UserMessage("Hello")],
-                            model=model_name,
-                            max_tokens=10  # Just request 10 tokens for quick test
-                        )
-                                
+                        # Check if we need a new client with specific API version
+                        if config["api_version"] and config["api_version"] != st.session_state.ai_models.get("api_version"):
+                            # Create version-specific client if needed
+                            token = get_secret("github_token_api")
+                            from azure.ai.inference import ChatCompletionsClient
+                            from azure.core.credentials import AzureKeyCredential
+                            
+                            client = ChatCompletionsClient(
+                                endpoint=st.session_state.ai_models["endpoint"],
+                                credential=AzureKeyCredential(token),
+                                api_version=config["api_version"]
+                            )
+                            response = client.complete(**api_params)
+                            
+                            # Update session state with the new client
+                            st.session_state.ai_models["client"] = client
+                            st.session_state.ai_models["api_version"] = config["api_version"]
+                        else:
+                            response = st.session_state.ai_models["client"].complete(**api_params)
+                            
                         st.success(f"✅ Connection to {model_name} successful!")
                         st.session_state.ai_models["model_name"] = model_name
                         
@@ -2301,35 +2307,26 @@ class MyScene(Scene):
                                 client = st.session_state.ai_models["client"]
                                 model_name = st.session_state.ai_models["model_name"]
                                 
-                                # Get configuration for this model
-                                config = MODEL_CONFIGS.get(model_name, MODEL_CONFIGS["default"])
-                                
-                                # Check if this model requires max_completion_tokens
-                                if config["param_name"] == "max_completion_tokens":
-                                    st.warning(f"Model {model_name} requires special handling that isn't fully supported. Using gpt-4o instead.")
-                                    model_name = "gpt-4o"
-                                    config = MODEL_CONFIGS["gpt-4o"]
-                                
                                 # 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:
-"""
+                {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 the API call with proper parameters
+                                # Prepare API parameters
                                 from azure.ai.inference.models import UserMessage
-                                response = client.complete(
-                                    messages=[UserMessage(prompt)],
-                                    model=model_name,
-                                    max_tokens=config["max_tokens"]
-                                )
+                                messages = [UserMessage(prompt)]
+                                api_params, config = prepare_api_params(messages, model_name)
+                                
+                                # Make the API call with proper parameters
+                                response = client.complete(**api_params)
                                 
                                 # Process the response
                                 if response and response.choices and len(response.choices) > 0:
@@ -2344,10 +2341,10 @@ Here's the complete Manim code:
                                     # 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}"""
+                
+                class MyScene(Scene):
+                    def construct(self):
+                        {completed_code}"""
                                     
                                     # Store the generated code
                                     st.session_state.generated_code = completed_code
@@ -2359,8 +2356,7 @@ class MyScene(Scene):
                                 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: