Sub-15ms Semantic Routers: Enforcing Strict Policy Without Burning LLM Tokens

import numpy as np
from typing import Dict, List, Optional
import onnxruntime as ort
from transformers import AutoTokenizer

class FastSemanticRouter:
    def __init__(self, onnx_model_path: str, tokenizer_name: str):
        self.session = ort.InferenceSession(onnx_model_path)
        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
        self.routes: Dict[str, np.ndarray] = {}

    def add_route(self, name: str, sample_embeddings: List[List[float]], threshold: float = 0.82):
        # Store normalized centroid vector for each route
        centroid = np.mean(sample_embeddings, axis=0)
        centroid = centroid / np.linalg.norm(centroid)
        self.routes[name] = {"centroid": centroid, "threshold": threshold}

    def route_query(self, query: str) -> Optional[str]:
        # Fast 5ms tokenization & ONNX forward pass on CPU
        inputs = self.tokenizer(query, return_tensors="np", padding=True, truncation=True)
        outputs = self.session.run(None, dict(inputs))
        query_vec = outputs[0][:, 0, :] # CLS token
        query_vec = query_vec / np.linalg.norm(query_vec)

        best_route = None
        highest_sim = -1.0

        for name, route_data in self.routes.items():
            sim = np.dot(query_vec, route_data["centroid"])[0]
            if sim > route_data["threshold"] and sim > highest_sim:
                highest_sim = sim
                best_route = name

        return best_route