Skip to content

Commit 061514a

Browse files
lzelllzell
authored
Add support for OpenAI embeddings (#106)
1 parent 6c1f567 commit 061514a

7 files changed

+287
-2
lines changed

‎README.md

+61-1
Original file line numberDiff line numberDiff line change
@@ -817,18 +817,78 @@ This example it taken from OpenAI's [function calling guide](https://platform.op
817817
}
818818
```
819819

820+
### How to get embeddings using OpenAI
821+
822+
```swift
823+
import AIProxy
824+
825+
/* Uncomment for BYOK use cases */
826+
// let openAIService = AIProxy.openAIDirectService(
827+
// unprotectedAPIKey: "your-openai-key"
828+
// )
829+
830+
/* Uncomment for all other production use cases */
831+
// let openAIService = AIProxy.openAIService(
832+
// partialKey: "partial-key-from-your-developer-dashboard",
833+
// serviceURL: "service-url-from-your-developer-dashboard"
834+
// )
835+
836+
let requestBody = OpenAIEmbeddingRequestBody(
837+
input: .text("hello world"),
838+
model: "text-embedding-3-small"
839+
)
840+
841+
// Or, for multiple embeddings from strings:
842+
843+
/*
844+
let requestBody = OpenAIEmbeddingRequestBody(
845+
input: .textArray([
846+
"hello world",
847+
"hola mundo"
848+
]),
849+
model: "text-embedding-3-small"
850+
)
851+
*/
852+
853+
// Or, for multiple embeddings from tokens:
854+
855+
/*
856+
let requestBody = OpenAIEmbeddingRequestBody(
857+
input: .intArray([0,1,2]),
858+
model: "text-embedding-3-small"
859+
)
860+
*/
861+
862+
do {
863+
let response = try await openAIService.embeddingRequest(body: requestBody)
864+
print(
865+
"""
866+
The response contains \(response.embeddings.count) embeddings.
867+
868+
The first vector starts with \(response.embeddings.first?.vector.prefix(10) ?? [])
869+
"""
870+
)
871+
} catch AIProxyError.unsuccessfulRequest(let statusCode, let responseBody) {
872+
print("Received \(statusCode) status code with response body: \(responseBody)")
873+
} catch {
874+
print("Could not perform embedding request to OpenAI: \(error.localizedDescription)")
875+
}
876+
```
877+
878+
820879
### How to use OpenAI through an Azure deployment
821880

822881
You can use all of the OpenAI snippets aboves with one change. Initialize the OpenAI service with:
823882

883+
```swift
824884
import AIProxy
825885

