Add CachedGraph model and repository for NetworkX graph caching

packages/django-app/app/discordbot/bacon.py

@@ -1,9 +1,7 @@
 import discord
 import networkx as nx
-import pandas as pd
-from asgiref.sync import sync_to_async
 from discord import app_commands
-from plex.models import PlexMovie
+from plex.commands import GetActorGraphCommand
 
 
 @app_commands.command(name="bacon", description="Shows hops between actors.")
@@ -16,80 +14,28 @@ async def bacon(
     with_producers: bool = False,
     with_writers: bool = False,
 ):
-    # munge input
+    """Find the shortest path between two people through movies."""
+    # Normalize input
     from_actor = from_actor.strip().lower()
     to_actor = to_actor.strip().lower()
 
-    movies = await sync_to_async(list)(PlexMovie.objects.all().values())
-    df = pd.DataFrame(movies)
+    # Defer response since this might take a moment
+    await interaction.response.defer()
 
-    graph = nx.Graph()
-    added_actors = []
-    added_directors = []
-    added_producers = []
-    added_writers = []
-
-    def add_movie_and_actors_to_graph(movie):
-        year = int(movie.year) if movie.year > 0 else None
-        year_string = f" ({year})" if year else None
-        movie_title = f"{movie.title}{year_string}"
-        graph.add_node(movie_title, type="movie", color="red")
-
-        for actor in movie.actors:
-            if actor not in added_actors:
-                actor = actor.lower()
-                graph.add_node(
-                    actor, type="actor", color="blue" if actor == from_actor else "green"
-                )
-                added_actors.append(actor)
-            graph.add_edge(movie_title, actor)
-
-        if with_directors:
-            for director in movie.directors:
-                if director not in added_directors:
-                    director = director.lower()
-                    graph.add_node(
-                        director,
-                        type="directors",
-                        color="yellow" if director == from_actor else "green",
-                    )
-                    added_directors.append(director)
-                graph.add_edge(movie_title, director)
-
-        if with_producers:
-            for producer in movie.producers:
-                if producer not in added_producers:
-                    producer = producer.lower()
-                    graph.add_node(
-                        producer,
-                        type="producers",
-                        color="purple" if producer == from_actor else "green",
-                    )
-                    added_producers.append(producer)
-                graph.add_edge(movie_title, producer)
-
-        if with_writers:
-            for writer in movie.writers:
-                if writer not in added_writers:
-                    writer = writer.lower()
-                    graph.add_node(
-                        writer, type="writers", color="red" if writer == from_actor else "green"
-                    )
-                    added_writers.append(writer)
-                graph.add_edge(movie_title, writer)
-
-    _ = df.apply(lambda m: add_movie_and_actors_to_graph(m), axis=1)
+    # Get actor graph (from cache or build if needed)
+    command = GetActorGraphCommand()
+    graph = await command.execute()
 
     try:
         path = nx.shortest_path(graph, source=from_actor, target=to_actor)
 
-        # build message
+        # Build message
         words_list = []
         hops = 0
 
         for idx, entry in enumerate(path):
             if idx == 0:
-                # from actor
+                # from person
                 words_list.append(entry.title())
                 words_list.append("worked on")
             elif idx % 2 != 0:
@@ -98,20 +44,18 @@ def add_movie_and_actors_to_graph(movie):
                 words_list.append(entry)
                 words_list.append("with")
             elif idx % 2 == 0 and idx != len(path) - 1:
-                # an actor
+                # an intermediary person
                 words_list.append(entry.title())
                 words_list.append("who worked on")
             elif idx == len(path) - 1:
-                # last actor
+                # target person
                 words_list.append(f"{entry.title()}.")
 
         message = f"{from_actor} and {to_actor} are connected by {hops} hops.\n" + " ".join(
             words_list
         )
