Fix infinite loop: implement GetOpRequest proxy for Veo polling

agent.py

@@ -13,6 +13,12 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 
+# Fix class to satisfy SDK requirements for polling
+class GetOpRequest:
+    def __init__(self, name):
+        self.name = name
+
+
 def get_file_hash(filepath):
     hash_md5 = hashlib.md5()
     with open(filepath, "rb") as f:
@@ -30,153 +36,96 @@ def get_or_upload_file(client, filepath):
                 return f
     except Exception:
         pass
-    logger.info(f"⬆️ Uploading new file: {file_hash}")
     return client.files.upload(file=filepath, config={'display_name': file_hash})
 
 
 def analyze_only(path_a, path_c, job_id=None):
-    update_job_status(job_id, "analyzing", 10, "Director checking file cache...")
+    update_job_status(job_id, "analyzing", 10, "Analyzing scenes...")
     client = genai.Client(api_key=Settings.GOOGLE_API_KEY)
-
     try:
         file_a = get_or_upload_file(client, path_a)
         file_c = get_or_upload_file(client, path_c)
-        
         while file_a.state.name == "PROCESSING" or file_c.state.name == "PROCESSING":
-            update_job_status(job_id, "analyzing", 20, "Google processing video...")
             time.sleep(2)
             file_a = client.files.get(name=file_a.name)
             file_c = client.files.get(name=file_c.name)
 
-        prompt = "You are a VFX Director. Analyze Video A and Video C. Return JSON with keys: analysis_a, analysis_c, visual_prompt_b."
-        
         res = client.models.generate_content(
             model="gemini-2.0-flash", 
-            contents=[prompt, file_a, file_c],
+            contents=["Analyze Video A and C. Return JSON with analysis_a, analysis_c, visual_prompt_b.", file_a, file_c],
             config=types.GenerateContentConfig(response_mime_type="application/json")
         )
-        
         text = res.text.strip()
-        if text.startswith("```json"): text = text[7:]
-        elif text.startswith("```"): text = text[3:]
-        if text.endswith("```"): text = text[:-3]
-        
-        try: data = json.loads(text.strip())
-        except: data = {}
-        if isinstance(data, list): data = data[0] if len(data) > 0 else {}
-        return {
-            "analysis_a": data.get("analysis_a", ""),
-            "analysis_c": data.get("analysis_c", ""),
-            "prompt": data.get("visual_prompt_b", text),
-            "status": "success"
-        }
+        if "```json" in text:
+            text = text.split("```json")[1].split("```")[0]
+        data = json.loads(text)
+        return {"prompt": data.get("visual_prompt_b", text), "status": "success"}
     except Exception as e:
-        logger.error(f"Analysis failed: {e}")
         return {"detail": str(e), "status": "error"}
 
 
 def generate_only(prompt, path_a, path_c, job_id, style, audio, neg, guidance, motion):
-    update_job_status(job_id, "generating", 50, "Production started (Veo 3.1)...")
+    update_job_status(job_id, "generating", 50, "Production started...")
     full_prompt = f"{style} style. {prompt} Soundtrack: {audio}"
-    if neg:
-        full_prompt += f" --no {neg}"
-
-    job_failed = False
     try:
-        if Settings.GCP_PROJECT_ID:
-            client = genai.Client(vertexai=True, project=Settings.GCP_PROJECT_ID, location=Settings.GCP_LOCATION)
-            
-            # 1. Start Job
-            op = client.models.generate_videos(
-                model='veo-3.1-generate-preview', 
-                prompt=full_prompt, 
-                config=types.GenerateVideosConfig(number_of_videos=1)
-            )
+        if not Settings.GCP_PROJECT_ID:
+            raise Exception("GCP_PROJECT_ID missing.")
+        client = genai.Client(vertexai=True, project=Settings.GCP_PROJECT_ID, location=Settings.GCP_LOCATION)
+
+        # Start Job
+        op = client.models.generate_videos(
+            model='veo-3.1-generate-preview', 
+            prompt=full_prompt, 
+            config=types.GenerateVideosConfig(number_of_videos=1)
+        )
+        
+        # Extract ID and create the proxy object for the SDK
+        op_name = op.name if hasattr(op, 'name') else str(op)
+        request_proxy = GetOpRequest(op_name)
+        logger.info(f"Polling Job: {op_name}")
+        start_time = time.time()
+        while True:
+            if time.time() - start_time > 600:
+                raise Exception("Timeout (10m).")
             
