Building a Contextual Bandit AI Agent in TypeScript

Most AI agent workflows use static rules: “use the balanced profile” or “always review before publishing.” That works fine until the tradeoffs shift — when latency matters more than cost, or when quality suddenly justifies a deeper research pass. A contextual bandit replaces those static rules with a learned routing policy that adapts to your objective.

This guide builds on the multi-agent article workflow from the LangGraph Database Agent TypeScript example, replacing the fixed agent profile with a contextual bandit that learns which path through the workflow delivers the best results for your current goal — whether that’s budget, speed, quality, or a balanced tradeoff.

1. What Is a Contextual Bandit?

A contextual bandit is a reinforcement learning algorithm that learns to pick the best action from a set of discrete choices. Unlike A/B testing (which treats every choice equally forever), a bandit actively explores unknown options while exploiting known good ones. Unlike full reinforcement learning (which models sequences of decisions across time), a bandit treats each decision independently.

In plain terms: the bandit has a set of “arms” it can pull — think slot machines in a casino. Each arm has an unknown reward distribution. The bandit’s job is to maximize total reward by trying arms it hasn’t tried much (exploration) while favoring arms that have paid off well (exploitation). Add “context” — information about the current situation, like remaining budget or quality target — and you have a contextual bandit that adapts its choices to the circumstance.

In this project, the arms are edges in a workflow graph: “start → research” vs “start → research-lite”, “writer → review” vs “writer → publish”, and so on. Each edge carries a utility score based on quality gain, cost, latency, retry risk, and success probability — weighted by your chosen objective.

2. Project Overview

The repository is on GitHub: github.com/juustesout/contextual-bandit-ai-agent-typescript

The workflow produces blog articles through five stages — research, writing, review, publishing — but instead of picking one static profile (budget/balanced/quality/speed) at the start, the contextual bandit re-evaluates at every decision point. It chooses the next node based on historical performance stored in PostgreSQL, using either the epsilon-greedy or UCB1 strategy.

3. Project Setup

git clone https://github.com/juustesout/contextual-bandit-ai-agent-typescript.git tsbandit
cd tsbandit
npm install

The dependencies are minimal: pg for PostgreSQL, typescript and @types/node for compilation. No LangChain, no LangGraph, no Zod — just plain TypeScript and a raw fetch-based MiniAgent class.

4. Environment Variables

Create a .env file in the project root:

OPENAI_API_KEY=sk-...
OPENAI_MODEL=gpt-4o-mini
DATAFORSEO_API_KEY=base64-of-login:password
DATAFORSEO_BASE_URL=https://api.dataforseo.com
WP_BASE_URL=https://yourblog.com
WP_USERNAME=admin
WP_APPLICATION_PASSWORD=xxxx
JINA_API_KEY=jina_...
DESKTOP_DB_HOST=host.docker.internal
DESKTOP_DB_PORT=5433
DESKTOP_DB_NAME=xxxx
DESKTOP_DB_USER=xxxx
DESKTOP_DB_PASSWORD=...

5. The Workflow Graph

The contextual bandit models the workflow as a directed graph. Each node is an execution step, and each edge is a possible transition. The bandit chooses at every branching point:

start → [research, research-lite]
research → [writer, fallback/retry]
research-lite → [writer]
writer → [review]
review → [publish, fallback/retry]
fallback/retry → [writer]

The graph is defined in src/bandit/types.ts. The entry decision (start) is a virtual node so the first edge is scored exactly like every other edge.

6. The Utility Function

Every edge is scored with a weighted utility function:

utility = δ × Qₙₒᵣₘ  −  α × Cₙₒᵣₘ  −  β × Lₙₒᵣₘ  −  γ × retryRisk  +  δ × successProb

Each raw value is normalized to a 0–1 scale so the weights are objective-relative, not dominated by absolute magnitudes:

• Qₙₒᵣₘ = qualityGain / 100 (review score, 0–100)
• Cₙₒᵣₘ = costUsd / 0.05 (max ~$0.05 per edge)
• Lₙₒᵣₘ = latencyMs / 120,000 (120s timeout)

The weights per objective live in src/bandit/utility.ts:

const DEFAULT_UTILITY_WEIGHTS = {
  budget:   { alpha: 0.5, beta: 0.3, gamma: 0.2, delta: 0.6 },
  balanced: { alpha: 0.2, beta: 0.15, gamma: 0.2, delta: 0.8 },
  quality:  { alpha: 0.05, beta: 0.05, gamma: 0.2, delta: 1.0 },
  speed:    { alpha: 0.1, beta: 0.4, gamma: 0.15, delta: 0.6 },
};

Quality mode pushes δ (delta) to 1.0 and cuts α/β to near zero — quality dominates the decision. Budget mode raises α to penalize cost. Speed mode raises β to penalize latency.

7. Exploration Strategies

The contextual bandit supports two strategies, both in src/bandit/strategies.ts:

Epsilon-Greedy

With probability ε, pick a random arm. Otherwise, pick the arm with the highest mean reward. Simple, effective, and the ε parameter tunes the exploration/exploitation tradeoff directly. Default ε = 0.1.

UCB1 (Upper Confidence Bound)

Pick the arm that maximizes meanReward + √(2 × ln(totalPulls) / armPulls). This automatically balances exploration and exploitation: arms with few pulls get a confidence bonus, while well-tested arms are chosen by their proven mean. No separate ε parameter needed.

Both strategies enforce a cold-start phase: arms with fewer than minSamples observations are explored uniformly before exploitation kicks in. This ensures the bandit has baseline data for every option before it starts optimizing.

8. Learning from Quality: Backpropagation

