Answer Agent with Query Classification

Overview

An answer agent with query classification combines the intelligence of intent detection with the depth of retrieval-augmented generation (RAG) to produce contextually appropriate, high-quality responses. Rather than treating every user input identically, this architecture first determines whether a query is a factual question requiring a reasoned answer, a search request needing document discovery, or a conversational prompt requiring a more open-ended response. This classification step is the cornerstone of the system — it routes queries down radically different processing paths, ensuring the right tools are applied to each request type.

Building an effective answer agent involves orchestrating several moving parts: a lightweight classifier (often a fine-tuned small model or a zero-shot LLM prompt), a retrieval layer backed by a vector store or keyword index, a reasoning chain that synthesizes retrieved documents into a coherent response, and a streaming output layer that delivers tokens incrementally via Server-Sent Events (SSE). Each component must hand off cleanly to the next while tracking metadata such as source documents, confidence scores, and context window utilization. In production systems, these components are wired together into a pipeline that handles hundreds of concurrent queries with low latency.

What sets advanced answer agents apart is their ability to communicate uncertainty and guide users toward better queries. Confidence scores surface when the retrieved context is thin or contradictory, prompting the agent to hedge its response appropriately. Follow-up questions are generated alongside the primary answer to deepen the conversation, improve recall in subsequent turns, and surface related topics the user may not have considered. This combination of streaming delivery, transparent confidence signaling, and proactive follow-up generation creates a user experience that feels intelligent and trustworthy rather than opaque and brittle.

Core Concepts

How It Works

Query Classification Pipeline

When a user submits a query, the first stage is rapid intent classification. The raw query string is passed to a classifier that returns one of several intent labels: factual_question, search, comparison, procedural, or conversational. In practice, this can be implemented as a zero-shot LLM call with a structured JSON output schema, which avoids the need to maintain a separate classification model while still achieving sub-100ms classification latency thanks to a small, fast model (e.g., GPT-4o-mini or Claude Haiku). The intent label is attached to the query object and flows through the rest of the pipeline as routing metadata.

For search intents, the pipeline skips multi-step reasoning and goes directly to retrieval, returning a ranked list of document excerpts with snippets and scores. For factual_question and procedural intents, the pipeline enters the full reasoning chain: query decomposition, per-step retrieval, evidence synthesis, and streaming answer generation. For conversational intents with no strong factual component, the pipeline falls back to a direct LLM completion without retrieval, keeping latency low for casual interactions.

Retrieval-Augmented Generation

Once the intent is known, the retrieval stage embeds the query using a text-embedding model (e.g., text-embedding-3-small) and queries the vector store for the top-k most similar document chunks. The chunks are then re-ranked using a cross-encoder that scores each (query, chunk) pair more precisely than the approximate nearest-neighbor lookup. The top-ranked chunks are assembled into the context window, with each chunk tagged with its source document ID, chunk index, and similarity score. If the total token count exceeds the context budget, lower-ranked chunks are truncated first. The assembled context is then passed to the LLM along with the original query and a system prompt that instructs the model to answer only from the provided context and to cite sources explicitly.

Multi-hop questions require a second retrieval pass: after generating an intermediate answer for the first sub-question, that answer is re-embedded and used to retrieve additional chunks that may not have been surface by the original query embedding. This iterative retrieval loop runs for a configurable number of steps (typically 2-3) before the final synthesis prompt is assembled.

Streaming Answer Assembly

The final LLM prompt includes the assembled context, the reasoning chain so far, and instructions to generate the answer in a structured format: answer text, confidence score (0-1), cited source IDs, and 2-3 follow-up questions. The LLM streams its response token by token. The server-side SSE handler parses the stream in real time, detecting when it enters structured fields (confidence, citations, follow-ups) by watching for known JSON delimiters, and emits typed SSE events accordingly. The React client subscribes to the event stream and maintains a state object that accumulates the answer text in one field while updating confidence and follow-up question fields as their corresponding SSE events arrive. This gives the UI the ability to show a confidence indicator and follow-up chips the moment they are available, rather than waiting for the full response to complete.

Implementation Patterns

Query Classification and Routing

The classifier is implemented as a typed async function that calls the LLM with a structured output schema. The result is used to route the query to the appropriate handler.

// types.ts
export type QueryIntent =
  | 'factual_question'
  | 'search'
  | 'comparison'
  | 'procedural'
  | 'conversational';

export interface ClassifiedQuery {
  raw: string;
  intent: QueryIntent;
  confidence: number;
  decomposedSteps?: string[];
}

