router.py

"""
R2-Router: LLM Router with Joint Model-Budget Optimization

Self-contained inference module. Routes queries to the optimal (model, token_budget)
pair by predicting per-query quality and cost using KNN.

Usage:
    from router import R2Router
    router = R2Router.from_pretrained("jqxue1999/r2-router")
    result = router.route(embedding)  # embedding: np.ndarray (1024,)

    # Or train from scratch:
    router = R2Router.from_training_data("jqxue1999/r2-router")
"""

import os
import json
import numpy as np
import joblib
from typing import Dict, List, Optional, Union
from sklearn.neighbors import KNeighborsRegressor


class R2Router:
    """
    R2-Router: Routes queries to optimal (LLM, token_budget) pair.

    Uses KNN to predict quality for each (model, budget) combination,
    then selects the pair that maximizes:
        risk = (1 - lambda) * quality - lambda * tokens * price / 1e6
    """

    def __init__(
        self,
        quality_knns: Dict[str, Dict[str, KNeighborsRegressor]],
        token_knns: Dict[str, KNeighborsRegressor],
        model_prices: Dict[str, float],
        model_names: Dict[str, str],
        budgets: Dict[str, int],
        lambda_val: float = 0.999,
    ):
        self.quality_knns = quality_knns  # {model: {budget: KNN}}
        self.token_knns = token_knns      # {model: KNN}
        self.model_prices = model_prices  # {model: price_per_million_output_tokens}
        self.model_names = model_names    # {short_name: full_name}
        self.budgets = budgets            # {budget_name: token_limit}
        self.lambda_val = lambda_val

    @classmethod
    def from_pretrained(cls, path: str, lambda_val: float = 0.999) -> "R2Router":
        """
        Load pre-trained KNN checkpoints.

        Args:
            path: Local directory or HuggingFace repo ID (e.g., "jqxue1999/r2-router")
            lambda_val: Cost-accuracy tradeoff (higher = more cost-sensitive)
        """
        # If HF repo ID, download first
        if not os.path.isdir(path):
            path = cls._download_from_hf(path)

        with open(os.path.join(path, "config.json")) as f:
            config = json.load(f)

        ckpt_dir = os.path.join(path, "checkpoints")
        quality_knns = {}
        token_knns = {}

        for model_name in config["models"]:
            quality_knns[model_name] = {}
            for budget_name in config["budgets"]:
                ckpt_path = os.path.join(ckpt_dir, f"quality_knn_{model_name}_{budget_name}.joblib")
                if os.path.exists(ckpt_path):
                    quality_knns[model_name][budget_name] = joblib.load(ckpt_path)

            tok_path = os.path.join(ckpt_dir, f"token_knn_{model_name}.joblib")
            if os.path.exists(tok_path):
                token_knns[model_name] = joblib.load(tok_path)

        model_prices = {
            mn: cfg["output_price_per_million"]
            for mn, cfg in config["models"].items()
        }
        model_names = {
            mn: cfg["full_name"]
            for mn, cfg in config["models"].items()
        }

        return cls(
            quality_knns=quality_knns,
            token_knns=token_knns,
            model_prices=model_prices,
            model_names=model_names,
            budgets=config["budgets"],
            lambda_val=lambda_val,
        )

    @classmethod
    def from_training_data(
        cls,
        path: str,
        k: int = 80,
        lambda_val: float = 0.999,
    ) -> "R2Router":
        """
        Train KNN from scratch using the provided training data.

        Args:
            path: Local directory or HuggingFace repo ID
            k: Number of KNN neighbors
            lambda_val: Cost-accuracy tradeoff
        """
        if not os.path.isdir(path):
            path = cls._download_from_hf(path)

        with open(os.path.join(path, "config.json")) as f:
            config = json.load(f)

        X_train = np.load(os.path.join(path, "training_data", "embeddings.npy"))
        with open(os.path.join(path, "training_data", "labels.json")) as f:
            labels = json.load(f)

        quality_knns = {}
        token_knns = {}

        for model_name, model_labels in labels.items():
            quality_knns[model_name] = {}
            for budget_name, bdata in model_labels.items():
                acc = np.array([x if x is not None else np.nan for x in bdata["accuracy"]])
                valid = ~np.isnan(acc)
                if valid.sum() < 3:
                    continue
                knn = KNeighborsRegressor(
                    n_neighbors=min(k, int(valid.sum()) - 1),
                    metric="cosine",
                    weights="distance",
                )
                knn.fit(X_train[valid], acc[valid])
                quality_knns[model_name][budget_name] = knn

            # Token predictor (use concise budget's output_tokens)
            if "concise" in model_labels and "output_tokens" in model_labels["concise"]:
                tok = np.array([x if x is not None else np.nan for x in model_labels["concise"]["output_tokens"]])
                valid = ~np.isnan(tok)
                if valid.sum() >= 3:
                    tknn = KNeighborsRegressor(
                        n_neighbors=min(k, int(valid.sum()) - 1),
                        metric="cosine",
                        weights="distance",
                    )
                    tknn.fit(X_train[valid], tok[valid])
                    token_knns[model_name] = tknn

        model_prices = {
            mn: cfg["output_price_per_million"]
            for mn, cfg in config["models"].items()
        }
        model_names = {
            mn: cfg["full_name"]
            for mn, cfg in config["models"].items()
        }

        return cls(
            quality_knns=quality_knns,
            token_knns=token_knns,
            model_prices=model_prices,
            model_names=model_names,
            budgets=config["budgets"],
            lambda_val=lambda_val,
        )

    @staticmethod
    def _download_from_hf(repo_id: str) -> str:
        """Download model from Hugging Face Hub."""
        try:
            from huggingface_hub import snapshot_download
            return snapshot_download(repo_id)
        except ImportError:
            raise ImportError(
                "huggingface_hub is required to download from HF. "
                "Install with: pip install huggingface_hub"
            )

    def route(
        self,
        embedding: np.ndarray,
        lambda_val: Optional[float] = None,
    ) -> Dict:
        """
        Route a query to the optimal (model, token_budget) pair.

        Args:
            embedding: Query embedding vector, shape (1024,) or (1, 1024)
            lambda_val: Override default lambda (higher = more cost-sensitive)

        Returns:
            Dict with keys: model, model_full_name, budget, token_limit,
                           predicted_quality, predicted_cost, risk, all_options
        """
        if embedding.ndim == 1:
            embedding = embedding.reshape(1, -1)

        lam = lambda_val if lambda_val is not None else self.lambda_val
        all_options = []

        for mn in self.quality_knns:
            price = self.model_prices.get(mn, 0)

            # Predict output tokens
            if mn in self.token_knns:
                tok = max(1.0, float(self.token_knns[mn].predict(embedding)[0]))
            else:
                tok = 50.0

            for budget_name, knn in self.quality_knns[mn].items():
                q = float(knn.predict(embedding)[0])
                risk = (1 - lam) * q - lam * tok * price / 1e6

                all_options.append({
                    "model": mn,
                    "model_full_name": self.model_names.get(mn, mn),
                    "budget": budget_name,
                    "token_limit": self.budgets.get(budget_name, budget_name),
                    "predicted_quality": q,
                    "predicted_tokens": tok,
                    "predicted_cost": tok * price / 1e6,
                    "risk": risk,
                })

        if not all_options:
            raise RuntimeError("No valid routing options")

        best = max(all_options, key=lambda x: x["risk"])
        best["all_options"] = all_options
        return best

    def route_batch(
        self,
        embeddings: np.ndarray,
        lambda_val: Optional[float] = None,
    ) -> List[Dict]:
        """
        Route a batch of queries.

        Args:
            embeddings: Query embeddings, shape (N, 1024)
            lambda_val: Override default lambda

        Returns:
            List of routing decisions
        """
        return [self.route(embeddings[i], lambda_val) for i in range(len(embeddings))]