Govern access to multiple AI models

Put one gateway function in front of AI Inference and a third-party provider to authenticate teams, route by policy, fall back, cache, and log every call.

Last updated

View as Markdown Agent setup

A platform team supports product teams that call AI models from several providers and from AI Inference. Each team holds its own provider keys and its own fallback code, so nobody can route a request to another model when a provider fails, stop a team that spends past its budget, or say which model answered a given request. This page puts a gateway built on Functions in front of every model. Applications call one endpoint with a team key, and the gateway authenticates the team, picks a model by policy, falls back to the next model when a call fails, caches responses the caller allows, and logs every call. The result is measured by requests answered despite a provider failure, spend per team within budget, the share of requests answered from cache, and audit coverage of model calls.

This use case does not cover building the applications or agents that call the models. For that, refer to Build AI agents.

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.
  • An HTTPS endpoint of your analytics or log platform that accepts POST requests, for the usage and audit logs.
  • The values of your own setup. This page uses 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 each value with yours in every step.

Required products

The gateway needsWhich meansProductDocumented in
One endpoint that holds the routing, fallback, and authentication logicA function on the gateway’s application, run on every path by one ruleFunctionsFunctions quickstart
Platform-hosted models among the routesAzion.AI.run with a model id from the AI Inference model catalogAI InferenceCall a model on AI Inference from a function
Team keys, the models each team may call, and whether a team is still within budgetOne record per team, read by key on every requestKV StoreKV Store API
Repeated prompts answered without a model call, where the caller allows itResponses stored and matched with the Cache APICacheCache a function’s response with the Cache API
A usage and audit record of every model callA log line per call, streamed from the Functions data sourceData StreamSend logs to an HTTP endpoint
One request investigated after the factThe function’s log lines for that requestReal-Time EventsReal-Time Events data sources
A rule that runs the functionThe Run Function behavior, which requires Application Accelerator on the applicationApplication AcceleratorHow Functions works

Reference architecture

This page builds the Multi-model AI gateway: one function in front of every model, with the policy, the keys, and the log in one place.

Read the diagram from the function outward. Every model call passes through it, so it is the one place where a team is identified, a model is chosen, and a call is recorded. The arrows to the models are ordered: the policy names a first route and the routes after it, and the function moves to the next one only when a call fails. The arrows to KV Store and the Cache API are reads that come before any model call, so a refused team or a stored answer costs no model call at all.

Dataflow

  1. A team application sends an OpenAI-compatible request to /v1/chat/completions, with its team key as a bearer token and a model alias in model rather than a model id.
  2. The gateway hashes the key and reads the team’s record from KV Store. An unknown key answers 401, a team over budget answers 429, and an alias the team may not call answers 403.
  3. When the caller allows caching, the gateway looks the request up in the cache and returns a stored response on a match.
  4. Otherwise the gateway calls the routes of the alias in order: a model on AI Inference first, then the third-party provider when the first call fails.
  5. The gateway writes one log line with the team, the alias, the model that answered, whether it fell back, the cache result, and the tokens used. Data Stream sends it to your analytics platform, and Real-Time Events holds the lines of each request for investigation.
  6. The gateway returns the model’s response, and stores it in the cache when the caller allowed it.

Components

  • Functions: runs authentication, routing, and fallback. The policy is code in the ai-gateway function, so a change to a route reaches every team at once, and no team holds a provider key or fallback logic of its own.
  • AI Inference: runs the platform-hosted models, called with Azion.AI.run and a model id. A call to it stays inside Azion and needs no provider key.
  • third-party LLM providers: the integrations that run external models. The function calls them with keys it reads from environment variables, so their latency and failures enter only the routes that name them.
  • Cache: holds responses to repeated prompts, through the Cache API of the function. A stored response is returned without a model call, only for requests whose caller allows it.
  • KV Store: holds the keys, quotas, and budgets: one record per team in the ai-gateway namespace, keyed by a hash of the team key, with the models the team may call and whether it is within budget. The function reads it on every request.
  • Data Stream: sends the usage and audit log lines the function writes, from the Functions data source, to your analytics platform, where spend per team is summed.
  • Real-Time Events: holds the function’s log lines grouped by request, for investigating one call after the fact.
  • application: the Platform Resource that is the gateway endpoint. A rule runs the function on its paths, so every application points at one domain.

Configure the gateway function

