Add adult content filtering option and update related functionality in TMDBClient and settings

.env.example

@@ -11,4 +11,7 @@ API_TOKEN=your_api_token_here
 ENV=dev
 
 # Logging level
-LOG_LEVEL=INFO
+LOG_LEVEL=INFO
+
+# Remove adult content from TMDB results
+FILTER_ADULT_CONTENT=true  # Set to true to filter out adult content

app/build_index.py

@@ -110,7 +110,7 @@ class TMDBClient:
                     
         return None
     
-    def get_popular_movies(self, max_pages: int = 100) -> List[int]:
+    def get_popular_movies(self, max_pages: int = 100, filter_adult: bool = True) -> List[int]:
         """Get movie IDs from popular movies pagination"""
         movie_ids = []
         
@@ -127,14 +127,18 @@ class TMDBClient:
                 logger.info(f"Reached last page ({data.get('total_pages')})")
                 break
                 
-            # Extract movie IDs
+            # 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")
+        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]:
@@ -285,10 +289,16 @@ def get_embeddings_batch(texts: List[str], client: OpenAI, model: str = "text-em
             else:
                 raise
 
-def build_index(max_pages: int = 10, model: str = "text-embedding-3-small", use_faiss: bool = True):
+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
+    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
     tmdb_client = TMDBClient(settings.tmdb_api_key)
     openai_client = OpenAI(api_key=settings.openai_api_key)
@@ -321,7 +331,10 @@ def build_index(max_pages: int = 10, model: str = "text-embedding-3-small", use_
     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)
+        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")
@@ -335,6 +348,14 @@ def build_index(max_pages: int = 10, model: str = "text-embedding-3-small", use_
             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)
     
@@ -530,11 +551,14 @@ if __name__ == "__main__":
                        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
+        use_faiss=not args.no_faiss,
+        override_adult_filter=args.include_adult
     )

app/main.py

@@ -10,6 +10,15 @@ from typing import List, Optional
 import logging
 import time
 
+# 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=os.getenv("LOG_LEVEL", "INFO").upper())
 logger = logging.getLogger(__name__)
@@ -208,6 +217,25 @@ async def health_check():
     """Health check endpoint"""
     return {"status": "healthy", "vectors_loaded": vectors is not None}
 
+async def get_movie_from_tmdb(tmdb_id: int):
+    """Fetch a single movie from TMDB API if not in local index"""
+    try:
+        settings = get_settings()
+        import requests
+        
+        url = f"https://api.themoviedb.org/3/movie/{tmdb_id}"
+        params = {"api_key": settings.tmdb_api_key}
+        
+        response = requests.get(url, params=params, timeout=10)
+        if response.status_code == 200:
+            return response.json()
+        else:
+            logger.warning(f"TMDB API returned {response.status_code} for movie {tmdb_id}")
+            return None
+    except Exception as e:
+        logger.error(f"Error fetching movie {tmdb_id} from TMDB: {e}")
+        return None
+
 @app.post("/explore", response_model=ExploreResponse)
 async def explore(
     request: ExploreRequest,
@@ -219,15 +247,32 @@ async def explore(
     start_time = time.time()
     
     try:
+        # Ensure top_k doesn't exceed available movies
+        total_movies = len(vectors) if vectors is not None else 0
+        actual_top_k = min(request.top_k, total_movies)
+        
+        if actual_top_k <= 0:
+            raise HTTPException(status_code=400, detail="No movies available")
+        
         # Convert TMDB IDs to internal indices
         liked_indices = []
         disliked_indices = []
+        missing_movies = []
         
         for tmdb_id in request.liked_ids:
             if str(tmdb_id) in id_map:
                 liked_indices.append(id_map[str(tmdb_id)])
             else:
                 logger.warning(f"TMDB ID {tmdb_id} not found in index")
+                # Optionally fetch movie info for debugging
+                movie_info = await get_movie_from_tmdb(tmdb_id)
+                if movie_info:
+                    missing_movies.append({
+                        "id": tmdb_id,
+                        "title": movie_info.get("title", "Unknown"),
+                        "release_date": movie_info.get("release_date", "Unknown")
+                    })
+                    logger.info(f"Missing movie: {movie_info.get('title')} ({movie_info.get('release_date', 'Unknown')})")
         
         for tmdb_id in request.disliked_ids:
             if str(tmdb_id) in id_map:
@@ -235,6 +280,10 @@ async def explore(
             else:
                 logger.warning(f"TMDB ID {tmdb_id} not found in index")
         
+        # Log missing movies for debugging
+        if missing_movies:
+            logger.info(f"Missing {len(missing_movies)} movies from index: {[m['title'] for m in missing_movies]}")
+        
         # Get embedding vectors
         liked_vectors = vectors[liked_indices] if liked_indices else None
         disliked_vectors = vectors[disliked_indices] if disliked_indices else None
@@ -251,9 +300,14 @@ async def explore(
         # Compute distances to subspace (residuals)
         residuals = np.linalg.norm(vectors - reconstructed, axis=1)
         
-        # Get top-k closest movies
-        top_k_indices = np.argpartition(residuals, min(request.top_k, len(residuals)))[:request.top_k]
-        top_k_indices = top_k_indices[np.argsort(residuals[top_k_indices])]
+        # Get top-k closest movies - use proper bounds checking
+        if actual_top_k >= len(residuals):
+            # If we want all movies, just sort them
+            top_k_indices = np.argsort(residuals)
+        else:
+            # Use argpartition for efficiency when we want a subset
+            top_k_indices = np.argpartition(residuals, actual_top_k-1)[:actual_top_k]
+            top_k_indices = top_k_indices[np.argsort(residuals[top_k_indices])]
         
         # Assign spiral coordinates
         spiral_coords = assign_spiral_coords(len(top_k_indices))
@@ -290,7 +344,7 @@ async def explore(
         )
         
         elapsed = time.time() - start_time
-        logger.info(f"Explore request processed in {elapsed:.3f}s - {len(request.liked_ids)} likes, {len(request.disliked_ids)} dislikes, {len(movies)} results")
+        logger.info(f"Explore request processed in {elapsed:.3f}s - {len(request.liked_ids)} likes ({len(liked_indices)} found), {len(request.disliked_ids)} dislikes ({len(disliked_indices)} found), {len(movies)} results")
         
         return response
         

app/settings.py

@@ -24,6 +24,9 @@ class Settings(BaseSettings):
     # Logging level
     log_level: str = "INFO"
     
+    # Filter adult content (True = exclude adult films, False = include all)
+    filter_adult_content: bool = True
+
     class Config:
         env_file = ".env"
         env_file_encoding = "utf-8"