Skip to content

Commit 4172be3

Browse files
intitniintitni
committed
Merge branch 'feature/ollama-suffix' into develop
2 parents 3edd92d + 0164426 commit 4172be3

12 files changed

Lines changed: 227 additions & 21 deletions

File tree

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

Lines changed: 30 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ public actor OllamaService {
2121
public enum Endpoint {
2222
case completion
2323
case chatCompletion
24+
case completionWithSuffix
2425
}
2526

2627
init(
@@ -38,7 +39,7 @@ public actor OllamaService {
3839
switch endpoint {
3940
case .chatCompletion:
4041
URL(string: "https://127.0.0.1:11434/api/chat")!
41-
case .completion:
42+
case .completion, .completionWithSuffix:
4243
URL(string: "https://127.0.0.1:11434/api/generate")!
4344
}
4445
}()
@@ -74,7 +75,24 @@ extension OllamaService: CodeCompletionServiceType {
7475
case .completion:
7576
let prompt = createPrompt(from: request)
7677
CodeCompletionLogger.logger.logPrompt([(prompt, "user")])
77-
let stream = try await sendPrompt(prompt)
78+
let stream = try await sendPrompt(prompt, raw: request.promptIsRaw)
79+
return stream.compactMap { $0.response }
80+
case .completionWithSuffix:
81+
let strategy = DefaultTruncateStrategy(maxTokenLimit: max(
82+
contextWindow / 3 * 2,
83+
contextWindow - maxToken - 20
84+
))
85+
let prompts = strategy.createTruncatedPrompt(promptStrategy: request)
86+
87+
let prefix = prompts.first { $0.role == .prefix }?.content ?? ""
88+
let suffix = prompts.last { $0.role == .suffix }?.content ?? ""
89+
90+
CodeCompletionLogger.logger.logPrompt([
91+
(prefix, "prefix"),
92+
(suffix, "suffix"),
93+
])
94+
95+
let stream = try await sendPrompt(prefix, suffix: suffix)
7896
return stream.compactMap { $0.response }
7997
}
8098
}
@@ -215,6 +233,8 @@ extension OllamaService {
215233
var options: ChatCompletionRequestBody.Options
216234
var keep_alive: String?
217235
var format: String?
236+
var raw: Bool?
237+
var suffix: String?
218238
}
219239

220240
func createPrompt(from request: PromptStrategy) -> String {
@@ -227,7 +247,11 @@ extension OllamaService {
227247
.trimmingCharacters(in: .whitespacesAndNewlines)
228248
}
229249

230-
func sendPrompt(_ prompt: String) async throws -> ResponseStream<ChatCompletionResponseChunk> {
250+
func sendPrompt(
251+
_ prompt: String,
252+
raw: Bool? = nil,
253+
suffix: String? = nil
254+
) async throws -> ResponseStream<ChatCompletionResponseChunk> {
231255
let requestBody = CompletionRequestBody(
232256
model: modelName,
233257
prompt: prompt,
@@ -238,7 +262,9 @@ extension OllamaService {
238262
num_predict: maxToken
239263
),
240264
keep_alive: keepAlive.isEmpty ? nil : keepAlive,
241-
format: format == .none ? nil : format.rawValue
265+
format: format == .none ? nil : format.rawValue,
266+
raw: raw,
267+
suffix: suffix
242268
)
243269

244270
var request = URLRequest(url: url)

‎Core/Sources/CodeCompletionService/CodeCompletionService.swift‎

Lines changed: 22 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,24 @@ 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
259277
case .unknown:
260278
throw Error.unknownFormat
261279
}

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

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

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

2324
case unknown
2425
}
@@ -34,19 +35,24 @@ public struct FIMModel: Codable, Equatable, Identifiable {
3435
public var maxTokens: Int
3536
@FallbackDecoding<EmptyString>
3637
public var modelName: String
38+
39+
@FallbackDecoding<EmptyChatModelOllamaInfo>
40+
public var ollamaInfo: ChatModel.Info.OllamaInfo
3741

3842
public init(
3943
apiKeyName: String = "",
4044
baseURL: String = "",
4145
isFullURL: Bool = false,
4246
maxTokens: Int = 4000,
43-
modelName: String = ""
47+
modelName: String = "",
48+
ollamaInfo: ChatModel.Info.OllamaInfo = ChatModel.Info.OllamaInfo()
4449
) {
4550
self.apiKeyName = apiKeyName
4651
self.baseURL = baseURL
4752
self.isFullURL = isFullURL
4853
self.maxTokens = maxTokens
4954
self.modelName = modelName
55+
self.ollamaInfo = ollamaInfo
5056
}
5157
}
5258

@@ -57,6 +63,10 @@ public struct FIMModel: Codable, Equatable, Identifiable {
5763
if baseURL.isEmpty { return "https://api.mistral.ai/v1/fim/completions" }
5864
if info.isFullURL { return baseURL }
5965
return "\(baseURL)/v1/fim/completions"
66+
case .ollama:
67+
let baseURL = info.baseURL
68+
if baseURL.isEmpty { return "http://localhost:11434/api/generate" }
69+
return "\(baseURL)/api/generate"
6070
case .unknown:
6171
return ""
6272
}

‎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 */,

‎CustomSuggestionService/ChatModelManagement/FIMModelEdit.swift‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ struct FIMModelEdit {
2020
var suggestedMaxTokens: Int?
2121
var apiKeySelection: APIKeySelection.State = .init()
2222
var baseURLSelection: BaseURLSelection.State = .init()
23+
var ollamaKeepAlive: String = ""
2324
}
2425

2526
enum Action: Equatable, BindableAction {
@@ -140,7 +141,8 @@ extension FIMModel {
140141
apiKeyName: info.apiKeyName,
141142
apiKeyManagement: .init(availableAPIKeyNames: [info.apiKeyName])
142143
),
143-
baseURLSelection: .init(baseURL: info.baseURL, isFullURL: info.isFullURL)
144+
baseURLSelection: .init(baseURL: info.baseURL, isFullURL: info.isFullURL),
145+
ollamaKeepAlive: info.ollamaInfo.keepAlive
144146
)
145147
}
146148

@@ -154,7 +156,8 @@ extension FIMModel {
154156
baseURL: state.baseURL.trimmingCharacters(in: .whitespacesAndNewlines),
155157
isFullURL: state.baseURLSelection.isFullURL,
156158
maxTokens: state.maxTokens,
157-
modelName: state.modelName.trimmingCharacters(in: .whitespacesAndNewlines)
159+
modelName: state.modelName.trimmingCharacters(in: .whitespacesAndNewlines),
160+
ollamaInfo: .init(keepAlive: state.ollamaKeepAlive)
158161
)
159162
)
160163
}

