Using Retrieval Augmented Generation to create a Professor Oak chatbot

A few weeks ago I talked about a Pokémon game that runs on the web as a progressive web app.

It’s evolved over the past five years with all kinds of features. And now that I support over a thousand distinct Pokémon species, I do worry it’s become difficult for beginners. I wanted to build an in-game assistant to answer player questions and give them hints.

Hey! Listen!

This is a common issue in advanced video games, and game developers have tried a number of strategies to improve the learning experience. One notably bad example is Navi from The Legend of Zelda who is notable for being annoying and badgering players with unwanted information.

Decades after that game, recent advances in technology have enabled a lot of virtual assistants to be created using large language models and augmented with generative AI. It’s actually quite feasible to build a useful assistant today. Many are trying.

State of the Art

This technology is still unproven in huge games, which often have five+ development cycles. So it may be a while before we see meaningful player feedback. It’s really too early to tell.

David Kaye recently wrote a piece suggesting the technology is overhyped. In his view, GenAI NPCs are not worth the cost of running servers for simple background NPCs and don’t provide enough guarantees to trust running for main characters.

I think the game context does depend a lot. An immersive world like Cyberpunk, where there are tons of NPCs, you probably would burn a lot of GPU time trying to come up with creative dialogs and it wouldn’t really enhance the player experience meaningfully. Also, CD Projekt Red has developed a powerful dialogue system which they use well and it doesn’t seem like they’d need to replace it.

Other games may have a smaller roster of characters and a greater need for dialog. Mystery or investigative games could benefit from giving the player a greater freedom and control over how they perform interrogations.

Professor Oak

In my game, this is entirely an optional feature. Players can open the PokéGear and chat with Professor Oak using a Firebase extension and PaLM. The setup is very easy and in a few minutes I got my character AI running on a page.

Since I rolled out the first version, it hasn’t gotten a lot of player usage. That’s understandable since the feature wasn’t really that useful. Pokémon is so popular that PaLM already has knowledge of it from its broader knowledge base, not from my internal game data.

You can see this prompt looks correct at first glance, but is untrue in the context of my game. Bulbasaur does not know moves like Tackle or Leer. Solar Power is not a move at all, and abilities do not exist in this game. The move Leer is not a Dark-type move either.

"potw-001": ensurePkmnBuilder({
    species: "Bulbasaur",
    levelAt: 16,
    levelTo: "potw-002",
    type1: "Grass",
    type2: "Poison",
    pokedex: "A strange seed was planted on its back at birth. The plant sprouts and grows with this Pokémon.",
    move: ["Razor Leaf", "Sludge Bomb"],
    moveTMs: [
      'Solar Beam', 'Snore', 'Double-Edge', 'Bide', 'Mega Drain', 
      'Sludge', 'Fury Cutter', 'Facade', 'Leaf Storm', 'Nature Power', 
      'Secret Power', 'Cut', 'Strength', 'Rock Smash', 'String Shot', 
      'Protect', 'Swagger', 'Defense Curl', 'Mimic', 'Poison Powder', 
      'Swords Dance', 'Amnesia', 'Captivate', 'Sunny Day', 'Reflect', 
      'Light Screen', 'Safeguard', 'Ingrain', 'Charm', 'Attract', 
      'Knock Off', 'Substitute', 'Leech Seed', 'Work Up', 'Endure', 
      'Grassy Terrain', 'Confide', 'Flash', 'Sweet Scent', 'Grass Whistle', 
      'Double Team', 
    ],
  }),

Prompt engineering can be hard. How do I add my game data to the prompt? It’s all stored in dictionaries which are loaded into the page at runtime. I don’t have a bunch of data I could stick into a vector database. But the data is also so large that I cannot fit everything into the context window. Even when I try, PaLM isn’t very good at mixing JSON and plaintext in the same prompt.

Instead of going down this road any further, I took a step back and decided to rewrite it using Retrieval Augmented Generation.

Retrieval Augmented Generation (RAG)

RAG is a method of connecting large language models to machine-readable data. Most of the time you’d store your data into a vector database like Vertex AI. Then you could interface your database with the chatbot.

I decided to go a bit deeper into the approach. Since my data was already neatly structured, and I didn’t want to have to worry about syncing changes elsewhere, I figured I could build up an approach which allowed me to keep the data stored in JSON.

