wehub-resource-sync
520 lines
15 KiB
Lua
520 lines
15 KiB
Lua
local helpers = require "spec.helpers"
|
|
local cjson = require "cjson"
|
|
local pl_file = require "pl.file"
|
|
|
|
local strip = require("kong.tools.string").strip
|
|
|
|
local PLUGIN_NAME = "ai-request-transformer"
|
|
|
|
local FILE_LOG_PATH_STATS_ONLY = os.tmpname()
|
|
local FILE_LOG_PATH_GEMINI_STATS_ONLY = os.tmpname()
|
|
|
|
local function wait_for_json_log_entry(FILE_LOG_PATH)
|
|
local json
|
|
|
|
assert
|
|
.with_timeout(10)
|
|
.ignore_exceptions(true)
|
|
.eventually(function()
|
|
local data = assert(pl_file.read(FILE_LOG_PATH))
|
|
|
|
data = strip(data)
|
|
assert(#data > 0, "log file is empty")
|
|
|
|
data = data:match("%b{}")
|
|
assert(data, "log file does not contain JSON")
|
|
|
|
json = cjson.decode(data)
|
|
end)
|
|
.has_no_error("log file contains a valid JSON entry")
|
|
|
|
return json
|
|
end
|
|
|
|
|
|
local REQUEST_BODY = [[
|
|
{
|
|
"persons": [
|
|
{
|
|
"name": "Kong A",
|
|
"age": 31
|
|
},
|
|
{
|
|
"name": "Kong B",
|
|
"age": 42
|
|
}
|
|
]
|
|
}
|
|
]]
|
|
|
|
local EXPECTED_RESULT_FLAT = {
|
|
persons = {
|
|
[1] = {
|
|
age = 62,
|
|
name = "Kong A"
|
|
},
|
|
[2] = {
|
|
age = 84,
|
|
name = "Kong B"
|
|
},
|
|
}
|
|
}
|
|
|
|
local _EXPECTED_CHAT_STATS = {
|
|
["ai-request-transformer"] = {
|
|
meta = {
|
|
plugin_id = '71083e79-4921-4f9f-97a4-ee7810b6cd8a',
|
|
provider_name = 'openai',
|
|
request_model = 'UNSPECIFIED',
|
|
response_model = 'gpt-3.5-turbo-0613',
|
|
llm_latency = 1
|
|
},
|
|
usage = {
|
|
prompt_tokens = 25,
|
|
completion_tokens = 12,
|
|
total_tokens = 37,
|
|
time_per_token = 1,
|
|
cost = 0.00037,
|
|
},
|
|
cache = {}
|
|
},
|
|
}
|
|
|
|
|
|
local _EXPECTED_CHAT_STATS_GEMINI = {
|
|
["ai-request-transformer"] = {
|
|
meta = {
|
|
plugin_id = '71083e79-4921-4f9f-97a4-ee7810b6cd8b',
|
|
provider_name = 'gemini',
|
|
request_model = 'UNSPECIFIED',
|
|
response_model = 'gemini-1.5-flash',
|
|
llm_latency = 1
|
|
},
|
|
usage = {
|
|
prompt_tokens = 2,
|
|
completion_tokens = 11,
|
|
total_tokens = 13,
|
|
time_per_token = 1,
|
|
cost = 0.00026,
|
|
},
|
|
cache = {}
|
|
},
|
|
}
|
|
|
|
|
|
local SYSTEM_PROMPT = "You are a mathematician. "
|
|
.. "Multiply all numbers in my JSON request, by 2."
|
|
|
|
|
|
local client
|
|
|
|
for _, strategy in helpers.all_strategies() do
|
|
describe(PLUGIN_NAME .. ": (access) [#" .. strategy .. "]", function()
|
|
local MOCK_PORT
|
|
local OPENAI_FLAT_RESPONSE
|
|
local GEMINI_GOOD
|
|
local GEMINI_GOOD_FAILS_SAFETY
|
|
local OPENAI_BAD_REQUEST
|
|
local OPENAI_INTERNAL_SERVER_ERROR
|
|
|
|
lazy_setup(function()
|
|
MOCK_PORT = helpers.get_available_port()
|
|
OPENAI_FLAT_RESPONSE = {
|
|
route_type = "llm/v1/chat",
|
|
logging = {
|
|
log_payloads = false,
|
|
log_statistics = true,
|
|
},
|
|
model = {
|
|
name = "gpt-4",
|
|
provider = "openai",
|
|
options = {
|
|
max_tokens = 512,
|
|
temperature = 0.5,
|
|
upstream_url = "http://"..helpers.mock_upstream_host..":"..MOCK_PORT.."/flat",
|
|
input_cost = 10.0,
|
|
output_cost = 10.0,
|
|
},
|
|
},
|
|
auth = {
|
|
header_name = "Authorization",
|
|
header_value = "Bearer openai-key",
|
|
},
|
|
}
|
|
|
|
GEMINI_GOOD = {
|
|
route_type = "llm/v1/chat",
|
|
logging = {
|
|
log_payloads = false,
|
|
log_statistics = true,
|
|
},
|
|
model = {
|
|
name = "gemini-1.5-flash",
|
|
provider = "gemini",
|
|
options = {
|
|
max_tokens = 512,
|
|
temperature = 0.6,
|
|
upstream_url = "http://"..helpers.mock_upstream_host..":"..MOCK_PORT.."/geminiflat",
|
|
input_cost = 20.0,
|
|
output_cost = 20.0,
|
|
},
|
|
},
|
|
auth = {
|
|
header_name = "x-goog-api-key",
|
|
header_value = "123",
|
|
},
|
|
}
|
|
|
|
GEMINI_GOOD_FAILS_SAFETY = {
|
|
route_type = "llm/v1/chat",
|
|
logging = {
|
|
log_payloads = false,
|
|
log_statistics = true,
|
|
},
|
|
model = {
|
|
name = "gemini-1.5-flash",
|
|
provider = "gemini",
|
|
options = {
|
|
max_tokens = 512,
|
|
temperature = 0.5,
|
|
upstream_url = "http://"..helpers.mock_upstream_host..":"..MOCK_PORT.."/failssafety",
|
|
input_cost = 10.0,
|
|
output_cost = 10.0,
|
|
},
|
|
},
|
|
auth = {
|
|
header_name = "x-goog-api-key",
|
|
header_value = "123",
|
|
},
|
|
}
|
|
|
|
OPENAI_BAD_REQUEST = {
|
|
route_type = "llm/v1/chat",
|
|
model = {
|
|
name = "gpt-4",
|
|
provider = "openai",
|
|
options = {
|
|
max_tokens = 512,
|
|
temperature = 0.5,
|
|
upstream_url = "http://"..helpers.mock_upstream_host..":"..MOCK_PORT.."/badrequest"
|
|
},
|
|
},
|
|
auth = {
|
|
header_name = "Authorization",
|
|
header_value = "Bearer openai-key",
|
|
},
|
|
}
|
|
|
|
OPENAI_INTERNAL_SERVER_ERROR = {
|
|
route_type = "llm/v1/chat",
|
|
model = {
|
|
name = "gpt-4",
|
|
provider = "openai",
|
|
options = {
|
|
max_tokens = 512,
|
|
temperature = 0.5,
|
|
upstream_url = "http://"..helpers.mock_upstream_host..":"..MOCK_PORT.."/internalservererror"
|
|
},
|
|
},
|
|
auth = {
|
|
header_name = "Authorization",
|
|
header_value = "Bearer openai-key",
|
|
},
|
|
}
|
|
|
|
local bp = helpers.get_db_utils(strategy == "off" and "postgres" or strategy, nil, { PLUGIN_NAME })
|
|
|
|
-- set up provider fixtures
|
|
local fixtures = {
|
|
http_mock = {},
|
|
}
|
|
|
|
fixtures.http_mock.openai = [[
|
|
server {
|
|
server_name llm;
|
|
listen ]]..MOCK_PORT..[[;
|
|
|
|
default_type 'application/json';
|
|
|
|
location ~/flat {
|
|
content_by_lua_block {
|
|
local pl_file = require "pl.file"
|
|
ngx.print(pl_file.read("spec/fixtures/ai-proxy/openai/request-transformer/response-in-json.json"))
|
|
}
|
|
}
|
|
|
|
location ~/geminiflat {
|
|
content_by_lua_block {
|
|
local pl_file = require "pl.file"
|
|
ngx.print(pl_file.read("spec/fixtures/ai-proxy/gemini/request-transformer/response-in-json.json"))
|
|
}
|
|
}
|
|
|
|
location = "/badrequest" {
|
|
content_by_lua_block {
|
|
local pl_file = require "pl.file"
|
|
|
|
ngx.status = 400
|
|
ngx.print(pl_file.read("spec/fixtures/ai-proxy/openai/llm-v1-chat/responses/bad_request.json"))
|
|
}
|
|
}
|
|
|
|
location = "/failssafety" {
|
|
content_by_lua_block {
|
|
local pl_file = require "pl.file"
|
|
|
|
ngx.status = 200
|
|
ngx.print(pl_file.read("spec/fixtures/ai-proxy/gemini/llm-v1-chat/responses/fails_safety.json"))
|
|
}
|
|
}
|
|
|
|
location = "/internalservererror" {
|
|
content_by_lua_block {
|
|
local pl_file = require "pl.file"
|
|
|
|
ngx.status = 500
|
|
ngx.header["content-type"] = "text/html"
|
|
ngx.print(pl_file.read("spec/fixtures/ai-proxy/openai/llm-v1-chat/responses/internal_server_error.html"))
|
|
}
|
|
}
|
|
}
|
|
]]
|
|
|
|
-- echo server via 'openai' LLM
|
|
local without_response_instructions = assert(bp.routes:insert {
|
|
paths = { "/echo-flat" }
|
|
})
|
|
bp.plugins:insert {
|
|
name = PLUGIN_NAME,
|
|
id = "71083e79-4921-4f9f-97a4-ee7810b6cd8a",
|
|
route = { id = without_response_instructions.id },
|
|
config = {
|
|
prompt = SYSTEM_PROMPT,
|
|
llm = OPENAI_FLAT_RESPONSE,
|
|
},
|
|
}
|
|
bp.plugins:insert {
|
|
name = "file-log",
|
|
route = { id = without_response_instructions.id },
|
|
config = {
|
|
path = FILE_LOG_PATH_STATS_ONLY,
|
|
},
|
|
}
|
|
|
|
-- echo server via 'non-openai' LLM
|
|
local gemini_without_response_instructions = assert(bp.routes:insert {
|
|
paths = { "/gemini-echo-flat" }
|
|
})
|
|
bp.plugins:insert {
|
|
name = PLUGIN_NAME,
|
|
id = "71083e79-4921-4f9f-97a4-ee7810b6cd8b",
|
|
route = { id = gemini_without_response_instructions.id },
|
|
config = {
|
|
prompt = SYSTEM_PROMPT,
|
|
llm = GEMINI_GOOD,
|
|
},
|
|
}
|
|
bp.plugins:insert {
|
|
name = "file-log",
|
|
route = { id = gemini_without_response_instructions.id },
|
|
config = {
|
|
path = FILE_LOG_PATH_GEMINI_STATS_ONLY,
|
|
},
|
|
}
|
|
|
|
local bad_request = assert(bp.routes:insert {
|
|
paths = { "/echo-bad-request" }
|
|
})
|
|
bp.plugins:insert {
|
|
name = PLUGIN_NAME,
|
|
route = { id = bad_request.id },
|
|
config = {
|
|
prompt = SYSTEM_PROMPT,
|
|
llm = OPENAI_BAD_REQUEST,
|
|
},
|
|
}
|
|
|
|
local fails_safety = assert(bp.routes:insert {
|
|
paths = { "/echo-fails-safety" }
|
|
})
|
|
bp.plugins:insert {
|
|
name = PLUGIN_NAME,
|
|
route = { id = fails_safety.id },
|
|
config = {
|
|
prompt = SYSTEM_PROMPT,
|
|
llm = GEMINI_GOOD_FAILS_SAFETY,
|
|
},
|
|
}
|
|
|
|
local internal_server_error = assert(bp.routes:insert {
|
|
paths = { "/echo-internal-server-error" }
|
|
})
|
|
bp.plugins:insert {
|
|
name = PLUGIN_NAME,
|
|
route = { id = internal_server_error.id },
|
|
config = {
|
|
prompt = SYSTEM_PROMPT,
|
|
llm = OPENAI_INTERNAL_SERVER_ERROR,
|
|
},
|
|
}
|
|
--
|
|
|
|
-- 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("openai response transformer integration", function()
|
|
it("transforms properly from LLM", function()
|
|
local r = client:get("/echo-flat", {
|
|
headers = {
|
|
["content-type"] = "application/json",
|
|
["accept"] = "application/json",
|
|
},
|
|
body = REQUEST_BODY,
|
|
})
|
|
|
|
local body = assert.res_status(200 , r)
|
|
local body_table, err = cjson.decode(body)
|
|
|
|
assert.is_nil(err)
|
|
assert.same(EXPECTED_RESULT_FLAT, body_table.post_data.params)
|
|
end)
|
|
|
|
it("logs statistics - openai format", function()
|
|
local r = client:get("/echo-flat", {
|
|
headers = {
|
|
["content-type"] = "application/json",
|
|
["accept"] = "application/json",
|
|
},
|
|
body = REQUEST_BODY,
|
|
})
|
|
|
|
local body = assert.res_status(200 , r)
|
|
local _, err = cjson.decode(body)
|
|
|
|
assert.is_nil(err)
|
|
|
|
local log_message = wait_for_json_log_entry(FILE_LOG_PATH_STATS_ONLY)
|
|
assert.same("127.0.0.1", log_message.client_ip)
|
|
assert.is_number(log_message.request.size)
|
|
assert.is_number(log_message.response.size)
|
|
|
|
-- test ai-request-transformer stats
|
|
local actual_chat_stats = log_message.ai
|
|
local actual_llm_latency = actual_chat_stats["ai-request-transformer"].meta.llm_latency
|
|
local actual_time_per_token = actual_chat_stats["ai-request-transformer"].usage.time_per_token
|
|
local time_per_token = math.floor(actual_llm_latency / actual_chat_stats["ai-request-transformer"].usage.completion_tokens)
|
|
|
|
log_message.ai["ai-request-transformer"].meta.llm_latency = 1
|
|
log_message.ai["ai-request-transformer"].usage.time_per_token = 1
|
|
|
|
assert.same(_EXPECTED_CHAT_STATS, log_message.ai)
|
|
assert.is_true(actual_llm_latency >= 0)
|
|
assert.same(actual_time_per_token, time_per_token)
|
|
end)
|
|
|
|
it("logs statistics - non-openai format", function()
|
|
local r = client:get("/gemini-echo-flat", {
|
|
headers = {
|
|
["content-type"] = "application/json",
|
|
["accept"] = "application/json",
|
|
},
|
|
body = REQUEST_BODY,
|
|
})
|
|
|
|
local body = assert.res_status(200 , r)
|
|
local _, err = cjson.decode(body)
|
|
|
|
assert.is_nil(err)
|
|
|
|
local log_message = wait_for_json_log_entry(FILE_LOG_PATH_GEMINI_STATS_ONLY)
|
|
assert.same("127.0.0.1", log_message.client_ip)
|
|
assert.is_number(log_message.request.size)
|
|
assert.is_number(log_message.response.size)
|
|
|
|
-- test ai-request-transformer stats
|
|
local actual_chat_stats = log_message.ai
|
|
local actual_llm_latency = actual_chat_stats["ai-request-transformer"].meta.llm_latency
|
|
local actual_time_per_token = actual_chat_stats["ai-request-transformer"].usage.time_per_token
|
|
local time_per_token = math.floor(actual_llm_latency / actual_chat_stats["ai-request-transformer"].usage.completion_tokens)
|
|
|
|
log_message.ai["ai-request-transformer"].meta.llm_latency = 1
|
|
log_message.ai["ai-request-transformer"].usage.time_per_token = 1
|
|
|
|
assert.same(_EXPECTED_CHAT_STATS_GEMINI, log_message.ai)
|
|
assert.is_true(actual_llm_latency >= 0)
|
|
assert.same(actual_time_per_token, time_per_token)
|
|
end)
|
|
|
|
it("bad request from LLM", function()
|
|
local r = client:get("/echo-bad-request", {
|
|
headers = {
|
|
["content-type"] = "application/json",
|
|
["accept"] = "application/json",
|
|
},
|
|
body = REQUEST_BODY,
|
|
})
|
|
|
|
local body = assert.res_status(400 , r)
|
|
local body_table, err = cjson.decode(body)
|
|
|
|
assert.is_nil(err)
|
|
assert.same({ error = { message = "failed to introspect request with AI service: status code 400" }}, body_table)
|
|
end)
|
|
|
|
it("fails Gemini content-safety", function()
|
|
local r = client:get("/echo-fails-safety", {
|
|
headers = {
|
|
["content-type"] = "application/json",
|
|
["accept"] = "application/json",
|
|
},
|
|
body = REQUEST_BODY,
|
|
})
|
|
|
|
local body = assert.res_status(400 , r)
|
|
local body_table, err = cjson.decode(body)
|
|
|
|
assert.is_nil(err)
|
|
assert.match_re(body_table.error.message, ".*transformation generation candidate breached Gemini content safety.*")
|
|
end)
|
|
|
|
it("internal server error from LLM", function()
|
|
local r = client:get("/echo-internal-server-error", {
|
|
headers = {
|
|
["content-type"] = "application/json",
|
|
["accept"] = "application/json",
|
|
},
|
|
body = REQUEST_BODY,
|
|
})
|
|
|
|
local body = assert.res_status(400 , r)
|
|
local body_table, err = cjson.decode(body)
|
|
|
|
assert.is_nil(err)
|
|
assert.same({ error = { message = "failed to introspect request with AI service: status code 500" }}, body_table)
|
|
end)
|
|
end)
|
|
end)
|
|
end
|