migrate the chat completions api details into separate module

Author: leafo

lua-openai-dev-1.rockspec

@@ -24,7 +24,7 @@ build = {
    type = "builtin",
    modules = {
       ["openai"] = "openai/init.lua",
-      ["openai.chat_session"] = "openai/chat_session.lua",
+      ["openai.chat_completions"] = "openai/chat_completions.lua",
       ["openai.responses"] = "openai/responses.lua"
    }
 }

openai/chat_session.lua openai/chat_completions.luaopenai/chat_session.lua renamed to openai/chat_completions.lua

@@ -65,6 +65,65 @@ local parse_error_message = types.partial({
     code = empty + types.string:tag("code")
   })
 })
+local parse_completion_chunk = types.partial({
+  object = "chat.completion.chunk",
+  choices = types.shape({
+    types.partial({
+      delta = types.partial({
+        ["content"] = types.string:tag("content")
+      }),
+      index = types.number:tag("index")
+    })
+  })
+})
+local consume_json_head
+do
+  local C, S, P
+  do
+    local _obj_0 = require("lpeg")
+    C, S, P = _obj_0.C, _obj_0.S, _obj_0.P
+  end
+  local consume_json = P(function(str, pos)
+    local str_len = #str
+    for k = pos + 1, str_len do
+      local candidate = str:sub(pos, k)
+      local parsed = false
+      pcall(function()
+        parsed = cjson.decode(candidate)
+      end)
+      if parsed then
+        return k + 1
+      end
+    end
+    return nil
+  end)
+  consume_json_head = S("\t\n\r ") ^ 0 * P("data: ") * C(consume_json) * C(P(1) ^ 0)
+end
+local create_chat_stream_filter
+create_chat_stream_filter = function(chunk_callback)
+  assert(types["function"](chunk_callback), "Must provide chunk_callback function when streaming response")
+  local accumulation_buffer = ""
+  return function(...)
+    local chunk = ...
+    if type(chunk) == "string" then
+      accumulation_buffer = accumulation_buffer .. chunk
+      while true do
+        local json_blob, rest = consume_json_head:match(accumulation_buffer)
+        if not (json_blob) then
+          break
+        end
+        accumulation_buffer = rest
+        do
+          chunk = parse_completion_chunk(cjson.decode(json_blob))
+          if chunk then
+            chunk_callback(chunk)
+          end
+        end
+      end
+    end
+    return ...
+  end
+end
 local ChatSession
 do
   local _class_0
@@ -125,7 +184,7 @@ do
       if stream_callback then
         assert(type(response) == "string", "Expected string response from streaming output")
         local parts = { }
-        local f = self.client:create_stream_filter(function(c)
+        local f = create_chat_stream_filter(function(c)
           return table.insert(parts, c.content)
         end)
         f(response)
@@ -185,5 +244,11 @@ do
 end
 return {
   ChatSession = ChatSession,
-  test_message = test_message
+  test_message = test_message,
+  test_function = test_function,
+  parse_chat_response = parse_chat_response,
+  parse_error_message = parse_error_message,
+  parse_completion_chunk = parse_completion_chunk,
+  consume_json_head = consume_json_head,
+  create_chat_stream_filter = create_chat_stream_filter
 }

openai/chat_session.moon openai/chat_completions.moonopenai/chat_session.moon renamed to openai/chat_completions.moon

@@ -1,3 +1,5 @@
+-- This is the legacy API https://platform.openai.com/docs/api-reference/chat
+
 cjson = require "cjson"
 import types from require "tableshape"
 
@@ -66,6 +68,82 @@ parse_error_message = types.partial {
   }
 }
 
+-- sse streaming chunk format from chat completions API
+-- {
+--   "id": "chatcmpl-XXX",
+--   "object": "chat.completion.chunk",
+--   "created": 1682979397,
+--   "model": "gpt-3.5-turbo-0301",
+--   "choices": [
+--     {
+--       "delta": {
+--         "content": " hello"
+--       },
+--       "index": 0,
+--       "finish_reason": null
+--     }
+--   ]
+-- }
+
+parse_completion_chunk = types.partial {
+  object: "chat.completion.chunk"
+  -- not sure of the whole range of chunks, so for now we strictly parse an append
+  choices: types.shape {
+    types.partial {
+      delta: types.partial {
+        "content": types.string\tag "content"
+      }
+      index: types.number\tag "index"
+    }
+  }
+}
+
+-- lpeg pattern to read a json data block from the front of a string, returns
+-- the json blob and the rest of the string if it could parse one
+consume_json_head = do
+  import C, S, P from require "lpeg"
+
+  -- this pattern reads from the front just enough characters to consume a
+  -- valid json object
+  consume_json = P (str, pos) ->
+    str_len = #str
+    for k=pos+1,str_len
+      candidate = str\sub pos, k
+      parsed = false
+      pcall -> parsed = cjson.decode candidate
+      if parsed
+        return k + 1
+
+    return nil -- fail
+
+  S("\t\n\r ")^0 * P("data: ") * C(consume_json) * C(P(1)^0)
+
+
+-- creates a ltn12 compatible filter function that will call chunk_callback
+-- for each parsed json chunk from the server-sent events api response
+create_chat_stream_filter = (chunk_callback) ->
+  assert types.function(chunk_callback), "Must provide chunk_callback function when streaming response"
+
+  accumulation_buffer = ""
+
+  (...) ->
+    chunk = ...
+
+    if type(chunk) == "string"
+      accumulation_buffer ..= chunk
+
+      while true
+        json_blob, rest = consume_json_head\match accumulation_buffer
+        unless json_blob
+          break
+
+        accumulation_buffer = rest
+        if chunk = parse_completion_chunk cjson.decode json_blob
+          chunk_callback chunk
+
+    ...
+
+
 -- handles appending response for each call to chat
 -- TODO: hadle appending the streaming response to the output
 class ChatSession
