Limit the suggestion to match the line limit

Author: intitni

Core/Sources/SuggestionService/RawSuggestionPostProcessing/DefaultRawSuggestionPostProcessingStrategy.swift

@@ -5,24 +5,14 @@ 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 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 {

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
     }
 }
 

Core/Sources/SuggestionService/Service.swift

@@ -54,7 +54,7 @@ actor Service {
                     case let .chatModel(model):
                         CodeCompletionLogger.logger.logModel(model)
                         suggestedCodeSnippets = try await service.getCompletions(
-                            prompt, 
+                            prompt,
                             streamStopStrategy: stopStream,
                             model: model,
                             count: 1
@@ -82,13 +82,21 @@ actor Service {
                     return suggestedCodeSnippets
                         .filter { !$0.allSatisfy { $0.isWhitespace || $0.isNewline } }
                         .map {
-                            CodeSuggestion(
-                                id: UUID().uuidString,
-                                text: postProcessor.postProcess(
+                            let suggestionText = postProcessor
+                                .postProcess(
                                     rawSuggestion: $0,
                                     infillPrefix: prompt.suggestionPrefix.prependingValue,
                                     suffix: prompt.suffix
-                                ),
+                                )
+                                .keepLines(
+                                    count: UserDefaults.shared
+                                        .value(for: \.maxNumberOfLinesOfSuggestion)
+                                )
+                                .removeTrailingNewlinesAndWhitespace()
+
+                            return CodeSuggestion(
+                                id: UUID().uuidString,
+                                text: suggestionText,
                                 position: request.cursorPosition,
                                 range: .init(
                                     start: .init(
@@ -186,3 +194,19 @@ actor Service {
     }
 }
 
+extension String {
+    func keepLines(count: Int) -> String {
+        if count <= 0 { return self }
+        let lines = breakLines()
+        return lines.prefix(count).joined()
+    }
+
+    func removeTrailingNewlinesAndWhitespace() -> String {
+        var text = self[...]
+        while let last = text.last, last.isNewline || last.isWhitespace {
+            text = text.dropLast(1)
+        }
+        return String(text)
+    }
+}
+