Route model calls through a gateway function

Put one function in front of AI Inference and a third-party provider that authenticates teams, routes by alias, falls back, caches, and logs each call.

Last updated

View as Markdown Agent setup

You route the model calls of several teams through one function: it reads each team’s record from KV Store, calls the routes of a model alias in order, a model on AI Inference first and a third-party provider on failure, caches the responses a caller allows, and writes one log line per call. To call one model from a function first, refer to Call a model on AI Inference from a function.

Prerequisites

  • An application and a workload that serve the gateway’s domain, with Application Accelerator turned on, which the Run Function behavior requires. To create them, refer to Applications quickstart.
  • KV Store enabled on the account. The product is in Preview and is not enabled by default, so request access through Technical Support.
  • A personal token, for the KV Store call. To create one, refer to Personal tokens.
  • The Azion CLI, installed and authorized, to store the environment variables.
  • A third-party provider whose chat endpoint accepts the OpenAI chat completions format, with its URL, a model name, and an API key.

The Cache API is not defined under azion dev, so test the function once it is deployed.

The examples use ai-gateway for the KV Store namespace, the function, and the cache, general and long-context for the two model aliases the gateway offers, checkout-team for one team, and gateway.example.com for the domain. Replace them with your values.

Create the namespace

The function opens the ai-gateway namespace on every request, so the namespace exists before the function runs.

To create the namespace, send its name to the KV Store API:

Shell
01curl --request POST \
02 --url https://api.azion.com/v4/workspace/kv/namespaces \
03 --header 'Accept: application/json' \
04 --header 'Authorization: Token [TOKEN VALUE]' \
05 --header 'Content-Type: application/json' \
06 --data '{"name": "ai-gateway"}'

The API answers 201 with the namespace. A namespace cannot be renamed or deleted, so check the name before you send it:

JSON
01{
02 "name": "ai-gateway",
03 "created_at": "2026-01-01T12:00:00.000000",
04 "last_modified": "2026-01-01T12:00:00.000000"
05}

The account holds the empty ai-gateway namespace.

The Govern access to multiple AI models use case uses the values of this example.

Store the values the function reads

The provider key and the admin secret stay out of the code, as environment variables.

To store the values the function reads, run these commands with the Azion CLI. A key that contains key or secret is stored as a secret by default:

Shell
01azion create variables --key "PROVIDER_URL" --value "<provider-chat-completions-url>" --secret false
02azion create variables --key "PROVIDER_MODEL" --value "<provider-model-name>" --secret false
03azion create variables --key "PROVIDER_API_KEY" --value "<provider-api-key>"
04azion create variables --key "ADMIN_SECRET" --value "<admin-secret>"

The account holds the four variables the gateway function reads with Azion.env.get().

The Govern access to multiple AI models use case uses the values of this example.

Create the gateway function

The function answers /v1/chat/completions for the teams and /admin/teams for the operator, and every other path answers 404. The fallback-test alias names a model id that does not exist, so every call to it falls back to the provider.

Create a function named ai-gateway with this code. The provider call sends the key as Authorization: Bearer; change that header to the one your provider requires:

