|
| 1 | +local tablex = require "pl.tablex" |
| 2 | + |
| 3 | +describe("stream decoder", function() |
| 4 | + setup(function() |
| 5 | + AWS = require "resty.aws" |
| 6 | + end) |
| 7 | + |
| 8 | + teardown(function() |
| 9 | + AWS = nil |
| 10 | + package.loaded["resty.aws"] = nil |
| 11 | + end) |
| 12 | + |
| 13 | + |
| 14 | + it("decodes multiple messages from a single chunk", function() |
| 15 | + local chunk_in_hex = "000000760000005296d5fade0b3a6576656e742d7479706507000c6d65737361676553746172740d3a636f6e74656e742d747970650700106170706c69636174696f6e2f6a736f6e0d3a6d6573736167652d747970650700056576656e747b22726f6c65223a22617373697374616e74227d6fa8c599000001110000005706f176a90b3a6576656e742d74797065070011636f6e74656e74426c6f636b44656c74610d3a636f6e74656e742d747970650700106170706c69636174696f6e2f6a736f6e0d3a6d6573736167652d747970650700056576656e747b22636f6e74656e74426c6f636b496e646578223a302c2264656c7461223a7b2274657874223a2250692028cf80292069732061206d617468656d61746963616c20636f6e7374616e74207468617420726570726573656e7473207468652063697263756d666572656e63652d746f2d6469616d6574657220726174696f206f66206120636972636c652e204974277320617070726f78696d6174656c7920332e31343135392e227d7d7d4e31940000007d00000056e6680fd60b3a6576656e742d74797065070010636f6e74656e74426c6f636b53746f700d3a636f6e74656e742d747970650700106170706c69636174696f6e2f6a736f6e0d3a6d6573736167652d747970650700056576656e747b22636f6e74656e74426c6f636b496e646578223a307db32e4d340000007a00000051ca2c46650b3a6576656e742d7479706507000b6d65737361676553746f700d3a636f6e74656e742d747970650700106170706c69636174696f6e2f6a736f6e0d3a6d6573736167652d747970650700056576656e747b2273746f70526561736f6e223a22656e645f7475726e227d5fbf09fc000000b90000004ee991d99b0b3a6576656e742d747970650700086d657461646174610d3a636f6e74656e742d747970650700106170706c69636174696f6e2f6a736f6e0d3a6d6573736167652d747970650700056576656e747b226d657472696373223a7b226c6174656e63794d73223a313530367d2c227573616765223a7b22696e707574546f6b656e73223a382c226f7574707574546f6b656e73223a33392c22746f74616c546f6b656e73223a34377d7d5ad2dd7b" |
| 16 | + local _STREAM = require("resty.aws.stream") |
| 17 | + local parser, err = _STREAM:new(chunk_in_hex, true) |
| 18 | + |
| 19 | + if err then |
| 20 | + assert.equal(err, nil) |
| 21 | + return |
| 22 | + end |
| 23 | + |
| 24 | + local messages = {} |
| 25 | + |
| 26 | + while true do |
| 27 | + local msg = parser:next_message() |
| 28 | + |
| 29 | + if not msg then |
| 30 | + break |
| 31 | + end |
| 32 | + |
| 33 | + messages[#messages+1] = msg |
| 34 | + end |
| 35 | + |
| 36 | + assert.same(messages, { |
| 37 | + { |
| 38 | + body = "{\"role\":\"assistant\"}", |
| 39 | + headers = { |
| 40 | + { |
| 41 | + key = ":event-type", |
| 42 | + value = "messageStart" |
| 43 | + }, |
| 44 | + { |
| 45 | + key = ":content-type", |
| 46 | + value = "application/json" |
| 47 | + }, |
| 48 | + { |
| 49 | + key = ":message-type", |
| 50 | + value = "event" |
| 51 | + } |
| 52 | + } |
| 53 | + }, |
| 54 | + { |
| 55 | + body = "{\"contentBlockIndex\":0,\"delta\":{\"text\":\"Pi (π) is a mathematical constant that represents the circumference-to-diameter ratio of a circle. It's approximately 3.14159.\"}}", |
| 56 | + headers = { |
| 57 | + { |
| 58 | + key = ":event-type", |
| 59 | + value = "contentBlockDelta" |
| 60 | + }, |
| 61 | + { |
| 62 | + key = ":content-type", |
| 63 | + value = "application/json" |
| 64 | + }, |
| 65 | + { |
| 66 | + key = ":message-type", |
| 67 | + value = "event" |
| 68 | + } |
| 69 | + } |
| 70 | + }, |
| 71 | + { |
| 72 | + body = "{\"contentBlockIndex\":0}", |
| 73 | + headers = { |
| 74 | + { |
| 75 | + key = ":event-type", |
| 76 | + value = "contentBlockStop" |
| 77 | + }, |
| 78 | + { |
| 79 | + key = ":content-type", |
| 80 | + value = "application/json" |
| 81 | + }, |
| 82 | + { |
| 83 | + key = ":message-type", |
| 84 | + value = "event" |
| 85 | + } |
| 86 | + } |
| 87 | + }, |
| 88 | + { |
| 89 | + body = "{\"stopReason\":\"end_turn\"}", |
| 90 | + headers = { |
| 91 | + { |
| 92 | + key = ":event-type", |
| 93 | + value = "messageStop" |
| 94 | + }, |
| 95 | + { |
| 96 | + key = ":content-type", |
| 97 | + value = "application/json" |
| 98 | + }, |
| 99 | + { |
| 100 | + key = ":message-type", |
| 101 | + value = "event" |
| 102 | + } |
| 103 | + } |
| 104 | + }, |
| 105 | + { |
| 106 | + body = "{\"metrics\":{\"latencyMs\":1506},\"usage\":{\"inputTokens\":8,\"outputTokens\":39,\"totalTokens\":47}}", |
| 107 | + headers = { |
| 108 | + { |
| 109 | + key = ":event-type", |
| 110 | + value = "metadata" |
| 111 | + }, |
| 112 | + { |
| 113 | + key = ":content-type", |
| 114 | + value = "application/json" |
| 115 | + }, |
| 116 | + { |
| 117 | + key = ":message-type", |
| 118 | + value = "event" |
| 119 | + } |
| 120 | + } |
| 121 | + } |
| 122 | + }) |
| 123 | + end) |
| 124 | +end) |
0 commit comments