Update app.py

app.py

@@ -349,8 +349,11 @@ Here's the complete Manim code:
             # Get the current model name and base URL
             model_name = models["model_name"]
             
+            # Handle model name - extract base name if it has a prefix
+            base_model_name = model_name.split('/')[-1] if '/' in model_name else model_name
+            
             # Convert message to the appropriate format based on model category
-            config = MODEL_CONFIGS.get(model_name, MODEL_CONFIGS["default"])
+            config = MODEL_CONFIGS.get(base_model_name, MODEL_CONFIGS["default"])
             category = config.get("category", "Other")
             
             if category == "OpenAI":
@@ -370,16 +373,22 @@ Here's the complete Manim code:
                 else:
                     client = models["openai_client"]
                 
-                # For OpenAI models, we need role-based messages
+                # Add openai/ prefix if not present
+                if "/" not in model_name:
+                    full_model_name = f"openai/{model_name}"
+                else:
+                    full_model_name = model_name
+                
+                # For OpenAI models, use developer role instead of system
                 messages = [
-                    {"role": "system", "content": "You are an expert in Manim animations."},
+                    {"role": "developer", "content": "You are an expert in Manim animations."},
                     {"role": "user", "content": prompt}
                 ]
                 
                 # Create params
                 params = {
                     "messages": messages,
-                    "model": model_name
+                    "model": full_model_name
                 }
                 
                 # Add token parameter
@@ -2148,7 +2157,8 @@ class MyScene(Scene):
                             
                             # Get model details
                             model_name = st.session_state.custom_model
-                            config = MODEL_CONFIGS.get(model_name, MODEL_CONFIGS["default"])
+                            base_model_name = model_name.split('/')[-1] if '/' in model_name else model_name
+                            config = MODEL_CONFIGS.get(base_model_name, MODEL_CONFIGS["default"])
                             category = config.get("category", "Other")
                             
                             if category == "OpenAI":
@@ -2176,7 +2186,7 @@ class MyScene(Scene):
                                 # Prepare parameters based on model configuration
                                 params = {
                                     "messages": [
-                                        {"role": "system", "content": "You are a helpful assistant."},
+                                        {"role": "developer", "content": "You are a helpful assistant."},
                                         {"role": "user", "content": "Hello, this is a connection test."}
                                     ],
                                     "model": full_model_name
@@ -2318,45 +2328,88 @@ class MyScene(Scene):
 
             # 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:
+                if st.session_state.ai_models:
                     try:
-                        # Test connection with minimal prompt
-                        from azure.ai.inference.models import UserMessage
+                        # Get model details
                         model_name = st.session_state.custom_model
+                        base_model_name = model_name.split('/')[-1] if '/' in model_name else model_name
+                        config = MODEL_CONFIGS.get(base_model_name, MODEL_CONFIGS["default"])
+                        category = config.get("category", "Other")
                         
-                        # Prepare parameters
-                        messages = [UserMessage("Hello")]
-                        api_params, config = prepare_api_params(messages, model_name)
-                        
-                        # 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
+                        if category == "OpenAI":
+                            # Get token
                             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"]
+                            # Create new OpenAI client
+                            from openai import OpenAI
+                            client = OpenAI(
+                                base_url="https://models.github.ai/inference",
+                                api_key=token
                             )
-                            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"]
+                            # Add openai/ prefix if not present
+                            if "/" not in model_name:
+                                full_model_name = f"openai/{model_name}"
+                            else:
+                                full_model_name = model_name
+                            
+                            # Test with minimal prompt
+                            response = client.chat.completions.create(
+                                messages=[
+                                    {"role": "developer", "content": "You are a helpful assistant."},
+                                    {"role": "user", "content": "Hello, this is a test."}
+                                ],
+                                model=full_model_name,
+                                **{config["param_name"]: config[config["param_name"]]}
+                            )
+                            
+                            # Update session state
+                            st.session_state.ai_models = {
+                                "openai_client": client,
+                                "model_name": full_model_name,
+                                "endpoint": "https://models.github.ai/inference",
+                                "last_loaded": datetime.now().isoformat(),
+                                "category": category
+                            }
+                            
+                            st.success(f"✅ Connection to {full_model_name} refreshed successfully!")
                         else:
-                            response = st.session_state.ai_models["client"].complete(**api_params)
+                            # Test connection with minimal prompt for Azure models
+                            from azure.ai.inference.models import UserMessage
+                            
+                            # Prepare parameters
+                            messages = [UserMessage("Hello")]
+                            api_params, config = prepare_api_params(messages, model_name)
+                            
+                            # 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
                             
-                        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:
+            if st.session_state.ai_models:
                 st.markdown("<div class='card'>", unsafe_allow_html=True)
                 st.markdown("#### Generate Animation from Description")
                 st.write("Describe the animation you want to create, or provide partial code to complete.")
@@ -2389,35 +2442,60 @@ class MyScene(Scene):
                     if code_input:
                         with st.spinner("AI is generating your animation code..."):
                             try:
-                                # Get the client and model name
-                                client = st.session_state.ai_models["client"]
+                                # Get model details
                                 model_name = st.session_state.ai_models["model_name"]
+                                base_model_name = model_name.split('/')[-1] if '/' in model_name else model_name
+                                config = MODEL_CONFIGS.get(base_model_name, MODEL_CONFIGS["default"])
+                                category = config.get("category", "Other")
                                 
                                 # 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:
-                """
-                                
-                                # Prepare API parameters
-                                from azure.ai.inference.models import UserMessage
-                                messages = [UserMessage(prompt)]
-                                api_params, config = prepare_api_params(messages, model_name)
-                                
-                                # Make the API call with proper parameters
-                                response = client.complete(**api_params)
+{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:
+"""
                                 
-                                # Process the response
-                                if response and response.choices and len(response.choices) > 0:
+                                if category == "OpenAI":
+                                    # Use OpenAI client
+                                    client = st.session_state.ai_models["openai_client"]
+                                    
+                                    # Use developer role instead of system
+                                    messages = [
+                                        {"role": "developer", "content": "You are an expert in Manim animations."},
+                                        {"role": "user", "content": prompt}
+                                    ]
+                                    
+                                    # Create params
+                                    params = {
+                                        "messages": messages,
+                                        "model": model_name,
+                                        config["param_name"]: config[config["param_name"]]
+                                    }
+                                    
+                                    # Call API
+                                    response = client.chat.completions.create(**params)
                                     completed_code = response.choices[0].message.content
+                                else:
+                                    # Use Azure client for non-OpenAI models
+                                    from azure.ai.inference.models import UserMessage
+                                    client = st.session_state.ai_models["client"]
                                     
+                                    # Convert message format for Azure
+                                    messages = [UserMessage(prompt)]
+                                    api_params, _ = prepare_api_params(messages, model_name)
+                                    
+                                    # Make API call with Azure client
+                                    response = client.complete(**api_params)
+                                    completed_code = response.choices[0].message.content
+                                
+                                # Process the response
+                                if completed_code:
                                     # Extract code from markdown if present
                                     if "```python" in completed_code:
                                         completed_code = completed_code.split("```python")[1].split("```")[0]
@@ -2427,10 +2505,10 @@ class MyScene(Scene):
                                     # 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