recommendation-systems.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Implementing an LLM-powered recommendation system"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>MAL_ID</th>\n",
       "      <th>Name</th>\n",
       "      <th>Score</th>\n",
       "      <th>Genres</th>\n",
       "      <th>sypnopsis</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>Cowboy Bebop</td>\n",
       "      <td>8.78</td>\n",
       "      <td>Action, Adventure, Comedy, Drama, Sci-Fi, Space</td>\n",
       "      <td>In the year 2071, humanity has colonized sever...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>5</td>\n",
       "      <td>Cowboy Bebop: Tengoku no Tobira</td>\n",
       "      <td>8.39</td>\n",
       "      <td>Action, Drama, Mystery, Sci-Fi, Space</td>\n",
       "      <td>other day, another bounty—such is the life of ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>6</td>\n",
       "      <td>Trigun</td>\n",
       "      <td>8.24</td>\n",
       "      <td>Action, Sci-Fi, Adventure, Comedy, Drama, Shounen</td>\n",
       "      <td>Vash the Stampede is the man with a $$60,000,0...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>7</td>\n",
       "      <td>Witch Hunter Robin</td>\n",
       "      <td>7.27</td>\n",
       "      <td>Action, Mystery, Police, Supernatural, Drama, ...</td>\n",
       "      <td>ches are individuals with special powers like ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>8</td>\n",
       "      <td>Bouken Ou Beet</td>\n",
       "      <td>6.98</td>\n",
       "      <td>Adventure, Fantasy, Shounen, Supernatural</td>\n",
       "      <td>It is the dark century and the people are suff...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   MAL_ID                             Name Score  \\\n",
       "0       1                     Cowboy Bebop  8.78   \n",
       "1       5  Cowboy Bebop: Tengoku no Tobira  8.39   \n",
       "2       6                           Trigun  8.24   \n",
       "3       7               Witch Hunter Robin  7.27   \n",
       "4       8                   Bouken Ou Beet  6.98   \n",
       "\n",
       "                                              Genres  \\\n",
       "0    Action, Adventure, Comedy, Drama, Sci-Fi, Space   \n",
       "1              Action, Drama, Mystery, Sci-Fi, Space   \n",
       "2  Action, Sci-Fi, Adventure, Comedy, Drama, Shounen   \n",
       "3  Action, Mystery, Police, Supernatural, Drama, ...   \n",
       "4          Adventure, Fantasy, Shounen, Supernatural   \n",
       "\n",
       "                                           sypnopsis  \n",
       "0  In the year 2071, humanity has colonized sever...  \n",
       "1  other day, another bounty—such is the life of ...  \n",
       "2  Vash the Stampede is the man with a $$60,000,0...  \n",
       "3  ches are individuals with special powers like ...  \n",
       "4  It is the dark century and the people are suff...  "
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import tiktoken\n",
    "\n",
    "import lancedb\n",
    "\n",
    "from openai import OpenAI\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "from langchain.chains import RetrievalQA\n",
    "from langchain.prompts import PromptTemplate\n",
    "\n",
    "from langchain_community.callbacks import get_openai_callback\n",
    "from langchain_community.vectorstores import LanceDB\n",
    "\n",
    "openai_api_key = \"\"\n",
    "client = OpenAI(api_key=openai_api_key)\n",
    "\n",
    "anime = pd. read_csv('data/anime_with_synopsis.csv')\n",
    "anime.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "anime = anime.dropna()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"Title: Monster. Overview: Dr. Kenzou Tenma, an elite neurosurgeon recently engaged to his hospital director's daughter, is well on his way to ascending the hospital hierarchy. That is until one night, a seemingly small event changes Dr. Tenma's life forever. While preparing to perform surgery on someone, he gets a call from the hospital director telling him to switch patients and instead perform life-saving brain surgery on a famous performer. His fellow doctors, fiancée, and the hospital director applaud his accomplishment; but because of the switch, a poor immigrant worker is dead, causing Dr. Tenma to have a crisis of conscience. So when a similar situation arises, Dr. Tenma stands his ground and chooses to perform surgery on the young boy Johan Liebert instead of the town's mayor. Unfortunately, this choice leads to serious ramifications for Dr. Tenma—losing his social standing being one of them. However, with the mysterious death of the director and two other doctors, Dr. Tenma's position is restored. With no evidence to convict him, he is released and goes on to attain the position of hospital director. Nine years later when Dr. Tenma saves the life of a criminal, his past comes back to haunt him—once again, he comes face to face with the monster he operated on. He must now embark on a quest of pursuit to make amends for the havoc spread by the one he saved. Genres: Drama, Horror, Mystery, Police, Psychological, Seinen, Thriller\""
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "anime['combined_info'] = anime.apply(lambda row: f\"Title: {row['Name']}. Overview: {row['sypnopsis']} Genres: {row['Genres']}\", axis=1)\n",
    "anime['combined_info'][0]\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "16206"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embedding_model = \"text-embedding-ada-002\"\n",
    "embedding_encoding = \"cl100k_base\"  # this the encoding for text-embedding-ada-002\n",
    "max_tokens = 8000  # the maximum for text-embedding-ada-002 is 8191\n",
    "\n",
    "encoding = tiktoken.get_encoding(embedding_encoding)\n",
    "\n",
    "# omit descriptions that are too long to embed\n",
    "anime[\"n_tokens\"] = anime.combined_info.apply(lambda x: len(encoding.encode(x)))\n",
    "anime = anime[anime.n_tokens <= max_tokens]\n",
    "len(anime)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>MAL_ID</th>\n",
       "      <th>Name</th>\n",
       "      <th>Score</th>\n",
       "      <th>Genres</th>\n",
       "      <th>sypnopsis</th>\n",
       "      <th>combined_info</th>\n",
       "      <th>n_tokens</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>Cowboy Bebop</td>\n",
       "      <td>8.78</td>\n",
       "      <td>Action, Adventure, Comedy, Drama, Sci-Fi, Space</td>\n",
       "      <td>In the year 2071, humanity has colonized sever...</td>\n",
       "      <td>Title: Cowboy Bebop. Overview: In the year 207...</td>\n",
       "      <td>245</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>5</td>\n",
       "      <td>Cowboy Bebop: Tengoku no Tobira</td>\n",
       "      <td>8.39</td>\n",
       "      <td>Action, Drama, Mystery, Sci-Fi, Space</td>\n",
       "      <td>other day, another bounty—such is the life of ...</td>\n",
       "      <td>Title: Cowboy Bebop: Tengoku no Tobira. Overvi...</td>\n",
       "      <td>199</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>6</td>\n",
       "      <td>Trigun</td>\n",
       "      <td>8.24</td>\n",
       "      <td>Action, Sci-Fi, Adventure, Comedy, Drama, Shounen</td>\n",
       "      <td>Vash the Stampede is the man with a $$60,000,0...</td>\n",
       "      <td>Title: Trigun. Overview: Vash the Stampede is ...</td>\n",
       "      <td>252</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>7</td>\n",
       "      <td>Witch Hunter Robin</td>\n",
       "      <td>7.27</td>\n",
       "      <td>Action, Mystery, Police, Supernatural, Drama, ...</td>\n",
       "      <td>ches are individuals with special powers like ...</td>\n",
       "      <td>Title: Witch Hunter Robin. Overview: ches are ...</td>\n",
       "      <td>125</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>8</td>\n",
       "      <td>Bouken Ou Beet</td>\n",
       "      <td>6.98</td>\n",
       "      <td>Adventure, Fantasy, Shounen, Supernatural</td>\n",
       "      <td>It is the dark century and the people are suff...</td>\n",
       "      <td>Title: Bouken Ou Beet. Overview: It is the dar...</td>\n",
       "      <td>188</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   MAL_ID                             Name Score  \\\n",
       "0       1                     Cowboy Bebop  8.78   \n",
       "1       5  Cowboy Bebop: Tengoku no Tobira  8.39   \n",
       "2       6                           Trigun  8.24   \n",
       "3       7               Witch Hunter Robin  7.27   \n",
       "4       8                   Bouken Ou Beet  6.98   \n",
       "\n",
       "                                              Genres  \\\n",
       "0    Action, Adventure, Comedy, Drama, Sci-Fi, Space   \n",
       "1              Action, Drama, Mystery, Sci-Fi, Space   \n",
       "2  Action, Sci-Fi, Adventure, Comedy, Drama, Shounen   \n",
       "3  Action, Mystery, Police, Supernatural, Drama, ...   \n",
       "4          Adventure, Fantasy, Shounen, Supernatural   \n",
       "\n",
       "                                           sypnopsis  \\\n",
       "0  In the year 2071, humanity has colonized sever...   \n",
       "1  other day, another bounty—such is the life of ...   \n",
       "2  Vash the Stampede is the man with a $$60,000,0...   \n",
       "3  ches are individuals with special powers like ...   \n",
       "4  It is the dark century and the people are suff...   \n",
       "\n",
       "                                       combined_info  n_tokens  \n",
       "0  Title: Cowboy Bebop. Overview: In the year 207...       245  \n",
       "1  Title: Cowboy Bebop: Tengoku no Tobira. Overvi...       199  \n",
       "2  Title: Trigun. Overview: Vash the Stampede is ...       252  \n",
       "3  Title: Witch Hunter Robin. Overview: ches are ...       125  \n",
       "4  Title: Bouken Ou Beet. Overview: It is the dar...       188  "
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "anime.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>MAL_ID</th>\n",
       "      <th>Name</th>\n",
       "      <th>Score</th>\n",
       "      <th>Genres</th>\n",
       "      <th>sypnopsis</th>\n",
       "      <th>combined_info</th>\n",
       "      <th>n_tokens</th>\n",
       "      <th>embedding</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>Cowboy Bebop</td>\n",
       "      <td>8.78</td>\n",
       "      <td>Action, Adventure, Comedy, Drama, Sci-Fi, Space</td>\n",
       "      <td>In the year 2071, humanity has colonized sever...</td>\n",
       "      <td>Title: Cowboy Bebop. Overview: In the year 207...</td>\n",
       "      <td>245</td>\n",
       "      <td>[0.00921056978404522, -0.012633174657821655, 0...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>5</td>\n",
       "      <td>Cowboy Bebop: Tengoku no Tobira</td>\n",
       "      <td>8.39</td>\n",
       "      <td>Action, Drama, Mystery, Sci-Fi, Space</td>\n",
       "      <td>other day, another bounty—such is the life of ...</td>\n",
       "      <td>Title: Cowboy Bebop: Tengoku no Tobira. Overvi...</td>\n",
       "      <td>199</td>\n",
       "      <td>[-0.008109764195978642, -0.028518257662653923,...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>6</td>\n",
       "      <td>Trigun</td>\n",
       "      <td>8.24</td>\n",
       "      <td>Action, Sci-Fi, Adventure, Comedy, Drama, Shounen</td>\n",
       "      <td>Vash the Stampede is the man with a $$60,000,0...</td>\n",
       "      <td>Title: Trigun. Overview: Vash the Stampede is ...</td>\n",
       "      <td>252</td>\n",
       "      <td>[0.0019446373917162418, -0.001545737381093204,...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>7</td>\n",
       "      <td>Witch Hunter Robin</td>\n",
       "      <td>7.27</td>\n",
       "      <td>Action, Mystery, Police, Supernatural, Drama, ...</td>\n",
       "      <td>ches are individuals with special powers like ...</td>\n",
       "      <td>Title: Witch Hunter Robin. Overview: ches are ...</td>\n",
       "      <td>125</td>\n",
       "      <td>[-0.014938411302864552, 0.007340028416365385, ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>8</td>\n",
       "      <td>Bouken Ou Beet</td>\n",
       "      <td>6.98</td>\n",
       "      <td>Adventure, Fantasy, Shounen, Supernatural</td>\n",
       "      <td>It is the dark century and the people are suff...</td>\n",
       "      <td>Title: Bouken Ou Beet. Overview: It is the dar...</td>\n",
       "      <td>188</td>\n",
       "      <td>[0.010889030061662197, 0.0069219209253787994, ...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   MAL_ID                             Name Score  \\\n",
       "0       1                     Cowboy Bebop  8.78   \n",
       "1       5  Cowboy Bebop: Tengoku no Tobira  8.39   \n",
       "2       6                           Trigun  8.24   \n",
       "3       7               Witch Hunter Robin  7.27   \n",
       "4       8                   Bouken Ou Beet  6.98   \n",
       "\n",
       "                                              Genres  \\\n",
       "0    Action, Adventure, Comedy, Drama, Sci-Fi, Space   \n",
       "1              Action, Drama, Mystery, Sci-Fi, Space   \n",
       "2  Action, Sci-Fi, Adventure, Comedy, Drama, Shounen   \n",
       "3  Action, Mystery, Police, Supernatural, Drama, ...   \n",
       "4          Adventure, Fantasy, Shounen, Supernatural   \n",
       "\n",
       "                                           sypnopsis  \\\n",
       "0  In the year 2071, humanity has colonized sever...   \n",
       "1  other day, another bounty—such is the life of ...   \n",
       "2  Vash the Stampede is the man with a $$60,000,0...   \n",
       "3  ches are individuals with special powers like ...   \n",
       "4  It is the dark century and the people are suff...   \n",
       "\n",
       "                                       combined_info  n_tokens  \\\n",
       "0  Title: Cowboy Bebop. Overview: In the year 207...       245   \n",
       "1  Title: Cowboy Bebop: Tengoku no Tobira. Overvi...       199   \n",
       "2  Title: Trigun. Overview: Vash the Stampede is ...       252   \n",
       "3  Title: Witch Hunter Robin. Overview: ches are ...       125   \n",
       "4  Title: Bouken Ou Beet. Overview: It is the dar...       188   \n",
       "\n",
       "                                           embedding  \n",
       "0  [0.00921056978404522, -0.012633174657821655, 0...  \n",
       "1  [-0.008109764195978642, -0.028518257662653923,...  \n",
       "2  [0.0019446373917162418, -0.001545737381093204,...  \n",
       "3  [-0.014938411302864552, 0.007340028416365385, ...  \n",
       "4  [0.010889030061662197, 0.0069219209253787994, ...  "
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def get_embedding(text, model=\"text-embedding-3-small\"):\n",
    "    text = text.replace(\"\\n\", \" \")\n",
    "    return client.embeddings.create(input=[text], model=model).data[0].embedding\n",
    "\n",
    "\n",
    "anime[\"embedding\"] = anime.combined_info.apply(lambda x: get_embedding(x, model=embedding_model))\n",
    "anime.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "anime.rename(columns = {'embedding': 'vector'}, inplace = True)\n",
    "anime.rename(columns = {'combined_info': 'text'}, inplace = True)\n",
    "anime.to_pickle('data/anime.pkl')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Start working with LLMs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "uri = \"dataset/sample-anime-lancedb\"\n",
    "\n",
    "db = lancedb.connect(uri)\n",
    "table = db.create_table(\"anime\", anime)\n",
    "\n",
    "# embeddings = OpenAIEmbeddings(engine=\"text-embedding-ada-002\")\n",
    "embeddings = OpenAIEmbeddings(\n",
    "    deployment=\"SL-document_embedder\",\n",
    "    model=\"text-embedding-ada-002\",\n",
    "    show_progress_bar=True,\n",
    "    openai_api_key=openai_api_key)\n",
    "\n",
    "docsearch = LanceDB(connection=table, embedding=embeddings)\n",
    "\n",
    "# simple similarity computation\n",
    "# query = \"I'm looking for an animated action movie. What could you suggest to me?\"\n",
    "# docs = docsearch.similarity_search(query, k=1)\n",
    "\n",
    "llm = ChatOpenAI(\n",
    "    model_name=\"gpt-3.5-turbo-1106\",\n",
    "    temperature=0,\n",
    "    api_key=openai_api_key\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prompt engineering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "1. Urikupen Kyuujo-tai: This adventure comedy follows a team of brave young animals that rescues others in peril. With a dog, a boar, a deer, a koala, a mouse, a seagull, and a lion, this show is sure to please those looking for an action anime featuring animals.\n",
      "\n",
      "2. Nekketsu Jinmen Inu: Life Is Movie: This parody follows a passionate human-faced dog NEET/would be detective in his adventures. Fans of action anime with animals are sure to be engaged by this mystery story.\n",
      "\n",
      "3. Daisetsusan no Yuusha Kibaou: The main character of this drama is Fang, who was born to a hunting dog and a circus-runaway European wolf. Fang returns from the circus to face his foe, a giant brown bear which killed his family, making this story a great pick for those looking for an action anime with animals.\n"
     ]
    }
   ],
   "source": [
    "# without prompt\n",
    "# qa_chain = RetrievalQA.from_chain_type(llm=llm, chain_type=\"stuff\", retriever=docsearch.as_retriever(),\n",
    "#                                        return_source_documents=True)\n",
    "\n",
    "# let’s say we are only interested in anime that, among their genres, are tagged as “Action”.\n",
    "# df_filtered = anime[anime['Genres'].apply(lambda x: 'Action' in x)]\n",
    "# qa_chain = RetrievalQA.from_chain_type(llm=llm, chain_type=\"stuff\",\n",
    "#                                        retriever=docsearch.as_retriever(search_kwargs={'data': df_filtered}),\n",
    "#                                        return_source_documents=True)\n",
    "\n",
    "# define custom prompt\n",
    "template = \"\"\"You are a movie recommender system that help users to find anime that match their preferences. \n",
    "Use the following pieces of context to answer the question at the end. \n",
    "For each question, suggest three anime, with a short description of the plot and the reason why the user migth like it.\n",
    "If you don't know the answer, just say that you don't know, don't try to make up an answer.\n",
    "\n",
    "{context}\n",
    "\n",
    "Question: {question}\n",
    "Your response:\"\"\"\n",
    "\n",
    "PROMPT = PromptTemplate(\n",
    "    template=template, input_variables=[\"context\", \"question\"])\n",
    "\n",
    "chain_type_kwargs = {\"prompt\": PROMPT}\n",
    "\n",
    "qa_chain = RetrievalQA.from_chain_type(llm=llm,\n",
    "                                       chain_type=\"stuff\",\n",
    "                                       retriever=docsearch.as_retriever(),\n",
    "                                       return_source_documents=True,\n",
    "                                       chain_type_kwargs=chain_type_kwargs)\n",
    "\n",
    "query = \"I'm looking for an action anime. What could you suggest to me?\"\n",
    "\n",
    "# Query and Response\n",
    "with get_openai_callback() as cb:\n",
    "    result = qa_chain({\"query\": query})\n",
    "\n",
    "print(result['result'])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}