NeuralUCB Router: An OpenAI-Compatible API Proxy That Routes LLM Requests Using a Multi-Armed Bandit Algorithm

Every team using LLMs in production hits the same wall: you don't need GPT-4 for every request, but you don't know in advance which ones actually need it. A user asking "what is 2+2?" and a user asking "find the bug in this recursive async Rust function" both hit the same /v1/chat/completions endpoint. Route them both to GPT-4o and you're paying $15 per million tokens for arithmetic. Route them both to a local model and quality collapses on complex work.

NeuralUCB Router solves this. It is a production-quality OpenAI-compatible API proxy that uses the NeuralUCB multi-armed bandit algorithm to dynamically route requests to the cheapest model while retaining 92%+ quality. Built autonomously using NEO.

The Problem

The naive solutions all fail:

Always use GPT-4o: $1,500 per month for 100K queries, most of it wasted on simple tasks.

Always use cheap or local model: Quality collapses on hard reasoning, code, and math.

Hard-coded if/else rules: Brittle. You cannot enumerate every query type, and rules do not adapt.

Random routing: No learning. Permanently suboptimal, ignores what is actually working.

Fine-tune a classifier: Requires labeled data, offline training, separate deployment, manual updates.

The root issue: routing is a sequential decision problem under uncertainty. Each query is different, the right model depends on context you can only partially observe, and you need to keep learning as usage patterns shift.

The Solution: NeuralUCB Bandit Routing

NeuralUCB Router frames LLM selection as a multi-armed bandit problem: the same class of problem used in ad auctions, clinical trials, and recommendation systems. Each LLM is an "arm". Each request is a "round". The goal: maximize cumulative reward (quality divided by cost) over time.

A classifier needs labeled training data telling it "this query → use GPT-4". A bandit generates its own training signal by trying models and observing outcomes. It learns online, from live traffic, with no pre-labeled dataset.

How NeuralUCB works for each incoming request:

1. Extract context features
   [query_length, token_entropy, is_code, is_math,
    time_of_day, session_cost, recent_quality, ...]

2. For each LLM arm, estimate:
   UCB score = predicted_reward(context) + exploration_bonus

   Where:
   - predicted_reward  = 2-layer MLP(context)   ← "exploitation"
   - exploration_bonus = λ · sqrt(gᵀ Z⁻¹ g)     ← "exploration"
     g = gradient of MLP w.r.t. last layer
     Z = covariance matrix (Sherman-Morrison updates, O(n²) not O(n³))

3. Route to arm with highest UCB score

4. Observe quality + cost → compute reward → update MLP weights + Z matrix

The UCB bonus is the key insight. A model that hasn't been tried much has high uncertainty, which means a high exploration bonus and more selections until its true value is known. Once well-characterized, the bonus shrinks and the MLP prediction dominates.

What happens after convergence (~500 requests):

Simple factual queries   →  local Ollama/Llama3    (free,   ~98% of the time)
Code generation          →  GPT-4o-mini            ($0.0006, ~87% of the time)
Complex reasoning/debug  →  GPT-4o                 ($0.015,  ~71% of the time)
Creative writing         →  Claude Haiku            ($0.0003, ~82% of the time)

The result: 88% cost reduction versus always-GPT-4o, while retaining 8.9/10 average quality.

Cost Savings

Key Features

• Dynamic Model Routing: Automatically selects the best model for each request based on context
• Cost Optimization: Routes 80% of requests to free/cheap models, saving 75%+ on API costs
• Quality Retention: Maintains 92%+ quality by using expensive models for complex tasks
• OpenAI-Compatible: Drop-in replacement for OpenAI API (/v1/chat/completions, /v1/models)
• Real-time Dashboard: Streamlit dashboard with auto-refresh, cost savings meter, routing heatmap
• Multiple Backends: Supports Ollama (local), OpenAI, Anthropic providers
• Audit Logging: SQLite audit log with CSV/Parquet export for analysis

Architecture