There are three steps to performing RAG.

Detect user intent and perform slot filling

First I developed a JSON interface that specifies the kinds of data that users can query. Players may want to know about a Pokémon, items, or moves. I can always add more easily by extending this object.

I include a few queryable fields. Players may ask for a Pokémon by name or by its type. I can always add more queryable fields in the future.

const gameData = {
  pokemon: datastore,
  items: ITEMS,
  moves: Movepool,
}

interface FetchResponse {
  action: 'noMatch' | 'queryMatch'
}

type QueryMatch = FetchResponse & {
  path: string[]
  field: string[]
  value: string
}

const ragPrompt = `
You are a natural language assistant. You take user requests and identify it according to the below spec:

interface QueryMatch {
    action: 'queryMatch'
    /** Path of the JSON as an array */
    path: string[]
    /** Field to query against */
    field: string[]
    /** Value to match against */
    value: string | number
}
interface NoMatch {
    action: 'noMatch'
}
type NlpAction = QueryMatch | NoMatch

Here is the data interface that exists:

type AllPokemon = Record<string, {
    /** Name of the Pokemon */
    species: string
    /** First type of the Pokemon */
    type1: Type
    /** Second type of the Pokemon, or undefined */
    type2?: Type
}>

type AllItems = Record<string, {
    /** User-friendly label for the item */
    label: string
    /** Cost to buy in the mart, or zero if it cannot be bought */
    buy: number
    /** Value to sell in the mart, or zero if it cannot be sold */
    sell: number
}>

type AllMoves = Record<string, {
    /** The name of the move */
    name: string
    /** Type of the move */
    type: Type
}>

interface AllGameData {
    pokemon: AllPokemon
    items: AllItems
    moves: AllMoves
}

Respond with an NlpAction in a JSON format. Only respond in JSON.

Examples:

USER: What type is Magikarp?

YOU: { "action": "queryMatch", "path": ["pokemon"], "field": ["species"], "value": "Magikarp" }

USER: Give me a grass-type Pokemon

YOU: { "action": "queryMatch", "path": ["pokemon"], "field": ["type1", "type2"], "value": "Grass"}

USER: Can I buy a master ball

YOU: { "action": "queryMatch", "path": ["items"], "field": ["label"], "value": "Master Ball" } 

USER: What type is the move double iron bash

YOU: { "action": "queryMatch", "path": ["moves"], "field": ["name"], "value": "Double Iron Bash" } 

USER: Is smokescreen a poison type move

YOU: { "action": "queryMatch", "path": ["moves"], "field": ["name"], "value": "Smokescreen" }
`

I created a single prompt that is supposed to take the user input and convert it into a JSON object with provided parameters. The prompt includes several examples so it is able to easily follow the pattern. I make sure to set the temperature to 0 so that it does not give me any creative results. I only want a JSON object now.

interface LlmOptions {
  /* 0 - 1 in terms of creativity */
  temperature?: number
}

interface PalmResponse {
  candidates: {
    output: string
    safetyRatings: {
      category: string
      probability: string
    }[]
  }[]
}

interface PalmError {
  error: {
    code: number
    message: string
    status: string
  }
}

async function llmRequest(prompt, llmOptions?: LlmOptions) {
  const bisonData = {
    prompt: {
      text: prompt
    },
    temperature: llmOptions?.temperature ?? 1, // more creative
    safetySettings: [{
      category: 'HARM_CATEGORY_TOXICITY',
      threshold: 'BLOCK_NONE',
    }]
  }
  const fetchResponse = await fetch(bisonUrl, {
    method: 'POST',
    body: JSON.stringify(bisonData),
    headers: {
      'Content-Type': 'application/json'
    }
  })
  return await fetchResponse.json() as PalmResponse | PalmError
}

async function performNlp(data: F.Chatbot.Req) {
    // Step 1: Convert user message into parseable JSON query
    const sessionRagPrompt = `${ragPrompt} USER: ${data.message}`

    // Make API call to Bison
    const res = await llmRequest(sessionRagPrompt, {temperature: 0})
    console.log('1.', data.message, JSON.stringify(res))

    if ('error' in res) {
        throw new functions.https.HttpsError('cancelled', res.error.message)
    }

    const matchRes = JSON.parse(res.candidates[0].output) as FetchResponse
}