@@ -133,7 +211,7 @@ class ChatSession
         "Expected string response from streaming output"
 
       parts = {}
-      f = @client\create_stream_filter (c) ->
+      f = create_chat_stream_filter (c) ->
         table.insert parts, c.content
 
       f response
@@ -159,4 +237,13 @@ class ChatSession
     -- response is missing for function_calls, so we return the entire message object
     out.response or out.message
 
-{:ChatSession, :test_message}
+{
+  :ChatSession
+  :test_message
+  :test_function
+  :parse_chat_response
+  :parse_error_message
+  :parse_completion_chunk
+  :consume_json_head
+  :create_chat_stream_filter
+}

openai/init.lua

@@ -4,88 +4,32 @@ local cjson = require("cjson")
 local unpack = table.unpack or unpack
 local types
 types = require("tableshape").types
-local ChatSession, test_message
-do
-  local _obj_0 = require("openai.chat_session")
-  ChatSession, test_message = _obj_0.ChatSession, _obj_0.test_message
-end
 local parse_url = require("socket.url").parse
-local parse_completion_chunk = types.partial({
-  object = "chat.completion.chunk",
-  choices = types.shape({
-    types.partial({
-      delta = types.partial({
-        ["content"] = types.string:tag("content")
-      }),
-      index = types.number:tag("index")
-    })
-  })
-})
-local consume_json_head
-do
-  local C, S, P
-  do
-    local _obj_0 = require("lpeg")
-    C, S, P = _obj_0.C, _obj_0.S, _obj_0.P
-  end
-  local consume_json = P(function(str, pos)
-    local str_len = #str
-    for k = pos + 1, str_len do
-      local candidate = str:sub(pos, k)
-      local parsed = false
-      pcall(function()
-        parsed = cjson.decode(candidate)
-      end)
-      if parsed then
-        return k + 1
-      end
-    end
-    return nil
-  end)
-  consume_json_head = S("\t\n\r ") ^ 0 * P("data: ") * C(consume_json) * C(P(1) ^ 0)
-end
 local OpenAI
 do
   local _class_0
   local _base_0 = {
     api_base = "https://api.openai.com/v1",
     default_model = "gpt-4.1",
     new_chat_session = function(self, ...)
+      local ChatSession
+      ChatSession = require("openai.chat_completions").ChatSession
       return ChatSession(self, ...)
     end,
     new_response_chat_session = function(self, ...)
       local ResponsesChatSession
       ResponsesChatSession = require("openai.responses").ResponsesChatSession
       return ResponsesChatSession(self, ...)
     end,
-    create_stream_filter = function(self, chunk_callback)
-      assert(types["function"](chunk_callback), "Must provide chunk_callback function when streaming response")
-      local accumulation_buffer = ""
-      return function(...)
-        local chunk = ...
-        if type(chunk) == "string" then
-          accumulation_buffer = accumulation_buffer .. chunk
-          while true do
-            local json_blob, rest = consume_json_head:match(accumulation_buffer)
-            if not (json_blob) then
-              break
-            end
-            accumulation_buffer = rest
-            do
-              chunk = parse_completion_chunk(cjson.decode(json_blob))
-              if chunk then
-                chunk_callback(chunk)
-              end
-            end
-          end
-        end
-        return ...
-      end
-    end,
     chat = function(self, messages, opts, chunk_callback)
       if chunk_callback == nil then
         chunk_callback = nil
       end
+      local test_message, create_chat_stream_filter
+      do
+        local _obj_0 = require("openai.chat_completions")
+        test_message, create_chat_stream_filter = _obj_0.test_message, _obj_0.create_chat_stream_filter
+      end
       local test_messages = types.array_of(test_message)
       assert(test_messages(messages))
       local payload = {
@@ -99,7 +43,7 @@ do
       end
       local stream_filter
       if payload.stream then
-        stream_filter = self:create_stream_filter(chunk_callback)
+        stream_filter = create_chat_stream_filter(chunk_callback)
       end
       return self:_request("POST", "/chat/completions", payload, nil, stream_filter)
     end,
@@ -294,7 +238,6 @@ do
 end
 return {
   OpenAI = OpenAI,
-  ChatSession = ChatSession,
   VERSION = VERSION,
   new = OpenAI
 }