-
Notifications
You must be signed in to change notification settings - Fork 4.8k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
refactor(ai-plugins): migrate ai transformer plugins to new framework
- Loading branch information
Showing
7 changed files
with
346 additions
and
279 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
113 changes: 113 additions & 0 deletions
113
kong/plugins/ai-request-transformer/filters/transform-request.lua
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,113 @@ | ||
| -- This software is copyright Kong Inc. and its licensors. | ||
| -- Use of the software is subject to the agreement between your organization | ||
| -- and Kong Inc. If there is no such agreement, use is governed by and | ||
| -- subject to the terms of the Kong Master Software License Agreement found | ||
| -- at https://konghq.com/enterprisesoftwarelicense/. | ||
| -- [ END OF LICENSE 0867164ffc95e54f04670b5169c09574bdbd9bba ] | ||
|
|
||
| local fmt = string.format | ||
| local ai_plugin_ctx = require("kong.llm.plugin.ctx") | ||
| local ai_plugin_o11y = require("kong.llm.plugin.observability") | ||
| local ai_shared = require("kong.llm.drivers.shared") | ||
| local llm = require("kong.llm") | ||
|
|
||
|
|
||
| local _M = { | ||
| NAME = "ai-request-transformer-transform-request", | ||
| STAGE = "REQ_TRANSFORMATION", | ||
| } | ||
|
|
||
| local FILTER_OUTPUT_SCHEMA = { | ||
| transformed = "boolean", | ||
| model = "table", | ||
| -- TODO: refactor this so they don't need to be duplicated | ||
| llm_prompt_tokens_count = "number", | ||
| llm_completion_tokens_count = "number", | ||
| llm_usage_cost = "number", | ||
| } | ||
|
|
||
| local _, set_ctx = ai_plugin_ctx.get_namespaced_accesors(_M.NAME, FILTER_OUTPUT_SCHEMA) | ||
|
|
||
| local _KEYBASTION = setmetatable({}, { | ||
| __mode = "k", | ||
| __index = ai_shared.cloud_identity_function, | ||
| }) | ||
|
|
||
|
|
||
| local function bad_request(msg) | ||
| kong.log.info(msg) | ||
| return kong.response.exit(400, { error = { message = msg } }) | ||
| end | ||
|
|
||
| local function internal_server_error(msg) | ||
| kong.log.err(msg) | ||
| return kong.response.exit(500, { error = { message = msg } }) | ||
| end | ||
|
|
||
| local function create_http_opts(conf) | ||
| local http_opts = {} | ||
|
|
||
| if conf.http_proxy_host then -- port WILL be set via schema constraint | ||
| http_opts.proxy_opts = http_opts.proxy_opts or {} | ||
| http_opts.proxy_opts.http_proxy = fmt("http://%s:%d", conf.http_proxy_host, conf.http_proxy_port) | ||
| end | ||
|
|
||
| if conf.https_proxy_host then | ||
| http_opts.proxy_opts = http_opts.proxy_opts or {} | ||
| http_opts.proxy_opts.https_proxy = fmt("http://%s:%d", conf.https_proxy_host, conf.https_proxy_port) | ||
| end | ||
|
|
||
| http_opts.http_timeout = conf.http_timeout | ||
| http_opts.https_verify = conf.https_verify | ||
|
|
||
| return http_opts | ||
| end | ||
|
|
||
|
|
||
|
|
||
| function _M:run(conf) | ||
| -- get cloud identity SDK, if required | ||
| local identity_interface = _KEYBASTION[conf.llm] | ||
|
|
||
| if identity_interface and identity_interface.error then | ||
| kong.log.err("error authenticating with ", conf.model.provider, " using native provider auth, ", identity_interface.error) | ||
| return kong.response.exit(500, "LLM request failed before proxying") | ||
| end | ||
|
|
||
| -- first find the configured LLM interface and driver | ||
| local http_opts = create_http_opts(conf) | ||
| conf.llm.__plugin_id = conf.__plugin_id | ||
| conf.llm.__key__ = conf.__key__ | ||
| local ai_driver, err = llm.new_driver(conf.llm, http_opts, identity_interface) | ||
|
|
||
| if not ai_driver then | ||
| return internal_server_error(err) | ||
| end | ||
|
|
||
| -- if asked, introspect the request before proxying | ||
| kong.log.debug("introspecting request with LLM") | ||
| local new_request_body, err = ai_driver:ai_introspect_body( | ||
| kong.request.get_raw_body(conf.max_request_body_size), | ||
| conf.prompt, | ||
| http_opts, | ||
| conf.transformation_extract_pattern | ||
| ) | ||
|
|
||
| if err then | ||
| return bad_request(err) | ||
| end | ||
|
|
||
| set_ctx("model", conf.llm.model) | ||
| set_ctx("llm_prompt_tokens_count", ai_plugin_o11y.metrics_get("llm_prompt_tokens_count") or 0) | ||
| set_ctx("llm_completion_tokens_count", ai_plugin_o11y.metrics_get("llm_completion_tokens_count") or 0) | ||
| set_ctx("llm_usage_cost", ai_plugin_o11y.metrics_get("llm_usage_cost") or 0) | ||
|
|
||
| -- set the body for later plugins | ||
| kong.service.request.set_raw_body(new_request_body) | ||
|
|
||
| set_ctx("transformed", true) | ||
| return true | ||
| end | ||
|
|
||
|
|
||
| return _M |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,95 +1,20 @@ | ||
| local _M = {} | ||
| local ai_plugin_base = require("kong.llm.plugin.base") | ||
|
|
||
| -- imports | ||
| local kong_meta = require "kong.meta" | ||
| local fmt = string.format | ||
| local llm = require("kong.llm") | ||
| local llm_state = require("kong.llm.state") | ||
| local ai_shared = require("kong.llm.drivers.shared") | ||
| -- | ||
| local NAME = "ai-request-transformer" | ||
| local PRIORITY = 777 | ||
|
|
||
| _M.PRIORITY = 777 | ||
| _M.VERSION = kong_meta.version | ||
| local AIPlugin = ai_plugin_base.define(NAME, PRIORITY) | ||
|
|
||
| local _KEYBASTION = setmetatable({}, { | ||
| __mode = "k", | ||
| __index = ai_shared.cloud_identity_function, | ||
| }) | ||
|
|
||
| local function bad_request(msg) | ||
| kong.log.info(msg) | ||
| return kong.response.exit(400, { error = { message = msg } }) | ||
| end | ||
| local SHARED_FILTERS = { | ||
| "enable-buffering", | ||
| } | ||
|
|
||
| local function internal_server_error(msg) | ||
| kong.log.err(msg) | ||
| return kong.response.exit(500, { error = { message = msg } }) | ||
| for _, filter in ipairs(SHARED_FILTERS) do | ||
| AIPlugin:enable(AIPlugin.register_filter(require("kong.llm.plugin.shared-filters." .. filter))) | ||
| end | ||
|
|
||
| local function create_http_opts(conf) | ||
| local http_opts = {} | ||
|
|
||
| if conf.http_proxy_host then -- port WILL be set via schema constraint | ||
| http_opts.proxy_opts = http_opts.proxy_opts or {} | ||
| http_opts.proxy_opts.http_proxy = fmt("http://%s:%d", conf.http_proxy_host, conf.http_proxy_port) | ||
| end | ||
|
|
||
| if conf.https_proxy_host then | ||
| http_opts.proxy_opts = http_opts.proxy_opts or {} | ||
| http_opts.proxy_opts.https_proxy = fmt("http://%s:%d", conf.https_proxy_host, conf.https_proxy_port) | ||
| end | ||
|
|
||
| http_opts.http_timeout = conf.http_timeout | ||
| http_opts.https_verify = conf.https_verify | ||
|
|
||
| return http_opts | ||
| end | ||
|
|
||
| function _M:access(conf) | ||
| llm_state.set_request_model(conf.llm.model and conf.llm.model.name) | ||
| local kong_ctx_shared = kong.ctx.shared | ||
|
|
||
| kong.service.request.enable_buffering() | ||
| llm_state.should_disable_ai_proxy_response_transform() | ||
|
|
||
| -- get cloud identity SDK, if required | ||
| local identity_interface = _KEYBASTION[conf.llm] | ||
|
|
||
| if identity_interface and identity_interface.error then | ||
| kong_ctx_shared.skip_response_transformer = true | ||
| kong.log.err("error authenticating with ", conf.model.provider, " using native provider auth, ", identity_interface.error) | ||
| return kong.response.exit(500, "LLM request failed before proxying") | ||
| end | ||
|
|
||
| -- first find the configured LLM interface and driver | ||
| local http_opts = create_http_opts(conf) | ||
| conf.llm.__plugin_id = conf.__plugin_id | ||
| conf.llm.__key__ = conf.__key__ | ||
| local ai_driver, err = llm.new_driver(conf.llm, http_opts, identity_interface) | ||
|
|
||
| if not ai_driver then | ||
| return internal_server_error(err) | ||
| end | ||
|
|
||
| -- if asked, introspect the request before proxying | ||
| kong.log.debug("introspecting request with LLM") | ||
| local new_request_body, err = ai_driver:ai_introspect_body( | ||
| kong.request.get_raw_body(conf.max_request_body_size), | ||
| conf.prompt, | ||
| http_opts, | ||
| conf.transformation_extract_pattern | ||
| ) | ||
|
|
||
| if err then | ||
| return bad_request(err) | ||
| end | ||
|
|
||
| -- set the body for later plugins | ||
| kong.service.request.set_raw_body(new_request_body) | ||
|
|
||
| -- continue into other plugins including ai-response-transformer, | ||
| -- which may exit early with a sub-request | ||
| end | ||
| AIPlugin:enable(AIPlugin.register_filter(require("kong.plugins.ai-request-transformer.filters.transform-request"))) | ||
|
|
||
|
|
||
| return _M | ||
| return AIPlugin:as_kong_plugin() |
Oops, something went wrong.