Skip to content

Commit

Permalink
feat(ai-proxy): stream format tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@tysoekong
tysoekong committed Jul 5, 2024
1 parent 65b0d8f commit 6c2cc3e
Show file tree
Hide file tree
Showing 8 changed files with 290 additions and 4 deletions.
5 changes: 4 additions & 1 deletion kong/llm/drivers/shared.lua
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
28 changes: 25 additions & 3 deletions spec/03-plugins/38-ai-proxy/01-unit_spec.lua
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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)
Expand Down
14 changes: 14 additions & 0 deletions ...ixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-beginning/expected-output.json
View file
Original file line number Diff line number Diff line change
@@ -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"
}
]
141 changes: 141 additions & 0 deletions spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-beginning/input.bin
View file
Original file line number Diff line number Diff line change
@@ -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
}
}
8 changes: 8 additions & 0 deletions spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-end/expected-output.json
View file
Original file line number Diff line number Diff line change
@@ -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"
}
]
80 changes: 80 additions & 0 deletions spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-end/input.bin
View file
Original file line number Diff line number Diff line change
@@ -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
}
}
]
11 changes: 11 additions & 0 deletions spec/fixtures/ai-proxy/unit/streaming-chunk-formats/text-event-stream/expected-output.json
View file
Original file line number Diff line number Diff line change
@@ -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}"
}
]
7 changes: 7 additions & 0 deletions spec/fixtures/ai-proxy/unit/streaming-chunk-formats/text-event-stream/input.bin
View file
Original file line number Diff line number Diff line change
@@ -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]

0 comments on commit 6c2cc3e

Please sign in to comment.