Update app.py

app.py

@@ -41,6 +41,29 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+# Model configuration mapping for different API requirements and limits
+MODEL_CONFIGS = {
+    "DeepSeek-V3-0324": {"max_tokens": 4000, "param_name": "max_tokens", "api_version": None, "category": "DeepSeek"},
+    "DeepSeek-R1": {"max_tokens": 4000, "param_name": "max_tokens", "api_version": None, "category": "DeepSeek"},
+    "Meta-Llama-3.1-405B-Instruct": {"max_tokens": 4000, "param_name": "max_tokens", "api_version": None, "category": "Meta"},
+    "Llama-3.2-90B-Vision-Instruct": {"max_tokens": 4000, "param_name": "max_tokens", "api_version": None, "category": "Meta"},
+    "Llama-3.3-70B-Instruct": {"max_tokens": 4000, "param_name": "max_tokens", "api_version": None, "category": "Meta"},
+    "Llama-4-Scout-17B-16E-Instruct": {"max_tokens": 4000, "param_name": "max_tokens", "api_version": None, "category": "Meta"},
+    "Llama-4-Maverick-17B-128E-Instruct-FP8": {"max_tokens": 4000, "param_name": "max_tokens", "api_version": None, "category": "Meta"},
+    "gpt-4o-mini": {"max_tokens": 16000, "param_name": "max_tokens", "api_version": None, "category": "OpenAI"},
+    "gpt-4.1": {"max_tokens": 16000, "param_name": "max_tokens", "api_version": None, "category": "OpenAI"},
+    "gpt-4o": {"max_tokens": 16000, "param_name": "max_tokens", "api_version": None, "category": "OpenAI"},
+    "o3-mini": {"max_tokens": 16000, "param_name": "max_tokens", "api_version": "2024-12-01-preview", "category": "Anthropic"},
+    "o1": {"max_tokens": 16000, "param_name": "max_completion_tokens", "api_version": "2024-12-01-preview", "category": "Anthropic"},
+    "o1-mini": {"max_tokens": 16000, "param_name": "max_completion_tokens", "api_version": "2024-12-01-preview", "category": "Anthropic"},
+    "o1-preview": {"max_tokens": 4000, "param_name": "max_tokens", "api_version": None, "category": "Anthropic"},
+    "Phi-4-multimodal-instruct": {"max_tokens": 4000, "param_name": "max_tokens", "api_version": None, "category": "Microsoft"},
+    "Mistral-large-2407": {"max_tokens": 4000, "param_name": "max_tokens", "api_version": None, "category": "Mistral"},
+    "Codestral-2501": {"max_tokens": 4000, "param_name": "max_tokens", "api_version": None, "category": "Mistral"},
+    # Default configuration for other models
+    "default": {"max_tokens": 4000, "param_name": "max_tokens", "api_version": None, "category": "Other"}
+}
+
 # Try to import Streamlit Ace
 try:
     from streamlit_ace import st_ace
@@ -306,16 +329,50 @@ Here's the complete Manim code:
         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"]
-            )
+            # 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"])
+            
+            # Prepare API call parameters based on model requirements
+            api_params = {
+                "messages": [UserMessage(prompt)],
+                "model": model_name
+            }
             
-            # Process the response exactly like in your example
+            # Add the appropriate token parameter
+            api_params[config["param_name"]] = config["max_tokens"]
+            
+            # Check if we need to specify API version
+            if config["api_version"]:
+                # If we need a specific API version, we need to create a new client with that version
+                logger.info(f"Using API version {config['api_version']} for model {model_name}")
+                
+                # Get token from session state
+                token = get_secret("github_token_api")
+                if not token:
+                    st.error("GitHub token not found in secrets")
+                    return None
+                
+                # Import required modules for creating client with specific API version
+                from azure.ai.inference import ChatCompletionsClient
+                from azure.core.credentials import AzureKeyCredential
+                
+                # Create client with specific API version
+                version_specific_client = ChatCompletionsClient(
+                    endpoint=models["endpoint"],
+                    credential=AzureKeyCredential(token),
+                    api_version=config["api_version"]
+                )
+                
+                # Make the API call with the version-specific client
+                response = version_specific_client.complete(**api_params)
+            else:
+                # Use the existing client
+                response = models["client"].complete(**api_params)
+            
+            # Process the response
             completed_code = response.choices[0].message.content
             
             # Process the code
