diff --git a/kong/llm/drivers/shared.lua b/kong/llm/drivers/shared.lua index 25d9d5773149..0e1d0d18a962 100644 --- a/kong/llm/drivers/shared.lua +++ b/kong/llm/drivers/shared.lua @@ -267,7 +267,10 @@ function _M.frame_to_events(frame, raw_json_mode) -- test for truncated chunk on the last line (no trailing \r\n\r\n) if #dat > 0 and #event_lines == i then ngx.log(ngx.DEBUG, "[ai-proxy] truncated sse frame head") - kong.ctx.plugin.truncated_frame = dat + if kong then + kong.ctx.plugin.truncated_frame = dat + end + break -- stop parsing immediately, server has done something wrong end diff --git a/spec/03-plugins/38-ai-proxy/01-unit_spec.lua b/spec/03-plugins/38-ai-proxy/01-unit_spec.lua index 5173739af8f1..22fb1e668e34 100644 --- a/spec/03-plugins/38-ai-proxy/01-unit_spec.lua +++ b/spec/03-plugins/38-ai-proxy/01-unit_spec.lua @@ -676,8 +676,24 @@ describe(PLUGIN_NAME .. ": (unit)", function() describe("streaming transformer tests", function() - it("transforms truncated-json type", function() - + it("transforms truncated-json type (beginning of stream)", function() + local input = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-beginning/input.bin")) + local events = ai_shared.frame_to_events(input, true) + + local expected = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-beginning/expected-output.json")) + local expected_events = cjson.decode(expected) + + assert.same(events, expected_events, true) + end) + + it("transforms truncated-json type (end of stream)", function() + local input = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-end/input.bin")) + local events = ai_shared.frame_to_events(input, true) + + local expected = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-end/expected-output.json")) + local expected_events = cjson.decode(expected) + + assert.same(events, expected_events, true) end) it("transforms complete-json type", function() @@ -691,7 +707,13 @@ describe(PLUGIN_NAME .. ": (unit)", function() end) it("transforms text/event-stream type", function() - + local input = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/text-event-stream/input.bin")) + local events = ai_shared.frame_to_events(input, false) -- not "truncated json mode" like Gemini + + local expected = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/text-event-stream/expected-output.json")) + local expected_events = cjson.decode(expected) + + assert.same(events, expected_events) end) end) diff --git a/spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-beginning/expected-output.json b/spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-beginning/expected-output.json new file mode 100644 index 000000000000..5f3b0afa51d4 --- /dev/null +++ b/spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-beginning/expected-output.json @@ -0,0 +1,14 @@ +[ + { + "data": "{\n \"candidates\": [\n {\n \"content\": {\n \"parts\": [\n {\n \"text\": \"The\"\n }\n ],\n \"role\": \"model\"\n },\n \"finishReason\": \"STOP\",\n \"index\": 0\n }\n ],\n \"usageMetadata\": {\n \"promptTokenCount\": 6,\n \"candidatesTokenCount\": 1,\n \"totalTokenCount\": 7\n }\n}" + }, + { + "data": "\n{\n \"candidates\": [\n {\n \"content\": {\n \"parts\": [\n {\n \"text\": \" theory of relativity is actually two theories by Albert Einstein: **special relativity** and\"\n }\n ],\n \"role\": \"model\"\n },\n \"finishReason\": \"STOP\",\n \"index\": 0,\n \"safetyRatings\": [\n {\n \"category\": \"HARM_CATEGORY_SEXUALLY_EXPLICIT\",\n \"probability\": \"NEGLIGIBLE\"\n },\n {\n \"category\": \"HARM_CATEGORY_HATE_SPEECH\",\n \"probability\": \"NEGLIGIBLE\"\n },\n {\n \"category\": \"HARM_CATEGORY_HARASSMENT\",\n \"probability\": \"NEGLIGIBLE\"\n },\n {\n \"category\": \"HARM_CATEGORY_DANGEROUS_CONTENT\",\n \"probability\": \"NEGLIGIBLE\"\n }\n ]\n }\n ],\n \"usageMetadata\": {\n \"promptTokenCount\": 6,\n \"candidatesTokenCount\": 17,\n \"totalTokenCount\": 23\n }\n}" + }, + { + "data": "\n{\n \"candidates\": [\n {\n \"content\": {\n \"parts\": [\n {\n \"text\": \" **general relativity**. Here's a simplified breakdown:\\n\\n**Special Relativity (\"\n }\n ],\n \"role\": \"model\"\n },\n \"finishReason\": \"STOP\",\n \"index\": 0,\n \"safetyRatings\": [\n {\n \"category\": \"HARM_CATEGORY_SEXUALLY_EXPLICIT\",\n \"probability\": \"NEGLIGIBLE\"\n },\n {\n \"category\": \"HARM_CATEGORY_HATE_SPEECH\",\n \"probability\": \"NEGLIGIBLE\"\n },\n {\n \"category\": \"HARM_CATEGORY_HARASSMENT\",\n \"probability\": \"NEGLIGIBLE\"\n },\n {\n \"category\": \"HARM_CATEGORY_DANGEROUS_CONTENT\",\n \"probability\": \"NEGLIGIBLE\"\n }\n ]\n }\n ],\n \"usageMetadata\": {\n \"promptTokenCount\": 6,\n \"candidatesTokenCount\": 33,\n \"totalTokenCount\": 39\n }\n}" + }, + { + "data": "\n{\n \"candidates\": [\n {\n \"content\": {\n \"parts\": [\n {\n \"text\": \"1905):**\\n\\n* **Focus:** The relationship between space and time.\\n* **Key ideas:**\\n * **Speed of light\"\n }\n ],\n \"role\": \"model\"\n },\n \"finishReason\": \"STOP\",\n \"index\": 0,\n \"safetyRatings\": [\n {\n \"category\": \"HARM_CATEGORY_SEXUALLY_EXPLICIT\",\n \"probability\": \"NEGLIGIBLE\"\n },\n {\n \"category\": \"HARM_CATEGORY_HATE_SPEECH\",\n \"probability\": \"NEGLIGIBLE\"\n },\n {\n \"category\": \"HARM_CATEGORY_HARASSMENT\",\n \"probability\": \"NEGLIGIBLE\"\n },\n {\n \"category\": \"HARM_CATEGORY_DANGEROUS_CONTENT\",\n \"probability\": \"NEGLIGIBLE\"\n }\n ]\n }\n ],\n \"usageMetadata\": {\n \"promptTokenCount\": 6,\n \"candidatesTokenCount\": 65,\n \"totalTokenCount\": 71\n }\n}\n" + } +] \ No newline at end of file diff --git a/spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-beginning/input.bin b/spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-beginning/input.bin new file mode 100644 index 000000000000..8cef2a01fa8d --- /dev/null +++ b/spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-beginning/input.bin @@ -0,0 +1,141 @@ +[{ + "candidates": [ + { + "content": { + "parts": [ + { + "text": "The" + } + ], + "role": "model" + }, + "finishReason": "STOP", + "index": 0 + } + ], + "usageMetadata": { + "promptTokenCount": 6, + "candidatesTokenCount": 1, + "totalTokenCount": 7 + } +} +, +{ + "candidates": [ + { + "content": { + "parts": [ + { + "text": " theory of relativity is actually two theories by Albert Einstein: **special relativity** and" + } + ], + "role": "model" + }, + "finishReason": "STOP", + "index": 0, + "safetyRatings": [ + { + "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "NEGLIGIBLE" + } + ] + } + ], + "usageMetadata": { + "promptTokenCount": 6, + "candidatesTokenCount": 17, + "totalTokenCount": 23 + } +} +, +{ + "candidates": [ + { + "content": { + "parts": [ + { + "text": " **general relativity**. Here's a simplified breakdown:\n\n**Special Relativity (" + } + ], + "role": "model" + }, + "finishReason": "STOP", + "index": 0, + "safetyRatings": [ + { + "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "NEGLIGIBLE" + } + ] + } + ], + "usageMetadata": { + "promptTokenCount": 6, + "candidatesTokenCount": 33, + "totalTokenCount": 39 + } +} +, +{ + "candidates": [ + { + "content": { + "parts": [ + { + "text": "1905):**\n\n* **Focus:** The relationship between space and time.\n* **Key ideas:**\n * **Speed of light" + } + ], + "role": "model" + }, + "finishReason": "STOP", + "index": 0, + "safetyRatings": [ + { + "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "NEGLIGIBLE" + } + ] + } + ], + "usageMetadata": { + "promptTokenCount": 6, + "candidatesTokenCount": 65, + "totalTokenCount": 71 + } +} diff --git a/spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-end/expected-output.json b/spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-end/expected-output.json new file mode 100644 index 000000000000..ba6a64384d95 --- /dev/null +++ b/spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-end/expected-output.json @@ -0,0 +1,8 @@ +[ + { + "data": "{\n \"candidates\": [\n {\n \"content\": {\n \"parts\": [\n {\n \"text\": \" is constant:** No matter how fast you are moving, light always travels at the same speed (approximately 299,792,458\"\n }\n ],\n \"role\": \"model\"\n },\n \"finishReason\": \"STOP\",\n \"index\": 0,\n \"safetyRatings\": [\n {\n \"category\": \"HARM_CATEGORY_SEXUALLY_EXPLICIT\",\n \"probability\": \"NEGLIGIBLE\"\n },\n {\n \"category\": \"HARM_CATEGORY_HATE_SPEECH\",\n \"probability\": \"NEGLIGIBLE\"\n },\n {\n \"category\": \"HARM_CATEGORY_HARASSMENT\",\n \"probability\": \"NEGLIGIBLE\"\n },\n {\n \"category\": \"HARM_CATEGORY_DANGEROUS_CONTENT\",\n \"probability\": \"NEGLIGIBLE\"\n }\n ]\n }\n ],\n \"usageMetadata\": {\n \"promptTokenCount\": 6,\n \"candidatesTokenCount\": 97,\n \"totalTokenCount\": 103\n }\n}" + }, + { + "data": "\n{\n \"candidates\": [\n {\n \"content\": {\n \"parts\": [\n {\n \"text\": \" not a limit.\\n\\nIf you're interested in learning more about relativity, I encourage you to explore further resources online or in books. There are many excellent introductory materials available. \\n\"\n }\n ],\n \"role\": \"model\"\n },\n \"finishReason\": \"STOP\",\n \"index\": 0,\n \"safetyRatings\": [\n {\n \"category\": \"HARM_CATEGORY_SEXUALLY_EXPLICIT\",\n \"probability\": \"NEGLIGIBLE\"\n },\n {\n \"category\": \"HARM_CATEGORY_HATE_SPEECH\",\n \"probability\": \"NEGLIGIBLE\"\n },\n {\n \"category\": \"HARM_CATEGORY_HARASSMENT\",\n \"probability\": \"NEGLIGIBLE\"\n },\n {\n \"category\": \"HARM_CATEGORY_DANGEROUS_CONTENT\",\n \"probability\": \"NEGLIGIBLE\"\n }\n ]\n }\n ],\n \"usageMetadata\": {\n \"promptTokenCount\": 6,\n \"candidatesTokenCount\": 547,\n \"totalTokenCount\": 553\n }\n}\n" + } +] \ No newline at end of file diff --git a/spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-end/input.bin b/spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-end/input.bin new file mode 100644 index 000000000000..d6489e74d19d --- /dev/null +++ b/spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-end/input.bin @@ -0,0 +1,80 @@ +,{ + "candidates": [ + { + "content": { + "parts": [ + { + "text": " is constant:** No matter how fast you are moving, light always travels at the same speed (approximately 299,792,458" + } + ], + "role": "model" + }, + "finishReason": "STOP", + "index": 0, + "safetyRatings": [ + { + "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "NEGLIGIBLE" + } + ] + } + ], + "usageMetadata": { + "promptTokenCount": 6, + "candidatesTokenCount": 97, + "totalTokenCount": 103 + } +} +, +{ + "candidates": [ + { + "content": { + "parts": [ + { + "text": " not a limit.\n\nIf you're interested in learning more about relativity, I encourage you to explore further resources online or in books. There are many excellent introductory materials available. \n" + } + ], + "role": "model" + }, + "finishReason": "STOP", + "index": 0, + "safetyRatings": [ + { + "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "NEGLIGIBLE" + } + ] + } + ], + "usageMetadata": { + "promptTokenCount": 6, + "candidatesTokenCount": 547, + "totalTokenCount": 553 + } +} +] \ No newline at end of file diff --git a/spec/fixtures/ai-proxy/unit/streaming-chunk-formats/text-event-stream/expected-output.json b/spec/fixtures/ai-proxy/unit/streaming-chunk-formats/text-event-stream/expected-output.json new file mode 100644 index 000000000000..f515516c7ec8 --- /dev/null +++ b/spec/fixtures/ai-proxy/unit/streaming-chunk-formats/text-event-stream/expected-output.json @@ -0,0 +1,11 @@ +[ + { + "data": "{ \"choices\": [ { \"delta\": { \"content\": \"\", \"role\": \"assistant\" }, \"finish_reason\": null, \"index\": 0, \"logprobs\": null } ], \"created\": 1720136012, \"id\": \"chatcmpl-9hQFArK1oMZcRwaIa86RGwrjVNmeY\", \"model\": \"gpt-4-0613\", \"object\": \"chat.completion.chunk\", \"system_fingerprint\": null}" + }, + { + "data": "{ \"choices\": [ { \"delta\": { \"content\": \"2\" }, \"finish_reason\": null, \"index\": 0, \"logprobs\": null } ], \"created\": 1720136012, \"id\": \"chatcmpl-9hQFArK1oMZcRwaIa86RGwrjVNmeY\", \"model\": \"gpt-4-0613\", \"object\": \"chat.completion.chunk\", \"system_fingerprint\": null}" + }, + { + "data": "{ \"choices\": [ { \"delta\": {}, \"finish_reason\": \"stop\", \"index\": 0, \"logprobs\": null } ], \"created\": 1720136012, \"id\": \"chatcmpl-9hQFArK1oMZcRwaIa86RGwrjVNmeY\", \"model\": \"gpt-4-0613\", \"object\": \"chat.completion.chunk\", \"system_fingerprint\": null}" + } +] \ No newline at end of file diff --git a/spec/fixtures/ai-proxy/unit/streaming-chunk-formats/text-event-stream/input.bin b/spec/fixtures/ai-proxy/unit/streaming-chunk-formats/text-event-stream/input.bin new file mode 100644 index 000000000000..efe2ad50c657 --- /dev/null +++ b/spec/fixtures/ai-proxy/unit/streaming-chunk-formats/text-event-stream/input.bin @@ -0,0 +1,7 @@ +data: { "choices": [ { "delta": { "content": "", "role": "assistant" }, "finish_reason": null, "index": 0, "logprobs": null } ], "created": 1720136012, "id": "chatcmpl-9hQFArK1oMZcRwaIa86RGwrjVNmeY", "model": "gpt-4-0613", "object": "chat.completion.chunk", "system_fingerprint": null} + +data: { "choices": [ { "delta": { "content": "2" }, "finish_reason": null, "index": 0, "logprobs": null } ], "created": 1720136012, "id": "chatcmpl-9hQFArK1oMZcRwaIa86RGwrjVNmeY", "model": "gpt-4-0613", "object": "chat.completion.chunk", "system_fingerprint": null} + +data: { "choices": [ { "delta": {}, "finish_reason": "stop", "index": 0, "logprobs": null } ], "created": 1720136012, "id": "chatcmpl-9hQFArK1oMZcRwaIa86RGwrjVNmeY", "model": "gpt-4-0613", "object": "chat.completion.chunk", "system_fingerprint": null} + +data: [DONE] \ No newline at end of file