feat(ai-proxy-multi): add semantic load-balancing algorithm

[AlinsRan]

Description

This PR adds a semantic load-balancing algorithm to ai-proxy-multi, so a route can pick an LLM instance by what the request is asking for rather than by weight (roundrobin) or by a hash key (chash).

Problem

ai-proxy-multi today can spread traffic across several LLM instances, but the choice is blind to the request. In practice the instances behind one route are rarely interchangeable — a frontier model, a small cheap model and a domain-tuned model differ by an order of magnitude in cost and latency. Today the only ways to send the right prompt to the right model are:

• make the client choose, by sending a different model (leaks backend topology into every caller, and every model swap becomes a client-side change), or
• put each model on its own route and select by URI/header (the caller must already know which model it wants — the gateway is not deciding anything), or
• run a classifier service in front of the gateway (an extra hop, an extra deployment, and its own failure mode).

None of these let a single endpoint accept "model": "auto" and route on meaning.

Proposal

Let each instance describe, in natural language, the kind of prompt it should serve. The gateway embeds those descriptions once, embeds the incoming prompt, and picks the closest instance by cosine similarity.

                       ┌──────────────────────────────────┐
 POST /chat            │ ai-proxy-multi (algorithm=semantic)
 {"model":"auto",      │                                  │
  "messages":[…]}  ──► │  1. extract last user message    │
                       │  2. embed it  ───────────────────┼──► embeddings endpoint
                       │  3. cosine vs each instance's    │    (reference vectors
                       │     reference vectors            │     embedded once,
                       │  4. best score ≥ threshold ?     │     cached per conf version)
                       │       yes → that instance        │
                       │       no  → catchall instance    │
                       └──────────────┬───────────────────┘
                                      ▼
                             chosen LLM instance

Key properties:

• One embedding call per request. Instance examples are embedded in a single batch and cached in an lrucache keyed by route + plugin config version, so a config change invalidates immediately and a steady-state request only embeds the prompt.
• No vector database. There are at most a handful of reference vectors per route; cosine similarity is computed in-process.
• Fail-open. Semantic routing is an optimisation, never a new failure mode. Every error path — embedding endpoint unreachable, non-200, malformed or null vector, non-numeric component, missing index, dimension mismatch (e.g. the embedding model was changed underneath), lrucache error — falls back to the catchall instance (or the first instance if none is marked). The routing path cannot return 500 because embeddings are unavailable.

Use cases

1. Cost routing — send hard prompts to the expensive model, everything else to the cheap one.

The single most common motivation. Coding and reasoning go to a frontier model; translation and summarisation go to a small one; anything that matches neither still gets served.

{
  "balancer": { "algorithm": "semantic", "threshold": 0.5 },
  "embeddings": {
    "provider": "openai",
    "model": "text-embedding-3-small",
    "auth": { "header": { "Authorization": "Bearer $EMBEDDING_KEY" } }
  },
  "instances": [
    {
      "name": "reasoning", "provider": "openai", "weight": 1,
      "options": { "model": "gpt-4o" },
      "auth": { "header": { "Authorization": "Bearer $KEY" } },
      "examples": ["write a python function", "debug this stack trace", "explain this algorithm"]
    },
    {
      "name": "cheap", "provider": "openai", "weight": 1,
      "options": { "model": "gpt-4o-mini" },
      "auth": { "header": { "Authorization": "Bearer $KEY" } },
      "examples": ["translate this sentence", "summarize this paragraph"]
    },
    {
      "name": "default", "provider": "openai", "weight": 1,
      "options": { "model": "gpt-4o-mini" },
      "auth": { "header": { "Authorization": "Bearer $KEY" } },
      "catchall": true
    }
  ]
}

Clients always POST "model": "auto". Swapping gpt-4o for a newer model, or re-balancing what counts as "hard", is a gateway config change — no client ships a new build.

2. Domain specialisation — route to fine-tuned or self-hosted models per subject.

A support assistant fronted by one endpoint, backed by a billing-tuned model, a technical-troubleshooting model, and a general model as catchall. The client does not need to know the taxonomy, and adding a fourth domain is one more instance with a few examples.

3. Data-residency / compliance routing.

Prompts that look like they carry regulated content ("examples": ["my social security number is", "here is the patient's diagnosis"]) route to a self-hosted instance with a tight threshold; everything else goes to the public API via catchall. Note this is best-effort routing, not a DLP control — it complements, and does not replace, a content-inspection plugin.

4. Progressive model migration.

Point a new model at a narrow set of examples and give it a high per-instance threshold, so it only receives prompts it is confidently a match for. Widen the examples (or lower the threshold) as confidence grows, and drop the old instance when its share reaches zero. expose_scores shows the score distribution while tuning.

Configuration

