Merge pull request alexrudall#248 from rmontgomery429/faraday-json-parsing

Improved stream parsing

Author: alexrudall

lib/openai/client.rb

@@ -93,16 +93,21 @@ def self.to_json(string)
       JSON.parse(string.gsub("}\n{", "},{").prepend("[").concat("]"))
     end
 
+    # Given a proc, returns a proc that can be used to iterate over a JSON stream of chunks.
+    # For each chunk, the proc is called giving it the JSON object. The JSON object could be a
+    # data object or an error object as described in the OpenAI API documentation.
+    #
+    # If the JSON object for a given data or error message is invalid, it is ignored.
+    #
+    # @param user_proc [Proc] The proc to call for each JSON object in the chunk.
+    # @return [Proc] A proc that can be used to iterate over the JSON stream.
     private_class_method def self.to_json_stream(user_proc:)
-      proc do |chunk, bytesize|
-        # Clean up response string of chunks.
-        chunk = chunk.gsub("\n\ndata: [DONE]\n\n", "").gsub("\n\ndata:", ",").gsub("data: ", "")
-
-        # Turn it into JSON.
-        chunk = to_json(chunk)
-
-        # Pass an array of JSONified chunk(s) to the user's Proc.
-        user_proc.call([chunk].flatten, bytesize)
+      proc do |chunk, _|
+        chunk.scan(/(?:data|error): (\{.*\})/i).flatten.each do |data|
+          user_proc.call(JSON.parse(data))
+        rescue JSON::ParserError
+          # Ignore invalid JSON.
+        end
       end
     end
 

spec/openai/client/chat_spec.rb

@@ -35,7 +35,7 @@
           it "succeeds" do
             VCR.use_cassette(cassette) do
               response
-              expect(chunks.dig(0, 0, "choices", 0, "index")).to eq(0)
+              expect(chunks.dig(0, "choices", 0, "index")).to eq(0)
             end
           end
         end

spec/openai/client/client_spec.rb

@@ -25,6 +25,88 @@
     end
   end
 
+  describe ".to_json_stream" do
+    context "with a proc" do
+      let(:user_proc) { proc { |x| x } }
+      let(:stream)    { OpenAI::Client.send(:to_json_stream, user_proc: user_proc) }
+
+      it "returns a proc" do
+        expect(stream).to be_a(Proc)
+      end
+
+      context "when called with a string containing a single JSON object" do
+        it "calls the user proc with the data parsed as JSON" do
+          expect(user_proc).to receive(:call).with(JSON.parse('{"foo": "bar"}'))
+          stream.call('data: { "foo": "bar" }')
+        end
+      end
+
+      context "when called with string containing more than one JSON object" do
+        it "calls the user proc for each data parsed as JSON" do
+          expect(user_proc).to receive(:call).with(JSON.parse('{"foo": "bar"}'))
+          expect(user_proc).to receive(:call).with(JSON.parse('{"baz": "qud"}'))
+
+          stream.call(<<-CHUNK)
+            data: { "foo": "bar" }
+
+            data: { "baz": "qud" }
+
+            data: [DONE]
+
+          CHUNK
+        end
+      end
+
+      context "when called with a string that does not even resemble a JSON object" do
+        let(:bad_examples) { ["", "foo", "data: ", "data: foo"] }
+
+        it "does not call the user proc" do
+          bad_examples.each do |chunk|
+            expect(user_proc).to_not receive(:call)
+            stream.call(chunk)
+          end
+        end
+      end
+
+      context "when called with a string containing that looks like a JSON object but is invalid" do
+        let(:chunk) do
+          <<-CHUNK
+            data: { "foo": "bar" }
+            data: { BAD ]:-> JSON }
+          CHUNK
+        end
+
+        it "does not raise an error" do
+          expect(user_proc).to receive(:call).with(JSON.parse('{"foo": "bar"}'))
+
+          expect do
+            stream.call(chunk)
+          end.to_not raise_error(JSON::ParserError)
+        end
+      end
+
+      context "when called with a string containing an error" do
+        let(:chunk) do
+          <<-CHUNK
+            data: { "foo": "bar" }
+            error: { "message": "A bad thing has happened!" }
+          CHUNK
+        end
+
+        it "does not raise an error" do
+          expect(user_proc).to receive(:call).with(JSON.parse('{ "foo": "bar" }'))
+          expect(user_proc).to receive(:call).with(
+            JSON.parse('{ "message": "A bad thing has happened!" }')
+          )
+
+          expect do
+            stream.call(chunk)
+          end.to_not raise_error(JSON::ParserError)
+        end
+      end
+    end
+  end
+
   describe ".to_json" do
     context "with a jsonl string" do
       let(:body) { "{\"prompt\":\":)\"}\n{\"prompt\":\":(\"}\n" }