826886
let openAIService = AIProxy.openAIService(
827887
partialKey: "partial-key-from-your-developer-dashboard",
828888
serviceURL: "service-url-from-your-developer-dashboard",
829889
requestFormat: .azureDeployment(apiVersion: "2024-06-01")
830890
)
831-
891+
```
832892

833893
***
834894

‎Sources/AIProxy/AIProxy.swift

+1-1
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ let aiproxyLogger = Logger(
1313
public struct AIProxy {
1414

1515
/// The current sdk version
16-
public static let sdkVersion = "0.70.0"
16+
public static let sdkVersion = "0.71.0"
1717

1818
/// - Parameters:
1919
/// - partialKey: Your partial key is displayed in the AIProxy dashboard when you submit your provider's key.

‎Sources/AIProxy/OpenAI/OpenAIDirectService.swift

+23
Original file line numberDiff line numberDiff line change
@@ -187,6 +187,29 @@ open class OpenAIDirectService: OpenAIService, DirectService {
187187
return try await self.makeRequestAndDeserializeResponse(request)
188188
}
189189

190+
/// Get a vector representation of a given input that can be easily consumed by machine learning models and algorithms. Related guide:
191+
/// https://platform.openai.com/docs/guides/embeddings
192+
/// - Parameters:
193+
/// - body: The request body to send to aiproxy and openai. See this reference:
194+
/// https://platform.openai.com/docs/api-reference/embeddings/create
195+
/// - Returns: An embedding response. See this reference:
196+
/// https://platform.openai.com/docs/api-reference/embeddings/object
197+
public func embeddingRequest(
198+
body: OpenAIEmbeddingRequestBody
199+
) async throws -> OpenAIEmbeddingResponseBody {
200+
let request = try AIProxyURLRequest.createDirect(
201+
baseURL: "https://api.openai.com",
202+
path: self.resolvedPath("embeddings"),
203+
body: try body.serialize(),
204+
verb: .post,
205+
contentType: "application/json",
206+
additionalHeaders: [
207+
"Authorization": "Bearer \(self.unprotectedAPIKey)"
208+
]
209+
)
210+
return try await self.makeRequestAndDeserializeResponse(request)
211+
}
212+
190213
private func resolvedPath(_ common: String) -> String {
191214
assert(common[common.startIndex] != "/")
192215
switch self.requestFormat {

‎Sources/AIProxy/OpenAI/OpenAIEmbeddingRequestBody.swift

+112
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
//
2+
// OpenAIEmbeddingRequestBody.swift
3+
// AIProxy
4+
//
5+
// Created by Lou Zell on 2/16/25.
6+
//
7+
8+
/// From OpenAI's docs:
9+
///
10+
/// > Embeddings are a numerical representation of text that can be used to measure the
11+
/// relatedness between two pieces of text. Embeddings are useful for search, clustering,
12+
/// recommendations, anomaly detection, and classification tasks. You can read more about our
13+
/// latest embedding models in the announcement blog post."
14+
///
15+
/// References:
16+
/// - https://openai.com/blog/new-embedding-models-and-api-updates
17+
/// - https://platform.openai.com/docs/api-reference/embeddings/create
18+
public struct OpenAIEmbeddingRequestBody: Encodable {
19+
20+
// Required
21+
22+
/// Input text to embed, encoded as a string or array of tokens. To embed multiple inputs
23+
/// in a single request, pass an array of strings or array of token arrays. The input must
24+
/// not exceed the max input tokens for the model (8192 tokens for text-embedding-ada-002),
25+
/// cannot be an empty string, and any array must be 2048 dimensions or less. Example
26+
/// Python code for counting tokens. Some models may also impose a limit on total number of
27+
/// tokens summed across inputs.
28+
public let input: EmbeddingInput
29+
30+
/// ID of the model to use.
31+
/// See the model list here: https://platform.openai.com/docs/models#embeddings
32+
///
33+
/// text-embedding-3-large
34+
/// Most capable embedding model for both english and non-english tasks.
35+
/// Output dimension: 3,072
36+
///
37+
/// text-embedding-3-small
38+
/// Increased performance over 2nd generation ada embedding model.
39+
/// Output dimension: 1,536
40+
///
41+
/// text-embedding-ada-002
42+
/// Most capable 2nd generation embedding model, replacing 16 first generation models.
43+
/// Output dimension: 1,536
44+
public let model: String
45+
46+
// Optional
47+
48+
/// The format to return the embeddings in.
49+
/// Defaults to float
50+
public let encodingFormat: EncodingFormat?
51+
52+
/// The number of dimensions the resulting output embeddings should have. Only supported in
53+
/// `text-embedding-3` and later models.
54+
public let dimensions: Int?
55+
56+
/// A unique identifier representing your end-user, which can help OpenAI to monitor and
57+
/// detect abuse.
58+
public let user: String?
59+
60+
private enum CodingKeys: String, CodingKey {
61+
case input
62+
case model
63+
64+
case encodingFormat = "encoding_format"
65+
case dimensions
66+
case user
67+
}
68+
69+
// This memberwise initializer is autogenerated.
70+
// To regenerate, use `cmd-shift-a` > Generate Memberwise Initializer
71+
// To format, place the cursor in the initializer's parameter list and use `ctrl-m`
72+
public init(
73+
input: OpenAIEmbeddingRequestBody.EmbeddingInput,
74+
model: String,
75+
encodingFormat: OpenAIEmbeddingRequestBody.EncodingFormat? = nil,
76+
dimensions: Int? = nil,
77+
user: String? = nil
78+
) {
79+
self.input = input
80+
self.model = model
81+
self.encodingFormat = encodingFormat
82+
self.dimensions = dimensions
83+
self.user = user
84+
}
85+
}
86+
87+
extension OpenAIEmbeddingRequestBody {
88+
public enum EmbeddingInput: Encodable {
89+
case text(String)
90+
case textArray([String])
91+
case intArray([Int])
92+
93+
public func encode(to encoder: any Encoder) throws {
94+
var container = encoder.singleValueContainer()
95+
switch self {
96+
case .text(let text):
97+
try container.encode(text)
98+
case .textArray(let arr):
99+
try container.encode(arr)
100+
case .intArray(let arr):
101+
try container.encode(arr)
102+
}
103+
}
104+
}
105+
}
106+
107+
extension OpenAIEmbeddingRequestBody {
108+
public enum EncodingFormat: String, Encodable {
109+
case float
110+
case base64
111+
}
112+
}

‎Sources/AIProxy/OpenAI/OpenAIEmbeddingResponseBody.swift

+56
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
//
2+
// OpenAIEmbeddingResponseBody.swift
3+
// AIProxy
4+
//
5+
// Created by Lou Zell on 2/16/25.
6+
//
7+
8+
/// https://platform.openai.com/docs/api-reference/embeddings/object
9+
public struct OpenAIEmbeddingResponseBody: Decodable {
10+
public let embeddings: [Embedding]
11+
public let model: String?
12+
public let usage: Usage?
13+
14+
private enum CodingKeys: CodingKey {
15+
case data
16+
case model
17+
case usage
18+
}
19+
20+
public init(from decoder: any Decoder) throws {
21+
let container = try decoder.container(keyedBy: CodingKeys.self)
22+
self.embeddings = try container.decode([Embedding].self, forKey: .data)
23+
self.model = try container.decodeIfPresent(String.self, forKey: .model)
24+
self.usage = try container.decodeIfPresent(Usage.self, forKey: .usage)
25+
}
26+
}
27+
28+
extension OpenAIEmbeddingResponseBody {
29+
public struct Embedding: Decodable {
30+
public let vector: [Double]
31+
public let index: Int?
32+
33+
private enum CodingKeys: CodingKey {
34+
case embedding
35+
case index
36+
}
37+
38+
public init(from decoder: any Decoder) throws {
39+
let container = try decoder.container(keyedBy: CodingKeys.self)
40+
self.vector = try container.decode([Double].self, forKey: .embedding)
41+
self.index = try container.decodeIfPresent(Int.self, forKey: .index)
42+
}
43+
}
44+
}
45+
46+
extension OpenAIEmbeddingResponseBody {
47+
public struct Usage: Decodable {
48+
public let promptTokens: Int
49+
public let totalTokens: Int
50+
51+
private enum CodingKeys: String, CodingKey {
52+
case promptTokens = "prompt_tokens"
53+
case totalTokens = "total_tokens"
54+
}
55+
}
56+
}

‎Sources/AIProxy/OpenAI/OpenAIProxiedService.swift

+23
Original file line numberDiff line numberDiff line change
@@ -188,6 +188,29 @@ open class OpenAIProxiedService: OpenAIService, ProxiedService {
188188
return try await self.makeRequestAndDeserializeResponse(request)
189189
}
190190

191+
/// Get a vector representation of a given input that can be easily consumed by machine learning models and algorithms. Related guide:
192+
/// https://platform.openai.com/docs/guides/embeddings
193+
/// - Parameters:
194+
/// - body: The request body to send to aiproxy and openai. See this reference:
195+
/// https://platform.openai.com/docs/api-reference/embeddings/create
196+
/// - Returns: An embedding response. See this reference:
197+
/// https://platform.openai.com/docs/api-reference/embeddings/object
198+
public func embeddingRequest(
199+
body: OpenAIEmbeddingRequestBody
200+
) async throws -> OpenAIEmbeddingResponseBody {
201+
let request = try await AIProxyURLRequest.create(
202+
partialKey: self.partialKey,
203+
serviceURL: self.serviceURL ?? legacyURL,
204+
clientID: self.clientID,
205+
proxyPath: self.resolvedPath("embeddings"),
206+
body: try body.serialize(),
207+
verb: .post,
208+
contentType: "application/json"
209+
)
210+
return try await self.makeRequestAndDeserializeResponse(request)
211+
}
212+
213+
191214
private func resolvedPath(_ common: String) -> String {
192215
assert(common[common.startIndex] != "/")
193216
switch self.requestFormat {

‎Sources/AIProxy/OpenAI/OpenAIService.swift

+11
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,17 @@ public protocol OpenAIService {
7878
func moderationRequest(
7979
body: OpenAIModerationRequestBody
8080
) async throws -> OpenAIModerationResponseBody
81+
82+
/// Get a vector representation of a given input that can be easily consumed by machine learning models and algorithms. Related guide:
83+
/// https://platform.openai.com/docs/guides/embeddings
84+
/// - Parameters:
85+
/// - body: The request body to send to aiproxy and openai. See this reference:
86+
/// https://platform.openai.com/docs/api-reference/embeddings/create
87+
/// - Returns: An embedding response. See this reference:
88+
/// https://platform.openai.com/docs/api-reference/embeddings/object
89+
func embeddingRequest(
90+
body: OpenAIEmbeddingRequestBody
91+
) async throws -> OpenAIEmbeddingResponseBody
8192
}
8293

8394
extension OpenAIService {

0 commit comments

Comments
 (0)