"""Pick the model that can serve a request, before the proxy dispatches it. LiteLLM loads this through ``get_instance_fn``, which resolves ``router_hook.instance`` against the directory of ``config.yaml``, so the file lives beside it in the container and the image stays the upstream one. The demand side is estimated from the request body. The supply side is read off LiteLLM's own model catalog, which ships 1595 entries in the image carrying the context window, the modality and tool support, the per-token prices and the deprecation date of every model it knows. A route may override any of that through ``model_info.traits``, and nothing here probes a backend. Two polarities govern the filters, and mixing them up routes silently wrong: * A capability the request *demands* must be proven. Vision, tools, a response schema, PDF or audio input are catalogued wherever they exist, so an absent flag means absent support and the route is dropped only when the request asks for it. * A *disqualifier* acts on explicit evidence only. Most entries carry no ``supports_system_messages`` and local tags are catalogued not at all, so treating silence as denial would drop nearly every route. A model is therefore excluded only when the catalog states a non-chat mode, denies system messages outright, or gives a deprecation date already past. """ from __future__ import annotations import asyncio import json import random import re import time import urllib.request from datetime import date import litellm from litellm.integrations.custom_logger import CustomLogger from infinito_decider import Decider, History, Sampler, record_of ROUTER_ALIAS = {{ LITELLM_ROUTER_ALIAS | to_json }} REMOTE_FALLBACK = {{ 'True' if LITELLM_ROUTER_REMOTE_FALLBACK | bool else 'False' }} ROUTER_WEIGHTS = {{ LITELLM_ROUTER_WEIGHTS | to_json }} ROUTER_STRATEGY = {{ LITELLM_ROUTER_STRATEGY | to_json }} S1_URL = {{ LITELLM_S1_URL | to_json }} S1_KEY = {{ LITELLM_S1_KEY | to_json }} S1_MODEL = {{ LITELLM_S1_MODEL | to_json }} S1_TIMEOUT = {{ LITELLM_S1_TIMEOUT | int }} S1_QUESTION = "route" S1_VERDICT_QUESTION = "verdict" S1_MAX_STATE_CHARS = {{ LITELLM_ROUTER_STATE_CHARS | int }} SAMPLE_ENABLED = {{ 'True' if LITELLM_ROUTER_SAMPLE_ENABLED | bool else 'False' }} SAMPLE_MODE = {{ LITELLM_ROUTER_SAMPLE_MODE | to_json }} SAMPLE_TRIGGER = {{ LITELLM_ROUTER_SAMPLE_TRIGGER | to_json }} SAMPLE_INTERVAL = {{ LITELLM_ROUTER_SAMPLE_INTERVAL | int }} SAMPLE_WIDTH = {{ LITELLM_ROUTER_SAMPLE_WIDTH | int }} HISTORY_TABLE = "infinito_router_history" TOKEN_BUCKETS = (256, 1024, 4096, 16384, 65536) SECRET_SHAPES = re.compile( r"-----BEGIN [A-Z ]*PRIVATE KEY-----" r"|\bsk-[A-Za-z0-9_-]{20,}" r"|\bgh[pousr]_[A-Za-z0-9]{20,}" r"|\bgithub_pat_[A-Za-z0-9_]{20,}" r"|\bAKIA[A-Z0-9]{16}\b" r"|\bxox[baprs]-[A-Za-z0-9-]{10,}" r"|\bAIza[A-Za-z0-9_-]{35}\b" ) ISO_DATE = re.compile(r"^\d{4}-\d{2}-\d{2}$") MIN_CHARS_PER_TOKEN = {{ LITELLM_ROUTER_MIN_CHARS_PER_TOKEN | int }} IMAGE_PARTS = ("image_url", "input_image") AUDIO_PARTS = ("input_audio", "audio_url") DOCUMENT_PARTS = ("file", "input_file", "document") SYSTEM_ROLES = ("system", "developer") UNPRICED_REMOTE_COST = 1.0 NOMINAL_OUTPUT_TOKENS = 512 CHAT_MODES = ("chat", "responses", None) def summarize(data): """What the request demands, in the terms a route declares. Args: data: the OpenAI-shaped request body the proxy received. Returns: A mapping with the estimated ``input_tokens`` and ``max_output`` plus one flag per capability the body asks for. """ characters = 0 need = { "vision": False, "audio": False, "documents": False, "images": 0, "system": False, } for message in data.get("messages") or []: if message.get("role") in SYSTEM_ROLES: need["system"] = True content = message.get("content") if isinstance(content, str): characters += len(content) continue for part in content or []: if not isinstance(part, dict): continue kind = part.get("type") if kind in IMAGE_PARTS: need["vision"] = True need["images"] += 1 elif kind in AUDIO_PARTS: need["audio"] = True elif kind in DOCUMENT_PARTS: need["documents"] = True else: characters += len(part.get("text") or "") response_format = data.get("response_format") or {} need.update( { "input_tokens": -(-characters // MIN_CHARS_PER_TOKEN), "max_output": int(data.get("max_tokens") or 0), "tools": bool(data.get("tools") or data.get("functions")), "tool_choice": data.get("tool_choice") not in (None, "none"), "parallel_tools": bool(data.get("parallel_tool_calls")), "schema": response_format.get("type") in ("json_schema", "json_object"), } ) return need def _catalogued(model, catalog): """The catalog entry for *model*, trying the prefixed and the bare key. ``ollama/llama3`` is a key as it stands, ``openai/gpt-4o-mini`` is filed under ``gpt-4o-mini``, and a local tag or a mock is in neither. """ for key in (model, model.split("/", 1)[-1]): found = catalog.get(key) if isinstance(found, dict): return found return {} def is_local(params): """Whether this deployment serves the route itself. An ``api_base`` points at a backend of this deployment and a ``mock_response`` never leaves the proxy. Anything else reaches a third party, which is why this is read twice: once as the factor the operator weights, and once as the filter that no weight can reach. """ return params.get("mock_response") is not None or bool(params.get("api_base")) def _deprecated(entry, today=None): """Whether the catalog gives a deprecation date already past. The catalog ships a ``sample_spec`` documentation stub whose date field is prose, so only an ISO value counts and anything else reads as undated. """ value = str(entry.get("deprecation_date") or "") if not ISO_DATE.match(value): return False return value < (today or date.today().isoformat()) def capabilities(entry, catalog, today=None): """What a route can do: the catalog's facts, overridden by declarations. Args: entry: one ``model_list`` entry. catalog: ``litellm.model_cost`` or a stand-in keyed the same way. today: the ISO date deprecation is measured against. Returns: A mapping of every capability, limit and price the filters and the strategies read. A demanded capability counts as absent unless proven; a disqualifier counts as harmless unless the catalog states it. """ params = entry.get("litellm_params") or {} info = entry.get("model_info") or {} known = _catalogued(params.get("model") or "", catalog) local = is_local(params) caps = { "context": info.get("max_input_tokens") or known.get("max_input_tokens"), "max_output": known.get("max_output_tokens"), "vision": bool(known.get("supports_vision")), "tools": bool(known.get("supports_function_calling")), "tool_choice": bool(known.get("supports_tool_choice")), "parallel_tools": bool(known.get("supports_parallel_function_calling")), "schema": bool(known.get("supports_response_schema")), "documents": bool(known.get("supports_pdf_input")), "audio": bool(known.get("supports_audio_input")), "images": known.get("max_images_per_prompt"), "system": known.get("supports_system_messages") is not False, "servable": known.get("mode") in CHAT_MODES or not known, "deprecated": _deprecated(known, today), "input_cost": known.get( "input_cost_per_token", 0.0 if local else UNPRICED_REMOTE_COST ), "output_cost": known.get( "output_cost_per_token", 0.0 if local else UNPRICED_REMOTE_COST ), "caching": bool(known.get("supports_prompt_caching")), "throughput": known.get("tpm") or 0, "local": local, "speed": 0.0, } caps.update( {key: value for key, value in (info.get("traits") or {}).items() if key in caps} ) return caps def expected_cost(caps, need): """What this route would bill for this request, input plus output. Ranking on the input price alone reads the cheaper half of the bill: an output token commonly costs several times an input one, and a short prompt with a long answer inverts the order the input price suggests. """ output_tokens = need["max_output"] or NOMINAL_OUTPUT_TOKENS return ( float(caps["input_cost"]) * need["input_tokens"] + float(caps["output_cost"]) * output_tokens ) def request_text(data): """Everything the caller wrote, as one string.""" parts = [] for message in data.get("messages") or []: content = message.get("content") if isinstance(content, str): parts.append(content) continue for part in content or []: if isinstance(part, dict) and part.get("text"): parts.append(str(part["text"])) return "\n".join(parts) def carries_secret(data): """Whether the request shows a structured secret.""" return bool(SECRET_SHAPES.search(request_text(data))) def local_reason(data, deployments, catalog): """Why this request must stay local, or the empty string. Returns empty when no local route is deployed at all, because then there is nothing to keep it on and refusing would only turn a working request into an error without protecting anything. """ if not any( capabilities(entry, catalog)["local"] for entry in deployments or [] if entry.get("model_name") not in (None, ROUTER_ALIAS) ): return "" if carries_secret(data): return "the request shows a structured secret" if not REMOTE_FALLBACK: return "services.litellm.router.remote_fallback is off" return "" def serves(caps, need): """Whether one route can take this request, or the reason it cannot.""" context = caps["context"] if context and need["input_tokens"] + need["max_output"] > int(context): return "its context window is too small" output_cap = caps["max_output"] if need["max_output"] and output_cap and need["max_output"] > int(output_cap): return "it caps the output below the request" for demand in ("vision", "audio", "documents", "tools", "tool_choice", "parallel_tools", "schema"): if need[demand] and not caps[demand]: return f"it does not support {demand.replace('_', ' ')}" if need["images"] and caps["images"] and need["images"] > int(caps["images"]): return "it takes fewer images than the request carries" if need["system"] and not caps["system"]: return "it does not take system messages" if not caps["servable"]: return "it is not a chat model" if caps["deprecated"]: return "it is past its deprecation date" return "" def eligible(deployments, need, exclude, catalog, today=None): """The routes that can serve this request, in publication order. Args: deployments: the proxy's own ``model_list`` entries. need: the mapping ``summarize`` returned. exclude: the alias the request named, so the router never picks itself. catalog: the model catalog ``capabilities`` reads. today: the ISO date deprecation is measured against. Returns: The surviving entries, each carrying its resolved capabilities under ``_caps`` and this request's expected price under ``_cost``. """ survivors = [] for entry in deployments or []: alias = entry.get("model_name") if not alias or alias == exclude: continue caps = capabilities(entry, catalog, today) if serves(caps, need): continue survivors.append( {**entry, "_caps": caps, "_cost": expected_cost(caps, need)} ) return survivors def _spread(values): """The low and high of *values*, for scaling one factor against another.""" numbers = [float(value) for value in values] return min(numbers), max(numbers) def _scaled(value, spread, *, higher_is_better): """*value* as 0..1 within *spread*, where 1 is always the better end. A spread of zero means the factor cannot separate this candidate set, so every route scores 1 and the factor drops out instead of deciding on a rounding artefact. """ low, high = spread if high == low: return 1.0 position = (float(value) - low) / (high - low) return position if higher_is_better else 1.0 - position def score(entry, spreads, weights): """How well a route serves the operator's stated priorities, 0..1 each. Each factor is scaled inside the candidate set rather than against an absolute, because the units do not compare: a price is a fraction of a cent, a rate is tens of tokens per second and a quota is millions. """ caps = entry["_caps"] return ( float(weights.get("locality", 0.0)) * (1.0 if caps["local"] else 0.0) + float(weights.get("cost", 0.0)) * _scaled(entry["_cost"], spreads["cost"], higher_is_better=False) + float(weights.get("speed", 0.0)) * _scaled(caps["speed"], spreads["speed"], higher_is_better=True) + float(weights.get("caching", 0.0)) * (1.0 if caps["caching"] else 0.0) + float(weights.get("throughput", 0.0)) * _scaled(caps["throughput"], spreads["throughput"], higher_is_better=True) ) def decide_weighted(candidates, weights): """The highest scoring route, ties broken by alias so a rerun repeats. Weights come from ``services.litellm.router_weights`` and say how much each factor matters, not whether it may be overruled. Confidentiality is not in here: a request that must stay local was already restricted before this ran, because a weight is a preference and a credential leaving the cluster is not something a preference should be able to buy. """ spreads = { "cost": _spread(entry["_cost"] for entry in candidates), "speed": _spread(entry["_caps"]["speed"] for entry in candidates), "throughput": _spread(entry["_caps"]["throughput"] for entry in candidates), } return min( candidates, key=lambda entry: (-score(entry, spreads, weights), entry["model_name"]), )["model_name"] def bucket(tokens): """The token band *tokens* falls in, as the label history groups under.""" for edge in TOKEN_BUCKETS: if tokens <= edge: return str(edge) return f"{TOKEN_BUCKETS[-1]}+" def signature(need): """The request class routing history accumulates under. Two requests share a class when they ask a route for the same kind of work: the same token band and the same capability demands. Grouping finer spreads the evidence too thin to ever decide; grouping coarser lets a verdict about a vision request speak for a plain chat. """ demands = "+".join( flag for flag in ("vision", "tools", "schema", "audio", "documents") if need[flag] ) return f"{bucket(need['input_tokens'] + need['max_output'])}:{demands or 'plain'}" def headroom(entry, need): """How much context this route would leave unused on this request. A route whose window is unknown sorts last rather than first: no window means no evidence of fit, and reading it as unlimited would hand every request to the least described route. """ context = entry["_caps"]["context"] if not context: return float("inf") return int(context) - (need["input_tokens"] + need["max_output"]) def describe(entry, need, tally=None): """One line telling the decider what this route is, in its own terms. Only measured or catalogued facts go in. A figure the catalog does not state is left out rather than defaulted, because a decider cannot tell an absent number from a zero one. """ caps = entry["_caps"] window = caps["context"] or "unstated" room = headroom(entry, need) facts = [ "runs in this cluster" if caps["local"] else "runs at a provider", f"context window {window}", "unknown headroom" if room == float("inf") else f"{int(room)} tokens spare", f"costs about {entry['_cost']:.4f} for this request", ] if caps["speed"]: facts.append(f"{caps['speed']:.1f} tokens per second measured") if caps["throughput"]: facts.append(f"{int(caps['throughput'])} tokens per minute quota") if caps["caching"]: facts.append("caches prompts") earlier = record_of(entry["model_name"], tally or {}) if earlier: facts.append(earlier) return ", ".join(facts) def ask_s1(text, candidates, need, tally=None): """Put the routing choice to the decider as one typed choice question. The decider answers typed questions rather than generating, so the candidates become the options of a single ``choice`` and the prompt becomes the state it classifies. Returns ``(alias, confidence)``. Raises: OSError: when the decider cannot be reached. ValueError: when it answers a shape or an option this router did not offer, which must not be read as a route. """ criteria = { entry["model_name"]: describe(entry, need, tally) for entry in candidates } answer = _systemone( text, S1_QUESTION, "Which of these models should answer the request? Prefer the smallest " "one that can serve it well.", criteria, ) alias = answer.get("choice") if alias not in criteria: raise ValueError( f"the decider chose {alias!r}, which was not among {sorted(criteria)}" ) return alias, float(answer.get("confidence") or 0.0) DECIDER = Decider(S1_URL, S1_KEY, S1_MODEL, S1_TIMEOUT, S1_MAX_STATE_CHARS) def _systemone(state, question, instructions, criteria): """Put one typed choice to the decider and return the answer block it sent. Raises: OSError: when the decider cannot be reached. ValueError: when it answers a different question than the one asked. """ return DECIDER.ask(state, question, instructions, criteria) def judge_answers(text, replies): """The alias whose reply the decider calls the best answer to *text*.""" return DECIDER.judge_answers(text, replies, S1_VERDICT_QUESTION) def decide_system_one(candidates, need, text, tally=None): """The route the decider picks for this prompt. The weighted strategy ranks on standing preferences and stops reading the prompt once the filter has run. This one hands the prompt and the surviving candidates to a System One model, which answers a typed choice without generating. An unreachable decider is not routed around: the caller named an alias precisely so it would not have to know which model answered, and quietly substituting a different decider would make that answer mean something else. """ alias, _confidence = ask_s1(text, candidates, need, tally) return alias def decide(candidates, need, text="", weights=None, strategy=None, tally=None): """Dispatch to the configured strategy. Args: candidates: the surviving entries from ``eligible``. need: the mapping ``summarize`` returned. text: what the caller wrote, which only ``system_one`` reads. weights: the weighted strategy's factor weights. strategy: ``system_one`` or ``weighted``. tally: what earlier comparisons said, which only ``system_one`` reads. Returns: The winning alias. """ strategy = ROUTER_STRATEGY if strategy is None else strategy if strategy == "system_one": return decide_system_one(candidates, need, text, tally) return decide_weighted(candidates, ROUTER_WEIGHTS if weights is None else weights) async def collect(router, data, candidates, width): """Ask the first *width* candidates the same question. Returns ``{alias: reply}`` for those that answered. A route that raises is left out rather than recorded as a loss: failing to answer is a fact about the moment, not evidence that its answers are worse. """ chosen = candidates[:width] async def one(entry): payload = {key: value for key, value in data.items() if key != "mock_response"} payload["model"] = entry["model_name"] reply = await router.acompletion(**payload) return entry["model_name"], reply.choices[0].message.content or "" replies = {} for outcome in await asyncio.gather( *(one(entry) for entry in chosen), return_exceptions=True ): if isinstance(outcome, BaseException): continue alias, content = outcome if content: replies[alias] = content return replies class SystemOneRouter(CustomLogger): """Rewrites the router alias to a model that can actually serve it. Every path that cannot route returns a string, which the proxy turns into a RejectedRequestError carrying that text. Answering from some default instead would hide a broken router behind a plausible reply, and the alias exists precisely so callers do not have to know which model answered. """ def __init__(self): super().__init__() self._sampler = Sampler(SAMPLE_TRIGGER, SAMPLE_INTERVAL) self._store = None async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): # noqa: ARG002 - LiteLLM hook signature if data.get("model") != ROUTER_ALIAS: return None from litellm.proxy.proxy_server import llm_router if llm_router is None: return ( f"'{ROUTER_ALIAS}' cannot route: the proxy has no router, so no " f"model list is available to choose from" ) need = summarize(data) candidates = eligible( llm_router.model_list, need, ROUTER_ALIAS, litellm.model_cost ) demand = ( f"about {need['input_tokens']} input plus {need['max_output']} " f"output tokens with vision={need['vision']} tools={need['tools']} " f"schema={need['schema']} audio={need['audio']} " f"documents={need['documents']}" ) if not candidates: return f"'{ROUTER_ALIAS}' cannot route: no served model can take {demand}" reason = local_reason(data, llm_router.model_list, litellm.model_cost) if reason: candidates = [entry for entry in candidates if entry["_caps"]["local"]] if not candidates: return ( f"'{ROUTER_ALIAS}' cannot route: {reason}, and no local model " f"can take {demand}. A remote model could, so either shrink " f"the request or set services.litellm.router.remote_fallback" ) text = request_text(data) request_class = signature(need) tally = await self._tally(request_class) try: data["model"] = decide(candidates, need, text, tally=tally) except (OSError, ValueError) as error: return ( f"'{ROUTER_ALIAS}' cannot route: the {ROUTER_STRATEGY} strategy " f"could not reach a decision ({type(error).__name__}: {error}). " f"Answering from a substitute decider would make the alias mean " f"something the caller did not ask for" ) if SAMPLE_ENABLED and len(candidates) > 1 and self._sampler.due(): await self._compare(llm_router, data, candidates, text, request_class) return data async def _tally(self, request_class): """What earlier comparisons said, or nothing when they cannot be read.""" store = self._history() if store is None: return {} try: return await store.wins(request_class) except Exception as error: # noqa: BLE001 - history must never fail a route print( f"[router] routing history unreadable, deciding without it: " f"{type(error).__name__}: {error}", flush=True, ) return {} def _history(self): """The store, or None when this proxy exposes no database to write to.""" try: from litellm.proxy.proxy_server import prisma_client except ImportError: return None if prisma_client is None: return None if self._store is None: self._store = History(prisma_client.db, HISTORY_TABLE) return self._store async def _compare(self, router, data, candidates, text, request_class): """Answer this request from several routes and record which won. In ``serve`` mode the winning reply is handed back to the caller through ``mock_response``, so the fan-out is the answer path and no further call is made. In ``shadow`` mode the caller keeps the route already decided and only the verdict is kept. """ try: replies = await collect(router, data, candidates, SAMPLE_WIDTH) if len(replies) < 2: return winner = judge_answers(text, replies) except Exception as error: # noqa: BLE001 - a comparison must never fail a route print( f"[router] comparison abandoned, serving the decided route: " f"{type(error).__name__}: {error}", flush=True, ) return store = self._history() if store is not None: try: await store.record(request_class, winner, sorted(replies)) except Exception as error: # noqa: BLE001 - see above print( f"[router] comparison not recorded: " f"{type(error).__name__}: {error}", flush=True, ) if SAMPLE_MODE == "serve": data["model"] = winner data["mock_response"] = replies[winner] instance = SystemOneRouter()