restapi.moon.1

ltn12 = require "ltn12"
json = json or require "cjson"

STREAMED = "Streamed successfully"

class Plain
  new: (@api_key, apibase, @headers) =>
    if apibase
      @api_base = apibase

  create_stream_sink: (delta_callback, ctx) =>
    streamed = false
    (chunk, err) ->
      if chunk==nil
        if err
          delta_callback nil, err
        else
          ctx.response = streamed and STREAMED or ""
      else
        delta_callback chunk
        streamed = true

      chunk, err

  authorization: { "authorization", "Bearer %s" }

  _request: (method, path, payload, delta_callback, extra_headers) =>
    headers = "accept": delta_callback and "text/event-stream" or "application/json"
    local source
    if payload
      body = payload
      if type(payload)=="table"
        body = json.encode payload
        headers["content-type"] = "application/json"
      headers["content-length"] = #body
      source = ltn12.source.string body

    if @api_key
      {name, template} = @authorization
      headers[name or "authorization"] = template and template\format(@api_key) or @api_key

    for _headers in *{ @headers, extra_headers }
      if _headers then for k,v in pairs _headers
        headers[k\lower()] = v or nil -- false becomes nil

    out = {}
    sink = if delta_callback
      @create_stream_sink delta_callback, out
    else
      ltn12.sink.table out

    success, status_code, response_headers, status_line = @get_provider!.request {
      url: @api_base .. path
      :sink
      :source
      :method
      :headers
    }
    if success
      response = if delta_callback then out.response else table.concat out
      pcall -> response = assert json.decode response
      response, status_code, response_headers, status_line
    else
      nil, status_code

  request: (method, endpoint, payload, delta_callback, headers) =>
    if method\match"^[/?]"
      endpoint, payload, delta_callback, headers = method, endpoint, payload, delta_callback
      method = payload and "POST" or "GET"

    response, status_code, headers, statusline = @_request method, endpoint, payload, delta_callback, headers
    if not response
      nil, status_code
    elseif status_code<200 or status_code>299
      nil, statusline, response~=STREAMED and response or nil
    else
      response, statusline, headers

  generate: (endpoint, payload, delta_callback, headers) =>
    @request "POST", endpoint, payload, delta_callback, headers

  providers: {
    http: "socket.http"
    https: "ssl.https"
  }

  get_provider: =>
    unless @provider
      scheme = @api_base\match"^(.-)://"
      provider = scheme and @providers[scheme]
      if provider
        @provider = require provider
      else
        error "unknown scheme: "..@api_base

    @provider


