add testst p

kong/llm/drivers/shared.lua

@@ -848,4 +848,7 @@ function _M.calculate_cost(query_body, tokens_models, tokens_factor)
   return query_cost, nil
 end
 
+-- for unit tests
+_M._count_words = count_words
+
 return _M

spec/03-plugins/38-ai-proxy/01-unit_spec.lua

@@ -732,4 +732,47 @@ describe(PLUGIN_NAME .. ": (unit)", function()
 
   end)
 
+  describe("count_words", function()
+    local c = ai_shared._count_words
+
+    it("normal prompts", function()
+      assert.same(10, c(string.rep("apple ", 10)))
+    end)
+
+    it("multi-modal prompts", function()
+      assert.same(10, c({
+        {
+          type = "text",
+          text = string.rep("apple ", 10),
+        },
+      }))
+
+      assert.same(20, c({
+        {
+          type = "text",
+          text = string.rep("apple ", 10),
+        },
+        {
+          type = "text",
+          text = string.rep("banana ", 10),
+        },
+      }))
+
+      assert.same(10, c({
+        {
+          type = "not_text",
+          text = string.rep("apple ", 10),
+        },
+        {
+          type = "text",
+          text = string.rep("banana ", 10),
+        },
+        {
+          type = "text",
+          -- somehow malformed
+        },
+      }))
+    end)
+  end)
+
 end)