Merge branch 'feature/custom-prompt-to-code-model' into develop

intitni · intitni · commit 202450cd5eae · 2023-09-26T19:00:49.000+08:00
diff --git a/Core/Sources/HostApp/AccountSettings/ChatModelManagement/ChatModelEditView.swift b/Core/Sources/HostApp/AccountSettings/ChatModelManagement/ChatModelEditView.swift
@@ -115,6 +115,13 @@ struct ChatModelEditView: View {
                 "Supports Function Calling",
                 isOn: viewStore.$supportsFunctionCalling
             )
+
+            Text(
+                "Function calling is required by some features, if this model doesn't support function calling, you should turn it off to avoid undefined behaviors."
+            )
+            .foregroundColor(.secondary)
+            .font(.callout)
+            .dynamicHeightTextInFormWorkaround()
         }
     }
 
diff --git a/Core/Sources/HostApp/FeatureSettings/ChatSettingsView.swift b/Core/Sources/HostApp/FeatureSettings/ChatSettingsView.swift
@@ -44,7 +44,7 @@ struct ChatSettingsView: View {
     var chatSettingsForm: some View {
         Form {
             Picker(
-                "Chat Feature Provider",
+                "Chat Model",
                 selection: $settings.defaultChatFeatureChatModelId
             ) {
                 if !settings.chatModels
@@ -63,7 +63,7 @@ struct ChatSettingsView: View {
             }
 
             Picker(
-                "Embedding Feature Provider",
+                "Embedding Model",
                 selection: $settings.defaultChatFeatureEmbeddingModelId
             ) {
                 if !settings.embeddingModels
diff --git a/Core/Sources/HostApp/FeatureSettings/PromptToCodeSettingsView.swift b/Core/Sources/HostApp/FeatureSettings/PromptToCodeSettingsView.swift
@@ -10,6 +10,13 @@ struct PromptToCodeSettingsView: View {
         var promptToCodeGenerateDescription
         @AppStorage(\.promptToCodeGenerateDescriptionInUserPreferredLanguage)
         var promptToCodeGenerateDescriptionInUserPreferredLanguage
+        @AppStorage(\.promptToCodeChatModelId)
+        var promptToCodeChatModelId
+        @AppStorage(\.promptToCodeEmbeddingModelId)
+        var promptToCodeEmbeddingModelId
+
+        @AppStorage(\.chatModels) var chatModels
+        @AppStorage(\.embeddingModels) var embeddingModels
         init() {}
     }
 
@@ -18,6 +25,50 @@ struct PromptToCodeSettingsView: View {
     var body: some View {
         VStack(alignment: .center) {
             Form {
+                Picker(
+                    "Chat Model",
+                    selection: $settings.promptToCodeChatModelId
+                ) {
+                    Text("Same as Chat Feature").tag("")
+                    
+                    if !settings.chatModels
+                        .contains(where: { $0.id == settings.promptToCodeChatModelId }),
+                        !settings.promptToCodeChatModelId.isEmpty
+                    {
+                        Text(
+                            (settings.chatModels.first?.name).map { "\($0) (Default)" }
+                                ?? "No Model Found"
+                        )
+                        .tag(settings.promptToCodeChatModelId)
+                    }
+
+                    ForEach(settings.chatModels, id: \.id) { chatModel in
+                        Text(chatModel.name).tag(chatModel.id)
+                    }
+                }
+
+                Picker(
+                    "Embedding Model",
+                    selection: $settings.promptToCodeEmbeddingModelId
+                ) {
+                    Text("Same as Chat Feature").tag("")
+                    
+                    if !settings.embeddingModels
+                        .contains(where: { $0.id == settings.promptToCodeEmbeddingModelId }),
+                        !settings.promptToCodeEmbeddingModelId.isEmpty
+                    {
+                        Text(
+                            (settings.embeddingModels.first?.name).map { "\($0) (Default)" }
+                                ?? "No Model Found"
+                        )
+                        .tag(settings.promptToCodeEmbeddingModelId)
+                    }
+
+                    ForEach(settings.embeddingModels, id: \.id) { embeddingModel in
+                        Text(embeddingModel.name).tag(embeddingModel.id)
+                    }
+                }
+
                 Toggle(isOn: $settings.promptToCodeGenerateDescription) {
                     Text("Generate Description")
                 }
diff --git a/Core/Sources/PromptToCodeService/OpenAIPromptToCodeService.swift b/Core/Sources/PromptToCodeService/OpenAIPromptToCodeService.swift
@@ -171,8 +171,9 @@ public final class OpenAIPromptToCodeService: PromptToCodeServiceType {
         What is your requirement?
         """
 
-        let configuration = UserPreferenceChatGPTConfiguration()
-            .overriding(.init(temperature: 0))
+        let configuration =
+            UserPreferenceChatGPTConfiguration(chatModelKey: \.promptToCodeChatModelId)
+                .overriding(.init(temperature: 0))
         let memory = AutoManagedChatGPTMemory(
             systemPrompt: systemPrompt,
             configuration: configuration,
diff --git a/Tool/Sources/OpenAIService/ChatGPTService.swift b/Tool/Sources/OpenAIService/ChatGPTService.swift
@@ -10,12 +10,18 @@ public protocol ChatGPTServiceType {
 }
 
 public enum ChatGPTServiceError: Error, LocalizedError {
+    case chatModelNotAvailable
+    case embeddingModelNotAvailable
     case endpointIncorrect
     case responseInvalid
     case otherError(String)
 
     public var errorDescription: String? {
         switch self {
+        case .chatModelNotAvailable:
+            return "Chat model is not available, please add a model in the settings."
+        case .embeddingModelNotAvailable:
+            return "Embedding model is not available, please add a model in the settings."
         case .endpointIncorrect:
             return "ChatGPT endpoint is incorrect"
         case .responseInvalid:
@@ -180,8 +186,12 @@ extension ChatGPTService {
 
     /// Send the memory as prompt to ChatGPT, with stream enabled.
     func sendMemory() async throws -> AsyncThrowingStream<StreamContent, Error> {
-        guard let url = URL(string: configuration.endpoint)
-        else { throw ChatGPTServiceError.endpointIncorrect }
+        guard let model = configuration.model else {
+            throw ChatGPTServiceError.chatModelNotAvailable
+        }
+        guard let url = URL(string: configuration.endpoint) else {
+            throw ChatGPTServiceError.endpointIncorrect
+        }
 
         await memory.refresh()
 
@@ -197,8 +207,6 @@ extension ChatGPTService {
         }
         let remainingTokens = await memory.remainingTokens
 
-        let model = configuration.model
-
         let requestBody = CompletionRequestBody(
             model: model.info.modelName,
             messages: messages,
@@ -287,8 +295,12 @@ extension ChatGPTService {
 
     /// Send the memory as prompt to ChatGPT, with stream disabled.
     func sendMemoryAndWait() async throws -> ChatMessage? {
-        guard let url = URL(string: configuration.endpoint)
-        else { throw ChatGPTServiceError.endpointIncorrect }
+        guard let model = configuration.model else {
+            throw ChatGPTServiceError.chatModelNotAvailable
+        }
+        guard let url = URL(string: configuration.endpoint) else {
+            throw ChatGPTServiceError.endpointIncorrect
+        }
 
         await memory.refresh()
 
@@ -304,8 +316,6 @@ extension ChatGPTService {
         }
         let remainingTokens = await memory.remainingTokens
 
-        let model = configuration.model
-
         let requestBody = CompletionRequestBody(
             model: model.info.modelName,
             messages: messages,
@@ -357,7 +367,7 @@ extension ChatGPTService {
     /// When a function call is detected, but arguments are not yet ready, we can call this
     /// to insert a message placeholder in memory.
     func prepareFunctionCall(_ call: ChatMessage.FunctionCall, messageId: String) async {
-        guard var function = functionProvider.function(named: call.name) else { return }
+        guard let function = functionProvider.function(named: call.name) else { return }
         let responseMessage = ChatMessage(
             id: messageId,
             role: .function,
@@ -380,7 +390,7 @@ extension ChatGPTService {
     ) async -> String {
         let messageId = messageId ?? uuidGenerator()
 
-        guard var function = functionProvider.function(named: call.name) else {
+        guard let function = functionProvider.function(named: call.name) else {
             return await fallbackFunctionCall(call, messageId: messageId)
         }
 
diff --git a/Tool/Sources/OpenAIService/Configuration/ChatGPTConfiguration.swift b/Tool/Sources/OpenAIService/Configuration/ChatGPTConfiguration.swift
@@ -4,7 +4,7 @@ import Preferences
 import Keychain
 
 public protocol ChatGPTConfiguration {
-    var model: ChatModel { get }
+    var model: ChatModel? { get }
     var temperature: Double { get }
     var apiKey: String { get }
     var stop: [String] { get }
@@ -15,11 +15,12 @@ public protocol ChatGPTConfiguration {
 
 public extension ChatGPTConfiguration {
     var endpoint: String {
-        model.endpoint
+        model?.endpoint ?? ""
     }
     
     var apiKey: String {
-        (try? Keychain.apiKey.get(model.info.apiKeyName)) ?? ""
+        guard let name = model?.info.apiKeyName else { return "" }
+        return (try? Keychain.apiKey.get(name)) ?? ""
     }
 
     func overriding(
diff --git a/Tool/Sources/OpenAIService/Configuration/EmbeddingConfiguration.swift b/Tool/Sources/OpenAIService/Configuration/EmbeddingConfiguration.swift
@@ -4,19 +4,20 @@ import Keychain
 import Preferences
 
 public protocol EmbeddingConfiguration {
-    var model: EmbeddingModel { get }
+    var model: EmbeddingModel? { get }
     var apiKey: String { get }
     var maxToken: Int { get }
     var dimensions: Int { get }
 }
 
 public extension EmbeddingConfiguration {
     var endpoint: String {
-        model.endpoint
+        model?.endpoint ?? ""
     }
     
     var apiKey: String {
-        (try? Keychain.apiKey.get(model.info.apiKeyName)) ?? ""
+        guard let name = model?.info.apiKeyName else { return "" }
+        return (try? Keychain.apiKey.get(name)) ?? ""
     }
 
     func overriding(
diff --git a/Tool/Sources/OpenAIService/Configuration/UserPreferenceChatGPTConfiguration.swift b/Tool/Sources/OpenAIService/Configuration/UserPreferenceChatGPTConfiguration.swift
@@ -3,19 +3,29 @@ import Foundation
 import Preferences
 
 public struct UserPreferenceChatGPTConfiguration: ChatGPTConfiguration {
+    public var chatModelKey: KeyPath<UserDefaultPreferenceKeys, PreferenceKey<String>>?
+    
     public var temperature: Double {
         min(max(0, UserDefaults.shared.value(for: \.chatGPTTemperature)), 2)
     }
 
-    public var model: ChatModel {
+    public var model: ChatModel? {
         let models = UserDefaults.shared.value(for: \.chatModels)
+        
+        if let chatModelKey {
+            let id = UserDefaults.shared.value(for: chatModelKey)
+            if let model = models.first(where: { $0.id == id }) {
+                return model
+            }
+        }
+        
         let id = UserDefaults.shared.value(for: \.defaultChatFeatureChatModelId)
         return models.first { $0.id == id }
-            ?? models.first ?? .init(id: "", name: "", format: .openAI, info: .init())
+            ?? models.first
     }
 
     public var maxTokens: Int {
-        model.info.maxTokens
+        model?.info.maxTokens ?? 0
     }
 
     public var stop: [String] {
@@ -30,7 +40,9 @@ public struct UserPreferenceChatGPTConfiguration: ChatGPTConfiguration {
         true
     }
 
-    public init() {}
+    public init(chatModelKey: KeyPath<UserDefaultPreferenceKeys, PreferenceKey<String>>? = nil) {
+        self.chatModelKey = chatModelKey
+    }
 }
 
 public class OverridingChatGPTConfiguration: ChatGPTConfiguration {
@@ -77,7 +89,7 @@ public class OverridingChatGPTConfiguration: ChatGPTConfiguration {
         overriding.temperature ?? configuration.temperature
     }
 
-    public var model: ChatModel {
+    public var model: ChatModel? {
         if let model = overriding.model { return model }
         let models = UserDefaults.shared.value(for: \.chatModels)
         guard let id = overriding.modelId,
diff --git a/Tool/Sources/OpenAIService/Configuration/UserPreferenceEmbeddingConfiguration.swift b/Tool/Sources/OpenAIService/Configuration/UserPreferenceEmbeddingConfiguration.swift
@@ -3,26 +3,40 @@ import Foundation
 import Preferences
 
 public struct UserPreferenceEmbeddingConfiguration: EmbeddingConfiguration {
-    public var model: EmbeddingModel {
+    public var embeddingModelKey: KeyPath<UserDefaultPreferenceKeys, PreferenceKey<String>>?
+
+    public var model: EmbeddingModel? {
         let models = UserDefaults.shared.value(for: \.embeddingModels)
+        
+        if let embeddingModelKey {
+            let id = UserDefaults.shared.value(for: embeddingModelKey)
+            if let model = models.first(where: { $0.id == id }) {
+                return model
+            }
+        }
+        
         let id = UserDefaults.shared.value(for: \.defaultChatFeatureEmbeddingModelId)
         return models.first { $0.id == id }
-            ?? models.first ?? .init(id: "", name: "", format: .openAI, info: .init())
+            ?? models.first
     }
 
     public var maxToken: Int {
-        model.info.maxTokens
+        model?.info.maxTokens ?? 0
     }
 
     public var dimensions: Int {
-        let dimensions = model.info.dimensions
+        let dimensions = model?.info.dimensions ?? 0
         if dimensions <= 0 {
             return 1536
         }
         return dimensions
     }
 
-    public init() {}
+    public init(
+        embeddingModelKey: KeyPath<UserDefaultPreferenceKeys, PreferenceKey<String>>? = nil
+    ) {
+        self.embeddingModelKey = embeddingModelKey
+    }
 }
 
 public class OverridingEmbeddingConfiguration<
@@ -55,7 +69,7 @@ public class OverridingEmbeddingConfiguration<
         self.configuration = configuration
     }
 
-    public var model: EmbeddingModel {
+    public var model: EmbeddingModel? {
         if let model = overriding.model { return model }
         let models = UserDefaults.shared.value(for: \.embeddingModels)
         guard let id = overriding.modelId,
@@ -67,7 +81,7 @@ public class OverridingEmbeddingConfiguration<
     public var maxToken: Int {
         overriding.maxTokens ?? configuration.maxToken
     }
-    
+
     public var dimensions: Int {
         overriding.dimensions ?? configuration.dimensions
     }
diff --git a/Tool/Sources/OpenAIService/EmbeddingService.swift b/Tool/Sources/OpenAIService/EmbeddingService.swift
@@ -41,10 +41,12 @@ public struct EmbeddingService {
     }
 
     public func embed(text: [String]) async throws -> EmbeddingResponse {
+        guard let model = configuration.model else {
+            throw ChatGPTServiceError.embeddingModelNotAvailable
+        }
         guard let url = URL(string: configuration.endpoint) else {
             throw ChatGPTServiceError.endpointIncorrect
         }
-        let model = configuration.model
         var request = URLRequest(url: url)
         request.httpMethod = "POST"
         let encoder = JSONEncoder()
@@ -90,10 +92,12 @@ public struct EmbeddingService {
     }
 
     public func embed(tokens: [[Int]]) async throws -> EmbeddingResponse {
+        guard let model = configuration.model else {
+            throw ChatGPTServiceError.embeddingModelNotAvailable
+        }
         guard let url = URL(string: configuration.endpoint) else {
             throw ChatGPTServiceError.endpointIncorrect
         }
-        let model = configuration.model
         var request = URLRequest(url: url)
         request.httpMethod = "POST"
         let encoder = JSONEncoder()
diff --git a/Tool/Sources/Preferences/Keys.swift b/Tool/Sources/Preferences/Keys.swift

Original file line number	Diff line number	Diff line change
`@@ -115,6 +115,13 @@ struct ChatModelEditView: View {`
`115`	`115`	`"Supports Function Calling",`
`116`	`116`	`isOn: viewStore.$supportsFunctionCalling`
`117`	`117`	`)`
	`118`	`+`
	`119`	`+ Text(`
	`120`	`+ "Function calling is required by some features, if this model doesn't support function calling, you should turn it off to avoid undefined behaviors."`
	`121`	`+ )`
	`122`	`+ .foregroundColor(.secondary)`
	`123`	`+ .font(.callout)`
	`124`	`+ .dynamicHeightTextInFormWorkaround()`
`118`	`125`	`}`
`119`	`126`	`}`
`120`	`127`