From 33a7bec43dadf591dae52f67965f396def3077f4 Mon Sep 17 00:00:00 2001 From: Kevin Hermawan <84965338+kevinhermawan@users.noreply.github.com> Date: Sun, 31 Mar 2024 15:29:40 +0700 Subject: [PATCH] fix: response serialization failed --- Sources/OllamaKit/OllamaKit+Chat.swift | 56 ++++++++++++++-------- Sources/OllamaKit/OllamaKit+Generate.swift | 56 ++++++++++++++-------- Sources/OllamaKit/OllamaKit.swift | 40 ++++++++++++++++ 3 files changed, 110 insertions(+), 42 deletions(-) diff --git a/Sources/OllamaKit/OllamaKit+Chat.swift b/Sources/OllamaKit/OllamaKit+Chat.swift index b3cd541..fd771c7 100644 --- a/Sources/OllamaKit/OllamaKit+Chat.swift +++ b/Sources/OllamaKit/OllamaKit+Chat.swift @@ -34,22 +34,29 @@ extension OllamaKit { public func chat(data: OKChatRequestData) -> AsyncThrowingStream { AsyncThrowingStream { continuation in let request = AF.streamRequest(router.chat(data: data)).validate() + var buffer = Data() - request.responseStreamDecodable(of: OKChatResponse.self, using: decoder) { stream in + request.responseStream { stream in switch stream.event { case .stream(let result): switch result { - case .success(let response): - continuation.yield(response) + case .success(let data): + buffer.append(data) + + while let chunk = extractNextJSON(from: &buffer) { + do { + let response = try decoder.decode(OKChatResponse.self, from: chunk) + continuation.yield(response) + } catch { + continuation.finish(throwing: error) + return + } + } case .failure(let error): continuation.finish(throwing: error) } - case .complete(let completion): - if let error = completion.error { - continuation.finish(throwing: error) - } else { - continuation.finish() - } + case .complete(_): + continuation.finish() } } } @@ -73,26 +80,33 @@ extension OllamaKit { /// ``` /// /// - Parameter data: The ``OKChatRequestData`` used to initiate the chat streaming from the Ollama API. - /// - Returns: An `AnyPublisher` emitting the live stream of chat responses from the Ollama API. - public func chat(data: OKChatRequestData) -> AnyPublisher { - let subject = PassthroughSubject() + /// - Returns: An `AnyPublisher` emitting the live stream of chat responses from the Ollama API. + public func chat(data: OKChatRequestData) -> AnyPublisher { + let subject = PassthroughSubject() let request = AF.streamRequest(router.chat(data: data)).validate() + var buffer = Data() - request.responseStreamDecodable(of: OKChatResponse.self, using: decoder) { stream in + request.responseStream { stream in switch stream.event { case .stream(let result): switch result { - case .success(let response): - subject.send(response) + case .success(let data): + buffer.append(data) + + while let chunk = extractNextJSON(from: &buffer) { + do { + let response = try decoder.decode(OKChatResponse.self, from: chunk) + subject.send(response) + } catch { + subject.send(completion: .failure(error)) + return + } + } case .failure(let error): subject.send(completion: .failure(error)) } - case .complete(let completion): - if completion.error != nil { - subject.send(completion: .failure(completion.error!)) - } else { - subject.send(completion: .finished) - } + case .complete(_): + subject.send(completion: .finished) } } diff --git a/Sources/OllamaKit/OllamaKit+Generate.swift b/Sources/OllamaKit/OllamaKit+Generate.swift index fb32426..012582a 100644 --- a/Sources/OllamaKit/OllamaKit+Generate.swift +++ b/Sources/OllamaKit/OllamaKit+Generate.swift @@ -34,22 +34,29 @@ extension OllamaKit { public func generate(data: OKGenerateRequestData) -> AsyncThrowingStream { AsyncThrowingStream { continuation in let request = AF.streamRequest(router.generate(data: data)).validate() + var buffer = Data() - request.responseStreamDecodable(of: OKGenerateResponse.self, using: decoder) { stream in + request.responseStream { stream in switch stream.event { case .stream(let result): switch result { - case .success(let response): - continuation.yield(response) + case .success(let data): + buffer.append(data) + + while let chunk = extractNextJSON(from: &buffer) { + do { + let response = try decoder.decode(OKGenerateResponse.self, from: chunk) + continuation.yield(response) + } catch { + continuation.finish(throwing: error) + return + } + } case .failure(let error): continuation.finish(throwing: error) } - case .complete(let completion): - if let error = completion.error { - continuation.finish(throwing: error) - } else { - continuation.finish() - } + case .complete(_): + continuation.finish() } } } @@ -73,26 +80,33 @@ extension OllamaKit { /// ``` /// /// - Parameter data: The ``OKGenerateRequestData`` used to initiate the streaming from the Ollama API. - /// - Returns: An `AnyPublisher` emitting the live stream of responses from the Ollama API. - public func generate(data: OKGenerateRequestData) -> AnyPublisher { - let subject = PassthroughSubject() + /// - Returns: An `AnyPublisher` emitting the live stream of responses from the Ollama API. + public func generate(data: OKGenerateRequestData) -> AnyPublisher { + let subject = PassthroughSubject() let request = AF.streamRequest(router.generate(data: data)).validate() + var buffer = Data() - request.responseStreamDecodable(of: OKGenerateResponse.self, using: decoder) { stream in + request.responseStream { stream in switch stream.event { case .stream(let result): switch result { - case .success(let response): - subject.send(response) + case .success(let data): + buffer.append(data) + + while let chunk = extractNextJSON(from: &buffer) { + do { + let response = try decoder.decode(OKGenerateResponse.self, from: chunk) + subject.send(response) + } catch { + subject.send(completion: .failure(error)) + return + } + } case .failure(let error): subject.send(completion: .failure(error)) } - case .complete(let completion): - if completion.error != nil { - subject.send(completion: .failure(completion.error!)) - } else { - subject.send(completion: .finished) - } + case .complete(_): + subject.send(completion: .finished) } } diff --git a/Sources/OllamaKit/OllamaKit.swift b/Sources/OllamaKit/OllamaKit.swift index 4312913..5b337f3 100644 --- a/Sources/OllamaKit/OllamaKit.swift +++ b/Sources/OllamaKit/OllamaKit.swift @@ -38,4 +38,44 @@ public struct OllamaKit { self.router = router } + + internal func extractNextJSON(from buffer: inout Data) -> Data? { + var isEscaped = false + var isWithinString = false + var nestingDepth = 0 + var objectStartIndex = buffer.startIndex + + for (index, byte) in buffer.enumerated() { + let character = Character(UnicodeScalar(byte)) + + if isEscaped { + isEscaped = false + } else if character == "\\" { + isEscaped = true + } else if character == "\"" { + isWithinString.toggle() + } else if !isWithinString { + switch character { + case "{": + nestingDepth += 1 + if nestingDepth == 1 { + objectStartIndex = index + } + case "}": + nestingDepth -= 1 + if nestingDepth == 0 { + let range = objectStartIndex..