diff --git a/Sources/OllamaKit/OllamaKit+PullModel.swift b/Sources/OllamaKit/OllamaKit+PullModel.swift new file mode 100644 index 0000000..d694cc7 --- /dev/null +++ b/Sources/OllamaKit/OllamaKit+PullModel.swift @@ -0,0 +1,78 @@ +// +// OllamaKit+PullModel.swift +// +// +// Created by Lukas Pistrol on 25.11.24. +// + +import Combine +import Foundation + +extension OllamaKit { + /// Establishes an asynchronous stream for pulling a model from the ollama library. + /// + /// This method will periodically yield ``OKPullModelResponse`` structures as the model is being pulled. + /// Depending on the size of the model and the speed of the internet connection, this process may take a while. + /// The stream will complete once the model has been fully pulled. + /// + /// ```swift + /// let ollamaKit = OllamaKit() + /// let reqData = OKPullModelRequestData(model: "llama3.2") + /// + /// for try await response in ollamaKit.pullModel(data: reqData) { + /// print(response.status) + /// if let progress = response.completed, let total = response.total { + /// print("Progress: \(progress)/\(total) bytes") + /// } + /// } + /// ``` + /// + /// - Parameter data: The ``OKPullModelRequestData`` used to initiate the chat streaming from the ollama library. + /// - Returns: An asynchronous stream that emits ``OKPullModelResponse``. + public func pullModel(data: OKPullModelRequestData) -> AsyncThrowingStream { + do { + let request = try OKRouter.pullModel(data: data).asURLRequest() + + return OKHTTPClient.shared.stream(request: request, with: OKPullModelResponse.self) + } catch { + return AsyncThrowingStream { continuation in + continuation.finish(throwing: error) + } + } + } + + /// Establishes a Combine publisher for pulling a model from the ollama library. + /// + /// This method will periodically yield ``OKPullModelResponse`` structures as the model is being pulled. + /// Depending on the size of the model and the speed of the internet connection, this process may take a while. + /// The stream will complete once the model has been fully pulled. + /// + /// ```swift + /// let ollamaKit = OllamaKit() + /// let reqData = OKPullModelRequestData(model: "llama3.2") + /// + /// ollamaKit.pullModel(data: reqData) + /// .sink { completion in + /// // Handle completion + /// } receiveValue: { response in + /// print(response.status) + /// if let progress = response.completed, let total = response.total { + /// print("Progress: \(progress)/\(total) bytes") + /// } + /// } + /// } + /// ``` + /// + /// - Parameter data: The ``OKPullModelRequestData`` used to initiate the chat streaming from the ollama library. + /// - Returns: A combine publisher that emits ``OKPullModelResponse``. + public func pullModel(data: OKPullModelRequestData) -> AnyPublisher { + do { + let request = try OKRouter.pullModel(data: data).asURLRequest() + + return OKHTTPClient.shared.stream(request: request, with: OKPullModelResponse.self) + } catch { + return Fail(error: error).eraseToAnyPublisher() + } + } + +} diff --git a/Sources/OllamaKit/RequestData/OKPullModelRequestData.swift b/Sources/OllamaKit/RequestData/OKPullModelRequestData.swift new file mode 100644 index 0000000..9a275d5 --- /dev/null +++ b/Sources/OllamaKit/RequestData/OKPullModelRequestData.swift @@ -0,0 +1,22 @@ +// +// OKPullModelRequestData.swift +// +// +// Created by Lukas Pistrol on 25.11.24. +// + +import Foundation + +/// A structure that encapsulates the data necessary for pulling a new model from the ollama library using the Ollama API. +public struct OKPullModelRequestData: Encodable { + private let stream: Bool + + /// A string representing the identifier of the model for which information is requested. + public let model: String + + public init(model: String) { + self.stream = true + self.model = model + } +} + diff --git a/Sources/OllamaKit/Responses/OKModelResponse.swift b/Sources/OllamaKit/Responses/OKModelResponse.swift index 2661bc1..6cecd9f 100644 --- a/Sources/OllamaKit/Responses/OKModelResponse.swift +++ b/Sources/OllamaKit/Responses/OKModelResponse.swift @@ -25,5 +25,26 @@ public struct OKModelResponse: Decodable { /// A `Date` representing the last modification date of the model. public let modifiedAt: Date + + /// The details about the model. + public let details: ModelDetails + + /// A structure that represents the details of the model. + public struct ModelDetails: Decodable { + /// The format of the model. E.g. "gguf". + public let format: String + + /// The family of the model. E.g. "llama". + public let family: String + + /// The parameter size of the model. E.g. "8.0B". + public let parameterSize: String + + /// The quantization level of the model. E.g. "Q4_0". + public let quantizationLevel: String + + /// All the families of the model. E.g. ["llama", "phi3"]. + public let families: [String]? + } } } diff --git a/Sources/OllamaKit/Responses/OKPullModelResponse.swift b/Sources/OllamaKit/Responses/OKPullModelResponse.swift new file mode 100644 index 0000000..94369ea --- /dev/null +++ b/Sources/OllamaKit/Responses/OKPullModelResponse.swift @@ -0,0 +1,23 @@ +// +// OKPullModelResponse.swift +// +// +// Created by Lukas Pistrol on 25.11.24. +// + +import Foundation + +/// The response model for pulling a new model from the ollama library. +public struct OKPullModelResponse: Decodable { + /// The current status. + public let status: String + + /// The digest hash of the current file. + public let digest: String? + + /// The size of the current file. + public let total: Int? + + /// The number of bytes that have been completed. + public let completed: Int? +} diff --git a/Sources/OllamaKit/Utils/OKRouter.swift b/Sources/OllamaKit/Utils/OKRouter.swift index a3e9e8b..90bc737 100644 --- a/Sources/OllamaKit/Utils/OKRouter.swift +++ b/Sources/OllamaKit/Utils/OKRouter.swift @@ -17,6 +17,7 @@ internal enum OKRouter { case chat(data: OKChatRequestData) case copyModel(data: OKCopyModelRequestData) case deleteModel(data: OKDeleteModelRequestData) + case pullModel(data: OKPullModelRequestData) case embeddings(data: OKEmbeddingsRequestData) internal var path: String { @@ -35,6 +36,8 @@ internal enum OKRouter { return "/api/copy" case .deleteModel: return "/api/delete" + case .pullModel: + return "/api/pull" case .embeddings: return "/api/embeddings" } @@ -56,6 +59,8 @@ internal enum OKRouter { return "POST" case .deleteModel: return "DELETE" + case .pullModel: + return "POST" case .embeddings: return "POST" } @@ -85,6 +90,8 @@ extension OKRouter { request.httpBody = try JSONEncoder.default.encode(data) case .deleteModel(let data): request.httpBody = try JSONEncoder.default.encode(data) + case .pullModel(let data): + request.httpBody = try JSONEncoder.default.encode(data) case .embeddings(let data): request.httpBody = try JSONEncoder.default.encode(data) default: