Using AI as a judge with Apache Airflow

The Problem

How do you know an agent's answer is right before anyone reads it? For a while, we didn't. Astronomer Academy has a Slack bot we call the EduAgent: you mention it in a channel with a question like "how many people passed Apache Airflow 3 Fundamentals last month?" and it replies in the thread with numbers it pulled from Snowflake.

Most of its answers were fine, and that made the bad ones hard to spot. Someone asked for the top 100 organizations and got a tidy top-10 table (90 rows short), followed by a paragraph about a few notable names. Someone else asked whether a customer's employees had signed in to the Academy. The bot said 14 had. The query behind it only counted accounts with the customer's email domain, and it never looked at sign-ins.

Neither answer had a SQL error, and every number in them came from a real query. Each one answered a slightly different question from the one that was asked, in the same confident tone these models use for everything, and nobody in the thread had a reason to doubt it. If you've pointed an agent at a warehouse, you've probably shipped an answer like these too.

So we put a second agent in front of the first. Its only job is to grade each answer before it reaches Slack, a pattern usually called LLM-as-a-judge, and both agents run on Airflow with the Common AI provider.

TL;DR: every data answer is graded before Slack sees it

The data agent is an @task.agent task that queries Snowflake through the provider's SQLToolset. The judge is a second model call attached to that agent as a Pydantic AI capability. It reads the answer together with the SQL and rows behind it, then returns a structured verdict. An answer gets posted only when the judge passes it with a score of 0.8 or higher. If it doesn't, the agent gets one retry with the judge's feedback, and if that fails too (or the judge itself fails) the thread gets a short "blocked" notice.

What a judge is, and where it breaks

A judge is a model call that grades another model's output. You hand it the original request, the output and a rubric, and it hands back a verdict. Judges are often used offline, scoring thousands of saved outputs so you can compare two prompts or two models. Ours runs online, on every answer, before a person sees it.

The bet behind the pattern is that checking an answer is easier than writing one. The data agent has to pick tables, write joins, convert time zones and turn rows into prose. The judge has a much narrower question to answer: given these rows, does this answer match what the person asked?

There's research behind that bet, and behind its limits. In Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena (Zheng et al., 2023), strong judges agreed with human preferences over 80% of the time, which is about as often as humans agree with each other. The same paper documents where judges go wrong: position bias, verbosity bias, self-enhancement bias and limited reasoning ability.

Position bias doesn't apply to us, since our judge grades one answer at a time. The other two biases do. Our judge and data agent run on the same model today (Claude Sonnet 5.5), which makes a judge that goes easy on its own style of answer a real risk. Verbosity is a risk too. The top-10 table mentioned in the introduction was long and nicely formatted, and a judge that rewards fluent prose would have passed it.

A judge is just another model call. It can time out or fail, and a vague rubric like “Is this answer accurate and helpful?” can lead to vague, overly generous grades. The rest of this post explains how EduAgent handles these problems in code.

How the EduAgent's judge works

The data agent is an @task.agent task from the Common AI provider. It queries Snowflake through the provider's SQLToolset, and it has one small function tool of its own, starter_package_status, which reads a live roster from Skilljar (our learning platform). The judge has no tools at all. It runs inside the agent's own run, after the agent settles on an answer and before the task returns it. Here's the task, on Airflow 3.3 with Common AI 1.0.0 and Pydantic AI 2.55:

skilljar_tools = FunctionToolset(tools=[starter_package_status], id="skilljar")


@task.agent(
    llm_conn_id="pydanticai_default",
    system_prompt=DATA_AGENT_PROMPT,
    toolsets=[
        SQLToolset(
            db_conn_id="snowflake_default",
            allowed_tables=[
                "MODEL_EDU.SKILLJAR_COURSE_PROGRESS",
                "IN_SKILLJAR.STUDENT",
                # ...and the other Academy tables the agent may read
            ],
            max_rows=200,
            max_result_bytes=16_000,
        ),
        skilljar_tools,
    ],
    capabilities=[judge],
    agent_params={"retries": {"tools": 3, "output": JUDGE_RETRIES}},
    retries=1,
)
def answer(message: dict) -> str:
    return thread_prompt(message["text"], message.get("thread_context"), "Follow-up question")