JavaScript
01const ROUTES = {
02 "general": [
03 { kind: "azion", model: "Qwen/Qwen3-30B-A3B-Instruct-2507-FP8" },
04 { kind: "provider" }
05 ],
06 "long-context": [
07 { kind: "azion", model: "gpt-oss-20b" },
08 { kind: "provider" }
09 ],
10 "fallback-test": [
11 { kind: "azion", model: "no-such-model" },
12 { kind: "provider" }
13 ]
14};
15const PROVIDER_TIMEOUT_MS = 30000;
16const CACHE_SECONDS = 3600;
17
18async function sha256(text) {
19 const digest = await crypto.subtle.digest("SHA-256", new TextEncoder().encode(text));
20 return [...new Uint8Array(digest)].map((b) => b.toString(16).padStart(2, "0")).join("");
21}
22
23async function callRoute(route, body) {
24 if (route.kind === "azion") {
25 const response = await Azion.AI.run(route.model, body);
26 if (!response?.choices?.length) throw new Error("no choices in the response");
27 return response;
28 }
29 const controller = new AbortController();
30 const timer = setTimeout(() => controller.abort(), PROVIDER_TIMEOUT_MS);
31 try {
32 const response = await fetch(Azion.env.get("PROVIDER_URL"), {
33 method: "POST",
34 headers: {
35 "Authorization": `Bearer ${Azion.env.get("PROVIDER_API_KEY")}`,
36 "Content-Type": "application/json"
37 },
38 body: JSON.stringify({ ...body, model: Azion.env.get("PROVIDER_MODEL") }),
39 signal: controller.signal
40 });
41 if (!response.ok) throw new Error(`provider answered ${response.status}`);
42 return await response.json();
43 } finally {
44 clearTimeout(timer);
45 }
46}
47
48async function putTeam(request, kv) {
49 if (request.method !== "PUT") {
50 return new Response("Method not allowed", { status: 405 });
51 }
52 if (request.headers.get("Authorization") !== `Bearer ${Azion.env.get("ADMIN_SECRET")}`) {
53 return new Response("Unauthorized", { status: 401 });
54 }
55 const { key, team, models, status } = await request.json();
56 if (!key || !team || !Array.isArray(models) || !["active", "blocked"].includes(status)) {
57 return Response.json({ error: "key, team, models, and status are required" }, { status: 400 });
58 }
59 await kv.put(`team:${await sha256(key)}`, { team, models, status });
60 return Response.json({ team, models, status });
61}
62
63export default {
64 async fetch(request, env, ctx) {
65 const url = new URL(request.url);
66 const kv = await Azion.KV.open("ai-gateway");
67 if (url.pathname === "/admin/teams") {
68 return putTeam(request, kv);
69 }
70 if (request.method !== "POST" || url.pathname !== "/v1/chat/completions") {
71 return new Response("Not found", { status: 404 });
72 }
73
74 const started = Date.now();
75 const auth = request.headers.get("Authorization") ?? "";
76 const teamKey = auth.startsWith("Bearer ") ? auth.slice(7) : "";
77 const team = teamKey ? await kv.get(`team:${await sha256(teamKey)}`, "json") : null;
78 if (!team) {
79 return Response.json({ error: "unknown team key" }, { status: 401 });
80 }
81 if (team.status !== "active") {
82 return Response.json({ error: "team budget reached" }, { status: 429 });
83 }
84
85 const { model: alias, ...body } = await request.json();
86 if (!ROUTES[alias] || !team.models.includes(alias)) {
87 return Response.json({ error: `model ${alias} is not allowed for ${team.team}` }, { status: 403 });
88 }
89 body.stream = false;
90
91 const cacheAllowed = request.headers.get("x-gateway-cache") === "allow";
92 const cache = await caches.open("ai-gateway");
93 const cacheKey = `https://ai-gateway.cache/${await sha256(JSON.stringify({ alias, body }))}`;
94 if (cacheAllowed) {
95 const hit = await cache.match(cacheKey);
96 if (hit) {
97 console.log(JSON.stringify({ event: "model_call", team: team.team, alias, route: null, ok: true, fallback: false, cache: "hit", total_tokens: 0, ms: Date.now() - started }));
98 return new Response(await hit.text(), {
99 headers: { "Content-Type": "application/json", "x-gateway-cache": "hit" }
100 });
101 }
102 }
103
104 let result = null;
105 let answered = null;
106 let attempts = 0;
107 for (const route of ROUTES[alias]) {
108 attempts++;
109 try {
110 result = await callRoute(route, body);
111 answered = route.model ?? "provider";
112 break;
113 } catch (error) {
114 console.log(JSON.stringify({ event: "route_failed", team: team.team, alias, route: route.model ?? "provider", error: String(error?.message ?? error) }));
115 }
116 }
117
118 console.log(JSON.stringify({
119 event: "model_call", team: team.team, alias, route: answered, ok: result !== null,
120 fallback: attempts > 1, cache: cacheAllowed ? "miss" : "off",
121 total_tokens: result?.usage?.total_tokens ?? 0, ms: Date.now() - started
122 }));
123 if (!result) {
124 return Response.json({ error: `every route of ${alias} failed` }, { status: 502 });
125 }
126
127 const text = JSON.stringify(result);
128 if (cacheAllowed) {
129 await cache.put(cacheKey, new Response(text, {
130 headers: { "Content-Type": "application/json", "cache-control": `max-age=${CACHE_SECONDS}` }
131 }));
132 }
133 return new Response(text, {
134 headers: { "Content-Type": "application/json", "x-gateway-route": answered }
135 });
136 },
137};