┌─────────────────────────────────────────────────────────────────┐
│                    NeuralUCB Router                             │
├─────────────────────────────────────────────────────────────────┤
│                                                                 │
│  User Request → /v1/chat/completions                            │
│       │                                                         │
│       ▼                                                         │
│  ┌─────────────────────────────────────────────────────────┐    │
│  │  Context Extraction                                     │    │
│  │  • Prompt length, task type, temporal features          │    │
│  └─────────────────────────────────────────────────────────┘    │
│       │                                                         │
│       ▼                                                         │
│  ┌─────────────────────────────────────────────────────────┐    │
│  │  NeuralUCB Bandit                                       │    │
│  │  • 2-layer MLP per model (hidden=64)                    │    │
│  │  • UCB = f(x) + λ√(gᵀZ⁻¹g)                              │    │
│  │  • Sherman-Morrison rank-1 updates                      │    │
│  └─────────────────────────────────────────────────────────┘    │
│       │                                                         │
│       ▼                                                         │
│  ┌─────────────────────────────────────────────────────────┐    │
│  │  Model Selection                                        │    │
│  │  ┌─────┐ ┌─────┐ ┌─────┐                                │    │
│  │  │LLM A│ │LLM B│ │LLM C│ ← Cheapest + Quality           │    │
│  │  └─────┘ └─────┘ └─────┘                                │    │
│  └─────────────────────────────────────────────────────────┘    │
│       │                                                         │
│       ▼                                                         │
│  Response with audit logging                                    │
│                                                                 │
└─────────────────────────────────────────────────────────────────┘

Quick Start

Docker Deployment

# Clone repository
git clone https://github.com/dakshjain-1616/neuralucb-router.git
cd neuralucb-router

# Set API keys
export OPENAI_API_KEY=your-key
export ANTHROPIC_API_KEY=your-key

# Start services
docker-compose up -d

# Access router
curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"model": "auto", "messages": [{"role": "user", "content": "Hello"}]}'

# Access dashboard
open http://localhost:8501

Local Installation

# Install dependencies
pip install -r requirements.txt

# Run router
python -m router.proxy.server

# Run dashboard
streamlit run dashboard/app.py --port 8501

Configuration

Edit configs/example_config.yaml:

models:
  - name: llama3.2-1b
    provider: ollama
    base_url: http://localhost:11434
    cost_per_1k: 0.0
  - name: gpt-4o-mini
    provider: openai
    cost_per_1k: 0.00015
  - name: claude-haiku-4-5
    provider: anthropic
    cost_per_1k: 0.00025

bandit:
  hidden_dim: 64
  lambda_param: 1.0
  exploration_weight: 0.1
  learning_rate: 0.01

Bandit Convergence

Model selection distribution across 1,000 requests:

The bandit learns to route 80% of requests to free Llama3.2-1b while maintaining quality on complex tasks.

Technical Details

NeuralUCB Algorithm

• Architecture: 2-layer MLP (hidden=64, ReLU activation, Sigmoid output) per model arm
• UCB Formula: UCB(model) = f(x; θ_m) + λ * sqrt(gᵀ Z⁻¹ g)
• f(x; θ_m): MLP prediction for model m
• g: Gradient of f with respect to θ_m
• Z: Covariance matrix updated via Sherman-Morrison
• λ: Exploration weight
• Updates: Sherman-Morrison rank-1 updates for Z⁻¹ (O(n²) vs O(n³) for full inverse)

Context Features (15 dimensions)

• prompt_length: Normalized token count
• task_type: One-hot [code, math, creative, qa, chat, other]
• avg_token_length: Average token length
• has_system_prompt: Binary flag
• time_of_day: Sin/cos encoding
• latency_ema: Per-model exponential moving average

Reward Calculation

Two modes:

• Embedding Cosine Similarity: Compare response to reference using sentence-transformers
• LLM-as-Judge: Use LLM to score quality (1-5 scale)

Reward = quality / cost_per_1k (normalized)

Testing

# Run all tests
pytest tests/ -v

# Run specific test
pytest tests/test_neural_ucb.py -v
pytest tests/test_context.py -v
pytest tests/test_proxy.py -v

All tests mock LLM calls. No API keys required for testing.

Dashboard Features

Auto-refresh: Updates every 2 seconds
Cost Savings Meter: Cumulative cost tracking
Model Performance Bar Chart: Selections and average rewards
Routing Heatmap: Visualize routing decisions over time
Recent Decisions Table: Last 20 routing events with metrics

API Endpoints

Project Structure

