From b24ee1b83350cd72ff0f0f0e81212e1b9c0789f7 Mon Sep 17 00:00:00 2001 From: Kevin Hermawan <84965338+kevinhermawan@users.noreply.github.com> Date: Sat, 30 Mar 2024 17:03:35 +0700 Subject: [PATCH] refactor: renames `generateEmbeddings` to `embeddings` --- Sources/OllamaKit/OllamaKit+Embeddings.swift | 60 +++++++++++++++++++ .../OllamaKit+GenerateEmbeddings.swift | 60 ------------------- ...ta.swift => OKEmbeddingsRequestData.swift} | 4 +- .../RequestData/OKGenerateRequestData.swift | 4 +- .../Responses/OKEmbeddingsResponse.swift | 15 +++++ .../OKGenerateEmbeddingsResponse.swift | 15 ----- .../Responses/OKGenerateResponse.swift | 4 +- Sources/OllamaKit/Utils/OKRouter.swift | 8 +-- Tests/OllamaKitTests/OllamaKitTests.swift | 2 +- 9 files changed, 86 insertions(+), 86 deletions(-) create mode 100644 Sources/OllamaKit/OllamaKit+Embeddings.swift delete mode 100644 Sources/OllamaKit/OllamaKit+GenerateEmbeddings.swift rename Sources/OllamaKit/RequestData/{OKGenerateEmbeddingsRequestData.swift => OKEmbeddingsRequestData.swift} (88%) create mode 100644 Sources/OllamaKit/Responses/OKEmbeddingsResponse.swift delete mode 100644 Sources/OllamaKit/Responses/OKGenerateEmbeddingsResponse.swift diff --git a/Sources/OllamaKit/OllamaKit+Embeddings.swift b/Sources/OllamaKit/OllamaKit+Embeddings.swift new file mode 100644 index 0000000..321d94d --- /dev/null +++ b/Sources/OllamaKit/OllamaKit+Embeddings.swift @@ -0,0 +1,60 @@ +// +// OllamaKit+Embeddings.swift +// +// +// Created by Paul Thrasher on 02/09/24. +// + +import Alamofire +import Combine +import Foundation + +extension OllamaKit { + /// Asynchronously generates embeddings from a specific model from the Ollama API. + /// + /// This method accepts ``OKEmbeddingsRequestData`` and returns an ``OKEmbeddingsResponse`` containing embeddings from the requested model. + /// + /// ```swift + /// let ollamaKit = OllamaKit() + /// let requestData = OKEmbeddingsRequestData(/* parameters */) + /// let embeddings = try await ollamaKit.embeddings(data: requestData) + /// ``` + /// + /// - Parameter data: The ``OKEmbeddingsRequestData`` used to query the API for generating specific model embeddings. + /// - Returns: An ``OKEmbeddingsResponse`` containing the embeddings from the model. + /// - Throws: An error if the request fails or the response can't be decoded. + public func embeddings(data: OKEmbeddingsRequestData) async throws -> OKEmbeddingsResponse { + let request = AF.request(router.embeddings(data: data)).validate() + let response = request.serializingDecodable(OKEmbeddingsResponse.self, decoder: decoder) + let value = try await response.value + + return value + } + + /// Retrieves embeddings from a specific model from the Ollama API as a Combine publisher. + /// + /// This method provides a reactive approach to generate embeddings. It accepts ``OKEmbeddingsRequestData`` and returns a Combine publisher that emits an ``OKEmbeddingsResponse`` upon successful retrieval. + /// + /// ```swift + /// let ollamaKit = OllamaKit() + /// let requestData = OKEmbeddingsRequestData(/* parameters */) + /// + /// ollamaKit.embeddings(data: requestData) + /// .sink(receiveCompletion: { completion in + /// // Handle completion + /// }, receiveValue: { response in + /// // Handle the received embeddings response + /// }) + /// .store(in: &cancellables) + /// ``` + /// + /// - Parameter data: The ``OKEmbeddingsRequestData`` used to query the API for embeddings from a specific model. + /// - Returns: A `AnyPublisher` that emits embeddings. + public func embeddings(data: OKEmbeddingsRequestData) -> AnyPublisher { + let request = AF.request(router.embeddings(data: data)).validate() + + return request + .publishDecodable(type: OKEmbeddingsResponse.self, decoder: decoder).value() + .eraseToAnyPublisher() + } +} diff --git a/Sources/OllamaKit/OllamaKit+GenerateEmbeddings.swift b/Sources/OllamaKit/OllamaKit+GenerateEmbeddings.swift deleted file mode 100644 index cdeb0e0..0000000 --- a/Sources/OllamaKit/OllamaKit+GenerateEmbeddings.swift +++ /dev/null @@ -1,60 +0,0 @@ -// -// OllamaKit+GenerateEmbeddings.swift -// -// -// Created by Paul Thrasher on 02/09/24. -// - -import Alamofire -import Combine -import Foundation - -extension OllamaKit { - /// Asynchronously generates embeddings from a specific model from the Ollama API. - /// - /// This method accepts ``OKGenerateEmbeddingsRequestData`` and returns an ``OKGenerateEmbeddingsResponse`` containing embeddings from the requested model. - /// - /// ```swift - /// let ollamaKit = OllamaKit() - /// let requestData = OKGenerateEmbeddingsRequestData(/* parameters */) - /// let generateEmbeddings = try await ollamaKit.generateEmbeddings(data: requestData) - /// ``` - /// - /// - Parameter data: The ``OKGenerateEmbeddingsRequestData`` used to query the API for generating specific model embeddings. - /// - Returns: An ``OKGenerateEmbeddingsResponse`` containing the embeddings from the model. - /// - Throws: An error if the request fails or the response can't be decoded. - public func generateEmbeddings(data: OKGenerateEmbeddingsRequestData) async throws -> OKGenerateEmbeddingsResponse { - let request = AF.request(router.generateEmbeddings(data: data)).validate() - let response = request.serializingDecodable(OKGenerateEmbeddingsResponse.self, decoder: decoder) - let value = try await response.value - - return value - } - - /// Retrieves embeddings from a specific model from the Ollama API as a Combine publisher. - /// - /// This method provides a reactive approach to generate embeddings. It accepts ``OKGenerateEmbeddingsRequestData`` and returns a Combine publisher that emits an ``OKGenerateEmbeddingsResponse`` upon successful retrieval. - /// - /// ```swift - /// let ollamaKit = OllamaKit() - /// let requestData = OKGenerateEmbeddingsRequestData(/* parameters */) - /// - /// ollamaKit.generateEmbeddings(data: requestData) - /// .sink(receiveCompletion: { completion in - /// // Handle completion - /// }, receiveValue: { generateEmbeddingsResponse in - /// // Handle the received model info response - /// }) - /// .store(in: &cancellables) - /// ``` - /// - /// - Parameter data: The ``OKGenerateEmbeddingsRequestData`` used to query the API for embeddings from a specific model. - /// - Returns: A `AnyPublisher` that emits embeddings. - public func generateEmbeddings(data: OKGenerateEmbeddingsRequestData) -> AnyPublisher { - let request = AF.request(router.generateEmbeddings(data: data)).validate() - - return request - .publishDecodable(type: OKGenerateEmbeddingsResponse.self, decoder: decoder).value() - .eraseToAnyPublisher() - } -} diff --git a/Sources/OllamaKit/RequestData/OKGenerateEmbeddingsRequestData.swift b/Sources/OllamaKit/RequestData/OKEmbeddingsRequestData.swift similarity index 88% rename from Sources/OllamaKit/RequestData/OKGenerateEmbeddingsRequestData.swift rename to Sources/OllamaKit/RequestData/OKEmbeddingsRequestData.swift index 826ebc6..1aa8774 100644 --- a/Sources/OllamaKit/RequestData/OKGenerateEmbeddingsRequestData.swift +++ b/Sources/OllamaKit/RequestData/OKEmbeddingsRequestData.swift @@ -1,5 +1,5 @@ // -// OKGenerateEmbeddingsRequestData.swift +// OKEmbeddingsRequestData.swift // // // Created by Paul Thrasher on 02/09/24. @@ -8,7 +8,7 @@ import Foundation /// A structure that encapsulates the data required for generating embeddings using the Ollama API. -public struct OKGenerateEmbeddingsRequestData: Encodable { +public struct OKEmbeddingsRequestData: Encodable { /// A string representing the identifier of the model. public let model: String diff --git a/Sources/OllamaKit/RequestData/OKGenerateRequestData.swift b/Sources/OllamaKit/RequestData/OKGenerateRequestData.swift index bda80f1..6cab851 100644 --- a/Sources/OllamaKit/RequestData/OKGenerateRequestData.swift +++ b/Sources/OllamaKit/RequestData/OKGenerateRequestData.swift @@ -23,8 +23,8 @@ public struct OKGenerateRequestData: Encodable { /// An optional string specifying the system message. public var system: String? - /// An optional array of integers representing contextual information. - public var context: [Int]? + /// An optional array of doubles representing contextual information. + public var context: [Double]? /// Optional ``OKCompletionOptions`` providing additional configuration for the generation request. public var options: OKCompletionOptions? diff --git a/Sources/OllamaKit/Responses/OKEmbeddingsResponse.swift b/Sources/OllamaKit/Responses/OKEmbeddingsResponse.swift new file mode 100644 index 0000000..a85d82c --- /dev/null +++ b/Sources/OllamaKit/Responses/OKEmbeddingsResponse.swift @@ -0,0 +1,15 @@ +// +// OKEmbeddingsResponse.swift +// +// +// Created by Paul Thrasher on 02/09/24. +// + +import Foundation + +/// A structure that represents the response to an embedding request from the Ollama API. +public struct OKEmbeddingsResponse: Decodable { + + /// An array of doubles representing the embeddings of the input prompt. + public let embedding: [Double]? +} diff --git a/Sources/OllamaKit/Responses/OKGenerateEmbeddingsResponse.swift b/Sources/OllamaKit/Responses/OKGenerateEmbeddingsResponse.swift deleted file mode 100644 index 47af714..0000000 --- a/Sources/OllamaKit/Responses/OKGenerateEmbeddingsResponse.swift +++ /dev/null @@ -1,15 +0,0 @@ -// -// OKGenerateEmbeddingsResponse.swift -// -// -// Created by Paul Thrasher on 02/09/24. -// - -import Foundation - -/// A structure that represents the response to an embedding request from the Ollama API. -public struct OKGenerateEmbeddingsResponse: Decodable { - - /// An array of floats representing the embeddings of the input prompt. - public let embedding: [Float]? -} diff --git a/Sources/OllamaKit/Responses/OKGenerateResponse.swift b/Sources/OllamaKit/Responses/OKGenerateResponse.swift index fe52331..1d9afc9 100644 --- a/Sources/OllamaKit/Responses/OKGenerateResponse.swift +++ b/Sources/OllamaKit/Responses/OKGenerateResponse.swift @@ -18,8 +18,8 @@ public struct OKGenerateResponse: OKCompletionResponse, Decodable { /// A string containing the generated content. public let response: String - /// An optional array of integers representing contextual information used in the generation. - public let context: [Int]? + /// An optional array of doubles representing contextual information used in the generation. + public let context: [Double]? /// A boolean indicating whether the generation process is complete. public let done: Bool diff --git a/Sources/OllamaKit/Utils/OKRouter.swift b/Sources/OllamaKit/Utils/OKRouter.swift index 341c565..22aa198 100644 --- a/Sources/OllamaKit/Utils/OKRouter.swift +++ b/Sources/OllamaKit/Utils/OKRouter.swift @@ -18,7 +18,7 @@ internal enum OKRouter { case chat(data: OKChatRequestData) case copyModel(data: OKCopyModelRequestData) case deleteModel(data: OKDeleteModelRequestData) - case generateEmbeddings(data: OKGenerateEmbeddingsRequestData) + case embeddings(data: OKEmbeddingsRequestData) internal var path: String { switch self { @@ -36,7 +36,7 @@ internal enum OKRouter { return "/api/copy" case .deleteModel: return "/api/delete" - case .generateEmbeddings: + case .embeddings: return "/api/embeddings" } } @@ -57,7 +57,7 @@ internal enum OKRouter { return .post case .deleteModel: return .delete - case .generateEmbeddings: + case .embeddings: return .post } } @@ -85,7 +85,7 @@ extension OKRouter: URLRequestConvertible { request.httpBody = try JSONEncoder.default.encode(data) case .deleteModel(let data): request.httpBody = try JSONEncoder.default.encode(data) - case .generateEmbeddings(let data): + case .embeddings(let data): request.httpBody = try JSONEncoder.default.encode(data) default: break diff --git a/Tests/OllamaKitTests/OllamaKitTests.swift b/Tests/OllamaKitTests/OllamaKitTests.swift index f64514a..6269baa 100644 --- a/Tests/OllamaKitTests/OllamaKitTests.swift +++ b/Tests/OllamaKitTests/OllamaKitTests.swift @@ -51,7 +51,7 @@ final class OllamaKitTests: XCTestCase { } - func testGenerateEmbeddingsFailure() async throws { + func testEmbeddingsFailure() async throws { } }