Update FIM strategies to use the default post processing strategy

Author: intitni

Core/Sources/SuggestionService/RawSuggestionPostProcessing/DefaultRawSuggestionPostProcessingStrategy.swift

@@ -16,8 +16,7 @@ extension RawSuggestionPostProcessingStrategy {
 }
 
 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)
@@ -28,13 +27,16 @@ struct DefaultRawSuggestionPostProcessingStrategy: RawSuggestionPostProcessingSt
 
     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/RequestStrategies/CodeLlamaFillInTheMiddleRequestStrategy.swift

@@ -20,8 +20,8 @@ struct CodeLlamaFillInTheMiddleRequestStrategy: RequestStrategy {
         )
     }
 
-    func createRawSuggestionPostProcessor() -> NoOpRawSuggestionPostProcessingStrategy {
-        NoOpRawSuggestionPostProcessingStrategy()
+    func createRawSuggestionPostProcessor() -> some RawSuggestionPostProcessingStrategy {
+        DefaultRawSuggestionPostProcessingStrategy(codeWrappingTags: nil)
     }
 
     enum Tag {
@@ -93,7 +93,7 @@ struct CodeLlamaFillInTheMiddleWithSystemPromptRequestStrategy: RequestStrategy
     }
 
     func createRawSuggestionPostProcessor() -> some RawSuggestionPostProcessingStrategy {
-        DefaultRawSuggestionPostProcessingStrategy(openingCodeTag: "", closingCodeTag: "")
+        DefaultRawSuggestionPostProcessingStrategy(codeWrappingTags: nil)
     }
 }
 

Core/Sources/SuggestionService/RequestStrategies/ContinueRequestStrategy.swift

@@ -23,10 +23,10 @@ struct ContinueRequestStrategy: RequestStrategy {
     }
 
     func createRawSuggestionPostProcessor() -> DefaultRawSuggestionPostProcessingStrategy {
-        DefaultRawSuggestionPostProcessingStrategy(
-            openingCodeTag: Tag.openingCode,
-            closingCodeTag: Tag.closingCode
-        )
+        DefaultRawSuggestionPostProcessingStrategy(codeWrappingTags: (
+            Tag.openingCode,
+            Tag.closingCode
+        ))
     }
 
     enum Tag {

Core/Sources/SuggestionService/RequestStrategies/DefaultRequestStrategy.swift

@@ -10,7 +10,7 @@ struct DefaultRequestStrategy: RequestStrategy {
     var sourceRequest: SuggestionRequest
     var prefix: [String]
     var suffix: [String]
-    
+
     var shouldSkip: Bool {
         prefix.last?.trimmingCharacters(in: .whitespaces) == "}"
     }
@@ -22,12 +22,12 @@ struct DefaultRequestStrategy: RequestStrategy {
             suffix: suffix
         )
     }
-    
+
     func createRawSuggestionPostProcessor() -> DefaultRawSuggestionPostProcessingStrategy {
-        DefaultRawSuggestionPostProcessingStrategy(
-            openingCodeTag: Tag.openingCode,
-            closingCodeTag: Tag.closingCode
-        )
+        DefaultRawSuggestionPostProcessingStrategy(codeWrappingTags: (
+            Tag.openingCode,
+            Tag.closingCode
+        ))
     }
 
     enum Tag {
@@ -65,7 +65,7 @@ struct DefaultRequestStrategy: RequestStrategy {
         var relevantCodeSnippets: [RelevantCodeSnippet] { sourceRequest.relevantCodeSnippets }
         var stopWords: [String] { [Tag.closingCode, "\n\n"] }
         var language: CodeLanguage? { sourceRequest.language }
-        
+
         var suggestionPrefix: SuggestionPrefix {
             guard let prefix = prefix.last else { return .empty }
             return .unchanged(prefix).curlyBracesLineBreak()
@@ -103,7 +103,7 @@ struct DefaultRequestStrategy: RequestStrategy {
             File Path: \(filePath)
             Indentation: \
             \(sourceRequest.indentSize) \(sourceRequest.usesTabsForIndentation ? "tab" : "space")
-            
+
             ---
 
             Here is the code: