Skip to content

Commit 40e005e

Browse files
intitniintitni
committed
Merge branch 'feature/custom-header-openai' into develop
2 parents f0dcf32 + a8f5b0a commit 40e005e

9 files changed

Lines changed: 150 additions & 6 deletions

File tree

‎Core/Sources/HostApp/AccountSettings/ChatModelManagement/ChatModelEdit.swift‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ struct ChatModelEdit {
3131
var enforceMessageOrder: Bool = false
3232
var openAIOrganizationID: String = ""
3333
var openAIProjectID: String = ""
34+
var customHeaders: [ChatModel.Info.CustomHeaderInfo.HeaderField] = []
3435
}
3536

3637
enum Action: Equatable, BindableAction {
@@ -205,7 +206,8 @@ extension ChatModel {
205206
),
206207
ollamaInfo: .init(keepAlive: state.ollamaKeepAlive),
207208
googleGenerativeAIInfo: .init(apiVersion: state.apiVersion),
208-
openAICompatibleInfo: .init(enforceMessageOrder: state.enforceMessageOrder)
209+
openAICompatibleInfo: .init(enforceMessageOrder: state.enforceMessageOrder),
210+
customHeaderInfo: .init(headers: state.customHeaders)
209211
)
210212
)
211213
}
@@ -227,7 +229,8 @@ extension ChatModel {
227229
baseURLSelection: .init(baseURL: info.baseURL, isFullURL: info.isFullURL),
228230
enforceMessageOrder: info.openAICompatibleInfo.enforceMessageOrder,
229231
openAIOrganizationID: info.openAIInfo.organizationID,
230-
openAIProjectID: info.openAIInfo.projectID
232+
openAIProjectID: info.openAIInfo.projectID,
233+
customHeaders: info.customHeaderInfo.headers
231234
)
232235
}
233236
}

‎Core/Sources/HostApp/AccountSettings/ChatModelManagement/ChatModelEditView.swift‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -285,6 +285,7 @@ struct ChatModelEditView: View {
285285

286286
struct OpenAICompatibleForm: View {
287287
@Perception.Bindable var store: StoreOf<ChatModelEdit>
288+
@State var isEditingCustomHeader = false
288289

289290
var body: some View {
290291
WithPerceptionTracking {
@@ -320,6 +321,12 @@ struct ChatModelEditView: View {
320321
Toggle(isOn: $store.enforceMessageOrder) {
321322
Text("Enforce message order to be user/assistant alternated")
322323
}
324+
325+
Button("Custom Headers") {
326+
isEditingCustomHeader.toggle()
327+
}
328+
}.sheet(isPresented: $isEditingCustomHeader) {
329+
CustomHeaderSettingsView(headers: $store.customHeaders)
323330
}
324331
}
325332
}

‎Core/Sources/HostApp/AccountSettings/CustomHeaderSettingsView.swift‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
import AIModel
2+
import Foundation
3+
import SwiftUI
4+
5+
struct CustomHeaderSettingsView: View {
6+
@Binding var headers: [ChatModel.Info.CustomHeaderInfo.HeaderField]
7+
@Environment(\.dismiss) var dismiss
8+
@State private var newKey = ""
9+
@State private var newValue = ""
10+
11+
var body: some View {
12+
VStack {
13+
List {
14+
ForEach(headers.indices, id: \.self) { index in
15+
HStack {
16+
TextField("Key", text: Binding(
17+
get: { headers[index].key },
18+
set: { newKey in
19+
headers[index].key = newKey
20+
}
21+
))
22+
TextField("Value", text: Binding(
23+
get: { headers[index].value },
24+
set: { headers[index].value = $0 }
25+
))
26+
Button(action: {
27+
headers.remove(at: index)
28+
}) {
29+
Image(systemName: "trash")
30+
.foregroundColor(.red)
31+
}
32+
}
33+
}
34+
35+
HStack {
36+
TextField("New Key", text: $newKey)
37+
TextField("New Value", text: $newValue)
38+
Button(action: {
39+
if !newKey.isEmpty {
40+
headers.append(ChatModel.Info.CustomHeaderInfo.HeaderField(
41+
key: newKey,
42+
value: newValue
43+
))
44+
newKey = ""
45+
newValue = ""
46+
}
47+
}) {
48+
Image(systemName: "plus.circle.fill")
49+
.foregroundColor(.green)
50+
}
51+
}
52+
}
53+
54+
HStack {
55+
Spacer()
56+
Button("Done") {
57+
dismiss()
58+
}
59+
}.padding()
60+
}
61+
.frame(height: 500)
62+
}
63+
}
64+
65+
#Preview {
66+
struct V: View {
67+
@State var headers: [ChatModel.Info.CustomHeaderInfo.HeaderField] = [
68+
.init(key: "key", value: "value"),
69+
.init(key: "key2", value: "value2"),
70+
]
71+
var body: some View {
72+
CustomHeaderSettingsView(headers: $headers)
73+
}
74+
}
75+
76+
return V()
77+
}
78+

‎Core/Sources/HostApp/AccountSettings/EmbeddingModelManagement/EmbeddingModelEdit.swift‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ struct EmbeddingModelEdit {
2626
var suggestedMaxTokens: Int?
2727
var apiKeySelection: APIKeySelection.State = .init()
2828
var baseURLSelection: BaseURLSelection.State = .init()
29+
var customHeaders: [ChatModel.Info.CustomHeaderInfo.HeaderField] = []
2930
}
3031

3132
enum Action: Equatable, BindableAction {
@@ -168,7 +169,8 @@ extension EmbeddingModel {
168169
isFullURL: state.isFullURL,
169170
maxTokens: state.maxTokens,
170171
modelName: state.modelName.trimmingCharacters(in: .whitespacesAndNewlines),
171-
ollamaInfo: .init(keepAlive: state.ollamaKeepAlive)
172+
ollamaInfo: .init(keepAlive: state.ollamaKeepAlive),
173+
customHeaderInfo: .init(headers: state.customHeaders)
172174
)
173175
)
174176
}
@@ -188,7 +190,8 @@ extension EmbeddingModel {
188190
baseURLSelection: .init(
189191
baseURL: info.baseURL,
190192
isFullURL: info.isFullURL
191-
)
193+
),
194+
customHeaders: info.customHeaderInfo.headers
192195
)
193196
}
194197
}

