Update app.py

app.py

@@ -8,8 +8,9 @@ import matplotlib.pyplot as plt
 import tempfile
 from gradio_client import Client, handle_file
 from dataclasses import dataclass
-from typing import List, Optional
+from typing import List, Optional, Tuple
 import logging
+import time
 
 # Logging configuration
 logging.basicConfig(level=logging.INFO)
@@ -55,18 +56,27 @@ class MathSolver:
             raise
 
     @staticmethod
-    def query_qwen2(image_path: str, question: str) -> str:
-        try:
-            client = Client("Qwen/Qwen2.5-Math-Demo")
-            return client.predict(
-                image=handle_file(image_path),
-                sketchpad=None,
-                question=question,
-                api_name="/math_chat_bot"
-            )
-        except Exception as e:
-            logger.error(f"Qwen2 Error: {str(e)}")
-            raise
+    def query_qwen2(image_path: str, question: str, max_retries: int = 3) -> Tuple[str, bool]:
+        """
+        Query Qwen2 model with retry mechanism
+        Returns: (result, success)
+        """
+        for attempt in range(max_retries):
+            try:
+                client = Client("Qwen/Qwen2.5-Math-Demo")
+                result = client.predict(
+                    image=handle_file(image_path),
+                    sketchpad=None,
+                    question=question,
+                    api_name="/math_chat_bot"
+                )
+                return result, True
+            except Exception as e:
+                logger.error(f"Qwen2 Error (attempt {attempt + 1}/{max_retries}): {str(e)}")
+                if attempt < max_retries - 1:
+                    time.sleep(2 ** attempt)  # Exponential backoff
+                continue
+        return "Error: Unable to process with Qwen2 model", False
 
     @staticmethod
     def extract_and_execute_python_code(text: str) -> Optional[List[str]]:
@@ -98,7 +108,7 @@ app = Flask(__name__)
 
 token = os.environ.get("TOKEN")
 gemini_config = GeminiConfig(
-    token,  # Replace with your actual API key
+    token,
     generation_config={
         "temperature": 1,
         "max_output_tokens": 8192,
@@ -135,11 +145,16 @@ def upload_image():
         with tempfile.NamedTemporaryFile(delete=False) as temp_file:
             file.save(temp_file.name)
 
-            result = (
-                math_solver.query_gemini(temp_file.name, prompt)
-                if model_choice == "gemini"
-                else math_solver.query_qwen2(temp_file.name, prompt)
-            )
+            if model_choice == "gemini":
+                result = math_solver.query_gemini(temp_file.name, prompt)
+                success = True
+            else:
+                result, success = math_solver.query_qwen2(temp_file.name, prompt)
+                if not success:
+                    # Fallback to Gemini if Qwen2 fails
+                    logger.info("Falling back to Gemini model")
+                    result = math_solver.query_gemini(temp_file.name, prompt)
+                    model_choice = "gemini (fallback)"
 
             # Extract and generate graphs
             image_paths = math_solver.extract_and_execute_python_code(result)