@@ -338,6 +395,7 @@ class MyScene(Scene):
         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:
@@ -1635,21 +1693,86 @@ def main():
         margin-bottom: 1rem;
         text-align: center;
     }
-    
+
+    /* Improved Cards */
     .card {
+        background-color: #ffffff;
+        border-radius: 12px;
+        padding: 1.8rem;
+        box-shadow: 0 6px 12px rgba(0, 0, 0, 0.08);
+        margin-bottom: 1.8rem;
+        border-left: 5px solid #4F46E5;
+        transition: all 0.3s ease;
+    }
+    .card:hover {
+        box-shadow: 0 8px 16px rgba(0, 0, 0, 0.12);
+        transform: translateY(-2px);
+    }
+
+    /* Tab styling */
+    .stTabs [data-baseweb="tab-list"] {
+        gap: 2px;
+    }
+    .stTabs [data-baseweb="tab"] {
+        height: 45px;
+        white-space: pre-wrap;
+        border-radius: 4px 4px 0 0;
+        font-weight: 500;
+    }
+    .stTabs [aria-selected="true"] {
+        background-color: #f0f4fd;
+        border-bottom: 2px solid #4F46E5;
+    }
+
+    /* Buttons */
+    .stButton button {
+        border-radius: 6px;
+        font-weight: 500;
+        transition: all 0.2s ease;
+    }
+    .stButton button:hover {
+        transform: translateY(-1px);
+        box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
+    }
+
+    /* Model selection */
+    .model-group {
+        margin-bottom: 1.5rem;
+        padding: 15px;
+        border-radius: 8px;
+        background-color: #f8f9fa;
+    }
+
+    .model-card {
         background-color: #f8f9fa;
         border-radius: 10px;
-        padding: 1.5rem;
-        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
-        margin-bottom: 1.5rem;
+        padding: 15px;
+        margin-bottom: 10px;
+        border-left: 4px solid #4F46E5;
+        transition: all 0.3s ease;
     }
