Skip to content

Commit

Permalink
refactor: migrates generate method from completion handler to Combine
Browse files Browse the repository at this point in the history
  • Loading branch information
kevinhermawan committed Nov 18, 2023
1 parent 60fd29f commit 0e7b8f5
Showing 1 changed file with 26 additions and 11 deletions.
37 changes: 26 additions & 11 deletions Sources/OllamaKit/OllamaKit.swift
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
// Created by Kevin Hermawan on 10/11/23.
//

import Combine
import Alamofire
import Foundation

Expand All @@ -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 {
Expand All @@ -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 {
Expand Down

0 comments on commit 0e7b8f5

Please sign in to comment.