test.lua

local cjson = require("cjson")
local websocket = require("resty.websocket.client")

local retried = 0
local conn = assert(websocket:new({max_payload_len = 1024 * 1024 * 8}))
assert(conn:connect("ws://localhost:8046/cgemma/session"))
while true do
  local data, tp, err = conn:recv_frame()
  while err == "again" do
    local frag, ct
    frag, ct, err = conn:recv_frame()
    if ct ~= "continuation" then
      conn:close()
      error(err)
    end
    data = data..frag
  end
  if tp == "text" then
    print(data)
    io.flush()
    local msg = cjson.decode(data)
    if msg.op == "init" then
      local bytes, err = conn:send_text(cjson.encode({
        op = "create",
        image = ngx.encode_base64(io.read("*a")),
        lang = "zh",
        poet = "海子"
      }))
      if not bytes then
        error(err)
      end
    elseif msg.op == "stream" and not msg.token then
      if retried == 0 then
        assert(conn:send_text(cjson.encode({op = "retry"})))
      elseif retried == 1 then
        assert(conn:send_text(cjson.encode({
          op = "retry",
          lang = "en",
          poet = "Emily Dickinson"
        })))
      elseif retried == 2 then
        assert(conn:send_text(cjson.encode({op = "retry"})))
      else
        conn:close()
        return
      end
      retried = retried + 1
    end
  elseif tp == "ping" then
    local res, err = conn:send_pong()
    if not res then
      conn:close()
      error(err)
    end
  elseif tp == "close" then
    conn:close()
    return
  elseif tp ~= "pong" then
    if err then
      if conn.fatal then
        conn:close()
        error(err)
      end
    else
      local res, err = conn:send_close(1003, "unsupported data type")
      if not res then
        conn:close()
        error(err)
      end
      conn:close()
      error("unsupported data type")
    end
  end
end