Teaching Agents to Write Maintainable Code

Over the last few months we've been putting final polish on the BAML language and prepping for a 1.0 release. We've spent over a year designing the core language, and on that front consider BAML to be nearly complete. But, work remains in making sure BAML meshes well with coding agents, since they are a major part of engineering workflows today.

As it stands, agents can write BAML code that solves problems and passes tests. It is a real turing complete language with a compiler, bytecode VM, formatter, etc. The remaining problems with BAML lie not in the BAML language itself, but in the BAML code that the agents produce. We, and a lot of our testers, noticed that the quality of BAML code written by agents is noticeably worse than what the same agents produce in TypeScript or Rust. For instance, here's part of a chess move generator an agent (GPT-6.1-Sol with Codex) wrote with our current skill installed:

function chess_candidates(position: ChessPosition) -> ChessCandidate[] {
    let board = position.board;
    let candidates: ChessCandidate[] = [];
    let from = 0;
    while (from < 64) {
        let piece = board[from];
        if (chess_color(piece, position.white)) {
            let kind = piece.to_lower_case();
            if (kind == "p") {
                chess_pawn_candidates(position, from, candidates);
            } else {
                let directions = if (kind == "n") {
                    CHESS_KNIGHTS
                } else {
                    CHESS_RAYS
                };
                for (let direction in directions) {
                    let dx = direction[0];
                    let dy = direction[1];
                    let diagonal = dx != 0 && dy != 0;
                    if ((kind != "b" || diagonal) && (kind != "r" || !diagonal)) {
                        let x = from % 8 + dx;
                        let y = from / 8 + dy;
                        while (chess_inside(x, y)) {
                            let to = y * 8 + x;
                            if (chess_color(board[to], position.white)) {
                                break;
                            }
                            candidates.push(chess_candidate(from, to));
                            if (board[to] != "" || kind == "n" || kind == "k") {
                                break;
                            }
                            x += dx;
                            y += dy;
                        }
                    }
                }
                if (kind == "k") {
                    chess_castling_candidates(position, from, candidates);
                }
            }
        }
        from += 1;
    }
    candidates
}

This is a single function, nested seven levels deep. Chess pieces are stored as strings, so a knight is kind == "n", and a typo still compiles. It works, but it doesn't make good use of BAML's rich feature set and nobody would want to maintain code like this. We need some way to improve this, or else it is hard to justify the use of BAML when most people are writing their code with agents.

It stands to reason that existing languages don't suffer as much from this problem, as there is plenty of good code written by humans in the training corpus. Not to mention that the big AI labs are likely optimizing on code quality metrics in these languages. So, without BAML code in the training data, or the support of a large lab, our only real knob for improving code quality is to improve the skill that we ship with the BAML toolchain.

Our current skill mostly just lists out BAML's features: we have match, we have type unions, we have reflection, etc. It does have a small "best practices" section, but most entries just explain what a feature does, and say nothing about when or how to use it; an agent learns that enums exist without ever being told in what context it should use one.

This approach clearly wasn't working; out of ten Codex runs for this chess move generator in BAML, not a single agent opted to use an enum. So, we chose to rewrite most of our skill to focus on best practices and anti-patterns collected from human review of agent-generated BAML code.

What's in the new skill

The new skill describes how to (and how not to!) write BAML, instead of just what features it has. It adds 27 guidelines agents are told to always follow, each with a bad and a good example. These were modeled somewhat after the format of the C++ Core Guidelines. Here's a couple examples relevant to the chess move generator example:

Use a closed type for fixed sets of values. string accepts typos. Use an enum or a literal union.

// Bad
class Ticket {
    status: string,
}

// Good
enum Status {
    Open,
    Closed,
}

class Ticket {
    status: Status,
}

Use match instead of if chains. match is exhaustive, so a new variant is a compile error instead of a silent fallthrough.

// Bad
function hex(c: Color) -> string {
    if (c == Color.Red) {
        "#f00"
    } else {
        "#0f0"
    }
}

// Good
function hex(c: Color) -> string {
    match (c) {
        Color.Red => "#f00",
        Color.Green => "#0f0",
    }
}

With these new guidelines, nine of ten Codex runs with the new skill declared at least one enum in a chess move generator.

Measuring it

