From 80e916097b36281b68ab579e9e08cb39519bd472 Mon Sep 17 00:00:00 2001
From: Pallav Agarwal <agpallav@gmail.com>
Date: Fri, 13 Oct 2023 00:26:22 -0400
Subject: [PATCH 1/2] Correctly handle partial JSON at the end of the chunk

---
 Sources/OpenAI/Private/StreamingSession.swift | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/Sources/OpenAI/Private/StreamingSession.swift b/Sources/OpenAI/Private/StreamingSession.swift
index 55902f09..e2d52e57 100644
--- a/Sources/OpenAI/Private/StreamingSession.swift
+++ b/Sources/OpenAI/Private/StreamingSession.swift
@@ -28,6 +28,8 @@ final class StreamingSession<ResultType: Codable>: NSObject, Identifiable, URLSe
         return session
     }()
     
+    private var prevChunkBuffer = ""
+
     init(urlRequest: URLRequest) {
         self.urlRequest = urlRequest
     }
@@ -47,14 +49,16 @@ final class StreamingSession<ResultType: Codable>: NSObject, Identifiable, URLSe
             onProcessingError?(self, StreamingError.unknownContent)
             return
         }
-        let jsonObjects = stringContent
+        let jsonObjects = "\(prevChunkBuffer)\(stringContent)"
             .components(separatedBy: "data:")
             .filter { $0.isEmpty == false }
             .map { $0.trimmingCharacters(in: .whitespacesAndNewlines) }
+        prevChunkBuffer = ""
+        
         guard jsonObjects.isEmpty == false, jsonObjects.first != streamingCompletionMarker else {
             return
         }
-        jsonObjects.forEach { jsonContent  in
+        jsonObjects.enumerated().forEach { (index, jsonContent)  in
             guard jsonContent != streamingCompletionMarker else {
                 return
             }
@@ -77,7 +81,12 @@ final class StreamingSession<ResultType: Codable>: NSObject, Identifiable, URLSe
                     let decoded = try JSONDecoder().decode(APIErrorResponse.self, from: jsonData)
                     onProcessingError?(self, decoded)
                 } catch {
-                    onProcessingError?(self, apiError)
+                    if index == jsonObjects.count - 1 {
+                        // Chunk ends in a partial JSON
+                        prevChunkBuffer = "data: \(jsonContent)"
+                    } else {
+                        onProcessingError?(self, apiError)
+                    }
                 }
             }
         }

From 8f08006ae8391d8778a8228740d4616031f99762 Mon Sep 17 00:00:00 2001
From: Ihor Makhnyk <ingvarus.bootcamp@macpaw.com>
Date: Mon, 13 Nov 2023 12:18:00 +0200
Subject: [PATCH 2/2] Move json processing to separate method

---
 Sources/OpenAI/Private/StreamingSession.swift | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/Sources/OpenAI/Private/StreamingSession.swift b/Sources/OpenAI/Private/StreamingSession.swift
index e2d52e57..a69e46cf 100644
--- a/Sources/OpenAI/Private/StreamingSession.swift
+++ b/Sources/OpenAI/Private/StreamingSession.swift
@@ -28,7 +28,7 @@ final class StreamingSession<ResultType: Codable>: NSObject, Identifiable, URLSe
         return session
     }()
     
-    private var prevChunkBuffer = ""
+    private var previousChunkBuffer = ""
 
     init(urlRequest: URLRequest) {
         self.urlRequest = urlRequest
@@ -49,11 +49,20 @@ final class StreamingSession<ResultType: Codable>: NSObject, Identifiable, URLSe
             onProcessingError?(self, StreamingError.unknownContent)
             return
         }
-        let jsonObjects = "\(prevChunkBuffer)\(stringContent)"
+        processJSON(from: stringContent)
+    }
+    
+}
+
+extension StreamingSession {
+    
+    private func processJSON(from stringContent: String) {
+        let jsonObjects = "\(previousChunkBuffer)\(stringContent)"
             .components(separatedBy: "data:")
             .filter { $0.isEmpty == false }
             .map { $0.trimmingCharacters(in: .whitespacesAndNewlines) }
-        prevChunkBuffer = ""
+        
+        previousChunkBuffer = ""
         
         guard jsonObjects.isEmpty == false, jsonObjects.first != streamingCompletionMarker else {
             return
@@ -82,8 +91,7 @@ final class StreamingSession<ResultType: Codable>: NSObject, Identifiable, URLSe
                     onProcessingError?(self, decoded)
                 } catch {
                     if index == jsonObjects.count - 1 {
-                        // Chunk ends in a partial JSON
-                        prevChunkBuffer = "data: \(jsonContent)"
+                        previousChunkBuffer = "data: \(jsonContent)" // Chunk ends in a partial JSON
                     } else {
                         onProcessingError?(self, apiError)
                     }
@@ -91,4 +99,5 @@ final class StreamingSession<ResultType: Codable>: NSObject, Identifiable, URLSe
             }
         }
     }
+    
 }