fix(ai-proxy): model field was incorrect in Gemini responses with var used as model name (#14751)

Fix #14715

Signed-off-by: spacewander <spacewanderlzx@gmail.com>

Author: spacewander

changelog/unreleased/kong/fix-ai-proxy-gemini-incorrect-model-name.yml

@@ -0,0 +1,3 @@
+message: "**ai-proxy**: Fixed an issue where the `model` field was incorrect in Gemini responses when a variable was used as the model name."
+type: "bugfix"
+scope: "Plugin"

kong/llm/drivers/gemini.lua

@@ -106,6 +106,7 @@ local function handle_stream_event(event_t, model_info, route_type)
     metadata.prompt_tokens     = event.usageMetadata and event.usageMetadata.promptTokenCount or 0
 
     local new_event = {
+      model = model_info.name,
       choices = {
         [1] = {
           delta = {
@@ -128,6 +129,7 @@ local function handle_stream_event(event_t, model_info, route_type)
 
     if event.candidates and #event.candidates > 0 then
       local new_event = {
+        model = model_info.name,
         choices = {
           [1] = {
             delta = {
@@ -395,6 +397,7 @@ local function from_gemini_chat_openai(response, model_info, route_type)
   -- messages/choices table is only 1 size, so don't need to static allocate
   local messages = {}
   messages.choices = {}
+  messages.model = model_info.name -- openai format always contains the model name
 
   if response.candidates and #response.candidates > 0 then
     -- for transformer plugins only

kong/llm/plugin/shared-filters/normalize-json-response.lua

@@ -28,7 +28,8 @@ local function transform_body(conf)
     local adapter = get_global_ctx("llm_format_adapter")
     if not adapter then
       -- openai formats
-      response_body, err = ai_driver.from_format(response_body, conf.model, route_type)
+      local model_t = ai_plugin_ctx.get_request_model_table_inuse()
+      response_body, err = ai_driver.from_format(response_body, model_t, route_type)
     end
 
     if err then

kong/llm/plugin/shared-filters/normalize-sse-chunk.lua

@@ -85,7 +85,8 @@ local function handle_streaming_frame(conf, chunk, finished)
   for _, event in ipairs(events) do
     -- TODO: currently only subset of driver follow the body, err, metadata pattern
     -- unify this so that it was always extracted from the body
-    local formatted, _, metadata = ai_driver.from_format(event, conf.model, "stream/" .. conf.route_type)
+    local model_t = ai_plugin_ctx.get_request_model_table_inuse()
+    local formatted, _, metadata = ai_driver.from_format(event, model_t, "stream/" .. conf.route_type)
 
     if formatted then
       frame_buffer:put("data: ")

spec/03-plugins/38-ai-proxy/11-gemini_integration_spec.lua

@@ -142,6 +142,36 @@ for _, strategy in helpers.all_strategies() do
           },
         }
 
+        -- 200 chat good with variable
+        local chat_good_with_var = assert(bp.routes:insert({
+          service = empty_service,
+          protocols = { "http" },
+          strip_path = true,
+          paths = { "~/gemini/llm/v1/chat/good/(?<model>[^/]+)" },
+        }))
+        bp.plugins:insert({
+          name = PLUGIN_NAME,
+          route = { id = chat_good_with_var.id },
+          config = {
+            route_type = "llm/v1/chat",
+            auth = {
+              header_name = "Authorization",
+              header_value = "Bearer gemini-key",
+            },
+            logging = {
+              log_payloads = true,
+              log_statistics = true,
+            },
+            model = {
+              name = "$(uri_captures.model)",
+              provider = "gemini",
+              options = {
+                upstream_url = "http://" .. helpers.mock_upstream_host .. ":" .. MOCK_PORT .. "/v1/chat/completions",
+              },
+            },
+          },
+        })
+
         -- start kong
         assert(helpers.start_kong({
           -- set the strategy
@@ -218,6 +248,20 @@ for _, strategy in helpers.all_strategies() do
           assert.match_re(actual_request_log, [[.*contents.*What is 1 \+ 1.*]])
           assert.match_re(actual_response_log, [[.*content.*Everything is okay.*]])
         end)
+
+        it("good request with model name from variable", function()
+          local r = client:get("/gemini/llm/v1/chat/good/gemni-2.0-flash", {
+            headers = {
+              ["content-type"] = "application/json",
+              ["accept"] = "application/json",
+            },
+            body = pl_file.read("spec/fixtures/ai-proxy/openai/llm-v1-chat/requests/good.json"),
+          })
+          -- validate that the request succeeded, response status 200
+          local body = assert.res_status(200, r)
+          local json = cjson.decode(body)
+          assert.equals("gemni-2.0-flash", json.model)
+        end)
       end)
     end)
   end