Skip to content

Commit ecf740e

Browse files
committed
Migrate service specific info to their own structs
1 parent 37111d2 commit ecf740e

4 files changed

Lines changed: 90 additions & 33 deletions

File tree

Tool/Sources/AIModel/ChatModel.swift

Lines changed: 36 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,24 @@ public struct ChatModel: Codable, Equatable, Identifiable {
2525
}
2626

2727
public struct Info: Codable, Equatable {
28+
public struct OllamaInfo: Codable, Equatable {
29+
@FallbackDecoding<EmptyString>
30+
public var keepAlive: String
31+
32+
public init(keepAlive: String = "") {
33+
self.keepAlive = keepAlive
34+
}
35+
}
36+
37+
public struct OpenAIInfo: Codable, Equatable {
38+
@FallbackDecoding<EmptyString>
39+
public var organizationID: String
40+
41+
public init(organizationID: String = "") {
42+
self.organizationID = organizationID
43+
}
44+
}
45+
2846
@FallbackDecoding<EmptyString>
2947
public var apiKeyName: String
3048
@FallbackDecoding<EmptyString>
@@ -39,12 +57,11 @@ public struct ChatModel: Codable, Equatable, Identifiable {
3957
public var supportsOpenAIAPI2023_11: Bool
4058
@FallbackDecoding<EmptyString>
4159
public var modelName: String
42-
public var azureOpenAIDeploymentName: String {
43-
get { modelName }
44-
set { modelName = newValue }
45-
}
46-
@FallbackDecoding<EmptyString>
47-
public var ollamaKeepAlive: String
60+
61+
@FallbackDecoding<EmptyChatModelOpenAIInfo>
62+
public var openAIInfo: OpenAIInfo
63+
@FallbackDecoding<EmptyChatModelOllamaInfo>
64+
public var ollamaInfo: OllamaInfo
4865

4966
public init(
5067
apiKeyName: String = "",
@@ -54,7 +71,8 @@ public struct ChatModel: Codable, Equatable, Identifiable {
5471
supportsFunctionCalling: Bool = true,
5572
supportsOpenAIAPI2023_11: Bool = false,
5673
modelName: String = "",
57-
ollamaKeepAlive: String = ""
74+
openAIInfo: OpenAIInfo = OpenAIInfo(),
75+
ollamaInfo: OllamaInfo = OllamaInfo()
5876
) {
5977
self.apiKeyName = apiKeyName
6078
self.baseURL = baseURL
@@ -63,7 +81,8 @@ public struct ChatModel: Codable, Equatable, Identifiable {
6381
self.supportsFunctionCalling = supportsFunctionCalling
6482
self.supportsOpenAIAPI2023_11 = supportsOpenAIAPI2023_11
6583
self.modelName = modelName
66-
self.ollamaKeepAlive = ollamaKeepAlive
84+
self.openAIInfo = openAIInfo
85+
self.ollamaInfo = ollamaInfo
6786
}
6887
}
6988

@@ -80,7 +99,7 @@ public struct ChatModel: Codable, Equatable, Identifiable {
8099
return "\(baseURL)/v1/chat/completions"
81100
case .azureOpenAI:
82101
let baseURL = info.baseURL
83-
let deployment = info.azureOpenAIDeploymentName
102+
let deployment = info.modelName
84103
let version = "2024-02-15-preview"
85104
if baseURL.isEmpty { return "" }
86105
return "\(baseURL)/openai/deployments/\(deployment)/chat/completions?api-version=\(version)"
@@ -104,3 +123,11 @@ public struct EmptyChatModelFormat: FallbackValueProvider {
104123
public static var defaultValue: ChatModel.Format { .openAI }
105124
}
106125

126+
public struct EmptyChatModelOllamaInfo: FallbackValueProvider {
127+
public static var defaultValue: ChatModel.Info.OllamaInfo { .init() }
128+
}
129+
130+
public struct EmptyChatModelOpenAIInfo: FallbackValueProvider {
131+
public static var defaultValue: ChatModel.Info.OpenAIInfo { .init() }
132+
}
133+

Tool/Sources/AIModel/EmbeddingModel.swift

Lines changed: 16 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
1-
import Foundation
21
import CodableWrappers
2+
import Foundation
33

44
public struct EmbeddingModel: Codable, Equatable, Identifiable {
55
public var id: String
@@ -24,6 +24,9 @@ public struct EmbeddingModel: Codable, Equatable, Identifiable {
2424
}
2525

2626
public struct Info: Codable, Equatable {
27+
public typealias OllamaInfo = ChatModel.Info.OllamaInfo
28+
public typealias OpenAIInfo = ChatModel.Info.OpenAIInfo
29+
2730
@FallbackDecoding<EmptyString>
2831
public var apiKeyName: String
2932
@FallbackDecoding<EmptyString>
@@ -36,12 +39,11 @@ public struct EmbeddingModel: Codable, Equatable, Identifiable {
3639
public var dimensions: Int
3740
@FallbackDecoding<EmptyString>
3841
public var modelName: String
39-
public var azureOpenAIDeploymentName: String {
40-
get { modelName }
41-
set { modelName = newValue }
42-
}
43-
@FallbackDecoding<EmptyString>
44-
public var ollamaKeepAlive: String
42+
43+
@FallbackDecoding<EmptyChatModelOpenAIInfo>
44+
public var openAIInfo: OpenAIInfo
45+
@FallbackDecoding<EmptyChatModelOllamaInfo>
46+
public var ollamaInfo: OllamaInfo
4547

4648
public init(
4749
apiKeyName: String = "",
@@ -50,18 +52,20 @@ public struct EmbeddingModel: Codable, Equatable, Identifiable {
5052
maxTokens: Int = 8192,
5153
dimensions: Int = 1536,
5254
modelName: String = "",
53-
ollamaKeepAlive: String = ""
55+
openAIInfo: OpenAIInfo = OpenAIInfo(),
56+
ollamaInfo: OllamaInfo = OllamaInfo()
5457
) {
5558
self.apiKeyName = apiKeyName
5659
self.baseURL = baseURL
5760
self.isFullURL = isFullURL
5861
self.maxTokens = maxTokens
5962
self.dimensions = dimensions
6063
self.modelName = modelName
61-
self.ollamaKeepAlive = ollamaKeepAlive
64+
self.openAIInfo = openAIInfo
65+
self.ollamaInfo = ollamaInfo
6266
}
6367
}
64-
68+
6569
public var endpoint: String {
6670
switch format {
6771
case .openAI:
@@ -75,7 +79,7 @@ public struct EmbeddingModel: Codable, Equatable, Identifiable {
7579
return "\(baseURL)/v1/embeddings"
7680
case .azureOpenAI:
7781
let baseURL = info.baseURL
78-
let deployment = info.azureOpenAIDeploymentName
82+
let deployment = info.modelName
7983
let version = "2024-02-15-preview"
8084
if baseURL.isEmpty { return "" }
8185
return "\(baseURL)/openai/deployments/\(deployment)/embeddings?api-version=\(version)"
@@ -87,11 +91,11 @@ public struct EmbeddingModel: Codable, Equatable, Identifiable {
8791
}
8892
}
8993

90-
9194
public struct EmptyEmbeddingModelInfo: FallbackValueProvider {
9295
public static var defaultValue: EmbeddingModel.Info { .init() }
9396
}
9497

9598
public struct EmptyEmbeddingModelFormat: FallbackValueProvider {
9699
public static var defaultValue: EmbeddingModel.Format { .openAI }
97100
}
101+

Tool/Sources/OpenAIService/APIs/OpenAIChatCompletionsService.swift

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,15 @@ actor OpenAIChatCompletionsService: ChatCompletionsStreamAPI, ChatCompletionsAPI
4545
request.setValue("application/json", forHTTPHeaderField: "Content-Type")
4646
if !apiKey.isEmpty {
4747
switch model.format {
48-
case .openAI, .openAICompatible:
48+
case .openAI:
49+
if !model.info.openAIInfo.organizationID.isEmpty {
50+
request.setValue(
51+
"OpenAI-Organization",
52+
forHTTPHeaderField: model.info.openAIInfo.organizationID
53+
)
54+
}
55+
request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
56+
case .openAICompatible:
4957
request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
5058
case .azureOpenAI:
5159
request.setValue(apiKey, forHTTPHeaderField: "api-key")
@@ -108,7 +116,15 @@ actor OpenAIChatCompletionsService: ChatCompletionsStreamAPI, ChatCompletionsAPI
108116
request.setValue("application/json", forHTTPHeaderField: "Content-Type")
109117
if !apiKey.isEmpty {
110118
switch model.format {
111-
case .openAI, .openAICompatible:
119+
case .openAI:
120+
if !model.info.openAIInfo.organizationID.isEmpty {
121+
request.setValue(
122+
"OpenAI-Organization",
123+
forHTTPHeaderField: model.info.openAIInfo.organizationID
124+
)
125+
}
126+
request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
127+
case .openAICompatible:
112128
request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
113129
case .azureOpenAI:
114130
request.setValue(apiKey, forHTTPHeaderField: "api-key")

Tool/Sources/OpenAIService/APIs/OpenAIEmbeddingService.swift

Lines changed: 20 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -33,11 +33,16 @@ struct OpenAIEmbeddingService: EmbeddingAPI {
3333
request.setValue("application/json", forHTTPHeaderField: "Content-Type")
3434
if !apiKey.isEmpty {
3535
switch model.format {
36-
case .openAI, .openAICompatible:
37-
request.setValue(
38-
"Bearer \(apiKey)",
39-
forHTTPHeaderField: "Authorization"
40-
)
36+
case .openAI:
37+
if model.info.openAIInfo.organizationID.isEmpty {
38+
request.setValue(
39+
"OpenAI-Organization",
40+
forHTTPHeaderField: model.info.openAIInfo.organizationID
41+
)
42+
}
43+
request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
44+
case .openAICompatible:
45+
request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
4146
case .azureOpenAI:
4247
request.setValue(apiKey, forHTTPHeaderField: "api-key")
4348
case .ollama:
@@ -84,11 +89,16 @@ struct OpenAIEmbeddingService: EmbeddingAPI {
8489
request.setValue("application/json", forHTTPHeaderField: "Content-Type")
8590
if !apiKey.isEmpty {
8691
switch model.format {
87-
case .openAI, .openAICompatible:
88-
request.setValue(
89-
"Bearer \(apiKey)",
90-
forHTTPHeaderField: "Authorization"
91-
)
92+
case .openAI:
93+
if model.info.openAIInfo.organizationID.isEmpty {
94+
request.setValue(
95+
"OpenAI-Organization",
96+
forHTTPHeaderField: model.info.openAIInfo.organizationID
97+
)
98+
}
99+
request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
100+
case .openAICompatible:
101+
request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
92102
case .azureOpenAI:
93103
request.setValue(apiKey, forHTTPHeaderField: "api-key")
94104
case .ollama:

0 commit comments

Comments
 (0)