janiussyafiq commented on code in PR #13308: URL: https://github.com/apache/apisix/pull/13308#discussion_r3170882292
########## apisix/plugins/ai-cache/semantic.lua: ########## @@ -0,0 +1,180 @@ +-- +-- 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 apisix_redis = require("apisix.utils.redis") +local uuid = require("resty.jit-uuid") +local ffi = require("ffi") + +local ffi_new = ffi.new +local ffi_string = ffi.string +local ngx_time = ngx.time +local tostring = tostring +local tonumber = tonumber +local type = type + +local _M = {} + + +local function index_name(dim) + return "ai-cache-idx-" .. dim +end + + +local function key_prefix(dim) + return "ai-cache:l2:" .. dim .. ":" +end + +local function pack_vector(vec) + local n = #vec + local buf = ffi_new("float[?]", n) + for i = 0, n - 1 do + buf[i] = vec[i + 1] + end + return ffi_string(buf, n * 4) +end + +local index_ready = {} + +local function ensure_index(red, dim) + if index_ready[dim] then + return true + end + + local _, err = red["FT.CREATE"](red, + index_name(dim), + "ON", "HASH", + "PREFIX", "1", key_prefix(dim), + "SCHEMA", + "embedding", "VECTOR", "HNSW", "6", + "TYPE", "FLOAT32", + "DIM", tostring(dim), + "DISTANCE_METRIC", "COSINE", + "scope", "TAG", + "created_at", "NUMERIC" + ) + + if err and not err:find("already exists") then + return nil, "FT.CREATE failed: " .. err + end + + index_ready[dim] = true + return true +end + + +function _M.search(conf, scope_hash, embedding_vec, threshold) + local red, err = apisix_redis.new(conf) + if not red then + return nil, nil, err + end + + local ok, init_err = ensure_index(red, #embedding_vec) + if not ok then + red:set_keepalive(conf.redis_keepalive_timeout, conf.redis_keepalive_pool) + return nil, nil, init_err + end + + local binary_vec = pack_vector(embedding_vec) + + local query + if scope_hash == "" then + query = "*=>[KNN 1 @embedding $vec AS dist]" + else + query = "@scope:{" .. scope_hash .. "}=>[KNN 1 @embedding $vec AS dist]" + end + + local res, search_err = red["FT.SEARCH"](red, + index_name(#embedding_vec), + query, + "PARAMS", "2", "vec", binary_vec, + "SORTBY", "dist", "ASC", + "LIMIT", "0", "1", + "RETURN", "2", "response", "dist", + "DIALECT", "2" + ) Review Comment: fixed -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
