Build a Multi-Provider AI Router for Cost and Latency Optimization

request -> classify_task -> ROUTING_RULES[task] -> provider 1 -> (error) -> provider 2 ...
   coding     Groq (gpt-oss-120b) -> Codestral -> Mistral Small
   research   Perplexity Agent API (fast effort level: web search, cited answers)
   enterprise Azure OpenAI -> Amazon Bedrock (only vendors your compliance review approved)
   general    Mistral Small -> Groq (gpt-oss-20b)

python3 -m venv .venv && source .venv/bin/activate
pip install groq mistralai openai perplexityai boto3
export GROQ_API_KEY=... MISTRAL_API_KEY=... PERPLEXITY_API_KEY=...
export AZURE_OPENAI_API_KEY=... AZURE_OPENAI_ENDPOINT=https://YOUR-RESOURCE.openai.azure.com
export AZURE_OPENAI_DEPLOYMENT=your-deployment-name
# Amazon Bedrock: no key variable. boto3 reads the standard AWS credential chain (AWS_PROFILE, SSO, instance or task role).
export AWS_DEFAULT_REGION=eu-central-1 BEDROCK_MODEL_ID=eu.anthropic.claude-sonnet-5-5

from dataclasses import dataclass
from enum import Enum

class TaskType(Enum):
    CODING = "coding"
    RESEARCH = "research"
    ENTERPRISE = "enterprise"
    GENERAL = "general"

@dataclass(frozen=True)
class ProviderConfig:
    name: str                                    # groq | mistral | perplexity | azure_openai | bedrock
    model: str                                   # model id, effort level, or env var holding the Azure deployment / Bedrock model id
    usd_per_m: tuple[float, float] | None = None  # (input, output) per 1M tokens; None = provider reports cost or unknown

# Prices copied from each vendor's pricing page; re-check them, they change.
PROVIDERS = {
    "groq_gpt_oss_120b": ProviderConfig("groq", "openai/gpt-oss-120b", (0.15, 0.60)),
    "groq_gpt_oss_20b": ProviderConfig("groq", "openai/gpt-oss-20b", (0.075, 0.30)),
    "codestral": ProviderConfig("mistral", "codestral-latest", (0.30, 0.90)),
    "mistral_small": ProviderConfig("mistral", "mistral-small-latest", (0.15, 0.60)),
    "perplexity_fast": ProviderConfig("perplexity", "fast"),  # Agent API effort level; the response includes its cost
    "azure_openai": ProviderConfig("azure_openai", "AZURE_OPENAI_DEPLOYMENT"),  # env var holding your deployment name
    "bedrock": ProviderConfig("bedrock", "BEDROCK_MODEL_ID"),  # env var holding the model or inference profile id
}

# Task type -> providers in order of preference. Enterprise lists only providers your compliance review
# approved; failing over to one it did not cover would defeat the point. Drop "bedrock" if Azure is the only one.
ROUTING_RULES = {
    TaskType.CODING: ["groq_gpt_oss_120b", "codestral", "mistral_small"],
    TaskType.RESEARCH: ["perplexity_fast"],
    TaskType.ENTERPRISE: ["azure_openai", "bedrock"],
    TaskType.GENERAL: ["mistral_small", "groq_gpt_oss_20b"],
}

import logging, os, re, time
import boto3
from dataclasses import dataclass
from groq import Groq
from mistralai.client import Mistral
from openai import OpenAI
from perplexity import Perplexity
from config import TaskType, PROVIDERS, ROUTING_RULES

log = logging.getLogger("router")

@dataclass
class RouteResult:
    content: str
    provider: str
    model: str
    input_tokens: int
    output_tokens: int
    latency_ms: float
    cost_usd: float | None

def _usage(resp):  # Groq, Mistral and Azure (OpenAI SDK) share this shape
    return resp.choices[0].message.content, resp.usage.prompt_tokens, resp.usage.completion_tokens

def call_groq(model, messages):
    return _usage(Groq(api_key=os.environ["GROQ_API_KEY"]).chat.completions.create(model=model, messages=messages)) + (None,)

def call_mistral(model, messages):
    return _usage(Mistral(api_key=os.environ["MISTRAL_API_KEY"]).chat.complete(model=model, messages=messages)) + (None,)

def call_azure(model, messages):
    # Azure's v1 API: the stock OpenAI client, /openai/v1/ base URL, no api-version; `model` is your deployment name
    client = OpenAI(api_key=os.environ["AZURE_OPENAI_API_KEY"],
                    base_url=os.environ["AZURE_OPENAI_ENDPOINT"].rstrip("/") + "/openai/v1/")
    return _usage(client.chat.completions.create(model=os.environ[model], messages=messages)) + (None,)