-    
-    .stButton > button {
-        width: 100%;
-        border-radius: 5px;
-        height: 2.5rem;
+    .model-card:hover {
+        box-shadow: 0 4px 8px rgba(0,0,0,0.1);
+        transform: translateY(-2px);
     }
-    
+    .model-category {
+        font-size: 1.2rem;
+        font-weight: 600;
+        padding: 10px 5px;
+        margin-top: 15px;
+        border-bottom: 2px solid #e9ecef;
+        color: #333;
+    }
+    .model-details {
+        font-size: 0.8rem;
+        color: #666;
+        margin-top: 5px;
+    }
+    .selected-model {
+        background-color: #e8f4fe;
+        border-left: 4px solid #0d6efd;
+    }
+
     .preview-container {
         border: 1px solid #e0e0e0;
         border-radius: 10px;
@@ -1657,7 +1780,7 @@ def main():
         margin-bottom: 1rem;
         min-height: 200px;
     }
-    
+
     .latex-preview {
         background-color: #f8f9fa;
         border-radius: 5px;
@@ -1665,12 +1788,12 @@ def main():
         margin-top: 0.5rem;
         min-height: 100px;
     }
-    
+
     .small-text {
         font-size: 0.8rem;
         color: #6c757d;
     }
-    
+
     .asset-card {
         background-color: #f0f2f5;
         border-radius: 8px;
@@ -1678,14 +1801,14 @@ def main():
         margin-bottom: 1rem;
         border-left: 4px solid #4F46E5;
     }
-    
+
     .timeline-container {
         background-color: #f8f9fa;
         border-radius: 10px;
         padding: 1.5rem;
         margin-bottom: 1.5rem;
     }
-    
+
     .keyframe {
         width: 12px;
         height: 12px;
@@ -1695,7 +1818,7 @@ def main():
         transform: translate(-50%, -50%);
         cursor: pointer;
     }
-    
+
     .educational-export-container {
         background-color: #f0f7ff;
         border-radius: 10px;
@@ -1703,7 +1826,7 @@ def main():
         margin-bottom: 1.5rem;
         border: 1px solid #c2e0ff;
     }
-    
+
     .code-output {
         background-color: #f8f9fa;
         border-radius: 8px;
@@ -1713,7 +1836,7 @@ def main():
         max-height: 400px;
         overflow-y: auto;
     }
-    
+
     .error-output {
         background-color: #fef2f2;
         border-radius: 8px;
@@ -2004,22 +2127,36 @@ class MyScene(Scene):
                             
                             # Define endpoint
                             endpoint = "https://models.inference.ai.azure.com"
-                            model_name = "gpt-4o"
+                            model_name = st.session_state.custom_model
                             
-                            # Create client directly following example
-                            client = ChatCompletionsClient(
-                                endpoint=endpoint,
-                                credential=AzureKeyCredential(token),
-                            )
+                            # Get model configuration
+                            config = MODEL_CONFIGS.get(model_name, MODEL_CONFIGS["default"])
+                            
+                            # Create client with appropriate API version
+                            api_version = config.get("api_version")
+                            if api_version:
+                                client = ChatCompletionsClient(
+                                    endpoint=endpoint,
+                                    credential=AzureKeyCredential(token),
+                                    api_version=api_version
+                                )
+                            else:
+                                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=1000,
-                                model=model_name
-                            )
+                            api_params = {
+                                "messages": [UserMessage("Hello, this is a connection test.")],
+                                "model": model_name
+                            }
+                            
+                            # Use appropriate parameter name
+                            api_params[config["param_name"]] = 100  # Just enough for a short response
+                            
+                            # Make the API call
+                            response = client.complete(**api_params)
                             
                             # Check if response is valid
                             if response and response.choices and len(response.choices) > 0:
@@ -2043,44 +2180,87 @@ class MyScene(Scene):
                             import traceback
                             st.code(traceback.format_exc())
             
-            # Model selection
-            st.markdown("#### Model Selection")
-            
-            # Predefined Azure models
-            popular_models = [
-                "DeepSeek-V3-0324",
-                "DeepSeek-R1",
-                "Meta-Llama-3.1-405B-Instruct",
-                "Llama-3.2-90B-Vision-Instruct",
-                "Llama-3.3-70B-Instruct"
-                "Llama-4-Scout-17B-16E-Instruct",
-                "Llama-4-Maverick-17B-128E-Instruct-FP8",
-                "gpt-4o-mini",
-                "gpt-4o",
-                "gpt-4.1",
-                "o3-mini",
-                "o1",
-                "o1-mini",
-                "o1-preview",
-                "Phi-4-multimodal-instruct",
-                "Mistral-large-2407",
-                "Codestral-2501",
-            ]
-            
-            selected_model = st.selectbox(
-                "Select a model:", 
-                options=popular_models,
-                index=0
-            )
-            
-            st.session_state.custom_model = selected_model
-            st.info(f"Currently selected model: {st.session_state.custom_model}")
-            
-            # 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}")
+            # Model selection with enhanced UI
+            st.markdown("### 🤖 Model Selection")
+            st.markdown("Select an AI model for generating animation code:")
+
+            # Group models by category for better organization
+            model_categories = {}
+            for model_name in MODEL_CONFIGS:
+                if model_name != "default":
+                    category = MODEL_CONFIGS[model_name].get("category", "Other")
+                    if category not in model_categories:
+                        model_categories[category] = []
+                    model_categories[category].append(model_name)
+
+            # Create tabbed interface for model categories
+            category_tabs = st.tabs(sorted(model_categories.keys()))
+
+            for i, category in enumerate(sorted(model_categories.keys())):
+                with category_tabs[i]:
+                    for model_name in sorted(model_categories[category]):
+                        config = MODEL_CONFIGS[model_name]
+                        is_selected = model_name == st.session_state.custom_model
+                        
+                        # Create styled card for each model
+                        st.markdown(f"""
+                        <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>API Version: {config['api_version'] if config['api_version'] else 'Default'}</p>
+                            </div>
+                        </div>
+                        """, unsafe_allow_html=True)
+                        
+                        # Button to select this model
+                        button_label = "Selected ✓" if is_selected else "Select Model"
+                        if st.button(button_label, key=f"model_{model_name}", disabled=is_selected):
+                            st.session_state.custom_model = model_name
+                            if st.session_state.ai_models and 'model_name' in st.session_state.ai_models:
+                                st.session_state.ai_models['model_name'] = model_name
+                            st.rerun()
+
+            # Display current model selection
+            st.info(f"🤖 **Currently using: {st.session_state.custom_model}**")
+
+            # Add a refresh button to update model connection
+            if st.button("🔄 Refresh Model Connection", key="refresh_model_connection"):
+                if st.session_state.ai_models and 'client' in st.session_state.ai_models:
+                    try:
+                        # 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"])
+                        
+                        # Use appropriate parameters based on model configuration
+                        api_params = {
+                            "messages": [UserMessage("Hello")],
+                            "model": model_name
+                        }
+                        api_params[config["param_name"]] = 10  # Just request 10 tokens for quick test
+                        
+                        if config["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)
+                        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
+                        
+                    except Exception as e:
+                        st.error(f"❌ Connection error: {str(e)}")
+                        st.info("Please try the Debug Connection section to re-initialize the API connection.")
             
             # AI code generation
             if st.session_state.ai_models and "client" in st.session_state.ai_models:
@@ -2120,6 +2300,9 @@ 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"])
+                                
                                 # Create the prompt
                                 prompt = f"""Write a complete Manim animation scene based on this code or idea:
 {code_input}
@@ -2133,15 +2316,42 @@ The code should be a complete, working Manim animation that includes:
 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=100000,
-                                    model=model_name
-                                )
+                                # Prepare API call parameters based on model requirements
+                                api_params = {
+                                    "messages": [UserMessage(prompt)],
+                                    "model": model_name
+                                }
+                                
+                                # Add the appropriate token parameter
+                                api_params[config["param_name"]] = config["max_tokens"]
+                                
+                                # Check if we need to specify API version
+                                if config["api_version"]:
+                                    # If we need a specific API version, create a new client with that version
+                                    logger.info(f"Using API version {config['api_version']} for model {model_name}")
+                                    
+                                    # Get token from session state
+                                    token = get_secret("github_token_api")
+                                    if not token:
+                                        st.error("GitHub token not found in secrets")
+                                        return None
+                                    
+                                    # Import required modules for creating client with specific API version
+                                    from azure.ai.inference import ChatCompletionsClient
+                                    from azure.core.credentials import AzureKeyCredential
+                                    
+                                    # Create client with specific API version
+                                    version_specific_client = ChatCompletionsClient(
+                                        endpoint=st.session_state.ai_models["endpoint"],
+                                        credential=AzureKeyCredential(token),
+                                        api_version=config["api_version"]
+                                    )
+                                    
+                                    # Make the API call with the version-specific client
+                                    response = version_specific_client.complete(**api_params)
+                                else:
+                                    # Use the existing client
+                                    response = client.complete(**api_params)
                                 
                                 # Process the response
                                 if response and response.choices and len(response.choices) > 0: