Merge branch 'feature/stream-response' into develop

Author: intitni

Core/Sources/CodeCompletionService/AzureOpenAIService.swift

@@ -39,18 +39,38 @@ public actor AzureOpenAIService {
 }
 
 extension AzureOpenAIService: CodeCompletionServiceType {
-    func getCompletion(_ request: PromptStrategy) async throws -> String {
+    func getCompletion(_ request: PromptStrategy) async throws -> AsyncStream<String> {
         switch endpoint {
         case .chatCompletion:
             let messages = createMessages(from: request)
             CodeCompletionLogger.logger.logPrompt(messages.map {
                 ($0.content, $0.role.rawValue)
             })
-            return try await sendMessages(messages)
+            return AsyncStream<String> { continuation in
+                let task = Task {
+                    let result = try await sendMessages(messages)
+                    try Task.checkCancellation()
+                    continuation.yield(result)
+                    continuation.finish()
+                }
+                continuation.onTermination = { _ in
+                    task.cancel()
+                }
+            }
         case .completion:
             let prompt = createPrompt(from: request)
             CodeCompletionLogger.logger.logPrompt([(prompt, "user")])
-            return try await sendPrompt(prompt)
+            return AsyncStream<String> { continuation in
+                let task = Task {
+                    let result = try await sendPrompt(prompt)
+                    try Task.checkCancellation()
+                    continuation.yield(result)
+                    continuation.finish()
+                }
+                continuation.onTermination = { _ in
+                    task.cancel()
+                }
+            }
         }
     }
 }

Core/Sources/CodeCompletionService/CodeCompletionService.swift

@@ -5,7 +5,7 @@ import Storage
 protocol CodeCompletionServiceType {
     func getCompletion(
         _ request: PromptStrategy
-    ) async throws -> String
+    ) async throws -> AsyncStream<String>
 }
 
 extension CodeCompletionServiceType {
@@ -16,7 +16,12 @@ extension CodeCompletionServiceType {
         try await withThrowingTaskGroup(of: String.self) { group in
             for _ in 0..<max(1, count) {
                 _ = group.addTaskUnlessCancelled {
-                    try await getCompletion(request)
+                    var result = ""
+                    let stream = try await getCompletion(request)
+                    for try await response in stream {
+                        result.append(response)
+                    }
+                    return result
                 }
             }
 

Core/Sources/CodeCompletionService/GoogleGeminiService.swift

@@ -25,12 +25,22 @@ public struct GoogleGeminiService {
 }
 
 extension GoogleGeminiService: CodeCompletionServiceType {
-    func getCompletion(_ request: PromptStrategy) async throws -> String {
+    func getCompletion(_ request: PromptStrategy) async throws -> AsyncStream<String> {
         let messages = createMessages(from: request)
         CodeCompletionLogger.logger.logPrompt(messages.map {
             ($0.parts.first?.text ?? "N/A", $0.role ?? "N/A")
         })
-        return try await sendMessages(messages)
+        return AsyncStream<String> { continuation in
+            let task = Task {
+                let result = try await sendMessages(messages)
+                try Task.checkCancellation()
+                continuation.yield(result)
+                continuation.finish()
+            }
+            continuation.onTermination = { _ in
+                task.cancel()
+            }
+        }
     }
 }
 

Core/Sources/CodeCompletionService/OpenAIService.swift

@@ -44,18 +44,38 @@ public actor OpenAIService {
 }
 
 extension OpenAIService: CodeCompletionServiceType {
-    func getCompletion(_ request: PromptStrategy) async throws -> String {
+    func getCompletion(_ request: PromptStrategy) async throws -> AsyncStream<String> {
         switch endpoint {
         case .chatCompletion:
             let messages = createMessages(from: request)
             CodeCompletionLogger.logger.logPrompt(messages.map {
                 ($0.content, $0.role.rawValue)
             })
-            return try await sendMessages(messages)
+            return AsyncStream<String> { continuation in
+                let task = Task {
+                    let result = try await sendMessages(messages)
+                    try Task.checkCancellation()
+                    continuation.yield(result)
+                    continuation.finish()
+                }
+                continuation.onTermination = { _ in
+                    task.cancel()
+                }
+            }
         case .completion:
             let prompt = createPrompt(from: request)
             CodeCompletionLogger.logger.logPrompt([(prompt, "user")])
-            return try await sendPrompt(prompt)
+            return AsyncStream<String> { continuation in
+                let task = Task {
+                    let result = try await sendPrompt(prompt)
+                    try Task.checkCancellation()
+                    continuation.yield(result)
+                    continuation.finish()
+                }
+                continuation.onTermination = { _ in
+                    task.cancel()
+                }
+            }
         }
     }
 }

Core/Sources/CodeCompletionService/TabbyService.swift

@@ -27,7 +27,7 @@ actor TabbyService {
 }
 
 extension TabbyService: CodeCompletionServiceType {
-    func getCompletion(_ request: PromptStrategy) async throws -> String {
+    func getCompletion(_ request: PromptStrategy) async throws -> AsyncStream<String> {
         let prefix = request.prefix.joined()
         let suffix = request.suffix.joined()
         let clipboard = request.relevantCodeSnippets.map(\.content).joined(separator: "\n\n")
@@ -46,7 +46,17 @@ extension TabbyService: CodeCompletionServiceType {
             (suffix, "suffix"),
             (clipboard, "clipboard"),
         ])
-        return try await send(requestBody)
+        return AsyncStream<String> { continuation in
+            let task = Task {
+                let result = try await send(requestBody)
+                try Task.checkCancellation()
+                continuation.yield(result)
+                continuation.finish()
+            }
+            continuation.onTermination = { _ in
+                task.cancel()
+            }
+        }
     }
 }