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

Add a `semantic` balancer algorithm to `ai-proxy-multi` that picks an LLM
instance by the semantic similarity between the request prompt and each
instance's configured `examples`, instead of by weight or hash.

Each non-catchall instance declares `examples` describing the prompts it
should serve. On the request path the plugin embeds those examples and the
incoming prompt via a configurable OpenAI-compatible embedding endpoint,
computes cosine similarity, aggregates per-instance scores, and routes to
the highest-scoring instance that clears its threshold. When no instance
clears its threshold — or embedding fails for any reason — the request
fails open to the `catchall` instance (else the first instance), so a
request always has a target and never 500s on the routing path.

Reference vectors are embedded once per config version and cached; only the
prompt is embedded per request. No vector database is required.

New config: `balancer.algorithm=semantic`, `balancer.threshold`,
`balancer.aggregation`, `balancer.expose_scores`, per-instance `examples` /
`threshold` / `catchall`, and an `embeddings` block.

Author: AlinsRan

apisix/plugins/ai-providers/azure-openai.lua

@@ -32,6 +32,7 @@ return require("apisix.plugins.ai-providers.base").new(
                 path = "/completions",
                 rewrite_request_body = rewrite_chat_request_body,
             },
+            ["openai-embeddings"] = { path = "/embeddings" },
         },
     }
 )

apisix/plugins/ai-providers/base.lua

@@ -188,7 +188,8 @@ function _M.build_request(self, ctx, conf, request_body, opts)
             .. "includes a path", 400
     end
 
-    local headers = transport_http.construct_forward_headers(auth.header or {}, ctx)
+    local headers = transport_http.construct_forward_headers(auth.header or {}, ctx,
+                                                             opts.skip_client_headers)
     if opts.host_header then
         headers["Host"] = opts.host_header
     end

apisix/plugins/ai-proxy-multi.lua

@@ -33,12 +33,16 @@ local ngx_now = ngx.now
 
 local require = require
 local pcall = pcall
+local error = error
+local tostring = tostring
 local ipairs = ipairs
 local type = type
 local string = string
 local url = require("socket.url")
 
 local priority_balancer = require("apisix.balancer.priority")
+local semantic = require("apisix.plugins.ai-proxy.semantic")
+local embedding = require("apisix.plugins.ai-proxy.embedding")
 local endpoint_regex = "^(https?)://([^:/]+):?(%d*)/?.*$"
 
 local pickers = {}
