""" Build FAISS index from movie embeddings This script should be run once to create the data files needed by the API """ import os import json import numpy as np import faiss from openai import OpenAI import requests from typing import Dict, List, Optional import time import argparse from concurrent.futures import ThreadPoolExecutor, as_completed import logging import pickle # Try different import patterns to handle both direct execution and module execution try: from .settings import get_settings except ImportError: try: from app.settings import get_settings except ImportError: from settings import get_settings # Configure logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) # Checkpoint file paths - use temp directory or disable for production import tempfile CHECKPOINT_DIR = os.environ.get('CHECKPOINT_DIR', tempfile.gettempdir()) MOVIE_DATA_CHECKPOINT = f"{CHECKPOINT_DIR}/movie_data.pkl" EMBEDDINGS_CHECKPOINT = f"{CHECKPOINT_DIR}/embeddings_progress.pkl" METADATA_CHECKPOINT = f"{CHECKPOINT_DIR}/metadata_progress.pkl" def save_checkpoint(data, filepath: str): """Save checkpoint data to file - skip if permissions denied""" try: os.makedirs(os.path.dirname(filepath), exist_ok=True) with open(filepath, 'wb') as f: pickle.dump(data, f) logger.info(f"Checkpoint saved: {filepath}") except PermissionError: logger.warning(f"Cannot save checkpoint due to permissions: {filepath}") except Exception as e: logger.warning(f"Failed to save checkpoint {filepath}: {e}") def load_checkpoint(filepath: str): """Load checkpoint data from file""" try: if os.path.exists(filepath): with open(filepath, 'rb') as f: data = pickle.load(f) logger.info(f"Checkpoint loaded: {filepath}") return data except Exception as e: logger.warning(f"Failed to load checkpoint {filepath}: {e}") return None def cleanup_checkpoints(): """Remove checkpoint files after successful completion""" try: import shutil if os.path.exists(CHECKPOINT_DIR) and CHECKPOINT_DIR != tempfile.gettempdir(): shutil.rmtree(CHECKPOINT_DIR) logger.info("Checkpoint files cleaned up") except Exception as e: logger.warning(f"Failed to cleanup checkpoints: {e}") class TMDBClient: """Client for TMDB API with retry and backoff""" def __init__(self, api_key: str): self.api_key = api_key self.base_url = "https://api.themoviedb.org/3" self.session = requests.Session() def _make_request(self, endpoint: str, params: dict = None, max_retries: int = 3) -> Optional[dict]: """Make API request with retry and backoff""" if params is None: params = {} params['api_key'] = self.api_key url = f"{self.base_url}{endpoint}" for attempt in range(max_retries): try: response = self.session.get(url, params=params, timeout=10) if response.status_code == 200: return response.json() elif response.status_code == 429: # Rate limit - wait and retry wait_time = 2 ** attempt logger.warning(f"Rate limited, waiting {wait_time}s before retry...") time.sleep(wait_time) continue elif response.status_code == 404: logger.warning(f"Resource not found: {url}") return None else: logger.error(f"API error {response.status_code}: {response.text}") except requests.exceptions.RequestException as e: logger.error(f"Request failed (attempt {attempt + 1}): {e}") if attempt < max_retries - 1: time.sleep(2 ** attempt) return None def get_popular_movies(self, max_pages: int = 100, filter_adult: bool = True) -> List[int]: """Get movie IDs from popular movies pagination""" movie_ids = [] for page in range(1, max_pages + 1): logger.info(f"Fetching popular movies page {page}/{max_pages}") data = self._make_request("/movie/popular", {"page": page}) if not data: logger.error(f"Failed to fetch page {page}") break # Check if we've exceeded total pages if page > data.get('total_pages', 0): logger.info(f"Reached last page ({data.get('total_pages')})") break # Extract movie IDs, filtering adult content if requested for movie in data.get('results', []): # Skip adult movies if filtering is enabled if filter_adult and movie.get('adult', False): logger.debug(f"Skipping adult movie: {movie.get('title', 'Unknown')} (ID: {movie.get('id')})") continue movie_ids.append(movie['id']) # Rate limiting time.sleep(0.25) # 4 requests per second max logger.info(f"Collected {len(movie_ids)} movie IDs from {page} pages (adult filter: {'ON' if filter_adult else 'OFF'})") return movie_ids def get_movie_details(self, movie_id: int) -> Optional[dict]: """Get detailed movie information""" return self._make_request(f"/movie/{movie_id}") def get_movie_credits(self, movie_id: int) -> Optional[dict]: """Get movie cast and crew""" return self._make_request(f"/movie/{movie_id}/credits") def fetch_movie_data(tmdb_client: TMDBClient, movie_ids: List[int], max_workers: int = 5) -> Dict[int, dict]: """Fetch detailed data for all movies with controlled parallelization""" movies_data = {} def fetch_single_movie(movie_id: int) -> tuple: """Fetch details and credits for a single movie""" try: # Get basic details details = tmdb_client.get_movie_details(movie_id) if not details: return movie_id, None # Get credits credits = tmdb_client.get_movie_credits(movie_id) if credits: details['credits'] = credits return movie_id, details except Exception as e: logger.error(f"Error fetching movie {movie_id}: {e}") return movie_id, None # Process movies in batches with controlled parallelization batch_size = 50 total_movies = len(movie_ids) for i in range(0, total_movies, batch_size): batch = movie_ids[i:i + batch_size] logger.info(f"Processing batch {i//batch_size + 1}/{(total_movies-1)//batch_size + 1} ({len(batch)} movies)") with ThreadPoolExecutor(max_workers=max_workers) as executor: futures = {executor.submit(fetch_single_movie, movie_id): movie_id for movie_id in batch} for future in as_completed(futures): movie_id, movie_data = future.result() if movie_data: movies_data[movie_id] = movie_data # Sleep between batches to be respectful to API time.sleep(1) logger.info(f"Successfully fetched data for {len(movies_data)}/{total_movies} movies") return movies_data def create_composite_text(movie_data: Dict) -> str: """Create composite text for embedding from movie data""" parts = [] # Title if movie_data.get('title'): parts.append(f"Title: {movie_data['title']}") # Tagline if movie_data.get('tagline'): parts.append(f"Tagline: {movie_data['tagline']}") # Overview if movie_data.get('overview'): parts.append(f"Overview: {movie_data['overview']}") # Release date if movie_data.get('release_date'): parts.append(f"Release Date: {movie_data['release_date']}") # Original language if movie_data.get('original_language'): parts.append(f"Language: {movie_data['original_language']}") # Spoken languages if movie_data.get('spoken_languages'): languages = [lang.get('iso_639_1', '') for lang in movie_data['spoken_languages'] if lang.get('iso_639_1')] if languages: parts.append(f"Spoken Languages: {', '.join(languages)}") # Genres if movie_data.get('genres'): genres = [genre['name'] for genre in movie_data['genres']] parts.append(f"Genres: {', '.join(genres)}") # Production companies if movie_data.get('production_companies'): companies = [company['name'] for company in movie_data['production_companies']] if companies: parts.append(f"Production Companies: {', '.join(companies)}") # Production countries if movie_data.get('production_countries'): countries = [country['name'] for country in movie_data['production_countries']] if countries: parts.append(f"Production Countries: {', '.join(countries)}") # Budget (only if > 0) if movie_data.get('budget') and movie_data['budget'] > 0: parts.append(f"Budget: ${movie_data['budget']:,}") # Popularity if movie_data.get('popularity'): parts.append(f"Popularity: {movie_data['popularity']}") # Vote average if movie_data.get('vote_average'): parts.append(f"Vote Average: {movie_data['vote_average']}") # Vote count if movie_data.get('vote_count'): parts.append(f"Vote Count: {movie_data['vote_count']}") # Director(s) if movie_data.get('credits', {}).get('crew'): directors = [person['name'] for person in movie_data['credits']['crew'] if person['job'] == 'Director'] if directors: parts.append(f"Director: {', '.join(directors)}") # Top 5 cast if movie_data.get('credits', {}).get('cast'): top_cast = [person['name'] for person in movie_data['credits']['cast'][:5]] if top_cast: parts.append(f"Cast: {', '.join(top_cast)}") return " / ".join(parts) def get_embeddings_batch(texts: List[str], client: OpenAI, model: str = "text-embedding-3-small") -> List[List[float]]: """Get embeddings for a batch of texts with retry""" max_retries = 3 for attempt in range(max_retries): try: response = client.embeddings.create( input=texts, model=model ) return [item.embedding for item in response.data] except Exception as e: logger.error(f"Error getting embeddings (attempt {attempt + 1}): {e}") if attempt < max_retries - 1: time.sleep(2 ** attempt) else: raise def build_index(max_pages: int = 10, model: str = "text-embedding-3-small", use_faiss: bool = True, override_adult_filter: bool = None): """Main function to build the FAISS index and data files""" settings = get_settings() # Determine adult filtering setting filter_adult = settings.filter_adult_content_bool if hasattr(settings, 'filter_adult_content_bool') else settings.filter_adult_content if override_adult_filter is not None: filter_adult = not override_adult_filter # --include-adult means don't filter logger.info(f"Adult filter override: {'DISABLED' if override_adult_filter else 'ENABLED'}") # Initialize clients with error handling for version compatibility tmdb_client = TMDBClient(settings.tmdb_api_key) try: # Try to create OpenAI client with different approaches for version compatibility try: openai_client = OpenAI(api_key=settings.openai_api_key) except TypeError as e: if "proxies" in str(e): # Fallback for version compatibility issues logger.warning(f"OpenAI client compatibility issue: {e}") logger.info("Trying alternative OpenAI client initialization...") import httpx # Create a basic httpx client without proxies http_client = httpx.Client(timeout=60.0) openai_client = OpenAI(api_key=settings.openai_api_key, http_client=http_client) else: raise except Exception as e: logger.error(f"❌ Failed to initialize OpenAI client: {e}") logger.error("Please check your OpenAI API key and ensure compatible versions are installed") return # Create data directory with absolute path script_dir = os.path.dirname(os.path.abspath(__file__)) data_dir = os.path.join(script_dir, "data") try: os.makedirs(data_dir, exist_ok=True) # Test write permissions test_file = os.path.join(data_dir, ".write_test") with open(test_file, 'w') as f: f.write("test") os.remove(test_file) logger.info(f"Data directory ready: {data_dir}") except PermissionError as e: logger.error(f"❌ Permission denied when creating data directory: {e}") logger.error("Make sure the data directory has write permissions") return except Exception as e: logger.error(f"❌ Failed to create or write to data directory: {e}") return # Check for existing movie data checkpoint movies_data = load_checkpoint(MOVIE_DATA_CHECKPOINT) if movies_data is not None: logger.info(f"🔄 Resuming from checkpoint: {len(movies_data)} movies data found") else: # Step 1: Get movie IDs logger.info(f"Fetching movie IDs from TMDB (max {max_pages} pages)...") movie_ids = tmdb_client.get_popular_movies( max_pages=max_pages, filter_adult=filter_adult ) if not movie_ids: logger.error("❌ No movie IDs retrieved from TMDB") return # Step 2: Fetch detailed movie data logger.info(f"Fetching detailed data for {len(movie_ids)} movies...") movies_data = fetch_movie_data(tmdb_client, movie_ids) if not movies_data: logger.error("❌ No movie data retrieved") return # Additional filtering at the detail level (double-check) if filter_adult: original_count = len(movies_data) movies_data = {k: v for k, v in movies_data.items() if not v.get('adult', False)} filtered_count = original_count - len(movies_data) if filtered_count > 0: logger.info(f"Filtered out {filtered_count} adult movies at detail level") # Save movie data checkpoint save_checkpoint(movies_data, MOVIE_DATA_CHECKPOINT) # Step 3: Create composite texts and process embeddings in batches logger.info("Creating embeddings...") embeddings = [] id_map = {} movie_metadata = {} processed_movie_ids = set() batch_size = 20 # Process 20 texts at a time # Check for existing embedding progress embedding_checkpoint = load_checkpoint(EMBEDDINGS_CHECKPOINT) metadata_checkpoint = load_checkpoint(METADATA_CHECKPOINT) if embedding_checkpoint is not None and metadata_checkpoint is not None: embeddings = embedding_checkpoint['embeddings'] id_map = embedding_checkpoint['id_map'] processed_movie_ids = set(embedding_checkpoint['processed_movie_ids']) movie_metadata = metadata_checkpoint logger.info(f"🔄 Resuming embeddings from checkpoint: {len(embeddings)} embeddings found") else: logger.info("Starting embeddings from scratch") # Process remaining movies remaining_movies = {k: v for k, v in movies_data.items() if k not in processed_movie_ids} logger.info(f"Processing {len(remaining_movies)} remaining movies") composite_texts = [] current_movie_ids = [] for movie_id, movie_data in remaining_movies.items(): # Create composite text composite_text = create_composite_text(movie_data) composite_texts.append(composite_text) current_movie_ids.append(movie_id) # Store metadata release_year = 0 if movie_data.get("release_date"): try: release_year = int(movie_data["release_date"][:4]) except (ValueError, IndexError): release_year = 0 movie_metadata[str(movie_id)] = { "id": movie_id, "title": movie_data.get("title", ""), "year": release_year, "poster_path": movie_data.get("poster_path"), "release_date": movie_data.get("release_date"), "genres": [g["name"] for g in movie_data.get("genres", [])] } # Process batch when full if len(composite_texts) >= batch_size: logger.info(f"Processing embedding batch ({len(embeddings)} done, {len(composite_texts)} in batch)") try: batch_embeddings = get_embeddings_batch(composite_texts, openai_client, model) embeddings.extend(batch_embeddings) # Update ID mapping and processed set for i, mid in enumerate(current_movie_ids): id_map[str(mid)] = len(id_map) processed_movie_ids.add(mid) # Save progress checkpoints embedding_data = { 'embeddings': embeddings, 'id_map': id_map, 'processed_movie_ids': list(processed_movie_ids) } save_checkpoint(embedding_data, EMBEDDINGS_CHECKPOINT) save_checkpoint(movie_metadata, METADATA_CHECKPOINT) # Clear batch composite_texts = [] current_movie_ids = [] # Sleep between batches time.sleep(0.5) except Exception as e: logger.error(f"Failed to process batch: {e}") logger.info("Progress has been saved, you can restart the script to resume") return # Process remaining texts if composite_texts: logger.info(f"Processing final embedding batch ({len(composite_texts)} texts)") try: batch_embeddings = get_embeddings_batch(composite_texts, openai_client, model) embeddings.extend(batch_embeddings) for i, mid in enumerate(current_movie_ids): id_map[str(mid)] = len(id_map) processed_movie_ids.add(mid) # Save final progress embedding_data = { 'embeddings': embeddings, 'id_map': id_map, 'processed_movie_ids': list(processed_movie_ids) } save_checkpoint(embedding_data, EMBEDDINGS_CHECKPOINT) save_checkpoint(movie_metadata, METADATA_CHECKPOINT) except Exception as e: logger.error(f"Failed to process final batch: {e}") logger.info("Progress has been saved, you can restart the script to resume") return if not embeddings: logger.error("❌ No embeddings generated") return logger.info(f"Generated {len(embeddings)} embeddings") # Step 4: Save embeddings as numpy array embeddings_array = np.array(embeddings, dtype=np.float32) embeddings_path = os.path.join(data_dir, "movies.npy") try: np.save(embeddings_path, embeddings_array) logger.info(f"Saved embeddings matrix: {embeddings_array.shape}") except Exception as e: logger.error(f"❌ Failed to save embeddings: {e}") return # Step 5: Build and save FAISS index if use_faiss: logger.info("Building FAISS index...") dimension = embeddings_array.shape[1] # Choose index type based on size if len(embeddings) < 10000: # For smaller datasets, use flat index index = faiss.IndexFlatL2(dimension) else: # For larger datasets, use IVF index nlist = min(int(np.sqrt(len(embeddings))), 1000) quantizer = faiss.IndexFlatL2(dimension) index = faiss.IndexIVFFlat(quantizer, dimension, nlist) # Train the index index.train(embeddings_array) index.add(embeddings_array) index_path = os.path.join(data_dir, "faiss.index") try: faiss.write_index(index, index_path) logger.info(f"FAISS index saved (type: {type(index).__name__}, dimension: {dimension})") except Exception as e: logger.error(f"❌ Failed to save FAISS index: {e}") return # Step 6: Save metadata files id_map_path = os.path.join(data_dir, "id_map.json") metadata_path = os.path.join(data_dir, "movie_metadata.json") try: with open(id_map_path, "w") as f: json.dump(id_map, f) with open(metadata_path, "w") as f: json.dump(movie_metadata, f) logger.info("✅ Index built successfully!") logger.info(f" - {len(embeddings)} movies indexed") logger.info(f" - Embedding model: {model}") logger.info(f" - Files saved in {data_dir}") logger.info(f" * movies.npy: embeddings matrix") logger.info(f" * id_map.json: TMDB ID to matrix position mapping") logger.info(f" * movie_metadata.json: movie metadata") if use_faiss: logger.info(f" * faiss.index: FAISS search index") # Cleanup checkpoints cleanup_checkpoints() except Exception as e: logger.error(f"❌ Failed to save metadata files: {e}") return # Remove the old functions that are no longer needed # create_movie_embedding and load_movie_data are replaced by the new implementation if __name__ == "__main__": parser = argparse.ArgumentParser(description="Build movie embeddings index from TMDB data") parser.add_argument("--max-pages", type=int, default=10, help="Maximum pages to fetch from TMDB popular movies (default: 10)") parser.add_argument("--model", type=str, default="text-embedding-3-small", help="OpenAI embedding model to use (default: text-embedding-3-small)") parser.add_argument("--no-faiss", action="store_true", help="Skip building FAISS index") parser.add_argument("--include-adult", action="store_true", help="Include adult movies (overrides FILTER_ADULT_CONTENT setting)") args = parser.parse_args() build_index( max_pages=args.max_pages, model=args.model, use_faiss=not args.no_faiss, override_adult_filter=args.include_adult )