refactor: migrates chat and generate from Future to `Passthroug…

…hSubject` (#8)
kevinhermawan · Jan 15, 2024 · 982815a · 982815a
1 parent a975d73
commit 982815a
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 30 deletions.
diff --git a/Sources/OllamaKit/OllamaKit+Chat.swift b/Sources/OllamaKit/OllamaKit+Chat.swift
@@ -30,22 +30,23 @@ extension OllamaKit {
     /// - Parameter data: The ``OKChatRequestData`` used to initiate the chat streaming from the Ollama API.
     /// - Returns: An `AnyPublisher<OKChatResponse, AFError>` emitting the live stream of chat responses from the Ollama API.
     public func chat(data: OKChatRequestData) -> AnyPublisher<OKChatResponse, AFError> {
-        return Future<OKChatResponse, AFError> { promise in
-            let request = AF.streamRequest(router.chat(data: data)).validate()
-
-            request.responseStreamDecodable(of: OKChatResponse.self, using: decoder) { stream in
-                switch stream.event {
-                case .stream(let result):
-                    switch result {
-                    case .success(let response):
-                        promise(.success(response))
-                    case .failure(let error):
-                        promise(.failure(error))
-                    }
-                case .complete(_):
-                    break
+        let subject = PassthroughSubject<OKChatResponse, AFError>()
+        let request = AF.streamRequest(router.chat(data: data)).validate()
+
+        request.responseStreamDecodable(of: OKChatResponse.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)
             }
-        }.eraseToAnyPublisher()
+        }
+
+        return subject.eraseToAnyPublisher()
     }
 }
diff --git a/Sources/OllamaKit/OllamaKit+Generate.swift b/Sources/OllamaKit/OllamaKit+Generate.swift
@@ -30,22 +30,23 @@ extension OllamaKit {
     /// - Parameter data: The ``OKGenerateRequestData`` used to initiate the streaming from the Ollama API.
     /// - Returns: An `AnyPublisher<OKGenerateResponse, AFError>` emitting the live stream of responses from the Ollama API.
     public func generate(data: OKGenerateRequestData) -> AnyPublisher<OKGenerateResponse, AFError> {
-        return Future<OKGenerateResponse, AFError> { promise in
-            let request = AF.streamRequest(router.generate(data: data)).validate()
-
-            request.responseStreamDecodable(of: OKGenerateResponse.self, using: decoder) { stream in
-                switch stream.event {
-                case .stream(let result):
-                    switch result {
-                    case .success(let response):
-                        promise(.success(response))
-                    case .failure(let error):
-                        promise(.failure(error))
-                    }
-                case .complete(_):
-                    break
+        let subject = PassthroughSubject<OKGenerateResponse, AFError>()
+        let request = AF.streamRequest(router.generate(data: data)).validate()
+
+        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)
             }
-        }.eraseToAnyPublisher()
+        }
+
+        return subject.eraseToAnyPublisher()
     }
 }