‎CustomSuggestionService/ChatModelManagement/FIMModelEditView.swift‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@ struct FIMModelEditView: View {
1919
switch store.format {
2020
case .mistral:
2121
mistralForm
22+
case .ollama:
23+
ollama
2224
case .unknown:
2325
EmptyView()
2426
}
@@ -65,6 +67,8 @@ struct FIMModelEditView: View {
6567
switch format {
6668
case .mistral:
6769
Text("Mistral").tag(format)
70+
case .ollama:
71+
Text("Ollama").tag(format)
6872
case .unknown:
6973
EmptyView()
7074
}
@@ -191,6 +195,31 @@ struct FIMModelEditView: View {
191195

192196
maxTokensTextField
193197
}
198+
199+
@ViewBuilder
200+
var ollama: some View {
201+
baseURLTextField(
202+
title: "",
203+
prompt: Text("https://127.0.0.1:11434/api/generate")
204+
) {
205+
Text("/api/generate")
206+
}
207+
208+
TextField("Model Name", text: $store.modelName)
209+
210+
maxTokensTextField
211+
212+
TextField(text: $store.ollamaKeepAlive, prompt: Text("Default Value")) {
213+
Text("Keep Alive")
214+
}
215+
216+
VStack(alignment: .leading, spacing: 8) {
217+
Text(Image(systemName: "exclamationmark.triangle.fill")) + Text(
218+
" For more details, please visit [https://ollama.com](https://ollama.com)"
219+
)
220+
}
221+
.padding(.vertical)
222+
}
194223
}
195224

196225
#Preview("Mistral") {

‎CustomSuggestionService/ContentView.swift‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ struct ContentView: View {
5353
) {
5454
Text("Suggestion Line Limit (0 for unlimited)")
5555
}
56-
56+
5757
NumberInput(
5858
value: settings.$maxGenerationToken,
5959
range: 0...Int.max,
@@ -125,6 +125,7 @@ struct ContentView: View {
125125
}
126126
}
127127
}
128+
.handleToast()
128129
}
129130
}
130131

@@ -180,6 +181,7 @@ struct RequestStrategyPicker: View {
180181
final class Settings: ObservableObject {
181182
@AppStorage(\.requestStrategyId) var requestStrategyId
182183
@AppStorage(\.fimTemplate) var fimTemplate
184+
@AppStorage(\.fimPromptIsRaw) var fimPromptIsRaw
183185
@AppStorage(\.fimStopToken) var fimStopToken
184186
}
185187

@@ -232,6 +234,7 @@ struct RequestStrategyPicker: View {
232234
text: $settings.fimTemplate,
233235
prompt: Text(UserDefaults.shared.defaultValue(for: \.fimTemplate))
234236
) { Text("FIM Template") }
237+
Toggle(isOn: $settings.fimPromptIsRaw) { Text("Raw Prompt") }
235238
TextField(text: $settings.fimStopToken) { Text("FIM Stop Token") }
236239
}
237240
}

0 commit comments

Comments
 (0)