def call_bedrock(model, messages):
    # Converse API: one request shape for every Bedrock model. Region and credentials come from the standard AWS chain.
    resp = boto3.client("bedrock-runtime").converse(
        modelId=os.environ[model],
        messages=[{"role": m["role"], "content": [{"text": m["content"]}]} for m in messages],
        inferenceConfig={"maxTokens": 1024},
    )
    blocks = resp["output"]["message"]["content"]  # reasoning models may also return reasoningContent blocks
    text = "".join(b["text"] for b in blocks if "text" in b)
    return text, resp["usage"]["inputTokens"], resp["usage"]["outputTokens"], None

def call_perplexity(effort, messages):
    # Sonar chat completions are retired; web-grounded answers now come from the Agent API
    resp = Perplexity(api_key=os.environ["PERPLEXITY_API_KEY"]).responses.create(effort=effort, input=messages[-1]["content"])
    cost = resp.usage.cost.total_cost if resp.usage and resp.usage.cost else None
    return resp.output_text, resp.usage.input_tokens, resp.usage.output_tokens, cost

CALLERS = {"groq": call_groq, "mistral": call_mistral, "perplexity": call_perplexity, "azure_openai": call_azure, "bedrock": call_bedrock}

class Stats:
    """Per-provider spend, call counts and an exponential moving average of observed latency."""
    def __init__(self):
        self.cost, self.calls, self.latency = {}, {}, {}
    def record(self, key, cost, ms):
        self.calls[key] = self.calls.get(key, 0) + 1
        self.cost[key] = self.cost.get(key, 0.0) + (cost or 0.0)
        self.latency[key] = ms if key not in self.latency else 0.8 * self.latency[key] + 0.2 * ms
    def summary(self):
        return {"calls": self.calls, "cost_usd": {k: round(v, 6) for k, v in self.cost.items()},
                "avg_latency_ms": {k: round(v) for k, v in self.latency.items()}}

stats = Stats()

def call_provider(key, messages) -> RouteResult:
    cfg = PROVIDERS[key]
    start = time.perf_counter()
    content, tin, tout, cost = CALLERS[cfg.name](cfg.model, messages)
    ms = (time.perf_counter() - start) * 1000
    if cost is None and cfg.usd_per_m:
        cost = (tin * cfg.usd_per_m[0] + tout * cfg.usd_per_m[1]) / 1e6
    stats.record(key, cost, ms)
    return RouteResult(content, key, cfg.model, tin, tout, round(ms, 1), None if cost is None else round(cost, 6))

def route(messages, task_type=TaskType.GENERAL, prefer_fast=False) -> RouteResult:
    keys = list(ROUTING_RULES[task_type])
    if prefer_fast:  # providers we have not measured yet keep their configured order, after measured ones
        keys.sort(key=lambda k: stats.latency.get(k, float("inf")))
    errors = []
    for key in keys:
        try:
            return call_provider(key, messages)
        except Exception as e:  # rate limit, timeout, 5xx, missing credentials
            log.warning("%s failed: %s", key, e)
            errors.append(f"{key}: {e}")
    raise RuntimeError("All providers failed: " + "; ".join(errors))

RULES = [
    (TaskType.CODING, r"\b(code|function|debug|implement|refactor|python|typescript|sql|regex)\b"),
    (TaskType.RESEARCH, r"\b(latest|current|recent|news|today|this week|search|look up)\b"),
    (TaskType.ENTERPRISE, r"\b(compliance|audit|policy|regulation|hipaa|soc ?2|gdpr)\b"),
]

def classify_task(text: str) -> TaskType:
    return next((t for t, pat in RULES if re.search(pat, text.lower())), TaskType.GENERAL)

import json
from router import route, classify_task, stats
from config import TaskType

def ask(text: str) -> str:
    task = classify_task(text)
    r = route([{"role": "user", "content": text}], task, prefer_fast=task == TaskType.CODING)
    print(f"[{task.value}] {r.provider} ({r.model}) {r.latency_ms} ms, cost: {r.cost_usd}")
    return r.content

if __name__ == "__main__":
    for q in ["Write a Python function that parses JSON with error handling",
              "What changed in the news about EU AI regulation this week?",
              "Draft a SOC 2 evidence checklist for our SaaS",
              "Explain the CAP theorem"]:
        print(ask(q), "\n")
    print(json.dumps(stats.summary(), indent=2))

$ python main.py
[coding] groq_gpt_oss_120b (openai/gpt-oss-120b) 410.2 ms, cost: 0.00031
...
{
  "calls": {"groq_gpt_oss_120b": 1, "perplexity_fast": 1, "azure_openai": 1, "mistral_small": 1},
  "cost_usd": {"groq_gpt_oss_120b": 0.00031, "perplexity_fast": 0.0094, "azure_openai": 0.0, "mistral_small": 0.00012},
  "avg_latency_ms": {"groq_gpt_oss_120b": 410, "perplexity_fast": 5100, "azure_openai": 900, "mistral_small": 700}
}