Code Example: Bandits in your LLM Gateway

This folder contains the code for the blog post Bandits in your LLM Gateway. ¹

Running the Experiment

Prerequisites

Make sure you have the following environment variables set:

export ANTHROPIC_API_KEY=your_api_key_here

Setup

1. Run Postgres migrations (required on first run):

docker compose run --rm gateway --run-postgres-migrations

2. Start all services:

docker compose up

This will start:

• ClickHouse: Database for inference results and feedback (port 8123)
• Postgres: Database for TensorZero metadata (port 5432)
• Gateway: TensorZero Gateway (port 3000)
• UI: TensorZero observability UI (port 4000)

Running the Experiment

Once the services are running, execute the experiment script:

uv run main.py

This will:

• Load NER (Named Entity Recognition) data from the CoNLL++ dataset
• Send inference requests to the TensorZero Gateway
• Submit feedback for each inference
• The Track-and-Stop algorithm will adaptively adjust sampling probabilities every 15 seconds

Viewing Results

• Real-time monitoring: Open http://localhost:4000/observability/functions/extract_entities to view the TensorZero UI