Field	Default	Notes
balancer.algorithm	roundrobin	new value: semantic
balancer.threshold	—	global minimum cosine similarity
balancer.expose_scores	false	debug: emit X-AI-Semantic-Route / X-AI-Semantic-Scores response headers
instances[].examples	—	1–64 utterances; required unless catchall
instances[].threshold	—	overrides balancer.threshold for this instance
instances[].catchall	—	at most one; the fallback target; needs no examples
embeddings	—	required when algorithm=semantic
embeddings.provider	—	openai / azure-openai
embeddings.model	—	e.g. text-embedding-3-small
embeddings.endpoint	—	optional for openai; for azure-openai the full URL incl. deployment path + api-version
embeddings.auth	—	header / query; both are in encrypt_fields
embeddings.timeout	10000	ms
embeddings.ssl_verify	true

The embedding request is issued through the existing ai-provider layer (openai-embeddings capability), so it reuses the same endpoint resolution, authentication and HTTP transport as the chat path rather than a second HTTP client. It is a self-contained sidecar call: the client's own request headers are deliberately not forwarded to the embedding provider.

Debugging the routing decision

balancer.expose_scores (default false) makes the plugin report how each instance scored and which one it picked:

X-AI-Semantic-Route: code
X-AI-Semantic-Scores: code:0.8213,translate:0.1904

X-AI-Semantic-Scores lists every instance that has examples, highest score first — not just the winner, so you can see how close the runners-up came. The catchall has no examples, hence no similarity score, and does not appear.

When nothing clears the threshold, the pick becomes fallback while the scores still show how close each one got — which is exactly what you need to calibrate threshold:

X-AI-Semantic-Route: fallback
X-AI-Semantic-Scores: code:0.3057,translate:0.2884

If the headers are absent while expose_scores is on, the embedding step itself failed and the request was served by the catchall before any score was computed.

Since these headers reveal instance names to the caller, they are meant for tuning, not production. The same score list is always logged at warn on the fallback path, with no configuration and nothing exposed to clients:

semantic routing: no instance cleared threshold (scores: code:0.3057,translate:0.2884), falling back

Alternatives considered

• LLM-based intent classification (ask a small model to label the prompt). Higher accuracy on nuanced intents, but adds a full generation round-trip on every request — typically hundreds of ms and a per-request token cost, versus a single embedding call. Also gives a non-deterministic label, which is awkward to threshold.
• Vector database. Unnecessary at this cardinality: a route has a handful of instances and at most 64 examples each. An in-process cosine over a cached, normalized matrix avoids a new dependency and a network hop.
• Local embedding model in the gateway. Would remove the network call, but pulls a model runtime into the data plane. Deferred; the embeddings block already allows pointing at a self-hosted OpenAI-compatible embedding server.

Known limitation

semantic selects an instance but does not participate in health checks or fallback_strategy / retry: if the chosen instance's upstream then fails, that failure is returned to the client. It only falls back when routing is inconclusive (no instance clears its threshold) or embedding fails. Combining semantic selection with health-aware failover is a follow-up; the behaviour is stated in the schema description and the docs so it is not a surprise.

Tests

t/plugin/ai-proxy-semantic.t (similarity math, incl. zero-vector and empty-aggregate edges), t/plugin/ai-proxy-semantic-schema.t (6 negative schema cases), t/plugin/ai-proxy-semantic-routing.t (14 cases): routing by intent, catchall fallback, and one case per fail-open guard — malformed (null) vector, non-table 200 body, non-200 response, dimension mismatch, reference-embedding failure — plus multimodal content arrays and a regression test asserting the client's Cookie never reaches the embedding endpoint.

Which issue(s) this PR fixes:

N/A — new feature.

Checklist

• I have explained the need for this PR and the problem it solves
• I have explained the changes or the new features added to this PR
• I have added tests corresponding to this change
• I have updated the documentation to reflect this change (docs/en and docs/zh)
• I have verified that this change is backward compatible (existing roundrobin / chash routes are unaffected; every new field is optional and semantic must be opted into)

[membphis]

Is there a specific reason to log the full embedding response body here? If not, please avoid logging the entire body, since it may contain sensitive information. Logging only the status and a bounded, sanitized error summary would be safer.

[Copilot]

Pull request overview

This PR introduces a new semantic load-balancing algorithm for ai-proxy-multi, allowing APISIX to select an LLM instance based on semantic similarity between the request prompt and per-instance example prompts, using an embeddings sidecar call and cached reference vectors.

Changes:

• Add semantic vector math + an embeddings batch client, and integrate semantic instance selection into ai-proxy-multi with fail-open fallback behavior.
• Extend schema/docs to support balancer.algorithm=semantic, per-instance examples/threshold/catchall, and an embeddings configuration block (including encrypted auth fields).
• Add test coverage for semantic math, schema validation, routing behavior, multimodal prompt extraction, and “do not forward client headers to embeddings”.

