Merge branch 'feature/google-gemini-support' into develop

Author: intitni

Core/Sources/HostApp/AccountSettings/ChatModelManagement/ChatModelEdit.swift

@@ -114,15 +114,26 @@ struct ChatModelEdit: ReducerProtocol {
                 return .none
 
             case .checkSuggestedMaxTokens:
-                guard state.format == .openAI,
-                      let knownModel = ChatGPTModel(rawValue: state.modelName)
-                else {
+                switch state.format {
+                case .openAI:
+                    if let knownModel = ChatGPTModel(rawValue: state.modelName) {
+                        state.suggestedMaxTokens = knownModel.maxToken
+                    } else {
+                        state.suggestedMaxTokens = nil
+                    }
+                    return .none
+                case .googleAI:
+                    if let knownModel = GoogleGenerativeAIModel(rawValue: state.modelName) {
+                        state.suggestedMaxTokens = knownModel.maxToken
+                    } else {
+                        state.suggestedMaxTokens = nil
+                    }
+                    return .none
+                default:
                     state.suggestedMaxTokens = nil
                     return .none
                 }
-                state.suggestedMaxTokens = knownModel.maxToken
-                return .none
-
+                
             case .apiKeySelection:
                 return .none
 

Core/Sources/HostApp/AccountSettings/ChatModelManagement/ChatModelEditView.swift

@@ -22,6 +22,8 @@ struct ChatModelEditView: View {
                             azureOpenAI
                         case .openAICompatible:
                             openAICompatible
+                        case .googleAI:
+                            googleAI
                         }
                     }
                 }
@@ -88,6 +90,8 @@ struct ChatModelEditView: View {
                             Text("Azure OpenAI").tag(format)
                         case .openAICompatible:
                             Text("OpenAI Compatible").tag(format)
+                        case .googleAI:
+                            Text("Google Generative AI").tag(format)
                         }
                     }
                 },
@@ -269,6 +273,35 @@ struct ChatModelEditView: View {
         maxTokensTextField
         supportsFunctionCallingToggle
     }
+    
+    @ViewBuilder
+    var googleAI: some View {
+        apiKeyNamePicker
+        
+        WithViewStore(
+            store,
+            removeDuplicates: { $0.modelName == $1.modelName }
+        ) { viewStore in
+            TextField("Model Name", text: viewStore.$modelName)
+                .overlay(alignment: .trailing) {
+                    Picker(
+                        "",
+                        selection: viewStore.$modelName,
+                        content: {
+                            if GoogleGenerativeAIModel(rawValue: viewStore.state.modelName) == nil {
+                                Text("Custom Model").tag(viewStore.state.modelName)
+                            }
+                            ForEach(GoogleGenerativeAIModel.allCases, id: \.self) { model in
+                                Text(model.rawValue).tag(model.rawValue)
+                            }
+                        }
+                    )
+                    .frame(width: 20)
+                }
+        }
+
+        maxTokensTextField
+    }
 }
 
 #Preview("OpenAI") {

Core/Sources/HostApp/AccountSettings/ChatModelManagement/ChatModelManagement.swift

@@ -10,6 +10,7 @@ extension ChatModel: ManageableAIModel {
         case .openAI: return "OpenAI"
         case .azureOpenAI: return "Azure OpenAI"
         case .openAICompatible: return "OpenAI Compatible"
+        case .googleAI: return "Google Generative AI"
         }
     }
 

Core/Sources/PromptToCodeService/OpenAIPromptToCodeService.swift

