Update app.py

app.py

@@ -3,20 +3,13 @@ import tempfile
 import json
 import re
 import gradio as gr
-# import google.generativeai as genai # Layton Luan - edit
-from google import genai              # Layton Luan - edit
+from google import genai
 
-# ==== DEBUG VERSION ====
 print("Google Generative AI SDK version:", genai.__version__)
 
-# ==== API KEY ROTATION ====
-# API_KEYS = [
-#        "AIzaSyAzoMIf_QK4piBlu779vxBo_CnNb6YH2Cc",
-#     "AIzaSyAyaOfYKS3b4Em0HmTOi0aLjumrXsURi80",
-#     "AIzaSyAjxoDOuio_Q30ErOk6gDx9S1zQVGDzGkw"
-# ]
-API_KEYS = ["AIzaSyBAZ1Zm2KCZHlmVKPgVf2Cch0c-0YJfJVg"
-                   ]
+API_KEYS = [
+    "AIzaSyBAZ1Zm2KCZHlmVKPgVf2Cch0c-0YJfJVg"
+]
 
 key_index = 0
 def get_next_key():
@@ -25,6 +18,7 @@ def get_next_key():
     key_index += 1
     return key
 
+
 # ==== JSON EXTRACTION ====
 def extract_json(text):
     match = re.search(r"```json\s*(.*?)\s*```", text, re.IGNORECASE | re.DOTALL)
@@ -44,6 +38,7 @@ def extract_json(text):
                 pass
         return {"raw_response": text}
 
+
 # ==== MAIN PROCESS ====
 def process_image(image, prompt):
     temp_file = None
@@ -53,21 +48,19 @@ def process_image(image, prompt):
         print(f"Received image: {type(image)}")
         print(f"Received prompt: {prompt}")
 
-        # === API key rotation ===
         api_key = get_next_key()
 
-        # Layton Luan - edit
-        # genai.configure(api_key=api_key)
+        # INIT CLIENT
         client = genai.Client(api_key=api_key)
 
-        # === Ensure RAG store exists ===
+        # Create RAG store
         RAG_STORE_NAME = "default_rag_store"
         try:
             genai.create_rag_store(name=RAG_STORE_NAME)
         except Exception as e:
-            print(f"(Info) RAG store may already exist: {e}")
+            print("(Info) RAG store exists:", e)
 
-        # === Create temp file ===
+        # Create temp file
         with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp:
             if hasattr(image, "save"):
                 image.save(tmp.name)
@@ -77,60 +70,47 @@ def process_image(image, prompt):
             else:
                 raise ValueError(f"Unsupported image type: {type(image)}")
             temp_file = tmp.name
-        print(f"Created temp file: {temp_file}")
+        print("Temp file:", temp_file)
 
-        # === Upload to Gemini with RAG ===
-
-        # Layton Luan - edit
-        # uploaded_file = genai.upload_file(
-        #     path=temp_file,
-        #     display_name=os.path.basename(temp_file),
-        #     mime_type="image/png",
-        #     rag_store_name=RAG_STORE_NAME,  # ✅ thêm dòng này
-        # )
+        # Upload file
         uploaded_file = client.files.upload(file=temp_file)
-        print(f"Uploaded to Gemini RAG store: {uploaded_file.name}")
-
-        # === Create model ===
-        # Layton Luan - edit
-        # model = genai.GenerativeModel("gemini-2.5-flash")
+        print("Uploaded:", uploaded_file.name)
 
-        # === Generate content using RAG context ===
+        # === CREATE MODEL USING genai.GenerativeModel ====
+        model = genai.GenerativeModel(
+            "gemini-2.0-flash",
+            client=client
+        )
 
-        # Layton Luan - edit
-        # response = model.generate_content(
-        #     [prompt],
-        #     request_options={"rag_store_name": RAG_STORE_NAME},  # ✅ cho phép truy vấn dựa theo store
-        # )
-        response = client.models.generate_content(
-            model='gemini-2.0-flash',
+        # === GENERATE CONTENT ====
+        response = model.generate_content(
             contents=[
-                [prompt],
-                # request_options={"rag_store_name": RAG_STORE_NAME},  # ✅ cho phép truy vấn dựa theo store
-                uploaded_file
-            ]
+                prompt,
+                uploaded_file   # file reference
+            ],
+            request_options={
+                "temperature": 0.2
+            }
         )
-        
-        print(f"Got response: {response.text[:200]}...")
 
-        result = extract_json(response.text) 
-        return result
+        print("Raw response:", response.text[:200], "...")
+
+        return extract_json(response.text)
 
     except Exception as e:
-        print(f"Error: {e}")
+        print("Error:", e)
         import traceback; traceback.print_exc()
         return {"error": str(e)}
 
     finally:
         if temp_file and os.path.exists(temp_file):
             os.remove(temp_file)
-            print(f"Cleaned temp file: {temp_file}")
         if uploaded_file:
             try:
                 genai.delete_file(uploaded_file.name)
-                print(f"Deleted from Gemini: {uploaded_file.name}")
-            except Exception as del_err:
-                print(f"Warning: Could not delete file: {del_err}")
+            except:
+                pass
+
 
 # ==== GRADIO UI ====
 demo = gr.Interface(