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()