The gateway is one function. Every decision it applies is in the code, so a policy change is a code change that every team gets at once.

  • Teams send an alias, not a model id. general routes to Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 on AI Inference, and long-context to gpt-oss-20b, whose 131k-token context is the longest among the models AI Inference runs. Both fall back to the provider. An alias lets the platform team change the model behind it without a change in any team’s code. The gateway calls an AI Inference route as Call a model on AI Inference from a function describes, with these values: the route’s model id, and the team’s request body with stream set to false.
  • The key is stored as a hash. The record key is team: followed by the SHA-256 of the team key, so KV Store never holds a usable key. SHA-256 is one of the digests crypto.subtle.digest supports.
  • A failed call moves to the next route. A call that throws, returns no choices, answers with an error status, or runs past 30 seconds counts as failed. Thirty seconds bounds how long a team waits on one route before the next one runs, inside the 5-minute wall-clock limit of a function.
  • Caching is the caller’s decision. The gateway caches only requests that carry x-gateway-cache: allow, because only the team knows whether one answer can stand for every identical prompt. The gateway caches as Cache a function’s response with the Cache API describes, with the cache ai-gateway, the key https://ai-gateway.cache/<sha-256 of the alias and the request body>, and max-age=3600, one hour.
  • Every call writes one log line. The line is the usage and audit record: a team’s spend is the sum of its total_tokens, and the line names the model that answered.
  • The fallback-test alias proves the fallback. Its first route names a model id that does not exist, so every call to it falls back. Remove it after Verify the setup.

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

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};

Run the function with these values, following Functions quickstart:

  • Function instance: ai-gateway, with no Args.
  • Rule: a Request Phase rule named gateway - all paths, with 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. The Cache API is not defined under azion dev, so test the function once it is deployed.

Configure the team records

A team record is a JSON object under team:<sha256 of the team key> in the ai-gateway namespace: the team’s name, the aliases it may call, and its status. The gateway reads the record on every request, and admits a request only while status is active.

Budget enforcement reads that status rather than counting each call. KV Store accepts one write per second to the same key and offers no atomic increment, so a counter written on every request would lose updates under concurrent traffic. Spend is computed instead in your analytics platform from the total_tokens of each team’s log lines. When a team reaches its budget, an operator sets its status to blocked, and the gateway answers 429 to that team from then on.

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}

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.

Configure the usage and audit stream

Every model_call line is a usage and audit record, and Data Stream delivers the lines from the Functions data source to your analytics platform, where spend and fallbacks are summed per team.

Create the stream as Send logs to an HTTP endpoint shows, with these values:

  • Data Source: Functions. In the API, functions_console.
  • Template: Functions Event Collector, whose $log_message variable carries each line the function writes.
  • Option: Filter Workloads, with the gateway’s workload as the only chosen workload. A filter keeps this stream from deactivating the other streams on the account, which saving an active stream with sampling does.
  • Connector: Standard HTTP/HTTPS POST, with your analytics platform’s URL and the header it requires to accept the request.

The stream becomes active within one to two minutes of saving it. Your analytics platform then receives one model_call line per request, and one route_failed line per route that failed.

Verify the setup

Each check sends a chat request with the team key from the team records step.

  • A team reaches a model through an alias. 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.

  • A failed route falls back. 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.

  • A repeated prompt the caller allows is answered from cache. Send the general request twice with the header x-gateway-cache: allow. The second response carries x-gateway-cache: hit.

  • The gateway refuses what the policy refuses. 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.

  • Every call is logged. Your analytics platform receives a model_call line for each request above, with "team":"checkout-team". To read the lines of one request, open the Functions Console data source in Real-Time Events, where the ID variable groups the lines of a single request.

Remove the fallback-test alias from ROUTES and from the team record once the fallback check passes. A new rule takes a few minutes to propagate.

Measuring results

MetricWhere to read itWhat working looks like
Requests answered despite a provider failureThe model_call lines with "ok":true and "fallback":true, in your analytics platformEvery route failure in a route_failed line is followed by a successful model_call for the same request, unless every route failed
Spend per team within budgetThe sum of total_tokens per team over the budget period, in your analytics platformEach team’s sum stays under its budget, and a team that reaches it has status set to blocked
Share of requests answered from cacheThe model_call lines with "cache":"hit" over the lines with "cache":"hit" or "cache":"miss"Rises for teams that allow caching on prompts that repeat
Audit coverage of model callsThe count of model_call lines against the gateway’s invocations in the Functions tab of Real-Time Metrics, less the /admin/teams calls and the refused requestsThe two counts match, so every model call has a record

Best practices

  • Give each team its own key, and rotate it by replacing the record. The record key is the hash of the team key, so a new key is a new record. Set the old record to blocked once the team switches, because a KV Store key cannot be listed and a forgotten record stays readable.
  • Keep the provider key in an environment variable. The function reads PROVIDER_API_KEY with Azion.env.get(), so the key never appears in the code, and teams never hold it. A changed variable reaches the function only after it is redeployed.
  • Let teams opt in to caching per request. A cached answer is returned for every identical prompt, whichever team sends it. Caching a prompt whose answer must change, such as one that asks about the current state of an account, returns a stale answer for an hour.
  • Account for KV Store reads at the gateway’s request rate. Every request reads one team record, and KV Store includes 100,000 keys read per day before a charge applies. For the included amounts and the rate past them, refer to KV Store limits.

Guides in this use case

Cache a function's response with the Cache APIStores and matches the responses that a caller allows the gateway to cache.Call a model on AI Inference from a functionCalls the first route of an alias, a model on AI Inference, with Azion.AI.run.