Update app.py

app.py

@@ -9,14 +9,15 @@ import base64
 from concurrent.futures import ThreadPoolExecutor
 from PIL import Image, ImageOps
 
-
 from fastapi import FastAPI, UploadFile, File, Form, HTTPException
 from fastapi.staticfiles import StaticFiles
 from fastapi.middleware.cors import CORSMiddleware
-from PIL import Image
 import rarfile
 import zipfile
-import google.generativeai as genai
+
+# --- MIGRATION: New SDK Imports ---
+from google import genai
+from google.genai import types
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -154,12 +155,9 @@ def parse_srt(content: str):
     return parsed
 
 
-logger = logging.getLogger(__name__)
-
 def compress_image(image_bytes, max_width=800, quality=80):
     """
     Compresses an image to WebP (best) or optimized JPEG.
-    Renamed back to 'compress_image' to fix your error.
     """
     try:
         img = Image.open(io.BytesIO(image_bytes))
@@ -170,7 +168,6 @@ def compress_image(image_bytes, max_width=800, quality=80):
         buffer = io.BytesIO()
 
         # 2. Try WebP first (Best quality/size ratio)
-        # If you strictly need JPEG, change use_webp to False
         use_webp = True 
         
         if use_webp:
@@ -202,14 +199,14 @@ def compress_image(image_bytes, max_width=800, quality=80):
 
     except Exception as e:
         logger.error(f"Image compression failed: {e}")
-        # If logging isn't setup, print the error so you can see it
-        print(f"Error: {e}")
         return None
 
+# --- MIGRATION: Updated Gemini Processing Function ---
 def process_batch_gemini(api_key, items, model_name):
     try:
