From d0e01b294057f33c4697c2a032c641db29663d44 Mon Sep 17 00:00:00 2001 From: Kevin Veen-Birkenbach Date: Thu, 24 Sep 2026 23:35:47 +0200 Subject: [PATCH] feat(svc-ai-litellm): compare routes on an interval and remember which won The router settings become one mapping. Nine scalars sharing a router_ prefix were a dict spelled out in key names, and the sampling section about to arrive would have added five more. Carried along: fifteen lookups in vars, three variant overrides, four test files, and the two rejection messages that name the settings path, which would otherwise have pointed a reader at a key that no longer exists. On a configurable trigger, seconds or requests or probability, the same request goes to several eligible routes at once. jeff is shown the replies rather than the routes that wrote them, with the aliases hidden behind option labels so the verdict rests on the text. In shadow mode the caller keeps the route already decided; in serve mode the winning reply is handed back through mock_response, so the fan-out is the answer path and no further call follows. Each comparison writes one row per offered route to the proxy's own database, grouped by a request class of token band and capability demands. Grouping finer would spread the evidence too thin to decide, coarser would let a verdict about a vision request speak for a plain chat. A route that raises is left out rather than recorded as a loss: failing to answer is a fact about the moment. The decider now also sees the price of this request, the measured rate, the quota and prompt caching, plus how each route fared before. 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. Sampling ships off. Every history and comparison path is caught wide and logged: neither may ever fail a route. Co-Authored-By: Claude Opus 5 (1M context) --- roles/svc-ai-litellm/meta/services.yml | 38 +- roles/svc-ai-litellm/meta/variants.yml | 12 +- .../templates/router_hook.py.j2 | 348 ++++++++++++++++-- roles/svc-ai-litellm/vars/main.yml | 23 +- .../files/test_probe_routing.py | 2 +- .../templates/test_router_hook.py | 210 ++++++++++- 6 files changed, 568 insertions(+), 65 deletions(-) diff --git a/roles/svc-ai-litellm/meta/services.yml b/roles/svc-ai-litellm/meta/services.yml index da86de75ab..f3f91caa9d 100644 --- a/roles/svc-ai-litellm/meta/services.yml +++ b/roles/svc-ai-litellm/meta/services.yml @@ -15,21 +15,29 @@ litellm: name: litellm max_output_tokens: 512 upstream_timeout: 240 - router_alias: auto - router_remote_fallback: false - # nocheck: dynamic-flag - router_strategy: "{{ 'system_one' if (lookup('config', 'svc-ai-litellm', 'services.jeff.enabled') | bool) else 'weighted' }}" - router_weights: - locality: 1.0 - cost: 0.4 - speed: 0.25 - caching: 0.15 - throughput: 0.1 - router_min_chars_per_token: 2 - router_state_chars: 512 - router_speed_samples: 12 - router_speed_timeout: 60 - router_speed_remeasure: false + router: + alias: auto + remote_fallback: false + # nocheck: dynamic-flag + strategy: "{{ 'system_one' if (lookup('config', 'svc-ai-litellm', 'services.jeff.enabled') | bool) else 'weighted' }}" + min_chars_per_token: 2 + state_chars: 512 + weights: + locality: 1.0 + cost: 0.4 + speed: 0.25 + caching: 0.15 + throughput: 0.1 + speed: + samples: 12 + timeout: 60 + remeasure: false + sample: + enabled: false + mode: shadow + trigger: seconds + interval: 900 + width: 3 remote_models: - alias: openai/gpt-4o-mini model: openai/gpt-4o-mini diff --git a/roles/svc-ai-litellm/meta/variants.yml b/roles/svc-ai-litellm/meta/variants.yml index dbfee12039..e413c82991 100644 --- a/roles/svc-ai-litellm/meta/variants.yml +++ b/roles/svc-ai-litellm/meta/variants.yml @@ -19,7 +19,9 @@ postgres: shared: true litellm: - router_speed_samples: 1 + router: + speed: + samples: 1 remote_models: - alias: mock/deterministic provider: mock @@ -40,7 +42,9 @@ shared: true flavor: torch litellm: - router_speed_samples: 1 + router: + speed: + samples: 1 remote_models: - alias: mock/deterministic provider: mock @@ -64,7 +68,9 @@ mem_reservation: 512m mem_limit: 1g litellm: - router_speed_samples: 1 + router: + speed: + samples: 1 remote_models: - alias: mock/deterministic provider: mock diff --git a/roles/svc-ai-litellm/templates/router_hook.py.j2 b/roles/svc-ai-litellm/templates/router_hook.py.j2 index 50d25c808f..cb56a1cd60 100644 --- a/roles/svc-ai-litellm/templates/router_hook.py.j2 +++ b/roles/svc-ai-litellm/templates/router_hook.py.j2 @@ -25,8 +25,11 @@ Two polarities govern the filters, and mixing them up routes silently wrong: from __future__ import annotations +import asyncio import json +import random import re +import time import urllib.request from datetime import date @@ -42,8 +45,17 @@ JEFF_KEY = {{ LITELLM_JEFF_KEY | to_json }} JEFF_MODEL = {{ LITELLM_JEFF_MODEL | to_json }} JEFF_TIMEOUT = {{ LITELLM_JEFF_TIMEOUT | int }} JEFF_QUESTION = "route" +JEFF_VERDICT_QUESTION = "verdict" JEFF_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,}" @@ -251,7 +263,7 @@ def local_reason(data, deployments, catalog): if carries_secret(data): return "the request shows a structured secret" if not REMOTE_FALLBACK: - return "services.litellm.router_remote_fallback is off" + return "services.litellm.router.remote_fallback is off" return "" @@ -366,6 +378,129 @@ def decide_weighted(candidates, weights): )["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'}" + + +class Sampler: + """Whether this request is answered by several routes instead of one. + + Args: + trigger: ``seconds`` spaces comparisons in time, ``requests`` counts + them out, ``probability`` draws one in *interval*. + interval: seconds, requests or odds, read according to *trigger*. + clock: monotonic source, injected so a test does not wait. + draw: uniform source in [0, 1), injected for the same reason. + """ + + def __init__(self, trigger, interval, clock=time.monotonic, draw=random.random): + self._trigger = trigger + self._interval = max(int(interval), 1) + self._clock = clock + self._draw = draw + self._last = None + self._seen = 0 + + def due(self): + """True when this request should be compared across routes.""" + if self._trigger == "requests": + self._seen += 1 + if self._seen < self._interval: + return False + self._seen = 0 + return True + if self._trigger == "probability": + return self._draw() < 1.0 / self._interval + now = self._clock() + if self._last is not None and now - self._last < self._interval: + return False + self._last = now + return True + + +class History: + """What earlier comparisons said about each route, per request class. + + The store is the proxy's own database, so a verdict outlives a restart. A + proxy without one keeps nothing at all: a process-local tally read as + history would let a fresh container claim experience it never gathered. + """ + + DDL = ( + f"CREATE TABLE IF NOT EXISTS {HISTORY_TABLE} (" + "id BIGSERIAL PRIMARY KEY," + "decided_at TIMESTAMPTZ NOT NULL DEFAULT now()," + "signature TEXT NOT NULL," + "alias TEXT NOT NULL," + "won BOOLEAN NOT NULL)" + ) + INDEX = ( + f"CREATE INDEX IF NOT EXISTS {HISTORY_TABLE}_signature_idx " + f"ON {HISTORY_TABLE} (signature)" + ) + + def __init__(self, client): + self._client = client + self._ready = False + + async def _ensure(self): + if not self._ready: + await self._client.execute_raw(self.DDL) + await self._client.execute_raw(self.INDEX) + self._ready = True + + async def record(self, request_class, winner, offered): + """Write one row per offered route, marking which one won.""" + await self._ensure() + for alias in offered: + await self._client.execute_raw( + f"INSERT INTO {HISTORY_TABLE} (signature, alias, won) " + f"VALUES ($1, $2, $3)", + request_class, + alias, + alias == winner, + ) + + async def wins(self, request_class): + """``{alias: (won, seen)}`` for *request_class*, empty when untried.""" + await self._ensure() + rows = await self._client.query_raw( + f"SELECT alias, SUM(CASE WHEN won THEN 1 ELSE 0 END) AS won, " + f"COUNT(*) AS seen FROM {HISTORY_TABLE} " + f"WHERE signature = $1 GROUP BY alias", + request_class, + ) + return {row["alias"]: (int(row["won"]), int(row["seen"])) for row in rows or []} + + +def record_of(alias, tally): + """How *alias* fared before, as a clause, or empty when it has no record.""" + won, seen = tally.get(alias, (0, 0)) + if not seen: + return "" + return f"won {won} of {seen} comparisons like this" + + def headroom(entry, need): """How much context this route would leave unused on this request. @@ -379,17 +514,35 @@ def headroom(entry, need): return int(context) - (need["input_tokens"] + need["max_output"]) -def describe(entry, need): - """One line telling the decider what this route is, in its own terms.""" +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" - place = "runs in this cluster" if caps["local"] else "runs at a provider" room = headroom(entry, need) - fit = "unknown headroom" if room == float("inf") else f"{int(room)} tokens spare" - return f"{place}, context window {window}, {fit}" + 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_jeff(text, candidates, need): +def ask_jeff(text, candidates, need, tally=None): """Put the routing choice to jeff as one typed choice question. jeff answers typed questions rather than generating, so the candidates @@ -401,19 +554,39 @@ def ask_jeff(text, candidates, need): ValueError: when jeff 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, + JEFF_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"jeff chose {alias!r}, which was not among {sorted(criteria)}" + ) + return alias, float(answer.get("confidence") or 0.0) + + +def _systemone(state, question, instructions, criteria): + """Put one typed choice to jeff and return the answer block it sent. + + Raises: + OSError: when jeff cannot be reached. + ValueError: when jeff answers a different question than the one asked. + """ payload = { - "state": text[:JEFF_MAX_STATE_CHARS], + "state": state[:JEFF_MAX_STATE_CHARS], "model": JEFF_MODEL, "questions": { - JEFF_QUESTION: { + question: { "type": "choice", - "instructions": ( - "Which of these models should answer the request? Prefer " - "the smallest one that can serve it well." - ), - "criteria": { - entry["model_name"]: describe(entry, need) for entry in candidates - }, + "instructions": instructions, + "criteria": criteria, } }, } @@ -428,15 +601,35 @@ def ask_jeff(text, candidates, need): ) with urllib.request.urlopen(request, timeout=JEFF_TIMEOUT) as response: # noqa: S310 configured internal origin body = json.loads(response.read()) - answer = (body.get("answers") or {}).get(JEFF_QUESTION) or {} - alias = answer.get("choice") - offered = {entry["model_name"] for entry in candidates} - if alias not in offered: - raise ValueError(f"jeff chose {alias!r}, which was not among {sorted(offered)}") - return alias, float(answer.get("confidence") or 0.0) + answer = (body.get("answers") or {}).get(question) + if not answer: + raise ValueError(f"jeff returned no answer to {question!r}") + return answer -def decide_system_one(candidates, need, text): +def judge_answers(text, replies): + """The alias whose reply jeff calls the best answer to *text*. + + The options are the replies themselves, not the routes that produced them, + so the verdict rests on what came back rather than on what the catalog + promised. Aliases are never shown to the decider for the same reason. + """ + labels = {f"option-{index}": alias for index, alias in enumerate(sorted(replies))} + answer = _systemone( + text, + JEFF_VERDICT_QUESTION, + "Which of these answers serves the request best?", + {label: replies[alias][:JEFF_MAX_STATE_CHARS] for label, alias in labels.items()}, + ) + label = answer.get("choice") + if label not in labels: + raise ValueError( + f"jeff chose {label!r}, which was not among {sorted(labels)}" + ) + return labels[label] + + +def decide_system_one(candidates, need, text, tally=None): """The route jeff picks for this prompt. The weighted strategy ranks on standing preferences and stops reading the @@ -447,11 +640,11 @@ def decide_system_one(candidates, need, text): quietly substituting a different decider would make that answer mean something else. """ - alias, _confidence = ask_jeff(text, candidates, need) + alias, _confidence = ask_jeff(text, candidates, need, tally) return alias -def decide(candidates, need, text="", weights=None, strategy=None): +def decide(candidates, need, text="", weights=None, strategy=None, tally=None): """Dispatch to the configured strategy. Args: @@ -460,16 +653,44 @@ def decide(candidates, need, text="", weights=None, strategy=None): 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) + 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. @@ -479,6 +700,11 @@ class SystemOneRouter(CustomLogger): 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 @@ -509,10 +735,13 @@ class SystemOneRouter(CustomLogger): 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" + 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, request_text(data)) + data["model"] = decide(candidates, need, text, tally=tally) except (OSError, ValueError) as error: return ( f"'{ROUTER_ALIAS}' cannot route: the {ROUTER_STRATEGY} strategy " @@ -520,7 +749,70 @@ class SystemOneRouter(CustomLogger): 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) + 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() diff --git a/roles/svc-ai-litellm/vars/main.yml b/roles/svc-ai-litellm/vars/main.yml index 3bc8d0de98..b2e5e9922b 100644 --- a/roles/svc-ai-litellm/vars/main.yml +++ b/roles/svc-ai-litellm/vars/main.yml @@ -7,19 +7,24 @@ LITELLM_PORT: "{{ lookup('config', application_id, 'services.litellm.ports.inter LITELLM_MASTER_KEY: "{{ lookup('config', application_id, 'secrets.credentials.master_key') }}" LITELLM_MAX_OUTPUT_TOKENS: "{{ lookup('config', application_id, 'services.litellm.max_output_tokens') }}" LITELLM_UPSTREAM_TIMEOUT: "{{ lookup('config', application_id, 'services.litellm.upstream_timeout') }}" -LITELLM_ROUTER_ALIAS: "{{ lookup('config', application_id, 'services.litellm.router_alias') }}" -LITELLM_ROUTER_REMOTE_FALLBACK: "{{ lookup('config', application_id, 'services.litellm.router_remote_fallback') }}" -LITELLM_ROUTER_WEIGHTS: "{{ lookup('config', application_id, 'services.litellm.router_weights') }}" -LITELLM_ROUTER_STRATEGY: "{{ lookup('config', application_id, 'services.litellm.router_strategy') }}" +LITELLM_ROUTER_ALIAS: "{{ lookup('config', application_id, 'services.litellm.router.alias') }}" +LITELLM_ROUTER_REMOTE_FALLBACK: "{{ lookup('config', application_id, 'services.litellm.router.remote_fallback') }}" +LITELLM_ROUTER_WEIGHTS: "{{ lookup('config', application_id, 'services.litellm.router.weights') }}" +LITELLM_ROUTER_STRATEGY: "{{ lookup('config', application_id, 'services.litellm.router.strategy') }}" LITELLM_JEFF_URL: "{{ JEFF_BASE_LOCAL_URL }}" LITELLM_JEFF_KEY: "{{ lookup('config', 'svc-ai-jeff', 'secrets.credentials.api_key') }}" LITELLM_JEFF_MODEL: "{{ lookup('config', 'svc-ai-jeff', 'services.jeff.model_alias') }}" LITELLM_JEFF_TIMEOUT: "{{ lookup('config', 'svc-ai-jeff', 'services.jeff.request_timeout') }}" -LITELLM_ROUTER_MIN_CHARS_PER_TOKEN: "{{ lookup('config', application_id, 'services.litellm.router_min_chars_per_token') }}" -LITELLM_ROUTER_STATE_CHARS: "{{ lookup('config', application_id, 'services.litellm.router_state_chars') }}" -LITELLM_ROUTER_SPEED_SAMPLES: "{{ lookup('config', application_id, 'services.litellm.router_speed_samples') }}" -LITELLM_ROUTER_SPEED_TIMEOUT: "{{ lookup('config', application_id, 'services.litellm.router_speed_timeout') }}" -LITELLM_ROUTER_SPEED_REMEASURE: "{{ lookup('config', application_id, 'services.litellm.router_speed_remeasure') }}" +LITELLM_ROUTER_MIN_CHARS_PER_TOKEN: "{{ lookup('config', application_id, 'services.litellm.router.min_chars_per_token') }}" +LITELLM_ROUTER_STATE_CHARS: "{{ lookup('config', application_id, 'services.litellm.router.state_chars') }}" +LITELLM_ROUTER_SPEED_SAMPLES: "{{ lookup('config', application_id, 'services.litellm.router.speed.samples') }}" +LITELLM_ROUTER_SPEED_TIMEOUT: "{{ lookup('config', application_id, 'services.litellm.router.speed.timeout') }}" +LITELLM_ROUTER_SPEED_REMEASURE: "{{ lookup('config', application_id, 'services.litellm.router.speed.remeasure') }}" +LITELLM_ROUTER_SAMPLE_ENABLED: "{{ lookup('config', application_id, 'services.litellm.router.sample.enabled') }}" +LITELLM_ROUTER_SAMPLE_MODE: "{{ lookup('config', application_id, 'services.litellm.router.sample.mode') }}" +LITELLM_ROUTER_SAMPLE_TRIGGER: "{{ lookup('config', application_id, 'services.litellm.router.sample.trigger') }}" +LITELLM_ROUTER_SAMPLE_INTERVAL: "{{ lookup('config', application_id, 'services.litellm.router.sample.interval') }}" +LITELLM_ROUTER_SAMPLE_WIDTH: "{{ lookup('config', application_id, 'services.litellm.router.sample.width') }}" LITELLM_ROUTER_HOOK_PATH_HOST: "{{ [ lookup('container', application_id, 'directories.config'), 'router_hook.py' ] | path_join }}" LITELLM_SPEED_PATH_HOST: "{{ [ lookup('container', application_id, 'directories.config'), 'measured_speed.json' ] | path_join }}" LITELLM_CONFIG_PATH_HOST: "{{ [ lookup('container', application_id, 'directories.config'), 'config.yaml' ] | path_join }}" diff --git a/tests/unit/python/roles/svc-ai-litellm/files/test_probe_routing.py b/tests/unit/python/roles/svc-ai-litellm/files/test_probe_routing.py index 9e016a920b..39660b45b9 100644 --- a/tests/unit/python/roles/svc-ai-litellm/files/test_probe_routing.py +++ b/tests/unit/python/roles/svc-ai-litellm/files/test_probe_routing.py @@ -306,7 +306,7 @@ class TestRouteVerdict(unittest.TestCase): SHIPPED_MIN_CHARS_PER_TOKEN = load_yaml( PROJECT_ROOT / "roles/svc-ai-litellm" / ROLE_FILE_META_SERVICES -)["litellm"]["router_min_chars_per_token"] +)["litellm"]["router"]["min_chars_per_token"] class TestOversizedPrompt(unittest.TestCase): diff --git a/tests/unit/python/roles/svc-ai-litellm/templates/test_router_hook.py b/tests/unit/python/roles/svc-ai-litellm/templates/test_router_hook.py index bf21fde851..5658061adf 100644 --- a/tests/unit/python/roles/svc-ai-litellm/templates/test_router_hook.py +++ b/tests/unit/python/roles/svc-ai-litellm/templates/test_router_hook.py @@ -36,9 +36,11 @@ JEFF_URL = "http://jeff:8000" SHIPPED = load_yaml(PROJECT_ROOT / "roles/svc-ai-litellm" / ROLE_FILE_META_SERVICES)[ "litellm" ] -SHIPPED_WEIGHTS = SHIPPED["router_weights"] -SHIPPED_MIN_CHARS_PER_TOKEN = SHIPPED["router_min_chars_per_token"] -SHIPPED_STATE_CHARS = SHIPPED["router_state_chars"] +SHIPPED_ROUTER = SHIPPED["router"] +SHIPPED_WEIGHTS = SHIPPED_ROUTER["weights"] +SHIPPED_MIN_CHARS_PER_TOKEN = SHIPPED_ROUTER["min_chars_per_token"] +SHIPPED_STATE_CHARS = SHIPPED_ROUTER["state_chars"] +SHIPPED_SAMPLE = SHIPPED_ROUTER["sample"] CATALOG = { "gpt-4o-mini": { @@ -96,7 +98,17 @@ def _stub_litellm(): sys.modules["litellm.integrations.custom_logger"] = custom_logger -def load(*, remote_fallback=False, weights=None, strategy="weighted"): +def load( + *, + remote_fallback=False, + weights=None, + strategy="weighted", + sample_enabled=False, + sample_mode=None, + sample_trigger=None, + sample_interval=None, + sample_width=None, +): """The rendered hook, imported as the container would import it.""" _stub_litellm() env = Environment(undefined=StrictUndefined, autoescape=False) # noqa: S701 - Python source, not markup @@ -113,6 +125,19 @@ def load(*, remote_fallback=False, weights=None, strategy="weighted"): LITELLM_JEFF_KEY="jeff-test-key", LITELLM_JEFF_MODEL="jev-latest", LITELLM_JEFF_TIMEOUT=20, + LITELLM_ROUTER_SAMPLE_ENABLED=sample_enabled, + LITELLM_ROUTER_SAMPLE_MODE=SHIPPED_SAMPLE["mode"] + if sample_mode is None + else sample_mode, + LITELLM_ROUTER_SAMPLE_TRIGGER=SHIPPED_SAMPLE["trigger"] + if sample_trigger is None + else sample_trigger, + LITELLM_ROUTER_SAMPLE_INTERVAL=SHIPPED_SAMPLE["interval"] + if sample_interval is None + else sample_interval, + LITELLM_ROUTER_SAMPLE_WIDTH=SHIPPED_SAMPLE["width"] + if sample_width is None + else sample_width, ) module = types.ModuleType("router_hook_under_test") module.__file__ = str(Path(TEMPLATE).with_suffix("")) @@ -473,7 +498,7 @@ class TestLocalOnly(HookCase, unittest.TestCase): routes = [local("tiny", context=10), remote("big", context=99999)] verdict = self.route(self.hook, self.BIG, routes) self.assertIsInstance(verdict, str) - self.assertIn("router_remote_fallback", verdict) + self.assertIn("router.remote_fallback", verdict) def test_the_switch_lets_a_big_prompt_reach_a_remote_model(self) -> None: hook = load(remote_fallback=True) @@ -602,9 +627,10 @@ class JeffStub: confidence: the confidence it reports alongside the choice. """ - def __init__(self, choice, confidence=0.9) -> None: + def __init__(self, choice, confidence=0.9, question="route") -> None: self.choice = choice self.confidence = confidence + self.question = question self.seen: dict = {} def urlopen(self, request, timeout=None): @@ -619,7 +645,7 @@ class JeffStub: body = { "model": "gliformer-large-v1", "answers": { - "route": { + self.question: { "type": "choice", "choice": self.choice, "confidence": self.confidence, @@ -740,7 +766,7 @@ class TestShippedConfiguration(HookCase, unittest.TestCase): def test_the_shipped_strategy_reads_the_jeff_service_flag(self) -> None: self.assertIn( "services.jeff.enabled", - SHIPPED["router_strategy"], + SHIPPED_ROUTER["strategy"], "the strategy follows whether the System One service is deployed, " "and it reads that from the one flag that decides it", ) @@ -748,7 +774,7 @@ class TestShippedConfiguration(HookCase, unittest.TestCase): def test_both_strategy_names_are_reachable_from_the_expression(self) -> None: env = Environment(undefined=StrictUndefined, autoescape=False) # noqa: S701 - a config value, not markup env.filters["bool"] = bool - template = env.from_string(SHIPPED["router_strategy"]) + template = env.from_string(SHIPPED_ROUTER["strategy"]) self.assertEqual(template.render(lookup=lambda *_args: True), "system_one") self.assertEqual( template.render(lookup=lambda *_args: False), @@ -789,5 +815,171 @@ class TestShippedConfiguration(HookCase, unittest.TestCase): ) +class TestRequestClass(unittest.TestCase): + """History only helps when comparable requests land in the same class.""" + + def setUp(self) -> None: + self.hook = load() + + def sign(self, tokens, **flags): + need = dict.fromkeys(("vision", "tools", "schema", "audio", "documents"), False) + need.update(flags) + need["input_tokens"] = tokens + need["max_output"] = 0 + return self.hook.signature(need) + + def test_two_sizes_in_one_band_share_a_class(self) -> None: + self.assertEqual(self.sign(300), self.sign(1000)) + + def test_crossing_a_band_splits_the_class(self) -> None: + self.assertNotEqual(self.sign(1000), self.sign(1100)) + + def test_a_capability_demand_splits_the_class(self) -> None: + self.assertNotEqual(self.sign(300), self.sign(300, vision=True)) + + def test_the_largest_band_still_names_a_class(self) -> None: + self.assertTrue(self.sign(10**9).endswith(":plain")) + + +class TestSampler(unittest.TestCase): + """Comparisons cost a multiple of a request, so the trigger must hold.""" + + def setUp(self) -> None: + self.hook = load() + + def test_seconds_fires_once_and_then_waits(self) -> None: + now = [1000.0] + sampler = self.hook.Sampler("seconds", 900, clock=lambda: now[0]) + self.assertTrue(sampler.due()) + now[0] += 899 + self.assertFalse(sampler.due()) + now[0] += 2 + self.assertTrue(sampler.due()) + + def test_requests_counts_out_the_interval(self) -> None: + sampler = self.hook.Sampler("requests", 3) + self.assertEqual( + [sampler.due() for _ in range(6)], + [False, False, True, False, False, True], + ) + + def test_probability_draws_one_in_the_interval(self) -> None: + draws = iter([0.0, 0.5]) + sampler = self.hook.Sampler("probability", 4, draw=lambda: next(draws)) + self.assertTrue(sampler.due()) + self.assertFalse(sampler.due()) + + +class TestDescribedMetrics(unittest.TestCase): + """What the decider is told about a route, and what it is not told.""" + + def setUp(self) -> None: + self.hook = load() + + def describe(self, route, tally=None): + need = self.hook.summarize({"messages": [{"role": "user", "content": "hi"}]}) + entry = self.hook.eligible([route], need, ALIAS, CATALOG, TODAY)[0] + return self.hook.describe(entry, need, tally) + + def test_a_measured_rate_is_offered(self) -> None: + line = self.describe(local("home", context=9999, traits={"speed": 42.5})) + self.assertIn("42.5 tokens per second measured", line) + + def test_an_unmeasured_rate_is_left_out_rather_than_zeroed(self) -> None: + line = self.describe(local("home", context=9999)) + self.assertNotIn("tokens per second", line) + + def test_the_price_of_this_request_is_offered(self) -> None: + self.assertIn("costs about", self.describe(local("home", context=9999))) + + def test_an_earlier_verdict_is_offered(self) -> None: + line = self.describe(local("home", context=9999), {"home": (7, 9)}) + self.assertIn("won 7 of 9 comparisons like this", line) + + def test_a_route_without_a_record_claims_none(self) -> None: + line = self.describe(local("home", context=9999), {"other": (3, 3)}) + self.assertNotIn("comparisons", line) + + +class TestAnswerVerdict(unittest.TestCase): + """Judging the replies, not the promises the catalog made about them.""" + + def setUp(self) -> None: + self.hook = load() + + def test_the_aliases_are_hidden_from_the_decider(self) -> None: + stub = JeffStub("option-0", question="verdict") + with unittest.mock.patch("urllib.request.urlopen", stub.urlopen): + self.hook.judge_answers("q", {"cheap": "a", "pricey": "b"}) + criteria = stub.seen["body"]["questions"]["verdict"]["criteria"] + self.assertEqual(sorted(criteria), ["option-0", "option-1"]) + self.assertEqual(sorted(criteria.values()), ["a", "b"]) + + def test_the_label_maps_back_to_the_alias_that_wrote_it(self) -> None: + stub = JeffStub("option-1", question="verdict") + with unittest.mock.patch("urllib.request.urlopen", stub.urlopen): + winner = self.hook.judge_answers("q", {"cheap": "a", "pricey": "b"}) + self.assertEqual(winner, "pricey") + + def test_an_answer_outside_the_options_is_not_read_as_a_verdict(self) -> None: + stub = JeffStub("option-9", question="verdict") + with ( + unittest.mock.patch("urllib.request.urlopen", stub.urlopen), + self.assertRaises(ValueError), + ): + self.hook.judge_answers("q", {"cheap": "a", "pricey": "b"}) + + +class TestCollect(unittest.IsolatedAsyncioTestCase): + """A route that cannot answer is absent from the comparison, not a loser.""" + + def setUp(self) -> None: + self.hook = load() + + class _Router: + def __init__(self, replies): + self.replies = replies + + async def acompletion(self, **kwargs): + reply = self.replies[kwargs["model"]] + if isinstance(reply, Exception): + raise reply + return types.SimpleNamespace( + choices=[ + types.SimpleNamespace(message=types.SimpleNamespace(content=reply)) + ] + ) + + async def test_a_raising_route_is_left_out(self) -> None: + router = self._Router({"a": "alpha", "b": RuntimeError("upstream down")}) + replies = await self.hook.collect( + router, + {"messages": []}, + [{"model_name": "a"}, {"model_name": "b"}], + 2, + ) + self.assertEqual(replies, {"a": "alpha"}) + + async def test_the_width_bounds_the_fan_out(self) -> None: + router = self._Router({"a": "alpha", "b": "beta", "c": "gamma"}) + replies = await self.hook.collect( + router, + {"messages": []}, + [{"model_name": "a"}, {"model_name": "b"}, {"model_name": "c"}], + 2, + ) + self.assertEqual(sorted(replies), ["a", "b"]) + + async def test_an_empty_reply_is_not_offered_for_judging(self) -> None: + router = self._Router({"a": "", "b": "beta"}) + replies = await self.hook.collect( + router, + {"messages": []}, + [{"model_name": "a"}, {"model_name": "b"}], + 2, + ) + self.assertEqual(replies, {"b": "beta"}) + + if __name__ == "__main__": unittest.main()