This new skill definitely made us feel better, but we wanted to see if we could find some quantitative measure of its effect. So, we brought in a few relatively simple metrics from the literature:

• Cognitive complexity (SonarSource): every if, loop, and match adds one point, plus one for each level of nesting it sits in, and else adds one
• Cyclomatic complexity (McCabe) is the number of independent paths through a function
• NPath (Nejmeh) is the number of possible execution paths. Branches in a row multiply
• Nesting is the deepest stack of ifs, loops, and lambdas
• Exits counts each return and throw, plus finishing normally

// Cognitive 9
function cost_bad(o: Order) -> int {
    if (o.paid) {
        if (o.in_stock) {
            if (o.express) {
                20
            } else {
                5
            }
        } else {
            0
        }
    } else {
        0
    }
}

// Cognitive 4
function cost_good(o: Order) -> int {
    if (!o.paid || !o.in_stock) {
        return 0;
    }
    if (o.express) {
        20
    } else {
        5
    }
}

// Cyclomatic 5
function days_bad(month: string) -> int {
    if (month == "jan") {
        31
    } else if (month == "feb") {
        28
    } else if (month == "mar") {
        31
    } else if (month == "apr") {
        30
    } else {
        0
    }
}

// Cyclomatic 2
const DAYS = { "jan": 31, "feb": 28, "mar": 31, "apr": 30 };

function days_good(month: string) -> int {
    DAYS.get(month) ?? 0
}

// NPath 8
function badges_bad(o: Order) -> string[] {
    let badges: string[] = [];
    if (o.paid) {
        badges.push("paid");
    }
    if (o.in_stock) {
        badges.push("in stock");
    }
    if (o.express) {
        badges.push("express");
    }
    badges
}

// NPath 1
function badges_good(o: Order) -> string[] {
    [
        Badge { name: "paid", on: o.paid },
        Badge { name: "in stock", on: o.in_stock },
        Badge { name: "express", on: o.express },
    ]
        .filter((b) -> {
            b.on
        })
        .map((b) -> {
            b.name
        })
}

// Nesting 3
function fragile_bad(carts: Cart[]) -> string[] {
    let names: string[] = [];
    for (let cart in carts) {
        for (let item in cart.items) {
            if (item.fragile) {
                names.push(item.name);
            }
        }
    }
    names
}

// Nesting 1
function fragile_good(carts: Cart[]) -> string[] {
    carts
        .flat_map((cart) -> {
            cart.items
        })
        .filter((item) -> {
            item.fragile
        })
        .map((item) -> {
            item.name
        })
}

// Exits 4
function label_bad(status: Status) -> string {
    if (status == Status.Open) {
        return "open";
    }
    if (status == Status.Paused) {
        return "paused";
    }
    return "closed";
}

// Exits 1
function label_good(status: Status) -> string {
    match (status) {
        Status.Open => "open",
        Status.Paused => "paused",
        Status.Closed => "closed",
    }
}

No metrics will ever objectively measure code quality, but we felt that these metrics encouraged the type of BAML code that we wanted to see: preference for match statements over if statements, no gigantic functions, a functional style, etc.

We built these into the toolchain with the baml codequality command. By default it prints a summary of the project, and --files and --functions break it down further. On the chess project above:

$ baml codequality --functions
Functions, worst cognitive first
  cognitive  cyclomatic  nesting  npath  exits  lines  function                   location
         50          17        7     59      1     46  chess_candidates           baml_src/chess_impl.baml:218
         44          22        5  1,472      5     51  chess_attacked             baml_src/chess_impl.baml:156

The Eval

With a way to quantify the code quality (or at least a metric related to it), we ran an eval. We had agents (GPT-6.1-Sol with Codex) implement four program skeletons in BAML, each with a public test suite for the agents to work against. Two are routine algorithmic problems:

• a POSIX extended regex engine
• chess move generation

And, the other two are intended to be more practical and make use of the standard library:

• the RealWorld "Conduit" blogging API
• next-departure lookups over GTFS transit feeds

We ran each task ten times for each of four different skill treatments:

1. No skill provided.
2. The old skill
3. The new skill (with guidelines)
4. The new skill, plus an instruction to run baml codequality while working

Results

Cognitive complexity

