合并代码

Author: enefry

Sources/GoogleAI/CountTokensRequest.swift

@@ -26,7 +26,10 @@ extension CountTokensRequest: GenerativeAIRequest {
   typealias Response = CountTokensResponse
 
   var url: URL {
-    URL(string: "\(GenerativeAISwift.baseURL)/\(options.apiVersion)/\(model):countTokens")!
+    return getUrl(nil)
+  }
+  func getUrl(_ baseUrl: String?) -> URL {
+      URL(string: "\(baseUrl ?? GenerativeAISwift.baseURL)/\(options.apiVersion)/\(model):countTokens")!
   }
 }
 

Sources/GoogleAI/GenerateContentRequest.swift

@@ -45,11 +45,14 @@ extension GenerateContentRequest: GenerativeAIRequest {
   typealias Response = GenerateContentResponse
 
   var url: URL {
-    let modelURL = "\(GenerativeAISwift.baseURL)/\(options.apiVersion)/\(model)"
-    if isStreaming {
-      return URL(string: "\(modelURL):streamGenerateContent?alt=sse")!
-    } else {
-      return URL(string: "\(modelURL):generateContent")!
-    }
+    return getUrl(nil)
+  }
+  func getUrl(_ baseUrl: String?) -> URL {
+        let modelURL = "\(baseUrl ?? GenerativeAISwift.baseURL)/\(options.apiVersion)/\(model)"
+        if isStreaming {
+          return URL(string: "\(modelURL):streamGenerateContent?alt=sse")!
+        } else {
+          return URL(string: "\(modelURL):generateContent")!
+        }
   }
 }

Sources/GoogleAI/GenerativeAIRequest.swift

@@ -19,7 +19,7 @@ protocol GenerativeAIRequest: Encodable {
   associatedtype Response: Decodable
 
   var url: URL { get }
-
+  func getUrl(_ baseUrl:String?) -> URL;
   var options: RequestOptions { get }
 }
 

Sources/GoogleAI/GenerativeAIService.swift

@@ -20,10 +20,13 @@ struct GenerativeAIService {
   private let apiKey: String
 
   private let urlSession: URLSession
-
-  init(apiKey: String, urlSession: URLSession) {
+    
+  private let baseUrl:String?
+    
+  init(apiKey: String, urlSession: URLSession, baseUrl: String? = nil) {
     self.apiKey = apiKey
     self.urlSession = urlSession
+    self.baseUrl = baseUrl
   }
 
   func loadRequest<T: GenerativeAIRequest>(request: T) async throws -> T.Response {
@@ -147,7 +150,7 @@ struct GenerativeAIService {
   // MARK: - Private Helpers
 
   private func urlRequest<T: GenerativeAIRequest>(request: T) throws -> URLRequest {
-    var urlRequest = URLRequest(url: request.url)
+    var urlRequest = URLRequest(url: request.getUrl(self.baseUrl))
     urlRequest.httpMethod = "POST"
     urlRequest.setValue(apiKey, forHTTPHeaderField: "x-goog-api-key")
     urlRequest.setValue("genai-swift/\(GenerativeAISwift.version)",

Sources/GoogleAI/GenerativeModel.swift

@@ -96,23 +96,12 @@ public final class GenerativeModel {
   public convenience init(name: String,
                           apiKey: String,
                           generationConfig: GenerationConfig? = nil,
-                          safetySettings: [SafetySetting]? = nil,
-                          tools: [Tool]? = nil,
-                          toolConfig: ToolConfig? = nil,
-                          systemInstruction: String...,
-                          requestOptions: RequestOptions = RequestOptions()) {
+                          safetySettings: [SafetySetting]? = nil) {
     self.init(
       name: name,
       apiKey: apiKey,
       generationConfig: generationConfig,
       safetySettings: safetySettings,
-      tools: tools,
-      toolConfig: toolConfig,
-      systemInstruction: ModelContent(
-        role: "system",
-        parts: systemInstruction.map { ModelContent.Part.text($0) }
-      ),
-      requestOptions: requestOptions,
       urlSession: .shared
     )
   }
@@ -126,9 +115,10 @@ public final class GenerativeModel {
        toolConfig: ToolConfig? = nil,
        systemInstruction: ModelContent? = nil,
        requestOptions: RequestOptions = RequestOptions(),
-       urlSession: URLSession) {
+       urlSession: URLSession,
+       baseUrl: String? = nil) {
     modelResourceName = GenerativeModel.modelResourceName(name: name)
-    generativeAIService = GenerativeAIService(apiKey: apiKey, urlSession: urlSession)
+    generativeAIService = GenerativeAIService(apiKey: apiKey, urlSession: urlSession, baseUrl: baseUrl)
     self.generationConfig = generationConfig
     self.safetySettings = safetySettings
     self.tools = tools