Update app.py

app.py

@@ -82,14 +82,12 @@ def prepare_api_params(messages, model_name):
         "model": model_name
     }
     
-    # Add the appropriate token parameter based on model
-    param_name = config["param_name"]
-    if param_name == "max_completion_tokens":
-        # Use the max_completion_tokens value from config
-        api_params["max_completion_tokens"] = config["max_completion_tokens"]
-    else:
-        # Use the max_tokens value for other models
-        api_params["max_tokens"] = config["max_tokens"]
+    # Add the appropriate token parameter based on model's parameter name
+    token_param = config["param_name"]
+    token_value = config[token_param]  # Get the actual value from the config
+    
+    # Add the parameter to the API params
+    api_params[token_param] = token_value
     
     return api_params, config
 
@@ -329,8 +327,8 @@ def init_ai_models_direct():
         
 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. Please use the Debug Connection section to test API connectivity.")
+    if not models:
+        st.error("AI models not properly initialized.")
         return None
     
     try:
@@ -348,45 +346,61 @@ Here's the complete Manim code:
 """
         
         with st.spinner("AI is generating your animation code..."):
-            from azure.ai.inference.models import UserMessage
-            
-            # Get the current model name
+            # Get the current model name and base URL
             model_name = models["model_name"]
             
-            # Prepare API parameters based on model
-            messages = [UserMessage(prompt)]
-            api_params, config = prepare_api_params(messages, model_name)
+            # Convert message to the appropriate format based on model category
+            config = MODEL_CONFIGS.get(model_name, MODEL_CONFIGS["default"])
+            category = config.get("category", "Other")
             
-            # 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}")
+            if category == "OpenAI":
+                # Import OpenAI client
+                from openai import OpenAI
                 
-                # Get token from session state
+                # Get token
                 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)
+                # Create or get client
+                if "openai_client" not in models:
+                    client = OpenAI(
+                        base_url="https://models.github.ai/inference",
+                        api_key=token
+                    )
+                    models["openai_client"] = client
+                else:
+                    client = models["openai_client"]
+                
+                # For OpenAI models, we need role-based messages
+                messages = [
+                    {"role": "system", "content": "You are an expert in Manim animations."},
+                    {"role": "user", "content": prompt}
+                ]
+                
+                # Create params
+                params = {
+                    "messages": messages,
+                    "model": model_name
+                }
+                
+                # Add token parameter
+                token_param = config["param_name"]
+                params[token_param] = config[token_param]
+                
+                # Make API call
+                response = client.chat.completions.create(**params)
+                completed_code = response.choices[0].message.content
+                
             else:
-                # Use the existing client
+                # Use Azure client
+                from azure.ai.inference.models import UserMessage
+                
+                # Convert message format for Azure
+                messages = [UserMessage(prompt)]
+                api_params, _ = prepare_api_params(messages, model_name)
+                
+                # Make API call with Azure client
                 response = models["client"].complete(**api_params)
-            
-            # Process the response
-            completed_code = response.choices[0].message.content
+                completed_code = response.choices[0].message.content
             
             # Process the code
             if "```python" in completed_code:
@@ -2132,54 +2146,114 @@ class MyScene(Scene):
                                 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"
+                            # Get model details
                             model_name = st.session_state.custom_model
+                            config = MODEL_CONFIGS.get(model_name, MODEL_CONFIGS["default"])
+                            category = config.get("category", "Other")
                             
-                            # 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")
-                            if api_version:
-                                client = ChatCompletionsClient(
-                                    endpoint=endpoint,
-                                    credential=AzureKeyCredential(token),
-                                    api_version=api_version
-                                )
-                            else:
-                                client = ChatCompletionsClient(
-                                    endpoint=endpoint,
-                                    credential=AzureKeyCredential(token),
+                            if category == "OpenAI":
+                                # Use OpenAI client for GitHub AI models
+                                try:
+                                    from openai import OpenAI
+                                except ImportError:
+                                    st.error("OpenAI package not installed. Please run 'pip install openai'")
+                                    st.stop()
+                                
+                                # Create OpenAI client with GitHub AI endpoint
+                                client = OpenAI(
+                                    base_url="https://models.github.ai/inference",
+                                    api_key=token,
                                 )
-                            
-                            # 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:
-                                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()
+                                # Prepare parameters based on model configuration
+                                params = {
+                                    "messages": [
+                                        {"role": "system", "content": "You are a helpful assistant."},
+                                        {"role": "user", "content": "Hello, this is a connection test."}
+                                    ],
+                                    "model": model_name
                                 }
+                                
+                                # Add appropriate token parameter
+                                token_param = config["param_name"]
+                                params[token_param] = config[token_param]
+                                
+                                # Make API call
+                                response = client.chat.completions.create(**params)
+                                
+                                # 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 = {
+                                        "openai_client": client,
+                                        "model_name": model_name,
+                                        "endpoint": "https://models.github.ai/inference",
+                                        "last_loaded": datetime.now().isoformat(),
+                                        "category": category
+                                    }
+                                else:
+                                    st.error("❌ API returned an empty response")
+                            
+                            elif category == "Azure" or category in ["DeepSeek", "Meta", "Microsoft", "Mistral", "Other"]:
+                                # Use Azure client for Azure API models
+                                try:
+                                    from azure.ai.inference import ChatCompletionsClient
+                                    from azure.ai.inference.models import SystemMessage, UserMessage
+                                    from azure.core.credentials import AzureKeyCredential
+                                except ImportError:
+                                    st.error("Azure AI packages not installed. Please run 'pip install azure-ai-inference azure-core'")
+                                    st.stop()
+                                
+                                # Define endpoint
+                                endpoint = "https://models.inference.ai.azure.com"
+                                
+                                # 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")
+                                if api_version:
+                                    client = ChatCompletionsClient(
+                                        endpoint=endpoint,
+                                        credential=AzureKeyCredential(token),
+                                        api_version=api_version
+                                    )
+                                else:
+                                    client = ChatCompletionsClient(
+                                        endpoint=endpoint,
+                                        credential=AzureKeyCredential(token),
+                                    )
+                                
+                                # 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:
+                                    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(),
+                                        "category": category,
+                                        "api_version": api_version
+                                    }
+                                else:
+                                    st.error("❌ API returned an empty response")
+                            
                             else:
-                                st.error("❌ API returned an empty response")
+                                st.error(f"Unsupported model category: {category}")
+                        
                         except ImportError as ie:
                             st.error(f"Module import error: {str(ie)}")
-                            st.info("Try installing required packages: azure-ai-inference and azure-core")
+                            st.info("Try installing required packages: openai, azure-ai-inference and azure-core")
                         except Exception as e:
                             st.error(f"❌ API test failed: {str(e)}")
                             import traceback
@@ -2215,12 +2289,13 @@ 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.get('max_tokens', config.get('max_completion_tokens', 'Unknown')):,}</p>
+                                <p>Max Tokens: {config.get(config['param_name'], 'Unknown')}</p>
+                                <p>Category: {config['category']}</p>
                                 <p>API Version: {config['api_version'] if config['api_version'] else 'Default'}</p>
                                 {warning_html}
                             </div>
                         </div>
-                        """, unsafe_allow_html=True)
+                    """, unsafe_allow_html=True)
                         
                         # Button to select this model
                         button_label = "Selected ✓" if is_selected else "Select Model"