The decorated function only builds the prompt: the person's message, with the earlier messages in the thread when it's a follow-up. The provider creates the agent from the pydanticai_default connection, runs it and pushes the result to XCom. agent_params goes straight through to Pydantic AI's Agent. Here it sets the one output retry the judge relies on, plus three tool retries, because SQLToolset hands every SQL error back to the agent to fix. The judge runs inside that run, which means whatever the task returns has already been judged: an approved answer or the blocked notice.

![][image1] The diagram leaves out one path. If the data agent's own run fails (say it keeps writing SQL the toolset rejects and runs out of tool retries), the task fails. Airflow retries it once, and after that nothing reaches the thread.

The verdict is a Pydantic model

You can't do much in code with a judge that replies in free text, so ours has to fill in this model:

class DataQueryJudgement(BaseModel):
    passed: bool = Field(
        description="True only when the response accurately satisfies the user request."
    )
    score: float = Field(ge=0.0, le=1.0)
    reason: str
    missing_requirements: list[str] = Field(default_factory=list)
    accuracy_concerns: list[str] = Field(default_factory=list)
    suggested_fix: str | None = None

The judge runs with this class as its output_type, and Pydantic AI validates every verdict. A score of 1.3 or a missing field never reaches our code. Only passed and score decide anything. The other four fields give a rejection to something the data agent can act on, and the retry is built on them.

Writing the rubric

Here's the judge's system prompt, trimmed a little:

You are a strict quality judge for Astronomer Academy data-query answers.
You do not answer the user's data question. You only judge the existing answer.

Pass only if the response:
1. Directly answers the user's actual request, including requested filters, dates, metrics, and grouping.
2. Preserves explicit requested output size and format. ...
3. Does not replace an explicit list request with a top-10 summary, notable examples, or narrative-only answer.
4. Does not make claims that are unsupported by the provided source evidence when source evidence is available.
5. Clearly states assumptions when the request is ambiguous.
6. Surfaces query/data errors instead of pretending the answer is complete.

Be especially strict with count/list requests such as "top N", "list N", "show N", and "all".

Rules 2 and 3 address the top-100 question from the opening. Rule 4 catches the sign-in answer, as long as the judge can see the query. Under the rules, four score bands give each number a meaning (0.90 to 1.00 is "no meaningful issues", 0.70 to 0.89 is "minor omissions only"). We also tell the judge not to answer the question. A judge that starts writing its own answer has produced a second answer, and nothing checks that one.

What the judge gets to see

An answer that says "14 employees signed in" looks fine by itself. It only looks wrong next to the SQL that produced it, so the judge gets that too. Pydantic AI keeps the run's full message history, and we pull every evidence tool call out of it, along with whatever came back:

EVIDENCE_TOOLS = {"query": "SQL", "starter_package_status": "Live Skilljar lookup"}


def evidence_from(messages, max_chars: int = 40_000) -> str:
    """Every evidence tool call the data agent made, and what it got back."""
    lines = []
    for message in messages:
        for part in message.parts:
            label = EVIDENCE_TOOLS.get(getattr(part, "tool_name", None))
            if label and isinstance(part, ToolCallPart):
                args = part.args_as_dict()
                lines.append(f"{label} call:\n{args.get('sql', args)}")
            elif label and isinstance(part, ToolReturnPart):
                lines.append(f"{label} result:\n{part.content}")
            elif label and isinstance(part, RetryPromptPart):
                lines.append(f"{label} error:\n{part.content}")
    text = "\n\n".join(lines)
    if len(text) > max_chars:
        # Keep the newest evidence: on a retry, that's the corrected query.
        log.warning("Judge evidence truncated: kept the last %d of %d chars", max_chars, len(text))
        text = "[Earlier evidence left out]\n\n" + text[-max_chars:]
    return text

query is the SQLToolset tool that runs a statement. The judge sees each query the agent ran and the rows it got back, which is enough to notice that a query with no sign-in column can't support a claim about sign-ins. A query that fails comes back from the toolset as a retry prompt (RetryPromptPart), and those go in too. Rule 6 can't flag an error the judge never saw.

If your agent runs a lot of queries, the 40,000-character cap keeps them from blowing up the judge's prompt. When the evidence runs over, the oldest part goes first, because on a retry the newest query is usually the corrected one. max_result_bytes=16_000 on the toolset keeps any single result well under the cap, so the judge sees the same rows the agent saw.

