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. 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 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 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.

Why run the model router locally?

I borrowed the architectural pattern from TypeSafe's hosted "System One" classifier, 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.

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, 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). 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:

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:

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:

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 boots laya-serve and LiteLLM, primes the model cache, and exports the proxy environment:

gcloud auth application-default login          # LiteLLM authenticates via ADC
export GOOGLE_CLOUD_PROJECT=<your-project-id>
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 <some repo>
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 <USER_REQUEST>…</USER_REQUEST><ADDITIONAL_METADATA>…, 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 <USER_REQUEST> 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. (The companion ADK agent demo from the DevFest talk is also up at google-adk-laya-router).