Consinder a State Monad for your Chatbot

I've recently been putting together a procedural text-based RPG backed by an LLM, partly because I couldn't find any existing one's I'd feel comfortable using, but mostly out of nostalgia for that time period around GPT2 when language models were a novelty people used to play nonsense Zork clones, rather than the industry-consuming inhuman bubble it is today. Creating this game presents an interesting challenge, because I both need to be able to ask questions with very particular answers and have a lot of context built up from previous answers.

As an example, say I want the LLM to generate an equippable item dropped after an encounter. The pseudocode would look something like:

conversation = ["This is a cyberpunk world that blah blah blah ..."]

name, description, conversation = ask(
    conversation, 
    "Name and describe a new item that {enemy} dropped"
)

for 0 → 3:
    a_name, a_desc, conversation = ask(
        conversation, 
        "Name and describe a new ability for {name}"
    )
    heals, conversation = ask(
        conversation, 
        "Does using {a_name} heal?"
    )
    etc...

The LLM has to output very particular things for this to work; name should be 1-3 words (preferably proper nouns), heals should be a boolean, etc. Not only that, but it needs to be aware of its previous answers, so that it doesn't make a knife that shoots bullets or near-duplicate abilities like "swing" & "attack".

The former problem can be somewhat tricky. In my case, I'm using llama.cpp's grammar parameter, which allows passing a context-free grammar to the completions endpoint to require a particular output format. This means we can, for example, pass root ::= "yes" | "no" and know the model will always respond "yes" or "no". I wish inference providers provided this option because it is incredibly useful, but I suspect they don't because it makes jailbreaking trivialhow?root ::= "Sure! " .+. I have also seen json schemas as an option with wider support which might also solve this issue, though I have not used it myself. In either case, this issue I'm giving less attention for this article.

The latter issue is, in a way, already solved in the pseudocode. conversation already contains the context we need including all the previous questions and answers. The problem is that it's a very noisy way of handling the context. Ideally, because it's obvious each question relates to the previous, we'd instead write something like:

given "This is a cyberpunk world that blah blah blah ...":
    name, description = ask(
        "Name and describe a new item that {enemy} dropped"
    )

    for 0 → 3:
        a_name, a_desc = ask("Name and describe a new ability for {name}")
        heals = ask("Does using {a_name} heal?")
        etc...

If we step back and consider what's happening when we call ask, we have as inputs:

• A string, with the question we want to ask
• A conversation that has the context to answer that question

and as outputs:

• The data that we asked for (description as a string, y/n as a boolean, etc.)
• What the conversation looks like after that's been asked and answered

So ask<T> should be of the form (string, Conversation) → (T, Conversation). We could bind the question ahead of time (such as via currying or function capturing), and simplify our question type to Conversation → (T, Conversation).

If you're familiar with the state monad, the latter type should stand out to youalsoYou can also likely already see the value of using one here. In that case, you don't really need to read the rest of this article.. If not, consider what a function that asks one question and then a follow-up question would look like:

howThatWorks(conversation):
    answer1, conversation = ask(conversation, "What is that?")
    answer2, conversation = ask(conversation, "Explain how {answer1} works.")
    return answer2, conversation

The signature of this function is also Conversation → (T, Conversation). If asking one question after another also looks like a question, then there should be some way of composing them. How do we compose any question with any follow-up?

Taking a look at the follow-up line, answer2, conversation = ask(conversation, "Explain how {answer1} works."), we could turn this into a function that converts what we're asking about into a question about that thing in particular:

explainHowThingWorks(thing):
    return (conversation) → ask(conversation, "Explain how {thing} works.")

So now we have a question type, Conversation → (T1, Conversation), and a follow-up type, T1 → Conversation → (T2, Conversation) The composition of these two into another question is easy to guess from the types:

bindFollowup(question, followup):
    return (conversation) →
        t1, conversation = question(conversation)
        t2, conversation = followup(t1)(conversation)
        return (t2, conversation)

So far, this might not appear to have helped reduce the verbosity of referring to conversation repeatedly. However, if we rearrange and curry ask from ask(Conversation, string) to ask(string)(Conversation), we can do this:

howThatWorks(conversation):
    return bindFollowup(
        ask("What is that?"),
        (that) → ask("Explain how {that} works")
    )(conversation)

Or even:

howThatWorks = bindFollowup(
    ask("What is that?"),
    (that) → ask("Explain how {that} works")
)

Which has effectively eliminated our need to handle the conversation context directly, outside of the implementation of ask. If you interpret bindFollowup to be the monadic "bind" operator (AKA >>=), then you get the benefit of tools built to handle monads, such as Haskell's do notation, which makes converting the original item generation example almost trivial.