feat(ai-cache): add semantic (L2) cache layer (#13632)

Author: janiussyafiq

Makefile

@@ -403,6 +403,10 @@ install: runtime
 
 	$(ENV_INSTALL) -d $(ENV_INST_LUADIR)/apisix/plugins/ai-cache
 	$(ENV_INSTALL) apisix/plugins/ai-cache/*.lua $(ENV_INST_LUADIR)/apisix/plugins/ai-cache
+	$(ENV_INSTALL) -d $(ENV_INST_LUADIR)/apisix/plugins/ai-cache/embeddings
+	$(ENV_INSTALL) apisix/plugins/ai-cache/embeddings/*.lua $(ENV_INST_LUADIR)/apisix/plugins/ai-cache/embeddings
+	$(ENV_INSTALL) -d $(ENV_INST_LUADIR)/apisix/plugins/ai-cache/vector-search
+	$(ENV_INSTALL) apisix/plugins/ai-cache/vector-search/*.lua $(ENV_INST_LUADIR)/apisix/plugins/ai-cache/vector-search
 
 	$(ENV_INSTALL) -d $(ENV_INST_LUADIR)/apisix/plugins/ai-lakera-guard
 	$(ENV_INSTALL) apisix/plugins/ai-lakera-guard/*.lua $(ENV_INST_LUADIR)/apisix/plugins/ai-lakera-guard

apisix/plugins/ai-cache.lua

@@ -20,16 +20,21 @@ local schema     = require("apisix.plugins.ai-cache.schema")
 local key_mod    = require("apisix.plugins.ai-cache.key")
 local binding    = require("apisix.plugins.ai-protocols.binding")
 local redis_util = require("apisix.utils.redis")
+local semantic   = require("apisix.plugins.ai-cache.semantic")
 
 local ngx        = ngx
 local ngx_null   = ngx.null
 local ipairs     = ipairs
+local pcall      = pcall
 local concat     = table.concat
+local str_format = string.format
 
-local CACHE_STATUS_HEADER = "X-AI-Cache-Status"
-local CACHE_AGE_HEADER    = "X-AI-Cache-Age"
-local DEFAULT_TTL         = 3600
-local DEFAULT_MAX_BODY    = 1048576
+local CACHE_STATUS_HEADER     = "X-AI-Cache-Status"
+local CACHE_AGE_HEADER        = "X-AI-Cache-Age"
+local CACHE_SIMILARITY_HEADER = "X-AI-Cache-Similarity"
+local DEFAULT_TTL             = 3600
+local DEFAULT_MAX_BODY        = 1048576
+local DEFAULT_SEMANTIC_TTL    = 86400
 
 local _M = {
     version  = 0.1,
@@ -39,8 +44,31 @@ local _M = {
 }
 
 
+local function has_layer(conf, name)
+    local layers = conf.layers or { "exact" }
+    for _, l in ipairs(layers) do
+        if l == name then
+            return true
+        end
+    end
+    return false
+end
+
+
 function _M.check_schema(conf)
-    return core.schema.check(schema, conf)
+    core.utils.check_https({
+        "semantic.embedding.openai.endpoint",
+        "semantic.embedding.azure_openai.endpoint",
+    }, conf, _M.name)
+    local ok, err = core.schema.check(schema, conf)
+    if not ok then
+        return false, err
+    end
+    if conf.semantic and not has_layer(conf, "semantic") then
+        core.log.warn("ai-cache: 'semantic' is configured but not listed in ",
+                      "'layers'; the semantic (L2) cache is inactive")
+    end
+    return true
 end
 
 
@@ -53,12 +81,17 @@ local function release(conf, red)
 end
 
 
-local function serve_hit(conf, ctx, cached)
-    ctx.ai_cache_status = "HIT"
+local function serve_hit(conf, ctx, cached, similarity)
+    local status = "HIT"
+    ctx.ai_cache_status = status
     if conf.cache_headers ~= false then
-        core.response.set_header(CACHE_STATUS_HEADER, "HIT")
+        core.response.set_header(CACHE_STATUS_HEADER, status)
         local age = ngx.time() - (cached.created_at or ngx.time())
         core.response.set_header(CACHE_AGE_HEADER, age < 0 and 0 or age)
+        if similarity then
+            core.response.set_header(CACHE_SIMILARITY_HEADER,
+                                     str_format("%.4f", similarity))
+        end
     end
     core.response.set_header("Content-Type", "application/json")
     return core.response.exit(200, cached.body)
@@ -126,16 +159,56 @@ function _M.access(conf, ctx)
         ctx.ai_cache_status = "MISS"
         return
     end
-    release(conf, red)
-
     if res ~= nil and res ~= ngx_null then
         local cached = core.json.decode(res)
         if cached and cached.body then
+            release(conf, red)
             return serve_hit(conf, ctx, cached)
         end
         core.log.warn("ai-cache: discarding malformed cache entry for ", ctx.ai_cache_key)
     end
 
+    -- L1 miss -> L2 semantic lookup. Release the L1 connection before
+    -- embed_query()'s HTTP call so the pool isn't pinned across the embedding
+    -- round-trip; re-acquire for the vector search. pcall keeps throws fail-open.
+    if has_layer(conf, "semantic") and conf.semantic then
+        release(conf, red)
+
+        local ok, vec = pcall(semantic.embed_query, conf, ctx, body)
+        if not ok then
+            core.log.warn("ai-cache: semantic embed error, fail-open as MISS: ", vec)
+            vec = nil
+            -- prevent log() from scheduling a write with partial/bad state
+            ctx.ai_cache_embedding = nil
+        end
+
+        if vec then
+            local sred
+            sred, err = redis_util.new(conf)
+            if not sred then
+                core.log.warn("ai-cache: redis unavailable for semantic search, ",
+                              "fail-open as MISS: ", err)
+            else
+                local sok, hit = pcall(semantic.search, sred, conf, ctx, vec)
+                if not sok then
+                    sred:close()
+                    core.log.warn("ai-cache: semantic search error, fail-open as MISS: ", hit)
+                    ctx.ai_cache_embedding = nil
+                else
+                    release(conf, sred)
+                    if hit then
+                        return serve_hit(conf, ctx,
+                            { body = hit.body, created_at = hit.created_at }, hit.similarity)
+                    end
+                end
+            end
+        end
+
+        ctx.ai_cache_status = "MISS"
+        return
+    end
+
+    release(conf, red)
     ctx.ai_cache_status = "MISS"
 end
 
@@ -174,7 +247,8 @@ end
 -- The response-capturing phases (body_filter / log) run in contexts where
 -- cosockets are disabled, so the Redis write is deferred to a 0-delay timer
 -- (timers run in a light thread where cosockets are allowed).
-local function write_to_cache(premature, conf, cache_key, response_body)
+-- l2 (optional) = { partition, embedding, dim, fingerprint, ttl } for L2 write.
+local function write_to_cache(premature, conf, cache_key, response_body, l2)
     if premature then
         return
     end
@@ -192,6 +266,13 @@ local function write_to_cache(premature, conf, cache_key, response_body)
         core.log.warn("ai-cache: redis set failed: ", err)
         return
     end
+    if l2 then
+        l2.created_at = ngx.time()
+        local wok, werr = pcall(semantic.write, red, conf, l2, response_body)
+        if not wok then
+            core.log.warn("ai-cache: semantic write error: ", werr)
+        end
+    end
     release(conf, red)
 end
 
@@ -218,7 +299,21 @@ function _M.log(conf, ctx)
     local response_body = concat(buf, "", 1, buf.n)
 
     local cache_key = key_mod.build(conf, ctx, ctx.ai_cache_fingerprint)
-    local ok, err = ngx.timer.at(0, write_to_cache, conf, cache_key, response_body)
+
+    -- Build the L2 doc from ctx fields stashed by semantic.embed_query(); the
+    -- embedding is only set on a successful embed, so a nil check guards the write.
+    local l2
+    if has_layer(conf, "semantic") and ctx.ai_cache_embedding then
+        l2 = {
+            partition  = ctx.ai_cache_partition,
+            embedding  = ctx.ai_cache_embedding,
+            dim        = ctx.ai_cache_dim,
+            fingerprint = ctx.ai_cache_fingerprint,
+            ttl        = (conf.semantic and conf.semantic.ttl) or DEFAULT_SEMANTIC_TTL,
+        }
+    end
+
+    local ok, err = ngx.timer.at(0, write_to_cache, conf, cache_key, response_body, l2)
     if not ok then
         core.log.warn("ai-cache: failed to schedule cache write: ", err)
     end

apisix/plugins/ai-cache/embeddings/azure_openai.lua

@@ -0,0 +1,39 @@
+--
+-- 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.
+--
+local base = require("apisix.plugins.ai-cache.embeddings.base")
+
+local _M = {}
+
+-- get_embeddings(conf, text, httpc, ssl_verify) -> (vector_table, err)
+function _M.get_embeddings(conf, text, httpc, ssl_verify)
+    local req = { input = text }
+    if conf.dimensions then
+        req.dimensions = conf.dimensions
+    end
+    return base.fetch({
+        endpoint = conf.endpoint,
+        headers = {
+            ["Content-Type"] = "application/json",
+            ["api-key"] = conf.api_key,
+        },
+        request = req,
+        httpc = httpc,
+        ssl_verify = ssl_verify,
+    })
+end
+
+return _M

apisix/plugins/ai-cache/embeddings/base.lua

@@ -0,0 +1,59 @@
+--
+-- 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.
+--
+
+-- Shared request/parse core for the embeddings drivers. Each provider differs
+-- only in its request shape, auth header, and default endpoint; the HTTP call,
+-- status check, and data[1].embedding extraction live here so a parsing fix
+-- lands in one place.
+local core    = require("apisix.core")
+local type    = type
+
+local HTTP_OK = ngx.HTTP_OK
+
+local _M = {}
+
+-- fetch(opts) where opts = { endpoint, headers, request, httpc, ssl_verify }
+-- -> (vector_table, err)
+function _M.fetch(opts)
+    local payload, err = core.json.encode(opts.request)
+    if not payload then
+        return nil, "encode embeddings request: " .. (err or "")
+    end
+
+    local res
+    res, err = opts.httpc:request_uri(opts.endpoint, {
+        method = "POST",
+        headers = opts.headers,
+        body = payload,
+        ssl_verify = opts.ssl_verify,
+    })
+    if not res or not res.body then
+        return nil, "embeddings request failed: " .. (err or "")
+    end
+    if res.status ~= HTTP_OK then
+        return nil, "embeddings endpoint returned " .. res.status
+    end
+
+    local decoded = core.json.decode(res.body)
+    if not decoded or type(decoded.data) ~= "table" or type(decoded.data[1]) ~= "table"
+       or type(decoded.data[1].embedding) ~= "table" then
+        return nil, "malformed embeddings response"
+    end
+    return decoded.data[1].embedding
+end
+
+return _M

apisix/plugins/ai-cache/embeddings/openai.lua

@@ -0,0 +1,41 @@
+--
+-- 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.
+--
+local base = require("apisix.plugins.ai-cache.embeddings.base")
+
+local _M = {}
+
+_M.DEFAULT_ENDPOINT = "https://api.openai.com/v1/embeddings"
+
+-- get_embeddings(conf, text, httpc, ssl_verify) -> (vector_table, err)
+function _M.get_embeddings(conf, text, httpc, ssl_verify)
+    local req = { model = conf.model, input = text }
+    if conf.dimensions then
+        req.dimensions = conf.dimensions
+    end
+    return base.fetch({
+        endpoint = conf.endpoint or _M.DEFAULT_ENDPOINT,
+        headers = {
+            ["Content-Type"] = "application/json",
+            ["Authorization"] = "Bearer " .. conf.api_key,
+        },
+        request = req,
+        httpc = httpc,
+        ssl_verify = ssl_verify,
+    })
+end
+
+return _M