Reset ChatCompletionsStreamAPI to return AsyncThrowingStream

Author: intitni

Pro

@@ -128,9 +128,28 @@ typealias ChatCompletionsStreamAPIBuilder = (
 ) -> any ChatCompletionsStreamAPI
 
 protocol ChatCompletionsStreamAPI {
-    associatedtype CompletionSequence: AsyncSequence
-        where CompletionSequence.Element == ChatCompletionsStreamDataChunk
-    func callAsFunction() async throws -> CompletionSequence
+    func callAsFunction() async throws -> AsyncThrowingStream<ChatCompletionsStreamDataChunk, Error>
+}
+
+extension AsyncSequence {
+    func toStream() -> AsyncThrowingStream<Element, Error> {
+        AsyncThrowingStream { continuation in
+            let task = Task {
+                do {
+                    for try await element in self {
+                        continuation.yield(element)
+                    }
+                    continuation.finish()
+                } catch {
+                    continuation.finish(throwing: error)
+                }
+            }
+
+            continuation.onTermination = { _ in
+                task.cancel()
+            }
+        }
+    }
 }
 
 struct ChatCompletionsStreamDataChunk: Codable {
@@ -159,7 +178,13 @@ struct ChatCompletionsStreamDataChunk: Codable {
 
 // MARK: - Non Stream API
 
-typealias ChatCompletionsAPIBuilder = (String, ChatModel, URL, ChatCompletionsRequestBody, ChatGPTPrompt)
+typealias ChatCompletionsAPIBuilder = (
+    String,
+    ChatModel,
+    URL,
+    ChatCompletionsRequestBody,
+    ChatGPTPrompt
+)
     -> any ChatCompletionsAPI
 
 protocol ChatCompletionsAPI {

Tool/Sources/OpenAIService/APIs/ChatCompletionsAPIDefinition.swift

@@ -33,9 +33,9 @@ extension OllamaService: ChatCompletionsAPI {
 }
 
 extension OllamaService: ChatCompletionsStreamAPI {
-    typealias CompletionSequence = AsyncMapSequence<ResponseStream<OllamaService.ChatCompletionResponseChunk>, ChatCompletionsStreamDataChunk>
-
-    func callAsFunction() async throws -> CompletionSequence {
+    func callAsFunction() async throws
+        -> AsyncThrowingStream<ChatCompletionsStreamDataChunk, Swift.Error>
+    {
         let requestBody = ChatCompletionRequestBody(
             model: model.info.modelName,
             messages: requestBody.messages.map { message in
@@ -115,7 +115,7 @@ extension OllamaService: ChatCompletionsStreamAPI {
             )
         }
 
-        return sequence
+        return sequence.toStream()
     }
 }
 

Tool/Sources/OpenAIService/APIs/OlamaService.swift

@@ -73,14 +73,14 @@ public class ChatGPTService: ChatGPTServiceType {
         apiKey, model, endpoint, requestBody, prompt in
         switch model.format {
         case .googleAI:
-            return GoogleCompletionStreamAPI(
+            return GoogleAIService(
                 apiKey: apiKey,
                 model: model,
                 requestBody: requestBody,
                 prompt: prompt
             )
         case .openAI, .openAICompatible, .azureOpenAI:
-            return OpenAICompletionStreamAPI(
+            return OpenAIService(
                 apiKey: apiKey,
                 model: model,
                 endpoint: endpoint,
@@ -93,14 +93,14 @@ public class ChatGPTService: ChatGPTServiceType {
         apiKey, model, endpoint, requestBody, prompt in
         switch model.format {
         case .googleAI:
-            return GoogleCompletionAPI(
+            return GoogleAIService(
                 apiKey: apiKey,
                 model: model,
                 requestBody: requestBody,
                 prompt: prompt
             )
         case .openAI, .openAICompatible, .azureOpenAI:
-            return OpenAICompletionAPI(
+            return OpenAIService(
                 apiKey: apiKey,
                 model: model,
                 endpoint: endpoint,

Tool/Sources/OpenAIService/ChatGPTService.swift

@@ -73,7 +73,10 @@ public struct EmbeddingService {
         }
 
         guard response.statusCode == 200 else {
-            let error = try? JSONDecoder().decode(CompletionAPIError.self, from: result)
+            let error = try? JSONDecoder().decode(
+                OpenAIService.CompletionAPIError.self,
+                from: result
+            )
             throw error ?? ChatGPTServiceError
                 .otherError(String(data: result, encoding: .utf8) ?? "Unknown Error")
         }
@@ -124,7 +127,10 @@ public struct EmbeddingService {
         }
 
         guard response.statusCode == 200 else {
-            let error = try? JSONDecoder().decode(CompletionAPIError.self, from: result)
+            let error = try? JSONDecoder().decode(
+                OpenAIService.CompletionAPIError.self,
+                from: result
+            )
             throw error ?? ChatGPTServiceError
                 .otherError(String(data: result, encoding: .utf8) ?? "Unknown Error")
         }