dnd-answers.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using GPT4 to answer D&D questions\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "D&D has a lot of rules. This notebook uses GPT4 and a vector database to answer D&D questions.\n",
    "\n",
    "To setup locally run:\n",
    "```bash\n",
    "python -m venv venv\n",
    "source ./venv/bin/activate\n",
    "pip install -r requirements.txt\n",
    "```\n",
    "The source line is shell and platform dependant. For example fish is `source ./venv/bin/activate.fish`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qdrant_client import QdrantClient\n",
    "from qdrant_client.models import VectorParams, Distance, PointStruct\n",
    "from sentence_transformers import SentenceTransformer\n",
    "from typing import List, Callable\n",
    "import uuid\n",
    "import os\n",
    "import re\n",
    "import glob\n",
    "from functools import partial\n",
    "import openai"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Data\n",
    "First we need the rules as text. In this example we load the SRD from a helpful GitHub repository that has converted it to Markdown. You could replace this with whatever rules you can get your hands on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Checkout github.com/OldManUmby/DND.SRD.Wiki shallow clone\n",
    "# Tested with 7c99bcc60bb67c066b1b59a3d20f52278371fd48\n",
    "datadir = \"./srd\"\n",
    "if not os.path.isdir(datadir):\n",
    "    os.system(f\"git clone --depth 1 https://github.com/OldManUmby/DND.SRD.Wiki {datadir}\")\n",
    "\n",
    "# Cleanup the stuff we dont' need like the (Alt folders and the readme)\n",
    "os.system(f\"rm -rf {datadir}/*Alt*\")\n",
    "os.system(f\"rm -rf {datadir}/.git\")\n",
    "os.system(f\"rm -rf {datadir}/.github\")\n",
    "os.system(f\"rm -rf {datadir}/*.md\")\n",
    "os.system(f\"rm -rf {datadir}/*.png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reading markdown files\n",
    "Next we just need to load the markdown files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def readMarkdownDirectoryRecursive(directory: str) -> List[str]:\n",
    "    markdown_files = []\n",
    "    for root, _, files in os.walk(directory):\n",
    "        for file in files:\n",
    "            if file.endswith(\".md\"):\n",
    "                markdown_files.append(os.path.join(root, file))\n",
    "    return list(map(readWholeFile, markdown_files))\n",
    "\n",
    "def readWholeFile(filename: str):\n",
    "        with open(filename, 'r') as file:\n",
    "                return file.read()\n",
    "\n",
    "srdMdFiles = readMarkdownDirectoryRecursive(datadir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chunking text\n",
    "Splitting the Markdown into chucks is a bit of a fiddly step. The size and quality of chunking are important for the embeddings model to be able to understand what is going on.\n",
    "\n",
    "Here we split on Markdown headings, then newlines, then we have to give up and split wherever. This seems to work well enough for the SRD text we are using here. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Splits a string closets to the middle as possble keeping the delimiter\n",
    "def split_middle(splitstr: str, item: str, maxlen: int) -> List[str]:\n",
    "    if len(item) < 8:\n",
    "        return [item]\n",
    "\n",
    "    half = len(item) // 2\n",
    "    left = item[:half]\n",
    "    right = item[half:]\n",
    "    try:\n",
    "        loc_left = left.rindex(splitstr)\n",
    "    except ValueError:\n",
    "        loc_left = None\n",
    "    try:\n",
    "        loc_right = right.index(splitstr)\n",
    "    except ValueError:\n",
    "        loc_right = None\n",
    "        \n",
    "    # No split found\n",
    "    if loc_left == None and loc_right == None:\n",
    "        return [item]\n",
    "        \n",
    "    # if left is closer than right or right is disqualified\n",
    "    if loc_right == None or (loc_left != None and (len(left) - loc_left) < loc_right):\n",
    "        split_left = left[:loc_left]\n",
    "        split_right = left[loc_left:] + right\n",
    "    else: # Otherwise right is closer than left\n",
    "        split_left = left + right[:loc_right]\n",
    "        split_right = right[loc_right:]\n",
    "    \n",
    "    # Prevent tiny splits.\n",
    "    if len(split_left) < 2 or len(split_right) < 2:\n",
    "        return [item]\n",
    "        \n",
    "    return [split_left, split_right]\n",
    "\n",
    "def split_recursive(items: List[str], maxlen: int, splitLogic: Callable[[str, int], List[str]]) -> List[str]:\n",
    "    result = []\n",
    "    for item in items:\n",
    "        if len(item) > maxlen:\n",
    "            split = splitLogic(item, maxlen)\n",
    "            # Base case, no possible separator.\n",
    "            if len(split) == 1:\n",
    "                result += [item]\n",
    "            else:\n",
    "                result += split_recursive(split, maxlen, splitLogic)\n",
    "        else:\n",
    "            result += [item]\n",
    "    return result\n",
    "    \n",
    "def make_splitter(split_pattern: str) -> Callable[[str, int], List[str]]:\n",
    "    return partial(split_middle, split_pattern)\n",
    "    \n",
    "def split_markdown_on_headings(items: List[str], maxlen: int) -> List[str]:\n",
    "    result = items\n",
    "    for depth in range(1, 7):\n",
    "        separator = \"\\n\" + \"#\"*depth + \" \"\n",
    "        result = split_recursive(result, maxlen, make_splitter(separator))\n",
    "    return result\n",
    "\n",
    "def split_markdown(text: List[str], maxlen: int) -> List[str]:\n",
    "    table_start = re.compile(r\"(\\|.*\\|)\")\n",
    "    \n",
    "    # Start with the whole text\n",
    "    splits = text\n",
    "    \n",
    "    # Split on headings\n",
    "    splits = split_markdown_on_headings(splits, maxlen)\n",
    "\n",
    "    # Split on Newlines\n",
    "    splits = split_recursive(splits, maxlen, make_splitter(\"\\n\"))\n",
    "    \n",
    "    # Trim whitespace\n",
    "    splits = list(map(str.strip, splits))\n",
    "    \n",
    "    # Last resort trim on whatever. This guarenttes no enteries over a certain value\n",
    "    splits = split_recursive(splits, maxlen, make_splitter(\"\"))\n",
    "    splits = list(map(str.strip, splits))\n",
    "\n",
    "    return splits\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "splitMd = split_markdown(srdMdFiles, 1024)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create Embeddings\n",
    "Next we create some embeddings. The embeddings allow us to do a symantic search for the users question within the rules. We will then store them in a vector database for easy retrieval in the next step. \n",
    "\n",
    "For this example we use sentance transformers embeddings. Picked because it was a good general purpose model. We could also use openai's embeddings model here if we don't want to wait around for them to compute locally.\n",
    "\n",
    "If you want to run this example and don't have a GPU. Remove the `device='cuda'` parameter. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Embeddor:\n",
    "    def embeddings(self, text: List[str]) -> List[List[float]]:\n",
    "        pass\n",
    "\n",
    "    def countTokens(self, text: List[str]) -> List[int]:\n",
    "        pass\n",
    "\n",
    "    def getEmbeddingSize(self) -> int:\n",
    "        pass\n",
    "        \n",
    "    def getModelName(self) -> str:\n",
    "        pass\n",
    "\n",
    "class SentenceTransformerEmbeddor(Embeddor):\n",
    "    def __init__(self, model_name: str):\n",
    "        self.model_name = model_name\n",
    "        self.model = SentenceTransformer(model_name, device='cuda')\n",
    "        self.embedding_size = self.model.get_sentence_embedding_dimension()\n",
    "\n",
    "    def embeddings(self, text: List[str]) -> List[List[float]]:\n",
    "        embeddings = self.model.encode(text)\n",
    "        return [embedding.tolist() for embedding in embeddings]\n",
    "\n",
    "    def countTokens(self, text: str) -> List[int]:\n",
    "        return len(self.model.tokenize([text])['input_ids'][0])\n",
    "    \n",
    "    def max_sequence_length(self) -> int:\n",
    "        return self.model.max_seq_length\n",
    "    \n",
    "    def count_oversized(self, text: List[str]) -> int:\n",
    "        count = 0\n",
    "        max_length = self.max_sequence_length()\n",
    "        for item in splitMd:\n",
    "            num_tokens = self.countTokens(item)\n",
    "            if num_tokens >= max_length:\n",
    "                count += 1\n",
    "\n",
    "    def getEmbeddingSize(self) -> int:\n",
    "        return self.embedding_size\n",
    "        \n",
    "    def getModelName(self) -> str:\n",
    "        return self.model_name\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddor = SentenceTransformerEmbeddor(\"sentence-transformers/multi-qa-mpnet-base-dot-v1\")\n",
    "embeddings = embeddor.embeddings(splitMd)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Store embeddings in vector database\n",
    "\n",
    "Next we store the embeddings we create in a vector database. The specifics of how to use the vector database are not important here and there are many options for vector databases. We use Qdrant here because it's convenient. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "UpdateResult(operation_id=0, status=<UpdateStatus.COMPLETED: 'completed'>)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client = QdrantClient(\":memory:\")\n",
    "srd_collection_name = \"srd\"\n",
    "client.create_collection(collection_name=srd_collection_name, vectors_config=VectorParams(size=embeddor.getEmbeddingSize(), distance=Distance.COSINE))\n",
    "points = [PointStruct(id=str(uuid.uuid4()), vector=embedding, payload={\"text\": text}) for embedding, text in zip(embeddings, splitMd)]\n",
    "client.upsert(collection_name=srd_collection_name, points=points)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Querying"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First we make a function to search the vector database for related chunks of text. \n",
    "\n",
    "We create an embedding from the query then search for similar embeddings. Returning the chunks of text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def search_srd(query: str, limit: int = 10):\n",
    "    # Create embedding from query\n",
    "    query_embedding = embeddor.embeddings([query])[0]\n",
    "    # Get the n most similar vectors\n",
    "    results = client.search(collection_name=srd_collection_name, query_vector=query_embedding, limit=limit, with_payload=True)\n",
    "    # Request the text for each result\n",
    "    return [result.payload[\"text\"] for result in results]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now for the main event the `ask_srd` function. We use the `search_srd` function above to get the top n chunks of text that might be relevant to the query. Then we incorporate those into our prompt to ask GPT-4 to look though them and answer our question. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ask_srd(query: str) -> str:\n",
    "    context = search_srd(query, 20)\n",
    "    context = \"\\n\".join(context)\n",
    "\n",
    "    prompt = f\"You are a helpful assistant that answers questions about D&D 5e. Given the following excerpts from the rules answer users questions. If you can't find the answer in the excerpts say you don't know.\\n{context}\"\n",
    "    system_message = {\"role\": \"system\", \"content\": prompt}\n",
    "    prompt_message = {\"role\": \"user\", \"content\": query}\n",
    "    srd_answer = openai.ChatCompletion.create(\n",
    "        model=\"gpt-4-0613\",\n",
    "        messages=[system_message, prompt_message],\n",
    "    )\n",
    "    \n",
    "    return srd_answer[\"choices\"][0][\"message\"][\"content\"]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can ask whatever questions we want! Your players will never argue with you again when your robot tells them they can't cast two fireballs on the same turn! No sorry that's not how quickened spell works. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "According to the rules regarding bonus action spells, if you cast a spell that has a casting time of 1 bonus action, you can't cast another spell during the same turn, except for a cantrip with a casting time of 1 action.\n"
     ]
    }
   ],
   "source": [
    "question = \"Can I cast two leveled spells on the same turn?\"\n",
    "print(ask_srd(question))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Bonus: function calling\n",
    "\n",
    "The function calling API can allow the model to decide when it needs to know something. For example, if it thinks it needs to lookup a specific spell.\n",
    "\n",
    "You could also use this to allow the lookup of specific information about your campaign such as your character's stats. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load spells from Spells folder and create a dictionary of spell name to spell text\n",
    "\n",
    "spell_files = glob.glob(\"./srd/Spells/*.md\")\n",
    "spell_files.sort()\n",
    "spell_names = [os.path.splitext(os.path.basename(file))[0].lower() for file in spell_files]\n",
    "spell_texts = [open(file, \"r\").read() for file in spell_files]\n",
    "spell_dict = dict(zip(spell_names, spell_texts))\n",
    "\n",
    "def get_spell(spell_name: str) -> str:\n",
    "    return spell_dict[spell_name]\n",
    "\n",
    "def resolve_functions(result, original_messages) -> str:\n",
    "    message = result[\"choices\"][0][\"message\"]\n",
    "    if message.get(\"function_call\"):\n",
    "        if message[\"function_call\"][\"name\"] == \"lookup_spell\":\n",
    "            spell_name = eval(message[\"function_call\"][\"arguments\"])[\"name\"]\n",
    "            spell_text = get_spell(spell_name.lower())\n",
    "            answer_with_function = openai.ChatCompletion.create(\n",
    "                model=\"gpt-4-0613\",\n",
    "                messages=original_messages + [{\"role\": \"function\", \"name\": \"lookup_spell\", \"content\": spell_text}],\n",
    "            )\n",
    "            return answer_with_function[\"choices\"][0][\"message\"][\"content\"]\n",
    "        else:\n",
    "            raise Exception(f\"Unknown function {message['function_call']['function']}\")\n",
    "    else:\n",
    "        return message[\"content\"]\n",
    "\n",
    "\n",
    "def ask_srd_functions(query: str) -> str:\n",
    "    context = search_srd(query, 20)\n",
    "    context = \"\\n\".join(context)\n",
    "\n",
    "    context = \"\"\n",
    "\n",
    "    prompt = f\"You are a helpful assistant that answers questions about D&D 5e. Given the following excerpts from the rules answer users questions. If you can't find the answer in the excerpts say you don't know.\\n{context}\"\n",
    "    system_message = {\"role\": \"system\", \"content\": prompt}\n",
    "    prompt_message = {\"role\": \"user\", \"content\": query}\n",
    "    \n",
    "    lookup_spell = {\n",
    "        \"name\": \"lookup_spell\",\n",
    "        \"description\": \"Retrieve the description of a 5e spell.\",\n",
    "        \"parameters\": {\n",
    "            \"type\": \"object\",\n",
    "            \"properties\": {\n",
    "                \"name\": {\n",
    "                    \"type\": \"string\",\n",
    "                    \"description\": f\"The name of the spell to lookup. Should be lowercase like 'fireball'.\",\n",
    "                },\n",
    "            },\n",
    "        },\n",
    "    }\n",
    "    functions=[lookup_spell]\n",
    "    messages=[system_message, prompt_message]\n",
    "    \n",
    "    answer = openai.ChatCompletion.create(\n",
    "        model=\"gpt-4-0613\",\n",
    "        messages=messages,\n",
    "        functions=functions,\n",
    "        function_call=\"auto\",\n",
    "    )\n",
    "    return resolve_functions(answer, messages)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "A fireball does 8d6 fire damage to each creature in a 20-foot radius sphere centered on a point you choose within range. If the creature makes a successful Dexterity saving throw, the damage is halved. If you cast this spell using a spell slot of 4th level or higher, the damage increases by 1d6 for each slot level above 3rd.\n"
     ]
    }
   ],
   "source": [
    "spell_question = \"How much damage does a fireball do?\"\n",
    "print(ask_srd_functions(spell_question))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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.11.3"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}