‎Core/Sources/HostApp/AccountSettings/EmbeddingModelManagement/EmbeddingModelEditView.swift‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -248,6 +248,7 @@ struct EmbeddingModelEditView: View {
248248

249249
struct OpenAICompatibleForm: View {
250250
@Perception.Bindable var store: StoreOf<EmbeddingModelEdit>
251+
@State var isEditingCustomHeader = false
251252

252253
var body: some View {
253254
WithPerceptionTracking {
@@ -278,6 +279,12 @@ struct EmbeddingModelEditView: View {
278279
TextField("Model Name", text: $store.modelName)
279280

280281
MaxTokensTextField(store: store)
282+
283+
Button("Custom Headers") {
284+
isEditingCustomHeader.toggle()
285+
}
286+
}.sheet(isPresented: $isEditingCustomHeader) {
287+
CustomHeaderSettingsView(headers: $store.customHeaders)
281288
}
282289
}
283290
}

‎Tool/Sources/AIModel/ChatModel.swift‎

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,25 @@ public struct ChatModel: Codable, Equatable, Identifiable {
6464
self.apiVersion = apiVersion
6565
}
6666
}
67+
68+
public struct CustomHeaderInfo: Codable, Equatable {
69+
public struct HeaderField: Codable, Equatable {
70+
public var key: String
71+
public var value: String
72+
73+
public init(key: String, value: String) {
74+
self.key = key
75+
self.value = value
76+
}
77+
}
78+
79+
@FallbackDecoding<EmptyArray>
80+
public var headers: [HeaderField]
81+
82+
public init(headers: [HeaderField] = []) {
83+
self.headers = headers
84+
}
85+
}
6786

6887
@FallbackDecoding<EmptyString>
6988
public var apiKeyName: String
@@ -86,6 +105,8 @@ public struct ChatModel: Codable, Equatable, Identifiable {
86105
public var googleGenerativeAIInfo: GoogleGenerativeAIInfo
87106
@FallbackDecoding<EmptyChatModelOpenAICompatibleInfo>
88107
public var openAICompatibleInfo: OpenAICompatibleInfo
108+
@FallbackDecoding<EmptyChatModelCustomHeaderInfo>
109+
public var customHeaderInfo: CustomHeaderInfo
89110

90111
public init(
91112
apiKeyName: String = "",
@@ -97,7 +118,8 @@ public struct ChatModel: Codable, Equatable, Identifiable {
97118
openAIInfo: OpenAIInfo = OpenAIInfo(),
98119
ollamaInfo: OllamaInfo = OllamaInfo(),
99120
googleGenerativeAIInfo: GoogleGenerativeAIInfo = GoogleGenerativeAIInfo(),
100-
openAICompatibleInfo: OpenAICompatibleInfo = OpenAICompatibleInfo()
121+
openAICompatibleInfo: OpenAICompatibleInfo = OpenAICompatibleInfo(),
122+
customHeaderInfo: CustomHeaderInfo = CustomHeaderInfo()
101123
) {
102124
self.apiKeyName = apiKeyName
103125
self.baseURL = baseURL
@@ -109,6 +131,7 @@ public struct ChatModel: Codable, Equatable, Identifiable {
109131
self.ollamaInfo = ollamaInfo
110132
self.googleGenerativeAIInfo = googleGenerativeAIInfo
111133
self.openAICompatibleInfo = openAICompatibleInfo
134+
self.customHeaderInfo = customHeaderInfo
112135
}
113136
}
114137

@@ -168,3 +191,7 @@ public struct EmptyChatModelGoogleGenerativeAIInfo: FallbackValueProvider {
168191
public struct EmptyChatModelOpenAICompatibleInfo: FallbackValueProvider {
169192
public static var defaultValue: ChatModel.Info.OpenAICompatibleInfo { .init() }
170193
}
194+
195+
public struct EmptyChatModelCustomHeaderInfo: FallbackValueProvider {
196+
public static var defaultValue: ChatModel.Info.CustomHeaderInfo { .init() }
197+
}

‎Tool/Sources/AIModel/EmbeddingModel.swift‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ public struct EmbeddingModel: Codable, Equatable, Identifiable {
2626
public struct Info: Codable, Equatable {
2727
public typealias OllamaInfo = ChatModel.Info.OllamaInfo
2828
public typealias OpenAIInfo = ChatModel.Info.OpenAIInfo
29+
public typealias CustomHeaderInfo = ChatModel.Info.CustomHeaderInfo
2930

3031
@FallbackDecoding<EmptyString>
3132
public var apiKeyName: String
@@ -44,6 +45,8 @@ public struct EmbeddingModel: Codable, Equatable, Identifiable {
4445
public var openAIInfo: OpenAIInfo
4546
@FallbackDecoding<EmptyChatModelOllamaInfo>
4647
public var ollamaInfo: OllamaInfo
48+
@FallbackDecoding<EmptyChatModelCustomHeaderInfo>
49+
public var customHeaderInfo: CustomHeaderInfo
4750

4851
public init(
4952
apiKeyName: String = "",
@@ -53,7 +56,8 @@ public struct EmbeddingModel: Codable, Equatable, Identifiable {
5356
dimensions: Int = 1536,
5457
modelName: String = "",
5558
openAIInfo: OpenAIInfo = OpenAIInfo(),
56-
ollamaInfo: OllamaInfo = OllamaInfo()
59+
ollamaInfo: OllamaInfo = OllamaInfo(),
60+
customHeaderInfo: CustomHeaderInfo = CustomHeaderInfo()
5761
) {
5862
self.apiKeyName = apiKeyName
5963
self.baseURL = baseURL
@@ -63,6 +67,7 @@ public struct EmbeddingModel: Codable, Equatable, Identifiable {
6367
self.modelName = modelName
6468
self.openAIInfo = openAIInfo
6569
self.ollamaInfo = ollamaInfo
70+
self.customHeaderInfo = customHeaderInfo
6671
}
6772
}
6873

‎Tool/Sources/OpenAIService/APIs/OpenAIChatCompletionsService.swift‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -259,6 +259,7 @@ actor OpenAIChatCompletionsService: ChatCompletionsStreamAPI, ChatCompletionsAPI
259259

260260
Self.setupAppInformation(&request)
261261
Self.setupAPIKey(&request, model: model, apiKey: apiKey)
262+
Self.setupExtraHeaderFields(&request, model: model)
262263

263264
let (result, response) = try await URLSession.shared.bytes(for: request)
264265
guard let response = response as? HTTPURLResponse else {
@@ -399,6 +400,12 @@ actor OpenAIChatCompletionsService: ChatCompletionsStreamAPI, ChatCompletionsAPI
399400
}
400401
}
401402
}
403+
404+
static func setupExtraHeaderFields(_ request: inout URLRequest, model: ChatModel) {
405+
for field in model.info.customHeaderInfo.headers where !field.key.isEmpty {
406+
request.setValue(field.value, forHTTPHeaderField: field.key)
407+
}
408+
}
402409
}
403410

404411
extension OpenAIChatCompletionsService.ResponseBody {

‎Tool/Sources/OpenAIService/APIs/OpenAIEmbeddingService.swift‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,7 @@ struct OpenAIEmbeddingService: EmbeddingAPI {
7575

7676
Self.setupAppInformation(&request)
7777
Self.setupAPIKey(&request, model: model, apiKey: apiKey)
78+
Self.setupExtraHeaderFields(&request, model: model)
7879

7980
let (result, response) = try await URLSession.shared.data(for: request)
8081
guard let response = response as? HTTPURLResponse else {
@@ -143,5 +144,11 @@ struct OpenAIEmbeddingService: EmbeddingAPI {
143144
}
144145
}
145146
}
147+
148+
static func setupExtraHeaderFields(_ request: inout URLRequest, model: EmbeddingModel) {
149+
for field in model.info.customHeaderInfo.headers where !field.key.isEmpty {
150+
request.setValue(field.value, forHTTPHeaderField: field.key)
151+
}
152+
}
146153
}
147154

0 commit comments

Comments
 (0)