wehub-resource-sync
454 lines
16 KiB
Lua
454 lines
16 KiB
Lua
local helpers = require("spec.helpers")
|
|
local cjson = require("cjson")
|
|
local pl_file = require("pl.file")
|
|
|
|
local PLUGIN_NAME = "ai-proxy"
|
|
|
|
for _, strategy in helpers.all_strategies() do
|
|
if strategy ~= "cassandra" then
|
|
describe(PLUGIN_NAME .. ": (access) [#" .. strategy .. "]", function()
|
|
local client
|
|
local MOCK_PORT
|
|
|
|
lazy_setup(function()
|
|
MOCK_PORT = helpers.get_available_port()
|
|
|
|
local bp = helpers.get_db_utils(strategy == "off" and "postgres" or strategy, nil, { PLUGIN_NAME })
|
|
|
|
-- set up huggingface mock fixtures
|
|
local fixtures = {
|
|
http_mock = {},
|
|
}
|
|
|
|
fixtures.http_mock.huggingface = [[
|
|
server {
|
|
server_name huggingface;
|
|
listen ]] .. MOCK_PORT .. [[;
|
|
|
|
default_type 'application/json';
|
|
|
|
location = "/v1/chat/completions" {
|
|
content_by_lua_block {
|
|
local pl_file = require "pl.file"
|
|
local json = require("cjson.safe")
|
|
|
|
local token = ngx.req.get_headers()["authorization"]
|
|
if token == "Bearer huggingface-key" then
|
|
ngx.req.read_body()
|
|
local body, err = ngx.req.get_body_data()
|
|
body, err = json.decode(body)
|
|
|
|
if err or (body.messages == ngx.null) then
|
|
ngx.status = 400
|
|
ngx.print(pl_file.read("spec/fixtures/ai-proxy/huggingface/llm-v1-chat/responses/bad_request.json"))
|
|
else
|
|
ngx.status = 200
|
|
ngx.print(pl_file.read("spec/fixtures/ai-proxy/huggingface/llm-v1-chat/responses/good.json"))
|
|
end
|
|
else
|
|
ngx.status = 401
|
|
ngx.print(pl_file.read("spec/fixtures/ai-proxy/huggingface/llm-v1-chat/responses/unauthorized.json"))
|
|
end
|
|
}
|
|
}
|
|
|
|
# completions is on the root of huggingface models
|
|
location = "/" {
|
|
content_by_lua_block {
|
|
local pl_file = require "pl.file"
|
|
local json = require("cjson.safe")
|
|
|
|
local token = ngx.req.get_headers()["authorization"]
|
|
if token == "Bearer huggingface-key" then
|
|
ngx.req.read_body()
|
|
local body, err = ngx.req.get_body_data()
|
|
body, err = json.decode(body)
|
|
|
|
if err or (body.prompt == ngx.null) then
|
|
ngx.status = 400
|
|
ngx.print(pl_file.read("spec/fixtures/ai-proxy/huggingface/llm-v1-completions/responses/bad_request.json"))
|
|
else
|
|
ngx.status = 200
|
|
ngx.print(pl_file.read("spec/fixtures/ai-proxy/huggingface/llm-v1-completions/responses/good.json"))
|
|
end
|
|
else
|
|
ngx.status = 401
|
|
ngx.print(pl_file.read("spec/fixtures/ai-proxy/huggingface/llm-v1-completions/responses/unauthorized.json"))
|
|
end
|
|
}
|
|
}
|
|
location = "/model-loading/v1/chat/completions" {
|
|
content_by_lua_block {
|
|
local pl_file = require "pl.file"
|
|
|
|
ngx.status = 503
|
|
ngx.print(pl_file.read("spec/fixtures/ai-proxy/huggingface/llm-v1-chat/responses/bad_response_model_load.json"))
|
|
}
|
|
}
|
|
location = "/model-timeout/v1/chat/completions" {
|
|
content_by_lua_block {
|
|
local pl_file = require "pl.file"
|
|
|
|
ngx.status = 504
|
|
ngx.print(pl_file.read("spec/fixtures/ai-proxy/huggingface/llm-v1-chat/responses/bad_response_timeout.json"))
|
|
}
|
|
}
|
|
}
|
|
]]
|
|
|
|
local empty_service = assert(bp.services:insert({
|
|
name = "empty_service",
|
|
host = "localhost", --helpers.mock_upstream_host,
|
|
port = 8080, --MOCK_PORT,
|
|
path = "/",
|
|
}))
|
|
|
|
-- 200 chat good with one option
|
|
local chat_good = assert(bp.routes:insert({
|
|
service = empty_service,
|
|
protocols = { "http" },
|
|
strip_path = true,
|
|
paths = { "/huggingface/llm/v1/chat/good" },
|
|
}))
|
|
bp.plugins:insert({
|
|
name = PLUGIN_NAME,
|
|
route = { id = chat_good.id },
|
|
config = {
|
|
route_type = "llm/v1/chat",
|
|
auth = {
|
|
header_name = "Authorization",
|
|
header_value = "Bearer huggingface-key",
|
|
},
|
|
model = {
|
|
name = "mistralai/Mistral-7B-Instruct-v0.2",
|
|
provider = "huggingface",
|
|
options = {
|
|
max_tokens = 256,
|
|
temperature = 1.0,
|
|
huggingface = {
|
|
use_cache = false,
|
|
wait_for_model = true,
|
|
},
|
|
upstream_url = "http://" .. helpers.mock_upstream_host .. ":" .. MOCK_PORT,
|
|
},
|
|
},
|
|
},
|
|
})
|
|
local completions_good = assert(bp.routes:insert({
|
|
service = empty_service,
|
|
protocols = { "http" },
|
|
strip_path = true,
|
|
paths = { "/huggingface/llm/v1/completions/good" },
|
|
}))
|
|
bp.plugins:insert({
|
|
name = PLUGIN_NAME,
|
|
route = { id = completions_good.id },
|
|
config = {
|
|
route_type = "llm/v1/completions",
|
|
auth = {
|
|
header_name = "Authorization",
|
|
header_value = "Bearer huggingface-key",
|
|
},
|
|
model = {
|
|
name = "mistralai/Mistral-7B-Instruct-v0.2",
|
|
provider = "huggingface",
|
|
options = {
|
|
max_tokens = 256,
|
|
temperature = 1.0,
|
|
huggingface = {
|
|
use_cache = false,
|
|
wait_for_model = true,
|
|
},
|
|
upstream_url = "http://" .. helpers.mock_upstream_host .. ":" .. MOCK_PORT,
|
|
},
|
|
},
|
|
},
|
|
})
|
|
-- 401 unauthorized
|
|
local chat_401 = assert(bp.routes:insert({
|
|
service = empty_service,
|
|
protocols = { "http" },
|
|
strip_path = true,
|
|
paths = { "/huggingface/llm/v1/chat/unauthorized" },
|
|
}))
|
|
bp.plugins:insert({
|
|
name = PLUGIN_NAME,
|
|
route = { id = chat_401.id },
|
|
config = {
|
|
route_type = "llm/v1/chat",
|
|
auth = {
|
|
header_name = "api-key",
|
|
header_value = "wrong-key",
|
|
},
|
|
model = {
|
|
name = "mistralai/Mistral-7B-Instruct-v0.2",
|
|
provider = "huggingface",
|
|
options = {
|
|
max_tokens = 256,
|
|
temperature = 1.0,
|
|
huggingface = {
|
|
use_cache = false,
|
|
wait_for_model = true,
|
|
},
|
|
upstream_url = "http://" .. helpers.mock_upstream_host .. ":" .. MOCK_PORT,
|
|
},
|
|
},
|
|
},
|
|
})
|
|
-- 401 unauthorized
|
|
local completions_401 = assert(bp.routes:insert({
|
|
service = empty_service,
|
|
protocols = { "http" },
|
|
strip_path = true,
|
|
paths = { "/huggingface/llm/v1/completions/unauthorized" },
|
|
}))
|
|
bp.plugins:insert({
|
|
name = PLUGIN_NAME,
|
|
route = { id = completions_401.id },
|
|
config = {
|
|
route_type = "llm/v1/completions",
|
|
auth = {
|
|
header_name = "api-key",
|
|
header_value = "wrong-key",
|
|
},
|
|
model = {
|
|
name = "mistralai/Mistral-7B-Instruct-v0.2",
|
|
provider = "huggingface",
|
|
options = {
|
|
max_tokens = 256,
|
|
temperature = 1.0,
|
|
huggingface = {
|
|
use_cache = false,
|
|
wait_for_model = true,
|
|
},
|
|
upstream_url = "http://" .. helpers.mock_upstream_host .. ":" .. MOCK_PORT,
|
|
},
|
|
},
|
|
},
|
|
})
|
|
-- 503 Service Temporarily Unavailable
|
|
local chat_503 = assert(bp.routes:insert({
|
|
service = empty_service,
|
|
protocols = { "http" },
|
|
strip_path = true,
|
|
paths = { "/huggingface/llm/v1/chat/bad-response/model-loading" },
|
|
}))
|
|
bp.plugins:insert({
|
|
name = PLUGIN_NAME,
|
|
route = { id = chat_503.id },
|
|
config = {
|
|
route_type = "llm/v1/chat",
|
|
auth = {
|
|
header_name = "api-key",
|
|
header_value = "huggingface-key",
|
|
},
|
|
model = {
|
|
name = "mistralai/Mistral-7B-Instruct-v0.2",
|
|
provider = "huggingface",
|
|
options = {
|
|
max_tokens = 256,
|
|
temperature = 1.0,
|
|
huggingface = {
|
|
use_cache = false,
|
|
wait_for_model = false,
|
|
},
|
|
upstream_url = "http://" .. helpers.mock_upstream_host .. ":" .. MOCK_PORT.."/model-loading",
|
|
},
|
|
},
|
|
},
|
|
})
|
|
-- 503 Service Timeout
|
|
local chat_503_to = assert(bp.routes:insert({
|
|
service = empty_service,
|
|
protocols = { "http" },
|
|
strip_path = true,
|
|
paths = { "/huggingface/llm/v1/chat/bad-response/model-timeout" },
|
|
}))
|
|
bp.plugins:insert({
|
|
name = PLUGIN_NAME,
|
|
route = { id = chat_503_to.id },
|
|
config = {
|
|
route_type = "llm/v1/chat",
|
|
auth = {
|
|
header_name = "api-key",
|
|
header_value = "huggingface-key",
|
|
},
|
|
model = {
|
|
name = "mistralai/Mistral-7B-Instruct-v0.2",
|
|
provider = "huggingface",
|
|
options = {
|
|
max_tokens = 256,
|
|
temperature = 1.0,
|
|
huggingface = {
|
|
use_cache = false,
|
|
wait_for_model = false,
|
|
},
|
|
upstream_url = "http://" .. helpers.mock_upstream_host .. ":" .. MOCK_PORT.."/model-timeout",
|
|
},
|
|
},
|
|
},
|
|
})
|
|
|
|
-- start kong
|
|
assert(helpers.start_kong({
|
|
-- set the strategy
|
|
database = strategy,
|
|
-- use the custom test template to create a local mock server
|
|
nginx_conf = "spec/fixtures/custom_nginx.template",
|
|
-- make sure our plugin gets loaded
|
|
plugins = "bundled," .. PLUGIN_NAME,
|
|
-- write & load declarative config, only if 'strategy=off'
|
|
declarative_config = strategy == "off" and helpers.make_yaml_file() or nil,
|
|
}, nil, nil, fixtures))
|
|
end)
|
|
|
|
lazy_teardown(function()
|
|
helpers.stop_kong()
|
|
end)
|
|
|
|
before_each(function()
|
|
client = helpers.proxy_client()
|
|
end)
|
|
|
|
after_each(function()
|
|
if client then
|
|
client:close()
|
|
end
|
|
end)
|
|
|
|
describe("huggingface llm/v1/chat", function()
|
|
it("good request", function()
|
|
local r = client:get("/huggingface/llm/v1/chat/good", {
|
|
headers = {
|
|
["content-type"] = "application/json",
|
|
["accept"] = "application/json",
|
|
},
|
|
body = pl_file.read("spec/fixtures/ai-proxy/huggingface/llm-v1-chat/requests/good.json"),
|
|
})
|
|
-- validate that the request succeeded, response status 200
|
|
local body = assert.res_status(200, r)
|
|
local json = cjson.decode(body)
|
|
|
|
-- check this is in the 'kong' response format
|
|
assert.equals(json.model, "mistralai/Mistral-7B-Instruct-v0.2")
|
|
assert.equals(json.object, "chat.completion")
|
|
|
|
assert.is_table(json.choices)
|
|
--print("json: ", inspect(json))
|
|
assert.is_string(json.choices[1].message.content)
|
|
assert.same(
|
|
" The sum of 1 + 1 is 2. This is a basic arithmetic operation and the answer is always the same: adding one to one results in having two in total.",
|
|
json.choices[1].message.content
|
|
)
|
|
end)
|
|
end)
|
|
describe("huggingface llm/v1/completions", function()
|
|
it("good request", function()
|
|
local r = client:get("/huggingface/llm/v1/completions/good", {
|
|
headers = {
|
|
["content-type"] = "application/json",
|
|
["accept"] = "application/json",
|
|
},
|
|
body = pl_file.read("spec/fixtures/ai-proxy/huggingface/llm-v1-completions/requests/good.json"),
|
|
})
|
|
|
|
-- validate that the request succeeded, response status 200
|
|
local body = assert.res_status(200, r)
|
|
local json = cjson.decode(body)
|
|
|
|
-- check this is in the 'kong' response format
|
|
assert.equals("mistralai/Mistral-7B-Instruct-v0.2", json.model)
|
|
assert.equals("llm/v1/completions", json.object)
|
|
|
|
assert.is_table(json.choices)
|
|
assert.is_table(json.choices[1])
|
|
assert.same("I am a language model AI created by Mistral AI", json.choices[1].message.content)
|
|
end)
|
|
end)
|
|
describe("huggingface no auth", function()
|
|
it("unauthorized request chat", function()
|
|
local r = client:get("/huggingface/llm/v1/chat/unauthorized", {
|
|
headers = {
|
|
["content-type"] = "application/json",
|
|
["accept"] = "application/json",
|
|
},
|
|
body = pl_file.read("spec/fixtures/ai-proxy/huggingface/llm-v1-chat/requests/good.json"),
|
|
})
|
|
|
|
local body = assert.res_status(401, r)
|
|
local json = cjson.decode(body)
|
|
assert.equals(json.error, "Authorization header is correct, but the token seems invalid")
|
|
end)
|
|
it("unauthorized request completions", function()
|
|
local r = client:get("/huggingface/llm/v1/completions/unauthorized", {
|
|
headers = {
|
|
["content-type"] = "application/json",
|
|
["accept"] = "application/json",
|
|
},
|
|
body = pl_file.read("spec/fixtures/ai-proxy/huggingface/llm-v1-completions/requests/good.json"),
|
|
})
|
|
|
|
local body = assert.res_status(401, r)
|
|
local json = cjson.decode(body)
|
|
assert.equals(json.error, "Authorization header is correct, but the token seems invalid")
|
|
end)
|
|
end)
|
|
describe("huggingface bad request", function()
|
|
it("bad chat request", function()
|
|
local r = client:get("/huggingface/llm/v1/chat/good", {
|
|
headers = {
|
|
["content-type"] = "application/json",
|
|
["accept"] = "application/json",
|
|
},
|
|
body = { messages = ngx.null },
|
|
})
|
|
|
|
local body = assert.res_status(400, r)
|
|
local json = cjson.decode(body)
|
|
assert.equals(json.error.message, "request body doesn't contain valid prompts")
|
|
end)
|
|
it("bad completions request", function()
|
|
local r = client:get("/huggingface/llm/v1/completions/good", {
|
|
headers = {
|
|
["content-type"] = "application/json",
|
|
["accept"] = "application/json",
|
|
},
|
|
body = { prompt = ngx.null },
|
|
})
|
|
|
|
local body = assert.res_status(400, r)
|
|
local json = cjson.decode(body)
|
|
assert.equals(json.error.message, "request body doesn't contain valid prompts")
|
|
end)
|
|
end)
|
|
describe("huggingface bad response", function()
|
|
it("bad chat response", function()
|
|
local r = client:get("/huggingface/llm/v1/chat/bad-response/model-loading", {
|
|
headers = {
|
|
["content-type"] = "application/json",
|
|
["accept"] = "application/json",
|
|
},
|
|
body = pl_file.read("spec/fixtures/ai-proxy/huggingface/llm-v1-chat/requests/good.json"),
|
|
})
|
|
|
|
local body = assert.res_status(503, r)
|
|
local json = cjson.decode(body)
|
|
assert.equals(json.error, "Model mistralai/Mistral-7B-Instruct-v0.2 is currently loading")
|
|
end)
|
|
it("bad completions request", function()
|
|
local r = client:get("/huggingface/llm/v1/chat/bad-response/model-timeout", {
|
|
headers = {
|
|
["content-type"] = "application/json",
|
|
["accept"] = "application/json",
|
|
},
|
|
body = pl_file.read("spec/fixtures/ai-proxy/huggingface/llm-v1-chat/requests/good.json"),
|
|
})
|
|
local body = assert.res_status(504, r)
|
|
local json = cjson.decode(body)
|
|
assert.equals(json.error, "Model mistralai/Mistral-7B-Instruct-v0.2 time out")
|
|
end)
|
|
end)
|
|
end)
|
|
end
|
|
end
|