Support for stream

Author: jamesrochabrun

Sources/PolyAI/Interfaces/Parameters/LLMMessageParameter.swift

@@ -7,8 +7,31 @@
 
 import Foundation
 
+// MARK: LLMMessageParameter
+
 public protocol LLMMessageParameter {
 
    var role: String { get }
    var content: String { get }
 }
+
+// MARK: LLMMessage
+
+public struct LLMMessage: LLMMessageParameter {
+   
+   public var role: String
+   public var content: String
+   
+   public enum Role: String {
+      case user
+      case assistant
+   }
+   
+   public init(
+      role: Role,
+      content: String)
+   {
+      self.role = role.rawValue
+      self.content = content
+   }
+}

Sources/PolyAI/Interfaces/Parameters/LLMParameter.swift

@@ -6,13 +6,15 @@
 //
 
 import Foundation
+import SwiftOpenAI
+import SwiftAnthropic
 
 public enum LLMParameter {
 
-   case openAI(model: String, messages: [LLMMessageParameter], maxTokens: Int? = nil, stream: Bool? = nil)
-   case anthropic(model: String, messages: [LLMMessageParameter], maxTokens: Int, stream: Bool)
+   case openAI(model: SwiftOpenAI.Model, messages: [LLMMessage], maxTokens: Int? = nil)
+   case anthropic(model: SwiftAnthropic.Model, messages: [LLMMessage], maxTokens: Int)
 
-   var llm: String {
+   var llmService: String {
       switch self {
       case .openAI: return "OpenAI"
       case .anthropic: return "Anthropic"

Sources/PolyAI/Interfaces/Response/LLMMessageResponse.swift

@@ -36,20 +36,12 @@ struct ChatUsageMetrics: UsageMetrics {
 public protocol ToolUsage {
    var toolId: String? { get }
    var toolName: String { get }
-   var toolInput: [String: String]? { get } // Assuming tools might have inputs. Adjust as necessary.
+   var toolInput: [String: String]? { get }
 }
 
 // MARK: OpenAI
 
 extension ChatCompletionObject: LLMMessageResponse {
-   public var tools: [ToolUsage] {
-      []
-   }
-   
-   public var role: String {
-      choices.first?.message.role ?? "unknown"
-   }
-   
    public var createdAt: Int? {
       created
    }
@@ -58,9 +50,18 @@ extension ChatCompletionObject: LLMMessageResponse {
       choices.first?.message.content ?? ""
    }
 
+   
    public var usageMetrics: UsageMetrics {
       ChatUsageMetrics(inputTokens: usage.promptTokens, outputTokens: usage.completionTokens, totalTokens: usage.totalTokens)
    }
+   
+   public var tools: [ToolUsage] {
+      []
+   }
+   
+   public var role: String {
+      choices.first?.message.role ?? "unknown"
+   }
 }
 
 // MARK: Anthropic

Sources/PolyAI/Interfaces/Response/LLMMessageStreamResponse.swift

@@ -1,8 +1,35 @@
 //
-//  File.swift
+//  LLMMessageStreamResponse.swift
 //  
 //
 //  Created by James Rochabrun on 4/14/24.
 //
 
 import Foundation
+import SwiftAnthropic
+import SwiftOpenAI
+
+// MARK: Interface
+
+public protocol LLMMessageStreamResponse {
+   var content: String? { get }
+}
+
+// MARK: OpenAI
+
+extension ChatCompletionChunkObject: LLMMessageStreamResponse {
+   
+   public var content: String? {
+      choices.first?.delta.content
+   }
+}
+
+// MARK: Anthropic
+
+extension MessageStreamResponse: LLMMessageStreamResponse {
+   
+   public var content: String? {
+      delta?.text
+   }
+}
+

Sources/PolyAI/Service/DefaultPolyAIService.swift

@@ -24,8 +24,8 @@ struct DefaultPolyAIService: PolyAIService {
          switch configuration {
          case .openAI(let apiKey, let organizationID, let configuration, let decoder):
             openAIService = OpenAIServiceFactory.service(apiKey: apiKey, organizationID: organizationID, configuration: configuration, decoder: decoder)
-         case .anthropic(let apiKey, let apiVersion, let configuration):
-            anthropicService = AnthropicServiceFactory.service(apiKey: apiKey, apiVersion: apiVersion, configuration: configuration)
+         case .anthropic(let apiKey, let configuration):
+            anthropicService = AnthropicServiceFactory.service(apiKey: apiKey, configuration: configuration)
          }
       }
    }
@@ -37,20 +37,62 @@ struct DefaultPolyAIService: PolyAIService {
       async throws -> LLMMessageResponse
    {
       switch parameter {
-      case .openAI(let model, let messages, let maxTokens, _):
+      case .openAI(let model, let messages, let maxTokens):
          guard let openAIService else {
-            throw PolyAIError.missingLLMConfiguration("You Must provide a valid configuration for the \(parameter.llm) API")
+            throw PolyAIError.missingLLMConfiguration("You Must provide a valid configuration for the \(parameter.llmService) API")
          }
          let messageParams: [SwiftOpenAI.ChatCompletionParameters.Message] = messages.map { .init(role: .init(rawValue: $0.role) ?? .user, content: .text($0.content)) }
-         let messageParameter = ChatCompletionParameters(messages: messageParams, model: .custom(model), maxTokens: maxTokens)
+         let messageParameter = ChatCompletionParameters(messages: messageParams, model: model, maxTokens: maxTokens)
          return try await openAIService.startChat(parameters: messageParameter)
-      case .anthropic(let model, let messages, let maxTokens, _):
+      case .anthropic(let model, let messages, let maxTokens):
          guard let anthropicService else {
-            throw PolyAIError.missingLLMConfiguration("You Must provide a valid configuration for the \(parameter.llm) API")
+            throw PolyAIError.missingLLMConfiguration("You Must provide a valid configuration for the \(parameter.llmService) API")
          }
          let messageParams: [SwiftAnthropic.MessageParameter.Message] = messages.map { MessageParameter.Message(role: SwiftAnthropic.MessageParameter.Message.Role(rawValue: $0.role) ?? .user, content: .text($0.content)) }
-         let messageParameter = MessageParameter(model: .other(model), messages: messageParams, maxTokens: maxTokens, stream: false)
+         let messageParameter = MessageParameter(model: model, messages: messageParams, maxTokens: maxTokens, stream: false)
          return try await anthropicService.createMessage(messageParameter, beta: nil)
       }
    }
+   
+   func streamMessage(
+      _ parameter: LLMParameter)
+      async throws -> AsyncThrowingStream<LLMMessageStreamResponse, Error>
+   {
+      switch parameter {
+      case .openAI(let model, let messages, let maxTokens):
+         guard let openAIService else {
+            throw PolyAIError.missingLLMConfiguration("You Must provide a valid configuration for the \(parameter.llmService) API")
+         }
+         let messageParams: [SwiftOpenAI.ChatCompletionParameters.Message] = messages.map { .init(role: .init(rawValue: $0.role) ?? .user, content: .text($0.content)) }
+         let messageParameter = ChatCompletionParameters(messages: messageParams, model: model, maxTokens: maxTokens)
+         let stream = try await openAIService.startStreamedChat(parameters: messageParameter)
+         return try mapToLLMMessageStreamResponse(stream: stream)
+      case .anthropic(let model, let messages, let maxTokens):
+         guard let anthropicService else {
+            throw PolyAIError.missingLLMConfiguration("You Must provide a valid configuration for the \(parameter.llmService) API")
+         }
+         let messageParams: [SwiftAnthropic.MessageParameter.Message] = messages.map { MessageParameter.Message(role: SwiftAnthropic.MessageParameter.Message.Role(rawValue: $0.role) ?? .user, content: .text($0.content)) }
+         let messageParameter = MessageParameter(model: model, messages: messageParams, maxTokens: maxTokens, stream: false)
+         let stream = try await anthropicService.streamMessage(messageParameter, beta: nil)
+         return try mapToLLMMessageStreamResponse(stream: stream)
+      }
+   }
+   
+   private func mapToLLMMessageStreamResponse<T: LLMMessageStreamResponse>(stream: AsyncThrowingStream<T, Error>)
+      throws -> AsyncThrowingStream<LLMMessageStreamResponse, Error>
+   {
+      let mappedStream = AsyncThrowingStream<LLMMessageStreamResponse, Error> { continuation in
+         Task {
+            do {
+               for try await chunk in stream {
+                  continuation.yield(chunk)
+               }
+               continuation.finish()
+            } catch {
+               continuation.finish(throwing: error)
+            }
+         }
+      }
+      return mappedStream
+   }
 }

Sources/PolyAI/Service/PolyAIService.swift

@@ -8,8 +8,8 @@
 import Foundation
 
 public enum LLMConfiguration {
-   case openAI(apiKey: String, organizationID: String?, configuration: URLSessionConfiguration = .default, decoder: JSONDecoder = .init())
-   case anthropic(apiKey: String, apiVersion: String, configuration: URLSessionConfiguration = .default)
+   case openAI(apiKey: String, organizationID: String? = nil, configuration: URLSessionConfiguration = .default, decoder: JSONDecoder = .init())
+   case anthropic(apiKey: String, configuration: URLSessionConfiguration = .default)
 }
 
 public protocol PolyAIService {
@@ -21,4 +21,8 @@ public protocol PolyAIService {
    func createMessage(
       _ parameter: LLMParameter)
    async throws -> LLMMessageResponse
+   
+   func streamMessage(
+      _ parameter: LLMParameter)
+   async throws -> AsyncThrowingStream<LLMMessageStreamResponse, Error>
 }