make create_chat_completion stream include all events, only types on the chat legacy alias

Author: leafo

openai/chat_completions.lua

@@ -114,12 +114,7 @@ create_chat_stream_filter = function(chunk_callback)
           break
         end
         accumulation_buffer = rest
-        do
-          chunk = parse_completion_chunk(cjson.decode(json_blob))
-          if chunk then
-            chunk_callback(chunk)
-          end
-        end
+        chunk_callback(cjson.decode(json_blob))
       end
     end
     return ...
@@ -159,7 +154,7 @@ do
       if stream_callback == nil then
         stream_callback = nil
       end
-      local status, response = self.client:create_chat_completion(self.messages, {
+      local status, response = self.client:chat(self.messages, {
         function_call = self.opts.function_call,
         functions = self.functions,
         model = self.opts.model,
@@ -186,7 +181,12 @@ do
         assert(type(response) == "string", "Expected string response from streaming output")
         local parts = { }
         local f = create_chat_stream_filter(function(c)
-          return table.insert(parts, c.content)
+          do
+            local parsed = parse_completion_chunk(c)
+            if parsed then
+              return table.insert(parts, parsed.content)
+            end
+          end
         end)
         f(response)
         local message = {
@@ -253,5 +253,6 @@ end
 return {
   ChatSession = ChatSession,
   test_message = test_message,
-  create_chat_stream_filter = create_chat_stream_filter
+  create_chat_stream_filter = create_chat_stream_filter,
+  parse_completion_chunk = parse_completion_chunk
 }

openai/chat_completions.moon

@@ -140,8 +140,9 @@ create_chat_stream_filter = (chunk_callback) ->
           break
 
         accumulation_buffer = rest
-        if chunk = parse_completion_chunk cjson.decode json_blob
-          chunk_callback chunk
+        chunk_callback cjson.decode json_blob
+        -- if chunk = parse_completion_chunk cjson.decode json_blob
+        --   chunk_callback chunk
 
     ...
 
@@ -186,7 +187,7 @@ class ChatSession
   -- append_response: should the response be appended to the chat history
   -- stream_callback: provide a function to enable streaming output. function will receive each chunk as it's generated
   generate_response: (append_response=true, stream_callback=nil) =>
-    status, response = @client\create_chat_completion @messages, {
+    status, response = @client\chat @messages, {
       function_call: @opts.function_call -- override the default function call behavior
       functions: @functions
       model: @opts.model
@@ -214,7 +215,8 @@ class ChatSession
 
       parts = {}
       f = create_chat_stream_filter (c) ->
-        table.insert parts, c.content
+        if parsed = parse_completion_chunk c
+          table.insert parts, parsed.content
 
       f response
       message = {
@@ -250,4 +252,5 @@ class ChatSession
   :ChatSession
   :test_message
   :create_chat_stream_filter
+  :parse_completion_chunk
 }

openai/init.lua

@@ -50,8 +50,26 @@ do
       end
       return self:_request("POST", "/chat/completions", payload, nil, stream_filter)
     end,
-    chat = function(self, ...)
-      return self:create_chat_completion(...)
+    chat = function(self, messages, opts, chunk_callback)
+      if chunk_callback == nil then
+        chunk_callback = nil
+      end
+      do
+        local cb = chunk_callback
+        if cb then
+          local parse_completion_chunk
+          parse_completion_chunk = require("openai.chat_completions").parse_completion_chunk
+          chunk_callback = function(chunk)
+            do
+              local delta = parse_completion_chunk(chunk)
+              if delta then
+                return cb(delta)
+              end
+            end
+          end
+        end
+      end
+      return self:create_chat_completion(messages, opts, chunk_callback)
     end,
     completion = function(self, prompt, opts)
       local payload = {

openai/init.moon

@@ -58,7 +58,17 @@ class OpenAI
     @_request "POST", "/chat/completions", payload, nil, stream_filter
 
   -- legacy alias for create_chat_completion (for backward compatibility)
-  chat: (...) => @create_chat_completion ...
+  -- the legacy method also has the filtered chunk responses instead of pushing
+  -- through every event through the callback
+  chat: (messages, opts, chunk_callback=nil) =>
+    if cb = chunk_callback
+      import parse_completion_chunk from require "openai.chat_completions"
+      chunk_callback = (chunk) ->
+        -- filter chunk to only pass through chat.completion.chunk with parsed delta
+        if delta = parse_completion_chunk chunk
+          cb delta
+
+    @create_chat_completion messages, opts, chunk_callback
 
   -- call /completions
   -- opts: additional parameters as described in https://platform.openai.com/docs/api-reference/completions

openai/responses.moon

@@ -130,8 +130,9 @@ parse_response_stream_chunk = (chunk) ->
       raw: chunk
     }
 
--- takes a stream of string chunks and extracts SSE json objects out of the
--- stream, calling chunk_callback when a parsed object is found
+
+-- creates a ltn12 compatible filter function that will call chunk_callback
+-- for each parsed json chunk from the server-sent events api response
 create_response_stream_filter = (chunk_callback) ->
   assert types.function(chunk_callback), "Must provide chunk_callback function when streaming response"
 

spec/openai_spec.moon

@@ -402,6 +402,25 @@ describe "OpenAI API Client", ->
 
       assert.same "This is a chat response.", response
 
+    it "processes streaming chunks with create_chat_completion (raw)", ->
+      client = OpenAI "test-api-key"
+
+      chunks_received = {}
+      stream_callback = (chunk) ->
+        table.insert chunks_received, chunk
+
+      status, response = client\create_chat_completion {
+        {role: "user", content: "tell me a joke"}
+      }, {stream: true}, stream_callback
+
+      assert.same 200, status
+      -- create_chat_completion passes raw JSON chunks
+      assert.same {
+        {object: "chat.completion.chunk", choices: {{delta: {content: "This is "}, index: 0}}}
+        {object: "chat.completion.chunk", choices: {{delta: {content: "a chat "}, index: 1}}}
+        {object: "chat.completion.chunk", choices: {{delta: {content: "response."}, index: 2}}}
+      }, chunks_received
+
   describe "responses", ->
     it "creates a response (raw API)", ->
       client = OpenAI "test-api-key"