Support Azure OpenAI as chat provider

intitni · intitni · commit a707645d86e9 · 2023-05-25T21:42:24.000+08:00
diff --git a/Core/Sources/HostApp/AccountSettings/AzureView.swift b/Core/Sources/HostApp/AccountSettings/AzureView.swift
@@ -0,0 +1,70 @@
+import AppKit
+import Client
+import OpenAIService
+import Preferences
+import SuggestionModel
+import SwiftUI
+
+final class AzureViewSettings: ObservableObject {
+    @AppStorage(\.azureOpenAIAPIKey) var azureOpenAIAPIKey: String
+    @AppStorage(\.azureOpenAIBaseURL) var azureOpenAIBaseURL: String
+    @AppStorage(\.azureChatGPTDeployment) var azureChatGPTDeployment: String
+    init() {}
+}
+
+struct AzureView: View {
+    @Environment(\.toast) var toast
+    @State var isTesting = false
+    @StateObject var settings = AzureViewSettings()
+
+    var body: some View {
+        Form {
+            SecureField(text: $settings.azureOpenAIAPIKey, prompt: Text("")) {
+                Text("OpenAI Service API Key")
+            }
+            .textFieldStyle(.roundedBorder)
+
+            TextField(
+                text: $settings.azureOpenAIBaseURL,
+                prompt: Text("https://XXXXXX.openai.azure.com")
+            ) {
+                Text("OpenAI Service Base URL")
+            }.textFieldStyle(.roundedBorder)
+
+            HStack {
+                TextField(
+                    text: $settings.azureChatGPTDeployment,
+                    prompt: Text("")
+                ) {
+                    Text("Chat Model Deployment Name")
+                }.textFieldStyle(.roundedBorder)
+
+                Button("Test") {
+                    Task { @MainActor in
+                        isTesting = true
+                        defer { isTesting = false }
+                        do {
+                            let reply = try await ChatGPTService(designatedProvider: .azureOpenAI)
+                                .sendAndWait(content: "Hello", summary: nil)
+                            toast(Text("ChatGPT replied: \(reply ?? "N/A")"), .info)
+                        } catch {
+                            toast(Text(error.localizedDescription), .error)
+                        }
+                    }
+                }
+                .disabled(isTesting)
+            }
+        }
+    }
+}
+
+struct AzureView_Previews: PreviewProvider {
+    static var previews: some View {
+        VStack(alignment: .leading, spacing: 8) {
+            AzureView()
+        }
+        .frame(height: 800)
+        .padding(.all, 8)
+    }
+}
+
diff --git a/Core/Sources/HostApp/AccountSettings/OpenAIView.swift b/Core/Sources/HostApp/AccountSettings/OpenAIView.swift
@@ -52,7 +52,7 @@ struct OpenAIView: View {
                 Button("Test") {
                     Task {
                         do {
-                            let reply = try await ChatGPTService()
+                            let reply = try await ChatGPTService(designatedProvider: .openAI)
                                 .sendAndWait(content: "Hello", summary: nil)
                             toast(Text("ChatGPT replied: \(reply ?? "N/A")"), .info)
                         } catch {
diff --git a/Core/Sources/HostApp/ServiceView.swift b/Core/Sources/HostApp/ServiceView.swift
@@ -30,6 +30,15 @@ struct ServiceView: View {
                 subtitle: "Chat, Prompt to Code",
                 image: "globe"
             )
+            
+            ScrollView {
+                AzureView().padding()
+            }.sidebarItem(
+                tag: 3,
+                title: "Azure",
+                subtitle: "Chat, Prompt to Code",
+                image: "globe"
+            )
         }
     }
 }
diff --git a/Core/Sources/OpenAIService/ChatGPTService.swift b/Core/Sources/OpenAIService/ChatGPTService.swift
@@ -17,13 +17,16 @@ public protocol ChatGPTServiceType: ObservableObject {
 public enum ChatGPTServiceError: Error, LocalizedError {
     case endpointIncorrect
     case responseInvalid
+    case otherError(String)
 
     public var errorDescription: String? {
         switch self {
         case .endpointIncorrect:
             return "ChatGPT endpoint is incorrect"
         case .responseInvalid:
             return "Response is invalid"
+        case let .otherError(content):
+            return content
         }
     }
 }
@@ -59,7 +62,7 @@ public actor ChatGPTService: ChatGPTServiceType {
     public var defaultTemperature: Double {
         min(max(0, UserDefaults.shared.value(for: \.chatGPTTemperature)), 2)
     }
-    
+
     var temperature: Double?
 
     public var model: String {
@@ -68,14 +71,30 @@ public actor ChatGPTService: ChatGPTServiceType {
         return value
     }
 
+    var designatedProvider: ChatFeatureProvider?
+
     public var endpoint: String {
-        var baseURL = UserDefaults.shared.value(for: \.openAIBaseURL)
-        if baseURL.isEmpty { return "https://api.openai.com/v1/chat/completions" }
-        return "\(baseURL)/v1/chat/completions"
+        switch designatedProvider ?? UserDefaults.shared.value(for: \.chatFeatureProvider) {
+        case .openAI:
+            let baseURL = UserDefaults.shared.value(for: \.openAIBaseURL)
+            if baseURL.isEmpty { return "https://api.openai.com/v1/chat/completions" }
+            return "\(baseURL)/v1/chat/completions"
+        case .azureOpenAI:
+            let baseURL = UserDefaults.shared.value(for: \.azureOpenAIBaseURL)
+            let deployment = UserDefaults.shared.value(for: \.azureChatGPTDeployment)
+            let version = "2023-05-15"
+            if baseURL.isEmpty { return "" }
+            return "\(baseURL)/openai/deployments/\(deployment)/chat/completions?api-version=\(version)"
+        }
     }
 
     public var apiKey: String {
-        UserDefaults.shared.value(for: \.openAIAPIKey)
+        switch designatedProvider ?? UserDefaults.shared.value(for: \.chatFeatureProvider) {
+        case .openAI:
+            return UserDefaults.shared.value(for: \.openAIAPIKey)
+        case .azureOpenAI:
+            return UserDefaults.shared.value(for: \.azureOpenAIAPIKey)
+        }
     }
 
     public var maxToken: Int {
@@ -97,10 +116,12 @@ public actor ChatGPTService: ChatGPTServiceType {
 
     public init(
         systemPrompt: String = "",
-        temperature: Double? = nil
+        temperature: Double? = nil,
+        designatedProvider: ChatFeatureProvider? = nil
     ) {
         self.systemPrompt = systemPrompt
         self.temperature = temperature
+        self.designatedProvider = designatedProvider
     }
 
     public func send(
@@ -129,7 +150,12 @@ public actor ChatGPTService: ChatGPTServiceType {
 
         isReceivingMessage = true
 
-        let api = buildCompletionStreamAPI(apiKey, url, requestBody)
+        let api = buildCompletionStreamAPI(
+            apiKey,
+            designatedProvider ?? UserDefaults.shared.value(for: \.chatFeatureProvider),
+            url,
+            requestBody
+        )
 
         return AsyncThrowingStream<String, Error> { continuation in
             Task {
@@ -210,7 +236,12 @@ public actor ChatGPTService: ChatGPTServiceType {
         isReceivingMessage = true
         defer { isReceivingMessage = false }
 
-        let api = buildCompletionAPI(apiKey, url, requestBody)
+        let api = buildCompletionAPI(
+            apiKey,
+            designatedProvider ?? UserDefaults.shared.value(for: \.chatFeatureProvider),
+            url,
+            requestBody
+        )
         let response = try await api()
 
         if let choice = response.choices.first {
@@ -296,3 +327,4 @@ func maxTokenForReply(model: String, remainingTokens: Int) -> Int {
     guard let model = ChatGPTModel(rawValue: model) else { return remainingTokens }
     return min(model.maxToken / 2, remainingTokens)
 }
+
diff --git a/Core/Sources/OpenAIService/CompletionAPI.swift b/Core/Sources/OpenAIService/CompletionAPI.swift
@@ -1,6 +1,8 @@
 import Foundation
+import Preferences
 
-typealias CompletionAPIBuilder = (String, URL, CompletionRequestBody) -> CompletionAPI
+typealias CompletionAPIBuilder = (String, ChatFeatureProvider, URL, CompletionRequestBody)
+    -> CompletionAPI
 
 protocol CompletionAPI {
     func callAsFunction() async throws -> CompletionResponseBody
@@ -12,13 +14,13 @@ struct CompletionResponseBody: Codable, Equatable {
         var role: ChatMessage.Role
         var content: String
     }
-    
+
     struct Choice: Codable, Equatable {
         var message: Message
         var index: Int
         var finish_reason: String
     }
-    
+
     struct Usage: Codable, Equatable {
         var prompt_tokens: Int
         var completion_tokens: Int
@@ -40,21 +42,29 @@ struct CompletionAPIError: Error, Codable, LocalizedError {
         var param: String
         var code: String
     }
+
     var error: E
-    
+
     var errorDescription: String? { error.message }
 }
 
 struct OpenAICompletionAPI: CompletionAPI {
     var apiKey: String
     var endpoint: URL
     var requestBody: CompletionRequestBody
+    var provider: ChatFeatureProvider
 
-    init(apiKey: String, endpoint: URL, requestBody: CompletionRequestBody) {
+    init(
+        apiKey: String,
+        provider: ChatFeatureProvider,
+        endpoint: URL,
+        requestBody: CompletionRequestBody
+    ) {
         self.apiKey = apiKey
         self.endpoint = endpoint
         self.requestBody = requestBody
         self.requestBody.stream = false
+        self.provider = provider
     }
 
     func callAsFunction() async throws -> CompletionResponseBody {
@@ -64,7 +74,11 @@ struct OpenAICompletionAPI: CompletionAPI {
         request.httpBody = try encoder.encode(requestBody)
         request.setValue("application/json", forHTTPHeaderField: "Content-Type")
         if !apiKey.isEmpty {
-            request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
+            if provider == .openAI {
+                request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
+            } else {
+                request.setValue(apiKey, forHTTPHeaderField: "api-key")
+            }
         }
 
         let (result, response) = try await URLSession.shared.data(for: request)
@@ -74,9 +88,11 @@ struct OpenAICompletionAPI: CompletionAPI {
 
         guard response.statusCode == 200 else {
             let error = try? JSONDecoder().decode(CompletionAPIError.self, from: result)
-            throw error ?? ChatGPTServiceError.responseInvalid
+            throw error ?? ChatGPTServiceError
+                .otherError(String(data: result, encoding: .utf8) ?? "Unknown Error")
         }
-        
+
         return try JSONDecoder().decode(CompletionResponseBody.self, from: result)
     }
 }
+
diff --git a/Core/Sources/OpenAIService/CompletionStreamAPI.swift b/Core/Sources/OpenAIService/CompletionStreamAPI.swift
@@ -1,7 +1,8 @@
 import AsyncAlgorithms
 import Foundation
+import Preferences
 
-typealias CompletionStreamAPIBuilder = (String, URL, CompletionRequestBody) -> CompletionStreamAPI
+typealias CompletionStreamAPIBuilder = (String, ChatFeatureProvider, URL, CompletionRequestBody) -> CompletionStreamAPI
 
 protocol CompletionStreamAPI {
     func callAsFunction() async throws -> (
@@ -54,12 +55,19 @@ struct OpenAICompletionStreamAPI: CompletionStreamAPI {
     var apiKey: String
     var endpoint: URL
     var requestBody: CompletionRequestBody
+    var provider: ChatFeatureProvider
 
-    init(apiKey: String, endpoint: URL, requestBody: CompletionRequestBody) {
+    init(
+        apiKey: String,
+        provider: ChatFeatureProvider,
+        endpoint: URL,
+        requestBody: CompletionRequestBody
+    ) {
         self.apiKey = apiKey
         self.endpoint = endpoint
         self.requestBody = requestBody
         self.requestBody.stream = true
+        self.provider = provider
     }
 
     func callAsFunction() async throws -> (
@@ -72,7 +80,11 @@ struct OpenAICompletionStreamAPI: CompletionStreamAPI {
         request.httpBody = try encoder.encode(requestBody)
         request.setValue("application/json", forHTTPHeaderField: "Content-Type")
         if !apiKey.isEmpty {
-            request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
+            if provider == .openAI {
+                request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
+            } else {
+                request.setValue(apiKey, forHTTPHeaderField: "api-key")
+            }
         }
 
         let (result, response) = try await URLSession.shared.bytes(for: request)
@@ -90,9 +102,9 @@ struct OpenAICompletionStreamAPI: CompletionStreamAPI {
             let error = try? decoder.decode(ChatGPTError.self, from: data)
             throw error ?? ChatGPTServiceError.responseInvalid
         }
-        
+
         var receivingDataTask: Task<Void, Error>?
-        
+
         let stream = AsyncThrowingStream<CompletionStreamDataTrunk, Error> { continuation in
             receivingDataTask = Task {
                 do {
@@ -122,3 +134,4 @@ struct OpenAICompletionStreamAPI: CompletionStreamAPI {
         )
     }
 }
+
diff --git a/Core/Sources/Preferences/ChatFeatureProvider.swift b/Core/Sources/Preferences/ChatFeatureProvider.swift
@@ -0,0 +1,4 @@
+public enum ChatFeatureProvider: String, CaseIterable {
+    case openAI
+    case azureOpenAI
+}
diff --git a/Core/Sources/Preferences/Keys.swift b/Core/Sources/Preferences/Keys.swift
@@ -99,6 +99,22 @@ public extension UserDefaultPreferenceKeys {
     }
 }
 
+// MARK: - Azure OpenAI Settings
+
+public extension UserDefaultPreferenceKeys {
+    var azureOpenAIAPIKey: PreferenceKey<String> {
+        .init(defaultValue: "", key: "AzureOpenAIAPIKey")
+    }
+    
+    var azureOpenAIBaseURL: PreferenceKey<String> {
+        .init(defaultValue: "", key: "AzureOpenAIBaseURL")
+    }
+    
+    var azureChatGPTDeployment: PreferenceKey<String> {
+        .init(defaultValue: "", key: "AzureChatGPTDeployment")
+    }
+}
+
 // MARK: - GitHub Copilot Settings
 
 public extension UserDefaultPreferenceKeys {
@@ -186,6 +202,10 @@ public extension UserDefaultPreferenceKeys {
 // MARK: - Chat
 
 public extension UserDefaultPreferenceKeys {
+    var chatFeatureProvider: PreferenceKey<ChatFeatureProvider> {
+        .init(defaultValue: .openAI, key: "ChatFeatureProvider")
+    }
+    
     var chatFontSize: PreferenceKey<Double> {
         .init(defaultValue: 12, key: "ChatFontSize")
     }

Original file line number	Diff line number	Diff line change
`@@ -30,6 +30,15 @@ struct ServiceView: View {`
`30`	`30`	`subtitle: "Chat, Prompt to Code",`
`31`	`31`	`image: "globe"`
`32`	`32`	`)`
	`33`	`+`
	`34`	`+ ScrollView {`
	`35`	`+ AzureView().padding()`
	`36`	`+ }.sidebarItem(`
	`37`	`+ tag: 3,`
	`38`	`+ title: "Azure",`
	`39`	`+ subtitle: "Chat, Prompt to Code",`
	`40`	`+ image: "globe"`
	`41`	`+ )`
`33`	`42`	`}`
`34`	`43`	`}`
`35`	`44`	`}`