feat(plugins): ai-prompt-template plugin

.github/labeler.yml

@@ -98,6 +98,10 @@ plugins/ai-prompt-decorator:
 - changed-files:
   - any-glob-to-any-file: kong/plugins/ai-prompt-decorator/**/*
 
+plugins/ai-prompt-template:
+- changed-files:
+  - any-glob-to-any-file: kong/plugins/ai-prompt-template/**/*
+
 plugins/aws-lambda:
 - changed-files:
   - any-glob-to-any-file: kong/plugins/aws-lambda/**/*

changelog/unreleased/kong/add-ai-prompt-template-plugin.yml

@@ -0,0 +1,3 @@
+message: Introduced the new **AI Prompt Template** which can offer consumers and array of LLM prompt templates, with variable substitutions.
+type: feature
+scope: Plugin

kong-3.6.0-0.rockspec

@@ -577,6 +577,10 @@ build = {
     ["kong.plugins.ai-prompt-decorator.handler"] = "kong/plugins/ai-prompt-decorator/handler.lua",
     ["kong.plugins.ai-prompt-decorator.schema"]  = "kong/plugins/ai-prompt-decorator/schema.lua",
 
+    ["kong.plugins.ai-prompt-template.handler"] = "kong/plugins/ai-prompt-template/handler.lua",
+    ["kong.plugins.ai-prompt-template.schema"]  = "kong/plugins/ai-prompt-template/schema.lua",
+    ["kong.plugins.ai-prompt-template.templater"]  = "kong/plugins/ai-prompt-template/templater.lua",
+
     ["kong.vaults.env"] = "kong/vaults/env/init.lua",
     ["kong.vaults.env.schema"] = "kong/vaults/env/schema.lua",
 

kong/constants.lua

@@ -38,6 +38,7 @@ local plugins = {
   "opentelemetry",
   "ai-proxy",
   "ai-prompt-decorator",
+  "ai-prompt-template",
 }
 
 local plugin_map = {}

kong/plugins/ai-prompt-template/handler.lua

@@ -0,0 +1,119 @@
+local _M = {}
+
+-- imports
+local kong_meta = require "kong.meta"
+local templater = require("kong.plugins.ai-prompt-template.templater"):new()
+local fmt       = string.format
+local parse_url = require("socket.url").parse
+local byte      = string.byte
+local sub       = string.sub
+local type      = type
+local byte      = byte
+--
+
+_M.PRIORITY = 773
+_M.VERSION = kong_meta.version
+
+
+local log_entry_keys = {
+  REQUEST_BODY = "ai.payload.original_request",
+}
+
+local function bad_request(msg)
+  kong.log.debug(msg)
+  return kong.response.exit(ngx.HTTP_BAD_REQUEST, { error = { message = msg } })
+end
+
+local BRACE_START = byte("{")
+local BRACE_END = byte("}")
+local COLON = byte(":")
+local SLASH = byte("/")
+
+---- BORROWED FROM `kong.pdk.vault`
+---
+-- Checks if the passed in reference looks like a reference.
+-- Valid references start with '{template://' and end with '}'.
+--
+-- @local
+-- @function is_reference
+-- @tparam string reference reference to check
+-- @treturn boolean `true` is the passed in reference looks like a reference, otherwise `false`
+local function is_reference(reference)
+  return type(reference)      == "string"
+     and byte(reference, 1)   == BRACE_START
+     and byte(reference, -1)  == BRACE_END
+     and byte(reference, 10)   == COLON
+     and byte(reference, 11)   == SLASH
+     and byte(reference, 12)   == SLASH
+     and sub(reference, 2, 9) == "template"
+end
+
+local function find_template(reference_string, templates)
+  local parts, err = parse_url(sub(reference_string, 2, -2))
+  if not parts then
+    return nil, fmt("template reference is not in format '{template://template_name}' (%s) [%s]", err, reference_string)
+  end
+
+  -- iterate templates to find it
+  for i, v in ipairs(templates) do
+    if v.name == parts.host then
+      return v, nil
+    end
+  end
+
+  return nil, fmt("could not find template name [%s]", parts.host)
+end
+
+function _M:access(conf)
+  kong.service.request.enable_buffering()
+  kong.ctx.shared.ai_prompt_templated = true
+
+  if conf.log_original_request then
+    kong.log.set_serialize_value(log_entry_keys.REQUEST_BODY, kong.request.get_raw_body())
+  end
+
+  local request, err = kong.request.get_body("application/json")
+  if err then
+    return bad_request("this LLM route only supports application/json requests")
+  end
+
+  if (not request.messages) and (not request.prompt) then
+    return bad_request("this LLM route only supports llm/chat or llm/completions type requests")
+  end
+
+  if request.messages and request.prompt then
+    return bad_request("cannot run 'messages' and 'prompt' templates at the same time")
+  end
+
+  local reference
+  if request.messages then
+    reference = request.messages
+
+  elseif request.prompt then
+    reference = request.prompt
+
+  else
+    return bad_request("only 'llm/v1/chat' and 'llm/v1/completions' formats are supported for templating")
+  end
+
+  if is_reference(reference) then
+    local requested_template, err = find_template(reference, conf.templates)
+    if not requested_template then
+      return bad_request(err)
+    end
+
+    -- try to render the replacement request
+    local rendered_template, err = templater:render(requested_template, request.properties or {})
+    if err then
+      return bad_request(err)
+    end
+
+    kong.service.request.set_raw_body(rendered_template)
+
+  elseif not (conf.allow_untemplated_requests) then
+    return bad_request("this LLM route only supports templated requests")
+  end
+end
+
+
+return _M

kong/plugins/ai-prompt-template/schema.lua

@@ -0,0 +1,51 @@
+local typedefs = require "kong.db.schema.typedefs"
+
+
+local template_schema = {
+  type = "record",
+  required = true,
+  fields = {
+    { name = {
+        type = "string",
+        description = "Unique name for the template, can be called with `{template://NAME}`",
+        required = true,
+    }},
+    { template = {
+        type = "string",
+        description = "Template string for this request, supports mustache-style `{{placeholders}}`",
+        required = true,
+    }},
+  }
+}
+
+
+return {
+  name = "ai-prompt-template",
+  fields = {
+    { protocols = typedefs.protocols_http },
+    { consumer = typedefs.no_consumer },
+    { config = {
+      type = "record",
+      fields = {
+        { templates = {
+            description = "Array of templates available to the request context.",
+            type = "array",
+            elements = template_schema,
+            required = true,
+        }},
+        { allow_untemplated_requests = {
+            description = "Set true to allow requests that don't call or match any template.",
+            type = "boolean",
+            required = true,
+            default = true,
+        }},
+        { log_original_request = {
+            description = "Set true to add the original request to the Kong log plugin(s) output.",
+            type = "boolean",
+            required = true,
+            default = false,
+        }},
+      }
+    }}
+  },
+}

kong/plugins/ai-prompt-template/templater.lua

@@ -0,0 +1,93 @@
+local _S = {}
+
+-- imports
+local fmt           = string.format
+--
+
+-- globals
+local GSUB_REPLACE_PATTERN = "{{([%w_]+)}}"
+--
+
+local function backslash_replacement_function(c)
+  if c == "\n" then
+     return "\\n"
+  elseif c == "\r" then
+     return "\\r"
+  elseif c == "\t" then
+     return "\\t"
+  elseif c == "\b" then
+     return "\\b"
+  elseif c == "\f" then
+     return "\\f"
+  elseif c == '"' then
+     return '\\"'
+  elseif c == '\\' then
+     return '\\\\'
+  else
+     return string.format("\\u%04x", c:byte())
+  end
+end
+
+local chars_to_be_escaped_in_JSON_string
+= '['
+..    '"'    -- class sub-pattern to match a double quote
+..    '%\\'  -- class sub-pattern to match a backslash
+..    '%z'   -- class sub-pattern to match a null
+..    '\001' .. '-' .. '\031' -- class sub-pattern to match control characters
+.. ']'
+
+-- borrowed from turbo-json
+local function sanitize_parameter(s)
+  if type(s) ~= "string" or s == "" then
+    return nil, nil, "only string arguments are supported"
+  end
+
+  -- check if someone is trying to inject JSON control characters to close the command
+  if s:sub(-1) == "," then
+    s = s:sub(1, -1)
+  end
+
+  return s:gsub(chars_to_be_escaped_in_JSON_string, backslash_replacement_function), nil
+end
+
+function _S:new(o)
+  local o = o or {}
+  setmetatable(o, self)
+  self.__index = self
+
+  return o
+end
+
+
+function _S:render(template, properties)
+  local sanitized_properties = {}
+  local err, _
+
+  for k, v in pairs(properties) do
+    sanitized_properties[k], _, err = sanitize_parameter(v)
+    if err then return nil, err end
+  end
+
+  local result = template.template:gsub(GSUB_REPLACE_PATTERN, sanitized_properties)
+
+  -- find any missing variables
+  local errors = {}
+  local error_string
+  for w in (result):gmatch(GSUB_REPLACE_PATTERN) do 
+    errors[w] = true
+  end
+
+  if next(errors) ~= nil then
+    for k, _ in pairs(errors) do
+      if not error_string then
+        error_string = fmt("missing template parameters: [%s]", k)
+      else
+        error_string = fmt("%s, [%s]", error_string, k)
+      end
+    end
+  end
+
+  return result, error_string
+end
+
+return _S

spec/01-unit/12-plugins_order_spec.lua

@@ -72,6 +72,7 @@ describe("Plugins", function()
       "response-ratelimiting",
       "request-transformer",
       "response-transformer",
+      "ai-prompt-template",
       "ai-prompt-decorator",
       "ai-proxy",
       "aws-lambda",