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 0f3b15bbdee3..009f079195d0 100644 --- a/spec/03-plugins/38-ai-proxy/01-unit_spec.lua +++ b/spec/03-plugins/38-ai-proxy/01-unit_spec.lua @@ -716,6 +716,20 @@ describe(PLUGIN_NAME .. ": (unit)", function() assert.same(events, expected_events) end) + it("transforms application/vnd.amazon.eventstream (AWS) type", function() + local input = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/aws/input.bin")) + local events = ai_shared.frame_to_events(input, "bedrock") + + local expected = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/aws/expected-output.json")) + local expected_events = cjson.decode(expected) + + assert.equal(#events, #expected_events) + for i, _ in ipairs(expected_events) do + -- tables are random ordered, so we need to compare each serialized event + assert.same(cjson.decode(events[i].data), cjson.decode(expected_events[i].data)) + end + end) + end) end) diff --git a/spec/fixtures/ai-proxy/unit/streaming-chunk-formats/aws/expected-output.json b/spec/fixtures/ai-proxy/unit/streaming-chunk-formats/aws/expected-output.json new file mode 100644 index 000000000000..8761c5593608 --- /dev/null +++ b/spec/fixtures/ai-proxy/unit/streaming-chunk-formats/aws/expected-output.json @@ -0,0 +1,20 @@ +[ + { + "data": "{\"body\":\"{\\\"p\\\":\\\"abcdefghijkl\\\",\\\"role\\\":\\\"assistant\\\"}\",\"headers\":{\":event-type\":\"messageStart\",\":content-type\":\"application\/json\",\":message-type\":\"event\"}}" + }, + { + "data": "{\"body\":\"{\\\"contentBlockIndex\\\":0,\\\"delta\\\":{\\\"text\\\":\\\"Hello! Relativity is a set of physical theories that are collectively known as special relativity and general relativity, proposed by Albert Einstein. These theories revolutionized our understanding of space, time, and gravity, and have had far-reach\\\"},\\\"p\\\":\\\"abcd\\\"}\",\"headers\":{\":event-type\":\"contentBlockDelta\",\":content-type\":\"application\\/json\",\":message-type\":\"event\"}}" + }, + { + "data": "{\"headers\":{\":event-type\":\"contentBlockDelta\",\":message-type\":\"event\",\":content-type\":\"application\\/json\"},\"body\":\"{\\\"contentBlockIndex\\\":0,\\\"delta\\\":{\\\"text\\\":\\\"ing implications in various scientific and technological fields. Special relativity applies to all physical phenomena in the absence of gravity, while general relativity explains the law of gravity and its effects on the nature of space, time, and matter.\\\"},\\\"p\\\":\\\"abcdefghijk\\\"}\"}" + }, + { + "data": "{\"body\":\"{\\\"contentBlockIndex\\\":0,\\\"p\\\":\\\"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQR\\\"}\",\"headers\":{\":content-type\":\"application\\/json\",\":event-type\":\"contentBlockStop\",\":message-type\":\"event\"}}" + }, + { + "data": "{\"body\":\"{\\\"p\\\":\\\"abcdefghijklm\\\",\\\"stopReason\\\":\\\"end_turn\\\"}\",\"headers\":{\":message-type\":\"event\",\":content-type\":\"application\\/json\",\":event-type\":\"messageStop\"}}" + }, + { + "data": "{\"headers\":{\":message-type\":\"event\",\":content-type\":\"application\\/json\",\":event-type\":\"metadata\"},\"body\":\"{\\\"metrics\\\":{\\\"latencyMs\\\":2613},\\\"p\\\":\\\"abcdefghijklmnopqrstuvwxyzABCDEF\\\",\\\"usage\\\":{\\\"inputTokens\\\":9,\\\"outputTokens\\\":97,\\\"totalTokens\\\":106}}\"}" + } +] \ No newline at end of file diff --git a/spec/fixtures/ai-proxy/unit/streaming-chunk-formats/aws/input.bin b/spec/fixtures/ai-proxy/unit/streaming-chunk-formats/aws/input.bin new file mode 100644 index 000000000000..8f9d03b4f7e0 Binary files /dev/null and b/spec/fixtures/ai-proxy/unit/streaming-chunk-formats/aws/input.bin differ