langchain/templates/neo4j-semantic-layer/neo4j_semantic_layer/recommendation_tool.py

from typing import Optional, Type

from langchain.callbacks.manager import (
    AsyncCallbackManagerForToolRun,
    CallbackManagerForToolRun,
)
from langchain.pydantic_v1 import BaseModel, Field
from langchain.tools import BaseTool

from neo4j_semantic_layer.utils import get_candidates, get_user_id, graph

recommendation_query_db_history = """
  MERGE (u:User {userId:$user_id})
  WITH u
  // get recommendation candidates
  OPTIONAL MATCH (u)-[r1:RATED]->()<-[r2:RATED]-()-[r3:RATED]->(recommendation)
  WHERE r1.rating > 3.5 AND r2.rating > 3.5 AND r3.rating > 3.5
        AND NOT EXISTS {(u)-[:RATED]->(recommendation)}
  // rank and limit recommendations
  WITH u, recommendation, count(*) AS count
  ORDER BY count DESC LIMIT 3
  RETURN recommendation.title AS movie
"""

recommendation_query_genre = """
MATCH (m:Movie)-[:IN_GENRE]->(g:Genre {name:$genre})
// filter out already seen movies by the user
WHERE NOT EXISTS {
  (m)<-[:RATED]-(:User {userId:$user_id})
}
// rank and limit recommendations
WITH m
ORDER BY m.imdbRating DESC LIMIT 3
RETURN m.title AS movie
"""


def recommendation_query_movie(genre: bool) -> str:
    return f"""
MATCH (m1:Movie)<-[r1:RATED]-()-[r2:RATED]->(m2:Movie)
WHERE r1.rating > 3.5 AND r2.rating > 3.5 and m1.title IN $movieTitles
// filter out already seen movies by the user
AND NOT EXISTS {{
  (m2)<-[:RATED]-(:User {{userId:$user_id}})
}}
{'AND EXISTS {(m2)-[:IN_GENRE]->(:Genre {name:$genre})}' if genre else ''}
// rank and limit recommendations
WITH m2, count(*) AS count
ORDER BY count DESC LIMIT 3
RETURN m2.title As movie
"""


def recommend_movie(movie: Optional[str] = None, genre: Optional[str] = None) -> str:
    """
    Recommends movies based on user's history and preference
    for a specific movie and/or genre.
    Returns:
        str: A string containing a list of recommended movies, or an error message.
    """
    user_id = get_user_id()
    params = {"user_id": user_id, "genre": genre}
    if not movie and not genre:
        # Try to recommend a movie based on the information in the db
        response = graph.query(recommendation_query_db_history, params)
        try:
            return ", ".join([el["movie"] for el in response])
        except Exception:
            return "Can you tell us about some of the movies you liked?"
    if not movie and genre:
        # Recommend top voted movies in the genre the user haven't seen before
        response = graph.query(recommendation_query_genre, params)
        try:
            return ", ".join([el["movie"] for el in response])
        except Exception:
            return "Something went wrong"

    candidates = get_candidates(movie, "movie")
    if not candidates:
        return "The movie you mentioned wasn't found in the database"
    params["movieTitles"] = [el["candidate"] for el in candidates]
    query = recommendation_query_movie(bool(genre))
    response = graph.query(query, params)
    try:
        return ", ".join([el["movie"] for el in response])
    except Exception:
        return "Something went wrong"


all_genres = [
    "Action",
    "Adventure",
    "Animation",
    "Children",
    "Comedy",
    "Crime",
    "Documentary",
    "Drama",
    "Fantasy",
    "Film-Noir",
    "Horror",
    "IMAX",
    "Musical",
    "Mystery",
    "Romance",
    "Sci-Fi",
    "Thriller",
    "War",
    "Western",
]


class RecommenderInput(BaseModel):
    movie: Optional[str] = Field(description="movie used for recommendation")
    genre: Optional[str] = Field(
        description=(
            "genre used for recommendation. Available options are:" f"{all_genres}"
        )
    )


class RecommenderTool(BaseTool):
    name = "Recommender"
    description = "useful for when you need to recommend a movie"
    args_schema: Type[BaseModel] = RecommenderInput

    def _run(
        self,
        movie: Optional[str] = None,
        genre: Optional[str] = None,
        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> str:
        """Use the tool."""
        return recommend_movie(movie, genre)

    async def _arun(
        self,
        movie: Optional[str] = None,
        genre: Optional[str] = None,
        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
    ) -> str:
        """Use the tool asynchronously."""
        return recommend_movie(movie, genre)