Skip to content

Commit

Permalink
refactor: renames generateEmbeddings to embeddings (#16)
Browse files Browse the repository at this point in the history
  • Loading branch information
kevinhermawan authored Mar 30, 2024
1 parent 1a6d2e9 commit 24202d0
Show file tree
Hide file tree
Showing 9 changed files with 86 additions and 86 deletions.
60 changes: 60 additions & 0 deletions Sources/OllamaKit/OllamaKit+Embeddings.swift
Original file line number Diff line number Diff line change
@@ -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<OKEmbeddingsResponse, AFError>` that emits embeddings.
public func embeddings(data: OKEmbeddingsRequestData) -> AnyPublisher<OKEmbeddingsResponse, AFError> {
let request = AF.request(router.embeddings(data: data)).validate()

return request
.publishDecodable(type: OKEmbeddingsResponse.self, decoder: decoder).value()
.eraseToAnyPublisher()
}
}
60 changes: 0 additions & 60 deletions Sources/OllamaKit/OllamaKit+GenerateEmbeddings.swift

This file was deleted.

Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
//
// OKGenerateEmbeddingsRequestData.swift
// OKEmbeddingsRequestData.swift
//
//
// Created by Paul Thrasher on 02/09/24.
Expand All @@ -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

Expand Down
4 changes: 2 additions & 2 deletions Sources/OllamaKit/RequestData/OKGenerateRequestData.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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?
Expand Down
15 changes: 15 additions & 0 deletions Sources/OllamaKit/Responses/OKEmbeddingsResponse.swift
Original file line number Diff line number Diff line change
@@ -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]?
}
15 changes: 0 additions & 15 deletions Sources/OllamaKit/Responses/OKGenerateEmbeddingsResponse.swift

This file was deleted.

4 changes: 2 additions & 2 deletions Sources/OllamaKit/Responses/OKGenerateResponse.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
8 changes: 4 additions & 4 deletions Sources/OllamaKit/Utils/OKRouter.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -36,7 +36,7 @@ internal enum OKRouter {
return "/api/copy"
case .deleteModel:
return "/api/delete"
case .generateEmbeddings:
case .embeddings:
return "/api/embeddings"
}
}
Expand All @@ -57,7 +57,7 @@ internal enum OKRouter {
return .post
case .deleteModel:
return .delete
case .generateEmbeddings:
case .embeddings:
return .post
}
}
Expand Down Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion Tests/OllamaKitTests/OllamaKitTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ final class OllamaKitTests: XCTestCase {

}

func testGenerateEmbeddingsFailure() async throws {
func testEmbeddingsFailure() async throws {

}
}

0 comments on commit 24202d0

Please sign in to comment.