Merge branch 'release/0.3.0'

Author: intitni

Core/Package.swift

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

…mpletionService/AzureOpenAIService.swift …tionService/API/AzureOpenAIService.swiftCore/Sources/CodeCompletionService/AzureOpenAIService.swift renamed to Core/Sources/CodeCompletionService/API/AzureOpenAIService.swift Core/Sources/CodeCompletionService/AzureOpenAIService.swift renamed to Core/Sources/CodeCompletionService/API/AzureOpenAIService.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:

…pletionService/GoogleGeminiService.swift …ionService/API/GoogleGeminiService.swiftCore/Sources/CodeCompletionService/GoogleGeminiService.swift renamed to Core/Sources/CodeCompletionService/API/GoogleGeminiService.swift Core/Sources/CodeCompletionService/GoogleGeminiService.swift renamed to Core/Sources/CodeCompletionService/API/GoogleGeminiService.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)
+    }
+}
+

…odeCompletionService/OllamaService.swift …ompletionService/API/OllamaService.swiftCore/Sources/CodeCompletionService/OllamaService.swift renamed to Core/Sources/CodeCompletionService/API/OllamaService.swift Core/Sources/CodeCompletionService/OllamaService.swift renamed to Core/Sources/CodeCompletionService/API/OllamaService.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
+    }
+}
+

…odeCompletionService/OpenAIService.swift …ompletionService/API/OpenAIService.swiftCore/Sources/CodeCompletionService/OpenAIService.swift renamed to Core/Sources/CodeCompletionService/API/OpenAIService.swift Core/Sources/CodeCompletionService/OpenAIService.swift renamed to Core/Sources/CodeCompletionService/API/OpenAIService.swift

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