Update app.py

app.py

@@ -1,6 +1,7 @@
 from flask import Flask, request, render_template, jsonify, send_from_directory
 from PIL import Image
-import io
+import google.genai as genai
+from google.genai import types
 import os
 import re
 import matplotlib.pyplot as plt
@@ -9,40 +10,46 @@ from gradio_client import Client, handle_file
 from dataclasses import dataclass
 from typing import List, Optional
 import logging
-from google import genai
-from google.genai import types
 
 # Logging configuration
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+@dataclass
+class GeminiConfig:
+    api_key: str
+    generation_config: dict
+    safety_settings: List[dict]
+    model_name: str = "gemini-2.0-flash-thinking-exp-01-21"
+
 class MathSolver:
-    def __init__(self):
-        plt.switch_backend('Agg')  # Non-interactive backend
+    def __init__(self, gemini_config: GeminiConfig):
+        self.gemini_config = gemini_config
+        self.client = genai.Client(
+            api_key=gemini_config.api_key,
+            http_options={'api_version': 'v1alpha'}
+        )
+        plt.switch_backend('Agg')
 
     def query_gemini(self, image_path: str, prompt: str) -> str:
         try:
             img = Image.open(image_path)
-            buffered = io.BytesIO()
-            img.save(buffered, format="PNG")
-            img_byte_arr = buffered.getvalue()
-
-            response = client.models.generate_content(
-                model="gemini-2.0-flash-thinking-exp-01-21",
-                config={'thinking_config': {'include_thoughts': True}, 'temperature': 1, 'max_output_tokens': 8192},
-                contents=[
-                    {'parts': [{'text': prompt}, {'inline_data': {'mime_type': 'image/png', 'data': img_byte_arr}}]}
-                ]
+            response = self.client.models.generate_content(
+                model=self.gemini_config.model_name,
+                config={
+                    'thinking_config': {'include_thoughts': True},
+                    'generation_config': self.gemini_config.generation_config,
+                    'safety_settings': self.gemini_config.safety_settings
+                },
+                contents=[prompt, img]
             )
-
-            full_response = ""
-            for candidate in response.candidates:
-                for part in candidate.content.parts:
-                    if part.thought:
-                        full_response += f"<br><b>Thought:</b><br> {part.text}<br>"
-                    else:
-                        full_response += f"<br><b>Answer:</b><br> {part.text}<br>"
-            return full_response
+            
+            result = ""
+            for part in response.candidates[0].content.parts:
+                if not part.thought:
+                    result += part.text + "\n"
+            return result.strip()
+            
         except Exception as e:
             logger.error(f"Gemini Error: {str(e)}")
             raise
@@ -89,13 +96,22 @@ class MathSolver:
 # Application configuration
 app = Flask(__name__)
 
-GOOGLE_API_KEY = os.environ.get("TOKEN")
-client = genai.Client(
-    api_key=GOOGLE_API_KEY,
-    http_options={'api_version': 'v1alpha'},
+token = os.environ.get("TOKEN")
+gemini_config = GeminiConfig(
+    token,  # Replace with your actual API key
+    generation_config={
+        "temperature": 1,
+        "max_output_tokens": 8192,
+    },
+    safety_settings=[
+        {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
+        {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
+        {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
+        {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
+    ]
 )
 
-math_solver = MathSolver()
+math_solver = MathSolver(gemini_config)
 
 @app.route('/')
 def index():
@@ -121,10 +137,11 @@ def upload_image():
 
             result = (
                 math_solver.query_gemini(temp_file.name, prompt)
-                if model_choice == "mariam's"
+                if model_choice == "gemini"
                 else math_solver.query_qwen2(temp_file.name, prompt)
             )
 
+            # Extract and generate graphs
             image_paths = math_solver.extract_and_execute_python_code(result)
             os.unlink(temp_file.name)