Update app.py

app.py

@@ -20,6 +20,8 @@ Major updates:
 import sys
 import subprocess
 from typing import Dict, Optional, Tuple, List
+import time
+import random
 
 def install(packages: List[str]):
     for package in packages:
@@ -37,7 +39,6 @@ install([
 # -----------------------------------------------------------------------------
 # 2. Static imports
 # -----------------------------------------------------------------------------
-import random
 import requests
 import json
 import tempfile
@@ -268,7 +269,7 @@ def get_technique_based_on_level(level: str) -> str:
     return random.choice(techniques.get(level, ["with slurs"]))
 
 # -----------------------------------------------------------------------------
-# 9. LLM Query Function (supports multiple models)
+# 9. LLM Query Function (with enhanced error handling)
 # -----------------------------------------------------------------------------
 def query_llm(model_name: str, prompt: str, instrument: str, level: str, key: str,
               time_sig: str, measures: int) -> str:
@@ -302,7 +303,90 @@ def query_llm(model_name: str, prompt: str, instrument: str, level: str, key: st
             "Sum must be exactly as specified. ONLY output the JSON array. No prose."
         )
 
-    if model_name == "Mistral":
+    # Retry up to 3 times for rate limited models
+    max_retries = 3
+    retry_delay = 5  # seconds
+    
+    for attempt in range(max_retries):
+        try:
+            if model_name == "Mistral":
+                headers = {
+                    "Authorization": f"Bearer {MISTRAL_API_KEY}",
+                    "Content-Type": "application/json",
+                }
+                payload = {
+                    "model": "mistral-medium",
+                    "messages": [
+                        {"role": "system", "content": system_prompt},
+                        {"role": "user", "content": user_prompt},
+                    ],
+                    "temperature": 0.7 if level == "Advanced" else 0.5,
+                    "max_tokens": 1000,
+                    "top_p": 0.95,
+                    "frequency_penalty": 0.2,
+                    "presence_penalty": 0.2,
+                }
+                response = requests.post(MISTRAL_API_URL, headers=headers, json=payload)
+                response.raise_for_status()
+                content = response.json()["choices"][0]["message"]["content"]
+                return content.replace("```json","").replace("```","").strip()
+            
+            elif model_name in ["DeepSeek", "Claude", "Gemma", "Kimi", "Llama 3.1"]:
+                client = OpenAI(
+                    base_url="https://openrouter.ai/api/v1",
+                    api_key=OPENROUTER_API_KEYS[model_name],
+                )
+                
+                model_map = {
+                    "DeepSeek": "deepseek/deepseek-chat-v3-0324:free",
+                    "Claude": "anthropic/claude-3.5-sonnet:beta",
+                    "Gemma": "google/gemma-3n-e2b-it:free",
+                    "Kimi": "moonshotai/kimi-dev-72b:free",
+                    "Llama 3.1": "meta-llama/llama-3.1-405b-instruct:free"
+                }
+                
+                # Special handling for Gemma API structure
+                if model_name == "Gemma":
+                    messages = [
+                        {"role": "user", "content": user_prompt}
+                    ]
+                else:
+                    messages = [
+                        {"role": "system", "content": system_prompt},
+                        {"role": "user", "content": user_prompt},
+                    ]
+                
+                completion = client.chat.completions.create(
+                    extra_headers={
+                        "HTTP-Referer": "https://github.com/AdaptiveMusicExerciseGenerator",
+                        "X-Title": "Music Exercise Generator",
+                    },
+                    model=model_map[model_name],
+                    messages=messages,
+                    temperature=0.7 if level == "Advanced" else 0.5,
+                    max_tokens=1000,
+                    top_p=0.95,
+                    frequency_penalty=0.2,
+                    presence_penalty=0.2,
+                )
+                content = completion.choices[0].message.content
+                return content.replace("```json","").replace("```","").strip()
+            
+            else:
+                return get_fallback_exercise(instrument, level, key, time_sig, measures)
+                
+        except Exception as e:
+            print(f"Error querying {model_name} API (attempt {attempt+1}): {e}")
+            if "429" in str(e) or "Rate limit" in str(e):
+                print(f"Rate limited, retrying in {retry_delay} seconds...")
+                time.sleep(retry_delay)
+                retry_delay *= 2  # Exponential backoff
+            else:
+                break
+    
+    # Fallback to Mistral if other APIs fail
+    print(f"All attempts failed for {model_name}, using Mistral fallback")
+    try:
         headers = {
             "Authorization": f"Bearer {MISTRAL_API_KEY}",
             "Content-Type": "application/json",
@@ -319,53 +403,12 @@ def query_llm(model_name: str, prompt: str, instrument: str, level: str, key: st
             "frequency_penalty": 0.2,
             "presence_penalty": 0.2,
         }
-        try:
-            response = requests.post(MISTRAL_API_URL, headers=headers, json=payload)
-            response.raise_for_status()
-            content = response.json()["choices"][0]["message"]["content"]
-            return content.replace("```json","").replace("```","").strip()
-        except Exception as e:
-            print(f"Error querying Mistral API: {e}")
-            return get_fallback_exercise(instrument, level, key, time_sig, measures)
-    
-    elif model_name in ["DeepSeek", "Claude", "Gemma", "Kimi", "Llama 3.1"]:
-        try:
-            client = OpenAI(
-                base_url="https://openrouter.ai/api/v1",
-                api_key=OPENROUTER_API_KEYS[model_name],
-            )
-            
-            model_map = {
-                "DeepSeek": "deepseek/deepseek-chat-v3-0324:free",
-                "Claude": "anthropic/claude-3.5-sonnet:beta",
-                "Gemma": "google/gemma-3n-e2b-it:free",
-                "Kimi": "moonshotai/kimi-dev-72b:free",
-                "Llama 3.1": "meta-llama/llama-3.1-405b-instruct:free"
-            }
-            
-            completion = client.chat.completions.create(
-                extra_headers={
-                    "HTTP-Referer": "https://github.com/AdaptiveMusicExerciseGenerator",
-                    "X-Title": "Music Exercise Generator",
-                },
-                model=model_map[model_name],
-                messages=[
-                    {"role": "system", "content": system_prompt},
-                    {"role": "user", "content": user_prompt},
-                ],
-                temperature=0.7 if level == "Advanced" else 0.5,
-                max_tokens=1000,
-                top_p=0.95,
-                frequency_penalty=0.2,
-                presence_penalty=0.2,
-            )
-            content = completion.choices[0].message.content
-            return content.replace("```json","").replace("```","").strip()
-        except Exception as e:
-            print(f"Error querying {model_name} API: {e}")
-            return get_fallback_exercise(instrument, level, key, time_sig, measures)
-    
-    else:
+        response = requests.post(MISTRAL_API_URL, headers=headers, json=payload)
+        response.raise_for_status()
+        content = response.json()["choices"][0]["message"]["content"]
+        return content.replace("```json","").replace("```","").strip()
+    except Exception as e:
+        print(f"Error querying Mistral fallback: {e}")
         return get_fallback_exercise(instrument, level, key, time_sig, measures)
 
 # -----------------------------------------------------------------------------
@@ -414,7 +457,7 @@ def generate_exercise(instrument: str, level: str, key: str, tempo: int, time_si
         return f"Error: {str(e)}", None, str(tempo), None, "0", time_signature, 0
 
 # -----------------------------------------------------------------------------
-# 12. AI chat assistant
+# 12. AI chat assistant with enhanced error handling
 # -----------------------------------------------------------------------------
 def handle_chat(message: str, history: List, instrument: str, level: str, ai_model: str):
     if not message.strip():
@@ -425,54 +468,81 @@ def handle_chat(message: str, history: List, instrument: str, level: str, ai_mod
         messages.append({"role": "assistant", "content": assistant_msg})
     messages.append({"role": "user", "content": message})
     
-    if ai_model == "Mistral":
-        headers = {"Authorization": f"Bearer {MISTRAL_API_KEY}", "Content-Type": "application/json"}
-        payload = {"model": "mistral-medium", "messages": messages, "temperature": 0.7, "max_tokens": 500}
-        try:
-            response = requests.post(MISTRAL_API_URL, headers=headers, json=payload)
-            response.raise_for_status()
-            content = response.json()["choices"][0]["message"]["content"]
-            history.append((message, content))
-            return "", history
-        except Exception as e:
-            history.append((message, f"Error: {str(e)}"))
-            return "", history
+    max_retries = 3
+    retry_delay = 3  # seconds
     
-    elif ai_model in ["DeepSeek", "Claude", "Gemma", "Kimi", "Llama 3.1"]:
+    for attempt in range(max_retries):
         try:
-            client = OpenAI(
-                base_url="https://openrouter.ai/api/v1",
-                api_key=OPENROUTER_API_KEYS[ai_model],
-            )
+            if ai_model == "Mistral":
+                headers = {"Authorization": f"Bearer {MISTRAL_API_KEY}", "Content-Type": "application/json"}
+                payload = {"model": "mistral-medium", "messages": messages, "temperature": 0.7, "max_tokens": 500}
+                response = requests.post(MISTRAL_API_URL, headers=headers, json=payload)
+                response.raise_for_status()
+                content = response.json()["choices"][0]["message"]["content"]
+                history.append((message, content))
+                return "", history
             
-            model_map = {
-                "DeepSeek": "deepseek/deepseek-chat-v3-0324:free",
-                "Claude": "anthropic/claude-3.5-sonnet:beta",
-                "Gemma": "google/gemma-3n-e2b-it:free",
-                "Kimi": "moonshotai/kimi-dev-72b:free",
-                "Llama 3.1": "meta-llama/llama-3.1-405b-instruct:free"
-            }
+            elif ai_model in ["DeepSeek", "Claude", "Gemma", "Kimi", "Llama 3.1"]:
+                client = OpenAI(
+                    base_url="https://openrouter.ai/api/v1",
+                    api_key=OPENROUTER_API_KEYS[ai_model],
+                )
+                
+                model_map = {
+                    "DeepSeek": "deepseek/deepseek-chat-v3-0324:free",
+                    "Claude": "anthropic/claude-3.5-sonnet:beta",
+                    "Gemma": "google/gemma-3n-e2b-it:free",
+                    "Kimi": "moonshotai/kimi-dev-72b:free",
+                    "Llama 3.1": "meta-llama/llama-3.1-405b-instruct:free"
+                }
+                
+                # Special handling for Gemma API structure
+                if ai_model == "Gemma":
+                    adjusted_messages = [{"role": "user", "content": msg["content"]} for msg in messages]
+                else:
+                    adjusted_messages = messages
+                
+                completion = client.chat.completions.create(
+                    extra_headers={
+                        "HTTP-Referer": "https://github.com/AdaptiveMusicExerciseGenerator",
+                        "X-Title": "Music Exercise Generator",
+                    },
+                    model=model_map[ai_model],
+                    messages=adjusted_messages,
+                    temperature=0.7,
+                    max_tokens=500,
+                )
+                content = completion.choices[0].message.content
+                history.append((message, content))
+                return "", history
             
-            completion = client.chat.completions.create(
-                extra_headers={
-                    "HTTP-Referer": "https://github.com/AdaptiveMusicExerciseGenerator",
-                    "X-Title": "Music Exercise Generator",
-                },
-                model=model_map[ai_model],
-                messages=messages,
-                temperature=0.7,
-                max_tokens=500,
-            )
-            content = completion.choices[0].message.content
-            history.append((message, content))
-            return "", history
+            else:
+                history.append((message, "Error: Invalid AI model selected"))
+                return "", history
+                
         except Exception as e:
-            history.append((message, f"Error: {str(e)}"))
-            return "", history
-    
-    else:
-        history.append((message, "Error: Invalid AI model selected"))
-        return "", history
+            print(f"Chat error with {ai_model} (attempt {attempt+1}): {e}")
+            if "429" in str(e) or "Rate limit" in str(e):
+                print(f"Rate limited, retrying in {retry_delay} seconds...")
+                time.sleep(retry_delay)
+                retry_delay *= 2  # Exponential backoff
+            else:
+                # Fallback to Mistral
+                print(f"Using Mistral fallback for chat")
+                try:
+                    headers = {"Authorization": f"Bearer {MISTRAL_API_KEY}", "Content-Type": "application/json"}
+                    payload = {"model": "mistral-medium", "messages": messages, "temperature": 0.7, "max_tokens": 500}
+                    response = requests.post(MISTRAL_API_URL, headers=headers, json=payload)
+                    response.raise_for_status()
+                    content = response.json()["choices"][0]["message"]["content"]
+                    history.append((message, content))
+                    return "", history
+                except Exception as e:
+                    history.append((message, f"Error: {str(e)}"))
+                    return "", history
+                    
+    history.append((message, "Error: All API attempts failed"))
+    return "", history
 
 # -----------------------------------------------------------------------------
 # 13. Gradio user interface definition