Update app.py

app.py

@@ -1,16 +1,15 @@
 import gradio as gr
-import os
 from llama_cpp import Llama
 from qdrant_client import QdrantClient
 from datasets import load_dataset
 from sentence_transformers import SentenceTransformer
 import cv2
+import os
 import tempfile
 import uuid
 import re
 import subprocess
 import time
-import traceback
 
 # Configuration
 QDRANT_COLLECTION_NAME = "video_frames"
@@ -53,8 +52,9 @@ except Exception as e:
 
 print("Loading dataset stream...")
 try:
+    # Load video dataset
     dataset = load_dataset("aegean-ai/ai-lectures-spring-24", split="train", streaming=True)
-    print(f"Dataset loaded.")
+    print(f"Dataset loaded. First item example: {next(iter(dataset))['__key__']}")
 except Exception as e:
     print(f"Error loading dataset: {e}")
     raise
@@ -85,7 +85,7 @@ def rag_query(client, collection_name, query_text, top_k=5, filter_condition=Non
         if filter_condition:
             search_params["filter"] = filter_condition
 
-        search_results = client.query_points(query_points=query_vector, **search_params)
+        search_results = client.search(**search_params)
 
         formatted_results = []
         for idx, result in enumerate(search_results):
@@ -105,69 +105,128 @@ def rag_query(client, collection_name, query_text, top_k=5, filter_condition=Non
         }
     except Exception as e:
         print(f"Error during RAG query: {e}")
-        traceback.print_exc()
         return {"error": str(e), "query": query_text, "results": []}
 
 
 def extract_video_segment(video_id, start_time, duration, dataset):
     """
-    Extracts a single video segment file path from the dataset stream.
-    Saves it to a temporary file and returns the path or None on failure.
-    Uses FFmpeg with -ss before -i and -t.
+    Generator function that extracts and yields a single video segment file path.
+    Modified to return a single path suitable for Gradio.
     """
     target_id = str(video_id)
-    target_key_pattern = re.compile(r"videos/" + re.escape(target_id) + r"/" + re.escape(target_id))
-
+    target_key = f"videos/{target_id}/{target_id}"
     start_time = float(start_time)
     duration = float(duration)
 
     unique_id = str(uuid.uuid4())
-    temp_dir = os.path.join(tempfile.gettempdir(), f"gradio_video_seg_{unique_id}")
+    temp_dir = os.path.join(tempfile.gettempdir(), f"gradio_video_{unique_id}")
     os.makedirs(temp_dir, exist_ok=True)
-    temp_video_path_full = os.path.join(temp_dir, f"{target_id}_full_{unique_id}.mp4")
+    temp_video_path = os.path.join(temp_dir, f"{target_id}_full_{unique_id}.mp4")
+    output_path_opencv = os.path.join(temp_dir, f"output_opencv_{unique_id}.mp4")
     output_path_ffmpeg = os.path.join(temp_dir, f"output_ffmpeg_{unique_id}.mp4")
 
-    print(f"Attempting to extract segment for video_id={target_id}, start={start_time:.2f}, duration={duration:.2f}")
-    print(f"Looking for dataset key matching pattern: {target_key_pattern.pattern}")
+    print(f"Attempting to extract segment for video_id={target_id}, start={start_time}, duration={duration}")
+    print(f"Looking for dataset key: {target_key}")
     print(f"Temporary directory: {temp_dir}")
 
-    found_sample = None
-    max_search_attempts = 1000 # Limit
-    print(f"Searching dataset stream for key matching pattern: {target_key_pattern.pattern}")
-
-    dataset_iterator = iter(dataset)
 
     try:
-        # Find and save the full video from the stream
-        for i in range(max_search_attempts):
-            try:
-                sample = next(dataset_iterator)
-                if '__key__' in sample and 'mp4' in sample and target_key_pattern.match(sample['__key__']):
-                    print(f"Found video key {sample['__key__']} after {i+1} iterations. Saving to {temp_video_path_full}...")
-                    with open(temp_video_path_full, 'wb') as f:
-                        f.write(sample['mp4'])
-                    print(f"Video saved successfully ({os.path.getsize(temp_video_path_full)} bytes).")
-                    found_sample = sample
-                    break # Found the video, exit loop
-            except StopIteration:
-                print("Reached end of dataset stream without finding the video within search limit.")
+        found = False
+        retries = 3
+        dataset_iterator = iter(dataset)
+
+        for _ in range(retries * 100):
+             try:
+                 sample = next(dataset_iterator)
+                 if '__key__' in sample and sample['__key__'] == target_key:
+                     found = True
+                     print(f"Found video key {target_key}. Saving to {temp_video_path}...")
+                     with open(temp_video_path, 'wb') as f:
+                         f.write(sample['mp4'])
+                     print(f"Video saved successfully ({os.path.getsize(temp_video_path)} bytes).")
+                     break
+             except StopIteration:
+                 print("Reached end of dataset stream without finding the video.")
+                 break
+             except Exception as e:
+                 print(f"Error iterating dataset: {e}")
+                 time.sleep(1)
+
+
+        if not found:
+             print(f"Could not find video with ID {target_id} (key: {target_key}) in the dataset stream after {_ + 1} attempts.")
+             return None
+
+        # Process the saved video
+        if not os.path.exists(temp_video_path) or os.path.getsize(temp_video_path) == 0:
+             print(f"Temporary video file {temp_video_path} is missing or empty.")
+             return None
+
+        cap = cv2.VideoCapture(temp_video_path)
+        if not cap.isOpened():
+            print(f"Error opening video file with OpenCV: {temp_video_path}")
+            return None
+
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        if fps <= 0:
+            print(f"Warning: Invalid FPS ({fps}) detected for {temp_video_path}. Assuming 30 FPS.")
+            fps = 30
+
+        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        total_vid_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        vid_duration = total_vid_frames / fps if fps > 0 else 0
+
+        print(f"Video properties: {width}x{height} @ {fps:.2f}fps, Total Duration: {vid_duration:.2f}s")
+
+        start_frame = int(start_time * fps)
+        end_frame = int((start_time + duration) * fps)
+
+        # Clamp frame numbers to valid range
+        start_frame = max(0, start_frame)
+        end_frame = min(total_vid_frames, end_frame)
+
+        if start_frame >= total_vid_frames or start_frame >= end_frame:
+             print(f"Calculated start frame ({start_frame}) is beyond video length ({total_vid_frames}) or segment is invalid.")
+             cap.release()
+             return None
+
+        cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
+        frames_to_write = end_frame - start_frame
+
+        print(f"Extracting frames from {start_frame} to {end_frame} ({frames_to_write} frames)")
+
+        # Try OpenCV first
+        fourcc_opencv = cv2.VideoWriter_fourcc(*'mp4v') # mp4v is often more compatible than avc1 with base OpenCV
+        out_opencv = cv2.VideoWriter(output_path_opencv, fourcc_opencv, fps, (width, height))
+
+        if not out_opencv.isOpened():
+             print("Error opening OpenCV VideoWriter with mp4v.")
+             cap.release()
+             return None
+
+        frames_written_opencv = 0
+        while frames_written_opencv < frames_to_write:
+            ret, frame = cap.read()
+            if not ret:
+                print("Warning: Ran out of frames before reaching target end frame.")
                 break
-            except Exception as e:
-                print(f"Warning: Error iterating dataset sample {i+1}: {e}")
+            out_opencv.write(frame)
+            frames_written_opencv += 1
 
-        if not found_sample or not os.path.exists(temp_video_path_full) or os.path.getsize(temp_video_path_full) == 0:
-            print(f"Could not find or save video with ID {target_id} from dataset stream.")
-            return None
+        out_opencv.release()
+        print(f"OpenCV finished writing {frames_written_opencv} frames to {output_path_opencv}")
 
-        # Process the saved video with FFmpeg
+        cap.release()
+
+        # FFmpeg
         final_output_path = None
         try:
             cmd = [
                 'ffmpeg',
-                '-y', # Overwrite output file if exists
-                '-ss', str(start_time), # Start time
-                '-i', temp_video_path_full, # Input file
-                '-t', str(duration), # Duration of the segment
+                '-ss', str(start_time),       # Start time
+                '-i', temp_video_path,        # Input file (original downloaded)
+                '-t', str(duration),          # Duration of the segment
                 '-c:v', 'libx264',
                 '-profile:v', 'baseline',
                 '-level', '3.0',
@@ -176,10 +235,11 @@ def extract_video_segment(video_id, start_time, duration, dataset):
                 '-movflags', '+faststart',
                 '-c:a', 'aac',
                 '-b:a', '128k',
+                '-y',
                 output_path_ffmpeg
             ]
             print(f"Running FFmpeg command: {' '.join(cmd)}")
-            result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
+            result = subprocess.run(cmd, capture_output=True, text=True, timeout=120) # Add timeout
 
             if result.returncode == 0 and os.path.exists(output_path_ffmpeg) and os.path.getsize(output_path_ffmpeg) > 0:
                 print(f"FFmpeg processing successful. Output: {output_path_ffmpeg}")
@@ -188,155 +248,167 @@ def extract_video_segment(video_id, start_time, duration, dataset):
                 print(f"FFmpeg error (Return Code: {result.returncode}):")
                 print(f"FFmpeg stdout:\n{result.stdout}")
                 print(f"FFmpeg stderr:\n{result.stderr}")
-                print("FFmpeg failed.")
-                final_output_path = None
+                print("Falling back to OpenCV output.")
+                if os.path.exists(output_path_opencv) and os.path.getsize(output_path_opencv) > 0:
+                     final_output_path = output_path_opencv
+                else:
+                     print("OpenCV output is also invalid or empty.")
+                     final_output_path = None
 
         except subprocess.TimeoutExpired:
-            print("FFmpeg command timed out.")
-            final_output_path = None
+             print("FFmpeg command timed out.")
+             print("Falling back to OpenCV output.")
+             if os.path.exists(output_path_opencv) and os.path.getsize(output_path_opencv) > 0:
+                 final_output_path = output_path_opencv
+             else:
+                 print("OpenCV output is also invalid or empty.")
+                 final_output_path = None
         except FileNotFoundError:
-            print("Error: ffmpeg command not found. Make sure FFmpeg is installed in the environment.")
-            final_output_path = None
+            print("Error: ffmpeg command not found. Make sure FFmpeg is installed and in your system's PATH.")
+            print("Falling back to OpenCV output.")
+            if os.path.exists(output_path_opencv) and os.path.getsize(output_path_opencv) > 0:
+                 final_output_path = output_path_opencv
+            else:
+                 print("OpenCV output is also invalid or empty.")
+                 final_output_path = None
         except Exception as e:
             print(f"An unexpected error occurred during FFmpeg processing: {e}")
-            traceback.print_exc()
-            final_output_path = None
-
-    finally:
-        # Clean up temporary files
-        print(f"Cleaning up temporary directory: {temp_dir}")
-        if os.path.exists(temp_video_path_full):
-            try:
-                os.remove(temp_video_path_full)
-                print(f"Cleaned up temporary full video: {temp_video_path_full}")
-            except Exception as e:
-                print(f"Warning: Could not remove temporary file {temp_video_path_full}: {e}")
+            print("Falling back to OpenCV output.")
+            if os.path.exists(output_path_opencv) and os.path.getsize(output_path_opencv) > 0:
+                 final_output_path = output_path_opencv
+            else:
+                 print("OpenCV output is also invalid or empty.")
+                 final_output_path = None
 
+        if os.path.exists(temp_video_path):
+             try:
+                 os.remove(temp_video_path)
+                 print(f"Cleaned up temporary full video: {temp_video_path}")
+             except Exception as e:
+                 print(f"Warning: Could not remove temporary file {temp_video_path}: {e}")
+
+        # If FFmpeg failed
         if final_output_path != output_path_ffmpeg and os.path.exists(output_path_ffmpeg):
-            try:
-                os.remove(output_path_ffmpeg)
-            except Exception as e:
-                print(f"Warning: Could not remove failed ffmpeg output {output_path_ffmpeg}: {e}")
+             try:
+                 os.remove(output_path_ffmpeg)
+             except Exception as e:
+                 print(f"Warning: Could not remove failed ffmpeg output {output_path_ffmpeg}: {e}")
 
-    if final_output_path and os.path.exists(final_output_path):
         print(f"Returning video segment path: {final_output_path}")
         return final_output_path
-    else:
-        print("Video segment extraction failed.")
+
+    except Exception as e:
+        print(f"Error processing video segment for {video_id}: {e}")
+        import traceback
+        traceback.print_exc()
+        if 'cap' in locals() and cap.isOpened(): cap.release()
+        if 'out_opencv' in locals() and out_opencv.isOpened(): out_opencv.release()
+        if os.path.exists(temp_video_path): os.remove(temp_video_path)
+        if os.path.exists(output_path_opencv): os.remove(output_path_opencv)
+        if os.path.exists(output_path_ffmpeg): os.remove(output_path_ffmpeg)
         return None
 
+QDRANT_COLLECTION_NAME = "video_frames"
+VIDEO_SEGMENT_DURATION = 40 # Extract 40 seconds around the timestamp
+
 
 def parse_llm_output(text):
     """
-    Parses the LLM's structured output using string manipulation.
-    Returns parsed data dictionary.
+    Parses the LLM's structured output using a mix of regex for simple
+    fields (video_id, timestamp) and string manipulation for reasoning
+    as a workaround for regex matching issues.
     """
     data = {}
-    print(f"\nDEBUG: Raw text input to parse_llm_output:\n---\n{text}\n---")
 
-    def extract_field(text, field_name):
-        start_marker_lower = "{" + field_name.lower() + ":"
-        start_index = text.lower().find(start_marker_lower)
+    # Parse video_id and timestamp with regex
+    simple_patterns = {
+        'video_id': r"\{Best Result:\s*\[?([^\]\}]+)\]?\s*\}",
+        'timestamp': r"\{Timestamp:\s*\[?([^\]\}]+)\]?\s*\}",
+    }
+    for key, pattern in simple_patterns.items():
+        match = re.search(pattern, text, re.IGNORECASE)
+        if match:
+            value = match.group(1).strip()
+            value = value.strip('\'"“”')
+            data[key] = value
+        else:
+            print(f"Warning: Could not parse '{key}' using regex pattern: {pattern}")
+            data[key] = None
+
+    # Parse reasoning
+    reasoning_value = None
+    try:
+        key_marker_lower = "{reasoning:"
+        start_index = text.lower().find(key_marker_lower)
 
         if start_index != -1:
-            actual_marker_end = start_index + len(start_marker_lower)
-            end_index = text.find('}', actual_marker_end)
+            search_start_for_brace = start_index + len(key_marker_lower)
+            end_index = text.find('}', search_start_for_brace)
 
             if end_index != -1:
+                actual_marker_end = start_index + len(key_marker_lower)
                 value = text[actual_marker_end : end_index]
+
                 value = value.strip()
                 if value.startswith('[') and value.endswith(']'):
-                    value = value[1:-1].strip()
+                    value = value[1:-1]
                 value = value.strip('\'"“”')
-                return value.strip()
+                value = value.strip()
+                reasoning_value = value
             else:
-                print(f"Warning: Found '{{{field_name}:' marker but no closing '}}' found afterwards.")
+                print("Warning: Found '{reasoning:' marker but no closing '}' found afterwards.")
         else:
-            print(f"Warning: Marker '{{{field_name}:' not found in text.")
-        return None
+            print("Warning: Marker '{reasoning:' not found in text.")
 
-    # Extract fields
-    data['video_id'] = extract_field(text, 'Best Result')
-    data['timestamp'] = extract_field(text, 'Timestamp')
-    data['content'] = extract_field(text, 'Content')
-    data['reasoning'] = extract_field(text, 'Reasoning')
+    except Exception as e:
+        print(f"Error during string manipulation parsing for reasoning: {e}")
+
+    data['reasoning'] = reasoning_value
 
-    # Validation
     if data.get('timestamp'):
         try:
             float(data['timestamp'])
         except ValueError:
             print(f"Warning: Parsed timestamp '{data['timestamp']}' is not a valid number.")
-            data['timestamp'] = None
 
-    print(f"Parsed LLM output: {data}")
+    print(f"Parsed LLM output (Using String Manipulation for Reasoning): {data}")
     return data
 
 
 def process_query_and_get_video(query_text):
     """
     Orchestrates RAG, LLM query, parsing, and video extraction.
-    Returns the path to the extracted video segment or None on failure.
-    Prints status and errors directly.
     """
     print(f"\n--- Processing query: '{query_text}' ---")
 
-    # Check if necessary components are initialized
-    if 'qdrant_client' not in globals() or qdrant_client is None:
-        print("Setup Error: Qdrant client is not initialized. Cannot proceed.")
-        return None
-    if 'llm' not in globals() or llm is None:
-         print("Setup Error: LLM is not initialized. Cannot proceed.")
-         return None
-    if 'embedding_model' not in globals() or embedding_model is None:
-         print("Setup Error: Embedding model is not initialized. Cannot proceed.")
-         return None
-    if 'dataset' not in globals() or dataset is None:
-         print("Setup Error: Dataset is not loaded. Cannot proceed.")
-         return None
-
-    # RAG Query
+    # 1. RAG Query
     print("Step 1: Performing RAG query...")
     rag_results = rag_query(qdrant_client, QDRANT_COLLECTION_NAME, query_text)
 
     if "error" in rag_results or not rag_results.get("results"):
         error_msg = rag_results.get('error', 'No relevant segments found by RAG.')
         print(f"RAG Error/No Results: {error_msg}")
-        return None
+        return f"Error during RAG search: {error_msg}", None
 
     print(f"RAG query successful. Found {len(rag_results['results'])} results.")
 
     # Format LLM Prompt
     print("Step 2: Formatting prompt for LLM...")
-    results_for_llm = "\n".join([
-        f"Rank: {r['rank']}, Score: {r['score']:.4f}, Video ID: {r['video_id']}, Timestamp: {r['timestamp']}, Subtitle: {r['subtitle']}"
-        for r in rag_results['results']
-    ])
-
     prompt = f"""You are tasked with selecting the most relevant information from a set of video subtitle segments to answer a query.
-
-QUERY: "{query_text}"
-
-Here are the relevant video segments found:
----
-{results_for_llm}
----
-
+QUERY (also seen below): "{query_text}"
 For each result provided, evaluate how well it directly addresses the definition or explanation related to the query. Pay attention to:
 1. Clarity of explanation
 2. Relevance to the query
 3. Completeness of information
-
 From the provided results, select the SINGLE BEST match that most directly answers the query.
-
 Format your response STRICTLY as follows, with each field on a new line:
 {{Best Result: [video_id]}}
 {{Timestamp: [timestamp]}}
-{{Content: [subtitle text from the selected result]}}
+{{Content: [subtitle text]}}
 {{Reasoning: [Brief explanation of why this result best answers the query]}}
-"""
+{rag_results}"""
 
-    # Call LLM
+    # 3. Call LLM
     print("Step 3: Querying the LLM...")
     try:
         output = llm.create_chat_completion(
@@ -347,56 +419,49 @@ Format your response STRICTLY as follows, with each field on a new line:
             temperature=0.1,
             max_tokens=300
         )
-        llm_response_text = output['choices'][0]['message']['content'].strip()
-        print(f"LLM Response:\n---\n{llm_response_text}\n---")
+        llm_response_text = output['choices'][0]['message']['content']
+        print(f"LLM Response:\n{llm_response_text}")
     except Exception as e:
         print(f"Error during LLM call: {e}")
-        traceback.print_exc()
-        return None
+        return f"Error calling LLM: {e}", None
 
-    # Parse LLM Response
+    # 4. Parse LLM Response
     print("Step 4: Parsing LLM response...")
     parsed_data = parse_llm_output(llm_response_text)
 
     video_id = parsed_data.get('video_id')
     timestamp_str = parsed_data.get('timestamp')
-    # Get reasoning/content
     reasoning = parsed_data.get('reasoning')
-    content = parsed_data.get('content')
-
-    if reasoning:
-        print(f"LLM Reasoning: {reasoning}")
-
-    if content:
-        print(f"LLM Selected Content: {content}")
-
 
     if not video_id or not timestamp_str:
         print("Error: Could not parse required video_id or timestamp from LLM response.")
-        print("Raw LLM response that failed parsing:\n---\n{llm_response_text}\n---")
-        return None
+        fallback_reasoning = reasoning if reasoning else "Could not determine the best segment."
+        error_msg = f"Failed to parse LLM response. LLM said:\n---\n{llm_response_text}\n---\nReasoning (if found): {fallback_reasoning}"
+        return error_msg, None
 
     try:
         timestamp = float(timestamp_str)
-        start_time = max(0.0, timestamp - (VIDEO_SEGMENT_DURATION / 4.0))
-        actual_duration = VIDEO_SEGMENT_DURATION
-        print(f"Calculated segment start time: {start_time:.2f}s")
-
+         # Adjust timestamp slightly - start a bit earlier if possible
+        start_time = max(0.0, timestamp - (VIDEO_SEGMENT_DURATION / 4))
     except ValueError:
         print(f"Error: Could not convert parsed timestamp '{timestamp_str}' to float.")
-        return None
+        error_msg = f"Invalid timestamp format from LLM ('{timestamp_str}'). LLM reasoning (if found): {reasoning}"
+        return error_msg, None
 
-    # 5. Extract Video Segment
-    print(f"Step 5: Extracting video segment (ID: {video_id}, Start: {start_time:.2f}s, Duration: {actual_duration:.2f}s)...")
-    video_path = extract_video_segment(video_id, start_time, actual_duration, dataset)
+    final_reasoning = reasoning if reasoning else "No reasoning provided by LLM."
+
+    # Extract Video Segment
+    print(f"Step 5: Extracting video segment (ID: {video_id}, Start: {start_time:.2f}s, Duration: {VIDEO_SEGMENT_DURATION}s)...")
+    global dataset
+    video_path = extract_video_segment(video_id, start_time, VIDEO_SEGMENT_DURATION, dataset)
 
     if video_path and os.path.exists(video_path):
         print(f"Video segment extracted successfully: {video_path}")
-        return video_path
+        return final_reasoning, video_path
     else:
         print("Failed to extract video segment.")
-        return None
-
+        error_msg = f"{final_reasoning}\n\n(However, failed to extract the corresponding video segment for ID {video_id} at timestamp {timestamp_str}.)"
+        return error_msg, None
 
 with gr.Blocks() as iface:
     gr.Markdown(
@@ -410,7 +475,7 @@ with gr.Blocks() as iface:
         query_input = gr.Textbox(label="Your Question", placeholder="e.g., What is a convolutional neural network?")
         submit_button = gr.Button("Ask & Find Video")
     with gr.Row():
-        video_output = gr.Video(label="Relevant Video Segment", format="mp4")
+        video_output = gr.Video(label="Relevant Video Segment")
 
     submit_button.click(
         fn=process_query_and_get_video,
@@ -425,7 +490,7 @@ with gr.Blocks() as iface:
             "Using only the videos, explain the the binary cross entropy loss function.",
         ],
         inputs=query_input,
-        outputs=video_output,
+        outputs= video_output,
         fn=process_query_and_get_video,
         cache_examples=False,
     )