Reviewed changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
t/plugin/ai-proxy-semantic.t	Unit tests for vector math helpers (normalize/dot/cosine/aggregate).
t/plugin/ai-proxy-semantic-schema.t	Schema validation tests for semantic config constraints.
t/plugin/ai-proxy-semantic-routing.t	End-to-end routing tests with mocked embeddings/LLM endpoints and fail-open cases.
docs/zh/latest/plugins/ai-proxy-multi.md	Chinese docs: add semantic algorithm config fields and usage examples.
docs/en/latest/plugins/ai-proxy-multi.md	English docs: add semantic algorithm config fields and usage examples.
apisix/plugins/ai-transport/http.lua	Add skip_client_headers support for sidecar calls (e.g., embeddings).
apisix/plugins/ai-proxy/semantic.lua	New pure-Lua vector math module used by semantic routing.
apisix/plugins/ai-proxy/schema.lua	Extend ai-proxy-multi schema with semantic/embeddings fields + encrypt_fields updates.
apisix/plugins/ai-proxy/embedding.lua	New embeddings batch client via ai-provider layer with bounded error handling.
apisix/plugins/ai-proxy-multi.lua	Implement semantic instance selection + reference-vector caching and prompt extraction.
apisix/plugins/ai-providers/base.lua	Wire skip_client_headers into header construction for provider requests.
apisix/plugins/ai-providers/azure-openai.lua	Add openai-embeddings capability path for Azure OpenAI provider.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[nic-6443]

Would it make sense to drop avg/min here and just keep max (ideally as the fixed behaviour, not a knob)?

max is what the closest prior art does — LiteLLM's auto-routing hardcodes MAX aggregation for exactly this reason: "a strong match on one keyword in a tier is not diluted by that tier's other utterances." avg reintroduces that dilution: an instance with one dead-on example plus a couple of loosely-related ones scores lower than an instance whose examples are all mediocre-but-consistent, which is usually the opposite of what you want for intent routing.

And since avg is the default, the out-of-the-box behaviour is the diluted one — a route that works great with two tight examples can start mis-routing the moment someone adds a third, broader example.

If there's a concrete use case for avg/min I'm happy to be wrong, but otherwise a single max (no enum, no default to reason about) is simpler and matches what people expect from semantic routing. What do you think?

[AlinsRan]

You are right, and the numbers back it up. examples are alternative phrasings of one intent, so the question is "does any of them match" — that is max. I ran the aggregation on a concrete case:

	A: one dead-on example + two broader ones (1.0 / 0.5 / 0.4)	B: uniformly mediocre (0.7 / 0.7)	picked
avg	0.633	0.700	B — despite A having a perfect match
max	1.000	0.700	A
min	0.400	0.700	B

And your "adding a third example breaks the route" point reproduces exactly: avg drops 0.750 → 0.633 when a broader example is appended, so a route that cleared a 0.7 threshold stops clearing it — for max nothing changes.

min is worse still: it demands that every phrasing match, which is nonsense for alternatives.

So there is no use case I can defend for either, and avg-as-default meant the out-of-the-box behaviour was the diluted one. Dropped the knob entirely in 9c58718 — no enum, no default: an instance is scored by its best-matching example. semantic.aggregate(scores, method) is now just semantic.max(scores).

[nic-6443]

Heads up that this /embeddings path probably won't take effect for azure-openai. In ai-providers/base.lua, build_request resolves the path as parsed_url.path or opts.target_path — so once the endpoint URL carries a path (which it must for azure, e.g. .../openai/deployments/{deployment}), parsed_url.path wins and this capability path is dropped. Azure also needs the ?api-version= query param, which nothing on this path adds.

The doc example for embeddings.endpoint stops at /openai/deployments/{deployment} (no /embeddings, no api-version), so following the docs a semantic route with provider: azure-openai would POST to the deployment root and likely 404. All the routing tests use provider: openai, so this branch isn't exercised.

Could you either (a) require the azure embeddings.endpoint to include the full /embeddings?api-version=... and document that, or (b) construct the azure embeddings URL the same way the chat path does — and add one azure embeddings test so it doesn't silently regress? openai works today, but azure looks unverified.

[AlinsRan]

Confirmed — this was broken, thanks for catching it. build_request resolves parsed_url.path or opts.target_path, and an Azure endpoint always carries a path, so the capability /embeddings was dead code. Worse, my doc example stopped at the deployment root, so following it would have POSTed there and 404'd — and since every routing test used provider: openai, nothing exercised this.

Went with (a), matching the convention the azure chat path and ai-cache-semantic already use (full URL, api-version in the query):

• azure-openai now declares the capability with no default path (["openai-embeddings"] = {}), so a path-less endpoint fails loudly in build_request instead of silently building a wrong URL.
• check_schema rejects an azure embeddings.endpoint without a deployment path — otherwise the route would fail open to the catchall on every request, i.e. silently never work.
• Docs (en + zh) and the schema description now show the full URL: https://{resource}.openai.azure.com/openai/deployments/{deployment}/embeddings?api-version=2024-02-01.
• Added an end-to-end test that drives a real azure-shaped URL and asserts the request lands on the deployment path with api-version, plus a negative schema test.