To create the function and its instance, follow Functions quickstart with the name ai-gateway, and name the instance ai-gateway, with no Args.

The application carries an ai-gateway instance that authenticates a team, routes its request by alias, and logs the call.

The Govern access to multiple AI models use case uses the values of this example.

Run the function on every path

Create a Request Phase rule as Add the rule that runs the function shows, with the name gateway - all paths, the criterion ${uri} starts with /, and the Run Function behavior selecting the ai-gateway instance.

The gateway answers on /v1/chat/completions and /admin/teams, and every other path answers 404. A new rule takes a few minutes to propagate.

Add a team record

Keys are written from a function rather than through the API, so the gateway’s /admin/teams path writes each record. To add a team that may call general and fallback-test, generate a random team key, then send it with the admin secret:

Shell
01curl -X PUT https://gateway.example.com/admin/teams \
02 -H 'Authorization: Bearer <admin-secret>' \
03 -H 'Content-Type: application/json' \
04 -d '{"key":"<checkout-team-key>","team":"checkout-team","models":["general","fallback-test"],"status":"active"}'

The gateway answers with the record it stored, without the key:

{"team":"checkout-team","models":["general","fallback-test"],"status":"active"}

Hand the team key to the team once, because the gateway keeps only its hash. To block the team, send the same request with "status":"blocked", and with "status":"active" to admit it again.

The ai-gateway namespace holds the checkout-team record under the hash of its key.

The Govern access to multiple AI models use case uses the values of this example.

Confirm the gateway routes, falls back, and refuses

Each check sends a chat request with the team key of Add a team record.

To send a request to general:

Shell
01curl -s -i -X POST https://gateway.example.com/v1/chat/completions \
02 -H 'Authorization: Bearer <checkout-team-key>' \
03 -H 'Content-Type: application/json' \
04 -d '{"model":"general","max_tokens":200,"messages":[{"role":"user","content":"Name three European capitals."}]}'

The response carries x-gateway-route: Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 and a chat.completion object whose generated text sits at choices[0].message.content.

Send the same request with "model":"fallback-test". The response carries x-gateway-route: provider, and the log lines for that request hold one route_failed line for no-such-model, then a model_call line with "fallback":true.

Send the general request twice with the header x-gateway-cache: allow. The second response carries x-gateway-cache: hit.

A request with no Authorization header answers 401. A request for long-context, which checkout-team may not call, answers 403. After you set the team’s status to blocked, any request with its key answers 429.

The gateway routes each team’s request by alias, falls back when a route fails, and refuses what the team record does not allow. Remove the fallback-test alias from ROUTES and from the team record once the fallback check passes.

These checks confirm the Govern access to multiple AI models use case.

Next steps

Cache a function's response with the Cache APIBuild a cache key for a request, and delete an entry after a write.Govern access to multiple AI modelsThe design this gateway serves: the aliases, the fallback timeout, the cache, and the budget model, with the reason for each.