# Wiring a local model router into Antigravity: An experiment with LiteLLM and Laya On 3 October, I gave a talk at GDG DevFest Melbourne called [*Token-Maxxing Is Not a Strategy: Getting More Value per Token on Google Cloud*](https://gdg-melbourne-devfest-2026.sessionize.com/session/1321310). One of the demos was a local model router sitting in front of a coding agent. This is how it works, how little code it actually takes, and the rough edges I ran into while building it. In August, Google [expanded Antigravity for enterprise customers](https://cloud.google.com/blog/products/ai-machine-learning/expanding-google-antigravity-for-enterprise-customers) by folding it into Gemini Enterprise. The big shift is billing: teams can finally run Antigravity under their central Google Cloud invoice rather than fiddling with standalone licenses or corporate credit cards. Combined with Application Default Credentials, budget caps, and pooled token quotas, the procurement friction is gone. The problem is that a coding agent is usually the heaviest token sink on an engineering team, and by default, it sends every task to whatever model you selected at launch. A one-character variable rename burns the same top-tier model as a multi-threaded deadlock. The obvious fix is dynamic routing: - Renames and typos drop down to **Gemini 3.5 Flash-Lite**. - Everyday edits, tests, and docstrings stay on **Gemini 3.8 Flash**. - Deep architectural refactors and gnarly debugging escalate to **Gemini 3.1 Pro**. The catch? You don't control the agent runtime. You can't just inject a callback into Antigravity’s execution loop. What you *can* intercept is network traffic. By slotting a [LiteLLM](https://docs.litellm.ai/) proxy between the Antigravity CLI (`agy`) and Vertex AI, we can let a tiny local model inspect the prompt and decide which Gemini tier should handle each ask. ![](https://cdn.hashnode.com/uploads/covers/6209aab37ee57a17fb59a894/cad0271e-1abe-48df-a0b8-23f4f9f9df8c.png align="center") ## Why run the model router locally? I borrowed the architectural pattern from TypeSafe's hosted "System One" classifier, [**Jev**](https://typesafe.ai/blog/introducing-system-one-models-and-jev): a tiny model scans the prompt and returns structured probabilities in a single fast pass before dispatching the real payload. LiteLLM even provides native support via `classifier_type: jev` in its [Auto Router](https://docs.litellm.ai/docs/auto_router). The catch is data boundaries. Hosted Jev runs on external infrastructure. For an enterprise coding agent, every routing decision carries snippets of proprietary code, directory trees, and internal bug reports. If your organization adopted Antigravity specifically to keep IP pinned inside your Google Cloud boundary, pinging an external SaaS just to route prompts defeats the purpose. That led me to [Laya](https://pypi.org/project/laya/), a 421M parameter Apache-2.0 decision model. It runs via `laya-serve` and grabs Apple silicon (`mps`) automatically on macOS: - **~160 ms overhead per decision:** Once warm, the latency (including the local HTTP roundtrip) is negligible compared to an upstream Gemini round-trip. The cold start took ~7 seconds, so the startup script handles warming it up. - **Zero data leakage:** The routing classification never leaves localhost. Only the final Gemini payload travels to Vertex AI inside your cloud project. - **Zero fine-tuning:** Tweaking routing behavior is just modifying plain-text classifier instructions. Conveniently, `laya-serve` implements Jev's `/v1/systemone` spec, making it drop-in compatible with tooling built for Jev. Hosted Jev typically clocks around 236–276 ms, so keeping it on-box actually shaved off latency. ## The implementation The entire router is contained in a single LiteLLM pre-call hook (`laya_`[`hook.py`](http://hook.py)). Because it subclasses `CustomLogger`, it intercepts calls across every dialect LiteLLM supports—including native Gemini endpoints (`streamGenerateContent`), OpenAI completions, and Anthropic messages. Laya evaluates two independent yes/no questions per incoming prompt: ```python QUESTIONS = { "trivial": { "type": "noul", "instructions": "Is this a small coding request, such as renaming something, fixing a typo, " "a one-line edit, or a quick factual question about the code?", }, "hard": { "type": "noul", "instructions": "Does this coding request need deep reasoning across many files, such as debugging " "a production issue, designing a feature, or changing the architecture?", }, } HARD_AT, TRIVIAL_AT = 0.5, 0.5 ``` If `hard ≥ 0.5`, LiteLLM upgrades the call to Pro. If `trivial ≥ 0.5`, it routes to Flash-Lite. Everything else defaults to Flash. The hook inspects the payload, asks Laya, and rewrites the target model in place: ```python async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): if data.get("model") not in ROUTED: # only route gemini-3.8-flash; let AGY background tasks pass return data text = last_human_text(data) # extract what the developer typed, ignore raw tool output if not text: return data try: tier, p, secs = await self.decide(text) # query Laya (memoised per ask) except Exception as e: # fail open: never break a workflow if the router hiccups print(f"[laya-route] Laya unavailable ({e!r}); keeping {data['model']}", flush=True) return data data["model"] = TIERS[tier] return data ``` Registering the hook in `litellm_config.yaml` is a one-liner alongside the Vertex model definitions: ```yaml litellm_settings: callbacks: laya_hook.proxy_handler_instance ``` A companion callback (`async_log_success_event`) logs a `[laya-cost]` entry after every call, reporting token consumption and running cost per task. ## Hooking into Antigravity Antigravity’s CLI (agy) has an undocumented flag for custom upstream routing: AGY_LLM_GATEWAY_URL. When set, agy diverts all traffic to that URL using standard Gemini wire format (POST /v1beta/models/:streamGenerateContent?alt=sse) and bypasses standard interactive login with "authenticated via gateway". (Note: Because this is undocumented, future CLI bumps might change or drop it. I verified this against AGY 1.2.14 and LiteLLM 1.103. The documented GOOGLE_GEMINI_BASE_URL likely behaves similarly if this flag ever changes.) A helper [start.sh](http://start.sh) boots laya-serve and LiteLLM, primes the model cache, and exports the proxy environment: ```bash gcloud auth application-default login # LiteLLM authenticates via ADC export GOOGLE_CLOUD_PROJECT= source start.sh # starts services, exports AGY_LLM_GATEWAY_* # In a separate pane: monitor routing decisions and live spend tail -F /tmp/litellm.log | grep --line-buffered -E "laya-(route|cost)" cd agy --model "Gemini 3.8 Flash (Low)" # point at Flash; the proxy handles dynamic upgrades/downgrades ``` From there, you interact with `agy` as usual while the proxy handles dispatch. During testing: - A quick rename scored `trivial: 0.90`, running 7 sub-requests on **Flash-Lite**. - Writing docstrings landed on **Flash** across 9 sub-requests. Tool calling worked cleanly on all tiers, including Flash-Lite. - A multi-file deadlock scenario scored `hard: 0.80` and was handed off to **Pro**. Since the proxy binds to `127.0.0.1`, I left LiteLLM’s master authentication key disabled. If you ever expose the port beyond loopback, turn on `master_key` in the config and set `AGY_LLM_GATEWAY_API_KEY`. ## Edge cases and lessons learned A few things broke during setup that are worth watching out for: - **Agent metadata drowns classifier signal.** `agy` doesn't pass raw developer input directly. It wraps prompts in `……`, appending timestamps, active buffers, and terminal state. Feeding that entire blob to Laya cratered classification accuracy: one complex prompt scored `hard: 0.70` on its own, but dropped to `0.22` when buried inside the wrapper—causing it to default to Flash and stall out across 10 unhelpful steps. A regex stripping everything outside `` solved it immediately. Always inspect raw agent wire payloads before tuning prompts. - **LiteLLM's Auto Router skips Gemini routes.** I originally tried LiteLLM’s built-in Auto Router targeting local Laya via `api_base`. It works fine over `/v1/chat/completions`, but `agy` sends native Gemini API payloads, where the router throws `strategy router marker… not a callable model`. A pre-call hook runs universally across all protocols, bypassing the issue. - **Small zero-shot models need binary choices.** Prompting Laya to pick between four tiers in a single prompt produced erratic results (both renames and deadlocks triggered `COMPLEX`). Breaking the logic into two focused binary questions stabilized the classifications. - **Keep one turn on one model.** A single user command triggers an agent loop of 7 to 10 sequential API calls, each carrying the same prompt text. Memoizing Laya's decision across the turn ensures the agent doesn't bounce between models mid-task, preserving Vertex prompt caching. Switching models on a new turn invalidates that cache—an operational trade-off to consider in cost modeling. - **Always fail open.** If the local Laya process crashes or hits latency spikes, the hook catches the exception and falls back to the requested default model (`gemini-3.8-flash`). A flaky router shouldn't break the developer's terminal. You don't need access to an agent's internal architecture to introduce smart routing. Dropping a proxy in front of the process and rewriting the payload via pre-call hooks lets a 421M parameter model handle the decision in ~160 ms on commodity laptop hardware—with zero prompt data leaving the machine. The hard part wasn't getting a small model to distinguish trivial edits from complex bugs. The actual work was handling prompt wrappers, ensuring cross-protocol proxy compatibility, and preserving cache affinity across multi-step agent loops. Keep in mind that this is a proof of concept tested against representative tasks, not a rigorous production benchmark. Before deploying something like this across an engineering org, you'll want an evaluation harness to measure classification error rates, verify whether token savings offset failed generation attempts, and replace the basic in-memory memoization cache with an eviction policy suited for shared infrastructure. The code is available on GitHub at [litellm-gemini-laya-router](https://github.com/derrickqin/litellm-gemini-laya-router). (The companion ADK agent demo from the DevFest talk is also up at [google-adk-laya-router](https://github.com/derrickqin/google-adk-laya-router)).