Every tool the judge should hear about has to be in EVIDENCE_TOOLS. If you add a data tool and forget to list it here, answers built on that tool reach the judge with nothing behind them. The repo's CLAUDE.md (which the coding agents working on the project read) has a line warning about it.

Plugging the judge into the agent

The judge plugs in through Pydantic AI's Hooks capability. You register a function on a point in the agent's lifecycle. after_output_process fires once the agent has a final answer, and raising ModelRetry from it sends the agent back to work with a message you choose. Here's the whole thing:

judge = Hooks()


@judge.on.after_output_process
async def judge_answer(ctx: RunContext, *, output_context, output: str) -> str:
    evidence = evidence_from(ctx.messages)
    try:
        judge_agent = PydanticAIHook(llm_conn_id="pydanticai_judge").create_agent(
            DataQueryJudgement, instructions=DATA_QUERY_JUDGE_PROMPT
        )
        verdict = (await judge_agent.run(judge_prompt(ctx.prompt, output, evidence))).output
    except Exception:
        log.exception("Data query judge failed; blocking the answer")
        return blocked_notice(None)

    log.info(
        "Judge on attempt %s: passed=%s score=%.2f reason=%s",
        ctx.retry + 1, verdict.passed, verdict.score, verdict.reason,
    )
    if verdict.passed and verdict.score >= MIN_PASSING_SCORE:  # 0.8
        return output
    if ctx.retry < JUDGE_RETRIES:  # 1
        raise ModelRetry(retry_feedback(verdict))
    return blocked_notice(verdict)

judge_prompt, retry_feedback and blocked_notice only format strings. The decisions are all in the lines above.

To pass, an answer needs passed=True and a score of at least 0.8. The rubric's 0.70 to 0.89 band means "minor omissions only", and a judge can reasonably call that a pass. We'd still hold that answer back. The threshold is a constant in code, which lets you move it without touching the prompt.

The except branch is the fail-closed rule. If the judge times out, or its output never validates, the hook returns the blocked notice. An answer nobody checked doesn't go out.

ModelRetry gives the data agent one more turn in the same conversation, with its earlier SQL and rows still in front of it. retry_feedback fills the message with the judge's reason, missing requirements, accuracy concerns and suggested fix. It also adds five rules of our own, like "Do not infer missing names, dates, organizations, or activity from email addresses or ordering." ctx.retry counts the retries so far, and after one the hook stops asking.

Then there's ctx.prompt. It holds the original prompt (the person's message, plus the thread for a follow-up) and stays the same through the retry, so both attempts get graded against what the person asked. We want it that way. The retry feedback is long and full of hints. Grade an agent against its own hints and it can pass without ever answering the question.

If the second attempt fails as well, the task returns the blocked notice, and that's what lands in the thread:

Data query blocked by QA check (score: 0.35)

The generated answer was withheld because it was not approved by the judge.
Check the Airflow task logs for the judge reason and suggested fix.

The verdict itself stays in the task log. It describes the rejected answer, which is the thing we decided not to post, and a unit test checks that none of it leaks into the Slack message.

Where Airflow fits

The Dag around the task is short. A custom SlackChannelTrigger polls the watched channels every 10 seconds. When a message mentions the bot, handle_slack_channel_messages has Claude Haiku 5.5 classify it, and data questions become events on the slack_data_query asset, which schedules this Dag:

@dag(schedule=slack_data_query_asset, tags=["edubot", "snowflake", "ai"])
def handle_data_query():

    @task
    def snowflake_questions(triggering_asset_events=None) -> list[dict]:
        return _messages_for_source(triggering_asset_events, "snowflake")
    # answer: the @task.agent task from earlier

    @task
    def reply_in_thread(message: dict, answer_text: str):
        post_message(message["channel"], markdown_to_slack(answer_text), thread_ts=message["ts"])

    @task_group
    def per_question(message: dict):
        reply_in_thread(message, answer(message))

    per_question.expand(message=snowflake_questions())

snowflake_questions skips the run when the batch has no Snowflake questions. Mapping the task group gives every question its own answer and reply_in_thread task instances, each with its own log. A reply waits only for its own answer, so a question that fails for good doesn't hold back the rest of the batch. (The repo's Dag also answers Zendesk questions, with no judge, and that branch is left out here.)

