src/recommendation_system.py

import chromadb
from chromadb.utils import embedding_functions
import json
import os
import re
from langchain_openai import ChatOpenAI
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from langchain.memory import ConversationBufferMemory
from tmdb_api_helper import TMDBHelper


class MovieRecommender:
    def __init__(self):
        # Initialize ChromaDB client
        self.chroma_client = chromadb.PersistentClient(path="data/embeddings")

        embedding_function = embedding_functions.SentenceTransformerEmbeddingFunction(
            model_name='all-MiniLM-L6-v2'
        )
        try:
            self.collection = self.chroma_client.get_collection(
                name="movie_collection",
                embedding_function=embedding_function,
            )
        except Exception:
            self._initialize_database()
            self.collection = self.chroma_client.get_collection(
                name="movie_collection",
                embedding_function=embedding_function,
            )

        # Initialize the language model and memory
        self.memory = ConversationBufferMemory(
            memory_key="chat_history",
            input_key="human_input"
        )
        self.llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0.7, max_tokens=1024)


        # User preferences storage
        self.user_preferences_file = "data/user_preferences.json"
        self.user_preferences = self._load_user_preferences()
        
        # Initialize TMDB helper
        self.tmdb_helper = TMDBHelper()

        # Setup prompt templates
        self._setup_prompts()

    def _initialize_database(self):
        """Download MovieLens data and build the ChromaDB collection on first run."""
        print("First run: downloading data and building movie database (this takes a few minutes)...")
        import os
        sys_path_entry = os.path.dirname(__file__)
        import sys
        if sys_path_entry not in sys.path:
            sys.path.insert(0, sys_path_entry)
        from movie_data_preparation import download_and_prepare_movielens
        from vector_database_setup import prepare_movie_descriptions, create_vector_database
        import ast

        movies_df = download_and_prepare_movielens()
        movies_df['genres'] = movies_df['genres'].apply(
            lambda g: ast.literal_eval(g) if isinstance(g, str) else g
        )
        movies_df = prepare_movie_descriptions(movies_df)
        create_vector_database(movies_df)
        print("Movie database ready.")

    def _setup_prompts(self):
        self.recommendation_template = """
        You are MovieMind, a friendly movie recommendation assistant.

        Previous conversation:
        {chat_history}

        User query: {human_input}

        Based on the query, these are the most relevant movies from our database:
        {movie_results}

        User preferences: {user_preferences}

        Your task:
        1. Analyze the user query and the provided movie results.
        2. Consider the user's preferences if available.
        3. Provide personalized movie recommendations from the list above.
        4. For each recommendation, include:
            - Title: (Movie Title)
            - Year: (Year)
            - Genre: (Genre)
            - Director: (Director)
            - Main actors: (Actors)
            - Short plot summary
            - Reason why you recommend it.
        5. Always provide this information in the exact format shown above.
        6. Make sure each movie is clearly separated from others.
        7. If appropriate, ask a follow-up question to refine future recommendations.

        Give at least 5 movie recommendations, not just a single movie.
        Respond in a conversational, helpful tone. Avoid phrases like "based on your query" or "personalized recommendations".

        AI Assistant:
        """

        self.recommendation_chain = LLMChain(
            llm=self.llm,
            prompt=PromptTemplate(
                input_variables=["chat_history", "human_input", "movie_results", "user_preferences"],
                template=self.recommendation_template
            ),
            memory=self.memory,
            verbose=False
        )

        self.general_template = """
        You are MovieMind, a friendly movie recommendation assistant.

        Previous conversation:
        {chat_history}

        Human: {human_input}

        Respond in a conversational, helpful tone.
        If the user is asking about movies or for recommendations,
        suggest they try asking more specifically about genres, actors, or the type of movie they're looking for.

        AI Assistant:
        """

        self.general_chain = LLMChain(
            llm=self.llm,
            prompt=PromptTemplate(
                input_variables=["chat_history", "human_input"],
                template=self.general_template
            ),
            memory=self.memory,
            verbose=False
        )

    def _load_user_preferences(self):
        """Load user preferences from file"""
        if os.path.exists(self.user_preferences_file):
            with open(self.user_preferences_file, 'r') as f:
                return json.load(f)
        else:
            os.makedirs('data', exist_ok=True)
            with open(self.user_preferences_file, 'w') as f:
                json.dump({}, f)
            return {}

    def _save_user_preferences(self):
        """Save user preferences to file"""
        with open(self.user_preferences_file, 'w') as f:
            json.dump(self.user_preferences, f, indent=2)

    """"
    def update_preferences(self, user_id, movie_id, liked=True):

        if user_id not in self.user_preferences:
            self.user_preferences[user_id] = {"favorites": []}

        prefs = self.user_preferences[user_id]

        self._save_user_preferences()
        """

    def get_response(self, user_id, message):
        """Generate a recommendation or general response based on user input"""
        # Step 1: Search for relevant movies
        results = self.collection.query(
            query_texts=[message],
            n_results=5
        )

        movie_results = results.get("documents", [[]])[0]

        # Step 2: Prepare movie descriptions
        movie_descriptions = ""
        for movie in movie_results:
            if isinstance(movie, str) and movie.strip():
                try:
                    movie_info = json.loads(movie)
                except json.JSONDecodeError:
                    continue
            else:
                movie_info = movie

            if movie_info:
                movie_descriptions += f"\nTitle: {movie_info.get('title', 'Unknown')}\n"
                movie_descriptions += f"Year: {movie_info.get('year', 'Unknown')}\n"
                movie_descriptions += f"Genre: {movie_info.get('genre', 'Unknown')}\n"
                movie_descriptions += f"Director: {movie_info.get('director', 'Unknown')}\n"
                actors = movie_info.get('actors', [])
                if isinstance(actors, list):
                    movie_descriptions += f"Actors: {', '.join(actors)}\n"
                else:
                    movie_descriptions += f"Actors: {actors}\n"
                movie_descriptions += f"Plot: {movie_info.get('plot', 'No plot available')}\n\n"

        # Step 3: Load user preferences
        user_prefs = self.user_preferences.get(user_id, {})
        favorites = user_prefs.get("favorites", [])
        
        if favorites:
            user_preferences_string = (
                f"Favorite movies: {', '.join(favorites)}."
            )
        else:
            user_preferences_string = "No preferences recorded yet."

        # Step 4: Create final response
        try:
            
            response = self.recommendation_chain.invoke({
                "chat_history": self.memory.buffer,
                "human_input": message,
                "movie_results": movie_descriptions,
                "user_preferences": user_preferences_string
            })
            
            # Extract the text response
            if isinstance(response, dict) and "text" in response:
                response = response["text"]
            
        except Exception as e:
            print(f"Error generating recommendation: {e}")
            # Fallback to general response
            try:
                response = self.general_chain.invoke({
                    "chat_history": self.memory.buffer,
                    "human_input": message
                })
                
                if isinstance(response, dict) and "text" in response:
                    response = response["text"]
                    
            except Exception as e2:
                print(f"Error generating general response: {e2}")
                response = "I'm having trouble generating a recommendation right now. Could you try again or ask in a different way?"

        # Step 5: Process the response to add poster data
        processed_response = self.process_response_with_posters(response)
        return processed_response
        
    def process_response_with_posters(self, response):
        """Process the response to add movie poster data"""
        # Split the response into paragraphs
        paragraphs = response.split("\n\n")
        result = ""
        
        # Process each paragraph
        for paragraph in paragraphs:
            # Skip empty paragraphs
            if not paragraph.strip():
                continue
                
            # Check if this paragraph contains movie information
            if "Title:" in paragraph or re.search(r'\b\(\d{4}\)\b', paragraph):
                # Extract movie title
                title_match = re.search(r"Title:\s*(.*?)(?:\n|$)", paragraph)
                if not title_match:
                    # Try to find title in format "Movie Title (Year)"
                    title_match = re.search(r"(.*?)\s*\(\d{4}\)", paragraph)
                
                if title_match:
                    movie_title = title_match.group(1).strip()
                    
                    # Extract year if available
                    year_match = re.search(r"Year:\s*(\d{4})", paragraph)
                    year = year_match.group(1) if year_match else None
                    
                    if not year:
                        # Try to find year in format "Movie Title (Year)"
                        year_match = re.search(r"\((\d{4})\)", paragraph)
                        year = year_match.group(1) if year_match else None
                    
                    # Get poster URL
                    poster_url = self.tmdb_helper.get_poster_url(movie_title, year)
                    
                    # Add poster URL to the response
                    if poster_url:
                        result += paragraph + f"\n[POSTER_URL: {poster_url}]\n\n"
                    else:
                        result += paragraph + "\n\n"
                else:
                    result += paragraph + "\n\n"
            else:
                result += paragraph + "\n\n"
        
        return result