Fix model name and error handling

Author: intitni

Tool/Sources/OpenAIService/APIs/GoogleAICompletionAPI.swift

@@ -10,7 +10,7 @@ struct GoogleCompletionAPI: CompletionAPI {
 
     func callAsFunction() async throws -> CompletionResponseBody {
         let aiModel = GenerativeModel(
-            name: model.name,
+            name: model.info.modelName,
             apiKey: apiKey,
             generationConfig: .init(GenerationConfig(
                 temperature: requestBody.temperature.map(Float.init),
@@ -30,30 +30,52 @@ struct GoogleCompletionAPI: CompletionAPI {
             )
         }
 
-        let response = try await aiModel.generateContent(history)
-
-        return .init(
-            object: "chat.completion",
-            model: model.name,
-            usage: .init(prompt_tokens: 0, completion_tokens: 0, total_tokens: 0),
-            choices: response.candidates.enumerated().map {
-                let (index, candidate) = $0
-                return .init(
-                    message: .init(
-                        role: .assistant,
-                        content: candidate.content.parts.first(where: { part in
-                            if let text = part.text {
-                                return !text.isEmpty
-                            } else {
-                                return false
-                            }
-                        })?.text ?? ""
-                    ),
-                    index: index,
-                    finish_reason: candidate.finishReason?.rawValue ?? ""
-                )
+        do {
+            let response = try await aiModel.generateContent(history)
+            
+            return .init(
+                object: "chat.completion",
+                model: model.info.modelName,
+                usage: .init(prompt_tokens: 0, completion_tokens: 0, total_tokens: 0),
+                choices: response.candidates.enumerated().map {
+                    let (index, candidate) = $0
+                    return .init(
+                        message: .init(
+                            role: .assistant,
+                            content: candidate.content.parts.first(where: { part in
+                                if let text = part.text {
+                                    return !text.isEmpty
+                                } else {
+                                    return false
+                                }
+                            })?.text ?? ""
+                        ),
+                        index: index,
+                        finish_reason: candidate.finishReason?.rawValue ?? ""
+                    )
+                }
+            )
+        } catch let error as GenerateContentError {
+            struct ErrorWrapper: Error, LocalizedError {
+                let error: Error
+                var errorDescription: String? {
+                    var s = ""
+                    dump(error, to: &s)
+                    return "Internal Error: \(s)"
+                }
             }
-        )
+            
+            switch error {
+            case let .internalError(underlying):
+                throw ErrorWrapper(error: underlying)
+            case .promptBlocked:
+                throw error
+            case .responseStoppedEarly:
+                throw error
+            }
+        } catch {
+            throw error
+        }
     }
 }
 

Tool/Sources/OpenAIService/APIs/GoogleAICompletionStreamAPI.swift

@@ -10,7 +10,7 @@ struct GoogleCompletionStreamAPI: CompletionStreamAPI {
 
     func callAsFunction() async throws -> AsyncThrowingStream<CompletionStreamDataChunk, Error> {
         let aiModel = GenerativeModel(
-            name: model.name,
+            name: model.info.modelName,
             apiKey: apiKey,
             generationConfig: .init(GenerationConfig(
                 temperature: requestBody.temperature.map(Float.init),
@@ -38,7 +38,7 @@ struct GoogleCompletionStreamAPI: CompletionStreamAPI {
                         if Task.isCancelled { break }
                         let chunk = CompletionStreamDataChunk(
                             object: "",
-                            model: model.name,
+                            model: model.info.modelName,
                             choices: response.candidates.map { candidate in
                                 .init(delta: .init(
                                     role: .assistant,
@@ -50,6 +50,24 @@ struct GoogleCompletionStreamAPI: CompletionStreamAPI {
                         continuation.yield(chunk)
                     }
                     continuation.finish()
+                } catch let error as GenerateContentError {
+                    struct ErrorWrapper: Error, LocalizedError {
+                        let error: Error
+                        var errorDescription: String? {
+                            var s = ""
+                            dump(error, to: &s)
+                            return "Internal Error: \(s)"
+                        }
+                    }
+                    
+                    switch error {
+                    case let .internalError(underlying):
+                        continuation.finish(throwing: ErrorWrapper(error: underlying))
+                    case .promptBlocked:
+                        continuation.finish(throwing: error)
+                    case .responseStoppedEarly:
+                        continuation.finish(throwing: error)
+                    }
                 } catch {
                     continuation.finish(throwing: error)
                 }

Tool/Sources/OpenAIService/Memory/AutoManagedChatGPTMemoryStrategy/AutoManagedChatGPTMemoryGoogleAIStrategy.swift

@@ -10,7 +10,7 @@ extension AutoManagedChatGPTMemory {
             guard let model = configuration.model else {
                 return 0
             }
-            let aiModel = GenerativeModel(name: model.name, apiKey: configuration.apiKey)
+            let aiModel = GenerativeModel(name: model.info.modelName, apiKey: configuration.apiKey)
             if message.isEmpty { return 0 }
             let modelMessage = ModelContent(message)
             return (try? await aiModel.countTokens([modelMessage]).totalTokens) ?? 0