Merge branch 'release/0.3.0'

Core/Package.swift

@@ -66,6 +66,10 @@ let package = Package(
             ]
         ),
 
+        .testTarget(
+            name: "CodeCompletionServiceTests",
+            dependencies: ["CodeCompletionService"]
+        ),
         .testTarget(
             name: "FundamentalTests",
             dependencies: ["Fundamental"]

Core/Sources/CodeCompletionService/CodeCompletionService.swift

@@ -11,17 +11,20 @@ protocol CodeCompletionServiceType {
 extension CodeCompletionServiceType {
     func getCompletions(
         _ request: PromptStrategy,
+        streamStopStrategy: StreamStopStrategy,
         count: Int
     ) async throws -> [String] {
         try await withThrowingTaskGroup(of: String.self) { group in
             for _ in 0..<max(1, count) {
                 _ = group.addTaskUnlessCancelled {
-                    var result = ""
+                    let limiter = StreamLineLimiter(strategy: streamStopStrategy)
                     let stream = try await getCompletion(request)
                     for try await response in stream {
-                        result.append(response)
+                        if case let .finish(result) = limiter.push(response) {
+                            return result
+                        }
                     }
-                    return result
+                    return limiter.result
                 }
             }
 
@@ -53,6 +56,7 @@ public struct CodeCompletionService {
 
     public func getCompletions(
         _ prompt: PromptStrategy,
+        streamStopStrategy: StreamStopStrategy,
         model: TabbyModel,
         count: Int
     ) async throws -> [String] {
@@ -71,13 +75,18 @@ public struct CodeCompletionService {
             }
         }())
 
-        let result = try await service.getCompletions(prompt, count: count)
+        let result = try await service.getCompletions(
+            prompt,
+            streamStopStrategy: streamStopStrategy,
+            count: count
+        )
         try Task.checkCancellation()
         return result
     }
 
     public func getCompletions(
         _ prompt: PromptStrategy,
+        streamStopStrategy: StreamStopStrategy,
         model: ChatModel,
         count: Int
     ) async throws -> [String] {
@@ -92,7 +101,11 @@ public struct CodeCompletionService {
                 stopWords: prompt.stopWords,
                 apiKey: apiKey
             )
-            let result = try await service.getCompletions(prompt, count: count)
+            let result = try await service.getCompletions(
+                prompt,
+                streamStopStrategy: streamStopStrategy,
+                count: count
+            )
             try Task.checkCancellation()
             return result
         case .azureOpenAI:
@@ -103,7 +116,11 @@ public struct CodeCompletionService {
                 stopWords: prompt.stopWords,
                 apiKey: apiKey
             )
-            let result = try await service.getCompletions(prompt, count: count)
+            let result = try await service.getCompletions(
+                prompt,
+                streamStopStrategy: streamStopStrategy,
+                count: count
+            )
             try Task.checkCancellation()
             return result
         case .googleAI:
@@ -112,7 +129,11 @@ public struct CodeCompletionService {
                 maxToken: model.info.maxTokens,
                 apiKey: apiKey
             )
-            let result = try await service.getCompletions(prompt, count: count)
+            let result = try await service.getCompletions(
+                prompt,
+                streamStopStrategy: streamStopStrategy,
+                count: count
+            )
             try Task.checkCancellation()
             return result
         case .ollama:
@@ -124,7 +145,11 @@ public struct CodeCompletionService {
                 keepAlive: model.info.ollamaInfo.keepAlive,
                 format: .none
             )
-            let result = try await service.getCompletions(prompt, count: count)
+            let result = try await service.getCompletions(
+                prompt,
+                streamStopStrategy: streamStopStrategy,
+                count: count
+            )
             try Task.checkCancellation()
             return result
         case .unknown:
@@ -134,6 +159,7 @@ public struct CodeCompletionService {
 
     public func getCompletions(
         _ prompt: PromptStrategy,
+        streamStopStrategy: StreamStopStrategy,
         model: CompletionModel,
         count: Int
     ) async throws -> [String] {
@@ -148,7 +174,11 @@ public struct CodeCompletionService {
                 stopWords: prompt.stopWords,
                 apiKey: apiKey
             )
-            let result = try await service.getCompletions(prompt, count: count)
+            let result = try await service.getCompletions(
+                prompt,
+                streamStopStrategy: streamStopStrategy,
+                count: count
+            )
             try Task.checkCancellation()
             return result
         case .azureOpenAI:
@@ -159,7 +189,11 @@ public struct CodeCompletionService {
                 stopWords: prompt.stopWords,
                 apiKey: apiKey
             )
-            let result = try await service.getCompletions(prompt, count: count)
+            let result = try await service.getCompletions(
+                prompt,
+                streamStopStrategy: streamStopStrategy,
+                count: count
+            )
             try Task.checkCancellation()
             return result
         case .ollama:
@@ -171,7 +205,11 @@ public struct CodeCompletionService {
                 keepAlive: model.info.ollamaInfo.keepAlive,
                 format: .none
             )
-            let result = try await service.getCompletions(prompt, count: count)
+            let result = try await service.getCompletions(
+                prompt,
+                streamStopStrategy: streamStopStrategy,
+                count: count
+            )
             try Task.checkCancellation()
             return result
         case .unknown:

Core/Sources/CodeCompletionService/StreamLineLimiter.swift

@@ -0,0 +1,62 @@
+import Foundation
+
+final class StreamLineLimiter {
+    public private(set) var result = ""
+    private var currentLine = ""
+    private var existedLines = [String]()
+    private let lineLimit: Int
+    private let strategy: any StreamStopStrategy
+
+    enum PushResult: Equatable {
+        case `continue`
+        case finish(String)
+    }
+
+    init(
+        lineLimit: Int = UserDefaults.shared.value(for: \.maxNumberOfLinesOfSuggestion),
+        strategy: any StreamStopStrategy
+    ) {
+        self.lineLimit = lineLimit
+        self.strategy = strategy
+    }
+
+    func push(_ token: String) -> PushResult {
+        currentLine.append(token)
+        if let newLine = currentLine.last(where: { $0.isNewline }) {
+            let lines = currentLine
+                .breakLines(proposedLineEnding: String(newLine), appendLineBreakToLastLine: false)
+            let (newLines, lastLine) = lines.headAndTail
+            existedLines.append(contentsOf: newLines)
+            currentLine = lastLine ?? ""
+        }
+
+        let stopResult = if lineLimit <= 0 {
+            StreamStopStrategyResult.continue
+        } else {
+            strategy.shouldStop(
+                existedLines: existedLines,
+                currentLine: currentLine,
+                proposedLineLimit: lineLimit
+            )
+        }
+
+        switch stopResult {
+        case .continue:
+            result.append(token)
+            return .continue
+        case let .stop(appendingNewContent):
+            if appendingNewContent {
+                result.append(token)
+            }
+            return .finish(result)
+        }
+    }
+}
+
+extension Array {
+    var headAndTail: ([Element], Element?) {
+        guard let tail = last else { return ([], nil) }
+        return (Array(dropLast()), tail)
+    }
+}
+

Core/Sources/CodeCompletionService/StreamStopStrategy/DefaultStreamStopStrategy.swift

@@ -0,0 +1,15 @@
+public struct DefaultStreamStopStrategy: StreamStopStrategy {
+    public init() {}
+
+    public func shouldStop(
+        existedLines: [String],
+        currentLine: String,
+        proposedLineLimit: Int
+    ) -> StreamStopStrategyResult {
+        if existedLines.count >= proposedLineLimit {
+            return .stop(appendingNewContent: true)
+        }
+        return .continue
+    }
+}
+

Core/Sources/CodeCompletionService/StreamStopStrategy/NeverStreamStopStrategy.swift

@@ -0,0 +1,12 @@
+public struct NeverStreamStopStrategy: StreamStopStrategy {
+    public init() {}
+
+    public func shouldStop(
+        existedLines: [String],
+        currentLine: String,
+        proposedLineLimit: Int
+    ) -> StreamStopStrategyResult {
+        .continue
+    }
+}
+

Core/Sources/CodeCompletionService/StreamStopStrategy/OpeningTagBasedStreamStopStrategy.swift

@@ -0,0 +1,31 @@
+import Foundation
+
+public struct OpeningTagBasedStreamStopStrategy: StreamStopStrategy {
+    public let openingTag: String
+    public let toleranceIfNoOpeningTagFound: Int
+
+    public init(openingTag: String, toleranceIfNoOpeningTagFound: Int) {
+        self.openingTag = openingTag
+        self.toleranceIfNoOpeningTagFound = toleranceIfNoOpeningTagFound
+    }
+
+    public func shouldStop(
+        existedLines: [String],
+        currentLine: String,
+        proposedLineLimit: Int
+    ) -> StreamStopStrategyResult {
+        if let index = existedLines.firstIndex(where: { $0.contains(openingTag) }) {
+            if existedLines.count - index - 1 >= proposedLineLimit {
+                return .stop(appendingNewContent: true)
+            }
+            return .continue
+        } else {
+            if existedLines.count >= proposedLineLimit + toleranceIfNoOpeningTagFound {
+                return .stop(appendingNewContent: true)
+            } else {
+                return .continue
+            }
+        }
+    }
+}
+

Core/Sources/CodeCompletionService/StreamStopStrategy/StreamStopStrategy.swift

@@ -0,0 +1,12 @@
+import Foundation
+
+public enum StreamStopStrategyResult {
+    case `continue`
+    case stop(appendingNewContent: Bool)
+}
+
+public protocol StreamStopStrategy {
+    func shouldStop(existedLines: [String], currentLine: String, proposedLineLimit: Int)
+        -> StreamStopStrategyResult
+}
+

Core/Sources/Storage/Preferences.swift

@@ -80,6 +80,13 @@ public extension UserDefaultPreferenceKeys {
             key: "CustomSuggestionService-TabbyModel"
         )
     }
+
+    var maxNumberOfLinesOfSuggestion: PreferenceKey<Int> {
+        .init(
+            defaultValue: 0,
+            key: "CustomSuggestionService-MaxNumberOfLinesOfSuggestion"
+        )
+    }
 
     var installBetaBuild: PreferenceKey<Bool> {
         .init(defaultValue: false, key: "CustomSuggestionService-InstallBetaBuild")

Core/Sources/SuggestionService/RawSuggestionPostProcessing/DefaultRawSuggestionPostProcessingStrategy.swift

@@ -5,36 +5,28 @@ protocol RawSuggestionPostProcessingStrategy {
     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
+    let codeWrappingTags: (opening: String, closing: String)?
 
     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)
+        return infillPrefix + suggestion
     }
 
     func extractSuggestion(from response: String) -> String {
         let escapedMarkdownCodeBlock = removeLeadingAndTrailingMarkdownCodeBlockMark(from: response)
-        let escapedTags = extractEnclosingSuggestion(
-            from: escapedMarkdownCodeBlock,
-            openingTag: openingCodeTag,
-            closingTag: closingCodeTag
-        )
-
-        return escapedTags
+        if let tags = codeWrappingTags {
+            let escapedTags = extractEnclosingSuggestion(
+                from: escapedMarkdownCodeBlock,
+                openingTag: tags.opening,
+                closingTag: tags.closing
+            )
+            return escapedTags
+        } else {
+            return escapedMarkdownCodeBlock
+        }
     }
 
     func removePrefix(from suggestion: inout String, infillPrefix: String) {

Core/Sources/SuggestionService/RawSuggestionPostProcessing/NoOpRawSuggestionPostProcessingStrategy.swift

@@ -2,7 +2,7 @@ import Foundation
 
 struct NoOpRawSuggestionPostProcessingStrategy: RawSuggestionPostProcessingStrategy {
     func postProcess(rawSuggestion: String, infillPrefix: String, suffix: [String]) -> String {
-        removeTrailingNewlinesAndWhitespace(from: infillPrefix + rawSuggestion)
+        infillPrefix + rawSuggestion
     }
 }