Update RequestStrategy to provide StreamStopStrategy

Author: intitni

Core/Sources/SuggestionService/RequestStrategies/CodeLlamaFillInTheMiddleRequestStrategy.swift

@@ -1,3 +1,4 @@
+import CodeCompletionService
 import CopilotForXcodeKit
 import Foundation
 import Fundamental
@@ -20,6 +21,10 @@ struct CodeLlamaFillInTheMiddleRequestStrategy: RequestStrategy {
         )
     }
 
+    func createStreamStopStrategy() -> some StreamStopStrategy {
+        FIMStreamStopStrategy(prefix: prefix)
+    }
+
     func createRawSuggestionPostProcessor() -> some RawSuggestionPostProcessingStrategy {
         DefaultRawSuggestionPostProcessingStrategy(codeWrappingTags: nil)
     }
@@ -92,8 +97,41 @@ struct CodeLlamaFillInTheMiddleWithSystemPromptRequestStrategy: RequestStrategy
         return prompt
     }
 
+    func createStreamStopStrategy() -> some StreamStopStrategy {
+        strategy.createStreamStopStrategy()
+    }
+
     func createRawSuggestionPostProcessor() -> some RawSuggestionPostProcessingStrategy {
-        DefaultRawSuggestionPostProcessingStrategy(codeWrappingTags: nil)
+        strategy.createRawSuggestionPostProcessor()
+    }
+}
+
+struct FIMStreamStopStrategy: StreamStopStrategy {
+    let prefix: [String]
+
+    func shouldStop(
+        existedLines: [String],
+        currentLine: String,
+        proposedLineLimit: Int
+    ) -> StreamStopStrategyResult {
+        if let prefixLastLine = prefix.last {
+            if let lastLineIndex = existedLines.lastIndex(of: prefixLastLine) {
+                if existedLines.count >= lastLineIndex + 1 + proposedLineLimit {
+                    return .stop(appendingNewContent: true)
+                }
+                return .continue
+            } else {
+                if existedLines.count >= proposedLineLimit {
+                    return .stop(appendingNewContent: true)
+                }
+                return .continue
+            }
+        } else {
+            if existedLines.count >= proposedLineLimit {
+                return .stop(appendingNewContent: true)
+            }
+            return .continue
+        }
     }
 }
 

Core/Sources/SuggestionService/RequestStrategies/ContinueRequestStrategy.swift

@@ -1,3 +1,4 @@
+import CodeCompletionService
 import CopilotForXcodeKit
 import Foundation
 import Fundamental
@@ -21,14 +22,21 @@ struct ContinueRequestStrategy: RequestStrategy {
             suffix: suffix
         )
     }
-    
+
     func createRawSuggestionPostProcessor() -> DefaultRawSuggestionPostProcessingStrategy {
         DefaultRawSuggestionPostProcessingStrategy(codeWrappingTags: (
             Tag.openingCode,
             Tag.closingCode
         ))
     }
 
+    func createStreamStopStrategy() -> some StreamStopStrategy {
+        OpeningTagBasedStreamStopStrategy(
+            openingTag: Tag.openingCode,
+            toleranceIfNoOpeningTagFound: 4
+        )
+    }
+
     enum Tag {
         public static let openingCode = "<Code3721>"
         public static let closingCode = "</Code3721>"

Core/Sources/SuggestionService/RequestStrategies/DefaultRequestStrategy.swift

@@ -1,3 +1,4 @@
+import CodeCompletionService
 import CopilotForXcodeKit
 import Foundation
 import Fundamental
@@ -23,6 +24,13 @@ struct DefaultRequestStrategy: RequestStrategy {
         )
     }
 
