diff --git a/FirebaseAI/CHANGELOG.md b/FirebaseAI/CHANGELOG.md index 4c2d482f8d5..109e8f83318 100644 --- a/FirebaseAI/CHANGELOG.md +++ b/FirebaseAI/CHANGELOG.md @@ -1,4 +1,7 @@ # 12.3.0 +- [feature] Added support for the Code Execution tool, which enables the model + to generate and run code to perform complex tasks like solving mathematical + equations or visualizing data. (#15280) - [fixed] Fixed a decoding error when generating images with the `gemini-2.5-flash-image-preview` model using `generateContentStream` or `sendMessageStream` with the Gemini Developer API. (#15262) diff --git a/FirebaseAI/Sources/AILog.swift b/FirebaseAI/Sources/AILog.swift index cae85a0ff0a..fe04716384a 100644 --- a/FirebaseAI/Sources/AILog.swift +++ b/FirebaseAI/Sources/AILog.swift @@ -63,6 +63,9 @@ enum AILog { case generateContentResponseUnrecognizedContentModality = 3012 case decodedUnsupportedImagenPredictionType = 3013 case decodedUnsupportedPartData = 3014 + case codeExecutionResultUnrecognizedOutcome = 3015 + case executableCodeUnrecognizedLanguage = 3016 + case fallbackValueUsed = 3017 // SDK State Errors case generateContentResponseNoCandidates = 4000 @@ -124,4 +127,32 @@ enum AILog { static func additionalLoggingEnabled() -> Bool { return ProcessInfo.processInfo.arguments.contains(enableArgumentKey) } + + /// Returns the unwrapped optional value if non-nil or returns the fallback value and logs. + /// + /// This convenience method is intended for use in place of `optionalValue ?? fallbackValue` with + /// the addition of logging on use of the fallback value. + /// + /// - Parameters: + /// - optionalValue: The value to unwrap. + /// - fallbackValue: The fallback (default) value to return when `optionalValue` is `nil`. + /// - level: The logging level to use for fallback messages; defaults to + /// `FirebaseLoggerLevel.warning`. + /// - code: The message code to use for fallback messages; defaults to + /// `MessageCode.fallbackValueUsed`. + /// - caller: The name of the unwrapped value; defaults to the name of the computed property or + /// function name from which the unwrapping occurred. + static func safeUnwrap(_ optionalValue: T?, + fallback fallbackValue: T, + level: FirebaseLoggerLevel = .warning, + code: MessageCode = .fallbackValueUsed, + caller: String = #function) -> T { + guard let unwrappedValue = optionalValue else { + AILog.log(level: level, code: code, """ + No value specified for '\(caller)' (\(T.self)); using fallback value '\(fallbackValue)'. + """) + return fallbackValue + } + return unwrappedValue + } } diff --git a/FirebaseAI/Sources/ModelContent.swift b/FirebaseAI/Sources/ModelContent.swift index a0dfe6eb937..cfe5b8ac3e5 100644 --- a/FirebaseAI/Sources/ModelContent.swift +++ b/FirebaseAI/Sources/ModelContent.swift @@ -39,6 +39,8 @@ struct InternalPart: Equatable, Sendable { case fileData(FileData) case functionCall(FunctionCall) case functionResponse(FunctionResponse) + case executableCode(ExecutableCode) + case codeExecutionResult(CodeExecutionResult) struct UnsupportedDataError: Error { let decodingError: DecodingError @@ -93,6 +95,16 @@ public struct ModelContent: Equatable, Sendable { return FunctionResponsePart( functionResponse, isThought: part.isThought, thoughtSignature: part.thoughtSignature ) + case let .executableCode(executableCode): + return ExecutableCodePart( + executableCode, isThought: part.isThought, thoughtSignature: part.thoughtSignature + ) + case let .codeExecutionResult(codeExecutionResult): + return CodeExecutionResultPart( + codeExecutionResult: codeExecutionResult, + isThought: part.isThought, + thoughtSignature: part.thoughtSignature + ) case .none: // Filter out parts that contain missing or unrecognized data return nil @@ -212,6 +224,8 @@ extension InternalPart.OneOfData: Codable { case fileData case functionCall case functionResponse + case executableCode + case codeExecutionResult } public func encode(to encoder: Encoder) throws { @@ -227,6 +241,10 @@ extension InternalPart.OneOfData: Codable { try container.encode(functionCall, forKey: .functionCall) case let .functionResponse(functionResponse): try container.encode(functionResponse, forKey: .functionResponse) + case let .executableCode(executableCode): + try container.encode(executableCode, forKey: .executableCode) + case let .codeExecutionResult(codeExecutionResult): + try container.encode(codeExecutionResult, forKey: .codeExecutionResult) } } @@ -242,6 +260,12 @@ extension InternalPart.OneOfData: Codable { self = try .functionCall(values.decode(FunctionCall.self, forKey: .functionCall)) } else if values.contains(.functionResponse) { self = try .functionResponse(values.decode(FunctionResponse.self, forKey: .functionResponse)) + } else if values.contains(.executableCode) { + self = try .executableCode(values.decode(ExecutableCode.self, forKey: .executableCode)) + } else if values.contains(.codeExecutionResult) { + self = try .codeExecutionResult( + values.decode(CodeExecutionResult.self, forKey: .codeExecutionResult) + ) } else { let unexpectedKeys = values.allKeys.map { $0.stringValue } throw UnsupportedDataError(decodingError: DecodingError.dataCorrupted( diff --git a/FirebaseAI/Sources/Tool.swift b/FirebaseAI/Sources/Tool.swift index 16c05b3a2e4..78dc8ef9443 100644 --- a/FirebaseAI/Sources/Tool.swift +++ b/FirebaseAI/Sources/Tool.swift @@ -71,17 +71,18 @@ public struct GoogleSearch: Sendable { public struct Tool: Sendable { /// A list of `FunctionDeclarations` available to the model. let functionDeclarations: [FunctionDeclaration]? + /// Specifies the Google Search configuration. let googleSearch: GoogleSearch? - init(functionDeclarations: [FunctionDeclaration]?) { - self.functionDeclarations = functionDeclarations - googleSearch = nil - } + let codeExecution: CodeExecution? - init(googleSearch: GoogleSearch) { + init(functionDeclarations: [FunctionDeclaration]? = nil, + googleSearch: GoogleSearch? = nil, + codeExecution: CodeExecution? = nil) { + self.functionDeclarations = functionDeclarations self.googleSearch = googleSearch - functionDeclarations = nil + self.codeExecution = codeExecution } /// Creates a tool that allows the model to perform function calling. @@ -126,6 +127,13 @@ public struct Tool: Sendable { public static func googleSearch(_ googleSearch: GoogleSearch = GoogleSearch()) -> Tool { return self.init(googleSearch: googleSearch) } + + /// Creates a tool that allows the model to execute code. + /// + /// For more details, see ``CodeExecution``. + public static func codeExecution() -> Tool { + return self.init(codeExecution: CodeExecution()) + } } /// Configuration for specifying function calling behavior. diff --git a/FirebaseAI/Sources/Types/Internal/InternalPart.swift b/FirebaseAI/Sources/Types/Internal/InternalPart.swift index bb62dd4c0b5..a8afe4439c3 100644 --- a/FirebaseAI/Sources/Types/Internal/InternalPart.swift +++ b/FirebaseAI/Sources/Types/Internal/InternalPart.swift @@ -63,6 +63,54 @@ struct FunctionResponse: Codable, Equatable, Sendable { } } +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) +struct ExecutableCode: Codable, Equatable, Sendable { + struct Language: CodableProtoEnum, Sendable, Equatable { + enum Kind: String { + case unspecified = "LANGUAGE_UNSPECIFIED" + case python = "PYTHON" + } + + let rawValue: String + + static let unrecognizedValueMessageCode = + AILog.MessageCode.executableCodeUnrecognizedLanguage + } + + let language: Language? + let code: String? + + init(language: Language, code: String) { + self.language = language + self.code = code + } +} + +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) +struct CodeExecutionResult: Codable, Equatable, Sendable { + struct Outcome: CodableProtoEnum, Sendable, Equatable { + enum Kind: String { + case unspecified = "OUTCOME_UNSPECIFIED" + case ok = "OUTCOME_OK" + case failed = "OUTCOME_FAILED" + case deadlineExceeded = "OUTCOME_DEADLINE_EXCEEDED" + } + + let rawValue: String + + static let unrecognizedValueMessageCode = + AILog.MessageCode.codeExecutionResultUnrecognizedOutcome + } + + let outcome: Outcome? + let output: String? + + init(outcome: Outcome, output: String) { + self.outcome = outcome + self.output = output + } +} + @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) struct ErrorPart: Part, Error { let error: Error diff --git a/FirebaseAI/Sources/Types/Public/Part.swift b/FirebaseAI/Sources/Types/Public/Part.swift index fb743d1025d..e0015901d61 100644 --- a/FirebaseAI/Sources/Types/Public/Part.swift +++ b/FirebaseAI/Sources/Types/Public/Part.swift @@ -202,3 +202,118 @@ public struct FunctionResponsePart: Part { self.thoughtSignature = thoughtSignature } } + +/// A part containing code that was executed by the model. +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) +public struct ExecutableCodePart: Part { + /// The language of the code in an ``ExecutableCodePart``. + public struct Language: Sendable, Equatable, CustomStringConvertible { + let internalLanguage: ExecutableCode.Language + + /// The Python programming language. + public static let python = ExecutableCodePart.Language(ExecutableCode.Language(kind: .python)) + + public var description: String { internalLanguage.rawValue } + + init(_ language: ExecutableCode.Language) { + internalLanguage = language + } + } + + let executableCode: ExecutableCode + let _isThought: Bool? + let thoughtSignature: String? + + /// The language of the code. + public var language: ExecutableCodePart.Language { + ExecutableCodePart.Language( + // Fallback to "LANGUAGE_UNSPECIFIED" if the value is ever omitted by the backend; this should + // never happen. + AILog.safeUnwrap( + executableCode.language, fallback: ExecutableCode.Language(kind: .unspecified) + ) + ) + } + + /// The code that was executed. + public var code: String { + // Fallback to empty string if `code` is ever omitted by the backend; this should never happen. + AILog.safeUnwrap(executableCode.code, fallback: "") + } + + public var isThought: Bool { _isThought ?? false } + + public init(language: ExecutableCodePart.Language, code: String) { + self.init( + ExecutableCode(language: language.internalLanguage, code: code), + isThought: nil, + thoughtSignature: nil + ) + } + + init(_ executableCode: ExecutableCode, isThought: Bool?, thoughtSignature: String?) { + self.executableCode = executableCode + _isThought = isThought + self.thoughtSignature = thoughtSignature + } +} + +/// The result of executing code. +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) +public struct CodeExecutionResultPart: Part { + /// The outcome of a code execution. + public struct Outcome: Sendable, Equatable, CustomStringConvertible { + let internalOutcome: CodeExecutionResult.Outcome + + /// The code executed without errors. + public static let ok = CodeExecutionResultPart.Outcome(CodeExecutionResult.Outcome(kind: .ok)) + + /// The code failed to execute. + public static let failed = + CodeExecutionResultPart.Outcome(CodeExecutionResult.Outcome(kind: .failed)) + + /// The code took too long to execute. + public static let deadlineExceeded = + CodeExecutionResultPart.Outcome(CodeExecutionResult.Outcome(kind: .deadlineExceeded)) + + public var description: String { internalOutcome.rawValue } + + init(_ outcome: CodeExecutionResult.Outcome) { + internalOutcome = outcome + } + } + + let codeExecutionResult: CodeExecutionResult + let _isThought: Bool? + let thoughtSignature: String? + + /// The outcome of the code execution. + public var outcome: CodeExecutionResultPart.Outcome { + CodeExecutionResultPart.Outcome( + // Fallback to "OUTCOME_UNSPECIFIED" if this value is ever omitted by the backend; this should + // never happen. + AILog.safeUnwrap( + codeExecutionResult.outcome, fallback: CodeExecutionResult.Outcome(kind: .unspecified) + ) + ) + } + + /// The output of the code execution. + public var output: String? { codeExecutionResult.output } + + public var isThought: Bool { _isThought ?? false } + + public init(outcome: CodeExecutionResultPart.Outcome, output: String) { + self.init( + codeExecutionResult: CodeExecutionResult(outcome: outcome.internalOutcome, output: output), + isThought: nil, + thoughtSignature: nil + ) + } + + init(codeExecutionResult: CodeExecutionResult, isThought: Bool?, thoughtSignature: String?) { + self.codeExecutionResult = codeExecutionResult + _isThought = isThought + self.thoughtSignature = thoughtSignature + } +} diff --git a/FirebaseAI/Sources/Types/Public/Tools/CodeExecution.swift b/FirebaseAI/Sources/Types/Public/Tools/CodeExecution.swift new file mode 100644 index 00000000000..5701b7523ce --- /dev/null +++ b/FirebaseAI/Sources/Types/Public/Tools/CodeExecution.swift @@ -0,0 +1,21 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +/// A tool that allows the model to execute code. +/// +/// This tool can be used to solve complex problems, for example, by generating and executing Python +/// code to solve a math problem. +public struct CodeExecution: Sendable, Encodable { + init() {} +} diff --git a/FirebaseAI/Tests/TestApp/Sources/Constants.swift b/FirebaseAI/Tests/TestApp/Sources/Constants.swift index be5c0c06891..bedd6a42053 100644 --- a/FirebaseAI/Tests/TestApp/Sources/Constants.swift +++ b/FirebaseAI/Tests/TestApp/Sources/Constants.swift @@ -26,6 +26,7 @@ public enum ModelNames { public static let gemini2FlashPreviewImageGeneration = "gemini-2.0-flash-preview-image-generation" public static let gemini2_5_FlashImagePreview = "gemini-2.5-flash-image-preview" public static let gemini2_5_Flash = "gemini-2.5-flash" + public static let gemini2_5_FlashLite = "gemini-2.5-flash-lite" public static let gemini2_5_Pro = "gemini-2.5-pro" public static let gemma3_4B = "gemma-3-4b-it" } diff --git a/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift b/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift index 5b70223ece4..d2fb589a432 100644 --- a/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift +++ b/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift @@ -424,6 +424,35 @@ struct GenerateContentIntegrationTests { } } + @Test(arguments: InstanceConfig.allConfigs) + func generateContent_codeExecution_succeeds(_ config: InstanceConfig) async throws { + let model = FirebaseAI.componentInstance(config).generativeModel( + modelName: ModelNames.gemini2_5_FlashLite, + generationConfig: generationConfig, + tools: [.codeExecution()] + ) + let prompt = """ + What is the sum of the first 5 prime numbers? Generate and run code for the calculation. + """ + + let response = try await model.generateContent(prompt) + + let candidate = try #require(response.candidates.first) + let executableCodeParts = candidate.content.parts.compactMap { $0 as? ExecutableCodePart } + #expect(executableCodeParts.count == 1) + let executableCodePart = try #require(executableCodeParts.first) + #expect(executableCodePart.language == .python) + #expect(executableCodePart.code.contains("sum")) + let codeExecutionResults = candidate.content.parts.compactMap { $0 as? CodeExecutionResultPart } + #expect(codeExecutionResults.count == 1) + let codeExecutionResultPart = try #require(codeExecutionResults.first) + #expect(codeExecutionResultPart.outcome == .ok) + let output = try #require(codeExecutionResultPart.output) + #expect(output.contains("28")) // 2 + 3 + 5 + 7 + 11 = 28 + let text = try #require(response.text) + #expect(text.contains("28")) + } + // MARK: Streaming Tests @Test(arguments: [ diff --git a/FirebaseAI/Tests/Unit/GenerativeModelGoogleAITests.swift b/FirebaseAI/Tests/Unit/GenerativeModelGoogleAITests.swift index b1ee49da6a1..c6335142959 100644 --- a/FirebaseAI/Tests/Unit/GenerativeModelGoogleAITests.swift +++ b/FirebaseAI/Tests/Unit/GenerativeModelGoogleAITests.swift @@ -308,6 +308,33 @@ final class GenerativeModelGoogleAITests: XCTestCase { XCTAssertTrue(thoughtSignature.hasPrefix("CtQOAVSoXO74PmYr9AFu")) } + func testGenerateContent_success_codeExecution() async throws { + MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler( + forResource: "unary-success-code-execution", + withExtension: "json", + subdirectory: googleAISubdirectory + ) + + let response = try await model.generateContent(testPrompt) + + XCTAssertEqual(response.candidates.count, 1) + let candidate = try XCTUnwrap(response.candidates.first) + let parts = candidate.content.parts + XCTAssertEqual(candidate.finishReason, .stop) + XCTAssertEqual(parts.count, 3) + let executableCodePart = try XCTUnwrap(parts[0] as? ExecutableCodePart) + XCTAssertFalse(executableCodePart.isThought) + XCTAssertEqual(executableCodePart.language, .python) + XCTAssertTrue(executableCodePart.code.starts(with: "prime_numbers = [2, 3, 5, 7, 11]")) + let codeExecutionResultPart = try XCTUnwrap(parts[1] as? CodeExecutionResultPart) + XCTAssertFalse(codeExecutionResultPart.isThought) + XCTAssertEqual(codeExecutionResultPart.outcome, .ok) + XCTAssertEqual(codeExecutionResultPart.output, "sum_of_primes=28\n") + let textPart = try XCTUnwrap(parts[2] as? TextPart) + XCTAssertFalse(textPart.isThought) + XCTAssertTrue(textPart.text.hasPrefix("The first 5 prime numbers are 2, 3, 5, 7, and 11.")) + } + func testGenerateContent_failure_invalidAPIKey() async throws { let expectedStatusCode = 400 MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler( @@ -526,6 +553,39 @@ final class GenerativeModelGoogleAITests: XCTestCase { } } + func testGenerateContentStream_success_codeExecution() async throws { + MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler( + forResource: "streaming-success-code-execution", + withExtension: "txt", + subdirectory: googleAISubdirectory + ) + + var parts = [any Part]() + let stream = try model.generateContentStream(testPrompt) + for try await response in stream { + if let responseParts = response.candidates.first?.content.parts { + parts.append(contentsOf: responseParts) + } + } + + let thoughtParts = parts.filter { $0.isThought } + XCTAssertEqual(thoughtParts.count, 0) + let textParts = parts.filter { $0 is TextPart } + XCTAssertGreaterThan(textParts.count, 0) + let executableCodeParts = parts.compactMap { $0 as? ExecutableCodePart } + XCTAssertEqual(executableCodeParts.count, 1) + let executableCodePart = try XCTUnwrap(executableCodeParts.first) + XCTAssertFalse(executableCodePart.isThought) + XCTAssertEqual(executableCodePart.language, .python) + XCTAssertTrue(executableCodePart.code.starts(with: "prime_numbers = [2, 3, 5, 7, 11]")) + let codeExecutionResultParts = parts.compactMap { $0 as? CodeExecutionResultPart } + XCTAssertEqual(codeExecutionResultParts.count, 1) + let codeExecutionResultPart = try XCTUnwrap(codeExecutionResultParts.first) + XCTAssertFalse(codeExecutionResultPart.isThought) + XCTAssertEqual(codeExecutionResultPart.outcome, .ok) + XCTAssertEqual(codeExecutionResultPart.output, "The sum of the first 5 prime numbers is: 28\n") + } + func testGenerateContentStream_failureInvalidAPIKey() async throws { MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler( forResource: "unary-failure-api-key", diff --git a/FirebaseAI/Tests/Unit/GenerativeModelVertexAITests.swift b/FirebaseAI/Tests/Unit/GenerativeModelVertexAITests.swift index 2b7a60ec0a8..847f5a8e643 100644 --- a/FirebaseAI/Tests/Unit/GenerativeModelVertexAITests.swift +++ b/FirebaseAI/Tests/Unit/GenerativeModelVertexAITests.swift @@ -457,6 +457,38 @@ final class GenerativeModelVertexAITests: XCTestCase { XCTAssertEqual(response.text, textPart.text) } + func testGenerateContent_success_codeExecution() async throws { + MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler( + forResource: "unary-success-code-execution", + withExtension: "json", + subdirectory: vertexSubdirectory + ) + + let response = try await model.generateContent(testPrompt) + + XCTAssertEqual(response.candidates.count, 1) + let candidate = try XCTUnwrap(response.candidates.first) + let parts = candidate.content.parts + XCTAssertEqual(candidate.finishReason, .stop) + XCTAssertEqual(parts.count, 4) + let textPart1 = try XCTUnwrap(parts[0] as? TextPart) + XCTAssertFalse(textPart1.isThought) + XCTAssertTrue(textPart1.text.hasPrefix("To find the sum of the first 5 prime numbers")) + let executableCodePart = try XCTUnwrap(parts[1] as? ExecutableCodePart) + XCTAssertFalse(executableCodePart.isThought) + XCTAssertEqual(executableCodePart.language, .python) + XCTAssertTrue(executableCodePart.code.starts(with: "prime_numbers = [2, 3, 5, 7, 11]")) + let codeExecutionResultPart = try XCTUnwrap(parts[2] as? CodeExecutionResultPart) + XCTAssertFalse(codeExecutionResultPart.isThought) + XCTAssertEqual(codeExecutionResultPart.outcome, .ok) + XCTAssertEqual(codeExecutionResultPart.output, "The sum of the first 5 prime numbers is: 28\n") + let textPart2 = try XCTUnwrap(parts[3] as? TextPart) + XCTAssertFalse(textPart2.isThought) + XCTAssertEqual( + textPart2.text, "The sum of the first 5 prime numbers (2, 3, 5, 7, and 11) is 28." + ) + } + func testGenerateContent_success_image_invalidSafetyRatingsIgnored() async throws { MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler( forResource: "unary-success-image-invalid-safety-ratings", @@ -1426,6 +1458,39 @@ final class GenerativeModelVertexAITests: XCTestCase { XCTAssertTrue(text.hasPrefix("The sky is blue due to a phenomenon")) } + func testGenerateContentStream_success_codeExecution() async throws { + MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler( + forResource: "streaming-success-code-execution", + withExtension: "txt", + subdirectory: vertexSubdirectory + ) + + var parts = [any Part]() + let stream = try model.generateContentStream(testPrompt) + for try await response in stream { + if let responseParts = response.candidates.first?.content.parts { + parts.append(contentsOf: responseParts) + } + } + + let thoughtParts = parts.filter { $0.isThought } + XCTAssertEqual(thoughtParts.count, 0) + let textParts = parts.filter { $0 is TextPart } + XCTAssertGreaterThan(textParts.count, 0) + let executableCodeParts = parts.compactMap { $0 as? ExecutableCodePart } + XCTAssertEqual(executableCodeParts.count, 1) + let executableCodePart = try XCTUnwrap(executableCodeParts.first) + XCTAssertFalse(executableCodePart.isThought) + XCTAssertEqual(executableCodePart.language, .python) + XCTAssertTrue(executableCodePart.code.starts(with: "prime_numbers = [2, 3, 5, 7, 11]")) + let codeExecutionResultParts = parts.compactMap { $0 as? CodeExecutionResultPart } + XCTAssertEqual(codeExecutionResultParts.count, 1) + let codeExecutionResultPart = try XCTUnwrap(codeExecutionResultParts.first) + XCTAssertFalse(codeExecutionResultPart.isThought) + XCTAssertEqual(codeExecutionResultPart.outcome, .ok) + XCTAssertEqual(codeExecutionResultPart.output, "The sum of the first 5 prime numbers is: 28\n") + } + func testGenerateContentStream_successWithInvalidSafetyRatingsIgnored() async throws { MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler( forResource: "streaming-success-image-invalid-safety-ratings", diff --git a/FirebaseAI/Tests/Unit/PartTests.swift b/FirebaseAI/Tests/Unit/PartTests.swift index 544e229e08b..f538586d439 100644 --- a/FirebaseAI/Tests/Unit/PartTests.swift +++ b/FirebaseAI/Tests/Unit/PartTests.swift @@ -86,6 +86,132 @@ final class PartTests: XCTestCase { XCTAssertEqual(functionResponse.response, [resultParameter: .string(resultValue)]) } + func testDecodeExecutableCodePart() throws { + let json = """ + { + "executableCode": { + "language": "PYTHON", + "code": "print('hello')" + } + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + + let part = try decoder.decode(ExecutableCodePart.self, from: jsonData) + + XCTAssertEqual(part.language, .python) + XCTAssertEqual(part.code, "print('hello')") + } + + func testDecodeExecutableCodePart_missingLanguage() throws { + let json = """ + { + "executableCode": { + "code": "print('hello')" + } + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + + let part = try decoder.decode(ExecutableCodePart.self, from: jsonData) + + XCTAssertEqual(part.language.description, "LANGUAGE_UNSPECIFIED") + XCTAssertEqual(part.code, "print('hello')") + } + + func testDecodeExecutableCodePart_missingCode() throws { + let json = """ + { + "executableCode": { + "language": "PYTHON" + } + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + + let part = try decoder.decode(ExecutableCodePart.self, from: jsonData) + + XCTAssertEqual(part.language, .python) + XCTAssertEqual(part.code, "") + } + + func testDecodeExecutableCodePart_missingLanguageAndCode() throws { + let json = """ + { + "executableCode": {} + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + + let part = try decoder.decode(ExecutableCodePart.self, from: jsonData) + + XCTAssertEqual(part.language.description, "LANGUAGE_UNSPECIFIED") + XCTAssertEqual(part.code, "") + } + + func testDecodeCodeExecutionResultPart() throws { + let json = """ + { + "codeExecutionResult": { + "outcome": "OUTCOME_OK", + "output": "hello" + } + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + + let part = try decoder.decode(CodeExecutionResultPart.self, from: jsonData) + + XCTAssertEqual(part.outcome, .ok) + XCTAssertEqual(part.output, "hello") + } + + func testDecodeCodeExecutionResultPart_missingOutcome() throws { + let json = """ + { + "codeExecutionResult": { + "output": "hello" + } + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + + let part = try decoder.decode(CodeExecutionResultPart.self, from: jsonData) + + XCTAssertEqual(part.outcome.description, "OUTCOME_UNSPECIFIED") + XCTAssertEqual(part.output, "hello") + } + + func testDecodeCodeExecutionResultPart_missingOutput() throws { + let json = """ + { + "codeExecutionResult": { + "outcome": "OUTCOME_OK" + } + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + + let part = try decoder.decode(CodeExecutionResultPart.self, from: jsonData) + + XCTAssertEqual(part.outcome, .ok) + XCTAssertNil(part.output) + } + + func testDecodeCodeExecutionResultPart_missingOutcomeAndOutput() throws { + let json = """ + { + "codeExecutionResult": {} + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + + let part = try decoder.decode(CodeExecutionResultPart.self, from: jsonData) + + XCTAssertEqual(part.outcome.description, "OUTCOME_UNSPECIFIED") + XCTAssertNil(part.output) + } + // MARK: - Part Encoding func testEncodeTextPart() throws { @@ -139,6 +265,38 @@ final class PartTests: XCTestCase { """) } + func testEncodeExecutableCodePart() throws { + let executableCodePart = ExecutableCodePart(language: .python, code: "print('hello')") + + let jsonData = try encoder.encode(executableCodePart) + + let json = try XCTUnwrap(String(data: jsonData, encoding: .utf8)) + XCTAssertEqual(json, """ + { + "executableCode" : { + "code" : "print('hello')", + "language" : "PYTHON" + } + } + """) + } + + func testEncodeCodeExecutionResultPart() throws { + let codeExecutionResultPart = CodeExecutionResultPart(outcome: .ok, output: "hello") + + let jsonData = try encoder.encode(codeExecutionResultPart) + + let json = try XCTUnwrap(String(data: jsonData, encoding: .utf8)) + XCTAssertEqual(json, """ + { + "codeExecutionResult" : { + "outcome" : "OUTCOME_OK", + "output" : "hello" + } + } + """) + } + // MARK: - Helpers private static func bundle() -> Bundle { diff --git a/FirebaseAI/Tests/Unit/Types/InternalPartTests.swift b/FirebaseAI/Tests/Unit/Types/InternalPartTests.swift index 2cd5c5fee2a..65121e913c2 100644 --- a/FirebaseAI/Tests/Unit/Types/InternalPartTests.swift +++ b/FirebaseAI/Tests/Unit/Types/InternalPartTests.swift @@ -283,4 +283,216 @@ final class InternalPartTests: XCTestCase { XCTAssertEqual(functionResponse.name, functionName) XCTAssertEqual(functionResponse.response, ["output": .string("someValue")]) } + + func testDecodeExecutableCodePartWithThought() throws { + let json = """ + { + "executableCode": { + "language": "PYTHON", + "code": "print('hello')" + }, + "thought": true + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + + let part = try decoder.decode(InternalPart.self, from: jsonData) + + XCTAssertEqual(part.isThought, true) + guard case let .executableCode(executableCode) = part.data else { + XCTFail("Decoded part is not an executableCode part.") + return + } + XCTAssertEqual(executableCode.language, .init(kind: .python)) + XCTAssertEqual(executableCode.code, "print('hello')") + } + + func testDecodeExecutableCodePartWithoutThought() throws { + let json = """ + { + "executableCode": { + "language": "PYTHON", + "code": "print('hello')" + } + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + + let part = try decoder.decode(InternalPart.self, from: jsonData) + + XCTAssertNil(part.isThought) + guard case let .executableCode(executableCode) = part.data else { + XCTFail("Decoded part is not an executableCode part.") + return + } + XCTAssertEqual(executableCode.language, .init(kind: .python)) + XCTAssertEqual(executableCode.code, "print('hello')") + } + + func testDecodeExecutableCodePart_missingLanguage() throws { + let json = """ + { + "executableCode": { + "code": "print('hello')" + } + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + + let part = try decoder.decode(InternalPart.self, from: jsonData) + + XCTAssertNil(part.isThought) + guard case let .executableCode(executableCode) = part.data else { + XCTFail("Decoded part is not an executableCode part.") + return + } + XCTAssertNil(executableCode.language) + XCTAssertEqual(executableCode.code, "print('hello')") + } + + func testDecodeExecutableCodePart_missingCode() throws { + let json = """ + { + "executableCode": { + "language": "PYTHON" + } + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + + let part = try decoder.decode(InternalPart.self, from: jsonData) + + XCTAssertNil(part.isThought) + guard case let .executableCode(executableCode) = part.data else { + XCTFail("Decoded part is not an executableCode part.") + return + } + XCTAssertEqual(executableCode.language, .init(kind: .python)) + XCTAssertNil(executableCode.code) + } + + func testDecodeExecutableCodePart_missingLanguageAndCode() throws { + let json = """ + { + "executableCode": {} + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + + let part = try decoder.decode(InternalPart.self, from: jsonData) + + XCTAssertNil(part.isThought) + guard case let .executableCode(executableCode) = part.data else { + XCTFail("Decoded part is not an executableCode part.") + return + } + XCTAssertNil(executableCode.language) + XCTAssertNil(executableCode.code) + } + + func testDecodeCodeExecutionResultPartWithThought() throws { + let json = """ + { + "codeExecutionResult": { + "outcome": "OUTCOME_OK", + "output": "hello" + }, + "thought": true + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + + let part = try decoder.decode(InternalPart.self, from: jsonData) + + XCTAssertEqual(part.isThought, true) + guard case let .codeExecutionResult(codeExecutionResult) = part.data else { + XCTFail("Decoded part is not a codeExecutionResult part.") + return + } + XCTAssertEqual(codeExecutionResult.outcome, .init(kind: .ok)) + XCTAssertEqual(codeExecutionResult.output, "hello") + } + + func testDecodeCodeExecutionResultPartWithoutThought() throws { + let json = """ + { + "codeExecutionResult": { + "outcome": "OUTCOME_OK", + "output": "hello" + } + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + + let part = try decoder.decode(InternalPart.self, from: jsonData) + + XCTAssertNil(part.isThought) + guard case let .codeExecutionResult(codeExecutionResult) = part.data else { + XCTFail("Decoded part is not a codeExecutionResult part.") + return + } + XCTAssertEqual(codeExecutionResult.outcome, .init(kind: .ok)) + XCTAssertEqual(codeExecutionResult.output, "hello") + } + + func testDecodeCodeExecutionResultPart_missingOutcome() throws { + let json = """ + { + "codeExecutionResult": { + "output": "hello" + } + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + + let part = try decoder.decode(InternalPart.self, from: jsonData) + + XCTAssertNil(part.isThought) + guard case let .codeExecutionResult(codeExecutionResult) = part.data else { + XCTFail("Decoded part is not a codeExecutionResult part.") + return + } + XCTAssertNil(codeExecutionResult.outcome) + XCTAssertEqual(codeExecutionResult.output, "hello") + } + + func testDecodeCodeExecutionResultPart_missingOutput() throws { + let json = """ + { + "codeExecutionResult": { + "outcome": "OUTCOME_OK" + } + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + + let part = try decoder.decode(InternalPart.self, from: jsonData) + + XCTAssertNil(part.isThought) + guard case let .codeExecutionResult(codeExecutionResult) = part.data else { + XCTFail("Decoded part is not a codeExecutionResult part.") + return + } + XCTAssertEqual(codeExecutionResult.outcome, .init(kind: .ok)) + XCTAssertNil(codeExecutionResult.output) + } + + func testDecodeCodeExecutionResultPart_missingOutcomeAndOutput() throws { + let json = """ + { + "codeExecutionResult": {} + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + + let part = try decoder.decode(InternalPart.self, from: jsonData) + + XCTAssertNil(part.isThought) + guard case let .codeExecutionResult(codeExecutionResult) = part.data else { + XCTFail("Decoded part is not a codeExecutionResult part.") + return + } + XCTAssertNil(codeExecutionResult.outcome) + XCTAssertNil(codeExecutionResult.output) + } } diff --git a/FirebaseAI/Tests/Unit/Types/ToolTests.swift b/FirebaseAI/Tests/Unit/Types/ToolTests.swift index b163894932d..9bfdf2313b7 100644 --- a/FirebaseAI/Tests/Unit/Types/ToolTests.swift +++ b/FirebaseAI/Tests/Unit/Types/ToolTests.swift @@ -27,7 +27,9 @@ final class ToolTests: XCTestCase { func testEncodeTool_googleSearch() throws { let tool = Tool.googleSearch() + let jsonData = try encoder.encode(tool) + let jsonString = try XCTUnwrap(String(data: jsonData, encoding: .utf8)) XCTAssertEqual(jsonString, """ { @@ -38,6 +40,21 @@ final class ToolTests: XCTestCase { """) } + func testEncodeTool_codeExecution() throws { + let tool = Tool.codeExecution() + + let jsonData = try encoder.encode(tool) + + let jsonString = try XCTUnwrap(String(data: jsonData, encoding: .utf8)) + XCTAssertEqual(jsonString, """ + { + "codeExecution" : { + + } + } + """) + } + func testEncodeTool_functionDeclarations() throws { let functionDecl = FunctionDeclaration( name: "test_function", @@ -45,11 +62,9 @@ final class ToolTests: XCTestCase { parameters: ["param1": .string()] ) let tool = Tool.functionDeclarations([functionDecl]) - - encoder.outputFormatting.insert(.withoutEscapingSlashes) let jsonData = try encoder.encode(tool) - let jsonString = try XCTUnwrap(String(data: jsonData, encoding: .utf8)) + let jsonString = try XCTUnwrap(String(data: jsonData, encoding: .utf8)) XCTAssertEqual(jsonString, """ { "functionDeclarations" : [