An earlier version of the project had a fundamental flaw: the quality gain per edge was calculated as a hardcoded delta between node scores (e.g., writer.score − research.score = 82 − 80 = 2). This was constant and identical for both research paths, so the bandit learned nothing about the actual article quality and defaulted to the cheapest route.

The fix: after the reviewer scores the final article with a quality_score (0–100) based on depth, source variety, and factual substantiation, that score is backpropagated to every edge that contributed to the run. Every observation in the run gets the same quality credit — the bandit learns that a run that produced a high-quality article (usually from deeper research) rewards all its edges equally. This is implemented in the observe() function and the post-review backpropagation step in src/workflow.ts.

9. PostgreSQL Persistence

The bandit stores per-edge statistics in a PostgreSQL table called edge_stats:

CREATE TABLE edge_stats (
  source_node      TEXT NOT NULL,
  target_node      TEXT NOT NULL,
  objective        TEXT NOT NULL,
  times_selected   INTEGER NOT NULL DEFAULT 0,
  cumulative_reward DOUBLE PRECISION NOT NULL DEFAULT 0,
  avg_reward       DOUBLE PRECISION NOT NULL DEFAULT 0,
  sum_quality_gain DOUBLE PRECISION NOT NULL DEFAULT 0,
  sum_cost_usd     DOUBLE PRECISION NOT NULL DEFAULT 0,
  sum_latency_ms   DOUBLE PRECISION NOT NULL DEFAULT 0,
  success_count    INTEGER NOT NULL DEFAULT 0,
  retry_count      INTEGER NOT NULL DEFAULT 0,
  rejection_count  INTEGER NOT NULL DEFAULT 0,
  last_updated     TIMESTAMPTZ NOT NULL DEFAULT NOW(),
  PRIMARY KEY (source_node, target_node, objective)
);

Each run upserts a row per traversed edge, including failed and retried ones. The aggregateEdgeStatsFromRows() function in src/database.ts turns these raw rows into the per-edge statistics the bandit consumes for cold-starting the next run.

10. Running the Workflow

Profile mode (static agent profile — no bandit):

npm run build && node dist/demo.js "AI sales agents for SMBs"

Bandit mode (learned routing):

npm run build && node dist/demo.js "AI sales agents" --bandit --objective quality --strategy ucb1

11. Interpreting the Output

A bandit run prints an edge-level decision log:

[bandit] start -> research-lite | EXPLORE | expectedUtility=0.0000 | Cold start: all 2 arm(s) below 5 samples
[bandit] writer -> review | EXPLORE | expectedUtility=0.0000 | Cold start: all 1 arm(s) below 5 samples
[bandit] review -> publish | EXPLORE | expectedUtility=0.0000 | Cold start: all 2 arm(s) below 5 samples

The routing object shows the full route, per-decision details, and the run reward:

"routing": {
  "mode": "bandit",
  "objective": "quality",
  "strategy": "ucb1",
  "route": ["research-lite", "writer", "review"],
  "runReward": 2.07,
  "edgeRewards": {
    "start->research-lite": 27.78,
    "research-lite->writer": 56.75,
    "writer->review": 16.56
  }
}

12. Training the Bandit

The included train-bandit.js script runs the workflow N times and generates an HTML report with Chart.js visualizations:

node train-bandit.js 15 "AI sales agents" quality

This runs 15 workflows, collects per-edge statistics after each run, and produces three charts:

• Average Reward per Edge — convergence of the utility estimate over time
• Cumulative Reward — which edges accumulate the most value
• Arm Selection Frequency — how often each edge was chosen after cold start

A typical 15-run training session completes in about 8–10 minutes and costs roughly $0.08 in API fees. The report is a self-contained HTML file — open it in any browser.

13. Architecture Notes

• Discrete arms, not continuous — the bandit works on discrete edges (research vs lite, review vs publish). Continuous parameters (depth, scrape limit) would need Gaussian Process bandits or policy gradients, which overcomplicate the architecture for an agent workflow.
• Objective-relative weights — the utility function normalizes all raw values to 0–1 before applying weights. This prevents latency (~90,000ms) from dominating quality (~80 points) purely by scale.
• Review is mandatory — the writer → publish shortcut was removed from the graph. Every article must pass through review, ensuring a quality gate before publication.
• Failures are learning signals — retried and rejected edges are recorded with negative rewards. The bandit learns to avoid unreliable paths, not just to favor successful ones.
• Quality is backpropagated — the reviewer’s quality_score is folded back into every edge in the run (except retry paths). Every edge that contributed to a high-quality article shares the credit, so the bandit learns which routes produce the best final output.

14. Going Further

This architecture is a foundation for adaptive agent routing. Natural extensions include:

• Continuous research depth — add 3–4 discrete depth arms (lite/standard/deep) instead of the binary research vs lite choice.
• Cross-objective transfer learning — share edge statistics between objectives with a Bayesian prior so quality runs benefit from data collected under balanced runs.
• Online retraining — run the bandit in a background cron job that periodically retrains on the accumulated edge_stats data and adjusts weights.
• Human-in-the-loop review — replace or augment the LLM reviewer with a human rating, feeding that score back through the same backpropagation mechanism.
• Multi-objective Pareto frontier — instead of a single objective, track the Pareto frontier of quality vs cost vs latency and let the operator choose the operating point.

15. Get the Code

The full source code is on GitHub: github.com/juustesout/contextual-bandit-ai-agent-typescript

Clone it, configure your API keys, point it at a PostgreSQL instance, and run npm run build && node dist/demo.js "your topic" --bandit --objective quality. The training script will show you what the bandit learns over 15 runs.