LLMs and Structured Conversation Graphs: The best of both worlds

from pydantic import BaseModel, Field
from graphviz import Digraph
from textwrap import wrap as twrap


class ChatGraph(BaseModel):
    """Stores the full graph for a given conversation"""

    nodes: list["ChatNode"]
    start_node: str = Field(description="Label of the node to begin from")

class ChatNode(BaseModel):
    """Individual nodes"""

    node_label: str = Field(description="A unique, informative ID for this node")
    instructions: str = Field(
        description=(
            "Instructions for the chatbot loosely outlining what it should say at this step. "
            "The chatbot will use these instructions plus background knowledge about the user to personalise its actual message."
        )
    )
    response_categories: list["ChatResponseCategory"] = Field(
        description="A list of possible categories of user responses"
    )

class ChatResponseCategory(BaseModel):
    """A possible classification for the user's message"""

    description: str = Field(description="A clear description of the category")
    label: str = Field(description="A simple, unambiguous label for this response category")
    goto: str = Field(description="The label of the ChatNode to go to if this response is given")

from dotenv import load_dotenv
from openai import OpenAI
load_dotenv()

client = OpenAI()

# This takes about 14 seconds
generated_chat_model_response = client.beta.chat.completions.parse(
    model="gpt-4o",
    response_format=ChatGraph,
    messages=[
        {
            "role": "system",
            "content": (
                "Use the schema provided to generate a short ChatModel for a basic CBT chatbot that helps the user deal with anxiety. "
                "Each node should include at least two response categories. "
                "Instructions for non-terminal nodes should always include questions for the user to keep the conversation moving. "
                "Aim for around 7 nodes in the conversation."
            ),
        }
    ],
)
generated_chat_model = generated_chat_model_response.choices[0].message.parsed

!mkdir -p chats
generated_chat_model.write_json("chats/cbt1.json")

## Uncomment to load from file
# generated_chat_model = ChatGraph.read_json("chats/cbt1.json")

print(
    generated_chat_model.model_dump_json(indent=2)
)

{
  "nodes": [
    {
      "node_label": "introduction",
      "instructions": "Greet the user and acknowledge that they are seeking help for anxiety. Ask them to describe their current feeling or situation.",
      "response_categories": [
        {
          "description": "User provides a general description or elaborates on their anxiety",
          "label": "User elaborates on anxiety",
          "goto": "explore_causes"
        },
        {
          "description": "User expresses uncertainty or is unable to specify what they're feeling",
          "label": "User uncertain about feelings",
          "goto": "explore_causes"
        }
      ]
    },
    {
      "node_label": "explore_causes",
      "instructions": "Ask the user about potential factors contributing to their anxiety. Encourage them to consider recent events, ongoing stressors, or specific triggers.",
      "response_categories": [
        {
          "description": "User identifies specific triggers or stressors",
          "label": "User identifies triggers",
          "goto": "discuss_thoughts"
        },
        {
          "description": "User is unsure or cannot identify specific triggers",
          "label": "User unsure of triggers",
          "goto": "suggest_journaling"
        }
      ]
    },

generated_chat_model.visualize()

from typing import Literal

bot_turn_prompt = """
You are a personalised, rule-based chatbot.
Follow the message instructions below to generate a personalised message for the user.
In writing this message, take into account the information you have available to you
from the previous messages from the user in this conversation,
and from the user profile below.

<general instructions>
Always try to keep the conversation moving forward.
Ask relevent follow-up questions.
</general instructions>

<message instructions>
{instructions}
</message instructions>

<user profile>
{user_profile}
</user profile>
"""

user_turn_prompt = """
You are a personalised, rule-based chatbot.
Your instructions at this stage in the conversation were as follows:

<message instructions>
{instructions}
</message instructions>

Based on what the user has just said in the conversation history below, categorise their response into one of the following classes:

<possible categories>
{classes}
</possible categories>

Respond with the class label only.
If it is not yet clear how to categorise the user's response, of if it is appropriate to send the user a follow-up message
before moving on to the next stage in the conversation, respond "UNKNOWN" instead.
"""


class ChatGraphRunner:
    """Run a chat graph"""
    def __init__(self, chat_graph: ChatGraph, user_profile: str = ""):
        self.nodes: list[ChatNode] = {
            node.node_label: node for node in chat_graph.nodes
        }
        self.current_node: str = chat_graph.start_node
        self.messages: list[dict] = []
        self.user_profile = user_profile

    def bot_turn(self, node_label: str):
        self.current_node = self.nodes[node_label]
        instructions = self.current_node.instructions
        classes = "\n".join([cat.label for cat in self.current_node.response_categories])
        prompt = bot_turn_prompt.format(
            instructions=instructions, user_profile=self.user_profile
        )
        print(f"\n---\nInstructions\n{instructions}\n---\n")
        # Generate a chatbot message based on the current instructions + conversation history
        bot_response = client.chat.completions.create(
            model="gpt-4o",
            messages=[{"role": "system", "content": prompt}] + self.messages,
        )
        bot_message = bot_response.choices[0].message.content
        self.messages.append({"role": "assistant", "content": bot_message})
        print(f"[Bot@{node_label}]: {bot_message}")
        self.user_turn()

    def user_turn(self):
        node = self.current_node
        classes = "\n".join(
            [
                f"- (Label: {cat.label}) {cat.description}"
                for cat in node.response_categories
            ]
        )
        class_labels = [cat.label for cat in node.response_categories] + ["UNKNOWN"]
        class ClassLabelType(BaseModel):
            label: Literal[tuple(class_labels)]
            
        instructions = self.current_node.instructions
        prompt = user_turn_prompt.format(instructions=instructions, classes=classes)
        user_input = input("[User]: ")
        self.messages.append({"role": "user", "content": user_input})
        classification_response = client.beta.chat.completions.parse(
            model="gpt-4o",
            response_format=ClassLabelType,
            messages=[{"role": "system", "content": prompt}] + self.messages,
        )
        classification_result: ClassLabelType = classification_response.choices[0].message.parsed.label
        if classification_result == "UNKNOWN":
            print(f"[staying on node {node.node_label}]")
            self.bot_turn(node.node_label)
        else:
            # Move to next node
            next_turn_label = [
                v.goto
                for v in node.response_categories
                if v.label == classification_result
            ][0]
            print(f"[moving to node {next_turn_label}]")
            self.bot_turn(next_turn_label)


    def start(self):
        return self.bot_turn(self.current_node)

runner = ChatGraphRunner(
    generated_chat_model, user_profile="User prefers a light, breezy tone, with emojis."
)
runner.start()