From 08c5636130e6d18e45af145fda9ee7e5b0a1c3f5 Mon Sep 17 00:00:00 2001 From: Andrew Heard Date: Wed, 17 Apr 2024 15:10:03 -0400 Subject: [PATCH 1/2] Add URI-based file data support --- .../ViewModels/FunctionCallingViewModel.swift | 2 +- Sources/GoogleAI/Chat.swift | 2 +- Sources/GoogleAI/ModelContent.swift | 22 ++++++++++++++++++- 3 files changed, 23 insertions(+), 3 deletions(-) diff --git a/Examples/GenerativeAISample/FunctionCallingSample/ViewModels/FunctionCallingViewModel.swift b/Examples/GenerativeAISample/FunctionCallingSample/ViewModels/FunctionCallingViewModel.swift index 0fd64a4..ec5f291 100644 --- a/Examples/GenerativeAISample/FunctionCallingSample/ViewModels/FunctionCallingViewModel.swift +++ b/Examples/GenerativeAISample/FunctionCallingSample/ViewModels/FunctionCallingViewModel.swift @@ -158,7 +158,7 @@ class FunctionCallingViewModel: ObservableObject { case let .functionCall(functionCall): messages.insert(functionCall.chatMessage(), at: messages.count - 1) functionCalls.append(functionCall) - case .data, .functionResponse: + case .data, .fileData, .functionResponse: fatalError("Unsupported response content.") } } diff --git a/Sources/GoogleAI/Chat.swift b/Sources/GoogleAI/Chat.swift index e443e18..a83947e 100644 --- a/Sources/GoogleAI/Chat.swift +++ b/Sources/GoogleAI/Chat.swift @@ -153,7 +153,7 @@ public class Chat { case let .text(str): combinedText += str - case .data, .functionCall, .functionResponse: + case .data, .fileData, .functionCall, .functionResponse: // Don't combine it, just add to the content. If there's any text pending, add that as // a part. if !combinedText.isEmpty { diff --git a/Sources/GoogleAI/ModelContent.swift b/Sources/GoogleAI/ModelContent.swift index 136fc1d..3e77472 100644 --- a/Sources/GoogleAI/ModelContent.swift +++ b/Sources/GoogleAI/ModelContent.swift @@ -25,6 +25,7 @@ public struct ModelContent: Codable, Equatable { enum CodingKeys: String, CodingKey { case text case inlineData + case fileData case functionCall case functionResponse } @@ -34,12 +35,24 @@ public struct ModelContent: Codable, Equatable { case bytes = "data" } + enum FileDataKeys: String, CodingKey { + case mimeType = "mime_type" + case url = "file_uri" + } + /// Text value. case text(String) - /// Data with a specified media type. Not all media types may be supported by the AI model. + /// Data with a specified media type. + /// + /// > Note: Supported media types depends on the model. case data(mimetype: String, Data) + /// URI-based data with a specified media type. + /// + /// > Note: Supported media types depends on the model. + case fileData(mimetype: String, URL) + /// A predicted function call returned from the model. case functionCall(FunctionCall) @@ -72,6 +85,13 @@ public struct ModelContent: Codable, Equatable { ) try inlineDataContainer.encode(mimetype, forKey: .mimeType) try inlineDataContainer.encode(bytes, forKey: .bytes) + case let .fileData(mimetype: mimetype, url): + var fileDataContainer = container.nestedContainer( + keyedBy: FileDataKeys.self, + forKey: .fileData + ) + try fileDataContainer.encode(mimetype, forKey: .mimeType) + try fileDataContainer.encode(url, forKey: .url) case let .functionCall(functionCall): try container.encode(functionCall, forKey: .functionCall) case let .functionResponse(functionResponse): From 962b40e7e60a68cef61a9ee1e86e4922b6f7dbad Mon Sep 17 00:00:00 2001 From: Andrew Heard Date: Wed, 17 Apr 2024 15:46:20 -0400 Subject: [PATCH 2/2] Switch to `uri` / `String` --- Sources/GoogleAI/ModelContent.swift | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/Sources/GoogleAI/ModelContent.swift b/Sources/GoogleAI/ModelContent.swift index 3e77472..9c96625 100644 --- a/Sources/GoogleAI/ModelContent.swift +++ b/Sources/GoogleAI/ModelContent.swift @@ -45,13 +45,19 @@ public struct ModelContent: Codable, Equatable { /// Data with a specified media type. /// - /// > Note: Supported media types depends on the model. + /// > Note: Supported media types depends on the model; see + /// > [supported file + /// > formats](https://ai.google.dev/tutorials/prompting_with_media#supported_file_formats) + /// > for details. case data(mimetype: String, Data) /// URI-based data with a specified media type. /// - /// > Note: Supported media types depends on the model. - case fileData(mimetype: String, URL) + /// > Note: Supported media types depends on the model; see + /// > [supported file + /// > formats](https://ai.google.dev/tutorials/prompting_with_media#supported_file_formats) + /// > for details. + case fileData(mimetype: String, uri: String) /// A predicted function call returned from the model. case functionCall(FunctionCall)