Adding gemini to the team

Author: jamesrochabrun

Package.resolved

@@ -19,6 +19,7 @@ let package = Package(
         // Dependencies declare other packages that this package depends on.
         .package(url: "https://github.com/jamesrochabrun/SwiftOpenAI", branch: "main"),
         .package(url: "https://github.com/jamesrochabrun/SwiftAnthropic", branch: "main"),
+        .package(url: "https://github.com/google/generative-ai-swift", branch: "main"),
     ],
     targets: [
         // Targets are the basic building blocks of a package, defining a module or a test suite.
@@ -27,7 +28,8 @@ let package = Package(
             name: "PolyAI",
             dependencies: [
                .product(name: "SwiftOpenAI", package: "SwiftOpenAI"),
-               .product(name: "SwiftAnthropic", package: "SwiftAnthropic")
+               .product(name: "SwiftAnthropic", package: "SwiftAnthropic"),
+               .product(name: "GoogleGenerativeAI", package: "generative-ai-swift")
             ]),
         .testTarget(
             name: "PolyAITests",

Package.swift

@@ -12,8 +12,10 @@ struct ApiKeyIntroView: View {
 
    @State private var anthropicAPIKey = ""
    @State private var openAIAPIKey = ""
+   @State private var geminiAPIKey = ""
    @State private var anthropicConfigAdded: Bool = false
    @State private var openAIConfigAdded: Bool = false
+   @State private var geminiConfigAdded: Bool = false
 
    @State private var configurations: [LLMConfiguration] = []
 
@@ -26,34 +28,24 @@ struct ApiKeyIntroView: View {
          VStack {
             Spacer()
             VStack(spacing: 24) {
-               VStack(alignment: .leading) {
-                  HStack {
-                     TextField("Enter Anthropic API Key", text: $anthropicAPIKey)
-                     Button {
-                        configurations.append(.anthropic(apiKey: anthropicAPIKey))
-                        anthropicConfigAdded = true
-                     } label: {
-                        Image(systemName: "plus")
-                     }
-                     .disabled(anthropicAPIKey.isEmpty)
+               LLMConfigurationView(
+                  provider: "Anthropic",
+                  configurationAdded: $anthropicConfigAdded,
+                  apiKey: $anthropicAPIKey) {
+                     configurations.append(.anthropic(apiKey: anthropicAPIKey))
                   }
-                  Text("Anthropic added to PolyAI 🚀")
-                     .opacity(anthropicConfigAdded ? 1 : 0)
-               }
-               VStack(alignment: .leading) {
-                  HStack {
-                     TextField("Enter OpenAI API Key", text: $openAIAPIKey)
-                     Button {
-                        configurations.append(.openAI(.api(key: openAIAPIKey)))
-                        openAIConfigAdded = true
-                     } label: {
-                        Image(systemName: "plus")
-                     }
-                     .disabled(openAIAPIKey.isEmpty)
+               LLMConfigurationView(
+                  provider: "OpenAI",
+                  configurationAdded: $openAIConfigAdded,
+                  apiKey: $openAIAPIKey) {
+                     configurations.append(.openAI(.api(key: openAIAPIKey)))
+                  }
+               LLMConfigurationView(
+                  provider: "Gemini",
+                  configurationAdded: $geminiConfigAdded,
+                  apiKey: $geminiAPIKey) {
+                     configurations.append(.gemini(apiKey: geminiAPIKey))
                   }
-                  Text("OpenAI added to PolyAI 🚀")
-                     .opacity(openAIConfigAdded ? 1 : 0)
-               }
             }
             .buttonStyle(.bordered)
             .padding()
@@ -76,6 +68,32 @@ struct ApiKeyIntroView: View {
    }
 }
 
+struct LLMConfigurationView: View {
+   
+   let provider: String
+   @Binding var configurationAdded: Bool
+   @Binding var apiKey: String
+   let addConfig: () -> Void
+   
+   
+   var body: some View {
+      VStack(alignment: .leading) {
+         HStack {
+            TextField("Enter \(provider) API Key", text: $apiKey)
+            Button {
+               addConfig()
+               configurationAdded = true
+            } label: {
+               Image(systemName: "plus")
+            }
+            .disabled(apiKey.isEmpty)
+         }
+         Text("\(provider) added to PolyAI 🚀")
+            .opacity(configurationAdded ? 1 : 0)
+      }
+   }
+}
+
 #Preview {
    ApiKeyIntroView()
 }

PolyAIExample/PolyAIExample.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved

@@ -23,6 +23,7 @@ struct MessageDemoView: View {
    enum LLM {
       case openAI
       case anthropic
+      case gemini
    }
 
    var body: some View {
@@ -74,6 +75,11 @@ struct MessageDemoView: View {
                         .init(role: .user, content: prompt)
                      ],
                      maxTokens: 1024)
+               case .gemini:
+                  parameters = .gemini(
+                     model: "gemini-pro", messages: [
+                        .init(role: .user, content: prompt)
+                  ], maxTokens: 2000)
                }
                try await observable.streamMessage(parameters: parameters)
             }
@@ -89,6 +95,7 @@ struct MessageDemoView: View {
       Picker("Options", selection: $selectedSegment) {
          Text("Anthropic").tag(LLM.anthropic)
          Text("OpenAI").tag(LLM.openAI)
+         Text("Gemini").tag(LLM.gemini)
       }
       .pickerStyle(SegmentedPickerStyle())
       .padding()

PolyAIExample/PolyAIExample/ApiKeysIntroView.swift

@@ -26,11 +26,19 @@ public enum LLMParameter {
    ///   - maxTokens: The maximum number of tokens to generate.
    case anthropic(model: SwiftAnthropic.Model, messages: [LLMMessage], maxTokens: Int)
 
+   /// Represents a configuration for interacting with Gemini's models.
+   /// - Parameters:
+   ///   - model: The specific model of Gemini to use.
+   ///   - messages: An array of messages to send to the model.
+   ///   - maxTokens: The maximum number of tokens to generate.
+   case gemini(model: String, messages: [LLMMessage], maxTokens: Int)
+   
    /// A computed property that returns the name of the LLM service based on the case.
    var llmService: String {
       switch self {
       case .openAI: return "OpenAI"
       case .anthropic: return "Anthropic"
+      case .gemini: return "Gemini"
       }
    }
 }

PolyAIExample/PolyAIExample/MessageDemoView.swift

@@ -0,0 +1,45 @@
+//
+//  LLMMessageResponse+Gemini.swift
+//
+//
+//  Created by James Rochabrun on 5/3/24.
+//
+
+import Foundation
+import GoogleGenerativeAI
+
+// MARK: Gemini
+
+extension GenerateContentResponse: LLMMessageResponse {
+  
+   public var id: String {
+      UUID().uuidString
+   }
+   
+   public var model: String {
+      ""
+   }
+   
+   public var role: String {
+      candidates.first?.content.role ?? "user"
+   }
+   
+   public var createdAt: Int? {
+      nil
+   }
+   
+   public var contentDescription: String {
+      text ?? ""
+   }
+   
+   public var usageMetrics: UsageMetrics {
+      ChatUsageMetrics(
+         inputTokens: usageMetadata?.promptTokenCount ?? 0,
+         outputTokens: usageMetadata?.candidatesTokenCount ?? 0,
+         totalTokens: usageMetadata?.totalTokenCount ?? 0)
+   }
+   
+   public var tools: [ToolUsage] {
+      []
+   }   
+}

Sources/PolyAI/Interfaces/Parameters/LLMParameter.swift

@@ -0,0 +1,18 @@
+//
+//  LLMMessageStreamResponse+Gemini.swift
+//  
+//
+//  Created by James Rochabrun on 5/4/24.
+//
+
+import Foundation
+import GoogleGenerativeAI
+
+// MARK: Gemini
+
+extension GenerateContentResponse: LLMMessageStreamResponse {
+   
+   public var content: String? {
+      text
+   }
+}

Sources/PolyAI/Interfaces/Response/Message/LLMMessageResponse+Gemini.swift

@@ -6,17 +6,29 @@
 //
 
 import Foundation
+import GoogleGenerativeAI
 import SwiftAnthropic
 import SwiftOpenAI
 
+// MARK: Error
+
 enum PolyAIError: Error {
    case missingLLMConfiguration(String)
 }
 
+// MARK: Gemini specific
+
+extension GenerativeModel {
+   public struct Configuration {
+      let apiKey: String
+   }
+}
+
 struct DefaultPolyAIService: PolyAIService {
 
    private var openAIService: OpenAIService?
    private var anthropicService: AnthropicService?
+   private var gemini: GenerativeModel.Configuration?
 
    init(configurations: [LLMConfiguration])
    {
@@ -36,6 +48,9 @@ struct DefaultPolyAIService: PolyAIService {
 
          case .anthropic(let apiKey, let configuration):
             anthropicService = AnthropicServiceFactory.service(apiKey: apiKey, configuration: configuration)
+            
+         case .gemini(let apiKey):
+            gemini = .init(apiKey: apiKey)
          }
       }
    }
@@ -54,6 +69,7 @@ struct DefaultPolyAIService: PolyAIService {
          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)
          return try await openAIService.startChat(parameters: messageParameter)
+         
       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")
@@ -71,6 +87,24 @@ struct DefaultPolyAIService: PolyAIService {
          let systemMessage  = messages.first { $0.role == "system" }
          let messageParameter = MessageParameter(model: model, messages: messageParams, maxTokens: maxTokens, system: systemMessage?.content, stream: false)
          return try await anthropicService.createMessage(messageParameter)
+         
+      case .gemini(let model, let messages, let maxTokens):
+         guard let gemini else {
+            throw PolyAIError.missingLLMConfiguration("You Must provide a valid configuration for the \(parameter.llmService) API")
+         }
+         let systemInstruction: ModelContent?
+         if let systemMessage = messages.first(where: { message in
+            message.role == "system"
+         })?.content {
+            systemInstruction = ModelContent(parts: [.text(systemMessage)])
+         } else {
+            systemInstruction = nil
+         }
+         let generativeModel = GenerativeModel(name: model, apiKey: gemini.apiKey, generationConfig: .init(GenerationConfig(maxOutputTokens: maxTokens)), systemInstruction: systemInstruction)
+         let userMessage = messages.first { message in
+            message.role == "user"
+         }?.content ?? ""
+         return try await generativeModel.generateContent(userMessage)
       }
    }
 
@@ -85,8 +119,11 @@ struct DefaultPolyAIService: PolyAIService {
          }
          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")
@@ -105,6 +142,25 @@ struct DefaultPolyAIService: PolyAIService {
          let messageParameter = MessageParameter(model: model, messages: messageParams, maxTokens: maxTokens, system: systemMessage?.content)
          let stream = try await anthropicService.streamMessage(messageParameter)
          return try mapToLLMMessageStreamResponse(stream: stream)
+         
+      case .gemini(let model, let messages, let maxTokens):
+         guard let gemini else {
+            throw PolyAIError.missingLLMConfiguration("You Must provide a valid configuration for the \(parameter.llmService) API")
+         }
+         let systemInstruction: ModelContent?
+         if let systemMessage = messages.first(where: { message in
+            message.role == "system"
+         })?.content {
+            systemInstruction = ModelContent(parts: [.text(systemMessage)])
+         } else {
+            systemInstruction = nil
+         }
+         let generativeModel = GenerativeModel(name: model, apiKey: gemini.apiKey, generationConfig: .init(GenerationConfig(maxOutputTokens: maxTokens)), systemInstruction: systemInstruction)
+         let userMessage = messages.first { message in
+            message.role == "user"
+         }?.content ?? ""
+         let stream = generativeModel.generateContentStream(userMessage)
+         return try mapToLLMMessageStreamResponse(stream: stream)
       }
    }
 

Sources/PolyAI/Interfaces/Response/Stream/LLMMessageStreamResponse+Gemini.swift

@@ -41,6 +41,11 @@ public enum LLMConfiguration {
    ///   - apiKey: The API key for authenticating requests to Anthropic.
    ///   - configuration: The URLSession configuration to use for network requests. Defaults to `.default`.
    case anthropic(apiKey: String, configuration: URLSessionConfiguration = .default)
+   
+   /// Configuration for accessing Gemini's API.
+   /// - Parameters:
+   ///   - apiKey: The API key for authenticating requests to Gemini.
+   case gemini(apiKey: String)
 }
 
 /// Defines the interface for a service that interacts with Large Language Models (LLMs).