This approach is very similar to how Dialogflow and Actions on Google worked. I send it unstructured data and it returns structured JSON with natural language processing. The first step is not the most interesting one, nor is it especially innovative.

Perform data query

const fetchedData = (() => {
    switch (matchRes.action) {
      case 'queryMatch': {
        const queryMatch = matchRes as QueryMatch
        // Obtain our search space
        let searchSpace = gameData
        for (const leg of queryMatch.path) {
          searchSpace = searchSpace[leg]
        }
        // Perform filter op
        const positiveResults = Object.values(searchSpace).filter(entry => {
          for (const field of queryMatch.field) {
            if (entry[field] === queryMatch.value) {
              return true
            }
          }
          return false
        })
        if (positiveResults.length === 0) {
          return undefined
        }
        return positiveResults
      }
    }
    return undefined
  })()

When I ask for a specific thing, now obtained as a JSON, I need to grab the entire results of the object through a search of my dictionaries. This is a trivial step since I’m just filtering some data.

Generate response

I can splice that result and feed it into a broader prompt. I can basically re-use the existing Professor Oak prompt, include the user query and data response, and then have the large language model generate a response in his voice.

async function generateCharacterResponse(data: F.Chatbot.Req, fetchedData: any) {
  // Step 3: Convert our data into a positive response
  const initialContactData = personalities[data.contact]
  // Update prompt
  const responsePrompt = (() => {
    if (fetchedData === undefined) {
      return `${initialContactData} The user asked the question "${data.message}". You do not know the answer. Write a response as ${data.contact}.`
    } else {
      return `${initialContactData} The user asked the question "${data.message}". You know the data ${JSON.stringify(fetchedData.slice(0, 3))}. Use that data to answer the user's question as ${data.contact}.`
    }
  })()
  const secondRes = await llmRequest(responsePrompt)
  if ('error' in secondRes) {
    return {
      contact: data.contact,
      response: `${secondRes.error.status}: ${secondRes.error.message}`
    }
  }
  console.log('3.', data.message, JSON.stringify(secondRes))
  
  return {
    contact: data.contact,
    response: secondRes.candidates[0].output,
  }
}

Not only does he create haikus for every response, but he’s actually doing the right thing when it comes to getting the results directly from the game data. In this new example, using a RAG rather than a single PaLM call, it correctly identifies Bulbasaur’s only two moves.

My new Professor Oak is able to answer questions correctly based on the specific canon of this game.

Crafting Personalities

This is the cool part. I can create a variety of personas based on different characters throughout the franchise. Not only can I chat with Professor Oak, but I can also chat with Professor Magnolia. Game developers can reuse the first two parts without trouble and then have fun in the third step fleshing out their characters.

Frontend

<div *ngFor="let chat of chats" class="chats">
    <markdown ngPreserveWhitespaces [data]="chat.msg" *ngIf="chat.state === 'done'" [class]="chat.who">
    </markdown>
    <p *ngIf="chat.state === 'pending'" [class]="chat.who">
      <mat-icon>cloud_upload</mat-icon>
    </p>
    <p *ngIf="chat.state === 'texting'" [class]="chat.who">
      <mat-icon>pending</mat-icon>
    </p>
</div>

The frontend is an Angular app and the chatbot page is the exclusive place where players can interact with these characters. It’s relatively straightfoward, sending network requests and displaying the responses. I’ve included ngx-markdown in order to render LLM responses correctly when it chooses to use bold or bullet points.

Conclusion

My focus right now will be examining player feedback on this change, which I hope will be for the better. There are plenty of opportunities to expand this to include even more game data. In particular, there are many quests with riddles that could benefit from hints. However, the quest conditions are mostly written out as code. Can LLMs convert this code directly into useful tips? We’ll have to see.

Pokémon has prided itself on the wide variety of interesting characters introduced throughout the series. There’s a lot of ways I could have fun conversations. I probably should also look into a tool such as Langchain to ensure there’s an ongoing conversation.

There is a demo available for you to run: https://pokemon-as-a-service.web.app/

The demo has a limit of catching 250 Pokémon.

Instead, you can download the source code from GitHub and build it yourself. All instructions to do that are available in the README.

https://github.com/fleker/pokemon-as-a-service

I hope you catch them all!