export interface RetrievedChunk {
  id: string;
  text: string;
  sourceUrl: string;
  sourceTitle: string;
  similarityScore: number;
}

export interface AnswerEvent {
  type: 'token' | 'citation' | 'confidence' | 'followup' | 'done' | 'error';
  payload: string | number | string[] | null;
}

// classifier.ts
import OpenAI from 'openai';

const client = new OpenAI();

export async function classifyQuery(query: string): Promise<ClassifiedQuery> {
  const response = await client.chat.completions.create({
    model: 'gpt-4o-mini',
    messages: [
      {
        role: 'system',
        content: `You are a query intent classifier. Analyze the user query and return JSON.
Intent types:
- factual_question: asking for a specific fact or explanation
- search: browsing/exploring for documents or examples
- comparison: comparing two or more things
- procedural: asking how to do something step-by-step
- conversational: casual chat with no strong factual need

Return JSON: { intent, confidence (0-1), decomposedSteps (array of sub-questions if complex, else null) }`
      },
      { role: 'user', content: query }
    ],
    response_format: { type: 'json_object' },
    temperature: 0,
  });

  const parsed = JSON.parse(response.choices[0].message.content ?? '{}');
  return {
    raw: query,
    intent: parsed.intent as QueryIntent,
    confidence: parsed.confidence ?? 0.5,
    decomposedSteps: parsed.decomposedSteps ?? undefined,
  };
}

// router.ts
import { classifyQuery, ClassifiedQuery } from './classifier';
import { runSearchPipeline } from './search';
import { runAnswerPipeline } from './answer';
import type { Response } from 'express';

export async function routeQuery(query: string, res: Response): Promise<void> {
  const classified = await classifyQuery(query);

  // Set SSE headers
  res.setHeader('Content-Type', 'text/event-stream');
  res.setHeader('Cache-Control', 'no-cache');
  res.setHeader('Connection', 'keep-alive');

  // Emit the classification result immediately so the UI can show intent
  res.write(`data: ${JSON.stringify({ type: 'classification', payload: classified })}\n\n`);

  switch (classified.intent) {
    case 'search':
      await runSearchPipeline(classified, res);
      break;
    case 'factual_question':
    case 'comparison':
    case 'procedural':
      await runAnswerPipeline(classified, res);
      break;
    case 'conversational':
      await runDirectCompletion(classified.raw, res);
      break;
    default:
      await runAnswerPipeline(classified, res);
  }

  res.write(`data: ${JSON.stringify({ type: 'done', payload: null })}\n\n`);
  res.end();
}

async function runDirectCompletion(query: string, res: Response): Promise<void> {
  const stream = await client.chat.completions.create({
    model: 'gpt-4o',
    messages: [{ role: 'user', content: query }],
    stream: true,
  });

  for await (const chunk of stream) {
    const token = chunk.choices[0]?.delta?.content ?? '';
    if (token) {
      res.write(`data: ${JSON.stringify({ type: 'token', payload: token })}\n\n`);
    }
  }
}

Streaming Answer Generation with Confidence Scoring

The answer pipeline retrieves context, runs multi-step reasoning, then streams the final answer while emitting structured events for citations, confidence, and follow-up questions.

// answer.ts
import OpenAI from 'openai';
import { embed, retrieveChunks } from './retrieval';
import type { ClassifiedQuery, RetrievedChunk, AnswerEvent } from './types';
import type { Response } from 'express';

const client = new OpenAI();
const MAX_CONTEXT_TOKENS = 6000;

function buildContextBlock(chunks: RetrievedChunk[]): string {
  let block = '';
  let tokenEstimate = 0;
  for (const chunk of chunks) {
    const entry = `[SOURCE:${chunk.id}] ${chunk.sourceTitle}\n${chunk.text}\n\n`;
    const entryTokens = Math.ceil(entry.length / 4); // rough char-to-token ratio
    if (tokenEstimate + entryTokens > MAX_CONTEXT_TOKENS) break;
    block += entry;
    tokenEstimate += entryTokens;
  }
  return block.trim();
}

async function retrieveForStep(stepQuery: string): Promise<RetrievedChunk[]> {
  const embedding = await embed(stepQuery);
  const chunks = await retrieveChunks(embedding, { topK: 8 });
  return chunks;
}

function emitEvent(res: Response, event: AnswerEvent): void {
  res.write(`data: ${JSON.stringify(event)}\n\n`);
}

