Add embeddings

Sources/SwiftOpenAI/OpenAI/DataModels/Embedding/EmbeddingDataModel.swift

@@ -0,0 +1,19 @@
+import Foundation
+
+public struct EmbeddingResponseDataModel: Decodable {
+    public let object: String
+    public let data: [EmbeddingDataModel]
+    public let model: String
+    public let usage: EmbeddingUsageDataModel
+}
+
+public struct EmbeddingDataModel: Decodable {
+    public let object: String
+    public let embedding: [Float]
+    public let index: Int
+}
+
+public struct EmbeddingUsageDataModel: Decodable {
+    public let promptTokens: Int
+    public let totalTokens: Int
+}

Sources/SwiftOpenAI/OpenAI/DataModels/OpenAIModelType.swift

@@ -4,6 +4,7 @@ public enum OpenAIModelType {
     case gpt4(GPT4)
     case gpt3_5(GPT3_5)
     case edit(EditModel)
+    case embedding(EmbeddingModel)
 
     var name: String {
         switch self {
@@ -13,6 +14,8 @@ public enum OpenAIModelType {
             return gpt3_5Model.rawValue
         case .edit(let editModel):
             return editModel.rawValue
+        case .embedding(let embeddingModel):
+            return embeddingModel.rawValue
         }
     }
 }
@@ -36,3 +39,7 @@ public enum EditModel: String {
     case text_davinci_edit_001 = "text-davinci-edit-001"
     case code_davinci_edit_001 = "code-davinci-edit-001"
 }
+
+public enum EmbeddingModel: String {
+    case text_embedding_ada_002 = "text-embedding-ada-002"
+}

Sources/SwiftOpenAI/OpenAI/OpenAIEndpoints/List/EmbeddingsEndpoint.swift

@@ -0,0 +1,24 @@
+import Foundation
+
+struct CreateEmbeddingsEndpoint: Endpoint {
+    private let model: OpenAIModelType
+    private let input: String
+
+    var method: HTTPMethod {
+        .POST
+    }
+
+    var path: String = "embeddings"
+
+    init(model: OpenAIModelType,
+         input: String) {
+        self.model = model
+        self.input = input
+    }
+
+    var parameters: [String : Any]? {
+        ["model": self.model.name as Any,
+         "input": self.input as Any
+        ]
+    }
+}

Sources/SwiftOpenAI/OpenAI/OpenAIEndpoints/OpenAIEndpoints.swift

@@ -18,6 +18,8 @@ enum OpenAIEndpoints {
                      numberOfImages: Int,
                      size: ImageSize)
 
+    case embeddings(model: OpenAIModelType, input: String)
+
     public var endpoint: Endpoint {
         switch self {
         case .listModels:
@@ -37,6 +39,9 @@ enum OpenAIEndpoints {
             return CreateImageEndpoint(prompt: prompt,
                                        numberOfImages: numberOfImages,
                                        size: size)
+        case .embeddings(model: let model, input: let input):
+            return CreateEmbeddingsEndpoint(model: model,
+                                            input: input)
         }
     }
 }

Sources/SwiftOpenAI/OpenAI/Requests/Embeddings/EmbeddingsRequest.swift

@@ -0,0 +1,40 @@
+import Foundation
+
+protocol EmbeddingsRequestProtocol {
+    func execute(api: API,
+                 apiKey: String,
+                 model: OpenAIModelType,
+                 input: String) async throws -> EmbeddingResponseDataModel?
+}
+
+final public class EmbeddingsRequest: EmbeddingsRequestProtocol  {
+    public typealias Init = (_ api: API,
+                             _ apiKey: String,
+                             _ model: OpenAIModelType,
+                             _ input: String) async throws -> EmbeddingResponseDataModel?
+
+    public init() { }
+
+    public func execute(api: API,
+                        apiKey: String,
+                        model: OpenAIModelType,
+                        input: String) async throws -> EmbeddingResponseDataModel? {
+        var endpoint = OpenAIEndpoints.embeddings(model: model, input: input).endpoint
+        api.routeEndpoint(&endpoint, environment: OpenAIEnvironmentV1())
+
+        var urlRequest = api.buildURLRequest(endpoint: endpoint)
+        api.addHeaders(urlRequest: &urlRequest,
+                       headers: ["Content-Type" : "application/json",
+                                 "Authorization" : "Bearer \(apiKey)"])
+
+        let result = await api.execute(with: urlRequest)
+
+        let jsonDecoder = JSONDecoder()
+        jsonDecoder.keyDecodingStrategy = .convertFromSnakeCase
+
+        return try api.parse(result,
+                             type: EmbeddingResponseDataModel.self,
+                             jsonDecoder: jsonDecoder,
+                             errorType: OpenAIAPIError.self)
+    }
+}

Sources/SwiftOpenAI/OpenAI/SwiftOpenAI.swift

@@ -19,6 +19,9 @@ protocol OpenAIProtocol {
                instruction: String) async throws -> EditsDataModel?
 
     func createImages(prompt: String, numberOfImages: Int, size: ImageSize) async throws -> CreateImageDataModel?
+
+    func embeddings(model: OpenAIModelType,
+                    input: String) async throws -> EmbeddingResponseDataModel?
 }
 
 public class SwiftOpenAI: OpenAIProtocol {
@@ -31,6 +34,7 @@ public class SwiftOpenAI: OpenAIProtocol {
     private let createChatCompletionsStreamRequest: CreateChatCompletionsStreamRequest.Init
     private let editsRequest: EditsRequest.Init
     private let createImagesRequest: CreateImagesRequest.Init
+    private let embeddingRequest: EmbeddingsRequest.Init
 
     public init(api: API = API(),
                 apiKey: String,
@@ -39,7 +43,8 @@ public class SwiftOpenAI: OpenAIProtocol {
                 createChatCompletionsRequest: @escaping CreateChatCompletionsRequest.Init = CreateChatCompletionsRequest().execute,
                 createChatCompletionsStreamRequest: @escaping CreateChatCompletionsStreamRequest.Init = CreateChatCompletionsStreamRequest().execute,
                 editsRequest: @escaping EditsRequest.Init = EditsRequest().execute,
-                createImagesRequest: @escaping CreateImagesRequest.Init = CreateImagesRequest().execute) {
+                createImagesRequest: @escaping CreateImagesRequest.Init = CreateImagesRequest().execute,
+                embeddingRequest: @escaping EmbeddingsRequest.Init = EmbeddingsRequest().execute) {
         self.api = api
         self.apiKey = apiKey
         self.listModelsRequest = listModelsRequest
@@ -48,6 +53,7 @@ public class SwiftOpenAI: OpenAIProtocol {
         self.createChatCompletionsStreamRequest = createChatCompletionsStreamRequest
         self.editsRequest = editsRequest
         self.createImagesRequest = createImagesRequest
+        self.embeddingRequest = embeddingRequest
     }
 
     public func listModels() async throws -> ModelListDataModel? {
@@ -79,4 +85,8 @@ public class SwiftOpenAI: OpenAIProtocol {
     public func createImages(prompt: String, numberOfImages: Int, size: ImageSize) async throws -> CreateImageDataModel? {
         try await createImagesRequest(api, apiKey, prompt, numberOfImages, size)
     }
+
+    public func embeddings(model: OpenAIModelType, input: String) async throws -> EmbeddingResponseDataModel? {
+        try await embeddingRequest(api, apiKey, model, input)
+    }
 }

Tests/SwiftOpenAITests/OpenAITests/Unit Tests/Embeddings/EmbeddingAPIClientSpec.swift

@@ -0,0 +1,52 @@
+import XCTest
+@testable import SwiftOpenAI
+
+final class EmbeddingAPIClientSpec: XCTestCase {
+    private var sut: EmbeddingsRequestProtocol!
+    private let model: OpenAIModelType = .embedding(.text_embedding_ada_002)
+    private let apiKey = "1234567890"
+    private let input = "What day of the wek is it?"
+
+    func testAsyncAPIRequest_ParsesValidJSONToChatCompletionsDataModel() async throws {
+        let json = loadJSON(name: "embeddings")
+
+        let api = API(requester: RequesterMock())
+        let endpoint = OpenAIEndpoints.embeddings(model: model, input: input).endpoint
+
+        sut = EmbeddingsRequest()
+
+        stubHTTP(model: model,
+                 endpoint: endpoint,
+                 json: json,
+                 statusCode: 200)
+
+        do {
+            let dataModel = try await sut.execute(api: api, apiKey: apiKey, model: model, input: input)
+            XCTAssertNotNil(dataModel)
+            XCTAssertEqual(dataModel?.object, "list")
+            XCTAssertEqual(dataModel?.model, "text-embedding-ada-002")
+            XCTAssertEqual(dataModel?.data[0].embedding.count, 1536)
+            XCTAssertEqual(dataModel?.data[0].object, "embedding")
+            XCTAssertEqual(dataModel?.data[0].embedding[0], 0.0023064255)
+            XCTAssertEqual(dataModel?.data[0].embedding[1], -0.009327292)
+            XCTAssertEqual(dataModel?.usage.promptTokens, 8)
+            XCTAssertEqual(dataModel?.usage.totalTokens, 8)
+        } catch {
+            XCTFail()
+        }
+    }
+
+    private func stubHTTP(model: OpenAIModelType,
+                          endpoint: Endpoint,
+                          json: Data,
+                          statusCode: Int) {
+
+        URLProtocolMock.completionHandler = { request in
+            let response = HTTPURLResponse(url: URL(string: endpoint.path)!,
+                                           statusCode: statusCode,
+                                           httpVersion: nil,
+                                           headerFields: [:])!
+            return (response, json)
+        }
+    }
+}

Tests/SwiftOpenAITests/OpenAITests/Unit Tests/Embeddings/EmbeddingEndpointSpec.swift

@@ -0,0 +1,21 @@
+import XCTest
+@testable import SwiftOpenAI
+
+final class EmbeddingEndpointSpec: XCTestCase {
+    func testEndpointCompletions_WithDavinciModelInput_CreatesCorrectEndpointParameters() throws {
+        let model: OpenAIModelType = .embedding(.text_embedding_ada_002)
+        let input = "What day of the wek is it?"
+
+        let sut = OpenAIEndpoints.embeddings(model: model, input: input).endpoint
+
+        let modelParameter = sut.parameters!["model"] as! String
+        let inputParameter = sut.parameters!["input"] as! String
+
+        XCTAssertEqual(sut.path, "embeddings")
+        XCTAssertEqual(sut.method, .POST)
+        XCTAssertEqual(sut.parameters?.count, 2)
+        XCTAssertEqual(modelParameter, model.name)
+        XCTAssertEqual(inputParameter, input)
+    }
+}
+

Tests/SwiftOpenAITests/OpenAITests/Unit Tests/Embeddings/EmbeddingParserSpec.swift

@@ -0,0 +1,25 @@
+import XCTest
+@testable import SwiftOpenAI
+
+final class EmbeddingParserSpec: XCTestCase {
+    private var api = API()
+
+    func testAsyncAPIRequest_ParsesValidJSONToEditsDataModel() async throws {
+        let jsonData = loadJSON(name: "embeddings")
+
+        let jsonDecoder = JSONDecoder()
+        jsonDecoder.keyDecodingStrategy = .convertFromSnakeCase
+
+        let dataModel = try! api.parse(.success(jsonData), type: EmbeddingResponseDataModel.self, jsonDecoder: jsonDecoder, errorType: OpenAIAPIError.self)
+
+        XCTAssertNotNil(dataModel)
+        XCTAssertEqual(dataModel?.object, "list")
+        XCTAssertEqual(dataModel?.model, "text-embedding-ada-002")
+        XCTAssertEqual(dataModel?.data[0].embedding.count, 1536)
+        XCTAssertEqual(dataModel?.data[0].object, "embedding")
+        XCTAssertEqual(dataModel?.data[0].embedding[0], 0.0023064255)
+        XCTAssertEqual(dataModel?.data[0].embedding[1], -0.009327292)
+        XCTAssertEqual(dataModel?.usage.promptTokens, 8)
+        XCTAssertEqual(dataModel?.usage.totalTokens, 8)
+    }
+}

Tests/SwiftOpenAITests/OpenAITests/Unit Tests/Embeddings/EmbeddingRequestSpec.swift

@@ -0,0 +1,25 @@
+import XCTest
+@testable import SwiftOpenAI
+
+final class EmbeddingRequestSpec: XCTestCase {
+    private let api = API()
+
+    func testRequest_CreatedWithCorrectHeaders() throws {
+        let apiKey = "1234567890"
+        let model: OpenAIModelType = .embedding(.text_embedding_ada_002)
+        let input = "What day of the wek is it?"
+
+        var endpoint = OpenAIEndpoints.embeddings(model: model, input: input).endpoint
+
+        api.routeEndpoint(&endpoint, environment: OpenAIEnvironmentV1())
+
+        var sut = api.buildURLRequest(endpoint: endpoint)
+        api.addHeaders(urlRequest: &sut,
+                       headers: ["Content-Type" : "application/json",
+                                 "Authorization" : "Bearer \(apiKey)"])
+
+        XCTAssertEqual(sut.allHTTPHeaderFields?.count, 2)
+        XCTAssertEqual(sut.allHTTPHeaderFields?["Content-Type"], "application/json")
+        XCTAssertEqual(sut.allHTTPHeaderFields?["Authorization"], "Bearer 1234567890")
+    }
+}