Skip to content

Commit

Permalink
feat(ai-proxy): add option to return model name in response header (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
fffonion authored Aug 9, 2024
1 parent 0019e4b commit 87d908f
Show file tree
Hide file tree
Showing 10 changed files with 24 additions and 0 deletions.
3 changes: 3 additions & 0 deletions changelog/unreleased/kong/ai-proxy-model-header.yml
Original file line number Diff line number Diff line change
@@ -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
1 change: 1 addition & 0 deletions kong/clustering/compat/removed_fields.lua
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
5 changes: 5 additions & 0 deletions kong/llm/proxy/handler.lua
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions kong/plugins/ai-proxy/schema.lua
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
8 changes: 8 additions & 0 deletions spec/02-integration/09-hybrid_mode/09-config-compat_spec.lua
Original file line number Diff line number Diff line change
Expand Up @@ -510,6 +510,7 @@ describe("CP/DP config compat transformations #" .. strategy, function()
},
},
max_request_body_size = 8192,
model_name_header = true,
},
}
-- ]]
Expand All @@ -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
Expand Down Expand Up @@ -695,6 +699,7 @@ describe("CP/DP config compat transformations #" .. strategy, function()
},
},
max_request_body_size = 8192,
model_name_header = true,
},
}
-- ]]
Expand All @@ -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
Expand Down
1 change: 1 addition & 0 deletions spec/03-plugins/38-ai-proxy/02-openai_integration_spec.lua
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
1 change: 1 addition & 0 deletions spec/03-plugins/38-ai-proxy/04-cohere_integration_spec.lua
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
1 change: 1 addition & 0 deletions spec/03-plugins/38-ai-proxy/05-azure_integration_spec.lua
Original file line number Diff line number Diff line change
Expand Up @@ -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])
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down

1 comment on commit 87d908f

@github-actions
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bazel Build

Docker image available kong/kong:87d908f559ae299f9a56f6fcc864413f9e5ef34c
Artifacts available https://github.com/Kong/kong/actions/runs/10315809078

Please sign in to comment.