Update app.py

app.py

@@ -5,7 +5,7 @@ import re
 import gradio as gr
 import google.generativeai as genai
 
-# In ra version để debug
+# ==== DEBUG VERSION ====
 print("Google Generative AI SDK version:", genai.__version__)
 
 # ==== API KEY ROTATION ====
@@ -14,7 +14,7 @@ API_KEYS = [
     "AIzaSyBfJeK_IAkfLpmnBsMNe6xwjcielrloSFY",
     "AIzaSyCprm_rLChQ7Rv7YkHKI_6tcbS213PiPto",
     "AIzaSyAPFCgH8uSjANmPRF9iHYIYcneTOod8Qi0",
-    "AIzaSyBbK-1P3JD6HPyE3QLhkOps6_-Xo3wUFbs"
+    "AIzaSyBbK-1P3JD6HPyE3QLhkOps6_-Xo3wUFbs",
 ]
 
 key_index = 0
@@ -34,59 +34,64 @@ def extract_json(text):
 
     try:
         return json.loads(json_text)
-    except:
-        first = json_text.find("{")
-        last = json_text.rfind("}")
+    except Exception:
+        first, last = json_text.find("{"), json_text.rfind("}")
         if first != -1 and last != -1 and last > first:
             try:
                 return json.loads(json_text[first:last+1])
-            except:
+            except Exception:
                 pass
         return {"raw_response": text}
 
 # ==== MAIN PROCESS ====
 def process_image(image, prompt):
+    temp_file = None
+    uploaded_file = None
+
     try:
         print(f"Received image: {type(image)}")
         print(f"Received prompt: {prompt}")
 
-        # Lấy API key
+        # === API key rotation ===
         api_key = get_next_key()
         genai.configure(api_key=api_key)
 
-        # Tạo model Gemini
-        model = genai.GenerativeModel("gemini-2.5-flash")
-
-        # Tạo file tạm
-        temp_file = None
-        uploaded_file = None
+        # === Ensure RAG store exists ===
+        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}")
 
-        with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as tmp:
-            if hasattr(image, 'save'):
+        # === Create temp file ===
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp:
+            if hasattr(image, "save"):
                 image.save(tmp.name)
             elif isinstance(image, str):
-                with open(image, 'rb') as f:
+                with open(image, "rb") as f:
                     tmp.write(f.read())
             else:
                 raise ValueError(f"Unsupported image type: {type(image)}")
             temp_file = tmp.name
-
         print(f"Created temp file: {temp_file}")
 
-        # === Upload ảnh lên Gemini ===
-        uploaded_file = genai.upload_file(temp_file)
+        # === Upload to Gemini with RAG ===
+        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
+        )
+        print(f"Uploaded to Gemini RAG store: {uploaded_file.name}")
 
-        #uploaded_file = genai.upload_file(
-        #    path=temp_file,
-        #    display_name=os.path.basename(temp_file),
-        #    mime_type="image/png",
-            
-            #rag_store_name ="default_rag_store"
-        #)
-        print(f"Uploaded to Gemini: {uploaded_file.name}")
+        # === Create model ===
+        model = genai.GenerativeModel("gemini-2.5-flash")
 
-        # === Gọi model ===
-        response = model.generate_content([prompt, uploaded_file])
+        # === Generate content using RAG context ===
+        response = model.generate_content(
+            [prompt],
+            request_options={"rag_store_name": RAG_STORE_NAME},  # ✅ cho phép truy vấn dựa theo store
+        )
         print(f"Got response: {response.text[:200]}...")
 
         result = extract_json(response.text)
@@ -94,12 +99,10 @@ def process_image(image, prompt):
 
     except Exception as e:
         print(f"Error: {e}")
-        import traceback
-        traceback.print_exc()
+        import traceback; traceback.print_exc()
         return {"error": str(e)}
 
     finally:
-        # Cleanup
         if temp_file and os.path.exists(temp_file):
             os.remove(temp_file)
             print(f"Cleaned temp file: {temp_file}")
@@ -118,8 +121,8 @@ demo = gr.Interface(
         gr.Textbox(lines=5, placeholder="Enter your prompt here...", label="Prompt"),
     ],
     outputs=gr.JSON(label="Response"),
-    title="OCR & Analyzer",
-    description="Upload an image + prompt, returns structured JSON extracted",
+    title="Gemini OCR & Analyzer (RAG Enhanced)",
+    description="Upload an image + prompt → analyze using Gemini 2.5 with RAG store",
     flagging_mode="never",
 )