@@ -48,6 +52,11 @@ local lrucache_server_picker = core.lrucache.new({
 local lrucache_health_status = core.lrucache.new({
     ttl = 300, count = 256
 })
+-- Keyed by route + conf version, so config changes invalidate immediately;
+-- the long ttl just avoids re-embedding references on unchanged config.
+local lrucache_semantic_vectors = core.lrucache.new({
+    ttl = 3600, count = 256
+})
 
 local plugin_name = "ai-proxy-multi"
 local _M = {
@@ -156,7 +165,41 @@ function _M.check_schema(conf)
         end
     end
 
-    return ok
+    if algo == "semantic" then
+        if not conf.embeddings then
+            return false, "must configure `embeddings` when balancer algorithm is semantic"
+        end
+        if conf.embeddings.provider == "azure-openai" and not conf.embeddings.endpoint then
+            return false, "must configure `embeddings.endpoint` when embeddings " ..
+                "provider is azure-openai"
+        end
+        local catchall_count = 0
+        for _, instance in ipairs(conf.instances) do
+            if instance.catchall then
+                catchall_count = catchall_count + 1
+            else
+                local has_example = false
+                if instance.examples then
+                    for _, ex in ipairs(instance.examples) do
+                        if type(ex) == "string" and ex ~= "" then
+                            has_example = true
+                            break
+                        end
+                    end
+                end
+                if not has_example then
+                    return false, "instance '" .. (instance.name or "?") ..
+                        "': must configure non-empty `examples` for the semantic " ..
+                        "algorithm unless `catchall` is set"
+                end
+            end
+        end
+        if catchall_count > 1 then
+            return false, "at most one instance may be marked `catchall`"
+        end
+    end
+
+    return true
 end
 
 
@@ -598,9 +641,159 @@ local function pick_target(ctx, conf, ups_tab)
 end
 
 
+local function extract_last_user_message()
+    local body = core.request.get_json_request_body_table()
+    if not body or type(body.messages) ~= "table" then
+        return nil
+    end
+    for i = #body.messages, 1, -1 do
+        local m = body.messages[i]
+        if type(m) == "table" and m.role == "user" then
+            local content = m.content
+            if type(content) == "string" then
+                return content
+            elseif type(content) == "table" then
+                -- multimodal content: concatenate the text parts so routing
+                -- still works for {type=text|image_url,...} arrays.
+                local parts = {}
+                for _, p in ipairs(content) do
+                    if type(p) == "table" and p.type == "text"
+                       and type(p.text) == "string" then
+                        parts[#parts + 1] = p.text
+                    end
+                end
+                if #parts > 0 then
+                    return table_concat(parts, " ")
+                end
+            end
+        end
+    end
+    return nil
+end
+
+
+-- Embed every instance's examples in one batch and group the normalized
+-- reference vectors by instance name. Raises on embedding failure so the
+-- lrucache below does not cache a bad result.
+local function build_instance_vectors(conf)
+    local texts = {}
+    local owners = {}
+    for _, inst in ipairs(conf.instances) do
+        if inst.examples then
+            for _, ex in ipairs(inst.examples) do
+                texts[#texts + 1] = ex
+                owners[#texts] = inst.name
+            end
+        end
+    end
+
+    local vecs, err = embedding.fetch(conf.embeddings, texts)
+    if not vecs then
+        error("failed to fetch reference embeddings: " .. tostring(err))
+    end
+
+    local by_instance = {}
+    for i, v in ipairs(vecs) do
+        local name = owners[i]
+        if name then
+            by_instance[name] = by_instance[name] or {}
+            core.table.insert(by_instance[name], semantic.normalize(v))
+        end
+    end
+    return by_instance
+end
+
+
+-- Guaranteed fallback: catchall instance if configured, else the first
+-- instance. Never fails, so a request always has a target.
+local function semantic_fallback(conf)
+    for _, inst in ipairs(conf.instances) do
+        if inst.catchall then
+            return inst.name, inst
+        end
+    end
+    local inst = conf.instances[1]
+    return inst.name, inst
+end
+
+
+local function pick_semantic_instance(ctx, conf)
+    local version = plugin.conf_version(conf)
+    local ok, by_instance = pcall(lrucache_semantic_vectors,
+                                  ctx.matched_route.key .. "#semantic", version,
+                                  build_instance_vectors, conf)
+    if not ok or not by_instance then
+        core.log.warn("semantic routing: ", by_instance, ", falling back")
+        return semantic_fallback(conf)
+    end
+
+    local prompt = extract_last_user_message()
+    if not prompt then
+        core.log.warn("semantic routing: no user message found, falling back")
+        return semantic_fallback(conf)
+    end
+
+    local qvecs, err = embedding.fetch(conf.embeddings, { prompt })
+    if not qvecs or not qvecs[1] then
+        core.log.warn("semantic routing: query embedding failed: ", err, ", falling back")
+        return semantic_fallback(conf)
+    end
+    local qvec = semantic.normalize(qvecs[1])
+    local qdim = #qvec
+
+    local aggregation = conf.balancer.aggregation or "avg"
+    local ranked = {}
+    for _, inst in ipairs(conf.instances) do
+        local refs = by_instance[inst.name]
+        if refs then
+            local scores = {}
+            for _, rv in ipairs(refs) do
+                -- guard against dimension drift (e.g. embedding model changed):
+                -- mismatched vectors would make dot() error, so fail open instead.
+                if #rv ~= qdim then
+                    core.log.warn("semantic routing: embedding dimension mismatch ",
+                                  "(query ", qdim, " vs reference ", #rv, "), falling back")
+                    return semantic_fallback(conf)
+                end
+                scores[#scores + 1] = semantic.dot(qvec, rv)
+            end
+            core.table.insert(ranked, {
+                name = inst.name,
+                score = semantic.aggregate(scores, aggregation),
+            })
+        end
+    end
+    core.table.sort(ranked, function(a, b) return a.score > b.score end)
+
+    -- Highest score first; pick the first instance that clears its own threshold
+    -- (per-instance override, else the global balancer.threshold).
+    for _, cand in ipairs(ranked) do
+        local inst = get_instance_conf(conf.instances, cand.name)
+        local thr = inst.threshold or conf.balancer.threshold or 0
+        if cand.score >= thr then
+            core.log.info("semantic routing picked instance: ", cand.name,
+                          ", score: ", cand.score)
+            return cand.name, inst
+        end
+    end
+
+    -- Only on the fallback path: surface why nothing matched. Cheap, because
+    -- this runs once per unmatched request rather than on every request.
+    local parts = {}
+    for _, c in ipairs(ranked) do
+        parts[#parts + 1] = c.name .. ":" .. string.format("%.4f", c.score)
+    end
+    core.log.warn("semantic routing: no instance cleared threshold (scores: ",
+                  table_concat(parts, ","), "), falling back")
+    return semantic_fallback(conf)
+end
+
+
 local function pick_ai_instance(ctx, conf, ups_tab)
     local instance_name, instance_conf, err
-    if #conf.instances == 1 then
+    if conf.balancer and conf.balancer.algorithm == "semantic" then
+        instance_name, instance_conf = pick_semantic_instance(ctx, conf)
+    elseif #conf.instances == 1 then
         instance_name = conf.instances[1].name
         instance_conf = conf.instances[1]
     else

apisix/plugins/ai-proxy/embedding.lua

@@ -0,0 +1,143 @@
+--
+-- Licensed to the Apache Software Foundation (ASF) under one or more
+-- contributor license agreements.  See the NOTICE file distributed with
+-- this work for additional information regarding copyright ownership.
+-- The ASF licenses this file to You under the Apache License, Version 2.0
+-- (the "License"); you may not use this file except in compliance with
+-- the License.  You may obtain a copy of the License at
+--
+--     http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing, software
+-- distributed under the License is distributed on an "AS IS" BASIS,
+-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+-- See the License for the specific language governing permissions and
+-- limitations under the License.
+--
+
+-- Batch embedding client for the semantic balancer. Routes the request through
+-- the shared ai-provider layer (apisix.plugins.ai-providers.*) so endpoint
+-- resolution, authentication and HTTP transport match the chat path, instead of
+-- re-implementing them here. One call embeds a list of texts (OpenAI-compatible
+-- /embeddings: input is an array), returning a vector per text. Keeps semantic
+-- routing free of a vector database.
+local core = require("apisix.core")
+local ipairs = ipairs
+local type = type
+local tostring = tostring
+local require = require
+local pcall = pcall
+
+local _M = {}
+
+local EMBEDDINGS_PROTOCOL = "openai-embeddings"
+
+
+-- Resolve a provider's embeddings endpoint path from its capabilities.
+local function embeddings_path(ai_provider, provider_name)
+    local cap = ai_provider.capabilities and ai_provider.capabilities[EMBEDDINGS_PROTOCOL]
+    if not cap then
+        return nil, "provider " .. tostring(provider_name) .. " does not support embeddings"
+    end
+    local path = cap.path
+    if type(path) == "function" then
+        path = path()
+    end
+    return path
+end
+
+
+-- Fetch embeddings for `texts` in a single batch request.
+-- Returns an array of vectors index-aligned with `texts`, or nil + err.
+function _M.fetch(conf, texts)
+    if not texts or #texts == 0 then
+        return {}
+    end
+
+    if not conf.model or conf.model == "" then
+        return nil, "embedding model is not configured"
+    end
+
+    local ok, ai_provider = pcall(require, "apisix.plugins.ai-providers." .. conf.provider)
+    if not ok then
+        return nil, "failed to load ai-provider: " .. tostring(conf.provider)
+    end
+
+    local target_path, perr = embeddings_path(ai_provider, conf.provider)
+    if not target_path then
+        return nil, perr
+    end
+
+    -- The embedding call is a self-contained sidecar request; use a throwaway
+    -- ctx so it never reads or mutates the in-flight request's ai_* fields.
+    -- ai_request_body_changed = true tells build_request to send our
+    -- { model, input } body instead of reusing the client's request body.
+    local ctx = { ai_request_body_changed = true }
+    local extra_opts = {
+        endpoint = conf.endpoint,
+        auth = conf.auth,
+        target_path = target_path,
+        -- This call carries its own credentials; never forward the client's
+        -- headers (Authorization, Cookie, ...) to the embedding provider.
+        skip_client_headers = true,
+    }
+    local req_conf = {
+        ssl_verify = conf.ssl_verify ~= false,
+        timeout = conf.timeout or 10000,
+        keepalive = true,
+    }
+
+    local status, raw_body, err = ai_provider:request(ctx, req_conf,
+        { model = conf.model, input = texts }, extra_opts)
+    if status ~= 200 then
+        if err then
+            return nil, "embedding request failed: " .. err
+        end
+        -- Keep the upstream body out of the returned error: that error is logged
+        -- once per failing request, and the body is untrusted third-party content
+        -- (it can be large, carry sensitive data, or forge log lines with
+        -- newlines). The status alone identifies the failure; the body is
+        -- available at info level when debugging.
+        core.log.info("embedding endpoint returned status ", status, ": ",
+                      raw_body or "")
+        return nil, "embedding endpoint returned status " .. tostring(status)
+    end
+
+    -- A 200 response whose body decodes to a scalar, boolean or null must not be
+    -- indexed (`data.data` on a number raises), and an object without `data` must
+    -- not have its body pasted into the error, which is logged per request.
+    local data, derr = core.json.decode(raw_body)
+    if type(data) ~= "table" or type(data.data) ~= "table" then
+        core.log.info("invalid embedding response: ", raw_body or "")
+        return nil, "invalid embedding response: " .. (derr or "unexpected shape")
+    end
+
+    local vectors = {}
+    for i, item in ipairs(data.data) do
+        -- OpenAI returns an `index`; fall back to positional order.
+        local idx = (type(item.index) == "number") and (item.index + 1) or i
+        local emb = item.embedding
+        if type(emb) ~= "table" or #emb == 0 then
+            return nil, "invalid embedding vector at index " .. (idx - 1)
+        end
+        for _, n in ipairs(emb) do
+            if type(n) ~= "number" then
+                return nil, "non-numeric embedding component at index " .. (idx - 1)
+            end
+        end
+        vectors[idx] = emb
+    end
+
+    -- Every input position must have produced a vector; a hole means a
+    -- dropped/duplicated index we must not silently route on. Checking each
+    -- slot is reliable where `#vectors` is not on a sparse table.
+    for i = 1, #texts do
+        if vectors[i] == nil then
+            return nil, "missing embedding for input index " .. (i - 1)
+        end
+    end
+    return vectors
+end
+
+
+return _M