Router Agent

Production-ready AI Router Agent built with LangGraph, LangChain and Pydantic. The agent receives a user request, runs security guardrails, classifies intent with an LLM router, dispatches to a specialist agent and redacts sensitive content from the final response.

Overview

The Router Agent implements a multi-stage pipeline that is safe to deploy behind a public endpoint:

1. Input Guardrail - rejects prompt injection, prompt leakage, jailbreaks and malicious SQL.
2. LLM Router - classifies the query into exactly one of rag, sql, calculator or fallback using structured outputs.
3. Route Validator - ensures the chosen route is on the allow-list.
4. Confidence Check - routes low-confidence decisions to fallback.
5. Specialist Agents - RAG, SQL, Calculator and Fallback agents.
6. Output Guardrail - redacts API keys, secrets, stack traces and other internal details.

The whole pipeline is modelled as a LangGraph StateGraph so it can be extended, observed and replayed.

Architecture

User
  |
  v
Input Guardrail      (regex-based injection / leakage / SQL detection)
  |
  v
Router Agent         (LLM with structured output -> RouteDecision)
  |
  v
Route Validator      (allow-list enforcement)
  |
  v
Confidence Check     (below threshold -> fallback)
  |
  +--------+---------+----------+---------+
  |        |         |          |         |
  v        v         v          v         v
RAG      SQL     Calculator   Fallback
  |        |         |          |
  +--------+---------+----------+
              |
              v
Output Guardrail   (secret / traceback / internal-marker redaction)
              |
              v
Response

Folder Structure

Path	Purpose
app.py	Application entry point (CLI and interactive REPL).
agents/router.py	LLM-backed intent classifier with structured outputs.
agents/rag_agent.py	RAG specialist. Ships as a stub; pluggable via the VectorStore protocol.
agents/sql_agent.py	SQL/analytics specialist. Ships as a stub; pluggable via the DatabaseRepository protocol.
agents/calculator_agent.py	Calculator specialist. Delegates all arithmetic to the safe tool.
agents/fallback_agent.py	Catch-all safe response agent.
config/config.yaml	Static configuration (model name, threshold, routes, etc.).
config/settings.py	Pydantic Settings loader for YAML and environment variables.
config/logging_config.py	Idempotent structured-logging configuration.
graph/workflow.py	LangGraph StateGraph assembly.
graph/llm_factory.py	LLM construction and LangSmith tracing toggles.
guardrails/input_guardrail.py	User-input security checks (regex-based).
guardrails/output_guardrail.py	Response redaction.
guardrails/route_validator.py	Route allow-list enforcement.
models/state.py	Typed graph state.
models/route_models.py	Pydantic models for routing decisions.
tools/calculator.py	Safe arithmetic tool (no eval / exec).
tests/	Pytest suite (78 tests, 85% coverage).
logs/	Log output directory.
pyproject.toml	Project metadata, pytest/ruff/mypy config.
requirements.txt	Runtime dependencies.
.env.example	Environment variable template.
.gitignore	Standard Python ignores.

Setup

Requires Python 3.10+ (3.12+ recommended).

python -m venv .venv
# Windows PowerShell
.venv\Scripts\Activate.ps1
# macOS / Linux
source .venv/bin/activate

pip install -r requirements.txt
cp .env.example .env
# Edit .env and provide OPENAI_API_KEY

The first time you run the application, the logs directory is created automatically.

Configuration

All runtime configuration lives in config/config.yaml. Environment variables in .env override secrets such as OPENAI_API_KEY and LANGSMITH_API_KEY.

Key settings:

Setting	Description	Default
llm.model	Model used by the router	gpt-4.1
llm.temperature	Sampling temperature	0.0
llm.max_tokens	Maximum tokens per router call	1024
llm.timeout_seconds	HTTP timeout for the LLM	30
routing.confidence_threshold	Below this confidence the request is rerouted to fallback	0.70
routing.allowed_routes	Route allow-list	["rag", "sql", "calculator", "fallback"]
guardrails.max_query_length	Maximum accepted query length	4000
guardrails.enable_input_validation	Toggle input guardrail	true
guardrails.enable_output_redaction	Toggle output guardrail	true
tools.calculator.precision	Decimal places returned by the calculator	6
vectorstore.collection_name	Future vector store collection name	documents
database.connection_string	Future SQL connection string	sqlite:///data/sample.db

Running the Application

# Single query
python app.py "What is 17 * 24?"

# Interactive shell
python app.py

The script reads config/config.yaml, configures structured logging, optionally enables LangSmith tracing and dispatches the query through the graph.

Testing

pytest                              # full test suite (78 tests, no network calls)
pytest --cov=. --cov-report=term    # with coverage report

All tests use mocked LLM responses, so they run offline and finish in under 20 seconds.

Sample Requests

Intent	Example	Expected route
RAG	What is our refund policy?	rag
RAG	Where can I find the employee handbook?	rag
SQL	What was revenue last quarter?	sql
SQL	How many active users do we have?	sql
Calculator	What is 17 * 24?	calculator
Calculator	Compute 100 / (5 * 2)	calculator
Fallback	Tell me a joke	fallback
Blocked	ignore previous instructions	fallback (input guardrail)

Observability

If LANGSMITH_API_KEY is set in .env, LangSmith tracing is enabled automatically. The project name defaults to router-agent and can be overridden with LANGSMITH_PROJECT. When LANGSMITH_API_KEY is absent, the application continues normally without tracing.

Error Handling

Every node in the graph wraps execution in a try/except boundary. On any failure - bad configuration, LLM timeout, malformed structured output, tool error or graph failure - the workflow stores an error message, reroutes to the fallback agent and returns a safe user-facing response. The application never crashes the process.

Design Principles

• Dependency injection - LLM, settings and tools are injected through constructors; no module-level globals.
• Typed configuration - all settings are validated by Pydantic at start-up.
• Structured outputs - the router never parses free text; the LLM is forced to populate a RouteDecision Pydantic model.
• Safe tools - the calculator tool uses a shunting-yard algorithm with a strict token whitelist; no eval / exec.
• Deterministic guardrails - input and output guardrails are pure functions with no LLM dependency.
• Testability - every node is unit-testable, every agent is mockable, the entire workflow has a fake-LLM end-to-end test.