From 0e7b8f58f642bde7f1c3927dd2f7994a41f30f2f Mon Sep 17 00:00:00 2001
From: Kevin Hermawan <84965338+kevinhermawan@users.noreply.github.com>
Date: Sun, 19 Nov 2023 03:32:51 +0700
Subject: [PATCH] refactor: migrates `generate` method from completion handler
 to Combine

---
 Sources/OllamaKit/OllamaKit.swift | 37 ++++++++++++++++++++++---------
 1 file changed, 26 insertions(+), 11 deletions(-)

diff --git a/Sources/OllamaKit/OllamaKit.swift b/Sources/OllamaKit/OllamaKit.swift
index 5745fe1..2af84c0 100644
--- a/Sources/OllamaKit/OllamaKit.swift
+++ b/Sources/OllamaKit/OllamaKit.swift
@@ -5,6 +5,7 @@
 //  Created by Kevin Hermawan on 10/11/23.
 //
 
+import Combine
 import Alamofire
 import Foundation
 
@@ -31,7 +32,8 @@ public struct OllamaKit {
 extension OllamaKit {
     /// Checks the reachability of the Ollama API.
     ///
-    /// This asynchronous method performs a network request to determine if the Ollama API is reachable from the current client. It can be used to verify network connectivity and API availability before attempting further API interactions.
+    /// This asynchronous method performs a network request to determine if the Ollama API is reachable from the current client.
+    /// It can be used to verify network connectivity and API availability before attempting further API interactions.
     ///
     /// - Returns: A Boolean value indicating whether the Ollama API is reachable (`true`) or not (`false`).
     public func reachable() async -> Bool {
@@ -49,25 +51,38 @@ extension OllamaKit {
 }
 
 extension OllamaKit {
-    /// Establishes a stream to the Ollama API for generating responses based on the provided data.
+    /// Establishes a Combine publisher for streaming responses from the Ollama API, based on the provided data.
     ///
-    /// This method continuously streams responses as they are generated by the Ollama API,
-    /// with the final response including detailed data about the generation process.
+    /// This method sets up a streaming connection using the Combine framework, allowing for real-time data handling as the responses are generated by the Ollama API.
     ///
-    /// - Parameters:
-    ///   - data: The data used to generate the stream.
-    ///   - stream: A closure that processes the streaming data.
-    public func generate(data: OKGenerateRequestData, stream: @escaping (DataStreamRequest.Stream<OKGenerateResponse, AFError>) -> Void) {
+    /// - Parameter data: The `OKGenerateRequestData` used to initiate the streaming from the Ollama API.
+    /// - Returns: An `AnyPublisher` emitting `OKGenerateResponse` and `AFError`, representing the live stream of responses from the Ollama API.
+    public func generate(data: OKGenerateRequestData) -> AnyPublisher<OKGenerateResponse, AFError> {
+        let subject = PassthroughSubject<OKGenerateResponse, AFError>()
         let request = AF.streamRequest(router.generate(data: data)).validate()
-        request.responseStreamDecodable(of: OKGenerateResponse.self, using: decoder, stream: stream)
+        
+        request.responseStreamDecodable(of: OKGenerateResponse.self, using: decoder) { stream in
+            switch stream.event {
+            case .stream(let result):
+                switch result {
+                case .success(let response):
+                    subject.send(response)
+                case .failure(let error):
+                    subject.send(completion: .failure(error))
+                }
+            case .complete(_):
+                subject.send(completion: .finished)
+            }
+        }
+        
+        return subject.eraseToAnyPublisher()
     }
 }
 
 extension OllamaKit {
     /// Asynchronously retrieves a list of available models from the Ollama API.
     ///
-    /// This method returns an `OKModelResponse` containing the details of the available models,
-    /// making it easy to understand what models are currently accessible.
+    /// This method returns an `OKModelResponse` containing the details of the available models, making it easy to understand what models are currently accessible.
     ///
     /// - Returns: An `OKModelResponse` object listing the available models.
     public func models() async throws -> OKModelResponse {