feat(ai-proxy): add option to return model name in response header (#…

…13472)

changelog/unreleased/kong/ai-proxy-model-header.yml

@@ -0,0 +1,3 @@
+message: '**ai-proxy**: Added a new response header X-Kong-LLM-Model that displays the name of the language model used in the AI-Proxy plugin.'
+type: feature
+scope: Plugin

kong/clustering/compat/removed_fields.lua

@@ -351,6 +351,7 @@ return {
       "model.options.bedrock",
       "auth.aws_access_key_id",
       "auth.aws_secret_access_key",
+      "model_name_header",
     },
     ai_prompt_decorator = {
       "max_request_body_size",

kong/llm/proxy/handler.lua

@@ -342,6 +342,11 @@ function _M:header_filter(conf)
     kong.response.clear_header(v)
   end
 
+  if ngx.var.http_kong_debug or conf.model_name_header then
+    local name = conf.model.provider .. "/" .. (kong.ctx.plugin.llm_model_requested or conf.model.name)
+    kong.response.set_header("X-Kong-LLM-Model", name)
+  end
+
   -- we use openai's streaming mode (SSE)
   if llm_state.is_streaming_mode() then
     -- we are going to send plaintext event-stream frames for ALL models

kong/plugins/ai-proxy-advanced/handler.lua

@@ -158,10 +158,7 @@ function _M:header_filter(conf)
     return
   end
 
-  if ngx.var.http_kong_debug or conf.model_name_header then
-    local name = target.model.provider .. "/" .. (kong.ctx.plugin.llm_model_requested or target.model.name)
-    kong.response.set_header("X-Kong-LLM-Model", name)
-  end
+  target.model_name_header = conf.model_name_header
 
   return proxy_handler:header_filter(target)
 end

kong/plugins/ai-proxy/schema.lua

@@ -27,6 +27,8 @@ local ai_proxy_only_config = {
       gt = 0,
       description = "max allowed body size allowed to be introspected",}
     },
+    { model_name_header = { description = "Display the model name selected in the X-Kong-LLM-Model response header",
+    type = "boolean", default = true, }},
 }
 
 for i, v in pairs(ai_proxy_only_config) do

spec/02-integration/09-hybrid_mode/09-config-compat_spec.lua

@@ -597,6 +597,7 @@ describe("CP/DP config compat transformations #" .. strategy, function()
               },
             },
             max_request_body_size = 8192,
+            model_name_header = true,
           },
         }
         -- ]]
@@ -606,6 +607,9 @@ describe("CP/DP config compat transformations #" .. strategy, function()
         -- max body size
         expected.config.max_request_body_size = nil
 
+        -- model name header
+        expected.config.model_name_header = nil
+
         -- gemini fields
         expected.config.auth.gcp_service_account_json = nil
         expected.config.auth.gcp_use_service_account = nil
@@ -794,6 +798,7 @@ describe("CP/DP config compat transformations #" .. strategy, function()
               },
             },
             max_request_body_size = 8192,
+            model_name_header = true,
           },
         }
         -- ]]
@@ -803,6 +808,9 @@ describe("CP/DP config compat transformations #" .. strategy, function()
         -- max body size
         expected.config.max_request_body_size = nil
 
+        -- model name header
+        expected.config.model_name_header = nil
+
         -- gemini fields
         expected.config.auth.gcp_service_account_json = nil
         expected.config.auth.gcp_use_service_account = nil

spec/03-plugins/38-ai-proxy/02-openai_integration_spec.lua

@@ -866,6 +866,7 @@ for _, strategy in helpers.all_strategies() do if strategy ~= "cassandra" then
         assert.equals(json.id, "chatcmpl-8T6YwgvjQVVnGbJ2w8hpOA17SeNy2")
         assert.equals(json.model, "gpt-3.5-turbo-0613")
         assert.equals(json.object, "chat.completion")
+        assert.equals(r.headers["X-Kong-LLM-Model"], "openai/gpt-3.5-turbo")
 
         assert.is_table(json.choices)
         assert.is_table(json.choices[1].message)

spec/03-plugins/38-ai-proxy/03-anthropic_integration_spec.lua

@@ -548,6 +548,7 @@ for _, strategy in helpers.all_strategies() do if strategy ~= "cassandra" then
         -- assert.equals(json.id, "chatcmpl-8T6YwgvjQVVnGbJ2w8hpOA17SeNy2")
         assert.equals(json.model, "claude-2.1")
         assert.equals(json.object, "chat.content")
+        assert.equals(r.headers["X-Kong-LLM-Model"], "anthropic/claude-2.1")
 
         assert.is_table(json.choices)
         assert.is_table(json.choices[1].message)

spec/03-plugins/38-ai-proxy/04-cohere_integration_spec.lua

@@ -423,6 +423,7 @@ for _, strategy in helpers.all_strategies() do if strategy ~= "cassandra" then
         -- check this is in the 'kong' response format
         assert.equals(json.model, "command")
         assert.equals(json.object, "chat.completion")
+        assert.equals(r.headers["X-Kong-LLM-Model"], "cohere/command")
 
         assert.is_table(json.choices)
         assert.is_table(json.choices[1].message)

spec/03-plugins/38-ai-proxy/05-azure_integration_spec.lua

@@ -500,6 +500,7 @@ for _, strategy in helpers.all_strategies() do if strategy ~= "cassandra" then
         assert.equals("cmpl-8TBeaJVQIhE9kHEJbk1RnKzgFxIqN", json.id)
         assert.equals("gpt-3.5-turbo-instruct", json.model)
         assert.equals("text_completion", json.object)
+        assert.equals(r.headers["X-Kong-LLM-Model"], "azure/gpt-3.5-turbo-instruct")
 
         assert.is_table(json.choices)
         assert.is_table(json.choices[1])

spec/03-plugins/38-ai-proxy/06-mistral_integration_spec.lua

@@ -345,6 +345,7 @@ for _, strategy in helpers.all_strategies() do if strategy ~= "cassandra" then
         assert.equals(json.id, "chatcmpl-8T6YwgvjQVVnGbJ2w8hpOA17SeNy2")
         assert.equals(json.model, "mistralai/Mistral-7B-Instruct-v0.1-instruct")
         assert.equals(json.object, "chat.completion")
+        assert.equals(r.headers["X-Kong-LLM-Model"], "mistral/mistralai/Mistral-7B-Instruct-v0.1-instruct")
 
         assert.is_table(json.choices)
         assert.is_table(json.choices[1].message)