@@ -222,7 +222,7 @@ extension OpenAIPromptToCodeService {
     {
         func extractCodeFromMarkdown(_ markdown: String) -> (code: String, endIndex: Int)? {
             let codeBlockRegex = try! NSRegularExpression(
-                pattern: #"```(?:\w+)?[\n]([\s\S]+?)[\n]```"#,
+                pattern: #"```(?:\w+)?\R([\s\S]+?)\R```"#,
                 options: .dotMatchesLineSeparators
             )
             let range = NSRange(markdown.startIndex..<markdown.endIndex, in: markdown)
@@ -232,7 +232,7 @@ extension OpenAIPromptToCodeService {
             }
 
             let incompleteCodeBlockRegex = try! NSRegularExpression(
-                pattern: #"```(?:\w+)?[\n]([\s\S]+?)$"#,
+                pattern: #"```(?:\w+)?\R([\s\S]+?)$"#,
                 options: .dotMatchesLineSeparators
             )
             let range2 = NSRange(markdown.startIndex..<markdown.endIndex, in: markdown)

Core/Sources/SuggestionWidget/FeatureReducers/PromptToCode.swift

@@ -128,7 +128,7 @@ public struct PromptToCode: ReducerProtocol {
         case revertButtonTapped
         case stopRespondingButtonTapped
         case modifyCodeFinished
-        case modifyCodeTrunkReceived(code: String, description: String)
+        case modifyCodeChunkReceived(code: String, description: String)
         case modifyCodeFailed(error: String)
         case modifyCodeCancelled
         case cancelButtonTapped
@@ -189,7 +189,7 @@ public struct PromptToCode: ReducerProtocol {
                         )
                         for try await fragment in stream {
                             try Task.checkCancellation()
-                            await send(.modifyCodeTrunkReceived(
+                            await send(.modifyCodeChunkReceived(
                                 code: fragment.code,
                                 description: fragment.description
                             ))
@@ -221,7 +221,7 @@ public struct PromptToCode: ReducerProtocol {
                 promptToCodeService.stopResponding()
                 return .cancel(id: CancellationKey.modifyCode(state.id))
 
-            case let .modifyCodeTrunkReceived(code, description):
+            case let .modifyCodeChunkReceived(code, description):
                 state.code = code
                 state.description = description
                 return .none

Pro

@@ -63,6 +63,7 @@ let package = Package(
         .package(url: "https://github.com/apple/swift-syntax.git", branch: "main"),
         .package(url: "https://github.com/GottaGetSwifty/CodableWrappers", from: "2.0.7"),
         .package(url: "https://github.com/krzyzanowskim/STTextView", from: "0.8.21"),
+        .package(url: "https://github.com/google/generative-ai-swift", from: "0.4.4"),
 
         // TreeSitter
         .package(url: "https://github.com/intitni/SwiftTreeSitter.git", branch: "main"),
@@ -130,6 +131,7 @@ let package = Package(
             name: "TokenEncoder",
             dependencies: [
                 .product(name: "Tiktoken", package: "Tiktoken"),
+                .product(name: "GoogleGenerativeAI", package: "generative-ai-swift"),
             ],
             resources: [
                 .copy("Resources/cl100k_base.tiktoken"),
@@ -313,6 +315,7 @@ let package = Package(
                 "Keychain",
                 .product(name: "JSONRPC", package: "JSONRPC"),
                 .product(name: "AsyncAlgorithms", package: "swift-async-algorithms"),
+                .product(name: "GoogleGenerativeAI", package: "generative-ai-swift"),
                 .product(
                     name: "ComposableArchitecture",
                     package: "swift-composable-architecture"

Tool/Package.swift

@@ -20,6 +20,7 @@ public struct ChatModel: Codable, Equatable, Identifiable {
         case openAI
         case azureOpenAI
         case openAICompatible
+        case googleAI
     }
 
     public struct Info: Codable, Equatable {
@@ -69,6 +70,10 @@ public struct ChatModel: Codable, Equatable, Identifiable {
             let version = "2023-07-01-preview"
             if baseURL.isEmpty { return "" }
             return "\(baseURL)/openai/deployments/\(deployment)/chat/completions?api-version=\(version)"
+        case .googleAI:
+            let baseURL = info.baseURL
+            if baseURL.isEmpty { return "https://generativelanguage.googleapis.com/v1" }
+            return "\(baseURL)/v1/chat/completions"
         }
     }
 }

Tool/Sources/AIModel/ChatModel.swift

@@ -38,9 +38,9 @@ public struct OpenAIChat: ChatModel {
         if stream {
             let stream = try await service.send(content: "")
             var message = ""
-            for try await trunk in stream {
-                message.append(trunk)
-                callbackManagers.send(CallbackEvents.LLMDidProduceNewToken(info: trunk))
+            for try await chunk in stream {
+                message.append(chunk)
+                callbackManagers.send(CallbackEvents.LLMDidProduceNewToken(info: chunk))
             }
             return await memory.history.last ?? .init(role: .assistant, content: "")
         } else {

Tool/Sources/LangChain/ChatModel/OpenAIChat.swift

@@ -24,9 +24,9 @@ public extension TextSplitter {
         let paddingLength = texts.count - metadata.count
         let metadata = metadata + .init(repeating: [:], count: paddingLength)
         for (text, metadata) in zip(texts, metadata) {
-            let trunks = try await split(text: text)
-            for trunk in trunks {
-                let document = Document(pageContent: trunk, metadata: metadata)
+            let chunks = try await split(text: text)
+            for chunk in chunks {
+                let document = Document(pageContent: chunk, metadata: metadata)
                 documents.append(document)
             }
         }