-        await interaction.response.send_message(message)
-    except nx.NetworkXNoPath as exception:
-        await interaction.response.send_message(exception)
-    except nx.NodeNotFound as exception:
-        await interaction.response.send_message(exception)
-
-    await interaction.response.send_message("gonna queue up something good!")
+        await interaction.followup.send(message)
+    except nx.NetworkXNoPath:
+        await interaction.followup.send(f"No path found between {from_actor} and {to_actor}.")
+    except nx.NodeNotFound as e:
+        await interaction.followup.send(f"Person not found: {str(e)}")

packages/django-app/app/plex/commands.py

@@ -1,13 +1,15 @@
 import asyncio
 import logging
 
+import networkx as nx
 from aiohttp import ClientSession
+from asgiref.sync import sync_to_async
 from common.commands.abstract_base_command import AbstractBaseCommand
 from services.plex import Plex
 from services.tmdb import TMDB
 
 from plex.forms import EnrichMovieActorsForm
-from plex.repositories import PlexMovieRepository
+from plex.repositories import CachedGraphRepository, PlexMovieRepository
 
 logger = logging.getLogger(__name__)
 
@@ -26,13 +28,14 @@ def execute(self) -> None:
         super().execute()
 
         latest_movie = PlexMovieRepository.get_latest()
+        synced_count = 0
 
         for movie in Plex.fetch_movies(sort="addedAt:desc", container_start=0, container_size=5):
             added_at = Plex.normalize_added_at(movie.addedAt)
             if latest_movie and added_at <= latest_movie.created_at:
                 # break out of loop if we start to get a movie
                 # added before the latest movie in the database
-                return
+                break
 
             try:
                 movie_details = Plex.extract_movie_details(movie)
@@ -59,10 +62,22 @@ def execute(self) -> None:
                         plex_movie.save()
 
                 print(f"Created PlexMovie: {plex_movie}")
+                synced_count += 1
             except Exception as e:
                 logger.exception(f"Failed to create PlexMovie: {movie}")
                 logger.exception(e)
 