class LinesStream extends Plain
  separatorPattern: "^(.-)\r?\n(.-)$"

  consumeBufferRest: false

  create_stream_sink: (delta_callback, ctx) =>
    acc_buffer, streamed, accumulate_only = "", false, false

    (chunk, err) ->
      if chunk==nil
        if streamed
          if not err and @consumeBufferRest and acc_buffer~=""
            @process_line acc_buffer, delta_callback, ctx
          elseif err = err or acc_buffer~="" and not acc_buffer\match"data: ?%[DONE%]" and "buffer not empty"
            delta_callback nil, err, acc_buffer

        unless err
          ctx.response = streamed and STREAMED or acc_buffer

      else
        acc_buffer ..= chunk
        unless accumulate_only
          line = true
          while line
            line, rest = string.match acc_buffer, @separatorPattern
            if line
              if @process_line line, delta_callback, ctx
                acc_buffer = rest
                streamed = true
              elseif streamed
                delta_callback nil, "Unable to parse line", line
                return nil, "Unable to parse line"
              else
                accumulate_only = true
                break

      chunk, err

  process_line: (line, delta_callback, ctx) =>
    if line=="" then return true
    @process delta_callback, line

  process: (delta_callback, data, extra) =>
    if data=="[DONE]" or data\match"^ -$" then return true
    success, parsed = pcall -> assert json.decode data
    if success
      if @is_valid parsed
        delta_callback parsed, nil, extra
      else
        delta_callback nil, "Unexpected json chunk", parsed

      true

  is_valid: (parsed) =>
    true -- stub

  defaultEndpoint: "/chat"

  generate: (endpoint, payload, delta_callback, headers) =>
    if type(endpoint)~="string"
      endpoint, payload, delta_callback, headers = @defaultEndpoint, endpoint, payload, delta_callback
    response, status_code, _, statusline = @_request "POST", endpoint, payload, delta_callback, headers
    if not response
      nil, status_code
    elseif status_code<200 or status_code>299
      nil, statusline, response~=STREAMED and response or nil
    elseif response~=STREAMED and (type(response)~="table" or not (@is_valid_final or @is_valid) @,response)
      nil, statusline, response
    else
      response, statusline

  --modelsKeys: {models:"data", id:"id"}

  models: (endpoint) =>
    response, status_code, _, statusline = @_request "GET", endpoint or "/models"
    if not response
      nil, status_code
    elseif status_code==200
      if type(response)~="table"
        return nil, response
      if @modelsKeys then
        keys = @modelsKeys
        models = keys.models and response[keys.models] or response
        models.keyId = keys.id
        return models
      -- else use heuristics
      models = response.data   or -- openai
               response.models or -- cohere, google gemini
               response.result or -- cloudflare
               response           -- together, github
      item = models[2]
      if not item
        return nil, "Unexpected response"
      keyId = item.name and not item.name\match" " and "name" or -- cloudflare, cohere, github, google gemini
              item.id and "id"  -- openai, copilot, openrouter
      if not keyId
        return nil, "Unexpected response"
      models.keyId = keyId
      models
    else
      nil, statusline, response or ""


class BasicSSE extends LinesStream
  process_line: (line, delta_callback) =>
    if line=="" or line\sub(1,1)==":" or line\match"^event:"
      true
    elseif data = line\match"^data: ?(.+)$"
      @process delta_callback, data


class SSE extends LinesStream
  process_line: (line, delta_callback, ctx) =>
    ctx.event = ctx.event or {}
    e = ctx.event

    -- https://html.spec.whatwg.org/multipage/server-sent-events.html#event-stream-interpretation
    if line==""
      if next e
        ctx.event = {} -- reset
        if e.data
          return @process delta_callback, e.data, e
        else
          unless e[""] -- comment
            delta_callback nil, "data not found", e
    else
      field, value = line\match"^(.-):%s?(.+)$"
      switch field
        when "data"
          e.data = e.data and e.data.."\n"..value or value
        when "event","id","retry","" --comment
          e[field] = value
        when nil
          unless ({event:1,id:1,retry:1,data:1})[line]
            return false -- do not try to parse SSE anymore
        else
          if field\match"^%w+$"
            delta_callback nil, 'Unknown field "%s" with value "%s"'\format(field, value), e
          return nil

    true


class OpenAI extends BasicSSE
  api_base: "https://api.openai.com/v1"

  defaultEndpoint: "/chat/completions"

  generate: (payload, delta_callback=nil) =>
    super payload, if payload.stream then delta_callback

  is_valid: (parsed) =>
    parsed.choices


class GoogleGemini extends BasicSSE
  api_base: "https://generativelanguage.googleapis.com/v1beta"

  authorization: { "x-goog-api-key" }

  generate: (payload, model, stream, delta_callback=nil) =>
    assert model, "Model not specified"
    task = stream and "streamGenerateContent" or "generateContent"
    endpoint = "/models/%s:%s%s"\format model, task, stream and "?alt=sse" or ""
    super endpoint, payload, if stream then delta_callback

  is_valid: (parsed) =>
    parsed.candidates

  modelsKeys: {models:"models", id:"name"}


{
  :STREAMED
  :Plain
  :LinesStream
  :BasicSSE
  :SSE
  :OpenAI
  :GoogleGemini
}