neuralucb-router/
├── router/bandit/
│   ├── neural_ucb.py     # NeuralUCB bandit algorithm
│   ├── context.py        # Context feature extraction
│   ├── reward.py         # Reward calculation
│   └── __init__.py
├── router/backends/
│   ├── base.py           # Backend provider ABC
│   ├── ollama.py         # Ollama integration
│   ├── openai_backend.py # OpenAI integration
│   ├── anthropic_backend.py # Anthropic integration
│   └── __init__.py
├── router/proxy/
│   ├── server.py         # FastAPI server
│   ├── middleware.py     # Auth/rate limiting
│   ├── openai_schema.py  # Pydantic schemas
│   └── __init__.py
├── router/audit/
│   ├── logger.py         # SQLite audit logging
│   ├── export.py         # CSV/Parquet export
│   └── __init__.py
├── router/config.py      # Configuration loader
├── dashboard/app.py      # Streamlit dashboard
├── tests/
│   ├── test_neural_ucb.py
│   ├── test_context.py
│   ├── test_proxy.py
│   └── __init__.py
├── configs/
│   ├── example_config.yaml
├── hf_export/
│   ├── README.md
│   ├── config.json
│   ├── router_state.json
│   ├── push_to_hub.py
├── infographics/
│   ├── routing_flow.txt
│   ├── cost_comparison.txt
│   ├── bandit_convergence.txt
├── docker-compose.yml
├── Dockerfile
├── requirements.txt
├── pyproject.toml
└── README.md

Export to Hugging Face

from hf_export.push_to_hub import export_to_hub

export_to_hub(
    repo_id="your-username/neuralucb-router",
    config_path="configs/example_config.yaml",
    state_path="hf_export/router_state.json"
)

Comparison

How I Built This Using NEO

This project was built using NEO. NEO is a fully autonomous AI engineering agent that can write code and build solutions for AI/ML tasks including AI model evals, prompt optimization and end to end AI pipeline development.

The requirement was a production-quality OpenAI-compatible API proxy that uses the NeuralUCB multi-armed bandit algorithm to dynamically route LLM requests based on context, cutting costs while retaining quality. NEO planned and produced the files in this repository: the bandit algorithm, context extraction, and reward calculation modules, three backend integrations for Ollama, OpenAI, and Anthropic, a FastAPI proxy server with middleware and Pydantic schemas, a SQLite audit logger with CSV and Parquet export, a Streamlit dashboard, a three-suite test setup, HuggingFace export tooling, Docker and compose files, and the configuration system.

The result is a fully working intelligent router that learns from live traffic, converges after roughly 500 requests, routes 80% of requests to free or cheap models, and delivers 8.9/10 quality at 12% of the cost of always using GPT-4o.

How You Can Use This With NEO

Drop it in as a zero-config OpenAI proxy and immediately start cutting costs.
Since the router exposes a fully OpenAI-compatible /v1/chat/completions endpoint, any existing application pointing at the OpenAI API needs only its base URL changed to start routing intelligently. No code changes required beyond that one configuration line.

Add your own models by editing a single YAML file.
Any provider supported by Ollama, OpenAI, or Anthropic can be added to configs/example_config.yaml with a name, provider, base URL, and cost per 1K tokens. The bandit picks it up on the next start and begins exploring it automatically alongside existing arms.

Use the audit log to understand exactly how your traffic is being routed.
Every routing decision is written to SQLite with full metrics. The /v1/router/audit endpoint and the CSV/Parquet export let teams pull this data into any analytics tool and see which models are handling which query types, what quality scores they are producing, and where cost savings are actually coming from.

Run the Streamlit dashboard as a live cost monitoring layer for your team.
The dashboard auto-refreshes every 2 seconds and shows cumulative cost savings, model selection distribution, and the last 20 routing decisions with metrics. Useful for keeping non-technical stakeholders informed without building a separate reporting tool.

Final Notes

Most teams end up on one of two bad equilibria: spending too much on a frontier model for every request, or accepting degraded quality from a cheap model across the board. NeuralUCB Router sits between those two by learning which requests actually need the expensive model and routing everything else to the cheapest option that can handle it.

The code is at https://github.com/dakshjain-1616/neuralucb-router
The model is on HuggingFace at https://huggingface.co/daksh-neo/neuralucb-router
Install via pip: pip install neuralucb-router
You can also build with NEO in your IDE using the VS Code extension or Cursor.
You can use NEO MCP with Claude Code: https://heyneo.com/claude-code