export async function runAnswerPipeline(
  classified: ClassifiedQuery,
  res: Response
): Promise<void> {
  // Multi-step retrieval: retrieve for each decomposed step, then the original query
  const steps = classified.decomposedSteps ?? [classified.raw];
  const allChunks: RetrievedChunk[] = [];
  const seenIds = new Set<string>();

  for (const step of steps) {
    const chunks = await retrieveForStep(step);
    for (const chunk of chunks) {
      if (!seenIds.has(chunk.id)) {
        seenIds.add(chunk.id);
        allChunks.push(chunk);
      }
    }
  }

  // Sort by similarity score descending, then build context
  allChunks.sort((a, b) => b.similarityScore - a.similarityScore);
  const contextBlock = buildContextBlock(allChunks);

  // Compute retrieval-based confidence from top-k scores
  const topScores = allChunks.slice(0, 5).map(c => c.similarityScore);
  const avgTopScore = topScores.length
    ? topScores.reduce((s, v) => s + v, 0) / topScores.length
    : 0;

  const systemPrompt = `You are a precise answer agent. Answer the user's question using ONLY the provided source documents.

Rules:
1. Cite sources inline as [SOURCE:id] after each claim.
2. After your answer, output a JSON block on its own line:
   {"confidence": 0.0-1.0, "citedIds": ["id1","id2"], "followUps": ["q1","q2","q3"]}
3. Confidence should reflect how well the sources support your answer (0=no support, 1=fully supported).
4. If sources are insufficient, state this explicitly and set confidence < 0.4.

Sources:
${contextBlock}`;

  const stream = await client.chat.completions.create({
    model: 'gpt-4o',
    messages: [
      { role: 'system', content: systemPrompt },
      { role: 'user', content: classified.raw },
    ],
    stream: true,
    temperature: 0.2,
  });

  let buffer = '';
  let inJsonBlock = false;
  let jsonBuffer = '';

  for await (const chunk of stream) {
    const token = chunk.choices[0]?.delta?.content ?? '';
    if (!token) continue;

    buffer += token;

    // Detect transition into the trailing JSON block
    if (!inJsonBlock && buffer.includes('\n{')) {
      inJsonBlock = true;
      // Emit everything up to the JSON block as tokens
      const splitIdx = buffer.lastIndexOf('\n{');
      const textPart = buffer.slice(0, splitIdx);
      if (textPart) {
        emitEvent(res, { type: 'token', payload: textPart });
      }
      jsonBuffer = buffer.slice(splitIdx + 1); // start of JSON
      continue;
    }

    if (inJsonBlock) {
      jsonBuffer += token;
      // Check if JSON block is complete
      if (jsonBuffer.trimEnd().endsWith('}')) {
        try {
          const meta = JSON.parse(jsonBuffer.trim());
          // Blend retrieval confidence with model self-assessed confidence
          const blendedConfidence = 0.5 * (meta.confidence ?? 0.5) + 0.5 * avgTopScore;
          emitEvent(res, { type: 'confidence', payload: Math.round(blendedConfidence * 100) / 100 });
          if (Array.isArray(meta.citedIds)) {
            const citations = allChunks
              .filter(c => meta.citedIds.includes(c.id))
              .map(c => `[${c.sourceTitle}](${c.sourceUrl})`);
            emitEvent(res, { type: 'citation', payload: citations.join('\n') });
          }
          if (Array.isArray(meta.followUps)) {
            emitEvent(res, { type: 'followup', payload: meta.followUps });
          }
        } catch {
          // JSON parse failed — emit as raw token so nothing is lost
          emitEvent(res, { type: 'token', payload: jsonBuffer });
        }
        inJsonBlock = false;
        jsonBuffer = '';
        buffer = '';
      }
    } else {
      emitEvent(res, { type: 'token', payload: token });
    }
  }

  // Flush any remaining buffer
  if (buffer && !inJsonBlock) {
    emitEvent(res, { type: 'token', payload: buffer });
  }
}

// React client hook (useAnswerStream.ts)
import { useState, useCallback } from 'react';

interface AnswerState {
  text: string;
  confidence: number | null;
  citations: string[];
  followUps: string[];
  intent: string | null;
  streaming: boolean;
  error: string | null;
}