The Common AI provider also replaced code we used to maintain. Our old execute_sql tool checked for SELECT by hand, wrapped every query in a LIMIT and serialized calls behind a lock. SQLToolset is read-only unless you pass allow_writes=True, and it caps results with max_rows and max_result_bytes. With allowed_tables set, it parses every query with sqlglot (subqueries, CTEs and joins included) and rejects one that reaches a table outside the list before it runs. The provider's docs call that an application-level guardrail and recommend a least-privilege database role behind it.

That allowlist is stricter than you might expect. Entries have no database part, so a query that names HQ.IN_SKILLJAR.STUDENT gets rejected even with IN_SKILLJAR.STUDENT on the list. Our agent's prompt named its tables that way, which means moving to the toolset includes dropping the prefix and setting the database on the connection. A table in a second database can't go on the list at all. Inline SQL comments are rejected too, and so is any function sqlglot doesn't recognize, unless you name it in allowed_functions (DATE_TRUNC and CONVERT_TIMEZONE pass fine). The agent gets each rejection back as an error and can rewrite the query, which is one more reason for the extra tool retries.

Tool logging is on by default, and each tool call shows up in the task log by name, with how long it took. The SQL itself only appears at DEBUG level. Next to those lines is the judge's line for each attempt: Judge on attempt 1: passed=False score=0.35 reason=.... When someone asks why the bot refused a question, you'll find the reason in the log of that task instance.

The judge has its own connection, pydanticai_judge. Right now it points at the same model as the data agent. To move the judge to a different model, you edit that connection's model field and leave the code alone. That's probably the cheapest answer to the self-enhancement risk from earlier, though we haven't measured whether it helps.

The task can also take durable=True. The provider then journals each model request and tool call, and when Airflow retries the task, the steps that already finished are replayed from the journal. If a worker dies while the judge's retry is running, the next attempt doesn't pay for the data agent's first answer again. This only helps a task with retries (ours has one), and toolsets need a unique id, which is why skilljar_tools has one. On Airflow 3.3 the journal lives in Airflow's task state store, so turning it on takes no extra configuration. Older versions need [common.ai] durable_cache_path pointed at object storage. Durable execution with Apache Airflow covers how the journal works. We haven't turned it on yet.

A person judges the ticket replies

The same bot drafts replies to Academy support tickets, and those skip the AI judge entirely. Each draft lands in Slack with Approve, Regenerate and Skip buttons, and nothing reaches the customer until someone clicks Approve. The clicks come back through SQS into a handle_slack_events Dag.

We picked the kind of judge by asking who reads the output and how long they're willing to wait:

A customer reply can wait for a person. A teammate asking for last month's numbers wants the answer while they're still looking at the thread, and nobody on the team has time to review every query. So every data answer gets an AI review, and only ticket replies wait for a person. If you'd rather keep human review inside Airflow, the provider has that too: require_approval=True on @task.llm pauses the task after it generates its output, until someone approves or rejects it in Airflow's human-in-the-loop interface. Our ticket flow does its review in Slack.

What catches what

The judge is one check among several, and it costs a model call every time. Give it only the problems the cheaper checks can't catch:

The code around the model is easy to test. Pydantic AI's FunctionModel lets you script both sides. The data agent calls query, answers, and answers again after a retry, while a scripted judge passes, fails or raises on each attempt. Our cases cover a pass, a fail then a pass, two fails, a pass with a score of 0.75, a judge that times out and a query that errors before the agent fixes it. None of them needs an API key. They tell you nothing about whether the judge grades well, though.

What we haven't solved

We don't measure the judge yet. The 80% agreement figure in the research is for those judges on those tasks and says nothing about ours, and we have no labeled set of past questions with a human verdict to compare against. The task logs hold every verdict the judge has given, so at least the raw material exists.

A strict judge also rejects some good answers. We don't know how many, and when it happens the blocked notice is the only sign anyone sees. If you build one of these, count your blocks from the start.

The same-model problem is still open, too. Fixing it is now a connection edit, which makes it a cheap experiment, but we haven't run it.

Start with the answers that fooled you

Most of the judge's code decides what it gets to see and what happens when it says no. The model call itself is a few lines. If your agent already runs as an @task.agent task, adding a judge means one Hooks object and one more entry in capabilities.

Spend your time on the rubric. Compare "accurate, relevant and helpful" with "does not replace an explicit list request with a top-10 summary." A judge can fail an answer on the second one. Before you write yours, go back through your agent's history and collect a handful of answers that looked right and weren't. Turn each one into a rubric rule and a FunctionModel test case.