+        # Rebuild actor graph cache after syncing
+        if synced_count > 0:
+            logger.info(f"Synced {synced_count} movies, rebuilding actor graph cache...")
+            try:
+                command = BuildActorGraphCommand()
+                graph = asyncio.run(command.execute())
+                CachedGraphRepository.save_actor_graph(graph)
+                logger.info("Actor graph cache rebuilt successfully")
+            except Exception as e:
+                logger.exception(f"Failed to rebuild actor graph cache: {e}")
+
 
 class EnrichMovieActorsCommand(AbstractBaseCommand):
     """
@@ -127,3 +142,81 @@ async def execute(self):
             except Exception as e:
                 logger.exception(f"Error enriching actors for {movie.title}: {e}")
                 raise
+
+
+class BuildActorGraphCommand(AbstractBaseCommand):
+    """
+    Command to build a NetworkX graph of all movies and their associated people.
+    Includes actors, directors, producers, and writers.
+    """
+
+    async def execute(self) -> nx.Graph:
+        """
+        Build and return a NetworkX graph of movie-person relationships.
+
+        Returns:
+            nx.Graph: Complete graph of all movie-person relationships
+        """
+        super().execute()
+
+        logger.info("Building actor graph from database...")
+
+        movies = await sync_to_async(list)(PlexMovieRepository.model.objects.all().values())
+
+        graph = nx.Graph()
+        added_people = set()
+
+        for movie_dict in movies:
+            # Add movie node
+            year = movie_dict.get("year")
+            year = int(year) if year and year > 0 else None
+            year_string = f" ({year})" if year else ""
+            movie_title = f"{movie_dict['title']}{year_string}"
+            graph.add_node(movie_title, type="movie")
+
+            # Add all people (actors, directors, producers, writers)
+            for person_type in ["actors", "directors", "producers", "writers"]:
+                people = movie_dict.get(person_type) or []
+                for person in people:
+                    person_lower = person.lower()
+                    if person_lower not in added_people:
+                        graph.add_node(person_lower, type="person")
+                        added_people.add(person_lower)
+                    graph.add_edge(movie_title, person_lower)
+
+        logger.info(
+            f"Built graph with {graph.number_of_nodes()} nodes and {graph.number_of_edges()} edges"
+        )
+
+        return graph
+
+
+class GetActorGraphCommand(AbstractBaseCommand):
+    """
+    Command to get the actor graph, either from cache or by building it.
+    """
+
+    async def execute(self) -> nx.Graph:
+        """
+        Get the actor graph from cache, or build and cache it if not available.
+
+        Returns:
+            nx.Graph: Complete graph of all movie-person relationships
+        """
+        super().execute()
+
+        # Try to load from cache
+        graph = await CachedGraphRepository.load_actor_graph_async()
+
+        if graph is None:
+            logger.info("Actor graph not in cache, building...")
+            # Build graph if not cached
+            command = BuildActorGraphCommand()
+            graph = await command.execute()
+            # Cache it for next time
+            await sync_to_async(CachedGraphRepository.save_actor_graph)(graph)
+            logger.info("Actor graph cached successfully")
+        else:
+            logger.info("Loaded actor graph from cache")
+
+        return graph

packages/django-app/app/plex/models.py

@@ -55,3 +55,33 @@ class Meta:
 
     def __str__(self):
         return f"{self.title} ({self.year})"
+
+
+class CachedGraph(models.Model):
+    """
+    Model for storing pickled NetworkX graphs in the database.
+    Generic storage for any graph type (actor, producer, director, etc.).
+    """
+
+    key = models.CharField(
+        max_length=255,
+        unique=True,
+        primary_key=True,
+        help_text=_("Unique identifier for the cached graph (e.g., 'actor_graph')"),  # type: ignore[arg-type]
+    )
+
+    data = models.BinaryField(
+        help_text=_("Pickled graph data stored as binary")  # type: ignore[arg-type]
+    )
+
+    updated_at = models.DateTimeField(
+        auto_now=True,
+        help_text=_("Timestamp of last cache update"),  # type: ignore[arg-type]
+    )
+
+    class Meta:
+        db_table = "cached_graphs"
+        default_permissions = ()
+
+    def __str__(self):
+        return f"CachedGraph: {self.key}"

packages/django-app/app/plex/repositories.py

@@ -1,9 +1,11 @@
+import pickle
+
 from asgiref.sync import sync_to_async
 from common.repositories.base_repository import BaseRepository
 from django.db.models import Q
 from django.utils import timezone
 
-from plex.models import PlexMovie
+from plex.models import CachedGraph, PlexMovie
 
 
 class PlexMovieRepository(BaseRepository):
@@ -147,3 +149,54 @@ async def update_movie_actors_async(cls, movie, actors, tmdb_id=None):
 
         await sync_to_async(movie.save)()
         return movie
+
+
+class CachedGraphRepository(BaseRepository):
+    """
+    Repository for managing cached NetworkX graphs.
+    Generic repository that can cache any type of graph (actors, producers, directors, etc.).
+    """
+
+    model = CachedGraph
+
+    @classmethod
+    def save_actor_graph(cls, graph):
+        """
+        Save actor relationship graph to cache.
+
+        Args:
+            graph: NetworkX Graph object containing actor-movie relationships
+
+        Returns:
+            CachedGraph instance
+        """
+        graph_data = pickle.dumps(graph)
+        cached_graph, _ = cls.model.objects.update_or_create(  # type: ignore[attr-defined]
+            key="actor_graph", defaults={"data": graph_data}
+        )
+        return cached_graph
+
+    @classmethod
+    def load_actor_graph(cls):
+        """
+        Load actor relationship graph from cache.
+
+        Returns:
+            NetworkX Graph object if found, None otherwise
+        """
+        try:
+            cached = cls.model.objects.get(key="actor_graph")  # type: ignore[attr-defined]
+            return pickle.loads(cached.data)
+        except cls.model.DoesNotExist:  # type: ignore[attr-defined]
+            return None
+
+    @classmethod
+    @sync_to_async
+    def load_actor_graph_async(cls):
+        """
+        Async version of load_actor_graph.
+
+        Returns:
+            NetworkX Graph object if found, None otherwise
+        """
+        return cls.load_actor_graph()