+    func createStreamStopStrategy() -> some StreamStopStrategy {
+        OpeningTagBasedStreamStopStrategy(
+            openingTag: Tag.openingCode,
+            toleranceIfNoOpeningTagFound: 4
+        )
+    }
+
     func createRawSuggestionPostProcessor() -> DefaultRawSuggestionPostProcessingStrategy {
         DefaultRawSuggestionPostProcessingStrategy(codeWrappingTags: (
             Tag.openingCode,

Core/Sources/SuggestionService/RequestStrategies/NaiveRequestStrategy.swift

@@ -1,3 +1,4 @@
+import CodeCompletionService
 import CopilotForXcodeKit
 import Foundation
 import Fundamental
@@ -8,7 +9,7 @@ struct NaiveRequestStrategy: RequestStrategy {
     var sourceRequest: SuggestionRequest
     var prefix: [String]
     var suffix: [String]
-    
+
     var shouldSkip: Bool {
         prefix.last?.trimmingCharacters(in: .whitespaces) == "}"
     }
@@ -20,11 +21,15 @@ struct NaiveRequestStrategy: RequestStrategy {
             suffix: suffix
         )
     }
-    
+
     func createRawSuggestionPostProcessor() -> some RawSuggestionPostProcessingStrategy {
         NoOpRawSuggestionPostProcessingStrategy()
     }
 
+    func createStreamStopStrategy() -> some StreamStopStrategy {
+        DefaultStreamStopStrategy()
+    }
+
     struct Request: PromptStrategy {
         let systemPrompt: String = ""
         var sourceRequest: SuggestionRequest
@@ -34,7 +39,7 @@ struct NaiveRequestStrategy: RequestStrategy {
         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()
@@ -71,9 +76,9 @@ struct NaiveRequestStrategy: RequestStrategy {
 
             return [.init(role: .user, content: """
             File path: \(filePath)
-            
+
             ---
-            
+
             \(code)
             """.trimmingCharacters(in: .whitespacesAndNewlines))]
         }

Core/Sources/SuggestionService/RequestStrategies/TabbyRequestStrategy.swift

@@ -1,3 +1,4 @@
+import CodeCompletionService
 import CopilotForXcodeKit
 import Foundation
 import Fundamental
@@ -24,6 +25,10 @@ struct TabbyRequestStrategy: RequestStrategy {
         NoOpRawSuggestionPostProcessingStrategy()
     }
 
+    func createStreamStopStrategy() -> some StreamStopStrategy {
+        NeverStreamStopStrategy()
+    }
+
     struct Prompt: PromptStrategy {
         let systemPrompt: String = ""
         var sourceRequest: SuggestionRequest

Core/Sources/SuggestionService/RequestStrategy.swift

@@ -1,3 +1,4 @@
+import CodeCompletionService
 import CopilotForXcodeKit
 import Foundation
 import Fundamental
@@ -8,6 +9,7 @@ import Parsing
 protocol RequestStrategy {
     associatedtype Prompt: PromptStrategy
     associatedtype RawSuggestionPostProcessor: RawSuggestionPostProcessingStrategy
+    associatedtype SomeStreamStopStrategy: StreamStopStrategy
 
     init(sourceRequest: SuggestionRequest, prefix: [String], suffix: [String])
 
@@ -17,6 +19,9 @@ protocol RequestStrategy {
     /// Create a prompt to generate code completion.
     func createPrompt() -> Prompt
 
+    /// Control how a stream should stop early.
+    func createStreamStopStrategy() -> SomeStreamStopStrategy
+
     /// The AI model may not return a suggestion in a ideal format. You can use it to reformat the
     /// suggestions.
     func createRawSuggestionPostProcessor() -> RawSuggestionPostProcessor

Core/Sources/SuggestionService/Service.swift

@@ -46,28 +46,32 @@ actor Service {
 
                     let prompt = strategy.createPrompt()
                     let postProcessor = strategy.createRawSuggestionPostProcessor()
+                    let stopStream = strategy.createStreamStopStrategy()
 
                     let suggestedCodeSnippets: [String]
 
                     switch getModel() {
                     case let .chatModel(model):
                         CodeCompletionLogger.logger.logModel(model)
                         suggestedCodeSnippets = try await service.getCompletions(
-                            prompt,
+                            prompt, 
+                            streamStopStrategy: stopStream,
                             model: model,
                             count: 1
                         )
                     case let .completionModel(model):
                         CodeCompletionLogger.logger.logModel(model)
                         suggestedCodeSnippets = try await service.getCompletions(
                             prompt,
+                            streamStopStrategy: stopStream,
                             model: model,
                             count: 1
                         )
                     case let .tabbyModel(model):
                         CodeCompletionLogger.logger.logModel(model)
                         suggestedCodeSnippets = try await service.getCompletions(
                             prompt,
+                            streamStopStrategy: stopStream,
                             model: model,
                             count: 1
                         )