Skip to content

Commit ecace5f

Browse files
intitniintitni
committed
Merge branch 'feature/more-fim-api' into develop
2 parents 4172be3 + e1c0202 commit ecace5f

File tree

6 files changed

+153
-8
lines changed

6 files changed

+153
-8
lines changed

‎Core/Sources/CodeCompletionService/API/MistralFIMService.swift‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ extension MistralFIMService: CodeCompletionServiceType {
3838

3939
func getCompletion(_ request: any PromptStrategy) async throws -> CompletionSequence {
4040
let result = try await send(request)
41-
return result.compactMap { $0.choices?.first?.delta?.content }
41+
return result.compactMap { $0.choices?.first?.delta?.content ?? $0.choices?.first?.text }
4242
}
4343
}
4444

@@ -75,6 +75,7 @@ extension MistralFIMService {
7575
struct Choice: Decodable {
7676
var index: Int
7777
var delta: Delta?
78+
var text: String?
7879
var finish_reason: String?
7980
}
8081

‎Core/Sources/CodeCompletionService/API/OllamaService.swift‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,12 @@ public actor OllamaService {
1212
let stopWords: [String]
1313
let keepAlive: String
1414
let format: ResponseFormat
15+
let authenticationMode: AuthenticationMode?
16+
17+
enum AuthenticationMode {
18+
case bearerToken(String)
19+
case header(name: String, value: String)
20+
}
1521

1622
public enum ResponseFormat: String {
1723
case none = ""
@@ -33,7 +39,8 @@ public actor OllamaService {
3339
temperature: Double = 0.2,
3440
stopWords: [String] = [],
3541
keepAlive: String = "",
36-
format: ResponseFormat = .none
42+
format: ResponseFormat = .none,
43+
authenticationMode: AuthenticationMode? = nil
3744
) {
3845
self.url = url.flatMap(URL.init(string:)) ?? {
3946
switch endpoint {
@@ -52,6 +59,7 @@ public actor OllamaService {
5259
self.keepAlive = keepAlive
5360
self.format = format
5461
self.contextWindow = contextWindow
62+
self.authenticationMode = authenticationMode
5563
}
5664
}
5765

@@ -272,6 +280,16 @@ extension OllamaService {
272280
let encoder = JSONEncoder()
273281
request.httpBody = try encoder.encode(requestBody)
274282
request.setValue("application/json", forHTTPHeaderField: "Content-Type")
283+
284+
switch authenticationMode{
285+
case .none:
286+
break
287+
case let .bearerToken(key):
288+
request.setValue("Bearer \(key)", forHTTPHeaderField: "Authorization")
289+
case let .header(name, value):
290+
request.setValue(value, forHTTPHeaderField: name)
291+
}
292+
275293
let (result, response) = try await URLSession.shared.bytes(for: request)
276294

277295
guard let response = response as? HTTPURLResponse else {

‎Core/Sources/CodeCompletionService/CodeCompletionService.swift‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -274,6 +274,35 @@ public struct CodeCompletionService {
274274
)
275275
try Task.checkCancellation()
276276
return result
277+
case .ollamaCompatible:
278+
let service = OllamaService(
279+
url: model.endpoint,
280+
endpoint: .completionWithSuffix,
281+
modelName: model.info.modelName,
282+
contextWindow: model.info.maxTokens,
283+
maxToken: UserDefaults.shared.value(for: \.maxGenerationToken),
284+
stopWords: prompt.stopWords,
285+
keepAlive: model.info.ollamaInfo.keepAlive,
286+
format: .none,
287+
authenticationMode: {
288+
switch model.info.authenticationMode {
289+
case .header:
290+
return .header(
291+
name: model.info.authenticationHeaderFieldName,
292+
value: apiKey
293+
)
294+
case .bearerToken:
295+
return .bearerToken(apiKey)
296+
}
297+
}()
298+
)
299+
let result = try await service.getCompletions(
300+
prompt,
301+
streamStopStrategy: streamStopStrategy,
302+
count: count
303+
)
304+
try Task.checkCancellation()
305+
return result
277306
case .unknown:
278307
throw Error.unknownFormat
279308
}

‎Core/Sources/Fundamental/Models/FIMModel.swift‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ public struct FIMModel: Codable, Equatable, Identifiable {
2020
public enum Format: String, Codable, Equatable, CaseIterable {
2121
case mistral
2222
case ollama
23+
case ollamaCompatible
2324

2425
case unknown
2526
}
@@ -35,6 +36,15 @@ public struct FIMModel: Codable, Equatable, Identifiable {
3536
public var maxTokens: Int
3637
@FallbackDecoding<EmptyString>
3738
public var modelName: String
39+
@FallbackDecoding<EmptyFIMModelAuthenticationMode>
40+
public var authenticationMode: AuthenticationMode
41+
@FallbackDecoding<EmptyString>
42+
public var authenticationHeaderFieldName: String
43+
44+
public enum AuthenticationMode: Codable, Equatable, CaseIterable {
45+
case header
46+
case bearerToken
47+
}
3848

3949
@FallbackDecoding<EmptyChatModelOllamaInfo>
4050
public var ollamaInfo: ChatModel.Info.OllamaInfo
@@ -45,6 +55,8 @@ public struct FIMModel: Codable, Equatable, Identifiable {
4555
isFullURL: Bool = false,
4656
maxTokens: Int = 4000,
4757
modelName: String = "",
58+
authenticationMode: AuthenticationMode = .bearerToken,
59+
authenticationHeaderFieldName: String = "",
4860
ollamaInfo: ChatModel.Info.OllamaInfo = ChatModel.Info.OllamaInfo()
4961
) {
5062
self.apiKeyName = apiKeyName
@@ -53,6 +65,8 @@ public struct FIMModel: Codable, Equatable, Identifiable {
5365
self.maxTokens = maxTokens
5466
self.modelName = modelName
5567
self.ollamaInfo = ollamaInfo
68+
self.authenticationMode = authenticationMode
69+
self.authenticationHeaderFieldName = authenticationHeaderFieldName
5670
}
5771
}
5872

@@ -67,6 +81,11 @@ public struct FIMModel: Codable, Equatable, Identifiable {
6781
let baseURL = info.baseURL
6882
if baseURL.isEmpty { return "http://localhost:11434/api/generate" }
6983
return "\(baseURL)/api/generate"
84+
case .ollamaCompatible:
85+
let baseURL = info.baseURL
86+
if baseURL.isEmpty { return "http://localhost:11434/api/generate" }
87+
if info.isFullURL { return baseURL }
88+
return "\(baseURL)/api/generate"
7089
case .unknown:
7190
return ""
7291
}
@@ -81,3 +100,6 @@ public struct EmptyFIMModelFormat: FallbackValueProvider {
81100
public static var defaultValue: FIMModel.Format { .unknown }
82101
}
83102

103+
public struct EmptyFIMModelAuthenticationMode: FallbackValueProvider {
104+
public static var defaultValue: FIMModel.Info.AuthenticationMode { .bearerToken }
105+
}

‎CustomSuggestionService/ChatModelManagement/FIMModelEdit.swift‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@ struct FIMModelEdit {
2121
var apiKeySelection: APIKeySelection.State = .init()
2222
var baseURLSelection: BaseURLSelection.State = .init()
2323
var ollamaKeepAlive: String = ""
24+
var authenticationMode: FIMModel.Info.AuthenticationMode = .bearerToken
25+
var authenticationHeaderFieldName: String = ""
2426
}
2527

2628
enum Action: Equatable, BindableAction {
@@ -142,7 +144,9 @@ extension FIMModel {
142144
apiKeyManagement: .init(availableAPIKeyNames: [info.apiKeyName])
143145
),
144146
baseURLSelection: .init(baseURL: info.baseURL, isFullURL: info.isFullURL),
145-
ollamaKeepAlive: info.ollamaInfo.keepAlive
147+
ollamaKeepAlive: info.ollamaInfo.keepAlive,
148+
authenticationMode: info.authenticationMode,
149+
authenticationHeaderFieldName: info.authenticationHeaderFieldName
146150
)
147151
}
148152

@@ -153,10 +157,12 @@ extension FIMModel {
153157
format: state.format,
154158
info: .init(
155159
apiKeyName: state.apiKeyName,
156-
baseURL: state.baseURL.trimmingCharacters(in: .whitespacesAndNewlines),
160+
baseURL: state.baseURL.trimmingCharacters(in: .whitespacesAndNewlines),
157161
isFullURL: state.baseURLSelection.isFullURL,
158162
maxTokens: state.maxTokens,
159163
modelName: state.modelName.trimmingCharacters(in: .whitespacesAndNewlines),
164+
authenticationMode: state.authenticationMode,
165+
authenticationHeaderFieldName: state.authenticationHeaderFieldName,
160166
ollamaInfo: .init(keepAlive: state.ollamaKeepAlive)
161167
)
162168
)

‎CustomSuggestionService/ChatModelManagement/FIMModelEditView.swift‎

Lines changed: 73 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@ struct FIMModelEditView: View {
2121
mistralForm
2222
case .ollama:
2323
ollama
24+
case .ollamaCompatible:
25+
ollamaCompatible
2426
case .unknown:
2527
EmptyView()
2628
}
@@ -69,6 +71,8 @@ struct FIMModelEditView: View {
6971
Text("Mistral").tag(format)
7072
case .ollama:
7173
Text("Ollama").tag(format)
74+
case .ollamaCompatible:
75+
Text("Ollama Compatible").tag(format)
7276
case .unknown:
7377
EmptyView()
7478
}
@@ -194,32 +198,97 @@ struct FIMModelEditView: View {
194198
}
195199

196200
maxTokensTextField
201+
202+
VStack(alignment: .leading, spacing: 8) {
203+
Text(Image(systemName: "exclamationmark.triangle.fill")) + Text(
204+
"""
205+
Compatible with:
206+
- Mistral FIM API
207+
- DeepSeek FIM API
208+
209+
or any API that takes a prompt and a suffix and responds with
210+
an OpenAI-compatible streaming response, using bearer token authentication.
211+
"""
212+
)
213+
}
214+
.padding(.vertical)
197215
}
198-
216+
199217
@ViewBuilder
200218
var ollama: some View {
201219
baseURLTextField(
202220
title: "",
203-
prompt: Text("https://127.0.0.1:11434/api/generate")
221+
prompt: Text("https://127.0.0.1:11434")
204222
) {
205223
Text("/api/generate")
206224
}
207225

208226
TextField("Model Name", text: $store.modelName)
209227

210228
maxTokensTextField
211-
229+
212230
TextField(text: $store.ollamaKeepAlive, prompt: Text("Default Value")) {
213231
Text("Keep Alive")
214232
}
215-
233+
216234
VStack(alignment: .leading, spacing: 8) {
217235
Text(Image(systemName: "exclamationmark.triangle.fill")) + Text(
218236
" For more details, please visit [https://ollama.com](https://ollama.com)"
219237
)
220238
}
221239
.padding(.vertical)
222240
}
241+
242+
@ViewBuilder
243+
var ollamaCompatible: some View {
244+
Picker(
245+
selection: $store.baseURLSelection.isFullURL,
246+
content: {
247+
Text("Base URL").tag(false)
248+
Text("Full URL").tag(true)
249+
},
250+
label: { Text("URL") }
251+
)
252+
.pickerStyle(.segmented)
253+
254+
baseURLTextField(
255+
title: "",
256+
prompt: store.baseURLSelection.isFullURL
257+
? Text("https://127.0.0.1:11434/api/generate")
258+
: Text("https://127.0.0.1:11434")
259+
) {
260+
if !store.baseURLSelection.isFullURL {
261+
Text("/api/generate")
262+
}
263+
}
264+
265+
Picker(
266+
selection: $store.authenticationMode,
267+
content: {
268+
Text("Bearer Token").tag(FIMModel.Info.AuthenticationMode.bearerToken)
269+
Text("Header Field").tag(FIMModel.Info.AuthenticationMode.header)
270+
},
271+
label: { Text("Authentication Mode") }
272+
)
273+
.pickerStyle(.segmented)
274+
275+
if store.authenticationMode == .header {
276+
TextField("Header Field Name", text: $store.authenticationHeaderFieldName)
277+
}
278+
279+
apiKeyNamePicker
280+
281+
TextField("Model Name", text: $store.modelName)
282+
283+
maxTokensTextField
284+
285+
VStack(alignment: .leading, spacing: 8) {
286+
Text(Image(systemName: "exclamationmark.triangle.fill")) + Text(
287+
" Use with any API that has the same format as Ollama. For more details, please visit [https://ollama.com](https://ollama.com)"
288+
)
289+
}
290+
.padding(.vertical)
291+
}
223292
}
224293

225294
#Preview("Mistral") {

0 commit comments

Comments
 (0)