Code from eointravers.com/blog/structured-conversation/

structured-conversation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "source": [
    "> Code from eointravers.com/blog/structured-conversation/ \n",
    "\n",
    "```yaml\n",
    "title: \"LLMs and Structured Conversation Graphs: The best of both worlds\"\n",
    "author: \"Eoin Travers\"\n",
    "date: \"2025-02-18\"\n",
    "tags:\n",
    "- AI\n",
    "- Engineering\n",
    "- Chatbots\n",
    "```\n",
    "\n",
    "## The Problem\n",
    "\n",
    "Creating open-ended LLM-powered chatbots is easy:\n",
    "you write a system prompt, stick it at the start of the conversation history,\n",
    "send the messages to an LLM, and away you go.\n",
    "Doing this gets you something that's very flexible, and quick to build,\n",
    "but lose a lot of control over the structure of your chatbot's conversations.\n",
    "By contrast, in traditional pre-LLM chatbots typically\n",
    "follow a an explicit, structured set of rules,\n",
    "powered by a [dialogue management system](https://en.wikipedia.org/wiki/Dialog_manager)\n",
    "to decide which pre-written message to show the user on each turn.\n",
    "These give you control, but not flexibility, and they require a lot of manual work.\n",
    "Can we have the best both worlds?\n",
    "\n",
    "Meanwhile, it's Q1 2025, and all the talk in AI engineering circles\n",
    "(at least the circles I like to spend time in)\n",
    "is about graphs/state machines,\n",
    "e.g. in [pydantic.ai](https://ai.pydantic.dev/graph/) (❤️)\n",
    "[LangGraph](https://www.langchain.com/langgraph) (🤔),\n",
    "and, at a lower level of abstraction,\n",
    "the state machines that power [Outlines](https://dottxt-ai.github.io/outlines/latest/reference/generation/structured_generation_explanation/) (🧠).\n",
    "Do you see where I'm going with this?\n",
    "\n",
    "\n",
    "## The Idea\n",
    "\n",
    "A structured conversation can be defined in terms of a *graph*.\n",
    "\n",
    "Each *node* of the graph consists of:\n",
    "\n",
    "- Instructions for what the chatbot should say.\n",
    "  In traditional chatbots, these would be pre-written, but with LLMs\n",
    "  we can use system prompts instead.\n",
    "- Conditional *edges* that that take you to the next node, depending on the user's response.\n",
    "  Traditional chatbots might use keywords, regular expressions, or standalone machine learning/intent detection models\n",
    "  to decide which edge to follow.\n",
    "  We get to just use LLMs. \n",
    "\n",
    "\n",
    "## Step 1: Define Structures\n",
    "\n",
    "To start, let's create pydantic models that can be used to define these conversation graphs.\n",
    "I'll also add some methods to the `ChatGraph` class to make it easier to work with.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseModel, Field\n",
    "from graphviz import Digraph\n",
    "from textwrap import wrap as twrap\n",
    "\n",
    "def wrap(s, w=20):\n",
    "    \"\"\"Wrap text to a given width\"\"\"\n",
    "    return \"\\n\".join(twrap(s, w))\n",
    "\n",
    "class ChatGraph(BaseModel):\n",
    "    \"\"\"Stores the full graph for a given conversation\"\"\"\n",
    "\n",
    "    nodes: list[\"ChatNode\"]\n",
    "    start_node: str = Field(description=\"Label of the node to begin from\")\n",
    "\n",
    "    def write_json(self, path: str):\n",
    "        with open(path, \"w\") as f:\n",
    "            f.write(self.model_dump_json(indent=2))\n",
    "\n",
    "    @classmethod\n",
    "    def read_json(cls, path: str):\n",
    "        with open(path, \"r\") as f:\n",
    "            return cls.model_validate_json(f.read())\n",
    "\n",
    "    def visualize(self, render: bool = False):\n",
    "        \"\"\"Visualise the graph using graphviz\"\"\"\n",
    "        dot = Digraph(graph_attr={\"rankdir\": \"LR\"})\n",
    "        for node in self.nodes:\n",
    "            src = node.node_label\n",
    "            dot.node(src, wrap(node.instructions, 20))\n",
    "            for rc in node.response_categories:\n",
    "                dot.edge(src, rc.goto, wrap(rc.description, 20))\n",
    "        if render:\n",
    "            dot.render(view=True, format=\"svg\")\n",
    "        return dot\n",
    "\n",
    "class ChatNode(BaseModel):\n",
    "    \"\"\"Individual nodes\"\"\"\n",
    "\n",
    "    node_label: str = Field(description=\"A unique, informative ID for this node\")\n",
    "    instructions: str = Field(\n",
    "        description=(\n",
    "            \"Instructions for the chatbot loosely outlining what it should say at this step. \"\n",
    "            \"The chatbot will use these instructions plus background knowledge about the user to personalise its actual message.\"\n",
    "        )\n",
    "    )\n",
    "    response_categories: list[\"ChatResponseCategory\"] = Field(\n",
    "        description=\"A list of possible categories of user responses\"\n",
    "    )\n",
    "\n",
    "class ChatResponseCategory(BaseModel):\n",
    "    \"\"\"A possible classification for the user's message\"\"\"\n",
    "\n",
    "    description: str = Field(description=\"A clear description of the category\")\n",
    "    label: str = Field(description=\"A simple, unambiguous label for this response category\")\n",
    "    goto: str = Field(description=\"The label of the ChatNode to go to if this response is given\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "source": [
    "\n",
    "## Step 2: Cheat by using AI to create my graph\n",
    "\n",
    "Now, in principle these conversation graphs can be created by hand,\n",
    "by editing the JSON directly, but that wouldn't be very AI.\n",
    "Instead, let's take advantage of the fact that we can get an LLM\n",
    "to generate data that is structured according to our schema,\n",
    "using either the Instructor package, or OpenAI's structured outputs.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dotenv import load_dotenv\n",
    "from openai import OpenAI\n",
    "load_dotenv()\n",
    "\n",
    "client = OpenAI()\n",
    "\n",
    "# This takes about 14 seconds\n",
    "generated_chat_model_response = client.beta.chat.completions.parse(\n",
    "    model=\"gpt-4o\",\n",
    "    response_format=ChatGraph,\n",
    "    messages=[\n",
    "        {\n",
    "            \"role\": \"system\",\n",
    "            \"content\": (\n",
    "                \"Use the schema provided to generate a short ChatModel for a basic CBT chatbot that helps the user deal with anxiety. \"\n",
    "                \"Each node should include at least two response categories. \"\n",
    "                \"Instructions for non-terminal nodes should always include questions for the user to keep the conversation moving. \"\n",
    "                \"Aim for around 7 nodes in the conversation.\"\n",
    "            ),\n",
    "        }\n",
    "    ],\n",
    ")\n",
    "generated_chat_model = generated_chat_model_response.choices[0].message.parsed\n",
    "\n",
    "!mkdir -p chats\n",
    "generated_chat_model.write_json(\"chats/cbt1.json\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Uncomment to load from file\n",
    "# generated_chat_model = ChatGraph.read_json(\"chats/cbt1.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\n",
    "    generated_chat_model.model_dump_json(indent=2)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generated_chat_model.visualize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is already pretty nice, if I do say so myself.\n",
    "\n",
    "## Step 3: A Chat Graph Runner\n",
    "\n",
    "Running this conversation is just a case of traversing the graph,\n",
    "generating the appropriate chatbot messages at each node,\n",
    "and deciding where to go next based on the user's response.\n",
    "Here's some code.\n",
    "This clearly isn't production-ready, but I think it illustrates the idea.\n",
    "I've also added support for additional personalisation through a \"user profile\",\n",
    "just because.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Literal\n",
    "\n",
    "bot_turn_prompt = \"\"\"\n",
    "You are a personalised, rule-based chatbot.\n",
    "Follow the message instructions below to generate a personalised message for the user.\n",
    "In writing this message, take into account the information you have available to you\n",
    "from the previous messages from the user in this conversation,\n",
    "and from the user profile below.\n",
    "\n",
    "<general instructions>\n",
    "Always try to keep the conversation moving forward.\n",
    "Ask relevent follow-up questions.\n",
    "</general instructions>\n",
    "\n",
    "<message instructions>\n",
    "{instructions}\n",
    "</message instructions>\n",
    "\n",
    "<user profile>\n",
    "{user_profile}\n",
    "</user profile>\n",
    "\"\"\"\n",
    "\n",
    "user_turn_prompt = \"\"\"\n",
    "You are a personalised, rule-based chatbot.\n",
    "Your instructions at this stage in the conversation were as follows:\n",
    "\n",
    "<message instructions>\n",
    "{instructions}\n",
    "</message instructions>\n",
    "\n",
    "Based on what the user has just said in the conversation history below, categorise their response into one of the following classes:\n",
    "\n",
    "<possible categories>\n",
    "{classes}\n",
    "</possible categories>\n",
    "\n",
    "Respond with the class label only.\n",
    "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\n",
    "before moving on to the next stage in the conversation, respond \"UNKNOWN\" instead.\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "class ChatGraphRunner:\n",
    "    \"\"\"Run a chat graph\"\"\"\n",
    "    def __init__(self, chat_graph: ChatGraph, user_profile: str = \"\"):\n",
    "        self.nodes: list[ChatNode] = {\n",
    "            node.node_label: node for node in chat_graph.nodes\n",
    "        }\n",
    "        self.current_node: str = chat_graph.start_node\n",
    "        self.messages: list[dict] = []\n",
    "        self.user_profile = user_profile\n",
    "\n",
    "    def bot_turn(self, node_label: str):\n",
    "        self.current_node = self.nodes[node_label]\n",
    "        instructions = self.current_node.instructions\n",
    "        classes = \"\\n\".join([cat.label for cat in self.current_node.response_categories])\n",
    "        prompt = bot_turn_prompt.format(\n",
    "            instructions=instructions, user_profile=self.user_profile\n",
    "        )\n",
    "        print(f\"\\n---\\nInstructions\\n{instructions}\\n---\\n\")\n",
    "        # Generate a chatbot message based on the current instructions + conversation history\n",
    "        bot_response = client.chat.completions.create(\n",
    "            model=\"gpt-4o\",\n",
    "            messages=[{\"role\": \"system\", \"content\": prompt}] + self.messages,\n",
    "        )\n",
    "        bot_message = bot_response.choices[0].message.content\n",
    "        self.messages.append({\"role\": \"assistant\", \"content\": bot_message})\n",
    "        print(f\"[Bot@{node_label}]: {bot_message}\")\n",
    "        self.user_turn()\n",
    "\n",
    "    def user_turn(self):\n",
    "        node = self.current_node\n",
    "        classes = \"\\n\".join(\n",
    "            [\n",
    "                f\"- (Label: {cat.label}) {cat.description}\"\n",
    "                for cat in node.response_categories\n",
    "            ]\n",
    "        )\n",
    "        class_labels = [cat.label for cat in node.response_categories] + [\"UNKNOWN\"]\n",
    "        class ClassLabelType(BaseModel):\n",
    "            label: Literal[tuple(class_labels)]\n",
    "            \n",
    "        instructions = self.current_node.instructions\n",
    "        prompt = user_turn_prompt.format(instructions=instructions, classes=classes)\n",
    "        user_input = input(\"[User]: \")\n",
    "        self.messages.append({\"role\": \"user\", \"content\": user_input})\n",
    "        classification_response = client.beta.chat.completions.parse(\n",
    "            model=\"gpt-4o\",\n",
    "            response_format=ClassLabelType,\n",
    "            messages=[{\"role\": \"system\", \"content\": prompt}] + self.messages,\n",
    "        )\n",
    "        classification_result: ClassLabelType = classification_response.choices[0].message.parsed.label\n",
    "        if classification_result == \"UNKNOWN\":\n",
    "            print(f\"[staying on node {node.node_label}]\")\n",
    "            self.bot_turn(node.node_label)\n",
    "        else:\n",
    "            # Move to next node\n",
    "            next_turn_label = [\n",
    "                v.goto\n",
    "                for v in node.response_categories\n",
    "                if v.label == classification_result\n",
    "            ][0]\n",
    "            print(f\"[moving to node {next_turn_label}]\")\n",
    "            self.bot_turn(next_turn_label)\n",
    "\n",
    "\n",
    "    def start(self):\n",
    "        return self.bot_turn(self.current_node)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, I've roleplayed a chat with this bot myself,\n",
    "just to illustrate how it works.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "runner = ChatGraphRunner(\n",
    "    generated_chat_model, user_profile=\"User prefers a light, breezy tone, with emojis.\"\n",
    ")\n",
    "runner.start()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "...at this point, I had had enough and ended the chat.\n",
    "It works!\n",
    "Clearly, there's room for improvement, but I think this is a nice pattern,\n",
    "with a lot of potential for use in cases where you want finer-grained control\n",
    "than can be achieved with the standard approach of plugging your users in to the LLM firehose.\n",
    "\n",
    "This code is available as a gist [here]().\n",
    "Like the rest of this website, it's MIT licensed,\n",
    "but if you end up using it I'd love to hear from you.\n"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}