From ede4d400000e8a70011ffdddc5f899ee9063d5a3 Mon Sep 17 00:00:00 2001 From: Shx Guo Date: Thu, 29 Feb 2024 15:31:11 +0800 Subject: [PATCH 01/27] Update CodeCompletionServiceType.getCompletion to return an AsyncStream --- .../AzureOpenAIService.swift | 26 ++++++++++++++++--- .../CodeCompletionService.swift | 9 +++++-- .../GoogleGeminiService.swift | 14 ++++++++-- .../CodeCompletionService/OpenAIService.swift | 26 ++++++++++++++++--- .../CodeCompletionService/TabbyService.swift | 14 ++++++++-- 5 files changed, 77 insertions(+), 12 deletions(-) diff --git a/Core/Sources/CodeCompletionService/AzureOpenAIService.swift b/Core/Sources/CodeCompletionService/AzureOpenAIService.swift index 88860df..516a459 100644 --- a/Core/Sources/CodeCompletionService/AzureOpenAIService.swift +++ b/Core/Sources/CodeCompletionService/AzureOpenAIService.swift @@ -39,18 +39,38 @@ public actor AzureOpenAIService { } extension AzureOpenAIService: CodeCompletionServiceType { - func getCompletion(_ request: PromptStrategy) async throws -> String { + func getCompletion(_ request: PromptStrategy) async throws -> AsyncStream { switch endpoint { case .chatCompletion: let messages = createMessages(from: request) CodeCompletionLogger.logger.logPrompt(messages.map { ($0.content, $0.role.rawValue) }) - return try await sendMessages(messages) + return AsyncStream { continuation in + let task = Task { + let result = try await sendMessages(messages) + try Task.checkCancellation() + continuation.yield(result) + continuation.finish() + } + continuation.onTermination = { _ in + task.cancel() + } + } case .completion: let prompt = createPrompt(from: request) CodeCompletionLogger.logger.logPrompt([(prompt, "user")]) - return try await sendPrompt(prompt) + return AsyncStream { continuation in + let task = Task { + let result = try await sendPrompt(prompt) + try Task.checkCancellation() + continuation.yield(result) + continuation.finish() + } + continuation.onTermination = { _ in + task.cancel() + } + } } } } diff --git a/Core/Sources/CodeCompletionService/CodeCompletionService.swift b/Core/Sources/CodeCompletionService/CodeCompletionService.swift index ea8a2de..96f3d4f 100644 --- a/Core/Sources/CodeCompletionService/CodeCompletionService.swift +++ b/Core/Sources/CodeCompletionService/CodeCompletionService.swift @@ -5,7 +5,7 @@ import Storage protocol CodeCompletionServiceType { func getCompletion( _ request: PromptStrategy - ) async throws -> String + ) async throws -> AsyncStream } extension CodeCompletionServiceType { @@ -16,7 +16,12 @@ extension CodeCompletionServiceType { try await withThrowingTaskGroup(of: String.self) { group in for _ in 0.. String { + func getCompletion(_ request: PromptStrategy) async throws -> AsyncStream { let messages = createMessages(from: request) CodeCompletionLogger.logger.logPrompt(messages.map { ($0.parts.first?.text ?? "N/A", $0.role ?? "N/A") }) - return try await sendMessages(messages) + return AsyncStream { continuation in + let task = Task { + let result = try await sendMessages(messages) + try Task.checkCancellation() + continuation.yield(result) + continuation.finish() + } + continuation.onTermination = { _ in + task.cancel() + } + } } } diff --git a/Core/Sources/CodeCompletionService/OpenAIService.swift b/Core/Sources/CodeCompletionService/OpenAIService.swift index 53fe7a1..8a25f4a 100644 --- a/Core/Sources/CodeCompletionService/OpenAIService.swift +++ b/Core/Sources/CodeCompletionService/OpenAIService.swift @@ -44,18 +44,38 @@ public actor OpenAIService { } extension OpenAIService: CodeCompletionServiceType { - func getCompletion(_ request: PromptStrategy) async throws -> String { + func getCompletion(_ request: PromptStrategy) async throws -> AsyncStream { switch endpoint { case .chatCompletion: let messages = createMessages(from: request) CodeCompletionLogger.logger.logPrompt(messages.map { ($0.content, $0.role.rawValue) }) - return try await sendMessages(messages) + return AsyncStream { continuation in + let task = Task { + let result = try await sendMessages(messages) + try Task.checkCancellation() + continuation.yield(result) + continuation.finish() + } + continuation.onTermination = { _ in + task.cancel() + } + } case .completion: let prompt = createPrompt(from: request) CodeCompletionLogger.logger.logPrompt([(prompt, "user")]) - return try await sendPrompt(prompt) + return AsyncStream { continuation in + let task = Task { + let result = try await sendPrompt(prompt) + try Task.checkCancellation() + continuation.yield(result) + continuation.finish() + } + continuation.onTermination = { _ in + task.cancel() + } + } } } } diff --git a/Core/Sources/CodeCompletionService/TabbyService.swift b/Core/Sources/CodeCompletionService/TabbyService.swift index 4b99423..c27c15e 100644 --- a/Core/Sources/CodeCompletionService/TabbyService.swift +++ b/Core/Sources/CodeCompletionService/TabbyService.swift @@ -27,7 +27,7 @@ actor TabbyService { } extension TabbyService: CodeCompletionServiceType { - func getCompletion(_ request: PromptStrategy) async throws -> String { + func getCompletion(_ request: PromptStrategy) async throws -> AsyncStream { let prefix = request.prefix.joined() let suffix = request.suffix.joined() let clipboard = request.relevantCodeSnippets.map(\.content).joined(separator: "\n\n") @@ -46,7 +46,17 @@ extension TabbyService: CodeCompletionServiceType { (suffix, "suffix"), (clipboard, "clipboard"), ]) - return try await send(requestBody) + return AsyncStream { continuation in + let task = Task { + let result = try await send(requestBody) + try Task.checkCancellation() + continuation.yield(result) + continuation.finish() + } + continuation.onTermination = { _ in + task.cancel() + } + } } } From a8bb261cb9036d007871da7bf17ade8136e5d881 Mon Sep 17 00:00:00 2001 From: Shx Guo Date: Thu, 29 Feb 2024 12:59:14 +0800 Subject: [PATCH 02/27] Add format ollama --- Core/Sources/Fundamental/Models/ChatModel.swift | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/Core/Sources/Fundamental/Models/ChatModel.swift b/Core/Sources/Fundamental/Models/ChatModel.swift index fc45989..582836d 100644 --- a/Core/Sources/Fundamental/Models/ChatModel.swift +++ b/Core/Sources/Fundamental/Models/ChatModel.swift @@ -22,6 +22,7 @@ public struct ChatModel: Codable, Equatable, Identifiable { case azureOpenAI case openAICompatible case googleAI + case ollama case unknown } @@ -86,6 +87,10 @@ public struct ChatModel: Codable, Equatable, Identifiable { let baseURL = info.baseURL if baseURL.isEmpty { return "https://generativelanguage.googleapis.com/v1" } return "\(baseURL)/v1/chat/completions" + case .ollama: + let baseURL = info.baseURL + if baseURL.isEmpty { return "http://localhost:11434/api/chat" } + return "\(baseURL)/api/chat" case .unknown: return "" } From 53d33f75342398a8af2b6714d03aa28c46ce6cff Mon Sep 17 00:00:00 2001 From: Shx Guo Date: Thu, 29 Feb 2024 17:09:25 +0800 Subject: [PATCH 03/27] Make getCompletion returns an AsyncSequence of String --- .../CodeCompletionService/CodeCompletionService.swift | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/Core/Sources/CodeCompletionService/CodeCompletionService.swift b/Core/Sources/CodeCompletionService/CodeCompletionService.swift index 96f3d4f..af5a6f9 100644 --- a/Core/Sources/CodeCompletionService/CodeCompletionService.swift +++ b/Core/Sources/CodeCompletionService/CodeCompletionService.swift @@ -3,9 +3,9 @@ import Fundamental import Storage protocol CodeCompletionServiceType { - func getCompletion( - _ request: PromptStrategy - ) async throws -> AsyncStream + associatedtype CompletionSequence: AsyncSequence where CompletionSequence.Element == String + + func getCompletion(_ request: PromptStrategy) async throws -> CompletionSequence } extension CodeCompletionServiceType { From ce29b359acf4f18ce5a816deed95006f2a1d41b1 Mon Sep 17 00:00:00 2001 From: Shx Guo Date: Thu, 29 Feb 2024 17:10:13 +0800 Subject: [PATCH 04/27] Add ollama to completion model --- Core/Sources/Fundamental/Models/CompletionModel.swift | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/Core/Sources/Fundamental/Models/CompletionModel.swift b/Core/Sources/Fundamental/Models/CompletionModel.swift index f38bfb9..52cbc77 100644 --- a/Core/Sources/Fundamental/Models/CompletionModel.swift +++ b/Core/Sources/Fundamental/Models/CompletionModel.swift @@ -21,7 +21,8 @@ public struct CompletionModel: Codable, Equatable, Identifiable { case openAI case azureOpenAI case openAICompatible - + case ollama + case unknown } @@ -73,6 +74,10 @@ public struct CompletionModel: Codable, Equatable, Identifiable { let version = "2023-07-01-preview" if baseURL.isEmpty { return "" } return "\(baseURL)/openai/deployments/\(deployment)/completions?api-version=\(version)" + case .ollama: + let baseURL = info.baseURL + if baseURL.isEmpty { return "http://localhost:11434/api/generate" } + return "\(baseURL)/api/generate" case .unknown: return "" } From 59b1399d760399d0fc62777ce88e9bfcdd1db0cb Mon Sep 17 00:00:00 2001 From: Shx Guo Date: Thu, 29 Feb 2024 17:12:25 +0800 Subject: [PATCH 05/27] Add OlamaService --- .../CodeCompletionService.swift | 20 ++ .../CodeCompletionService/OllamaService.swift | 240 ++++++++++++++++++ .../ResponseStream.swift | 41 +++ 3 files changed, 301 insertions(+) create mode 100644 Core/Sources/CodeCompletionService/OllamaService.swift create mode 100644 Core/Sources/CodeCompletionService/ResponseStream.swift diff --git a/Core/Sources/CodeCompletionService/CodeCompletionService.swift b/Core/Sources/CodeCompletionService/CodeCompletionService.swift index af5a6f9..281e0f0 100644 --- a/Core/Sources/CodeCompletionService/CodeCompletionService.swift +++ b/Core/Sources/CodeCompletionService/CodeCompletionService.swift @@ -115,6 +115,16 @@ public struct CodeCompletionService { let result = try await service.getCompletions(prompt, count: count) try Task.checkCancellation() return result + case .ollama: + let service = OllamaService( + url: model.endpoint, + endpoint: .chatCompletion, + modelName: model.info.modelName, + stopWords: prompt.stopWords + ) + let result = try await service.getCompletions(prompt, count: count) + try Task.checkCancellation() + return result case .unknown: throw Error.unknownFormat } @@ -150,6 +160,16 @@ public struct CodeCompletionService { let result = try await service.getCompletions(prompt, count: count) try Task.checkCancellation() return result + case .ollama: + let service = OllamaService( + url: model.endpoint, + endpoint: .completion, + modelName: model.info.modelName, + stopWords: prompt.stopWords + ) + let result = try await service.getCompletions(prompt, count: count) + try Task.checkCancellation() + return result case .unknown: throw Error.unknownFormat } diff --git a/Core/Sources/CodeCompletionService/OllamaService.swift b/Core/Sources/CodeCompletionService/OllamaService.swift new file mode 100644 index 0000000..aea819b --- /dev/null +++ b/Core/Sources/CodeCompletionService/OllamaService.swift @@ -0,0 +1,240 @@ +import CopilotForXcodeKit +import Foundation +import Fundamental + +public actor OllamaService { + let url: URL + let endpoint: Endpoint + let modelName: String + let maxToken: Int + let temperature: Double + let stopWords: [String] + + public enum Endpoint { + case completion + case chatCompletion + } + + init( + url: String? = nil, + endpoint: Endpoint, + modelName: String, + maxToken: Int? = nil, + temperature: Double = 0.2, + stopWords: [String] = [] + ) { + self.url = url.flatMap(URL.init(string:)) ?? { + switch endpoint { + case .chatCompletion: + URL(string: "https://127.0.0.1:11434/api/chat")! + case .completion: + URL(string: "https://127.0.0.1:11434/api/generate")! + } + }() + + self.endpoint = endpoint + self.modelName = modelName + self.maxToken = maxToken ?? 4096 + self.temperature = temperature + self.stopWords = stopWords + } +} + +extension OllamaService: CodeCompletionServiceType { + typealias CompletionSequence = AsyncThrowingCompactMapSequence< + ResponseStream, + String + > + + func getCompletion( + _ request: PromptStrategy + ) async throws -> CompletionSequence { + switch endpoint { + case .chatCompletion: + let messages = createMessages(from: request) + CodeCompletionLogger.logger.logPrompt(messages.map { + ($0.content, $0.role.rawValue) + }) + let stream = try await sendMessages(messages) + return stream.compactMap { $0.message?.content } + case .completion: + let prompt = createPrompt(from: request) + CodeCompletionLogger.logger.logPrompt([(prompt, "user")]) + let stream = try await sendPrompt(prompt) + return stream.compactMap { $0.response } + } + } +} + +extension OllamaService { + struct Message: Codable, Equatable { + public enum Role: String, Codable { + case user + case assistant + case system + } + + /// The role of the message. + public var role: Role + /// The content of the message. + public var content: String + } + + enum Error: Swift.Error, LocalizedError { + case decodeError(Swift.Error) + case otherError(String) + + public var errorDescription: String? { + switch self { + case let .decodeError(error): + return error.localizedDescription + case let .otherError(message): + return message + } + } + } +} + +// MARK: - Chat Completion API + +/// https://github.com/ollama/ollama/blob/main/docs/api.md#chat-request-streaming +extension OllamaService { + struct ChatCompletionRequestBody: Codable { + struct Options: Codable { + var temperature: Double + var stop: [String] + var num_predict: Int + var top_k: Int? + var top_p: Double? + } + + var model: String + var messages: [Message] + var stream: Bool + var options: Options + } + + struct ChatCompletionResponseChunk: Decodable { + var model: String + var message: Message? + var response: String? + var done: Bool + var total_duration: Int64? + var load_duration: Int64? + var prompt_eval_count: Int? + var prompt_eval_duration: Int64? + var eval_count: Int? + var eval_duration: Int64? + } + + func createMessages(from request: PromptStrategy) -> [Message] { + let strategy = DefaultTruncateStrategy(maxTokenLimit: max( + maxToken / 3 * 2, + maxToken - 300 - 20 + )) + let prompts = strategy.createTruncatedPrompt(promptStrategy: request) + return [ + .init(role: .system, content: request.systemPrompt), + ] + prompts.map { prompt in + switch prompt.role { + case .user: + return .init(role: .user, content: prompt.content) + case .assistant: + return .init(role: .assistant, content: prompt.content) + } + } + } + + func sendMessages(_ messages: [Message]) async throws + -> ResponseStream + { + let requestBody = ChatCompletionRequestBody( + model: modelName, + messages: messages, + stream: true, + options: .init( + temperature: temperature, + stop: stopWords, + num_predict: 300 + ) + ) + + var request = URLRequest(url: url) + request.httpMethod = "POST" + let encoder = JSONEncoder() + request.httpBody = try encoder.encode(requestBody) + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + let (result, response) = try await URLSession.shared.bytes(for: request) + + guard let response = response as? HTTPURLResponse else { + throw CancellationError() + } + + guard response.statusCode == 200 else { + let text = try await result.lines.reduce(into: "") { partialResult, current in + partialResult += current + } + throw Error.otherError(text) + } + + return ResponseStream(result: result) + } +} + +// MARK: - Completion API + +extension OllamaService { + struct CompletionRequestBody: Codable { + var model: String + var prompt: String + var stream: Bool + var options: ChatCompletionRequestBody.Options + } + + func createPrompt(from request: PromptStrategy) -> String { + let strategy = DefaultTruncateStrategy(maxTokenLimit: max( + maxToken / 3 * 2, + maxToken - 300 - 20 + )) + let prompts = strategy.createTruncatedPrompt(promptStrategy: request) + return ([request.systemPrompt] + prompts.map(\.content)).joined(separator: "\n\n") + } + + func sendPrompt(_ prompt: String) async throws -> ResponseStream { + let requestBody = CompletionRequestBody( + model: modelName, + prompt: prompt, + stream: true, + options: .init( + temperature: temperature, + stop: stopWords, + num_predict: 300 + ) + ) + + var request = URLRequest(url: url) + request.httpMethod = "POST" + let encoder = JSONEncoder() + request.httpBody = try encoder.encode(requestBody) + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + let (result, response) = try await URLSession.shared.bytes(for: request) + + guard let response = response as? HTTPURLResponse else { + throw CancellationError() + } + + guard response.statusCode == 200 else { + let text = try await result.lines.reduce(into: "") { partialResult, current in + partialResult += current + } + throw Error.otherError(text) + } + + return ResponseStream(result: result) + } + + func countToken(_ message: Message) -> Int { + message.content.count + } +} + diff --git a/Core/Sources/CodeCompletionService/ResponseStream.swift b/Core/Sources/CodeCompletionService/ResponseStream.swift new file mode 100644 index 0000000..41996e7 --- /dev/null +++ b/Core/Sources/CodeCompletionService/ResponseStream.swift @@ -0,0 +1,41 @@ +import Foundation + +struct ResponseStream: AsyncSequence { + func makeAsyncIterator() -> Stream.AsyncIterator { + stream.makeAsyncIterator() + } + + typealias Stream = AsyncThrowingStream + typealias AsyncIterator = Stream.AsyncIterator + typealias Element = Chunk + + let stream: Stream + + init(result: URLSession.AsyncBytes) { + stream = AsyncThrowingStream { continuation in + let task = Task { + do { + for try await line in result.lines { + if Task.isCancelled { break } + let prefix = "data: " + guard line.hasPrefix(prefix), + let content = line.dropFirst(prefix.count).data(using: .utf8), + let chunk = try? JSONDecoder() + .decode(Chunk.self, from: content) + else { continue } + continuation.yield(chunk) + } + continuation.finish() + } catch { + continuation.finish(throwing: error) + result.task.cancel() + } + } + continuation.onTermination = { _ in + task.cancel() + result.task.cancel() + } + } + } +} + From b35e453e02808c999992e4dae15856e427c0cb11 Mon Sep 17 00:00:00 2001 From: Shx Guo Date: Thu, 29 Feb 2024 17:29:04 +0800 Subject: [PATCH 06/27] Add UI for Ollama models --- .../ChatModelEditView.swift | 18 ++++++++++++++++++ .../CompletionModelEditView.swift | 18 ++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/CustomSuggestionService/ChatModelManagement/ChatModelEditView.swift b/CustomSuggestionService/ChatModelManagement/ChatModelEditView.swift index 71fb62b..2f1ff20 100644 --- a/CustomSuggestionService/ChatModelManagement/ChatModelEditView.swift +++ b/CustomSuggestionService/ChatModelManagement/ChatModelEditView.swift @@ -25,6 +25,8 @@ struct ChatModelEditView: View { openAICompatible case .googleAI: googleAI + case .ollama: + ollama case .unknown: EmptyView() } @@ -77,6 +79,8 @@ struct ChatModelEditView: View { Text("OpenAI Compatible").tag(format) case .googleAI: Text("Google Generative AI").tag(format) + case .ollama: + Text("Ollama").tag(format) case .unknown: EmptyView() } @@ -260,6 +264,20 @@ struct ChatModelEditView: View { maxTokensTextField } + + @ViewBuilder + var ollama: some View { + baseURLTextField( + title: "", + prompt: Text("https://127.0.0.1:11434/api/chat") + ) { + Text("/api/chat") + } + + TextField("Model Name", text: $store.modelName) + + maxTokensTextField + } } #Preview("OpenAI") { diff --git a/CustomSuggestionService/ChatModelManagement/CompletionModelEditView.swift b/CustomSuggestionService/ChatModelManagement/CompletionModelEditView.swift index d2286b3..289dee9 100644 --- a/CustomSuggestionService/ChatModelManagement/CompletionModelEditView.swift +++ b/CustomSuggestionService/ChatModelManagement/CompletionModelEditView.swift @@ -23,6 +23,8 @@ struct CompletionModelEditView: View { azureOpenAI case .openAICompatible: openAICompatible + case .ollama: + ollama case .unknown: EmptyView() } @@ -73,6 +75,8 @@ struct CompletionModelEditView: View { Text("Azure OpenAI").tag(format) case .openAICompatible: Text("OpenAI Compatible").tag(format) + case .ollama: + Text("Ollama").tag(format) case .unknown: EmptyView() } @@ -252,6 +256,20 @@ struct CompletionModelEditView: View { maxTokensTextField } + + @ViewBuilder + var ollama: some View { + baseURLTextField( + title: "", + prompt: Text("https://127.0.0.1:11434/api/generate") + ) { + Text("/api/generate") + } + + TextField("Model Name", text: $store.modelName) + + maxTokensTextField + } } #Preview("OpenAI") { From df23c0f5fad944f1014b5052fc8aefbdb29c4cc6 Mon Sep 17 00:00:00 2001 From: Shx Guo Date: Thu, 29 Feb 2024 18:06:34 +0800 Subject: [PATCH 07/27] Fix streaming --- Core/Sources/CodeCompletionService/ResponseStream.swift | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/Core/Sources/CodeCompletionService/ResponseStream.swift b/Core/Sources/CodeCompletionService/ResponseStream.swift index 41996e7..597334f 100644 --- a/Core/Sources/CodeCompletionService/ResponseStream.swift +++ b/Core/Sources/CodeCompletionService/ResponseStream.swift @@ -11,18 +11,15 @@ struct ResponseStream: AsyncSequence { let stream: Stream - init(result: URLSession.AsyncBytes) { + init(result: URLSession.AsyncBytes, lineExtractor: @escaping (String) -> String? = { $0 }) { stream = AsyncThrowingStream { continuation in let task = Task { do { for try await line in result.lines { if Task.isCancelled { break } - let prefix = "data: " - guard line.hasPrefix(prefix), - let content = line.dropFirst(prefix.count).data(using: .utf8), - let chunk = try? JSONDecoder() - .decode(Chunk.self, from: content) + guard let content = lineExtractor(line)?.data(using: .utf8) else { continue } + let chunk = try JSONDecoder().decode(Chunk.self, from: content) continuation.yield(chunk) } continuation.finish() From d0851698327cb45a29e132c60dbdb63cbf23a474 Mon Sep 17 00:00:00 2001 From: Shx Guo Date: Thu, 29 Feb 2024 22:35:32 +0800 Subject: [PATCH 08/27] Add UI form Ollama models --- .../CodeCompletionService.swift | 8 ++++-- .../CodeCompletionService/OllamaService.swift | 25 ++++++++++++++++--- .../Fundamental/Models/ChatModel.swift | 6 ++++- .../Fundamental/Models/CompletionModel.swift | 6 ++++- .../ChatModelManagement/ChatModelEdit.swift | 7 ++++-- .../ChatModelEditView.swift | 11 ++++++++ .../CompletionModelEdit.swift | 7 ++++-- .../CompletionModelEditView.swift | 11 ++++++++ 8 files changed, 70 insertions(+), 11 deletions(-) diff --git a/Core/Sources/CodeCompletionService/CodeCompletionService.swift b/Core/Sources/CodeCompletionService/CodeCompletionService.swift index 281e0f0..2214f58 100644 --- a/Core/Sources/CodeCompletionService/CodeCompletionService.swift +++ b/Core/Sources/CodeCompletionService/CodeCompletionService.swift @@ -120,7 +120,9 @@ public struct CodeCompletionService { url: model.endpoint, endpoint: .chatCompletion, modelName: model.info.modelName, - stopWords: prompt.stopWords + stopWords: prompt.stopWords, + keepAlive: model.info.ollamaKeepAlive, + format: .none ) let result = try await service.getCompletions(prompt, count: count) try Task.checkCancellation() @@ -165,7 +167,9 @@ public struct CodeCompletionService { url: model.endpoint, endpoint: .completion, modelName: model.info.modelName, - stopWords: prompt.stopWords + stopWords: prompt.stopWords, + keepAlive: model.info.ollamaKeepAlive, + format: .none ) let result = try await service.getCompletions(prompt, count: count) try Task.checkCancellation() diff --git a/Core/Sources/CodeCompletionService/OllamaService.swift b/Core/Sources/CodeCompletionService/OllamaService.swift index aea819b..d3e62b4 100644 --- a/Core/Sources/CodeCompletionService/OllamaService.swift +++ b/Core/Sources/CodeCompletionService/OllamaService.swift @@ -9,6 +9,13 @@ public actor OllamaService { let maxToken: Int let temperature: Double let stopWords: [String] + let keepAlive: String + let format: ResponseFormat + + public enum ResponseFormat: String { + case none = "" + case json = "json" + } public enum Endpoint { case completion @@ -21,7 +28,9 @@ public actor OllamaService { modelName: String, maxToken: Int? = nil, temperature: Double = 0.2, - stopWords: [String] = [] + stopWords: [String] = [], + keepAlive: String = "", + format: ResponseFormat = .none ) { self.url = url.flatMap(URL.init(string:)) ?? { switch endpoint { @@ -37,6 +46,8 @@ public actor OllamaService { self.maxToken = maxToken ?? 4096 self.temperature = temperature self.stopWords = stopWords + self.keepAlive = keepAlive + self.format = format } } @@ -112,6 +123,8 @@ extension OllamaService { var messages: [Message] var stream: Bool var options: Options + var keep_alive: String? + var format: String? } struct ChatCompletionResponseChunk: Decodable { @@ -156,7 +169,9 @@ extension OllamaService { temperature: temperature, stop: stopWords, num_predict: 300 - ) + ), + keep_alive: keepAlive.isEmpty ? nil : keepAlive, + format: format == .none ? nil : format.rawValue ) var request = URLRequest(url: url) @@ -189,6 +204,8 @@ extension OllamaService { var prompt: String var stream: Bool var options: ChatCompletionRequestBody.Options + var keep_alive: String? + var format: String? } func createPrompt(from request: PromptStrategy) -> String { @@ -209,7 +226,9 @@ extension OllamaService { temperature: temperature, stop: stopWords, num_predict: 300 - ) + ), + keep_alive: keepAlive.isEmpty ? nil : keepAlive, + format: format == .none ? nil : format.rawValue ) var request = URLRequest(url: url) diff --git a/Core/Sources/Fundamental/Models/ChatModel.swift b/Core/Sources/Fundamental/Models/ChatModel.swift index 582836d..3cb7d61 100644 --- a/Core/Sources/Fundamental/Models/ChatModel.swift +++ b/Core/Sources/Fundamental/Models/ChatModel.swift @@ -46,6 +46,8 @@ public struct ChatModel: Codable, Equatable, Identifiable { get { modelName } set { modelName = newValue } } + @FallbackDecoding + public var ollamaKeepAlive: String public init( apiKeyName: String = "", @@ -54,7 +56,8 @@ public struct ChatModel: Codable, Equatable, Identifiable { maxTokens: Int = 4000, supportsFunctionCalling: Bool = true, supportsOpenAIAPI2023_11: Bool = false, - modelName: String = "" + modelName: String = "", + ollamaKeepAlive: String = "" ) { self.apiKeyName = apiKeyName self.baseURL = baseURL @@ -63,6 +66,7 @@ public struct ChatModel: Codable, Equatable, Identifiable { self.supportsFunctionCalling = supportsFunctionCalling self.supportsOpenAIAPI2023_11 = supportsOpenAIAPI2023_11 self.modelName = modelName + self.ollamaKeepAlive = ollamaKeepAlive } } diff --git a/Core/Sources/Fundamental/Models/CompletionModel.swift b/Core/Sources/Fundamental/Models/CompletionModel.swift index 52cbc77..3e4af42 100644 --- a/Core/Sources/Fundamental/Models/CompletionModel.swift +++ b/Core/Sources/Fundamental/Models/CompletionModel.swift @@ -41,19 +41,23 @@ public struct CompletionModel: Codable, Equatable, Identifiable { get { modelName } set { modelName = newValue } } + @FallbackDecoding + public var ollamaKeepAlive: String public init( apiKeyName: String = "", baseURL: String = "", isFullURL: Bool = false, maxTokens: Int = 4000, - modelName: String = "" + modelName: String = "", + ollamaKeepAlive: String = "" ) { self.apiKeyName = apiKeyName self.baseURL = baseURL self.isFullURL = isFullURL self.maxTokens = maxTokens self.modelName = modelName + self.ollamaKeepAlive = ollamaKeepAlive } } diff --git a/CustomSuggestionService/ChatModelManagement/ChatModelEdit.swift b/CustomSuggestionService/ChatModelManagement/ChatModelEdit.swift index a8ec53c..13409a6 100644 --- a/CustomSuggestionService/ChatModelManagement/ChatModelEdit.swift +++ b/CustomSuggestionService/ChatModelManagement/ChatModelEdit.swift @@ -20,6 +20,7 @@ struct ChatModelEdit { var suggestedMaxTokens: Int? var apiKeySelection: APIKeySelection.State = .init() var baseURLSelection: BaseURLSelection.State = .init() + var ollamaKeepAlive: String = "" } enum Action: Equatable, BindableAction { @@ -152,7 +153,8 @@ extension ChatModel { apiKeyName: info.apiKeyName, apiKeyManagement: .init(availableAPIKeyNames: [info.apiKeyName]) ), - baseURLSelection: .init(baseURL: info.baseURL) + baseURLSelection: .init(baseURL: info.baseURL), + ollamaKeepAlive: info.ollamaKeepAlive ) } @@ -166,7 +168,8 @@ extension ChatModel { baseURL: state.baseURL.trimmingCharacters(in: .whitespacesAndNewlines), maxTokens: state.maxTokens, supportsFunctionCalling: false, - modelName: state.modelName.trimmingCharacters(in: .whitespacesAndNewlines) + modelName: state.modelName.trimmingCharacters(in: .whitespacesAndNewlines), + ollamaKeepAlive: state.ollamaKeepAlive ) ) } diff --git a/CustomSuggestionService/ChatModelManagement/ChatModelEditView.swift b/CustomSuggestionService/ChatModelManagement/ChatModelEditView.swift index 2f1ff20..a8134fe 100644 --- a/CustomSuggestionService/ChatModelManagement/ChatModelEditView.swift +++ b/CustomSuggestionService/ChatModelManagement/ChatModelEditView.swift @@ -277,6 +277,17 @@ struct ChatModelEditView: View { TextField("Model Name", text: $store.modelName) maxTokensTextField + + TextField(text: $store.ollamaKeepAlive, prompt: Text("Default Value")) { + Text("Keep Alive") + } + + VStack(alignment: .leading, spacing: 8) { + Text(Image(systemName: "exclamationmark.triangle.fill")) + Text( + " For more details, please visit [https://ollama.com](https://ollama.com)" + ) + } + .padding(.vertical) } } diff --git a/CustomSuggestionService/ChatModelManagement/CompletionModelEdit.swift b/CustomSuggestionService/ChatModelManagement/CompletionModelEdit.swift index 1931f09..78b31d7 100644 --- a/CustomSuggestionService/ChatModelManagement/CompletionModelEdit.swift +++ b/CustomSuggestionService/ChatModelManagement/CompletionModelEdit.swift @@ -20,6 +20,7 @@ struct CompletionModelEdit { var suggestedMaxTokens: Int? var apiKeySelection: APIKeySelection.State = .init() var baseURLSelection: BaseURLSelection.State = .init() + var ollamaKeepAlive: String = "" } enum Action: Equatable, BindableAction { @@ -143,7 +144,8 @@ extension CompletionModel { apiKeyName: info.apiKeyName, apiKeyManagement: .init(availableAPIKeyNames: [info.apiKeyName]) ), - baseURLSelection: .init(baseURL: info.baseURL) + baseURLSelection: .init(baseURL: info.baseURL), + ollamaKeepAlive: info.ollamaKeepAlive ) } @@ -156,7 +158,8 @@ extension CompletionModel { apiKeyName: state.apiKeyName, baseURL: state.baseURL.trimmingCharacters(in: .whitespacesAndNewlines), maxTokens: state.maxTokens, - modelName: state.modelName.trimmingCharacters(in: .whitespacesAndNewlines) + modelName: state.modelName.trimmingCharacters(in: .whitespacesAndNewlines), + ollamaKeepAlive: state.ollamaKeepAlive ) ) } diff --git a/CustomSuggestionService/ChatModelManagement/CompletionModelEditView.swift b/CustomSuggestionService/ChatModelManagement/CompletionModelEditView.swift index 289dee9..fa3e830 100644 --- a/CustomSuggestionService/ChatModelManagement/CompletionModelEditView.swift +++ b/CustomSuggestionService/ChatModelManagement/CompletionModelEditView.swift @@ -269,6 +269,17 @@ struct CompletionModelEditView: View { TextField("Model Name", text: $store.modelName) maxTokensTextField + + TextField(text: $store.ollamaKeepAlive, prompt: Text("Default Value")) { + Text("Keep Alive") + } + + VStack(alignment: .leading, spacing: 8) { + Text(Image(systemName: "exclamationmark.triangle.fill")) + Text( + " For more details, please visit [https://ollama.com](https://ollama.com)" + ) + } + .padding(.vertical) } } From bdfe08ee7ad7637c61d1cd509b50c07aef0f886e Mon Sep 17 00:00:00 2001 From: Shx Guo Date: Fri, 1 Mar 2024 14:30:01 +0800 Subject: [PATCH 09/27] Update ResponseStream accept a block that parses lines --- .../CodeCompletionService/OllamaService.swift | 20 +++++++++++++++---- .../ResponseStream.swift | 19 ++++++++++++------ 2 files changed, 29 insertions(+), 10 deletions(-) diff --git a/Core/Sources/CodeCompletionService/OllamaService.swift b/Core/Sources/CodeCompletionService/OllamaService.swift index d3e62b4..279e2d7 100644 --- a/Core/Sources/CodeCompletionService/OllamaService.swift +++ b/Core/Sources/CodeCompletionService/OllamaService.swift @@ -11,10 +11,10 @@ public actor OllamaService { let stopWords: [String] let keepAlive: String let format: ResponseFormat - + public enum ResponseFormat: String { case none = "" - case json = "json" + case json } public enum Endpoint { @@ -192,7 +192,13 @@ extension OllamaService { throw Error.otherError(text) } - return ResponseStream(result: result) + return ResponseStream(result: result) { + let chunk = try JSONDecoder().decode( + ChatCompletionResponseChunk.self, + from: $0.data(using: .utf8) ?? Data() + ) + return .init(chunk: chunk, done: chunk.done) + } } } @@ -249,7 +255,13 @@ extension OllamaService { throw Error.otherError(text) } - return ResponseStream(result: result) + return ResponseStream(result: result) { + let chunk = try JSONDecoder().decode( + ChatCompletionResponseChunk.self, + from: $0.data(using: .utf8) ?? Data() + ) + return .init(chunk: chunk, done: chunk.done) + } } func countToken(_ message: Message) -> Int { diff --git a/Core/Sources/CodeCompletionService/ResponseStream.swift b/Core/Sources/CodeCompletionService/ResponseStream.swift index 597334f..ce28b7f 100644 --- a/Core/Sources/CodeCompletionService/ResponseStream.swift +++ b/Core/Sources/CodeCompletionService/ResponseStream.swift @@ -1,6 +1,6 @@ import Foundation -struct ResponseStream: AsyncSequence { +struct ResponseStream: AsyncSequence { func makeAsyncIterator() -> Stream.AsyncIterator { stream.makeAsyncIterator() } @@ -8,19 +8,26 @@ struct ResponseStream: AsyncSequence { typealias Stream = AsyncThrowingStream typealias AsyncIterator = Stream.AsyncIterator typealias Element = Chunk + + struct LineContent { + let chunk: Chunk? + let done: Bool + } let stream: Stream - init(result: URLSession.AsyncBytes, lineExtractor: @escaping (String) -> String? = { $0 }) { + init(result: URLSession.AsyncBytes, lineExtractor: @escaping (String) throws -> LineContent) { stream = AsyncThrowingStream { continuation in let task = Task { do { for try await line in result.lines { if Task.isCancelled { break } - guard let content = lineExtractor(line)?.data(using: .utf8) - else { continue } - let chunk = try JSONDecoder().decode(Chunk.self, from: content) - continuation.yield(chunk) + let content = try lineExtractor(line) + if let chunk = content.chunk { + continuation.yield(chunk) + } + + if content.done { break } } continuation.finish() } catch { From 83df7c7e7963e869d24b95ee5f81a72bf9968ded Mon Sep 17 00:00:00 2001 From: Shx Guo Date: Fri, 1 Mar 2024 15:01:28 +0800 Subject: [PATCH 10/27] Add RawSuggestionPostProcessingStrategy --- ...tRawSuggestionPostProcessingStrategy.swift | 105 ++++++++++++++++++ ...pRawSuggestionPostProcessingStrategy.swift | 8 ++ .../ContinueRequestStrategy.swift | 23 ++-- .../DefaultRequestStrategy.swift | 23 ++-- .../NaiveRequestStrategy.swift | 4 + .../SuggestionService/RequestStrategy.swift | 90 +-------------- Core/Sources/SuggestionService/Service.swift | 13 ++- 7 files changed, 140 insertions(+), 126 deletions(-) create mode 100644 Core/Sources/SuggestionService/RawSuggestionPostProcessing/DefaultRawSuggestionPostProcessingStrategy.swift create mode 100644 Core/Sources/SuggestionService/RawSuggestionPostProcessing/NoOpRawSuggestionPostProcessingStrategy.swift diff --git a/Core/Sources/SuggestionService/RawSuggestionPostProcessing/DefaultRawSuggestionPostProcessingStrategy.swift b/Core/Sources/SuggestionService/RawSuggestionPostProcessing/DefaultRawSuggestionPostProcessingStrategy.swift new file mode 100644 index 0000000..d7384f9 --- /dev/null +++ b/Core/Sources/SuggestionService/RawSuggestionPostProcessing/DefaultRawSuggestionPostProcessingStrategy.swift @@ -0,0 +1,105 @@ +import Foundation +import Parsing + +protocol RawSuggestionPostProcessingStrategy { + func postProcessRawSuggestion(suggestionPrefix: String, suggestion: String) -> String +} + +struct DefaultRawSuggestionPostProcessingStrategy: RawSuggestionPostProcessingStrategy { + let openingCodeTag: String + let closingCodeTag: String + + func postProcessRawSuggestion(suggestionPrefix: String, suggestion: String) -> String { + let suggestion = extractEnclosingSuggestion( + from: removeLeadingAndTrailingMarkdownCodeBlockMark(from: suggestion), + openingTag: openingCodeTag, + closingTag: closingCodeTag + ) + + if suggestion.hasPrefix(suggestionPrefix) { + var processed = suggestion + processed.removeFirst(suggestionPrefix.count) + return processed + } + + return suggestionPrefix + suggestion + } + + /// Extract suggestions that is enclosed in tags. + func extractEnclosingSuggestion( + from response: String, + openingTag: String, + closingTag: String + ) -> String { + let case_openingTagAtTheStart_parseEverythingInsideTheTag = Parse(input: Substring.self) { + openingTag + + OneOf { // parse until tags or the end + Parse { + OneOf { + PrefixUpTo(openingTag) + PrefixUpTo(closingTag) + } + Skip { + Rest() + } + } + + Rest() + } + } + + let case_noTagAtTheStart_parseEverythingBeforeTheTag = Parse(input: Substring.self) { + OneOf { + PrefixUpTo(openingTag) + PrefixUpTo(closingTag) + } + + Skip { + Rest() + } + } + + let parser = Parse(input: Substring.self) { + OneOf { + case_openingTagAtTheStart_parseEverythingInsideTheTag + case_noTagAtTheStart_parseEverythingBeforeTheTag + Rest() + } + } + + var text = response[...] + do { + let suggestion = try parser.parse(&text) + return String(suggestion) + } catch { + return response + } + } + + /// If the response starts with markdown code block, we should remove it. + func removeLeadingAndTrailingMarkdownCodeBlockMark(from response: String) -> String { + let removePrefixMarkdownCodeBlockMark = Parse(input: Substring.self) { + Skip { + "```" + PrefixThrough("\n") + } + OneOf { + Parse { + PrefixUpTo("```") + Skip { Rest() } + } + Rest() + } + } + + do { + var response = response[...] + let suggestion = try removePrefixMarkdownCodeBlockMark.parse(&response) + return String(suggestion) + } catch { + return response + } + } +} + diff --git a/Core/Sources/SuggestionService/RawSuggestionPostProcessing/NoOpRawSuggestionPostProcessingStrategy.swift b/Core/Sources/SuggestionService/RawSuggestionPostProcessing/NoOpRawSuggestionPostProcessingStrategy.swift new file mode 100644 index 0000000..213aafd --- /dev/null +++ b/Core/Sources/SuggestionService/RawSuggestionPostProcessing/NoOpRawSuggestionPostProcessingStrategy.swift @@ -0,0 +1,8 @@ +import Foundation + +struct NoOpRawSuggestionPostProcessingStrategy: RawSuggestionPostProcessingStrategy { + func postProcessRawSuggestion(suggestionPrefix: String, suggestion: String) -> String { + suggestionPrefix + suggestion + } +} + diff --git a/Core/Sources/SuggestionService/RequestStrategies/ContinueRequestStrategy.swift b/Core/Sources/SuggestionService/RequestStrategies/ContinueRequestStrategy.swift index 466e7f8..b53eee7 100644 --- a/Core/Sources/SuggestionService/RequestStrategies/ContinueRequestStrategy.swift +++ b/Core/Sources/SuggestionService/RequestStrategies/ContinueRequestStrategy.swift @@ -21,6 +21,13 @@ struct ContinueRequestStrategy: RequestStrategy { suffix: suffix ) } + + func createRawSuggestionPostProcessor() -> DefaultRawSuggestionPostProcessingStrategy { + DefaultRawSuggestionPostProcessingStrategy( + openingCodeTag: Tag.openingCode, + closingCodeTag: Tag.closingCode + ) + } enum Tag { public static let openingCode = "" @@ -158,21 +165,5 @@ struct ContinueRequestStrategy: RequestStrategy { ) } } - - func postProcessRawSuggestion(suggestionPrefix: String, suggestion: String) -> String { - let suggestion = extractEnclosingSuggestion( - from: removeLeadingAndTrailingMarkdownCodeBlockMark(from: suggestion), - openingTag: Tag.openingCode, - closingTag: Tag.closingCode - ) - - if suggestion.hasPrefix(suggestionPrefix) { - var processed = suggestion - processed.removeFirst(suggestionPrefix.count) - return processed - } - - return suggestionPrefix + suggestion - } } diff --git a/Core/Sources/SuggestionService/RequestStrategies/DefaultRequestStrategy.swift b/Core/Sources/SuggestionService/RequestStrategies/DefaultRequestStrategy.swift index 5d41068..cab2ab5 100644 --- a/Core/Sources/SuggestionService/RequestStrategies/DefaultRequestStrategy.swift +++ b/Core/Sources/SuggestionService/RequestStrategies/DefaultRequestStrategy.swift @@ -22,6 +22,13 @@ struct DefaultRequestStrategy: RequestStrategy { suffix: suffix ) } + + func createRawSuggestionPostProcessor() -> DefaultRawSuggestionPostProcessingStrategy { + DefaultRawSuggestionPostProcessingStrategy( + openingCodeTag: Tag.openingCode, + closingCodeTag: Tag.closingCode + ) + } enum Tag { public static let openingCode = "" @@ -142,21 +149,5 @@ struct DefaultRequestStrategy: RequestStrategy { ) } } - - func postProcessRawSuggestion(suggestionPrefix: String, suggestion: String) -> String { - let suggestion = extractEnclosingSuggestion( - from: removeLeadingAndTrailingMarkdownCodeBlockMark(from: suggestion), - openingTag: Tag.openingCode, - closingTag: Tag.closingCode - ) - - if suggestion.hasPrefix(suggestionPrefix) { - var processed = suggestion - processed.removeFirst(suggestionPrefix.count) - return processed - } - - return suggestionPrefix + suggestion - } } diff --git a/Core/Sources/SuggestionService/RequestStrategies/NaiveRequestStrategy.swift b/Core/Sources/SuggestionService/RequestStrategies/NaiveRequestStrategy.swift index 73c26e9..1274c37 100644 --- a/Core/Sources/SuggestionService/RequestStrategies/NaiveRequestStrategy.swift +++ b/Core/Sources/SuggestionService/RequestStrategies/NaiveRequestStrategy.swift @@ -20,6 +20,10 @@ struct NaiveRequestStrategy: RequestStrategy { suffix: suffix ) } + + func createRawSuggestionPostProcessor() -> some RawSuggestionPostProcessingStrategy { + NoOpRawSuggestionPostProcessingStrategy() + } struct Request: PromptStrategy { let systemPrompt: String = "" diff --git a/Core/Sources/SuggestionService/RequestStrategy.swift b/Core/Sources/SuggestionService/RequestStrategy.swift index 7ada9fd..5be3d00 100644 --- a/Core/Sources/SuggestionService/RequestStrategy.swift +++ b/Core/Sources/SuggestionService/RequestStrategy.swift @@ -7,6 +7,7 @@ import Parsing /// This protocol allows for different strategies to be used to generate prompts. protocol RequestStrategy { associatedtype Prompt: PromptStrategy + associatedtype RawSuggestionPostProcessor: RawSuggestionPostProcessingStrategy init(sourceRequest: SuggestionRequest, prefix: [String], suffix: [String]) @@ -18,9 +19,7 @@ protocol RequestStrategy { /// The AI model may not return a suggestion in a ideal format. You can use it to reformat the /// suggestions. - /// - /// By default, it will return the prefix + suggestion. - func postProcessRawSuggestion(suggestionPrefix: String, suggestion: String) -> String + func createRawSuggestionPostProcessor() -> RawSuggestionPostProcessor } public enum RequestStrategyOption: String, CaseIterable, Codable { @@ -46,91 +45,6 @@ extension RequestStrategyOption { extension RequestStrategy { var shouldSkip: Bool { false } - - func postProcessRawSuggestion(suggestionPrefix: String, suggestion: String) -> String { - suggestionPrefix + suggestion - } -} - -// MARK: - Shared Implementations - -extension RequestStrategy { - /// Extract suggestions that is enclosed in tags. - func extractEnclosingSuggestion( - from response: String, - openingTag: String, - closingTag: String - ) -> String { - let case_openingTagAtTheStart_parseEverythingInsideTheTag = Parse(input: Substring.self) { - openingTag - - OneOf { // parse until tags or the end - Parse { - OneOf { - PrefixUpTo(openingTag) - PrefixUpTo(closingTag) - } - Skip { - Rest() - } - } - - Rest() - } - } - - let case_noTagAtTheStart_parseEverythingBeforeTheTag = Parse(input: Substring.self) { - OneOf { - PrefixUpTo(openingTag) - PrefixUpTo(closingTag) - } - - Skip { - Rest() - } - } - - let parser = Parse(input: Substring.self) { - OneOf { - case_openingTagAtTheStart_parseEverythingInsideTheTag - case_noTagAtTheStart_parseEverythingBeforeTheTag - Rest() - } - } - - var text = response[...] - do { - let suggestion = try parser.parse(&text) - return String(suggestion) - } catch { - return response - } - } - - /// If the response starts with markdown code block, we should remove it. - func removeLeadingAndTrailingMarkdownCodeBlockMark(from response: String) -> String { - let removePrefixMarkdownCodeBlockMark = Parse(input: Substring.self) { - Skip { - "```" - PrefixThrough("\n") - } - OneOf { - Parse { - PrefixUpTo("```") - Skip { Rest() } - } - Rest() - } - } - - do { - var response = response[...] - let suggestion = try removePrefixMarkdownCodeBlockMark.parse(&response) - return String(suggestion) - } catch { - return response - } - } } // MARK: - Suggestion Prefix Helpers diff --git a/Core/Sources/SuggestionService/Service.swift b/Core/Sources/SuggestionService/Service.swift index b6cca62..f987293 100644 --- a/Core/Sources/SuggestionService/Service.swift +++ b/Core/Sources/SuggestionService/Service.swift @@ -44,7 +44,8 @@ actor Service { let service = CodeCompletionService() - let promptStrategy = strategy.createPrompt() + let prompt = strategy.createPrompt() + let postProcessor = strategy.createRawSuggestionPostProcessor() let suggestedCodeSnippets: [String] @@ -52,21 +53,21 @@ actor Service { case let .chatModel(model): CodeCompletionLogger.logger.logModel(model) suggestedCodeSnippets = try await service.getCompletions( - promptStrategy, + prompt, model: model, count: 1 ) case let .completionModel(model): CodeCompletionLogger.logger.logModel(model) suggestedCodeSnippets = try await service.getCompletions( - promptStrategy, + prompt, model: model, count: 1 ) case let .tabbyModel(model): CodeCompletionLogger.logger.logModel(model) suggestedCodeSnippets = try await service.getCompletions( - promptStrategy, + prompt, model: model, count: 1 ) @@ -80,8 +81,8 @@ actor Service { CodeSuggestion( id: UUID().uuidString, text: Self.removeTrailingNewlinesAndWhitespace( - from: strategy.postProcessRawSuggestion( - suggestionPrefix: promptStrategy + from: postProcessor.postProcessRawSuggestion( + suggestionPrefix: prompt .suggestionPrefix.prependingValue, suggestion: $0 ) From 0d09f80f3760b739ef4bd0913c99c5781556895c Mon Sep 17 00:00:00 2001 From: Shx Guo Date: Fri, 1 Mar 2024 16:45:00 +0800 Subject: [PATCH 11/27] Update the API of RawSuggestionPostProcessingStrategy --- ...tRawSuggestionPostProcessingStrategy.swift | 59 +++++++++++++++---- ...pRawSuggestionPostProcessingStrategy.swift | 4 +- Core/Sources/SuggestionService/Service.swift | 18 ++---- 3 files changed, 54 insertions(+), 27 deletions(-) diff --git a/Core/Sources/SuggestionService/RawSuggestionPostProcessing/DefaultRawSuggestionPostProcessingStrategy.swift b/Core/Sources/SuggestionService/RawSuggestionPostProcessing/DefaultRawSuggestionPostProcessingStrategy.swift index d7384f9..bd8dc34 100644 --- a/Core/Sources/SuggestionService/RawSuggestionPostProcessing/DefaultRawSuggestionPostProcessingStrategy.swift +++ b/Core/Sources/SuggestionService/RawSuggestionPostProcessing/DefaultRawSuggestionPostProcessingStrategy.swift @@ -2,31 +2,66 @@ import Foundation import Parsing protocol RawSuggestionPostProcessingStrategy { - func postProcessRawSuggestion(suggestionPrefix: String, suggestion: String) -> String + func postProcess(rawSuggestion: String, infillPrefix: String, suffix: [String]) -> String +} + +extension RawSuggestionPostProcessingStrategy { + func removeTrailingNewlinesAndWhitespace(from string: String) -> String { + var text = string[...] + while let last = text.last, last.isNewline || last.isWhitespace { + text = text.dropLast(1) + } + return String(text) + } } struct DefaultRawSuggestionPostProcessingStrategy: RawSuggestionPostProcessingStrategy { let openingCodeTag: String let closingCodeTag: String - func postProcessRawSuggestion(suggestionPrefix: String, suggestion: String) -> String { - let suggestion = extractEnclosingSuggestion( - from: removeLeadingAndTrailingMarkdownCodeBlockMark(from: suggestion), + func postProcess(rawSuggestion: String, infillPrefix: String, suffix: [String]) -> String { + var suggestion = extractSuggestion(from: rawSuggestion) + removePrefix(from: &suggestion, infillPrefix: infillPrefix) + removeSuffix(from: &suggestion, suffix: suffix) + return removeTrailingNewlinesAndWhitespace(from: infillPrefix + suggestion) + } + + func extractSuggestion(from response: String) -> String { + let escapedMarkdownCodeBlock = removeLeadingAndTrailingMarkdownCodeBlockMark(from: response) + let escapedTags = extractEnclosingSuggestion( + from: escapedMarkdownCodeBlock, openingTag: openingCodeTag, closingTag: closingCodeTag ) - if suggestion.hasPrefix(suggestionPrefix) { - var processed = suggestion - processed.removeFirst(suggestionPrefix.count) - return processed + return escapedTags + } + + func removePrefix(from suggestion: inout String, infillPrefix: String) { + if suggestion.hasPrefix(infillPrefix) { + suggestion.removeFirst(infillPrefix.count) } + } - return suggestionPrefix + suggestion + /// Window-mapping the lines in suggestion and the suffix to remove the common suffix. + func removeSuffix(from suggestion: inout String, suffix: [String]) { + let suggestionLines = suggestion.breakLines(appendLineBreakToLastLine: true) + if let last = suggestionLines.last, let lastIndex = suffix.firstIndex(of: last) { + var i = lastIndex - 1 + var j = suggestionLines.endIndex - 2 + while i >= 0, j >= 0, suffix[i] == suggestionLines[j] { + i -= 1 + j -= 1 + } + if i < 0 { + let endIndex = max(j, 0) + suggestion = suggestionLines[...endIndex].joined() + } + } } /// Extract suggestions that is enclosed in tags. - func extractEnclosingSuggestion( + fileprivate func extractEnclosingSuggestion( from response: String, openingTag: String, closingTag: String @@ -78,7 +113,9 @@ struct DefaultRawSuggestionPostProcessingStrategy: RawSuggestionPostProcessingSt } /// If the response starts with markdown code block, we should remove it. - func removeLeadingAndTrailingMarkdownCodeBlockMark(from response: String) -> String { + fileprivate func removeLeadingAndTrailingMarkdownCodeBlockMark(from response: String) + -> String + { let removePrefixMarkdownCodeBlockMark = Parse(input: Substring.self) { Skip { "```" diff --git a/Core/Sources/SuggestionService/RawSuggestionPostProcessing/NoOpRawSuggestionPostProcessingStrategy.swift b/Core/Sources/SuggestionService/RawSuggestionPostProcessing/NoOpRawSuggestionPostProcessingStrategy.swift index 213aafd..c8b047d 100644 --- a/Core/Sources/SuggestionService/RawSuggestionPostProcessing/NoOpRawSuggestionPostProcessingStrategy.swift +++ b/Core/Sources/SuggestionService/RawSuggestionPostProcessing/NoOpRawSuggestionPostProcessingStrategy.swift @@ -1,8 +1,8 @@ import Foundation struct NoOpRawSuggestionPostProcessingStrategy: RawSuggestionPostProcessingStrategy { - func postProcessRawSuggestion(suggestionPrefix: String, suggestion: String) -> String { - suggestionPrefix + suggestion + func postProcess(rawSuggestion: String, infillPrefix: String, suffix: [String]) -> String { + removeTrailingNewlinesAndWhitespace(from: infillPrefix + rawSuggestion) } } diff --git a/Core/Sources/SuggestionService/Service.swift b/Core/Sources/SuggestionService/Service.swift index f987293..3243591 100644 --- a/Core/Sources/SuggestionService/Service.swift +++ b/Core/Sources/SuggestionService/Service.swift @@ -80,12 +80,10 @@ actor Service { .map { CodeSuggestion( id: UUID().uuidString, - text: Self.removeTrailingNewlinesAndWhitespace( - from: postProcessor.postProcessRawSuggestion( - suggestionPrefix: prompt - .suggestionPrefix.prependingValue, - suggestion: $0 - ) + text: postProcessor.postProcess( + rawSuggestion: $0, + infillPrefix: prompt.suggestionPrefix.prependingValue, + suffix: prompt.suffix ), position: request.cursorPosition, range: .init( @@ -173,13 +171,5 @@ actor Service { return (previousLines, nextLines) } - - static func removeTrailingNewlinesAndWhitespace(from string: String) -> String { - var text = string[...] - while let last = text.last, last.isNewline || last.isWhitespace { - text = text.dropLast(1) - } - return String(text) - } } From 4d33c983d7b554f47aeafae83090acb42f00cb34 Mon Sep 17 00:00:00 2001 From: Shx Guo Date: Fri, 1 Mar 2024 16:45:11 +0800 Subject: [PATCH 12/27] Add unit tests for DefaultRawSuggestionPostProcessingStrategy --- ...uggestionPostProcessingStrategyTests.swift | 272 ++++++++++++++++++ 1 file changed, 272 insertions(+) create mode 100644 Core/Tests/SuggestionServiceTests/DefaultRawSuggestionPostProcessingStrategyTests.swift diff --git a/Core/Tests/SuggestionServiceTests/DefaultRawSuggestionPostProcessingStrategyTests.swift b/Core/Tests/SuggestionServiceTests/DefaultRawSuggestionPostProcessingStrategyTests.swift new file mode 100644 index 0000000..eb256d0 --- /dev/null +++ b/Core/Tests/SuggestionServiceTests/DefaultRawSuggestionPostProcessingStrategyTests.swift @@ -0,0 +1,272 @@ +import Foundation +import XCTest + +@testable import SuggestionService + +class DefaultRawSuggestionPostProcessingStrategyTests: XCTestCase { + func test_whenSuggestionHasCodeTagAtTheFirstLine_shouldExtractCodeInside() { + let strategy = DefaultRawSuggestionPostProcessingStrategy( + openingCodeTag: "", + closingCodeTag: "" + ) + let result = strategy.extractSuggestion( + from: """ + suggestion + """ + ) + + XCTAssertEqual(result, "suggestion") + } + + func test_whenSuggestionHasCodeTagAtTheFirstLine_closingTagInOtherLines_shouldExtractCodeInside( + ) { + let strategy = DefaultRawSuggestionPostProcessingStrategy( + openingCodeTag: "", + closingCodeTag: "" + ) + let result = strategy.extractSuggestion( + from: """ + suggestion + yes + """ + ) + + XCTAssertEqual(result, "suggestion\nyes") + } + + func test_whenSuggestionHasCodeTag_butNoClosingTag_shouldExtractCodeAfterTheTag() { + let strategy = DefaultRawSuggestionPostProcessingStrategy( + openingCodeTag: "", + closingCodeTag: "" + ) + let result = strategy.extractSuggestion( + from: """ + suggestion + yes + """ + ) + + XCTAssertEqual(result, "suggestion\nyes") + } + + func test_whenOnlyLinebreaksOrSpacesBeforeOpeningTag_shouldExtractCodeInsideTheTags() { + let strategy = DefaultRawSuggestionPostProcessingStrategy( + openingCodeTag: "", + closingCodeTag: "" + ) + let result = strategy.extractSuggestion( + from: """ + + + suggestion + """ + ) + + XCTAssertEqual(result, "suggestion") + + let result2 = strategy.extractSuggestion( + from: """ + suggestion + """ + ) + + XCTAssertEqual(result2, "suggestion") + + let result3 = strategy.extractSuggestion( + from: """ + + + suggestion + """ + ) + + XCTAssertEqual(result3, "suggestion") + } + + func test_whenMultipleOpeningTagFound_shouldTreatTheNextOneAsClosing() { + let strategy = DefaultRawSuggestionPostProcessingStrategy( + openingCodeTag: "", + closingCodeTag: "" + ) + let result = strategy.extractSuggestion( + from: """ + suggestionhello + """ + ) + XCTAssertEqual(result, "suggestion") + } + + func test_whenMarkdownCodeBlockFound_shouldExtractCodeInside() { + let strategy = DefaultRawSuggestionPostProcessingStrategy( + openingCodeTag: "", + closingCodeTag: "" + ) + let result = strategy.extractSuggestion( + from: """ + ```language + suggestion + ``` + """ + ) + + XCTAssertEqual(result, "suggestion") + } + + func test_whenOnlyLinebreaksOrSpacesBeforeMarkdownCodeBlock_shouldExtractCodeInside() { + let strategy = DefaultRawSuggestionPostProcessingStrategy( + openingCodeTag: "", + closingCodeTag: "" + ) + let result = strategy.extractSuggestion( + from: """ + + + ``` + suggestion + ``` + """ + ) + + XCTAssertEqual(result, "suggestion") + + let result2 = strategy.extractSuggestion( + from: """ + ``` + suggestion + ``` + """ + ) + + XCTAssertEqual(result2, "suggestion") + + let result3 = strategy.extractSuggestion( + from: """ + + + ``` + suggestion + ``` + """ + ) + + XCTAssertEqual(result3, "suggestion") + } + + func test_whenMarkdownCodeBlockAndCodeTagFound_firstlyExtractCodeTag_thenCodeTag() { + let strategy = DefaultRawSuggestionPostProcessingStrategy( + openingCodeTag: "", + closingCodeTag: "" + ) + let result = strategy.extractSuggestion( + from: """ + ```language + suggestion + suggestion + ``` + """ + ) + XCTAssertEqual(result, "suggestion") + } + + func test_whenMarkdownCodeBlockAndCodeTagFound_butNoClosingTag_firstlyExtractCodeTag_thenCodeTag( + ) { + let strategy = DefaultRawSuggestionPostProcessingStrategy( + openingCodeTag: "", + closingCodeTag: "" + ) + let result = strategy.extractSuggestion( + from: """ + ```language + suggestion + suggestion + ``` + """ + ) + XCTAssertEqual(result, "suggestion\nsuggestion\n") + } + + func test_whenSuggestionHasTheSamePrefix_removeThePrefix() { + let strategy = DefaultRawSuggestionPostProcessingStrategy( + openingCodeTag: "", + closingCodeTag: "" + ) + let result = strategy.extractSuggestion( + from: "suggestion" + ) + + XCTAssertEqual(result, "suggestion") + } + + func test_whenSuggestionLooksLikeAMessage_parseItCorrectly() { + let strategy = DefaultRawSuggestionPostProcessingStrategy( + openingCodeTag: "", + closingCodeTag: "" + ) + let result = strategy.extractSuggestion( + from: """ + Here is the suggestion: + ```language + suggestion + ``` + """ + ) + + XCTAssertEqual(result, "suggestion") + } + + func test_whenSuggestionHasTheSamePrefix_inTags_removeThePrefix() { + let strategy = DefaultRawSuggestionPostProcessingStrategy( + openingCodeTag: "", + closingCodeTag: "" + ) + var suggestion = "prefix suggestion" + strategy.removePrefix(from: &suggestion, infillPrefix: "prefix") + + XCTAssertEqual(suggestion, " suggestion") + } + + func test_whenSuggestionHasTheSameSuffix_removeTheSuffix() { + let strategy = DefaultRawSuggestionPostProcessingStrategy( + openingCodeTag: "", + closingCodeTag: "" + ) + var suggestion = "suggestion\na\nb" + strategy.removeSuffix(from: &suggestion, suffix: [ + "a\n", + "b\n", + ]) + + XCTAssertEqual(suggestion, "suggestion\n") + + var suggestion2 = "suggestion\na\nb" + strategy.removeSuffix(from: &suggestion2, suffix: []) + + XCTAssertEqual(suggestion2, "suggestion\na\nb") + + var suggestion3 = "suggestion\na\nb" + strategy.removeSuffix(from: &suggestion3, suffix: ["b\n"]) + + XCTAssertEqual(suggestion3, "suggestion\na\n") + } + + func test_case_1() { + let strategy = DefaultRawSuggestionPostProcessingStrategy( + openingCodeTag: "", + closingCodeTag: "" + ) + let result = strategy.postProcess( + rawSuggestion: """ + ```language + prefix suggestion + a + b + c + ``` + """, + infillPrefix: "prefix", + suffix: ["a\n", "b\n", "c\n"] + ) + XCTAssertEqual(result, "prefix suggestion") + } +} + From 8a0de32d893c29542a30d1afd749c9be71187892 Mon Sep 17 00:00:00 2001 From: Shx Guo Date: Fri, 1 Mar 2024 16:55:12 +0800 Subject: [PATCH 13/27] Fix markdown code block handling --- ...tRawSuggestionPostProcessingStrategy.swift | 24 ++++++++++++++++++- ...uggestionPostProcessingStrategyTests.swift | 10 ++++---- 2 files changed, 28 insertions(+), 6 deletions(-) diff --git a/Core/Sources/SuggestionService/RawSuggestionPostProcessing/DefaultRawSuggestionPostProcessingStrategy.swift b/Core/Sources/SuggestionService/RawSuggestionPostProcessing/DefaultRawSuggestionPostProcessingStrategy.swift index bd8dc34..227183e 100644 --- a/Core/Sources/SuggestionService/RawSuggestionPostProcessing/DefaultRawSuggestionPostProcessingStrategy.swift +++ b/Core/Sources/SuggestionService/RawSuggestionPostProcessing/DefaultRawSuggestionPostProcessingStrategy.swift @@ -116,9 +116,31 @@ struct DefaultRawSuggestionPostProcessingStrategy: RawSuggestionPostProcessingSt fileprivate func removeLeadingAndTrailingMarkdownCodeBlockMark(from response: String) -> String { - let removePrefixMarkdownCodeBlockMark = Parse(input: Substring.self) { + let leadingMarkdownCodeBlockMarkParser = Parse(input: Substring.self) { Skip { + Many { + OneOf { + " " + "\n" + } + } "```" + } + } + + let messagePrefixingMarkdownCodeBlockMarkParser = Parse(input: Substring.self) { + Skip { + PrefixThrough(":") + "\n```" + } + } + + let removePrefixMarkdownCodeBlockMark = Parse(input: Substring.self) { + Skip { + OneOf { + leadingMarkdownCodeBlockMarkParser + messagePrefixingMarkdownCodeBlockMarkParser + } PrefixThrough("\n") } OneOf { diff --git a/Core/Tests/SuggestionServiceTests/DefaultRawSuggestionPostProcessingStrategyTests.swift b/Core/Tests/SuggestionServiceTests/DefaultRawSuggestionPostProcessingStrategyTests.swift index eb256d0..dc30c2e 100644 --- a/Core/Tests/SuggestionServiceTests/DefaultRawSuggestionPostProcessingStrategyTests.swift +++ b/Core/Tests/SuggestionServiceTests/DefaultRawSuggestionPostProcessingStrategyTests.swift @@ -109,7 +109,7 @@ class DefaultRawSuggestionPostProcessingStrategyTests: XCTestCase { """ ) - XCTAssertEqual(result, "suggestion") + XCTAssertEqual(result, "suggestion\n") } func test_whenOnlyLinebreaksOrSpacesBeforeMarkdownCodeBlock_shouldExtractCodeInside() { @@ -127,7 +127,7 @@ class DefaultRawSuggestionPostProcessingStrategyTests: XCTestCase { """ ) - XCTAssertEqual(result, "suggestion") + XCTAssertEqual(result, "suggestion\n") let result2 = strategy.extractSuggestion( from: """ @@ -137,7 +137,7 @@ class DefaultRawSuggestionPostProcessingStrategyTests: XCTestCase { """ ) - XCTAssertEqual(result2, "suggestion") + XCTAssertEqual(result2, "suggestion\n") let result3 = strategy.extractSuggestion( from: """ @@ -149,7 +149,7 @@ class DefaultRawSuggestionPostProcessingStrategyTests: XCTestCase { """ ) - XCTAssertEqual(result3, "suggestion") + XCTAssertEqual(result3, "suggestion\n") } func test_whenMarkdownCodeBlockAndCodeTagFound_firstlyExtractCodeTag_thenCodeTag() { @@ -211,7 +211,7 @@ class DefaultRawSuggestionPostProcessingStrategyTests: XCTestCase { """ ) - XCTAssertEqual(result, "suggestion") + XCTAssertEqual(result, "suggestion\n") } func test_whenSuggestionHasTheSamePrefix_inTags_removeThePrefix() { From 443d6b02db0223d2c88732d16a4d90e9d994ebd8 Mon Sep 17 00:00:00 2001 From: Shx Guo Date: Fri, 1 Mar 2024 16:59:32 +0800 Subject: [PATCH 14/27] Remove tests --- ...uggestionPostProcessingStrategyTests.swift | 34 ------------------- 1 file changed, 34 deletions(-) diff --git a/Core/Tests/SuggestionServiceTests/DefaultRawSuggestionPostProcessingStrategyTests.swift b/Core/Tests/SuggestionServiceTests/DefaultRawSuggestionPostProcessingStrategyTests.swift index dc30c2e..f8e92d2 100644 --- a/Core/Tests/SuggestionServiceTests/DefaultRawSuggestionPostProcessingStrategyTests.swift +++ b/Core/Tests/SuggestionServiceTests/DefaultRawSuggestionPostProcessingStrategyTests.swift @@ -49,40 +49,6 @@ class DefaultRawSuggestionPostProcessingStrategyTests: XCTestCase { XCTAssertEqual(result, "suggestion\nyes") } - func test_whenOnlyLinebreaksOrSpacesBeforeOpeningTag_shouldExtractCodeInsideTheTags() { - let strategy = DefaultRawSuggestionPostProcessingStrategy( - openingCodeTag: "", - closingCodeTag: "" - ) - let result = strategy.extractSuggestion( - from: """ - - - suggestion - """ - ) - - XCTAssertEqual(result, "suggestion") - - let result2 = strategy.extractSuggestion( - from: """ - suggestion - """ - ) - - XCTAssertEqual(result2, "suggestion") - - let result3 = strategy.extractSuggestion( - from: """ - - - suggestion - """ - ) - - XCTAssertEqual(result3, "suggestion") - } - func test_whenMultipleOpeningTagFound_shouldTreatTheNextOneAsClosing() { let strategy = DefaultRawSuggestionPostProcessingStrategy( openingCodeTag: "", From 4a7a6a7731fa17dc07f3b8b8ee8ce115d7b98a71 Mon Sep 17 00:00:00 2001 From: Shx Guo Date: Fri, 1 Mar 2024 17:30:02 +0800 Subject: [PATCH 15/27] Add CodeLlamaFillInTheMiddleRequestStrategy --- ...eLlamaFillInTheMiddleRequestStrategy.swift | 78 +++++++++++++++++++ 1 file changed, 78 insertions(+) create mode 100644 Core/Sources/SuggestionService/RequestStrategies/CodeLlamaFillInTheMiddleRequestStrategy.swift diff --git a/Core/Sources/SuggestionService/RequestStrategies/CodeLlamaFillInTheMiddleRequestStrategy.swift b/Core/Sources/SuggestionService/RequestStrategies/CodeLlamaFillInTheMiddleRequestStrategy.swift new file mode 100644 index 0000000..4cdc549 --- /dev/null +++ b/Core/Sources/SuggestionService/RequestStrategies/CodeLlamaFillInTheMiddleRequestStrategy.swift @@ -0,0 +1,78 @@ +import CopilotForXcodeKit +import Foundation +import Fundamental + +/// https://ollama.com/library/codellama +struct CodeLlamaFillInTheMiddleRequestStrategy: RequestStrategy { + var sourceRequest: SuggestionRequest + var prefix: [String] + var suffix: [String] + + var shouldSkip: Bool { + prefix.last?.trimmingCharacters(in: .whitespaces) == "}" + } + + func createPrompt() -> Prompt { + Prompt( + sourceRequest: sourceRequest, + prefix: prefix, + suffix: suffix + ) + } + + func createRawSuggestionPostProcessor() -> DefaultRawSuggestionPostProcessingStrategy { + DefaultRawSuggestionPostProcessingStrategy(openingCodeTag: "", closingCodeTag: "") + } + + enum Tag { + public static let prefix = "
"
+        public static let suffix = ""
+        public static let middle = ""
+    }
+
+    struct Prompt: PromptStrategy {
+        let systemPrompt: String = """
+        You are a senior programer who take the surrounding code and \
+        references from the codebase into account in order to write high-quality code to \
+        complete the code enclosed in the given code. \
+        You only respond with code that works and fits seamlessly with surrounding code. \
+        Don't include anything else beyond the code. \
+        The prefix will follow the PRE tag and the suffix will follow the SUF tag.
+        """
+        var sourceRequest: SuggestionRequest
+        var prefix: [String]
+        var suffix: [String]
+        var filePath: String { sourceRequest.relativePath ?? sourceRequest.fileURL.path }
+        var relevantCodeSnippets: [RelevantCodeSnippet] { sourceRequest.relevantCodeSnippets }
+        var stopWords: [String] { ["\n\n"] }
+        var language: CodeLanguage? { sourceRequest.language }
+
+        var suggestionPrefix: SuggestionPrefix {
+            guard let prefix = prefix.last else { return .empty }
+            return .unchanged(prefix).curlyBracesLineBreak()
+        }
+
+        func createPrompt(
+            truncatedPrefix: [String],
+            truncatedSuffix: [String],
+            includedSnippets: [RelevantCodeSnippet]
+        ) -> [PromptMessage] {
+            return [
+                .init(
+                    role: .user,
+                    content: """
+                    \(Tag.prefix) // File Path: \(filePath)
+                    // Indentation: \
+                    \(sourceRequest.indentSize) \
+                    \(sourceRequest.usesTabsForIndentation ? "tab" : "space")
+                    \(includedSnippets.map(\.content).joined(separator: "\n\n"))
+                    \(truncatedPrefix.joined()) \
+                    \(Tag.suffix)\(truncatedSuffix.joined()) \
+                    \(Tag.middle)
+                    """.trimmingCharacters(in: .whitespacesAndNewlines)
+                ),
+            ]
+        }
+    }
+}
+

From 74e5fad3a9bce7bb12bfcf361c45a3bc55c8d5b1 Mon Sep 17 00:00:00 2001
From: Shx Guo 
Date: Fri, 1 Mar 2024 17:31:43 +0800
Subject: [PATCH 16/27] Trim newlines and whitespaces in prompt

---
 Core/Sources/CodeCompletionService/OllamaService.swift    | 1 +
 .../DefaultRawSuggestionPostProcessingStrategy.swift      | 4 ++++
 .../RequestStrategies/ContinueRequestStrategy.swift       | 8 ++++----
 .../RequestStrategies/DefaultRequestStrategy.swift        | 4 ++--
 .../RequestStrategies/NaiveRequestStrategy.swift          | 2 +-
 5 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/Core/Sources/CodeCompletionService/OllamaService.swift b/Core/Sources/CodeCompletionService/OllamaService.swift
index 279e2d7..151524c 100644
--- a/Core/Sources/CodeCompletionService/OllamaService.swift
+++ b/Core/Sources/CodeCompletionService/OllamaService.swift
@@ -221,6 +221,7 @@ extension OllamaService {
         ))
         let prompts = strategy.createTruncatedPrompt(promptStrategy: request)
         return ([request.systemPrompt] + prompts.map(\.content)).joined(separator: "\n\n")
+            .trimmingCharacters(in: .whitespacesAndNewlines)
     }
 
     func sendPrompt(_ prompt: String) async throws -> ResponseStream {
diff --git a/Core/Sources/SuggestionService/RawSuggestionPostProcessing/DefaultRawSuggestionPostProcessingStrategy.swift b/Core/Sources/SuggestionService/RawSuggestionPostProcessing/DefaultRawSuggestionPostProcessingStrategy.swift
index 227183e..e7bbfd3 100644
--- a/Core/Sources/SuggestionService/RawSuggestionPostProcessing/DefaultRawSuggestionPostProcessingStrategy.swift
+++ b/Core/Sources/SuggestionService/RawSuggestionPostProcessing/DefaultRawSuggestionPostProcessingStrategy.swift
@@ -66,6 +66,10 @@ struct DefaultRawSuggestionPostProcessingStrategy: RawSuggestionPostProcessingSt
         openingTag: String,
         closingTag: String
     ) -> String {
+        guard !openingTag.isEmpty, !closingTag.isEmpty else {
+            return response
+        }
+        
         let case_openingTagAtTheStart_parseEverythingInsideTheTag = Parse(input: Substring.self) {
             openingTag
 
diff --git a/Core/Sources/SuggestionService/RequestStrategies/ContinueRequestStrategy.swift b/Core/Sources/SuggestionService/RequestStrategies/ContinueRequestStrategy.swift
index b53eee7..c136da8 100644
--- a/Core/Sources/SuggestionService/RequestStrategies/ContinueRequestStrategy.swift
+++ b/Core/Sources/SuggestionService/RequestStrategies/ContinueRequestStrategy.swift
@@ -54,7 +54,7 @@ struct ContinueRequestStrategy: RequestStrategy {
         ```
         o World)
         ```
-        """
+        """.trimmingCharacters(in: .whitespacesAndNewlines)
         var sourceRequest: SuggestionRequest
         var prefix: [String]
         var suffix: [String]
@@ -113,18 +113,18 @@ struct ContinueRequestStrategy: RequestStrategy {
             ```
 
             Complete code inside \(Tag.openingCode)
-            """)
+            """.trimmingCharacters(in: .whitespacesAndNewlines))
 
             let mockResponse = PromptMessage(role: .assistant, content: """
             \(Tag.openingCode)\(infillBlock)
-            """)
+            """.trimmingCharacters(in: .whitespacesAndNewlines))
 
             let continuePrompt = PromptMessage(role: .user, content: """
             Continue generating. \
             Don't duplicate existing implementations. \
             Don't try to fix what was written. \
             Don't worry about typos.
-            """)
+            """.trimmingCharacters(in: .whitespacesAndNewlines))
 
             return [
                 initialPrompt,
diff --git a/Core/Sources/SuggestionService/RequestStrategies/DefaultRequestStrategy.swift b/Core/Sources/SuggestionService/RequestStrategies/DefaultRequestStrategy.swift
index cab2ab5..448bf20 100644
--- a/Core/Sources/SuggestionService/RequestStrategies/DefaultRequestStrategy.swift
+++ b/Core/Sources/SuggestionService/RequestStrategies/DefaultRequestStrategy.swift
@@ -57,7 +57,7 @@ struct DefaultRequestStrategy: RequestStrategy {
         ###
          World")\(Tag.closingCode)
         ###
-        """
+        """.trimmingCharacters(in: .whitespacesAndNewlines)
         var sourceRequest: SuggestionRequest
         var prefix: [String]
         var suffix: [String]
@@ -114,7 +114,7 @@ struct DefaultRequestStrategy: RequestStrategy {
             Complete code inside \(Tag.openingCode):
 
             \(Tag.openingCode)\(infillBlock)
-            """
+            """.trimmingCharacters(in: .whitespacesAndNewlines)
         }
 
         static func createSnippetsPrompt(includedSnippets: [RelevantCodeSnippet]) -> String {
diff --git a/Core/Sources/SuggestionService/RequestStrategies/NaiveRequestStrategy.swift b/Core/Sources/SuggestionService/RequestStrategies/NaiveRequestStrategy.swift
index 1274c37..3d21111 100644
--- a/Core/Sources/SuggestionService/RequestStrategies/NaiveRequestStrategy.swift
+++ b/Core/Sources/SuggestionService/RequestStrategies/NaiveRequestStrategy.swift
@@ -75,7 +75,7 @@ struct NaiveRequestStrategy: RequestStrategy {
             ---
             
             \(code)
-            """)]
+            """.trimmingCharacters(in: .whitespacesAndNewlines))]
         }
     }
 }

From 2769c138ba1164f87aeefe497f0a78903a572933 Mon Sep 17 00:00:00 2001
From: Shx Guo 
Date: Fri, 1 Mar 2024 18:01:07 +0800
Subject: [PATCH 17/27] Adjust CodeLlama fitm strategy

---
 ...eLlamaFillInTheMiddleRequestStrategy.swift | 41 ++++++++++++++-----
 .../SuggestionService/RequestStrategy.swift   |  6 +++
 CustomSuggestionService/ContentView.swift     |  6 +++
 3 files changed, 43 insertions(+), 10 deletions(-)

diff --git a/Core/Sources/SuggestionService/RequestStrategies/CodeLlamaFillInTheMiddleRequestStrategy.swift b/Core/Sources/SuggestionService/RequestStrategies/CodeLlamaFillInTheMiddleRequestStrategy.swift
index 4cdc549..40799bb 100644
--- a/Core/Sources/SuggestionService/RequestStrategies/CodeLlamaFillInTheMiddleRequestStrategy.swift
+++ b/Core/Sources/SuggestionService/RequestStrategies/CodeLlamaFillInTheMiddleRequestStrategy.swift
@@ -31,20 +31,13 @@ struct CodeLlamaFillInTheMiddleRequestStrategy: RequestStrategy {
     }
 
     struct Prompt: PromptStrategy {
-        let systemPrompt: String = """
-        You are a senior programer who take the surrounding code and \
-        references from the codebase into account in order to write high-quality code to \
-        complete the code enclosed in the given code. \
-        You only respond with code that works and fits seamlessly with surrounding code. \
-        Don't include anything else beyond the code. \
-        The prefix will follow the PRE tag and the suffix will follow the SUF tag.
-        """
+        fileprivate(set) var systemPrompt: String = ""
         var sourceRequest: SuggestionRequest
         var prefix: [String]
         var suffix: [String]
         var filePath: String { sourceRequest.relativePath ?? sourceRequest.fileURL.path }
         var relevantCodeSnippets: [RelevantCodeSnippet] { sourceRequest.relevantCodeSnippets }
-        var stopWords: [String] { ["\n\n"] }
+        var stopWords: [String] { ["\n\n", ""] }
         var language: CodeLanguage? { sourceRequest.language }
 
         var suggestionPrefix: SuggestionPrefix {
@@ -57,6 +50,7 @@ struct CodeLlamaFillInTheMiddleRequestStrategy: RequestStrategy {
             truncatedSuffix: [String],
             includedSnippets: [RelevantCodeSnippet]
         ) -> [PromptMessage] {
+            let suffix = truncatedSuffix.joined()
             return [
                 .init(
                     role: .user,
@@ -67,7 +61,7 @@ struct CodeLlamaFillInTheMiddleRequestStrategy: RequestStrategy {
                     \(sourceRequest.usesTabsForIndentation ? "tab" : "space")
                     \(includedSnippets.map(\.content).joined(separator: "\n\n"))
                     \(truncatedPrefix.joined()) \
-                    \(Tag.suffix)\(truncatedSuffix.joined()) \
+                    \(Tag.suffix)\(suffix.isEmpty ? "\n// End of file" : suffix) \
                     \(Tag.middle)
                     """.trimmingCharacters(in: .whitespacesAndNewlines)
                 ),
@@ -76,3 +70,30 @@ struct CodeLlamaFillInTheMiddleRequestStrategy: RequestStrategy {
     }
 }
 
+struct CodeLlamaFillInTheMiddleWithSystemPromptRequestStrategy: RequestStrategy {
+    let strategy: CodeLlamaFillInTheMiddleRequestStrategy
+
+    init(sourceRequest: SuggestionRequest, prefix: [String], suffix: [String]) {
+        strategy = .init(sourceRequest: sourceRequest, prefix: prefix, suffix: suffix)
+    }
+
+    func createPrompt() -> some PromptStrategy {
+        var prompt = strategy.createPrompt()
+        prompt.systemPrompt = """
+        You are a senior programer who take the surrounding code and \
+        references from the codebase into account in order to write high-quality code to \
+        complete the code enclosed in the given code. \
+        You only respond with code that works and fits seamlessly with surrounding code. \
+        Don't include anything else beyond the code. \
+        The prefix will follow the PRE tag and the suffix will follow the SUF tag. \
+        You should write the code that fits seamlessly after the MID tag.
+        """.trimmingCharacters(in: .whitespacesAndNewlines)
+
+        return prompt
+    }
+
+    func createRawSuggestionPostProcessor() -> some RawSuggestionPostProcessingStrategy {
+        strategy.createRawSuggestionPostProcessor()
+    }
+}
+
diff --git a/Core/Sources/SuggestionService/RequestStrategy.swift b/Core/Sources/SuggestionService/RequestStrategy.swift
index 5be3d00..5f6060a 100644
--- a/Core/Sources/SuggestionService/RequestStrategy.swift
+++ b/Core/Sources/SuggestionService/RequestStrategy.swift
@@ -26,6 +26,8 @@ public enum RequestStrategyOption: String, CaseIterable, Codable {
     case `default` = ""
     case naive
     case `continue`
+    case codeLlamaFillInTheMiddle
+    case codeLlamaFillInTheMiddleWithSystemPrompt
 }
 
 extension RequestStrategyOption {
@@ -37,6 +39,10 @@ extension RequestStrategyOption {
             return NaiveRequestStrategy.self
         case .continue:
             return ContinueRequestStrategy.self
+        case .codeLlamaFillInTheMiddle:
+            return CodeLlamaFillInTheMiddleRequestStrategy.self
+        case .codeLlamaFillInTheMiddleWithSystemPrompt:
+            return CodeLlamaFillInTheMiddleWithSystemPromptRequestStrategy.self
         }
     }
 }
diff --git a/CustomSuggestionService/ContentView.swift b/CustomSuggestionService/ContentView.swift
index 7de29df..185b0ca 100644
--- a/CustomSuggestionService/ContentView.swift
+++ b/CustomSuggestionService/ContentView.swift
@@ -177,6 +177,12 @@ struct RequestStrategyPicker: View {
                         Text("Naive").tag(option.rawValue)
                     case .continue:
                         Text("Continue").tag(option.rawValue)
+                    case .codeLlamaFillInTheMiddle:
+                        Text("CodeLlama Fill-in-the-Middle (Good for Codellama:xb-code)")
+                            .tag(option.rawValue)
+                    case .codeLlamaFillInTheMiddleWithSystemPrompt:
+                        Text("CodeLlama Fill-in-the-Middle with System Prompt")
+                            .tag(option.rawValue)
                     }
                 }
             }

From e5168f9ae3ba74c3bfabf1c06816aaa5ca536802 Mon Sep 17 00:00:00 2001
From: Shx Guo 
Date: Fri, 1 Mar 2024 18:11:08 +0800
Subject: [PATCH 18/27] Update README.md

---
 CustomSuggestionService/ContentView.swift |  6 ++++--
 README.md                                 | 26 +++++++++++++++--------
 2 files changed, 21 insertions(+), 11 deletions(-)

diff --git a/CustomSuggestionService/ContentView.swift b/CustomSuggestionService/ContentView.swift
index 185b0ca..fa1fe18 100644
--- a/CustomSuggestionService/ContentView.swift
+++ b/CustomSuggestionService/ContentView.swift
@@ -178,8 +178,10 @@ struct RequestStrategyPicker: View {
                     case .continue:
                         Text("Continue").tag(option.rawValue)
                     case .codeLlamaFillInTheMiddle:
-                        Text("CodeLlama Fill-in-the-Middle (Good for Codellama:xb-code)")
-                            .tag(option.rawValue)
+                        Text(
+                            "CodeLlama Fill-in-the-Middle (Good for Codellama:xb-code and other models with Fill-in-the-Middle support)"
+                        )
+                        .tag(option.rawValue)
                     case .codeLlamaFillInTheMiddleWithSystemPrompt:
                         Text("CodeLlama Fill-in-the-Middle with System Prompt")
                             .tag(option.rawValue)
diff --git a/README.md b/README.md
index b8804c2..d6ec841 100644
--- a/README.md
+++ b/README.md
@@ -1,4 +1,4 @@
-#  Custom Suggestion Service for Copilot for Xcode
+# Custom Suggestion Service for Copilot for Xcode
 
 This extension offers a custom suggestion service for [Copilot for Xcode](https://github.com/intitni/CopilotForXcode), allowing you to leverage a chat model to enhance the suggestions provided as you write code.
 
@@ -27,31 +27,39 @@ The app supports three types of suggestion services:
 - Models with completions API
 - [Tabby](https://tabby.tabbyml.com)
 
-It is recommended to use Tabby since they have extensive experience in crafting prompts.
+If you are new to running a model locally, you can try [LM Studio](https://lmstudio.ai).
+
+### Recommended Settings
+
+- Use Tabby since they have extensive experience in code completion.
+- Use models with completions API with Fill-in-the-Middle support (for example, codellama:7b-code), and use the "Codellama Fill-in-the-Middle" strategy.
+
+### Others
 
-If you choose not to use Tabby, it is advisable to use a custom model with the completions API and employ the default request strategy.
+In other situations, it is advisable to use a custom model with the completions API over a chat completions API, and employ the default request strategy.
 
 Ensure that the prompt format remains as simple as the following:
 
-``` 
+```
 {System}
 {User}
 {Assistant}
 ```
 
-If you are new to running a model locally, you can try [LM Studio](https://lmstudio.ai).
-
 ## Strategies
 
 - Default: This strategy meticulously explains the context to the model, prompting it to generate a suggestion.
 - Naive: This strategy rearranges the code in a naive way to trick the model into believing it's appending code at the end of a file.
 - Continue: This strategy employs the "Please Continue" technique to persuade the model that it has started a suggestion and must continue to complete it. (Only effective with the chat completion API).
+- CodeLlama Fill-in-the-Middle: It uses special tokens to guide the models to generate suggestions. The models need to support FIM to use it (codellama:xb-code, startcoder, etc.). This strategy uses the special tokens documented by CodeLlama.
+- CodeLlama Fill-in-the-Middle with System Prompt: The previous one doesn't have a system prompt telling it what to do. You can try to use it in models that don't support FIM.
 
 ## Contribution
 
-Prompt engineering is a challenging task, and your assistance is invaluable. 
+Prompt engineering is a challenging task, and your assistance is invaluable.
 
-The most complex things are located within the `Core` package. 
+The most complex things are located within the `Core` package.
 
-- To add a new service, please refer to the `CodeCompletionService` folder. 
+- To add a new service, please refer to the `CodeCompletionService` folder.
 - To add new request strategies, check out the `SuggestionService` folder.
+

From 1356a58515c6f3b18d3eeadcf2029fb492942329 Mon Sep 17 00:00:00 2001
From: Shx Guo 
Date: Fri, 1 Mar 2024 18:13:58 +0800
Subject: [PATCH 19/27] Use no-op post processing strategy instead in CodeLlama
 FIM strategy

---
 .../CodeLlamaFillInTheMiddleRequestStrategy.swift           | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/Core/Sources/SuggestionService/RequestStrategies/CodeLlamaFillInTheMiddleRequestStrategy.swift b/Core/Sources/SuggestionService/RequestStrategies/CodeLlamaFillInTheMiddleRequestStrategy.swift
index 40799bb..72bd498 100644
--- a/Core/Sources/SuggestionService/RequestStrategies/CodeLlamaFillInTheMiddleRequestStrategy.swift
+++ b/Core/Sources/SuggestionService/RequestStrategies/CodeLlamaFillInTheMiddleRequestStrategy.swift
@@ -20,8 +20,8 @@ struct CodeLlamaFillInTheMiddleRequestStrategy: RequestStrategy {
         )
     }
 
-    func createRawSuggestionPostProcessor() -> DefaultRawSuggestionPostProcessingStrategy {
-        DefaultRawSuggestionPostProcessingStrategy(openingCodeTag: "", closingCodeTag: "")
+    func createRawSuggestionPostProcessor() -> NoOpRawSuggestionPostProcessingStrategy {
+        NoOpRawSuggestionPostProcessingStrategy()
     }
 
     enum Tag {
@@ -93,7 +93,7 @@ struct CodeLlamaFillInTheMiddleWithSystemPromptRequestStrategy: RequestStrategy
     }
 
     func createRawSuggestionPostProcessor() -> some RawSuggestionPostProcessingStrategy {
-        strategy.createRawSuggestionPostProcessor()
+        DefaultRawSuggestionPostProcessingStrategy(openingCodeTag: "", closingCodeTag: "")
     }
 }
 

From 9df752a2d4bd40970c5554fd85c45f1a5a2340e3 Mon Sep 17 00:00:00 2001
From: Shx Guo 
Date: Fri, 1 Mar 2024 21:24:15 +0800
Subject: [PATCH 20/27] Add TabbyStrategy

---
 .../CodeCompletionService/TabbyService.swift  |  6 +-
 .../TabbyRequestStrategy.swift                | 62 +++++++++++++++++++
 2 files changed, 64 insertions(+), 4 deletions(-)
 create mode 100644 Core/Sources/SuggestionService/RequestStrategies/TabbyRequestStrategy.swift

diff --git a/Core/Sources/CodeCompletionService/TabbyService.swift b/Core/Sources/CodeCompletionService/TabbyService.swift
index c27c15e..8ed28cf 100644
--- a/Core/Sources/CodeCompletionService/TabbyService.swift
+++ b/Core/Sources/CodeCompletionService/TabbyService.swift
@@ -30,13 +30,12 @@ extension TabbyService: CodeCompletionServiceType {
     func getCompletion(_ request: PromptStrategy) async throws -> AsyncStream {
         let prefix = request.prefix.joined()
         let suffix = request.suffix.joined()
-        let clipboard = request.relevantCodeSnippets.map(\.content).joined(separator: "\n\n")
         let requestBody = RequestBody(
             language: request.language?.rawValue,
             segments: .init(
-                prefix: clipboard + "\n\n" + prefix,
+                prefix: prefix,
                 suffix: suffix,
-                clipboard: clipboard // it's seems to be ignored by Tabby
+                clipboard: ""
             ),
             temperature: temperature,
             seed: nil
@@ -44,7 +43,6 @@ extension TabbyService: CodeCompletionServiceType {
         CodeCompletionLogger.logger.logPrompt([
             (prefix, "prefix"),
             (suffix, "suffix"),
-            (clipboard, "clipboard"),
         ])
         return AsyncStream { continuation in
             let task = Task {
diff --git a/Core/Sources/SuggestionService/RequestStrategies/TabbyRequestStrategy.swift b/Core/Sources/SuggestionService/RequestStrategies/TabbyRequestStrategy.swift
new file mode 100644
index 0000000..0570779
--- /dev/null
+++ b/Core/Sources/SuggestionService/RequestStrategies/TabbyRequestStrategy.swift
@@ -0,0 +1,62 @@
+import CopilotForXcodeKit
+import Foundation
+import Fundamental
+
+/// A special strategy for Tabby.
+struct TabbyRequestStrategy: RequestStrategy {
+    var sourceRequest: SuggestionRequest
+    var prefix: [String]
+    var suffix: [String]
+
+    var shouldSkip: Bool {
+        prefix.last?.trimmingCharacters(in: .whitespaces) == "}"
+    }
+
+    func createPrompt() -> Prompt {
+        Prompt(
+            sourceRequest: sourceRequest,
+            prefix: prefix,
+            suffix: suffix
+        )
+    }
+
+    func createRawSuggestionPostProcessor() -> some RawSuggestionPostProcessingStrategy {
+        NoOpRawSuggestionPostProcessingStrategy()
+    }
+
+    struct Prompt: PromptStrategy {
+        let systemPrompt: String = ""
+        var sourceRequest: SuggestionRequest
+        var prefix: [String]
+        var suffix: [String]
+        var filePath: String { sourceRequest.relativePath ?? sourceRequest.fileURL.path }
+        var relevantCodeSnippets: [RelevantCodeSnippet] { sourceRequest.relevantCodeSnippets }
+        var stopWords: [String] { [] }
+        var language: CodeLanguage? { sourceRequest.language }
+
+        var suggestionPrefix: SuggestionPrefix {
+            guard let prefix = prefix.last else { return .empty }
+            return .unchanged(prefix)
+        }
+
+        init(sourceRequest: SuggestionRequest, prefix: [String], suffix: [String]) {
+            self.sourceRequest = sourceRequest
+
+            let prefix = sourceRequest.relevantCodeSnippets.map { $0.content + "\n\n" }
+                + prefix
+
+            self.prefix = prefix
+            self.suffix = suffix
+        }
+
+        /// Not used by ``TabbyService``.
+        func createPrompt(
+            truncatedPrefix: [String],
+            truncatedSuffix: [String],
+            includedSnippets: [RelevantCodeSnippet]
+        ) -> [PromptMessage] {
+            []
+        }
+    }
+}
+

From af9419f920cf9559dd38f5c09b68f5091aeab1a7 Mon Sep 17 00:00:00 2001
From: Shx Guo 
Date: Fri, 1 Mar 2024 21:24:37 +0800
Subject: [PATCH 21/27] Use TabbyStrategy when using tabby service

---
 Core/Sources/SuggestionService/Service.swift | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/Core/Sources/SuggestionService/Service.swift b/Core/Sources/SuggestionService/Service.swift
index 3243591..9e37ae0 100644
--- a/Core/Sources/SuggestionService/Service.swift
+++ b/Core/Sources/SuggestionService/Service.swift
@@ -128,6 +128,15 @@ actor Service {
         suffix: [String]
     ) -> any RequestStrategy {
         let id = UserDefaults.shared.value(for: \.requestStrategyId)
+        if let type = CustomModelType(rawValue: UserDefaults.shared.value(for: \.chatModelId)),
+           type == .tabby
+        {
+            return TabbyRequestStrategy(
+                sourceRequest: sourceRequest,
+                prefix: prefix,
+                suffix: suffix
+            )
+        }
         let strategyOption = RequestStrategyOption(rawValue: id) ?? .default
         return strategyOption.strategy.init(
             sourceRequest: sourceRequest,

From 16bd70143ac43a9100453fd2f2869580d202d289 Mon Sep 17 00:00:00 2001
From: Shx Guo 
Date: Fri, 1 Mar 2024 21:29:12 +0800
Subject: [PATCH 22/27] Update README.md

---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index d6ec841..402a696 100644
--- a/README.md
+++ b/README.md
@@ -27,7 +27,7 @@ The app supports three types of suggestion services:
 - Models with completions API
 - [Tabby](https://tabby.tabbyml.com)
 
-If you are new to running a model locally, you can try [LM Studio](https://lmstudio.ai).
+If you are new to running a model locally, you can try [Ollama](https://ollama.com) and [LM Studio](https://lmstudio.ai).
 
 ### Recommended Settings
 

From 9ff813c3c877c5413d0a6e68e1b14dccd0e35bed Mon Sep 17 00:00:00 2001
From: Shx Guo 
Date: Mon, 4 Mar 2024 00:52:10 +0800
Subject: [PATCH 23/27] Bump CopilotForXcodeKit to 0.5.0

---
 Core/Package.swift                                          | 2 +-
 Core/Sources/SuggestionService/SuggestionService.swift      | 6 +++++-
 .../xcshareddata/swiftpm/Package.resolved                   | 4 ++--
 CustomSuggestionService/Dependency.swift                    | 6 +++++-
 CustomSuggestionService/TestField/TestField.swift           | 6 +++++-
 5 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/Core/Package.swift b/Core/Package.swift
index e73697a..184ee54 100644
--- a/Core/Package.swift
+++ b/Core/Package.swift
@@ -15,7 +15,7 @@ let package = Package(
     dependencies: [
         .package(
             url: "https://github.com/intitni/CopilotForXcodeKit",
-            from: "0.4.0"
+            from: "0.5.0"
         ),
         .package(
             url: "https://github.com/pointfreeco/swift-dependencies",
diff --git a/Core/Sources/SuggestionService/SuggestionService.swift b/Core/Sources/SuggestionService/SuggestionService.swift
index dc36b5a..c5e882c 100644
--- a/Core/Sources/SuggestionService/SuggestionService.swift
+++ b/Core/Sources/SuggestionService/SuggestionService.swift
@@ -8,7 +8,11 @@ public class SuggestionService: SuggestionServiceType {
     public init() {}
 
     public var configuration: SuggestionServiceConfiguration {
-        .init(acceptsRelevantCodeSnippets: true, mixRelevantCodeSnippetsInSource: false)
+        .init(
+            acceptsRelevantCodeSnippets: true,
+            mixRelevantCodeSnippetsInSource: false,
+            acceptsRelevantSnippetsFromOpenedFiles: true
+        )
     }
 
     public func notifyAccepted(_ suggestion: CodeSuggestion, workspace: WorkspaceInfo) async {}
diff --git a/CustomSuggestionService.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved b/CustomSuggestionService.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved
index 5f0e3c5..97fcb52 100644
--- a/CustomSuggestionService.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved
+++ b/CustomSuggestionService.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved
@@ -23,8 +23,8 @@
       "kind" : "remoteSourceControl",
       "location" : "https://github.com/intitni/CopilotForXcodeKit.git",
       "state" : {
-        "revision" : "91b252bd30691a5e02d600ba175eb4e69e4d4dfc",
-        "version" : "0.4.0"
+        "revision" : "05d5992a1dca572cf16019b23dfb4fca3be2109b",
+        "version" : "0.5.0"
       }
     },
     {
diff --git a/CustomSuggestionService/Dependency.swift b/CustomSuggestionService/Dependency.swift
index 9eedea5..94e836e 100644
--- a/CustomSuggestionService/Dependency.swift
+++ b/CustomSuggestionService/Dependency.swift
@@ -11,7 +11,11 @@ struct SuggestionServiceDependencyKey: DependencyKey {
 
 struct MockSuggestionService: SuggestionServiceType {
     var configuration: SuggestionServiceConfiguration {
-        .init(acceptsRelevantCodeSnippets: true, mixRelevantCodeSnippetsInSource: false)
+        .init(
+            acceptsRelevantCodeSnippets: true,
+            mixRelevantCodeSnippetsInSource: false,
+            acceptsRelevantSnippetsFromOpenedFiles: true
+        )
     }
 
     func getSuggestions(
diff --git a/CustomSuggestionService/TestField/TestField.swift b/CustomSuggestionService/TestField/TestField.swift
index 81b64ff..7cbf7a5 100644
--- a/CustomSuggestionService/TestField/TestField.swift
+++ b/CustomSuggestionService/TestField/TestField.swift
@@ -106,7 +106,11 @@ struct TestField {
                                 indentSize: 4,
                                 usesTabsForIndentation: false,
                                 relevantCodeSnippets: [
-                                    .init(content: relevantCodeSnippet, priority: 999),
+                                    .init(
+                                        content: relevantCodeSnippet,
+                                        priority: 999,
+                                        filePath: ""
+                                    ),
                                 ]
                             ),
                             workspace: workspace

From 9ac9515a9cdc7415be42b3acc644d69451c135d2 Mon Sep 17 00:00:00 2001
From: Shx Guo 
Date: Mon, 4 Mar 2024 00:52:27 +0800
Subject: [PATCH 24/27] Use Context Window instead of Max Tokens

---
 .../ChatModelManagement/ChatModelEditView.swift             | 6 +++---
 .../ChatModelManagement/CompletionModelEditView.swift       | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/CustomSuggestionService/ChatModelManagement/ChatModelEditView.swift b/CustomSuggestionService/ChatModelManagement/ChatModelEditView.swift
index a8134fe..53d7713 100644
--- a/CustomSuggestionService/ChatModelManagement/ChatModelEditView.swift
+++ b/CustomSuggestionService/ChatModelManagement/ChatModelEditView.swift
@@ -128,7 +128,7 @@ struct ChatModelEditView: View {
             )
 
             TextField(text: textFieldBinding) {
-                Text("Max Tokens (Including Reply)")
+                Text("Context Window")
                     .multilineTextAlignment(.trailing)
             }
             .overlay(alignment: .trailing) {
@@ -269,7 +269,7 @@ struct ChatModelEditView: View {
     var ollama: some View {
         baseURLTextField(
             title: "",
-            prompt: Text("https://127.0.0.1:11434/api/chat")
+            prompt: Text("http://127.0.0.1:11434")
         ) {
             Text("/api/chat")
         }
@@ -284,7 +284,7 @@ struct ChatModelEditView: View {
         
         VStack(alignment: .leading, spacing: 8) {
             Text(Image(systemName: "exclamationmark.triangle.fill")) + Text(
-                " For more details, please visit [https://ollama.com](https://ollama.com)"
+                " For more details, please visit [https://ollama.com](https://ollama.com)."
             )
         }
         .padding(.vertical)
diff --git a/CustomSuggestionService/ChatModelManagement/CompletionModelEditView.swift b/CustomSuggestionService/ChatModelManagement/CompletionModelEditView.swift
index fa3e830..edee33e 100644
--- a/CustomSuggestionService/ChatModelManagement/CompletionModelEditView.swift
+++ b/CustomSuggestionService/ChatModelManagement/CompletionModelEditView.swift
@@ -124,7 +124,7 @@ struct CompletionModelEditView: View {
             )
 
             TextField(text: textFieldBinding) {
-                Text("Max Tokens (Including Reply)")
+                Text("Context Window")
                     .multilineTextAlignment(.trailing)
             }
             .overlay(alignment: .trailing) {

From e2e6440df850399d0ad45a5258eb9d06ffc4f88a Mon Sep 17 00:00:00 2001
From: Shx Guo 
Date: Mon, 4 Mar 2024 00:56:35 +0800
Subject: [PATCH 25/27] Update model format to match CopilotForXcode

---
 .../CodeCompletionService.swift               |  4 +-
 .../Fundamental/Models/ChatModel.swift        | 51 ++++++++++++++-----
 .../Fundamental/Models/CompletionModel.swift  | 22 ++++----
 .../ChatModelManagement/ChatModelEdit.swift   |  4 +-
 .../CompletionModelEdit.swift                 |  4 +-
 5 files changed, 56 insertions(+), 29 deletions(-)

diff --git a/Core/Sources/CodeCompletionService/CodeCompletionService.swift b/Core/Sources/CodeCompletionService/CodeCompletionService.swift
index 2214f58..ad8c69b 100644
--- a/Core/Sources/CodeCompletionService/CodeCompletionService.swift
+++ b/Core/Sources/CodeCompletionService/CodeCompletionService.swift
@@ -121,7 +121,7 @@ public struct CodeCompletionService {
                 endpoint: .chatCompletion,
                 modelName: model.info.modelName,
                 stopWords: prompt.stopWords,
-                keepAlive: model.info.ollamaKeepAlive,
+                keepAlive: model.info.ollamaInfo.keepAlive,
                 format: .none
             )
             let result = try await service.getCompletions(prompt, count: count)
@@ -168,7 +168,7 @@ public struct CodeCompletionService {
                 endpoint: .completion,
                 modelName: model.info.modelName,
                 stopWords: prompt.stopWords,
-                keepAlive: model.info.ollamaKeepAlive,
+                keepAlive: model.info.ollamaInfo.keepAlive,
                 format: .none
             )
             let result = try await service.getCompletions(prompt, count: count)
diff --git a/Core/Sources/Fundamental/Models/ChatModel.swift b/Core/Sources/Fundamental/Models/ChatModel.swift
index 3cb7d61..9679027 100644
--- a/Core/Sources/Fundamental/Models/ChatModel.swift
+++ b/Core/Sources/Fundamental/Models/ChatModel.swift
@@ -23,11 +23,29 @@ public struct ChatModel: Codable, Equatable, Identifiable {
         case openAICompatible
         case googleAI
         case ollama
-        
+
         case unknown
     }
 
     public struct Info: Codable, Equatable {
+        public struct OllamaInfo: Codable, Equatable {
+            @FallbackDecoding
+            public var keepAlive: String
+
+            public init(keepAlive: String = "") {
+                self.keepAlive = keepAlive
+            }
+        }
+
+        public struct OpenAIInfo: Codable, Equatable {
+            @FallbackDecoding
+            public var organizationID: String
+
+            public init(organizationID: String = "") {
+                self.organizationID = organizationID
+            }
+        }
+
         @FallbackDecoding
         public var apiKeyName: String
         @FallbackDecoding
@@ -38,16 +56,13 @@ public struct ChatModel: Codable, Equatable, Identifiable {
         public var maxTokens: Int
         @FallbackDecoding
         public var supportsFunctionCalling: Bool
-        @FallbackDecoding
-        public var supportsOpenAIAPI2023_11: Bool
         @FallbackDecoding
         public var modelName: String
-        public var azureOpenAIDeploymentName: String {
-            get { modelName }
-            set { modelName = newValue }
-        }
-        @FallbackDecoding
-        public var ollamaKeepAlive: String
+
+        @FallbackDecoding
+        public var openAIInfo: OpenAIInfo
+        @FallbackDecoding
+        public var ollamaInfo: OllamaInfo
 
         public init(
             apiKeyName: String = "",
@@ -55,18 +70,18 @@ public struct ChatModel: Codable, Equatable, Identifiable {
             isFullURL: Bool = false,
             maxTokens: Int = 4000,
             supportsFunctionCalling: Bool = true,
-            supportsOpenAIAPI2023_11: Bool = false,
             modelName: String = "",
-            ollamaKeepAlive: String = ""
+            openAIInfo: OpenAIInfo = OpenAIInfo(),
+            ollamaInfo: OllamaInfo = OllamaInfo()
         ) {
             self.apiKeyName = apiKeyName
             self.baseURL = baseURL
             self.isFullURL = isFullURL
             self.maxTokens = maxTokens
             self.supportsFunctionCalling = supportsFunctionCalling
-            self.supportsOpenAIAPI2023_11 = supportsOpenAIAPI2023_11
             self.modelName = modelName
-            self.ollamaKeepAlive = ollamaKeepAlive
+            self.openAIInfo = openAIInfo
+            self.ollamaInfo = ollamaInfo
         }
     }
 
@@ -83,7 +98,7 @@ public struct ChatModel: Codable, Equatable, Identifiable {
             return "\(baseURL)/v1/chat/completions"
         case .azureOpenAI:
             let baseURL = info.baseURL
-            let deployment = info.azureOpenAIDeploymentName
+            let deployment = info.modelName
             let version = "2023-07-01-preview"
             if baseURL.isEmpty { return "" }
             return "\(baseURL)/openai/deployments/\(deployment)/chat/completions?api-version=\(version)"
@@ -109,3 +124,11 @@ public struct EmptyChatModelFormat: FallbackValueProvider {
     public static var defaultValue: ChatModel.Format { .unknown }
 }
 
+public struct EmptyChatModelOllamaInfo: FallbackValueProvider {
+    public static var defaultValue: ChatModel.Info.OllamaInfo { .init() }
+}
+
+public struct EmptyChatModelOpenAIInfo: FallbackValueProvider {
+    public static var defaultValue: ChatModel.Info.OpenAIInfo { .init() }
+}
+
diff --git a/Core/Sources/Fundamental/Models/CompletionModel.swift b/Core/Sources/Fundamental/Models/CompletionModel.swift
index 3e4af42..3ba7ed9 100644
--- a/Core/Sources/Fundamental/Models/CompletionModel.swift
+++ b/Core/Sources/Fundamental/Models/CompletionModel.swift
@@ -27,6 +27,9 @@ public struct CompletionModel: Codable, Equatable, Identifiable {
     }
 
     public struct Info: Codable, Equatable {
+        public typealias OllamaInfo = ChatModel.Info.OllamaInfo
+        public typealias OpenAIInfo = ChatModel.Info.OpenAIInfo
+        
         @FallbackDecoding
         public var apiKeyName: String
         @FallbackDecoding
@@ -37,12 +40,11 @@ public struct CompletionModel: Codable, Equatable, Identifiable {
         public var maxTokens: Int
         @FallbackDecoding
         public var modelName: String
-        public var azureOpenAIDeploymentName: String {
-            get { modelName }
-            set { modelName = newValue }
-        }
-        @FallbackDecoding
-        public var ollamaKeepAlive: String
+       
+        @FallbackDecoding
+        public var openAIInfo: OpenAIInfo
+        @FallbackDecoding
+        public var ollamaInfo: OllamaInfo
 
         public init(
             apiKeyName: String = "",
@@ -50,14 +52,16 @@ public struct CompletionModel: Codable, Equatable, Identifiable {
             isFullURL: Bool = false,
             maxTokens: Int = 4000,
             modelName: String = "",
-            ollamaKeepAlive: String = ""
+            openAIInfo: OpenAIInfo = OpenAIInfo(),
+            ollamaInfo: OllamaInfo = OllamaInfo()
         ) {
             self.apiKeyName = apiKeyName
             self.baseURL = baseURL
             self.isFullURL = isFullURL
             self.maxTokens = maxTokens
             self.modelName = modelName
-            self.ollamaKeepAlive = ollamaKeepAlive
+            self.openAIInfo = openAIInfo
+            self.ollamaInfo = ollamaInfo
         }
     }
 
@@ -74,7 +78,7 @@ public struct CompletionModel: Codable, Equatable, Identifiable {
             return "\(baseURL)/v1/completions"
         case .azureOpenAI:
             let baseURL = info.baseURL
-            let deployment = info.azureOpenAIDeploymentName
+            let deployment = info.modelName
             let version = "2023-07-01-preview"
             if baseURL.isEmpty { return "" }
             return "\(baseURL)/openai/deployments/\(deployment)/completions?api-version=\(version)"
diff --git a/CustomSuggestionService/ChatModelManagement/ChatModelEdit.swift b/CustomSuggestionService/ChatModelManagement/ChatModelEdit.swift
index 13409a6..def138d 100644
--- a/CustomSuggestionService/ChatModelManagement/ChatModelEdit.swift
+++ b/CustomSuggestionService/ChatModelManagement/ChatModelEdit.swift
@@ -154,7 +154,7 @@ extension ChatModel {
                 apiKeyManagement: .init(availableAPIKeyNames: [info.apiKeyName])
             ),
             baseURLSelection: .init(baseURL: info.baseURL),
-            ollamaKeepAlive: info.ollamaKeepAlive
+            ollamaKeepAlive: info.ollamaInfo.keepAlive
         )
     }
 
@@ -169,7 +169,7 @@ extension ChatModel {
                 maxTokens: state.maxTokens,
                 supportsFunctionCalling: false,
                 modelName: state.modelName.trimmingCharacters(in: .whitespacesAndNewlines),
-                ollamaKeepAlive: state.ollamaKeepAlive
+                ollamaInfo: .init(keepAlive: state.ollamaKeepAlive)
             )
         )
     }
diff --git a/CustomSuggestionService/ChatModelManagement/CompletionModelEdit.swift b/CustomSuggestionService/ChatModelManagement/CompletionModelEdit.swift
index 78b31d7..4e7cd12 100644
--- a/CustomSuggestionService/ChatModelManagement/CompletionModelEdit.swift
+++ b/CustomSuggestionService/ChatModelManagement/CompletionModelEdit.swift
@@ -145,7 +145,7 @@ extension CompletionModel {
                 apiKeyManagement: .init(availableAPIKeyNames: [info.apiKeyName])
             ),
             baseURLSelection: .init(baseURL: info.baseURL),
-            ollamaKeepAlive: info.ollamaKeepAlive
+            ollamaKeepAlive: info.ollamaInfo.keepAlive
         )
     }
 
@@ -159,7 +159,7 @@ extension CompletionModel {
                 baseURL: state.baseURL.trimmingCharacters(in: .whitespacesAndNewlines),
                 maxTokens: state.maxTokens,
                 modelName: state.modelName.trimmingCharacters(in: .whitespacesAndNewlines),
-                ollamaKeepAlive: state.ollamaKeepAlive
+                ollamaInfo: .init(keepAlive: state.ollamaKeepAlive)
             )
         )
     }

From 98f5577ac0967160cdd298ac38d877c1da6b313d Mon Sep 17 00:00:00 2001
From: Shx Guo 
Date: Mon, 4 Mar 2024 00:59:13 +0800
Subject: [PATCH 26/27] Bump version to 0.2.0

---
 Version.xcconfig | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/Version.xcconfig b/Version.xcconfig
index 16dd90b..166c909 100644
--- a/Version.xcconfig
+++ b/Version.xcconfig
@@ -1,3 +1,3 @@
-APP_VERSION = 0.1.0
-APP_BUILD = 12
+APP_VERSION = 0.2.0
+APP_BUILD = 20
  

From 95bd1e4400fb8f8a66fd58be828b72eec9d0e981 Mon Sep 17 00:00:00 2001
From: Shx Guo 
Date: Mon, 4 Mar 2024 12:56:43 +0800
Subject: [PATCH 27/27] Update appcast.xml

---
 appcast.xml | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/appcast.xml b/appcast.xml
index 84db9d9..aff937d 100644
--- a/appcast.xml
+++ b/appcast.xml
@@ -2,6 +2,18 @@
 
     
         Custom Suggestion Service
+        
+            0.2.0
+            Mon, 04 Mar 2024 12:50:34 +0800
+            20
+            0.2.0
+            13.0
+            
+                https://github.com/intitni/CustomSuggestionServiceForCopilotForXcode/releases/tag/0.2.0
+            
+            
+        
+        
         
             0.1.0
             Thu, 22 Feb 2024 16:20:36 +0800
@@ -11,7 +23,7 @@
             
                 https://github.com/intitni/CustomSuggestionServiceForCopilotForXcode/releases/tag/0.1.0
             
-            
+