Skip to content

Commit 1281aa0

Browse files
intitniintitni
committed
Merge branch 'release/0.5.0'
2 parents 2d2d8b7 + 1539bed commit 1281aa0

16 files changed

Lines changed: 398 additions & 47 deletions

File tree

‎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: 49 additions & 5 deletions
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 = ""
@@ -21,6 +27,7 @@ public actor OllamaService {
2127
public enum Endpoint {
2228
case completion
2329
case chatCompletion
30+
case completionWithSuffix
2431
}
2532

2633
init(
@@ -32,13 +39,14 @@ public actor OllamaService {
3239
temperature: Double = 0.2,
3340
stopWords: [String] = [],
3441
keepAlive: String = "",
35-
format: ResponseFormat = .none
42+
format: ResponseFormat = .none,
43+
authenticationMode: AuthenticationMode? = nil
3644
) {
3745
self.url = url.flatMap(URL.init(string:)) ?? {
3846
switch endpoint {
3947
case .chatCompletion:
4048
URL(string: "https://127.0.0.1:11434/api/chat")!
41-
case .completion:
49+
case .completion, .completionWithSuffix:
4250
URL(string: "https://127.0.0.1:11434/api/generate")!
4351
}
4452
}()
@@ -51,6 +59,7 @@ public actor OllamaService {
5159
self.keepAlive = keepAlive
5260
self.format = format
5361
self.contextWindow = contextWindow
62+
self.authenticationMode = authenticationMode
5463
}
5564
}
5665

@@ -74,7 +83,24 @@ extension OllamaService: CodeCompletionServiceType {
7483
case .completion:
7584
let prompt = createPrompt(from: request)
7685
CodeCompletionLogger.logger.logPrompt([(prompt, "user")])
77-
let stream = try await sendPrompt(prompt)
86+
let stream = try await sendPrompt(prompt, raw: request.promptIsRaw)
87+
return stream.compactMap { $0.response }
88+
case .completionWithSuffix:
89+
let strategy = DefaultTruncateStrategy(maxTokenLimit: max(
90+
contextWindow / 3 * 2,
91+
contextWindow - maxToken - 20
92+
))
93+
let prompts = strategy.createTruncatedPrompt(promptStrategy: request)
94+
95+
let prefix = prompts.first { $0.role == .prefix }?.content ?? ""
96+
let suffix = prompts.last { $0.role == .suffix }?.content ?? ""
97+
98+
CodeCompletionLogger.logger.logPrompt([
99+
(prefix, "prefix"),
100+
(suffix, "suffix"),
101+
])
102+
103+
let stream = try await sendPrompt(prefix, suffix: suffix)
78104
return stream.compactMap { $0.response }
79105
}
80106
}
@@ -215,6 +241,8 @@ extension OllamaService {
215241
var options: ChatCompletionRequestBody.Options
216242
var keep_alive: String?
217243
var format: String?
244+
var raw: Bool?
245+
var suffix: String?
218246
}
219247

220248
func createPrompt(from request: PromptStrategy) -> String {
@@ -227,7 +255,11 @@ extension OllamaService {
227255
.trimmingCharacters(in: .whitespacesAndNewlines)
228256
}
229257

230-
func sendPrompt(_ prompt: String) async throws -> ResponseStream<ChatCompletionResponseChunk> {
258+
func sendPrompt(
259+
_ prompt: String,
260+
raw: Bool? = nil,
261+
suffix: String? = nil
262+
) async throws -> ResponseStream<ChatCompletionResponseChunk> {
231263
let requestBody = CompletionRequestBody(
232264
model: modelName,
233265
prompt: prompt,
@@ -238,14 +270,26 @@ extension OllamaService {
238270
num_predict: maxToken
239271
),
240272
keep_alive: keepAlive.isEmpty ? nil : keepAlive,
241-
format: format == .none ? nil : format.rawValue
273+
format: format == .none ? nil : format.rawValue,
274+
raw: raw,
275+
suffix: suffix
242276
)
243277

244278
var request = URLRequest(url: url)
245279
request.httpMethod = "POST"
246280
let encoder = JSONEncoder()
247281
request.httpBody = try encoder.encode(requestBody)
248282
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+
249293
let (result, response) = try await URLSession.shared.bytes(for: request)
250294

251295
guard let response = response as? HTTPURLResponse else {

‎Core/Sources/CodeCompletionService/CodeCompletionService.swift‎

Lines changed: 51 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,7 @@ public struct CodeCompletionService {
9999
endpoint: .chatCompletion,
100100
modelName: model.info.modelName,
101101
contextWindow: model.info.maxTokens,
102-
maxToken: UserDefaults.shared.value(for: \.maxGenerationToken),
102+
maxToken: UserDefaults.shared.value(for: \.maxGenerationToken),
103103
stopWords: prompt.stopWords,
104104
apiKey: apiKey
105105
)
@@ -114,7 +114,7 @@ public struct CodeCompletionService {
114114
let service = AzureOpenAIService(
115115
url: model.endpoint,
116116
endpoint: .chatCompletion,
117-
modelName: model.info.modelName,
117+
modelName: model.info.modelName,
118118
contextWindow: model.info.maxTokens,
119119
maxToken: UserDefaults.shared.value(for: \.maxGenerationToken),
120120
stopWords: prompt.stopWords,
@@ -179,7 +179,7 @@ public struct CodeCompletionService {
179179
endpoint: .completion,
180180
modelName: model.info.modelName,
181181
contextWindow: model.info.maxTokens,
182-
maxToken: UserDefaults.shared.value(for: \.maxGenerationToken),
182+
maxToken: UserDefaults.shared.value(for: \.maxGenerationToken),
183183
stopWords: prompt.stopWords,
184184
apiKey: apiKey
185185
)
@@ -211,7 +211,7 @@ public struct CodeCompletionService {
211211
let service = OllamaService(
212212
url: model.endpoint,
213213
endpoint: .completion,
214-
modelName: model.info.modelName,
214+
modelName: model.info.modelName,
215215
contextWindow: model.info.maxTokens,
216216
maxToken: UserDefaults.shared.value(for: \.maxGenerationToken),
217217
stopWords: prompt.stopWords,
@@ -256,6 +256,53 @@ public struct CodeCompletionService {
256256
)
257257
try Task.checkCancellation()
258258
return result
259+
case .ollama:
260+
let service = OllamaService(
261+
url: model.endpoint,
262+
endpoint: .completionWithSuffix,
263+
modelName: model.info.modelName,
264+
contextWindow: model.info.maxTokens,
265+
maxToken: UserDefaults.shared.value(for: \.maxGenerationToken),
266+
stopWords: prompt.stopWords,
267+
keepAlive: model.info.ollamaInfo.keepAlive,
268+
format: .none
269+
)
270+
let result = try await service.getCompletions(
271+
prompt,
272+
streamStopStrategy: streamStopStrategy,
273+
count: count
274+
)
275+
try Task.checkCancellation()
276+
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
259306
case .unknown:
260307
throw Error.unknownFormat
261308
}

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

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@ public struct FIMModel: Codable, Equatable, Identifiable {
1919

2020
public enum Format: String, Codable, Equatable, CaseIterable {
2121
case mistral
22+
case ollama
23+
case ollamaCompatible
2224

2325
case unknown
2426
}
@@ -34,19 +36,37 @@ public struct FIMModel: Codable, Equatable, Identifiable {
3436
public var maxTokens: Int
3537
@FallbackDecoding<EmptyString>
3638
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+
}
48+
49+
@FallbackDecoding<EmptyChatModelOllamaInfo>
50+
public var ollamaInfo: ChatModel.Info.OllamaInfo
3751

3852
public init(
3953
apiKeyName: String = "",
4054
baseURL: String = "",
4155
isFullURL: Bool = false,
4256
maxTokens: Int = 4000,
43-
modelName: String = ""
57+
modelName: String = "",
58+
authenticationMode: AuthenticationMode = .bearerToken,
59+
authenticationHeaderFieldName: String = "",
60+
ollamaInfo: ChatModel.Info.OllamaInfo = ChatModel.Info.OllamaInfo()
4461
) {
4562
self.apiKeyName = apiKeyName
4663
self.baseURL = baseURL
4764
self.isFullURL = isFullURL
4865
self.maxTokens = maxTokens
4966
self.modelName = modelName
67+
self.ollamaInfo = ollamaInfo
68+
self.authenticationMode = authenticationMode
69+
self.authenticationHeaderFieldName = authenticationHeaderFieldName
5070
}
5171
}
5272

@@ -57,6 +77,15 @@ public struct FIMModel: Codable, Equatable, Identifiable {
5777
if baseURL.isEmpty { return "https://api.mistral.ai/v1/fim/completions" }
5878
if info.isFullURL { return baseURL }
5979
return "\(baseURL)/v1/fim/completions"
80+
case .ollama:
81+
let baseURL = info.baseURL
82+
if baseURL.isEmpty { return "http://localhost:11434/api/generate" }
83+
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"
6089
case .unknown:
6190
return ""
6291
}
@@ -71,3 +100,6 @@ public struct EmptyFIMModelFormat: FallbackValueProvider {
71100
public static var defaultValue: FIMModel.Format { .unknown }
72101
}
73102

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

‎Core/Sources/Fundamental/PromptStrategy.swift‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@ public protocol PromptStrategy {
1717
var stopWords: [String] { get }
1818
/// The language of the source code.
1919
var language: CodeLanguage? { get }
20+
/// If the prompt generated is raw.
21+
var promptIsRaw: Bool { get }
2022

2123
/// Creates a prompt about the source code and relevant code snippets to be sent to the AI
2224
/// model.
@@ -92,5 +94,7 @@ public extension PromptStrategy {
9294
guard let prefix = prefix.last else { return .empty }
9395
return .unchanged(prefix)
9496
}
97+
98+
var promptIsRaw: Bool { false }
9599
}
96100

‎Core/Sources/Storage/Preferences.swift‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,13 @@ public extension UserDefaultPreferenceKeys {
122122
)
123123
}
124124

125+
var fimPromptIsRaw: PreferenceKey<Bool> {
126+
.init(
127+
defaultValue: false,
128+
key: "CustomSuggestionService-FimPromptIsRaw"
129+
)
130+
}
131+
125132
var maxGenerationToken: PreferenceKey<Int> {
126133
.init(
127134
defaultValue: 200,

‎Core/Sources/SuggestionService/RequestStrategies/FillInTheMiddleRequestStrategy.swift‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ struct FillInTheMiddleRequestStrategy: RequestStrategy {
4242
var relevantCodeSnippets: [RelevantCodeSnippet] { sourceRequest.relevantCodeSnippets }
4343
var stopWords: [String] { ["\n\n", Tag.stop].filter { !$0.isEmpty } }
4444
var language: CodeLanguage? { sourceRequest.language }
45+
var promptIsRaw: Bool { UserDefaults.shared.value(for: \.fimPromptIsRaw) }
4546

4647
var suggestionPrefix: SuggestionPrefix {
4748
guard let prefix = prefix.last else { return .empty }

‎CustomSuggestionService.xcodeproj/project.pbxproj‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
/* Begin PBXBuildFile section */
1010
C80AC8F32C274ADD00669BDE /* FIMModelEdit.swift in Sources */ = {isa = PBXBuildFile; fileRef = C80AC8F12C274ADD00669BDE /* FIMModelEdit.swift */; };
1111
C80AC8F42C274ADD00669BDE /* FIMModelEditView.swift in Sources */ = {isa = PBXBuildFile; fileRef = C80AC8F22C274ADD00669BDE /* FIMModelEditView.swift */; };
12+
C813647A2CA29924000E2237 /* HandleToast.swift in Sources */ = {isa = PBXBuildFile; fileRef = C81364792CA29924000E2237 /* HandleToast.swift */; };
1213
C83B83AF2B7DD261007B4442 /* Dependency.swift in Sources */ = {isa = PBXBuildFile; fileRef = C83B83AE2B7DD261007B4442 /* Dependency.swift */; };
1314
C84697512B7B7B3700B8B840 /* TestFieldView.swift in Sources */ = {isa = PBXBuildFile; fileRef = C84697502B7B7B3700B8B840 /* TestFieldView.swift */; };
1415
C84697542B7B8B8300B8B840 /* STTextView in Frameworks */ = {isa = PBXBuildFile; productRef = C84697532B7B8B8300B8B840 /* STTextView */; };
@@ -71,6 +72,7 @@
7172
/* Begin PBXFileReference section */
7273
C80AC8F12C274ADD00669BDE /* FIMModelEdit.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = FIMModelEdit.swift; sourceTree = "<group>"; };
7374
C80AC8F22C274ADD00669BDE /* FIMModelEditView.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = FIMModelEditView.swift; sourceTree = "<group>"; };
75+
C81364792CA29924000E2237 /* HandleToast.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = HandleToast.swift; sourceTree = "<group>"; };
7476
C81547D92B8737DF002203B3 /* LICENSE */ = {isa = PBXFileReference; lastKnownFileType = text; path = LICENSE; sourceTree = "<group>"; };
7577
C81547DB2B873DC1002203B3 /* appcast.xml */ = {isa = PBXFileReference; lastKnownFileType = text.xml; path = appcast.xml; sourceTree = "<group>"; };
7678
C83B83AE2B7DD261007B4442 /* Dependency.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Dependency.swift; sourceTree = "<group>"; };
@@ -259,6 +261,7 @@
259261
isa = PBXGroup;
260262
children = (
261263
C8EA62D42B6BE4BE00E081FC /* Toast.swift */,
264+
C81364792CA29924000E2237 /* HandleToast.swift */,
262265
);
263266
path = Toast;
264267
sourceTree = "<group>";
@@ -397,6 +400,7 @@
397400
C8EA62D62B6BE5FD00E081FC /* ChatModelEdit.swift in Sources */,
398401
C83B83AF2B7DD261007B4442 /* Dependency.swift in Sources */,
399402
C8D0114C2B59912700219412 /* CustomSuggestionServiceApp.swift in Sources */,
403+
C813647A2CA29924000E2237 /* HandleToast.swift in Sources */,
400404
C8774F482B83A009008FF699 /* CompletionModelEdit.swift in Sources */,
401405
C89BA4BC2B861DE4008801C5 /* UpdaterChecker.swift in Sources */,
402406
C8EA62DC2B6BE60500E081FC /* APIKeySubmission.swift in Sources */,

0 commit comments

Comments
 (0)