-        genai.configure(api_key=api_key)
-        model = genai.GenerativeModel(model_name)
+        # 1. Instantiate the Client (New SDK pattern)
+        # This replaces genai.configure()
+        client = genai.Client(api_key=api_key)
 
         prompt_parts = [
             "You are a Subtitle Quality Control (QC) bot.",
@@ -224,172 +221,14 @@ def process_batch_gemini(api_key, items, model_name):
             prompt_parts.append(f"Index: {item['index']}")
             prompt_parts.append(f"Expected Text: \"{item['expected_text']}\"")
             prompt_parts.append(f"Image:")
+            
+            # The new SDK handles PIL images directly in the contents list just like the old one
             img = Image.open(io.BytesIO(item['image_data']))
             prompt_parts.append(img)
 
-        # Enforce JSON mode
-        response = model.generate_content(
-            prompt_parts,
-            generation_config={"response_mime_type": "application/json"}
-        )
-        
-        text = response.text.replace("```json", "").replace("```", "").strip()
-        
-        try:
-            return json.loads(text)
-        except json.JSONDecodeError as e:
-            # Handle Truncated JSON (Output Token Limit Exceeded)
-            # This happens if the batch size is too large for the model's output window
-            logger.warning(f"JSON Parse Error (likely truncated response): {e}. Attempting repair...")
-            
-            # Repair Strategy: Find the last closing brace '}', discard everything after, and close the array ']'
-            last_object_idx = text.rfind("}")
-            if last_object_idx != -1:
-                repaired_text = text[:last_object_idx+1] + "]"
-                try:
-                    repaired_data = json.loads(repaired_text)
-                    logger.info(f"Successfully repaired JSON. Recovered {len(repaired_data)}/{len(items)} items.")
-                    return repaired_data
-                except json.JSONDecodeError:
-                    logger.error("JSON repair failed.")
-            
-            return None # Fail gracefully if repair is impossible
-
-    except Exception as e:
-        logger.error(f"Gemini API Error with key ...{api_key[-4:]}: {e}")
-        return None 
-
-# --- Main Endpoint ---
-
-@app.post("/api/analyze")
-async def analyze_subtitles(
-    srt_file: UploadFile = File(...),
-    media_files: list[UploadFile] = File(...),
-    api_keys: str = Form(...),
-    batch_size: int = Form(20),
-    model_name: str = Form("gemini-3-flash-preview"),
-    compression_quality: float = Form(0.7)
-):
-    temp_dir = tempfile.mkdtemp()
-    try:
-        # Convert float quality (0.1-1.0) to integer (10-100) for PIL
-        pil_quality = max(10, min(100, int(compression_quality * 100)))
-
-        # 1. Read SRT
-        srt_content = (await srt_file.read()).decode('utf-8', errors='ignore')
-        srt_data = parse_srt(srt_content)
-        srt_data.sort(key=lambda x: x['startTimeMs'])
-
-        # 2. Process Media
-        images = []
-        for file in media_files:
-            file_path = os.path.join(temp_dir, file.filename)
-            with open(file_path, "wb") as f:
-                shutil.copyfileobj(file.file, f)
-            
-            if file.filename.lower().endswith('.rar'):
-                try:
-                    with rarfile.RarFile(file_path) as rf:
-                        rf.extractall(temp_dir)
-                except rarfile.RarCannotExec:
-                     raise HTTPException(status_code=500, detail="Unrar executable not found in container.")
-            elif file.filename.lower().endswith('.zip'):
-                with zipfile.ZipFile(file_path, 'r') as zf:
-                    zf.extractall(temp_dir)
-        
-        for root, _, files in os.walk(temp_dir):
-            for filename in files:
-                if filename.lower().endswith(('.jpg', '.jpeg', '.png', '.webp', '.bmp')):
-                    full_path = os.path.join(root, filename)
-                    ms = parse_filename_to_ms(filename)
-                    if ms is not None:
-                        with open(full_path, "rb") as f:
-                            raw_bytes = f.read()
-                            compressed = compress_image(raw_bytes, quality=pil_quality)
-                            if compressed:
-                                images.append({
-                                    "filename": filename,
-                                    "timeMs": ms,
-                                    "data": compressed
-                                })
-
-        images.sort(key=lambda x: x['timeMs'])
-        
-        # 3. Pair
-        pairs = []
-        for i in range(len(images)):
-            img = images[i]
-            srt = srt_data[i] if i < len(srt_data) else None
-            
-            if srt:
-                # Create Thumbnail (lower quality for UI speed)
-                thumb_bytes = compress_image(img['data'], quality=50, max_width=300)
-                thumb_b64 = base64.b64encode(thumb_bytes).decode('utf-8')
-                
-                pairs.append({
-                    "index": i,
-                    "image_data": img['data'],
-                    "expected_text": srt['text'],
-                    "srt_id": srt['id'],
-                    "srt_time": srt['time'],
-                    "filename": img['filename'],
-                    "thumb": f"data:image/jpeg;base64,{thumb_b64}",
-                    "status": "pending"
-                })
-
-        if not pairs:
-            return {"status": "error", "message": "No valid image/subtitle pairs found."}
-
-        # 4. Process Gemini
-        keys = [k.strip() for k in api_keys.split('\n') if k.strip()]
-        if not keys:
-            raise HTTPException(status_code=400, detail="No API Keys provided")
-
-        results_map = {}
-        batches = [pairs[i:i + batch_size] for i in range(0, len(pairs), batch_size)]
-        
-        def worker(batch_idx, batch):
-            key = keys[batch_idx % len(keys)]
-            return process_batch_gemini(key, batch, model_name)
-
-        with ThreadPoolExecutor(max_workers=len(keys)) as executor:
-            futures = [executor.submit(worker, i, b) for i, b in enumerate(batches)]
-            for future in futures:
-                res = future.result()
-                if res:
-                    for item in res:
-                        results_map[item['index']] = item
-
-        # 5. Build Output
-        final_output = []
-        for p in pairs:
-            analysis = results_map.get(p['index'])
-            status = "pending"
-            reason = ""
-            detected = ""
-            if analysis:
-                status = "match" if analysis['match'] else "mismatch"
-                reason = analysis.get('reason', '')
-                detected = analysis.get('detected_text', '')
-
-            final_output.append({
-                "id": p['index'],
-                "filename": p['filename'],
-                "thumb": p['thumb'],
-                "expected": p['expected_text'],
-                "detected": detected,
-                "status": status,
-                "reason": reason,
-                "srt_id": p['srt_id'],
-                "srt_time": p['srt_time']
-            })
-
-        return {"status": "success", "results": final_output}
-
-    except Exception as e:
-        logger.error(f"Server Error: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
-    finally:
-        shutil.rmtree(temp_dir)
-
-app.mount("/", StaticFiles(directory="static", html=True), name="static")
+        # 2. Call generate_content via the client
+        response = client.models.generate_content(
+            model=model_name,
+            contents=prompt_parts,
+            config=types.GenerateContentConfig(
+                response_mime_type="application/json"