Lesson 6: Building the QueryUnderstandingAgent

Today we build:

• A weighted, multi-category classifier replacing Day 1’s single-keyword stub

• Entity extraction that pulls structured facts out of a raw question

• An explicit confidence threshold that routes uncertain queries away from retrieval instead of guessing

Why This Matters

Day 1’s QueryUnderstandingStage picked the first keyword substring it found and called it done — no confidence, no way to say “I’m not sure.” That’s fine for a proof of concept; it’s a liability in production, where a system that confidently misclassifies a query has already made its first mistake before retrieval even runs. Today’s agent is the first component in the pipeline built to know the difference between “confident” and “guessing.”

Core Concept

Instead of a first-match rule, the agent scores every category against the query using weighted keywords, then normalizes those scores into a probability-like distribution that sums to 1.0. The category with the highest share wins — but only if that share clears a confidence threshold. Below the threshold, the agent doesn’t guess; it returns out_of_scope and the query never reaches retrieval at all.

This matters because ambiguity is real, not a bug to eliminate. Run today’s code on “Can I get a refund if I cancel my annual plan?” and the agent produces a genuine tie: refund_policy and cancellation score exactly 0.5 each. A more nuanced case — “What is the cost to cancel and get a refund on my purchase?” — splits three ways: 0.44 / 0.33 / 0.22. The confidence score isn’t a made-up number; it’s a direct, honest readout of how much the query actually leans toward one category over the others.

Entity extraction runs alongside classification, pulling out details like plan type (”annual,” “monthly,” “lifetime”) whenever they’re mentioned — structured facts the retrieval stage can use later, instead of re-parsing the raw text.

Architecture

QueryUnderstandingAgent.process() is the pipeline’s first real seam: raw text goes in, a structured Intent object comes out, carrying the winning category, its confidence, the full score breakdown across every category, extracted entities, and an explicit in_scope boolean. Everything downstream — retrieval, synthesis, the critic — can trust that boolean instead of re-deriving it.

Implementation

def process(self, query_text: str) -> Intent:
    category_scores = self._score_categories(query_text.lower())
    best_category = max(category_scores, key=category_scores.get)
    best_confidence = category_scores[best_category]

    if best_confidence < CONFIDENCE_THRESHOLD:
        return Intent(query_text=query_text, intent_label="out_of_scope",
                       confidence=best_confidence, in_scope=False, ...)

    return Intent(query_text=query_text, intent_label=best_category,
                   confidence=best_confidence, in_scope=True, ...)

The threshold (0.34) isn’t tuned against real data yet — there’s no eval harness to tune it against until Phase 1. Today it’s a reasonable default, made visible and adjustable in one place rather than buried in scattered if-statements.

Build, Test, Verify

chmod +x start.sh stop.sh
./start.sh

start.sh installs dependencies, builds the Docker image if available, runs the lesson, runs the tests, and generates dashboard.html. Five test queries render as cards — three clean classifications, one genuine tie, and one out-of-scope weather question correctly routed away from retrieval. Tests: 6 passed.

./stop.sh

Implementation Guide

GitHub Link

https://github.com/sysdr/ai-reliability-engineering/tree/main/lesson_06_package/lesson_06_package

Package contents

lesson_06_package/
├── lesson_code.py    # QueryUnderstandingAgent, Intent, generate_dashboard
├── test_lesson.py    # 6 tests verifying classification, entities, and scope routing
├── requirements.txt  # pytest only
├── Dockerfile         # minimal python:3.11-slim image
├── start.sh           # install deps, build, run, test, verify — one command
├── stop.sh            # cleanup — removes .venv, dashboard.html, caches, Docker image
└── README.md          # quick-start reference

Run it

cd lesson_06_package
chmod +x start.sh stop.sh
./start.sh

You should see dashboard.html generated, and 6 tests passed. Open the dashboard — five query cards, each with category confidence bars, extracted plan-type entities where relevant, and a scope badge.

Cleanup: ./stop.sh — also removes the generated dashboard.html.

Architecture, in more detail

_score_categories() sums keyword weights per category, then divides each category’s raw score by the total across all categories — this normalization is what makes category_scores sum to 1.0 and behave like a confidence distribution rather than an arbitrary number.

_extract_entities() runs independently of classification — a query can be out_of_scope and still have an entity extracted, or in_scope with no entities found at all. The two operations don’t depend on each other.

process() combines both, then makes the single most consequential decision in the agent: compare the winning category’s confidence against CONFIDENCE_THRESHOLD (0.34) and decide whether the query proceeds to retrieval or gets labeled out_of_scope and stops here.

Why the confidence threshold matters more than the classification itself

Anyone can write a classifier that always returns its best guess. The harder, more valuable behavior is knowing when not to guess. With only three categories, a genuinely ambiguous query can score as low as 0.33 per category in a perfect three-way split — which is why the threshold sits at 0.34, just above an even three-way tie. A query that’s genuinely torn between all three categories gets routed to out_of_scope rather than picking one arbitrarily; a query that clearly leans toward one category, even without total certainty, still gets through.

This threshold is a placeholder default, not a tuned value — there’s no eval harness yet (that starts in Phase 1) to measure whether 0.34 is actually the right cutoff for real queries. Making it a single named constant, rather than scattering the logic across the codebase, is what makes it possible to tune later without a rewrite.

Test coverage

Working demo link:

Troubleshooting

• ModuleNotFoundError: No module named 'lesson_code' — run tests from inside lesson_06_package.

• A query you try yourself classifies unexpectedly — check CATEGORY_KEYWORDS directly; this is a small, hand-written keyword list, not a trained model, so it only recognizes what’s explicitly listed.

• dashboard.html doesn’t open automatically — start.sh‘s auto-open is best-effort and silently skipped in headless environments; open the file manually.

Real-World Connection

Systems that skip explicit out-of-scope handling tend to fail the same way: a user asks something the system was never built to answer, and instead of saying so, it confidently retrieves the closest-sounding passage and answers anyway. That’s a harder failure to catch than an outright error, because nothing in the response signals anything went wrong. An explicit confidence floor is one of the cheapest reliability wins available, and it’s why it comes this early in the pipeline.

Next Steps

Day 7 builds the SynthesisAgent, which takes the retrieval results this agent’s in_scope decision permits and turns them into a drafted answer.