Each further treatment clearly improves the cognitive complexity of the final outcome. But, some tasks are more highly affected than others. GTFS barely changed between the old and new skills, until we brought in baml codequality.

Cyclomatic complexity

For Chess and Conduit, we see the same clear step-by-step improvement in cyclomatic complexity that we saw with cognitive complexity. Incremental improvement is less clear on Regex and GTFS, but the final new skill including codequality improved the average in every category, though on Regex only because of two unusually low runs.

NPath

For all but GTFS we see clear iterative improvements in scores with each further treatment, and the final runs with codequality always come out on top.

Nesting

Just as with cognitive complexity, there is a clear decrease in nesting as we introduce further treatments, and the final runs with the new skill alongside codequality score the best.

Exits

Exits definitely show the weakest trend, and are the only metric where the runs for the new skill with codequality are not consistently better than other treatments.

Conclusion

By these metrics, the skill improvements and baml codequality definitely had a measurable effect on the code. The worst function beats the old skill on cognitive complexity, cyclomatic complexity, NPath, and nesting on all four tasks.

Whether or not code quality truly increased is up to your trust in the metrics. It seems clear that some of the results are heavily swayed by the specific algorithm or program that you are implementing. For instance, the number of exits did not seem to be very affected by our skill changes in the GTFS runs at all, whereas we saw drastic changes in other programs. An agent may also make decisions that tend to decrease one metric while increasing another or vice versa.

At best these metrics are a proxy for code quality. A metric going up or down a few points is not necessarily good or bad. And, it is possible agents hill-climbed these metrics (even though the instructions specifically said not to). A lot of improvement on these metrics can boil down to just splitting large functions into multiple smaller functions. After all, you can decrease the average number of branches per function if you just split the branches between functions.

The stronger evidence for us is anecdotal. We chose these metrics specifically because we thought they would encourage more of the style of code that we wanted to read. And, when we read the code it just seemed overall better. Here's the same chess logic from an agent run using the new skill:

enum ChessKind {
    Pawn,
    Knight,
    Bishop,
    Rook,
    Queen,
    King,
}

const CHESS_DIAGONALS = [ChessStep { file: -1, rank: -1 }, ...];
const CHESS_STRAIGHTS = [ChessStep { file: -1, rank: 0 }, ...];
const CHESS_RAYS = CHESS_DIAGONALS.concat(CHESS_STRAIGHTS);

function candidates(self, from: int) -> ChessPly[] {
    let piece = self.board[from]
        ?? return [];
    if (piece.side != self.side) {
        return [];
    }
    match (piece.kind) {
        ChessKind.Pawn => self.pawn_moves(from),
        ChessKind.Knight => {
            CHESS_KNIGHT_STEPS.flat_map((step) -> {
                self.ray_moves(from, step, false)
            })
        },
        ChessKind.Bishop => {
            CHESS_DIAGONALS.flat_map((step) -> {
                self.ray_moves(from, step, true)
            })
        },
        ChessKind.Rook => {
            CHESS_STRAIGHTS.flat_map((step) -> {
                self.ray_moves(from, step, true)
            })
        },
        ChessKind.Queen => {
            CHESS_RAYS.flat_map((step) -> {
                self.ray_moves(from, step, true)
            })
        },
        ChessKind.King => {
            CHESS_RAYS
                .flat_map((step) -> {
                    self.ray_moves(from, step, false)
                })
                .concat(self.castle(from, true))
                .concat(self.castle(from, false))
        },
    }
}

function legal_moves(self) -> Move[] {
    CHESS_SQUARES
        .flat_map((from) -> {
            self.candidates(from)
        })
        .filter((ply) -> {
            self.after(ply).king_safe()
        })
        .map((ply) -> {
            ply.external()
        })
}

The difference from the code that was written with the old skill is night and day. The new code makes good use of BAML's features and type system: piece kinds are an enum, the movement directions are constant arrays, a single ray_moves function handles both sliding pieces and one-step pieces, etc.

Unfortunately, we can't put a number on our taste, or our feelings about the code. But, we do feel that code written with the new skill tends to be strictly of higher quality, and much closer to what we would've written by hand.

Try it

If you'd like to try it out, the new skill, including the baml codequality hint, has shipped in the latest version of BAML!