export function useAnswerStream() {
  const [state, setState] = useState<AnswerState>({
    text: '', confidence: null, citations: [], followUps: [],
    intent: null, streaming: false, error: null,
  });

  const ask = useCallback(async (query: string) => {
    setState({ text: '', confidence: null, citations: [], followUps: [],
               intent: null, streaming: true, error: null });

    const response = await fetch('/api/answer', {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({ query }),
    });

    if (!response.body) return;
    const reader = response.body.getReader();
    const decoder = new TextDecoder();

    while (true) {
      const { done, value } = await reader.read();
      if (done) break;

      const lines = decoder.decode(value).split('\n');
      for (const line of lines) {
        if (!line.startsWith('data: ')) continue;
        try {
          const event: AnswerEvent = JSON.parse(line.slice(6));
          setState(prev => {
            switch (event.type) {
              case 'token':
                return { ...prev, text: prev.text + (event.payload as string) };
              case 'confidence':
                return { ...prev, confidence: event.payload as number };
              case 'citation':
                return { ...prev, citations: (event.payload as string).split('\n') };
              case 'followup':
                return { ...prev, followUps: event.payload as string[] };
              case 'classification':
                return { ...prev, intent: (event.payload as any).intent };
              case 'done':
                return { ...prev, streaming: false };
              case 'error':
                return { ...prev, streaming: false, error: event.payload as string };
              default:
                return prev;
            }
          });
        } catch { /* skip malformed events */ }
      }
    }
  }, []);

  return { state, ask };
}

Practical Applications

In a production deployment, the answer agent powers the intelligent search and Q&A interface that sits at the core of the product. When a user types a query into the search bar, the classification step determines in under 100ms whether the user wants a direct answer (e.g., “What is the retention policy for Tier-1 accounts?”) or a browsable list of documents (e.g., “retention policy templates”). This routing decision shapes the entire UI response: question intents render a streaming prose answer with a confidence badge and follow-up chips, while search intents render a document grid with relevance scores and snippet highlights.

The multi-step reasoning chain is particularly valuable for the comparison intent class. When a user asks “How does the enterprise plan differ from the pro plan in terms of API limits?”, the agent decomposes this into three sub-questions (enterprise API limits, pro API limits, and the comparison itself), retrieves separate context chunks for each, and synthesizes a structured comparison table. This prevents the common failure mode where a single retrieval query surfaces only one side of a comparison.

Confidence scores are surfaced in the UI as a color-coded indicator (green above 0.7, yellow 0.4-0.7, red below 0.4). When confidence is low, the UI automatically expands the follow-up questions section and displays a prompt encouraging the user to refine their query. This feedback loop has measurably reduced dead-end searches by guiding users toward queries that yield higher-confidence answers, improving both user satisfaction and the quality of data in the query analytics pipeline.

The SSE streaming architecture also enables a key UX feature: showing source citations as they are identified, before the full answer is complete. Users see the source documents populating in a side panel as the agent generates its response, giving them the ability to open a primary source immediately rather than waiting for the full answer to render. This is implemented by parsing citation SSE events in real time and rendering document cards that link directly to the source material.

Common Pitfalls

• Over-relying on a single intent label: Real-world queries are often ambiguous. A query like “show me examples of RAG implementations” could be classified as either search or factual. Build confidence thresholds into your router so that low-confidence classifications trigger a hybrid response (document list plus a brief synthesized summary) rather than a hard either/or branch.
• Ignoring token budget in context assembly: Naively concatenating all retrieved chunks can overflow the LLM’s context window or push important content past the attention cliff. Always enforce a hard token budget, prefer re-ranking to top-k cutoffs, and log context utilization metrics so you can tune the budget over time.
• Streaming partial JSON events: When parsing SSE streams for structured data (confidence, citations), token chunking will often split JSON across multiple events. Always buffer incoming tokens until you detect a complete JSON object rather than attempting to parse each token individually — otherwise you will see frequent JSON parse errors in production.
• Miscalibrated confidence scores: If the LLM self-assessed confidence is not blended with retrieval-quality metrics, it will be systematically overconfident on poorly-retrieved contexts (the model “makes up” a confident answer from sparse context). Always blend model confidence with retrieval scores and validate the calibration against a labeled test set.
• Not handling SSE connection drops gracefully: Long-running streaming responses are vulnerable to client disconnects, proxy timeouts, and network interruptions. Implement a resume mechanism (cursor-based) that allows the client to reconnect and request only the tokens it has not yet received, and emit keepalive comments (: ping\n\n) every 15-20 seconds to prevent proxy timeouts.
• Generating unhelpful follow-up questions: Follow-up questions that are too similar to the original query, too vague, or unrelated to the answer content erode user trust. Include the generated answer text in the follow-up generation prompt (not just the original query) so that follow-ups reflect the specific gaps or tangents in the answer, and filter out near-duplicate follow-ups with a simple embedding similarity check.