Update app.py

app.py

@@ -1,43 +1,48 @@
 from flask import Flask, request, render_template, jsonify, send_from_directory
 from PIL import Image
-import google.generativeai as genai
+import io
 import os
 import re
 import matplotlib.pyplot as plt
 import tempfile
 from gradio_client import Client, handle_file
-# import subprocess  # Not used
 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-exp-1206"
-
 class MathSolver:
-    def __init__(self, gemini_config: GeminiConfig):
-        self.gemini_config = gemini_config
-        genai.configure(api_key=gemini_config.api_key)
+    def __init__(self):
         plt.switch_backend('Agg')  # Non-interactive backend
 
     def query_gemini(self, image_path: str, prompt: str) -> str:
         try:
             img = Image.open(image_path)
-            model = genai.GenerativeModel(
-                model_name=self.gemini_config.model_name,
-                generation_config=self.gemini_config.generation_config,
-                safety_settings=self.gemini_config.safety_settings
+            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 = model.generate_content([prompt, img], request_options={"timeout": 600})
-            return response.text
+
+            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
         except Exception as e:
             logger.error(f"Gemini Error: {str(e)}")
             raise
@@ -65,8 +70,7 @@ class MathSolver:
         image_paths = []
         for code in code_blocks:
             try:
-                code = "import numpy as np\n" + code 
-                # Replace single backslashes with double backslashes
+                code = "import numpy as np\n" + code
                 code = code.replace("\\", "\\\\")
 
                 with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmpfile:
@@ -85,22 +89,13 @@ class MathSolver:
 # Application configuration
 app = Flask(__name__)
 
-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"},
-    ]
+GOOGLE_API_KEY = os.environ.get("TOKEN")
+client = genai.Client(
+    api_key=GOOGLE_API_KEY,
+    http_options={'api_version': 'v1alpha'},
 )
 
-math_solver = MathSolver(gemini_config)
+math_solver = MathSolver()
 
 @app.route('/')
 def index():
@@ -130,7 +125,6 @@ def upload_image():
                 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)