TL;DR
Production conversational agents have to be both flexible (handle whatever the user says) and rigid (never recommend alcohol to a minor, never exceed a mobile screen’s character budget, always cite sources). LLMs are bad at holding onto a long list of hard constraints inside one system prompt — the longer the prompt, the worse the latency and the more constraints silently get dropped. Kakao’s fix is architectural: model the agent as a directed acyclic graph (DAG) where each node is a narrow LLM call with its own short system prompt, its own tools, and its own rule about which parts of the chat history it’s even allowed to see. They bootstrap training data by having annotators correct a prototype agent’s outputs, then fine-tune on that data with a “response masking” trick that only trains each node on the turns it actually produced — because in a graph agent, a single conversation’s history is stitched together from replies generated under totally different system prompts, and training naively on all of it teaches the model to blur its own rules. The payoff, measured on their deployed Korean e-commerce assistant (“AI Shopping Mate,” live on KakaoTalk since December 2024): a 27B–32B in-house model beats GPT-4o on format adherence, tool-call accuracy, and head-to-head human preference in every category except casual chat.
Problem & Motivation
The pain is concrete and it’s the one every team building a compliance-heavy agent eventually hits: LLMs generate the next token probabilistically, and “please don’t recommend cigarettes to a 16-year-old” is a soft nudge in a prompt, not a hard constraint on the sampler. Three things compound this in production:
- Business rules are strict, but LLMs are not. An e-commerce agent must retrieve real product metadata rather than answer from pretrained knowledge, or it can hallucinate a recommendation that violates age/content policy — and that’s a failure even if the user is happy with the result.
- Formatting requirements pile up per surface. A mobile messenger interface wants short responses, emoji bullets, and specific card layouts. Certain product categories add their own constraints on top — no hype language, mandatory source attribution, a required “brand story” before the pitch.
- The naive fix (write it all into the system prompt) doesn’t scale. More rules means a longer prompt, and longer prompts measurably hurt both latency and reasoning accuracy (this is a known LLM failure mode, not speculation — the paper cites Levy et al. 2024 on exactly this). So the more compliant you try to make one big prompt, the worse the underlying model gets at following it.
The paper’s own baseline makes this concrete: “Basic” (B) — one system prompt with every rule concatenated in, plus the vendor’s native tool-calling — is what most teams ship first. It’s also the worst performer on every metric they measure, including for GPT-4o.
What’s New (Core Contribution)
Two genuinely new pieces, plus a validated real-world deployment as the receipts.
-
Workflow-graph agent architecture (Multi-State DAG Framework). Before: one model, one system prompt, one flat toolset — or, in existing graph frameworks like LangGraph/Dify, a graph where nodes are just “does a computation and picks the next node,” with no opinion on prompt design or history handling. Now: each LLM-calling node gets its own short system prompt scoped to only the rules relevant in that state, its own few-shot examples, its own tools with typed schemas, and a custom routine that can rewrite the conversation history before the LLM sees it (e.g., strip everything except the last turn). This is the load-bearing idea: constraint-following gets easier when each individual LLM call only has to hold a handful of rules in its head, not the union of every rule in the product.
-
Node-aware fine-tuning via response masking. Before: standard multi-turn fine-tuning trains the loss on every assistant turn in a conversation, under whatever system prompt is attached to that training example. Now: because a single graph-agent conversation is stitched from turns generated by different nodes (different system prompts), training node A’s example on a turn that was actually produced by node B teaches the model that node A’s rules include node B’s rules. The fix is to mask the loss so each training example only supervises the turns generated by its own node. Simple to state, but it’s the detail that makes fine-tuning a graph agent work at all instead of quietly degrading it.
-
A real, load-bearing case study, not a toy benchmark. The framework isn’t validated on an academic leaderboard — it’s the architecture behind “AI Shopping Mate,” live on KakaoTalk and the web since December 2024, covering over a million products, with human A/B “battle” tests against GPT-4o run on real traffic categories (safety probes, product recommendation, messenger features, regular chat).
What’s not new: graphs-as-agents is the standard LangGraph/Dify pattern, and constrained decoding for schema-valid tool calls is a known technique (Willard & Louf 2023; XGrammar). The contribution is the specific combination — per-node prompt isolation + history rewriting + schema-constrained handoffs + node-aware loss masking — assembled and proven under real production traffic.
How It Works (Technically)
The graph. Formally the agent is $G = (V, E)$: a set of nodes $V$ and directed edges $E \subset V \times V$. Every node $v$ has a routine $f_v$ that either calls an external tool or calls an LLM, and returns a pair $(o_v, v_n)$ — its output, and the next node to visit. Running the agent is just graph traversal: start at an entry node $v_{init}$, keep moving to whatever successor node the current node’s output points to, stop at $v_{final}$, return the output there. This part is exactly what LangGraph/Dify already give you.
The paper’s addition is what happens inside an LLM-calling node:
- A scoped system prompt $s_v$. Not the whole rulebook — just the rules, formatting requirements, and few-shot examples relevant to what this node does.
recommend_reasoncarries Markdown/emoji formatting rules;purchase_messagecarries the purchase-confirmation format. - A history-rewriting routine (
modify_history). By default a node sees the full conversation so far. But some nodes explicitly don’t want that —purchase_messagestrips everything except the final turn (the purchase details) before generating its response. Why: irrelevant earlier chat (browsing, small talk, other products discussed) is exactly the kind of context that causes hallucination once the task has narrowed to “confirm this specific purchase.” Limiting the input is a cheap, deterministic way to cut hallucination risk that a bigger system prompt can’t buy you. - Typed schemas + constrained decoding at the LLM→tool boundary. Tool nodes declare an input schema and an output schema. When an LLM node’s output needs to become a tool call’s input, constrained decoding (grammar-constrained generation, à la Willard & Louf 2023 / XGrammar) forces the output to actually satisfy that schema, instead of hoping the model emits valid JSON.
One trace, start to finish (the gift-recommendation scenario from the paper’s Figure 2): a user opens the chat and says “recommend a wine for my friend’s birthday.”
- The conversation starts at
chat, the general-purpose entry node. Its job is to either answer directly (small talk, “who are you”) or recognize this needs a task-specific node and emit a tool-call-shaped hand-off. - Because the message implies a product search,
chatroutes to asearch_productstool node. Its input schema requires a query string and (optionally) category/price filters, extracted from the user’s message; constrained decoding guarantees the arguments the LLM produced are well-formed before the tool actually runs. - The tool returns product results, which flow into
recommend_reason, an LLM node. This node’s system prompt encodes response formatting (Markdown, emoji bullets, mobile-friendly length) and — by default — sees the full conversation history, because knowing what led to this recommendation is genuinely useful context here. - The user says “I’ll take the second one.” The graph routes to a
purchase_gifttool node (typed schema: product id, recipient, etc.), then topurchase_message, another LLM node — but this one’smodify_historyroutine deliberately discards everything except the last turn before generating the confirmation. The earlier wine-browsing chat is irrelevant now and only risks leaking into the confirmation text. purchase_messagereachesfinal. For the next user turn, the graph restarts fromchat(an entry point does not mean the process ends — the graph is re-entered per turn or per new task within the same multi-turn session).
Notice what happened: no single LLM call ever had to hold “formatting rules AND purchase-confirmation rules AND search-query-extraction rules” in one prompt. Each node held exactly what it needed, no more.
Architecture & data flow
flowchart TD chat["chat (LLM node)<br/>entry point, routes or free-chats"] search[["search_products (tool node)<br/>schema: query, filters"]] rec["recommend_reason (LLM node)<br/>sees full history, formatting rules"] purchase[["purchase_gift (tool node)<br/>schema: product id, recipient"]] pmsg["purchase_message (LLM node)<br/>modify_history: keep last turn only"] final(["final"]) chat -->|"routes via tool-call"| search search --> rec rec -->|"user confirms"| purchase purchase --> pmsg pmsg --> final chat -->|"out-of-scenario message"| chat final -.->|"next user turn"| chat
A message moving through the workflow graph, node by node. Each active node lights up with the slice of context it's actually allowed to see — notice how little that is compared to "everything said so far." Click to step through the trace above.
Fine-tuning with response masking — demystified
The training format itself is unremarkable: for each LLM node $v$, format its interactions as a chatbot-style sequence $(s_v, x_1, o_1, x_2, o_2, \dots, x_n, o_n)$ — system prompt, then alternating observations ($x_i$: user messages or tool results) and agent outputs ($o_i$).
The problem is what “the conversation history” means in a graph agent. Take a two-node graph with nodes $v_1$ and $v_2$. A real conversation might look like $(s_{v_1}, x_1, o_1, x_2, o_2, x_3, o_3)$ from $v_1$‘s point of view — except $o_2$ was actually generated by $v_2$, under $v_2$‘s system prompt, following $v_2$‘s rules. If you fine-tune naively (loss on every assistant turn), you’re telling the model “under $s_{v_1}$‘s instructions, the correct thing to say is $o_2$” — except $o_2$ was written to satisfy a different set of constraints. Do this across thousands of examples and you’ve trained the model to average together rules that were supposed to stay separate. This is exactly the kind of subtle data-contamination bug that would otherwise show up months later as “the model randomly ignores formatting rules in production” with no obvious root cause.
The fix: mask the loss so a training example for node $v$ only backpropagates through the tokens of turns that node $v$ actually generated. $o_2$ still appears in $v_1$‘s training example as context (the model needs to know it happened), but the loss for those tokens is zeroed out — the model is never asked to reproduce it. Practically, they implement this with Axolotl’s segment-level input masking, so it’s a labeling scheme on top of ordinary supervised fine-tuning, not a new loss function.
The mechanism, simplified
# One training example per LLM node v. mask=1 means "compute loss here",
# mask=0 means "context only, don't supervise this span" — because
# these tokens were generated by a DIFFERENT node's system prompt.
def build_training_example(node_id, turns):
"""
turns: ordered list of {speaker, node_id, text}
speaker in {"user", "tool", "assistant"}
"""
tokens, loss_mask = [], []
system_prompt = SYSTEM_PROMPTS[node_id] # s_v: short, node-scoped
tokens += tokenize(system_prompt)
loss_mask += [0] * len(tokens) # never train on the prompt itself
for turn in turns:
turn_tokens = tokenize(turn["text"])
tokens += turn_tokens
if turn["speaker"] == "assistant" and turn["node_id"] == node_id:
loss_mask += [1] * len(turn_tokens) # this node produced it -> supervise
else:
loss_mask += [0] * len(turn_tokens) # other node's output, or user/tool -> context only
return tokens, loss_mask
def masked_loss(logits, labels, loss_mask):
# standard cross-entropy, but zeroed wherever loss_mask == 0
per_token_loss = cross_entropy(logits, labels, reduction="none")
return (per_token_loss * loss_mask).sum() / loss_mask.sum()
A conversation stitched from two nodes' turns. Toggle masking on/off to see which tokens the fine-tuning loss actually touches — and which ones would otherwise leak $v_2$'s rules into $v_1$'s training signal.
Data collection, one level up. None of this fine-tuning works without training data shaped like real graph traversals, and generating that data is itself the paper’s second practical contribution: annotators can’t easily write a correct multi-step answer to “recommend a wine for sirloin steak” from scratch — it requires imagining tool calls and graph hops in their head. So the paper bootstraps: build a prototype agent (GPT-4o wired into the same workflow graph), have annotators interact with it as end users, log the entire graph traversal (tool calls, arguments, node sequence, outputs), then have annotators correct whatever the prototype got wrong — assisted by automated checkers (a static type-checker for tool-call JSON catches the most common annotator mistake: malformed arguments).
flowchart LR A["Prototype agent<br/>(GPT-4o + workflow graph)"] --> B["Annotator interacts<br/>as end user"] B --> C["Full traversal logged<br/>(nodes, tool calls, outputs)"] C --> D["Annotator corrects errors<br/>+ automated JSON checkers"] D --> E["Node-tagged dataset<br/>(sv, x1,o1,...,xn,on) per node"] E --> F["Fine-tune with<br/>response masking"] F --> G["Deployed model<br/>(AI Shopping Mate)"]
Built on Prior Work
| Prior idea | What it gave | What this paper changes |
|---|---|---|
| Rule-based dialog managers (Rasa, Talkamatic) | Reliable, interpretable state tracking via hand-authored rules | Keeps the “state” concept but replaces hand-coded transition logic with LLM reasoning inside each state |
| Tool-use agents (Toolformer, ToolLLM, ReAct, CoT/ToT) | The ability to reason and call tools at all | Doesn’t add new reasoning — structures where and under what constraints that reasoning happens, via per-node scoping |
| Graph-agent frameworks (LangGraph, Dify) | The abstraction: agent-as-graph-traversal, node returns (output, next-node) | Adds everything the abstraction is silent on: per-node prompt isolation, history-rewriting routines, schema-constrained handoffs, and a fine-tuning method for the resulting message histories |
| MARCO (guardrail + retry for output validity) | Explicitly measuring output validity, not just task success | Builds the constraint into the architecture (short, node-scoped prompts) instead of a post-hoc guardrail-and-retry loop, avoiding MARCO’s latency/accuracy cost |
| Amazon Bedrock Agents / Google Vertex AI + LangChain | Industrial precedent for post-processing and workflow-based control | Adds the specific fine-tuning recipe (response masking) these platforms don’t prescribe |
| Constrained decoding (Willard & Louf 2023, XGrammar) | Grammar-constrained generation for structured output | Applies it specifically at LLM-node → tool-node handoffs inside the graph |
Results & Evidence
Evaluated on 161 held-out conversations (2,100 turns) across three metrics: accuracy (right tool, right arguments, judged by an LLM-as-a-Judge given the flexibility of natural-language arguments), format adherence (binary, code-validated against the required message format), and response quality/validity (1–3 scale, LLM-as-a-Judge against a human reference).
Four architectures were compared per model: Basic (B) — one concatenated system prompt, vendor-native tool calling; Workflow Graph (WG) — the graph architecture with no fine-tuning; Workflow Graph + Fine-Tuning (WG-FT) — graph plus the response-masking fine-tune (only possible for open-weight/internal models, not GPT-4o).
| Qwen2.5 32B (B→WG→WG-FT) | Gemma 3 27B (B→WG→WG-FT) | Internal 27–32B (B→WG→WG-FT) | GPT-4o (B→WG) | |
|---|---|---|---|---|
| Accuracy | 0.578 → 0.616 → 0.884 | 0.622 → 0.711 → 0.887 | 0.744 → 0.790 → 0.890 | 0.864 → 0.888 |
| Format adherence | 0.734 → 0.813 → 0.969 | 0.692 → 0.882 → 0.966 | 0.655 → 0.951 → 0.987 | 0.778 → 0.964 |
| Response quality (1–3) | 2.816 → 2.831 → 2.880 | 2.821 → 2.849 → 2.911 | 2.893 → 2.874 → 2.953 | 2.856 → 2.882 |
What jumps out: the graph structure alone (WG, no fine-tuning) closes most of the gap to GPT-4o, and fine-tuning on top of it lets a 27–32B open/internal model beat GPT-4o’s own workflow-graph score on every metric. The internal model’s WG-FT numbers (0.890 accuracy, 0.987 format adherence) beat GPT-4o’s WG numbers (0.888, 0.964) outright. Format adherence sees the biggest jump from the graph structure itself — shorter, focused prompts are dramatically better at holding a formatting rule than one long prompt, regardless of model.
Beyond the offline benchmark, the paper reports a real production signal: deployed on live traffic as “AI Shopping Mate,” their internal model beat GPT-4o head-to-head (anonymized human preference, both models wired to identical tools/data) in Safety (60.5%), Product recommendation (82.4%), and Messenger features (60.6%), losing only Regular chat (42.4% — attributed to GPT-4o’s language fluency, an axis the LLM-as-a-Judge evaluation apparently didn’t capture well).
What this does and doesn’t establish. It’s a genuine, production-validated result — not a leaderboard flex. But treat it with real caveats:
- No ablation isolating the graph from the fine-tuning. WG-FT bundles two changes (architecture + training). We know the combination works; we don’t know the marginal contribution of response masking specifically versus “any fine-tuning on graph-shaped data would have helped.”
- The Basic baseline is a strawman by construction. Concatenating every rule into one prompt without any structure is close to a worst-case naive implementation, not necessarily what a careful prompt engineer would ship. The comparison is still useful — it’s the failure mode most teams actually hit — but “up to 14–27% improvement” is measured against that starting point, not against a well-tuned single-prompt baseline.
- No comparison to the closest prior work. MARCO (the paper’s own cited closest competitor) and vanilla LangGraph/Dify implementations are discussed but never benchmarked head-to-head.
- LLM-as-a-Judge has a documented blind spot here: it apparently under-weighted fluency, which is exactly why GPT-4o won “regular chat” in the human battle test despite comparable offline scores. The paper’s own Limitations section flags this — evaluation validity is an open problem, not a solved one.
- Single domain, single language, single company’s annotators. Korean e-commerce, one annotator pool with acknowledged demographic skew. Generalization to other domains/languages is asserted, not tested.
How You’d Use It
This maps almost directly onto agent work you’re already doing, and it’s the fix worth reaching for whenever your own “AI assistant” is one sprawling system prompt that keeps failing compliance checks.
- As an architecture pattern on LangGraph (or your own MAS routing layer) — your harness. The paper’s contribution is orthogonal to the graph library itself — it’s a set of conventions you can bolt onto LangGraph today: give every node a scoped prompt instead of a shared one, add an explicit
modify_historyhook per node (most graph frameworks let you control what state gets passed forward — this just makes it a deliberate design decision instead of an accident), and use structured-output/schema validation (Pydantic + Instructor, or grammar-constrained decoding) at every LLM→tool handoff. - As a diagnostic for “why is our agent inconsistent.” If your own single-prompt agent is randomly dropping formatting rules or compliance constraints as the prompt grows, this paper’s core claim — long prompts degrade instruction-following, and splitting the prompt across states fixes it — is a testable, cheap-to-try hypothesis before reaching for fine-tuning at all.
- As the missing half of the fine-tuning story if you’re fine-tuning agents. If you’re already fine-tuning a model behind a graph-based agent (or planning to), response masking is close to a free correctness fix — it’s a data-labeling change, not a new training pipeline, and it directly prevents the failure mode of one node’s rules leaking into another’s.
- As a data-collection methodology for your own fine-tuning work. The prototype-agent-plus-annotator-correction loop is a genuinely reusable recipe whenever you need bespoke fine-tuning data for a complex multi-step agent and don’t have it yet: bootstrap with a strong general model, log full traversals, correct with cheap automated checks (schema validators) before human review.
Build Your Own (Minimal Recipe)
You can prototype the core idea — graph structure + response masking — in a few days without touching a foundation model.
Components, in build order:
- Define the graph. A handful of nodes as config:
{id, type: llm|tool, system_prompt, tools/schema, modify_history_fn}. Start with 4–6 nodes covering your top user intents plus a generalchatentry node. - Router logic. The entry node’s job is dispatch: either answer directly or emit a tool-call-shaped signal that your orchestrator maps to the next node. LangGraph gives you this out of the box; a plain
dict-based state machine works for a toy version. modify_historyhooks. Default: pass full history through. Add an override per node where it matters (e.g., “keep only the last turn” for any confirmation/checkout-style node) — this is a five-line function, but deciding which nodes need it is the actual design work.- Schema-constrained handoffs. Validate (or grammar-constrain) any LLM output that becomes a tool call’s input. Pydantic + retry-on-validation-failure gets you 80% of the value without needing full grammar-constrained decoding.
- Bootstrap a dataset. Run a strong general model (GPT-4o/Claude/whatever you have) through your graph, have a few people interact with it and correct outputs, and log the full traversal — which node produced which turn is the one field you must not lose.
- Fine-tune with masking. Tag every training turn with its producing node id, then zero the loss on any turn not produced by the node the example is being trained for. Axolotl’s segment-level masking supports this directly; otherwise it’s a straightforward custom collator.
The genuinely hard parts:
- Deciding the graph shape. How many nodes, where the boundaries are, and which rules belong to which node is a design problem with no formula — get it wrong and you either end up back at one bloated prompt (too few nodes) or an unmaintainable spaghetti graph (too many). Expect iteration; this is where most of your real effort goes.
- Getting node provenance right in the training data. The masking trick is only as good as your bookkeeping — if you don’t reliably tag every historical turn with the node that generated it, you can’t compute the mask correctly, and the failure mode (silent supervision leakage) doesn’t show up until you’re debugging weird production behavior later.
Reach for: LangGraph for the graph/orchestration layer, Pydantic + Instructor (or Outlines/XGrammar) for schema-constrained tool calls, Axolotl + LoRA for the fine-tune, vLLM or SGLang for serving.
How to Improve It
The paper is honest about several open problems, and there’s more leverage than what’s listed:
- Ablate architecture from training. Run WG (no fine-tune) vs. a model fine-tuned on graph-shaped data without masking vs. WG-FT, to isolate how much of the gain is the graph structure, how much is having graph-shaped data at all, and how much is specifically the masking trick.
- Replace annotator-bottlenecked data collection with LLM-simulated users — the paper explicitly flags this as future work. An LLM playing the “user” role against your prototype agent could 10x your data volume; the open question is whether simulated-user data preserves enough realism to matter, which is directly testable by comparing downstream WG-FT performance on real vs. simulated data.
- Investigate the LLM-as-a-Judge / human-preference gap systematically, not just note it. The “regular chat” fluency miss is a concrete, measurable blind spot — worth a targeted study on what judge prompts or additional criteria (fluency, naturalness) would close it, since this evaluation gap will bite anyone using LLM-as-a-Judge for agent QA, not just this team.
- Automate graph-shape validation. Borrow the kind of static consistency checking used in other declarative-agent formalisms (e.g., flagging a node whose trigger conditions can never fire, or two nodes with contradictory formatting rules) so graph design stops being pure trial-and-error as the node count grows.
- Benchmark against the actual named competitors. MARCO and a vanilla LangGraph implementation are discussed but never run head-to-head against WG/WG-FT — that comparison is the most obvious missing piece of evidence, and cheap to produce given the framework already exists.
Glossary
- DAG (directed acyclic graph) — a graph where edges only point “forward,” so there’s no way to loop back to an earlier node; used here to model the agent’s possible conversation paths.
- Workflow graph / node / state — the paper’s term for one step in the DAG: either an LLM call with a scoped system prompt, or a tool call with typed input/output schemas.
- System prompt — the instructions given to an LLM before the conversation; here, scoped per node instead of one shared prompt for the whole agent.
- Constrained decoding — forcing an LLM’s output tokens to match a grammar/schema (e.g., valid JSON matching a specific shape) at generation time, instead of hoping and validating after the fact.
modify_history— a per-node function that edits the conversation history before the LLM sees it (e.g., drop everything except the last turn) to reduce hallucination from irrelevant context.- LLM-as-a-Judge — using an LLM (here, o3-mini) to score another model’s output against criteria or a reference answer, instead of (or alongside) human evaluation.
- Response masking / loss masking — during fine-tuning, zeroing the loss on specific tokens (here, any assistant turn not produced by the node currently being trained) so the model isn’t supervised on data outside that node’s rules.
- Format adherence — a binary metric: does the response match the required output format (length, structure, emoji usage) for its surface (here, mobile chat)?
- Response validity / quality — a 1–3 LLM-judged score for how clear, helpful, and relevant a response is versus a human-written reference.
- LoRA (Low-Rank Adaptation) — a parameter-efficient fine-tuning method that trains small added weight matrices instead of the full model; used here via the Axolotl framework, then merged into the base model before serving.
- Axolotl — an open-source fine-tuning framework; used here for its support of segment-level (per-span) loss masking, which is what makes response masking practical to implement.
- vLLM / SGLang — high-throughput LLM serving engines used to deploy the open-weight and internal models in production.