-            # 2. Extract ID String
-            op_name = str(op)
-            if hasattr(op, 'name'): op_name = op.name
-            elif isinstance(op, dict) and 'name' in op: op_name = op['name']
+            # REFRESH logic using the proxy object
+            try:
+                op = client.operations.get(request_proxy)
+            except Exception as e:
+                logger.warning(f"Polling error: {e}")
+                time.sleep(20)
+                continue
             
-            logger.info(f"Tracking ID: {op_name}")
+            # Check Status
+            is_done = getattr(op, 'done', False)
+            if is_done:
+                logger.info("Generation Done.")
+                break
             
-            # 3. Poll for Completion
-            start_time = time.time()
-            while True:
-                if time.time() - start_time > 300: 
-                    raise Exception("Generation timed out.")
-                
-                try:
-                    # FIX: Use positional argument for get()
-                    current_op = client.operations.get(op_name)
-                except Exception as e:
-                    logger.warning(f"Refresh failed: {e}")
-                    time.sleep(10)
-                    continue
-                
-                # Check Status
-                is_done = False
-                if hasattr(current_op, 'done'): is_done = current_op.done
-                elif isinstance(current_op, dict): is_done = current_op.get('done')
-                
-                if is_done:
-                    logger.info("Job DONE.")
-                    op = current_op 
-                    break
-                
-                logger.info("Waiting for Veo...")
-                time.sleep(10)
+            logger.info("Waiting for Veo...")
+            time.sleep(20)
+        
+        # Process Result
+        res_val = op.result
+        result = res_val() if callable(res_val) else res_val
+        
+        if result and (getattr(result, 'generated_videos', None) or 'generated_videos' in result):
+            vid = result.generated_videos[0] if hasattr(result, 'generated_videos') else result['generated_videos'][0]
+            bridge_path = tempfile.mktemp(suffix=".mp4")
             
-            # 4. Get Result
-            result = None
-            if hasattr(op, 'result'):
-                try: 
-                    res_val = op.result
-                    if callable(res_val): result = res_val()
-                    else: result = res_val
-                except: pass
-            elif isinstance(op, dict): 
-                result = op.get('result')
+            if hasattr(vid.video, 'uri') and vid.video.uri:
+                download_blob(vid.video.uri, bridge_path)
+            else:
+                bridge_path = save_video_bytes(vid.video.video_bytes)
             
-            generated_videos = None
-            if result:
-                if hasattr(result, 'generated_videos'): generated_videos = result.generated_videos
-                elif isinstance(result, dict): generated_videos = result.get('generated_videos')
+            update_job_status(job_id, "stitching", 85, "Stitching...")
+            final_cut = os.path.join("outputs", f"{job_id}_final.mp4")
+            merged_path = stitch_videos(path_a, bridge_path, path_c, final_cut)
             
-            if generated_videos:
-                vid = generated_videos[0]
-                bridge_path = None
-                
-                uri = getattr(vid.video, 'uri', None) if hasattr(vid, 'video') else vid.get('video', {}).get('uri')
-                video_bytes = getattr(vid.video, 'video_bytes', None) if hasattr(vid, 'video') else vid.get('video', {}).get('video_bytes')
-
-                if uri:
-                    bridge_path = tempfile.mktemp(suffix=".mp4")
-                    download_blob(uri, bridge_path)
-                elif video_bytes:
-                    bridge_path = save_video_bytes(video_bytes)
-                
-                if bridge_path:
-                    update_job_status(job_id, "stitching", 85, "Stitching...", video_url=bridge_path)
-                    final_cut = os.path.join("outputs", f"{job_id}_merged_temp.mp4")
-                    merged_path = stitch_videos(path_a, bridge_path, path_c, final_cut)
-                    
-                    msg = "Done! (Merged)" if merged_path else "Done! (Bridge Only)"
-                    update_job_status(job_id, "completed", 100, msg, video_url=bridge_path, merged_video_url=merged_path)
-                    return
-            else:
-                raise Exception("No videos returned.")
+            msg = "Done! (Merged)" if merged_path else "Done! (Bridge Only)"
+            update_job_status(job_id, "completed", 100, msg, video_url=bridge_path, merged_video_url=merged_path)
         else:
-             raise Exception("GCP_PROJECT_ID not set.")
-
+            raise Exception("No video output.")
     except Exception as e:
-        logger.error(f"Gen Fatal: {e}")
-        update_job_status(job_id, "error", 0, f"Error: {e}")
-        job_failed = True
-    finally:
-        if not job_failed:
-            try:
-                with open(f"outputs/{job_id}.json", "r") as f:
-                    if json.load(f).get("status") not in ["completed", "error"]:
-                        update_job_status(job_id, "error", 0, "Job timed out.")
-            except: pass
+        logger.error(f"Error: {e}")
+        update_job_status(job_id, "error", 0, str(e))