Update app.py

app.py

@@ -6,6 +6,7 @@ Generates custom musical exercises with LLM, perfectly fit to user-specified num
 AND time signature, guaranteeing exact durations in MIDI and in the UI!
 
 Major updates:
+- Added Gemma, Kimi Dev 72b, and Llama 3.1 AI model options
 - Added duration sum display in Exercise Data tab
 - Shows total duration units (16th notes) for verification
 - Added DeepSeek AI model option
@@ -19,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:
@@ -36,7 +39,6 @@ install([
 # -----------------------------------------------------------------------------
 # 2. Static imports
 # -----------------------------------------------------------------------------
-import random
 import requests
 import json
 import tempfile
@@ -55,14 +57,20 @@ import os
 import subprocess as sp
 import base64
 import shutil
-from openai import OpenAI  # For DeepSeek API
+from openai import OpenAI  # For API models
 
 # -----------------------------------------------------------------------------
 # 3. Configuration & constants
 # -----------------------------------------------------------------------------
 MISTRAL_API_URL = "https://api.mistral.ai/v1/chat/completions"
 MISTRAL_API_KEY = "yQdfM8MLbX9uhInQ7id4iUTwN4h4pDLX"  # Replace with your key
-DEEPSEEK_API_KEY = "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"
+OPENROUTER_API_KEYS = {
+    "DeepSeek": "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",
+    "Claude": "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",
+    "Gemma": "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",
+    "Kimi": "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",
+    "Llama 3.1": "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"
+}
 
 SOUNDFONT_URLS = {
     "Trumpet": "https://github.com/FluidSynth/fluidsynth/raw/master/sf2/Trumpet.sf2",
@@ -261,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 Mistral and DeepSeek)
+# 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:
@@ -295,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",
@@ -312,45 +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 == "DeepSeek":
-        try:
-            client = OpenAI(
-                base_url="https://openrouter.ai/api/v1",
-                api_key=DEEPSEEK_API_KEY,
-            )
-            
-            completion = client.chat.completions.create(
-                extra_headers={
-                    "HTTP-Referer": "https://github.com/AdaptiveMusicExerciseGenerator",
-                    "X-Title": "Music Exercise Generator",
-                },
-                model="deepseek/deepseek-chat-v3-0324:free",
-                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 DeepSeek 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)
 
 # -----------------------------------------------------------------------------
@@ -399,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. Simple AI chat assistant (optional, shares LLM)
+# 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():
@@ -410,49 +468,84 @@ 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 == "DeepSeek":
+    for attempt in range(max_retries):
         try:
-            client = OpenAI(
-                base_url="https://openrouter.ai/api/v1",
-                api_key=DEEPSEEK_API_KEY,
-            )
+            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
             
-            completion = client.chat.completions.create(
-                extra_headers={
-                    "HTTP-Referer": "https://github.com/AdaptiveMusicExerciseGenerator",
-                    "X-Title": "Music Exercise Generator",
-                },
-                model="deepseek/deepseek-chat-v3-0324:free",
-                messages=messages,
-                temperature=0.7,
-                max_tokens=500,
-            )
-            content = completion.choices[0].message.content
-            history.append((message, content))
-            return "", history
+            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
+            
+            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 (for humans!)
+# 13. Gradio user interface definition
 # -----------------------------------------------------------------------------
 def create_ui() -> gr.Blocks:
     with gr.Blocks(title="Adaptive Music Exercise Generator", theme="soft") as demo:
@@ -466,7 +559,7 @@ def create_ui() -> gr.Blocks:
                 with gr.Group(visible=True) as params_group:
                     gr.Markdown("### Exercise Parameters")
                     ai_model = gr.Radio(
-                        ["Mistral", "DeepSeek"],
+                        ["Mistral", "DeepSeek", "Claude", "Gemma", "Kimi", "Llama 3.1"],
                         value="Mistral",
                         label="AI Model"
                     )