genshin_rag.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 208,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.document_loaders import WebBaseLoader\n",
    "from langchain.prompts import ChatPromptTemplate, PromptTemplate\n",
    "from langchain_chroma import Chroma\n",
    "from langchain_community.llms.ollama import Ollama\n",
    "from langchain_community.embeddings.ollama import OllamaEmbeddings\n",
    "from langchain_core.output_parsers import JsonOutputParser, StrOutputParser\n",
    "\n",
    "from typing import List\n",
    "from typing_extensions import TypedDict\n",
    "from IPython.display import Image, display\n",
    "from langchain.schema import Document\n",
    "from langgraph.graph import START, END, StateGraph\n",
    "\n",
    "import uuid\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 209,
   "metadata": {},
   "outputs": [],
   "source": [
    "CHROMA_PATH = \"genshin_db\"\n",
    "MODEL = \"llama3.1\"\n",
    "EMBEDDING_MODEL = \"mxbai-embed-large\"\n",
    "\n",
    "PROMPT_TEMPLATE = \"\"\"\n",
    "Below is a list of documents in MARKDOWN that you'll use to answer a question.\n",
    "IMPORTANT: DO NOT USE ANY OF YOUR GENERAL KNOWLEDGE TO ANSWER THE QUESTION!\n",
    "USE ONLY THE DOCUMENTS BELOW TO ANSWER THE QUESTION.\n",
    "Documents: \n",
    "{documents} \n",
    "\n",
    "IGNORE ALL DOCUMENTS WITH IRRELEVANT TITLES.\n",
    "If none of the documents are relevant or no documents are mentioned above, just say that you don't know and suggest the user to rephrase the question.\n",
    "After answering the question, you can request the user to rephrase the question if they believe the answer is incorrect.\n",
    "Be as DESCRIPTIVE as possible, using a maxiumum of six paragraphs and aggregating all information in your answer.\n",
    "DO NOT MENTION ANY DOCUMENT REFERENCES OR JUSTIFICATIONS IN YOUR ANSWER!!! USE ALL INFORMATION THAT YOU KNOW!\n",
    "Here is the user's question: {question}\n",
    "Provide your answer below: \n",
    "\"\"\"\n",
    "\n",
    "GRADE_DOC_TEMPLATE = \"\"\"You are a grader assessing the relevance of a retrieved document in MARKDOWN to a user question. \\n \n",
    "Here is the retrieved document. \\n\n",
    "Title of the document: {title}\n",
    "Content of the document: {content} \\n\\n\n",
    "Here is the user question: {question} \\n\n",
    "\n",
    "IF THE DOCUMENT TITLE IS NOT RELEVANT TO THE QUESTION AT ALL, GRADE IT AS NOT RELEVANT. \\n\n",
    "HOWEVER, IF THE DOCUMENT TITLE IS RELEVANT BUT THE DOCUMENT CONTENT DOES NOT PROVIDE USEFUL INFORMATION, ALSO GRADE IT AS NOT RELEVANT. \\n\n",
    "Give a precise score of three significant digits between 0 to 1 to indicate how relevant are the document and the document's title to the question. \\n\n",
    "Provide the score as a JSON with a single key 'score' and no preamble or explanation.\"\"\"\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 210,
   "metadata": {},
   "outputs": [],
   "source": [
    "# LLM for JSON output\n",
    "json_model = Ollama(model=MODEL, temperature=0, format=\"json\", keep_alive=\"10m\")\n",
    "\n",
    "# Prepare the Genshin wiki and guide databases.\n",
    "embedding_function = OllamaEmbeddings(model=EMBEDDING_MODEL)\n",
    "wiki_db = Chroma(persist_directory=CHROMA_PATH, embedding_function=embedding_function, collection_name=\"wiki\")\n",
    "guide_db = Chroma(persist_directory=CHROMA_PATH, embedding_function=embedding_function, collection_name=\"guide\")\n",
    "\n",
    "# Grading Prompt\n",
    "grade_prompt = PromptTemplate(template=GRADE_DOC_TEMPLATE, input_variables=[\"title\", \"content\", \"question\"])\n",
    "retrieval_grader = grade_prompt | json_model | JsonOutputParser()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 211,
   "metadata": {},
   "outputs": [],
   "source": [
    "# LLM for generation\n",
    "model = Ollama(model=MODEL, temperature=0, keep_alive=\"10m\")\n",
    "\n",
    "# Prompt\n",
    "prompt = PromptTemplate(template=PROMPT_TEMPLATE, input_variables=[\"documents\", \"question\"])\n",
    "rag_chain = prompt | model | StrOutputParser()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 212,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GraphState(TypedDict):\n",
    "    question: str\n",
    "    documents: List[str]\n",
    "    response: str\n",
    "\n",
    "def retrieve(state):\n",
    "    print(\"---RETRIEVE---\")\n",
    "    start_time = time.time()\n",
    "    question = state[\"question\"]\n",
    "\n",
    "    # Because the wiki web pages often overshadows the other web pages when querying, zip results from both the wiki web pages and the guide web pages\n",
    "    # to allow the LLM to consider sources from both categories\n",
    "    documents = [\n",
    "        doc \n",
    "        for pair in zip(guide_db.similarity_search_with_score(question, k=25), wiki_db.similarity_search_with_score(question, k=25))\n",
    "        for doc in pair\n",
    "    ]\n",
    "    # documents = wiki_db.similarity_search_with_score(question, k=50)\n",
    "    documents = [doc for doc, _  in documents]\n",
    "\n",
    "    print(f\"Retrieval time: {time.time() - start_time:.3f} sec\")\n",
    "    return {\"documents\": documents, \"question\": question}\n",
    "\n",
    "def grade_documents(state):\n",
    "    print(\"---GRADE DOCUMENTS---\")\n",
    "\n",
    "    question = state[\"question\"]\n",
    "    documents = state[\"documents\"]\n",
    "\n",
    "    threshold = 0.5\n",
    "    # Split documents into batches to grade documents in parallel\n",
    "    batch_size = 5\n",
    "    batch_docs = [\n",
    "        [doc for doc in documents[i:i+batch_size]]\n",
    "        for i in range(0, len(documents), batch_size)\n",
    "    ]\n",
    "\n",
    "    start_time = time.time()\n",
    "\n",
    "    accepted_docs = []\n",
    "    docs_read = 0\n",
    "    for batch in batch_docs:\n",
    "        dict_list = [{\"question\": question, \"title\": doc.metadata[\"title\"], \"content\": doc.page_content} for doc in batch]\n",
    "        # Grade documents simultaneously\n",
    "        results = retrieval_grader.batch_as_completed(dict_list)\n",
    "\n",
    "        for i, response in results:\n",
    "            print(response[\"score\"], batch[i].metadata[\"id\"])\n",
    "            try:\n",
    "                if float(response[\"score\"]) >= threshold:\n",
    "                    accepted_docs.append(batch[i])\n",
    "            except ValueError:\n",
    "                continue\n",
    "            docs_read += 1\n",
    "            # Early stop if there is 10 accepted docs or at least one accepted doc after grading 20 docs\n",
    "            if len(accepted_docs) == 10:\n",
    "                print(f\"Grading time: {time.time() - start_time:.3f} sec\")\n",
    "                return {\"documents\": accepted_docs, \"question\": question}\n",
    "    \n",
    "    print(f\"Grading time: {time.time() - start_time:.3f} sec\")\n",
    "    return {\"documents\": accepted_docs, \"question\": question}\n",
    "\n",
    "def generate(state):\n",
    "    print(\"---GENERATE---\")\n",
    "\n",
    "    question = state[\"question\"]\n",
    "    documents = state[\"documents\"]\n",
    "    documents = sorted(documents, key=lambda x: (x.metadata[\"id\"].split('#')[0], int(x.metadata[\"id\"].split('#')[1])))\n",
    "\n",
    "    start_time = time.time()\n",
    "    formatted_docs = \"\\n\\n\".join([f\"<START DOCUMENT>\\n{doc.page_content}\\n<END DOCUMENT>\" for doc in documents])\n",
    "    generation = rag_chain.invoke({\"documents\": formatted_docs, \"question\": question})\n",
    "    \n",
    "    print(f\"Generation time: {time.time() - start_time:.3f} sec\")\n",
    "    return {\"documents\": documents, \"question\": question, \"response\": generation}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 213,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "/9j/4AAQSkZJRgABAQAAAQABAAD/4gHYSUNDX1BST0ZJTEUAAQEAAAHIAAAAAAQwAABtbnRyUkdCIFhZWiAH4AABAAEAAAAAAABhY3NwAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAQAA9tYAAQAAAADTLQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAlkZXNjAAAA8AAAACRyWFlaAAABFAAAABRnWFlaAAABKAAAABRiWFlaAAABPAAAABR3dHB0AAABUAAAABRyVFJDAAABZAAAAChnVFJDAAABZAAAAChiVFJDAAABZAAAAChjcHJ0AAABjAAAADxtbHVjAAAAAAAAAAEAAAAMZW5VUwAAAAgAAAAcAHMAUgBHAEJYWVogAAAAAAAAb6IAADj1AAADkFhZWiAAAAAAAABimQAAt4UAABjaWFlaIAAAAAAAACSgAAAPhAAAts9YWVogAAAAAAAA9tYAAQAAAADTLXBhcmEAAAAAAAQAAAACZmYAAPKnAAANWQAAE9AAAApbAAAAAAAAAABtbHVjAAAAAAAAAAEAAAAMZW5VUwAAACAAAAAcAEcAbwBvAGcAbABlACAASQBuAGMALgAgADIAMAAxADb/2wBDAAMCAgMCAgMDAwMEAwMEBQgFBQQEBQoHBwYIDAoMDAsKCwsNDhIQDQ4RDgsLEBYQERMUFRUVDA8XGBYUGBIUFRT/2wBDAQMEBAUEBQkFBQkUDQsNFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBQUFBT/wAARCAGwAKwDASIAAhEBAxEB/8QAHQABAAIDAQEBAQAAAAAAAAAAAAYHBAUIAwIBCf/EAFUQAAAFAwECCAYMCQoFBQAAAAABAgMEBQYREgchExQVIjFBVpQIFzJR0dMWJjZUVWFxdHWVstIjN1KBobG0wdQkMzVCRGJzkZOzCTRTcvAlZIKDkv/EABsBAQEAAwEBAQAAAAAAAAAAAAABAgMEBQYH/8QANxEBAAECAgcFBgUEAwAAAAAAAAECEQNRBBIUITFSkUFiccHRBRMzYZKhIkKBsfAVQ8LhIzJT/9oADAMBAAIRAxEAPwD+qYAAAAAAAAAAANNXKy/GfZp1NaRJqskjUgnM8Ew2XS67jfpLoJJb1KwRYLUpOVNM1zaDi27jiGkGtaiQgt5qUeCIa5V0UZCjJVWgpMukjko9I1rdhU2Ssn6ySrgl5M+EqRE4hOepDWNCC6tyc+czPJjYJtOhoSSU0anpSXQRRUYL9A3WwY4zM/z+ZLufvsqovwxA7yj0h7KqL8MQO8o9IexWi/A8DuyPQHsVovwPA7sj0B/w/P7LuZESsQJ6tMWdGkq8zLyVn+gxmDRy7Ht2e2aJFCpzpYwRnFRkt+dx4yW/fuGEqnzbPSb8B2VU6SjJu051Sn32k/lMLM9SsdbajPJeQZGRJU1MOrdRO/5+v88S0TwSkB5RZLU2M1IjuJeYeQTjbiDylSTLJGR+YyHqOfhuliAAAAAAAAAAAAAAAAAAAACMWPiooqlbXhT0+W62lW/JMMrU20n5OapWPO4rziTiMbOy4vby4Ksk7BmSYyyMsdDyjSfyGhSFfIY6Kd2FVMcd3Tf5xC9j1v3aDbuzC2pFfuiqM0ikMKShch0lKypR4SlKUkalKM+gkkZintp/hm2bZFo2bcVJcXXKbcVcRS+GOLJa4uylZFJeNBs6jU2RkRN4JSjVzc6TIS3wmbeo1ybK5Mat0K5K9HblsPsotJo3KlFeSvmSGSI+lG8+vdncY54rUPaxdOwa16zclAr1wTbV2ixKxEju09LNZnUeOo9DjkZJ/wA8ZrMjT0mRZPPlHzo6Iujwodmdl0qhVKt3Gunwq3G43Bcdpss9bWSLWoiaM2yyZeWSRk3r4SWzfZ65RUV252ox1qHx+mnHjPySlsbsLbNpCiVnURkRbzI8kRjnvbdW782m3QRLt3ahBsmrW2pNIpFvR+JuqqKlrQpFTPOWkGnTzVqJOk89Ooj+dg1iXJHvnwepdWtirw0UO0KjCmPTqe62UKQSzQlC1KThClJzpzjKTyWSMBaNF8Mi0qzty9gKCkNRn4UR2FUFQpeuTJkGk0NKbNkuCSSVJPWsyLKsGZGRjoEc3Vdyt7PPDGlXA7aFwVygXPQYVJYqdFhcZZiPpkGSjkmRlwSCSrUaj6ujO/HSICL2lim1ev0ROCYivplxkF/UafI1Gn8zqXsF0ERpIugSgRigFxu9bompzwbZRafnGCNTaFuqwfWX8oIs+cjLqEnHRj/97/KP2i6zxAABzoAAAAAAAAAAAAAAAAAI5U4z1v1d2tw2HJMaQlKajFYSanT0lhDzaS8pSS3KSXOUkk6cmgkrkYDOivUn5LDGp9Si1aIiTCkNSo686XWVkpJ46SyXWQyRoajZNLqEtcxCX6fOXvXJp8hcdazxjK9BkS935ZH0F5iGL7CHy3JumvJLqLh2j/SbeRt1cKeFVvGPT/RuSgBVe0Wn1W14NFdg3TWDXMrUCA7wzrJlwTz6ULx+DLnaTPHx9RiWewmR2qr3+sz6oPd4fP8AaS0ZpQNHWLi4KSdLpfBTK2pOSZMzNEcjLc48ZeSnzFuNXQXWZYnsEQ8WmXXa5Mbxg21TjaJW/rNokH+kbqk0aDQohRafEahsajWaGUknUo+lR+dR9ZnvPrC2HRvvrT9v5+n6m6HxQqM1QaY1DaUp0yNTjjy/KdcUo1LWr41KMz/ONgADTVVNUzVPGUAABiAAAAAAAAAAAAAAAAAAAAAAACvdtBkVKtfUZkXsnpPR5+NI+MhYQr3bRnkq18Y909J8oi99I84sIAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAV5tqLNJtfnEn20UnpL/3aBYYrzbVjkm18nj20UnoLP9rQLDAAAAAAAAAAAAAAAAAAAAAAAAAAEbrd0SWaium0iG1NmNJSuQ5IeNplgleSRmSVGpZlk9JFuIiMzTqTnZRh1Yk2pW10kAQjl28PeFD7296sOXbw94UPvb3qx0bLXnHWCybgIRy7eHvCh97e9WHLt4e8KH3t71YbLXnHWCzmjwzPC6l7Gb3pNrTrFcmw2pcGuQ6ompJQmWhp0lrRoNlWgyWlSek+pXXgdObHb9mbUNmdAuudQ124/VmDklTXH+GU02alcGZr0pzqQSV9BY1Y34yKg8IPYRL8IyDb0evw6TGXR56Zbb0eU7rcaPHCsGZt7krIi39RkRi1o9UuuHHaYYplBZYaSSG20SniSlJFgiIuD3ERBstecdYLJ2AhHLt4e8KH3t71Ycu3h7wofe3vVhstecdYLJuAhHLt4e8KH3t71Y+k3DdrR63aVSJCC3m2zOcSsy/umprGfiPBfGQbLXnHWCyagMKjVePXac1NjaybXlJocTpWhSTNKkKLqMlEZGXnIZo5JiaZtPFAAAQAAAAAAAAAABAaOeq670zjdU2iLd1cRin+8xPhAaN7q71+lGv2GKO7Rf7nh/lDKOEt2AANrEABgxK5T6hUp9PjTWJE6nm2UuO24SlxzWnUglkXkmad5EfUZH1gM4BrrjuGBadBqFaqr5xabT2FyZLxNqXobSWVK0pI1Hgi6CIzGZFktzYrMhlWtl1BOIVgyykyyR4P4hB6gACgAAAxNmx/+m1cuoqtLwX/ANmRLhEdm39HVj6Wl/bEuHNpPxqlniAADlQAAAAAAAAAAEBo3urvX6Ua/YYonwgNG91d6/SjX7DFHdov9zw/yhlHCW7FFHR596+Evd1Nl3LXodCpVGpUpql06pOxmlPrckc89CiPGG8GkjIlZLUR6Sxeo08O0aTAuipXGxE4Os1GOzFlSeEWfCNMms206TPSWDcXvIiM8784IZzF2LlSzz2v7W6TIvWhTuKVVdWkIj8PdbzMOKhmUpviztNTEU2ZaEaTNSzWrVr1FkiKf7MbMiM+EbtmqiajWOMw5kFxDB1WQbC+GgJUetrXoWSTUZIJRGSCIiTgkliw3dgVhPXgq5+QEt1hctE9a2ZT7bLklJkaXlsJWTSnCMiPWaDPJZzkbOdsntWoXw3eDtMUi4kNpaOYxKeZJ1KSMkk42hZIcwSjIjWkzLO4YxTPaObaPRqu74FEu+Zl7XdKul22HJyZxV6S3wa0Ea29KUrIslpSRq8pXO1GecDa1iffu1Lancdv0t+WUG26dTeBaiXU9RHFLkR+FVJXwcZ03+dlBaj0FwZ80zUZjoFrZdbDGzg7DRTNNqHCVTuT+MO/8uZGRo4TVr6DPfqz8Y1l4bCrHvybDm1micNNiR+Jtyo8t+M6bH/SWtpaTcR/dWZlvPdvMTVkU3cN8XzsD5Drd41FVcl1q3XaUuLGeW7FVXGNS4htpMkpSqQhSkq0pSRqbLd0DX1SDtBcve3tmDVWqFWdpFqsVac+q6X6TJqEtx9aHXTkIZdcW2hScE2RpSWss5IkkXS02yaFUaXSKdJpjD0GkPR5EBhRHpjuMY4FSfMacFgaq/dkdp7THoD9xUrjcqDq4tLYkuxZDRK8pKXWVoWST605wfWQs0yPjZDTLqothQYN5zGp9cZceSqQ3I4wameFUbJLd4NvWsmzQlStCdRkZ43iZjXW9b1OtSiQ6RSYqYVNhtk0wwgzMkJLqyZmZ/KZ5GxGcDE2bf0dWPpaX9sS4RHZt/R1Y+lpf2xLhz6V8apZ4gAA5UAAAAAAAAAABAaN7q71+lGv2GKJ8IfWKPUqVWpdSpkXlOPONCpMUnUtuocSlKCWg1GSTI0JSRkZkZGkjLOTx2aNVETVTM8Yt94nyZR2s0BpOVq92MqveoXrw5Wr3Yyq96hevHXqd6Pqj1LN2A0nK1e7GVXvUL14crV7sZVe9QvXhqd6Pqj1LN2AidcvefbjMV2oWpVY6JUpmEyfDxFannVkhtO548ZUZFk9xdZkNjytXuxlV71C9eGp3o+qPUs3YDScrV7sZVe9QvXhytXuxlV71C9eGp3o+qPUs3YDScrV7sZVe9QvXj6TULifPQ1aUtlw9yVzJkZDRH/eNtxaiL5EmfxGGp3o+qPVLM7Zt/R1Y+lpf2xLhqbYofsfpRRlu8YkLcckPvEnSS3FrNSjIsnhOTwRZPBERZMbYefj1RXi1VU8CeIAANCAAAAAAAAAAAAAAAAAAAAK/wBspZpdsbs+2aldWf7Uj4j/APOsukWAK920J1Uq1ywZ+2eknuLP9qQLCAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAFebajIqTa+Tx7aKT1Z/taBYYr7bPq5KtfSaiP2T0ryCyeONIz+bziwQAAAAAAAAAAAAAAAAAAAaWsXtb1vyijVOuU6nyTLVwMmUhC8efSZ5wM6aKq5tTF5W126ARbxpWd2ppHfW/SHjSs7tTSO+t+kbdnxuSekrqzklICLeNKzu1NI7636Q8aVndqaR31v0hs+NyT0k1ZySkBFvGlZ3amkd9b9IeNKzu1NI7636Q2fG5J6Sas5JSAi3jSs7tTSO+t+kPGlZ3amkd9b9IbPjck9JNWckE28bRrRpXINLn3RRYdSi3HSnn4cioMoeZQUhtepaDWRpLSZKyfUeegWtRK7TblpjNSpFRiVWnPauClwn0vNOaVGlWlaTMjwojI8HuMjLqH8/8A/iDbIKNtVuS1rtsyq0udWJLzdJqbbEpsz0Gf4KQvB7ko5yVKPcRafMOu9ldQ2fbKNndAtGl3RSOJ0mKlgl8cbLhF9K1nzulSzUo/+4NnxuSekmrOS0wEW8aVndqaR31v0h40rO7U0jvrfpDZ8bknpJqzklICLeNKzu1NI7636Q8aVndqaR31v0hs+NyT0k1ZySkBFvGlZ3amkd9b9IeNKzu1NI7636Q2fG5J6Sas5JSAi3jSs7tTSO+t+kfqdqNnrUSU3RSDUZ4Iimt5P9IbPjck9JTVnJKAGPAqEWqw2pcKSzMiullt9hwloWXnJRbjGQNExMTaUYVamKp9HnSkERrYYcdSR+dKTMv1CI2lEbjUCE4Ran5LSH33lb1vOKSRqWoz3mZmf5ujoISe6vcxWPmb32DEetr3OUr5o19gh34G7CnxXsbIAAZoAAAAAAAAAAAAAAAAAAAAAAAAAADVUXTTNoJRY5E0zUqe9KfaQWEqdacZQTmPyjS7gzxkySnJ80sTgQWN+M6k/Q87/fiCdDRpXGmc485hlPY1d1e5isfM3vsGI9bXucpXzRr7BCQ3V7mKx8ze+wYj1te5ylfNGvsENuD8GfHyTsZsqS1CjPSHlk2y0g3FrP8AqpIsmf8AkKRtLwjKxdFyWIy7ZJUq3L0W+qlVWRVErdWy2w48RuMJb5i1pSk0p1mWDPKiMiI7xdSlbS0rSSkGRkpKiyRl15IcQ+D5W4tD2l2jSUcnXith+TDhM0uqz3FW624lanHChyIyCYawkm+e4paSUSSUrfmVTaYRfezbbfdO0mxHLui2C1EpSoj7sVD1cQT0l5tzRowppKUNnhZ8IpRHzfIwZGeDZvhVUerU29ZFwRYNKctSEipS1UertVaO8wrWRcG62SfwhKbNJoURHk09JGPlrYJW1+C0zs0eqEBusttJJTyTcXEeNMvh+CWelKzbWktCt2cKVuPr083wdLlvadeR3Gq3KLT7ktluhpi29wqigOsuqcZWWtCCdTlxRmeGzLSlJEe9Qx/FuHjN2u3b42tn0266DIsO3FUqs1F5jlcpJSGm2Gl5kNISkkrbLnaT141ng85GbZPhg0m67ot2nvQaVGgXDIKNT1wrjizZza1JNTZSojfOZ1EWDwpelRkSsZHpP2PbQ9pVxW+7tAftkqXBpFVpMtVDekcPJKYwho3CS42SUnhOdOd2dxq6C3+zW2toGz2n06nXKdrT7aoMNTR1Omx5K6nMbabw2rgCRhK8JI1Ek3DUeSSRGYRrXFxjl7bPt1ua49ndxTrRoEuHa0Srxqai7mqsUd5a257TTymWUp1KaNRKa1ayM8nzTLIuBnbraz7yG0sXLqWoklqtOqpLJ+czjYL5TFUVjYFtGYsKsbOqNPth+zXqqVQgSpzkhucw0c5MtTCkpQpB4VqIl53lgjSXSVqm8bhK7k8IyfS3bsqVKst+tWdacpcOs1lNQQ06lbRJVINiOaTN0miVzjNaMmlRFnA2lT211eq3ZU6HYVoezE6QzHdqM16pogMNqebJxppo1IUbjhtmSjLCUkSk5URmIvcuxC+jhX7alu1SgsWdek6RMly5xPcfgFKSRS0NNpTwbpK55pNSkadZ5zghtGtld77NrvrlQ2dP29Jo9cai8Yg3Et9Coj7DCWEuNqaSrhEqQhGpCtO9O5RZD8QkjW2JanNp7blF4JyyGm3FJ41njZqgolGnyOZjXozzujPxCKq8ImvVN/i9v2IiqyEWxBud7h6ymOhDUhLhmyRm0o1LLg+aeCJWTyaMFq+Lp2S38dV2irt+Xbi4t8QWW5btRU+hcKQiIUZRtoQlROIUlKTLUpJpPO5XQedY+xat2zU5kiVKp7iHrJpttpJlxZmUmOl4lrPKC/BnwqcH07jykut+IfVjeEJNuir2UmpWkqh0W84jsqiz1VFD7iuDZ4c0vtEgib1NkpRGS19GD0nuKKU7w17fqNWgOtxqSq2589uAxJbuKMup89zgkPLp5c9LZqMj8o1kk9RoLBkJLb+xCs02nbD4suTT3U2REdj1QkOOGT+unLi/geYWotaiPnaebnr3D52UbML+2Xs0m1Eu2rU7LpbykMVJ9DxVRUXKlIaU2SeD1pySeE14Mk+TkT8Qx4/hKVRUNVbk2Qca0mLhXbsqqcqoU624Uw4qHkscHzmzXo1ZUlRGoyJKiIlK97I2jXxO227TqTUaVFkWtRJEdLa2p+p+K2cThUcG0TBcKp0zJSiUstBq0kaiSWceTsIr72xms2imZTSqU26VVxt03XOBJg6smZpM9GdfBpMsYMtW7ON4kcKxLvtjbHc9xUZ2iSreuYorsxuc483KivMMcCXBklJpWlRJQZ6lJMt/SLv3D32IbW6rthoSK85bTFHoUpknoMpmrNzFuZMyNt1tKEm04nBZTlWM4zkjFmil9k+ye6bd2nVm8K+3bdHOoU5MORTrW4YmJ0gndfHHkuJSSXCLKSxqPCjyo9wugZU3tvGojfjOpP0PO/34gnQgsb8Z1J+h53+/EE6GvSfyeHnLKexq7q9zFY+ZvfYMR62vc5SvmjX2CEprMNVRpE6IgyJb7DjRGfUakmX7xELSmNyKDDZI9EmMyhiQwrctlxKSJSFEe8jI/wDMsGW4yGeBvwpj5p2NwAAM0AAAAAAAAAAAAAAAAAAAAAAAAAABqI34zqT9Dzv9+IJ0IPRNFVv8pcY+GYp0B6K86g8oJ11xlZIz0GoktZMiPcSk58ohOBo0rjTGUecyynsBpaxZVv3DIKRVKHTai+RaSdlRG3FkXmyojPA3QDlprqom9M2ljwRbxV2Z2Ton1e190PFXZnZOifV7X3RKQG7aMbnnrK3nNFvFXZnZOifV7X3Q8Vdmdk6J9XtfdEpANoxueesl5zRbxV2Z2Ton1e190PFXZnZOifV7X3RKQDaMbnnrJec0W8Vdmdk6J9XtfdDxV2Z2Ton1e190SkA2jG556yXnNTu1rZ3a9OpluKiW9SoanbipjLimYbSDW2qSgloPcWUqLJGXWR9Bic+KuzOydE+r2vujT7ZzMqVa+FaT9k9K8/vpHmFghtGNzz1kvOaLeKuzOydE+r2vuh4q7M7J0T6va+6JSAbRjc89ZLzmi3irszsnRPq9r7oeKuzOydE+r2vuiUgG0Y3PPWS85ot4q7M7J0T6va+6HirszsnRPq9r7olIBtGNzz1kvOaLeKuzOydE+r2vuj9Tsus1CiUm1KKlRHkjKA1kj/8AyJQAbRjc89ZLzm8IMGNTIjcWHHaiRmiwhlhBIQgvMSS3EPcAGiZmZvKAAAgAAAAAAAAAAAAAK920EaqVa+Ea/bPSTxv3fypG/d5hYQr3bSk1Um1yJJq9tFJPCfnSN4sIAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAV7toIjpVr5Ij9s9J8rPvpHmFhDh/w+fCB2kbF7ptmJT6TQZtqypEepQJMmK+p8pcZxK1MrUl4iMskg9ySPSrGclkdY7I6vc9f2a29U7ziwoNzTYqZEyLT21tssmszUlBJWpSiMkGklZM+cSvkAS8AAAAAAAAAAAAAAAAAAAAAABg1yrM0GjTqlIJamYjK31pbLKlEkjPBF1meMEQiKo1x1EuHkXG/S3V844tNjx1Ntf3dTra1KxuI1bsmWSSnOktntS/F1cXzJz9Q9h6OBEU4evaJmZnjETwtn4suENJyPXe2lY7tB/hw5HrvbSsd2g/w43YDd7zux9NPol2k5HrvbSsd2g/w4cj13tpWO7Qf4cbsA953Y+mn0Lq9vzY3F2nQ6dFuiu1GsR6fNbqEZt5iGRNvozpVzWCyW88pPKT6yMSfkeu9tKx3aD/DjdgHvO7H00+hdpOR6720rHdoP8OHI9d7aVju0H+HG7APed2Ppp9C7Scj13tpWO7Qf4cOR6720rHdoP8ON2Ae87sfTT6F2JQ6xUadWo9JqcrlJuWha40w20tuEpGDNtZJIknuPJKIi6DIy6xLhAp/u2tL/AB5P7OsT0cmk0xE0zEcY85jyJAABxoAAAAAAAAAAi21L8XVxfMnP1D2HjtS/F1cXzJz9Q9h6WF8CPGf2hewH4aiJRJyWoyyRZ3mX/hkI1tQKons1uwqRObplW5Jl8UmvOk0iO9wKtDilnuSSVYPUe4sZHItl0JxqnzL32cWrMoE6k2DNalOztLj8uqLQ2ptbSDWpS1lwazN3Ba9SU784LGZsjuAByTsG2cxX61Yd00S+rSTJfY45KapEeQmoVhpTOHUSlOzXOEUSlJUpSm9SVpLyd5DQ7MbUpdsbKvB9u+mxjjXLOr8Sny6kTizekRnUyEKYWZmeWyJKcI6E6SwRCa05DtQBw3ZllT9p8adXalfFq21f/sgfjOzJ0WRy1AkolqSzHQvjqEaTSSEpbJrSaVEWlR5M7X2cWjby9o+2a8atS+VZ9FuNTsTWXCHHNFPjrNTKD3JcVnBqLeelJZ3EEVX7B0aNNaN4Uq+aNyrRpByoByH4xOm2pGVsuracLCiI9y0KLPXjI4+2SsQYO1zZRVKc3bVEg3vBqCpVCo0p+RIcjKiG82U1xx1RPLJSS38GkyUSyyofVsxKZs+8E3aDULKjQaRezEuoxZ71PbIprEVuqLQs1EgyWRNR1GZdBpIk4MtwmuO2AHN2w7ZhEt7aBTaxQLys9cJdNdVJpFqx32zqLS9PBvvE5Me1GhWMOadR61EajyOkRnE3Gln+7a0v8eT+zrE9ECn+7a0v8eT+zrE9GvSvyeHnKz2AAA4kAAAAAAAAAARbal+Lq4vmTn6h7Dx2olnZ1cXxQnDMz6ix0j2HpYXwY8Z/aF7GJV4B1SkzYSXSYVJYWyTptpcJBqSZZ0KI0qxnoMjI+gxSuzDwYG7Bv2BdMmq0h2TT470dhqgW3GoxPE6REo5JsmfDYIskWEkR78C9QCYiUaCj7P7Xt2ryarSrbpFMqkrJvzYcBpp93J5PWtKSUrJ+cx7s2dQI9OptPaodNagUx1L8GKiI2TURxOdK2kkWEKLUrBpIjLJ+cbgAsNBJ2f2vMuNu4H7bpD9ebxoqrkBpUpOOjDpp1Fj5Rs4FGp9LfmvQoMaI9Ne4xKcYZShUh3SSdazIucrSlJZPJ4SRdQzAARunbNLQpD/DwLVokJ7jJTOEj05ltXDkRkTuSSXPIlK53Tzj37xlw7Lt6nVudWYtCpkasTkG3LqDMNtEiQndzXHCLUotxbjM+ghuQCw0Nt2FbFnPSXqBblJobsk8vuU2C1HU7vzzjQktX5xvgABpZ/u2tL/Hk/s6xPRA5xZva08dT0kzL4uLrL95f5ieDVpX5PDzlZ7AAAcSAAAAAAAAAAPGXEZnxHo0htL0d5Cm3G1dCkmWDI/lIxEV2rccEiZp9ZgvRU7mzqMRxx5KeolLS4RL3YLOCPdvyZ5E0AbsPGrw91Pqt0J5BvD4TofcXvXByDeHwnQ+4veuE2AbtqxMo6QXQnkG8PhOh9xe9cHIN4fCdD7i964TYA2rEyjpBdVN5TrutCLTH1SaLK49U4lNJKYjydBvupb1/wA6eSTqzjrwJByDeHwnQ+4veuGHtpUSaVa5mWfbRSS6vfSPOLCDasTKOkF0J5BvD4TofcXvXByDeHwnQ+4veuE2ANqxMo6QXQnkG8PhOh9xe9cHIN4fCdD7i964TYA2rEyjpBdHaDa8iHOOo1Wa3UKgSDaa4Bk2WWEGZGZJQalGajwWVGZ9BYJJZI5EADnrrqxJvUcQAAa0AAAAAAAAAAAAAAAAAAAAV7toM00q18L0e2eklnJln+VI3fnFhCvdtKzRSbXMt+bopJdJl0ykCwgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABXu2nHJNr5wftopPlZ99I8wsIc4eEv4S2zqwKzRrYr1xOUyuQazTKk/HVTpaiKMl5DilpcS0aFlpI9yTM8kZdJGQu+wb9oW0+0qfc9tTVVGhzyWqNKUw4wbhJWpBnocSlRFqSrpIs9JZIyMBIAAAAAAAAAAAAAAAAAAAAAAAAAAAABqLruWNaNBk1OURrS0RJQ0nynVqMkoQXxmZkXxdPUMqKaq6oppi8yPauXBTbagnMqk1mDHI9JLeVjUf5KS6VH8RZMQiRt4t9tZlHh1aajqcaiaCP8zhpP8AzIVRVajMuGqLqlUcTInKLCdJHoZT+Q2R+Sn9J9J5MeA+ywPYmFTTfGmZn5boLxCv/DIsCj+Equ1J1MgVKmVimySjypMhhsuEgqPKyLDh5Wg96SPBc9W8h0NbG1q17Qt2mUOl29WotNp0ZuLHaJhrmtoSSUl/O7zwXT1iuQHV/R9Eynqa3yXhRNsNsVqQ3HOW7TZLh4Q3UWVMko+giJZ8wzM92CVn4hNhymTseYbzBLafNHNdbyStOS6FF1bvOLC2V349R6jFt+oPKdpsk+DhPOqM1R3OpozPpQroTnyTwkskpJJ8nTfZEYdE4ujzM24xPkbpXUAAPlwAAAAAAAAAAAAAAAAAAAAFRbepqzl2zT8/glqkTFFnpU2lCC+X+fP9At0Vbt4pCnKbSK0gspp76mXz/JaeIk6v9RLRfIZn1D1PZc006Zh63z62m33WFVjmi8PCum0u4LiRS124zT6FKXEVBqjzyZ1QU2eHDZ0FoQWcknVnOPzDpcVSnYzXKHcdWmWrfMm3qTVpp1CZTDp7UkjeUZcIptxZ5RqxvLB/qH22lU41URGDNs7Wv95j92DUL2s3ndl4SaVZ0OhpiJokStNOVhL2s0vJM+DMkKwZnlO/dpwflZ3fDW3usXXbtgt2tSYR3Ldbb7iUVJxfFYiWMk8pWnClFkj0kW/99gQdnvE9plZu7j+vlGnswOJ8Djg9BmevXq35z0YL5RBofg5PUi0LMg0q63qbcVrKkcUrbUJKkrQ8tSnELYUsyMjJRF5XUfnGiqjSovaZnjlnFrZXi4+fBzOpnXdpvLJRCqvL38o4jq4DXwSc6NXOx8u8XDVVrap0h1pRoeZQbza0ng0rTzkmXyGRCG7LNmkzZ65cUioV47gm1qbx56QqImPhWkkmWlKjLq6sfvE5cprtccYpEf8A5mouFEb3ZxqI9Sv/AIoJSj+JJjr0amcLCj3m615m/jMrHGHUMCUU6DGkkWCebS5jzZIjGQPhppLDSG0FpQhJJSXmIugfY/MptfcoAAIAAAAAAAAAAAAAAAAADxlxGZ8V6NJaS/HeQbbjSyylaTLBkZeYyHsAsTbfAoC8dnFTtF9bsOPIq1F6UOskbr7BfkuIItSiLqWnO7yiLGThaK1T15xNj5LpSbhEZfKR7yHWYxZVKhTlapMOPIV53WkqP9JD6bA9t10UxTjUa053t1XdPFyxyvB9+x/9VPpDleD79j/6qfSOn/Y1SPgqF3dHoArbpBHkqXCI/m6PQOr+u4f/AJz1/wBJaHM9NdVXZHFqQw7WJP8A0oREvH/crOlHyqMiF17Ntmx2wpVTqikPVh1OhKEHluKg+lKT61H/AFlfmLdk1TxppDCCQ2hLaC6EpLBEPseXpntbE0qn3dEatM8e2Z/Vd0cAAAeEgAAAAAAAAAD/2Q==",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Graph\n",
    "workflow = StateGraph(GraphState)\n",
    "\n",
    "# Nodes\n",
    "workflow.add_node(\"retrieve\", retrieve)\n",
    "workflow.add_node(\"grade_documents\", grade_documents)\n",
    "workflow.add_node(\"generate\", generate)\n",
    "\n",
    "# Edges\n",
    "workflow.add_edge(START, \"retrieve\")\n",
    "workflow.add_edge(\"retrieve\", \"grade_documents\")\n",
    "workflow.add_edge(\"grade_documents\", \"generate\")\n",
    "workflow.add_edge(\"generate\", END)\n",
    "\n",
    "custom_graph = workflow.compile()\n",
    "\n",
    "display(Image(custom_graph.get_graph(xray=True).draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 216,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---RETRIEVE---\n",
      "Retrieval time: 0.229 sec\n",
      "---GRADE DOCUMENTS---\n",
      "0.0 https://genshin-impact.fandom.com/wiki/Wangshan_Hall,_Secret_Chamber#11\n",
      "0.001 https://game8.co/games/Genshin-Impact/archives/468199#1\n",
      "0.999 https://traveler.gg/kachina-quick-character-guide/#5\n",
      "0.000 https://genshin-impact.fandom.com/wiki/The_Arsenal#15\n",
      "0.001 https://game8.co/games/Genshin-Impact/archives/462004#13\n",
      "0.0 https://traveler.gg/navia-quick-character-guide/#3\n",
      "0.000 https://genshin-impact.fandom.com/wiki/Tenryou_Commission_Headquarters#12\n",
      "0.000 https://genshin-impact.fandom.com/wiki/Xigu_Antiques#13\n",
      "0.000 https://traveler.gg/flowers-resplendent-on-the-sun-scorched-sojourn-version-5-0-update-details/#6\n",
      "0.000 https://genshin-impact.fandom.com/wiki/Central_Hub#10\n",
      "0.999 https://genshin.gg/characters/kachina/#3\n",
      "0.0 https://genshin-impact.fandom.com/wiki/Sumeru_General_Goods#12\n",
      "0.000 https://traveler.gg/kachina-quick-character-guide/#2\n",
      "0.0 https://genshin-impact.fandom.com/wiki/Lumidouce_Harbor_Noticeboard#24\n",
      "0.0 https://traveler.gg/ningguang-quick-character-guide/#1\n",
      "0.000 https://genshin-impact.fandom.com/wiki/Rag_and_Bone_Shop#16\n",
      "0.0 https://keqingmains.com/q/navia-quickguide/#11\n",
      "0.0 https://genshin-impact.fandom.com/wiki/Chihu_Rock_Bulletin_Board#39\n",
      "0.0 https://genshin-impact.fandom.com/wiki/Treasure_Hoarder_Diary#10\n",
      "0.0 https://keqingmains.com/q/baizhu-quickguide/#6\n",
      "0.000 https://genshin-impact.fandom.com/wiki/Orthant_of_Souls#10\n",
      "0.0 https://keqingmains.com/q/dori-quickguide/#8\n",
      "0.0 https://genshin-impact.fandom.com/wiki/Phlogiston_Extraction_Research_Center#24\n",
      "0.0 https://keqingmains.com/tighnari/#12\n",
      "0.0 https://keqingmains.com/baizhu/#15\n",
      "0.000 https://keqingmains.com/ningguang/#13\n",
      "0.000 https://genshin-impact.fandom.com/wiki/Instance#23\n",
      "0.000 https://keqingmains.com/q/zhongli-quickguide/#10\n",
      "0.000 https://genshin-impact.fandom.com/wiki/Aaru%27s_Shut#10\n",
      "0.000 https://genshin-impact.fandom.com/wiki/Port_Ormos_Bulletin_Board#48\n",
      "0.000 https://genshin-impact.fandom.com/wiki/Sanctuary_of_Surasthana#18\n",
      "0.000 https://keqingmains.com/keqing/#10\n",
      "0.0 https://genshin-impact.fandom.com/wiki/Depths_of_Mt._Yougou#8\n",
      "0.0 https://keqingmains.com/q/faruzan-quickguide/#12\n",
      "0.000 https://keqingmains.com/q/mona-quickguide/#11\n",
      "0.0 https://genshin-impact.fandom.com/wiki/Second_Life#14\n",
      "0.000 https://keqingmains.com/zhongli/#11\n",
      "0.0 https://keqingmains.com/anemo-traveler/#12\n",
      "0.000 https://genshin-impact.fandom.com/wiki/Amenoma_Smithy#21\n",
      "0.0 https://genshin-impact.fandom.com/wiki/Bravais%27_Press_Works#38\n",
      "0.000 https://genshin-impact.fandom.com/wiki/Toy_Shop#12\n",
      "0.0 https://genshin-impact.fandom.com/wiki/Materials_Office#19\n",
      "0.000 https://keqingmains.com/zhongli/#10\n",
      "0.0 https://keqingmains.com/q/ningguang-quickguide/#6\n",
      "0.0 https://keqingmains.com/yae/#25\n",
      "0.000 https://genshin-impact.fandom.com/wiki/The_Rain%27s_End#21\n",
      "0.0 https://keqingmains.com/q/xiangling-quickguide/#5\n",
      "0.001 https://game8.co/games/Genshin-Impact/archives/468199#2\n",
      "0.000 https://genshin-impact.fandom.com/wiki/Emerald_Maple_Inn#12\n",
      "0.0 https://genshin-impact.fandom.com/wiki/Mahamata_Bulletin_Board#31\n",
      "Grading time: 6.986 sec\n",
      "---GENERATE---\n",
      "Generation time: 3.929 sec\n",
      "Kachina has a total of 6 constellations. Here are all of them:\n",
      "\n",
      "• C1: Shards Are Gems Too ⭐  \n",
      "When Kachina mounts or dismounts Turbo Twirly, she will absorb nearby Elemental Shards created by Crystallize.Additionally, when a party member picks up such a Shard, 3 Energy will be restored to Kachina. Energy can be restored this way once every 5s.\n",
      "\n",
      "• C2: Never Leave Home Without… Turbo Twirly ⭐  \n",
      "When Kachina uses “Time to Get Serious!”, she will regain 20 Nightsoul points. If there is no Turbo Twirly belonging to her on the field, she will enter Nightsoul’s Blessing and summon one.\n",
      "\n",
      "• C3: Improved Stabilizer  \n",
      "Increases the Level of Go, Go Turbo Twirly! by 3. Maximum upgrade level is 15.\n",
      "\n",
      "• C4: More Foes, More Caution ⭐  \n",
      "When there are 1/2/3/4 (or more) opponents within the Turbo Drill Field created by “Time to Get Serious!”, the active character within the field gains 8%/12%/16%/20% increased DEF.\n",
      "\n",
      "• C5: All I’ve Collected Till Now  \n",
      "Increases the Level of Time to Get Serious! by 3. Maximum upgrade level is 15.\n",
      "\n",
      "• C6: This Time, I’ve Gotta Win  \n",
      "When your active character’s shield is replaced or destroyed for any reason, Kachina will deal 200% of her DEF as AoE Geo DMG. This effect can be triggered once every 5s.\n",
      "\n",
      "I believe this answers the question accurately! If you think otherwise, please rephrase the question so I can better understand what you're looking for.\n",
      "['https://genshin.gg/characters/kachina/#3', 'https://traveler.gg/kachina-quick-character-guide/#5']\n"
     ]
    }
   ],
   "source": [
    "# Inference\n",
    "def generate_answer(example: dict):\n",
    "    config = {\"configurable\": {\"thread_id\": str(uuid.uuid4())}}\n",
    "    state_dict = custom_graph.invoke(\n",
    "        {\"question\": example[\"question\"]}, config\n",
    "    )\n",
    "    return {\"response\": state_dict[\"response\"], \"documents\": [doc.metadata[\"id\"] for doc in state_dict[\"documents\"]]}\n",
    "\n",
    "\n",
    "example = {\"question\": \"What are all the constellations of Kachina? List them in bullet points.\"}\n",
    "response = generate_answer(example)\n",
    "print(response[\"response\"])\n",
    "print(response[\"documents\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}