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

@@ -175,6 +175,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

@@ -335,6 +335,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/schema.lua

@@ -20,6 +20,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

@@ -510,6 +510,7 @@ describe("CP/DP config compat transformations #" .. strategy, function()
               },
             },
             max_request_body_size = 8192,
+            model_name_header = true,
           },
         }
         -- ]]
@@ -519,6 +520,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
@@ -695,6 +699,7 @@ describe("CP/DP config compat transformations #" .. strategy, function()
               },
             },
             max_request_body_size = 8192,
+            model_name_header = true,
           },
         }
         -- ]]
@@ -704,6 +709,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

@@ -858,6 +858,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

@@ -541,6 +541,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

@@ -416,6 +416,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

@@ -493,6 +493,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

@@ -338,6 +338,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)