diff --git a/FirebaseVertexAI/Tests/Unit/ChatTests.swift b/FirebaseVertexAI/Tests/Unit/ChatTests.swift new file mode 100644 index 00000000000..668e3704122 --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/ChatTests.swift @@ -0,0 +1,67 @@ +// Copyright 2023 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. + +import Foundation +@testable import GoogleGenerativeAI +import XCTest + +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, *) +final class ChatTests: XCTestCase { + var urlSession: URLSession! + + override func setUp() { + let configuration = URLSessionConfiguration.default + configuration.protocolClasses = [MockURLProtocol.self] + urlSession = URLSession(configuration: configuration) + } + + override func tearDown() { + MockURLProtocol.requestHandler = nil + } + + func testMergingText() async throws { + let fileURL = try XCTUnwrap(Bundle.module.url( + forResource: "streaming-success-basic-reply-long", + withExtension: "txt" + )) + + MockURLProtocol.requestHandler = { request in + let response = HTTPURLResponse( + url: request.url!, + statusCode: 200, + httpVersion: nil, + headerFields: nil + )! + return (response, fileURL.lines) + } + + let model = GenerativeModel(name: "my-model", apiKey: "API_KEY", urlSession: urlSession) + let chat = Chat(model: model, history: []) + let input = "Test input" + let stream = chat.sendMessageStream(input) + + // Ensure the values are parsed correctly + for try await value in stream { + XCTAssertNotNil(value.text) + } + + XCTAssertEqual(chat.history.count, 2) + XCTAssertEqual(chat.history[0].parts[0].text, input) + + let finalText = "1 2 3 4 5 6 7 8 9 10" + let assembledExpectation = ModelContent(role: "model", parts: finalText) + XCTAssertEqual(chat.history[0].parts[0].text, input) + XCTAssertEqual(chat.history[1], assembledExpectation) + } +} diff --git a/FirebaseVertexAI/Tests/Unit/CountTokenResponses/failure-model-not-found.json b/FirebaseVertexAI/Tests/Unit/CountTokenResponses/failure-model-not-found.json new file mode 100644 index 00000000000..50fcb725667 --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/CountTokenResponses/failure-model-not-found.json @@ -0,0 +1,13 @@ +{ + "error": { + "code": 404, + "message": "models/test-model-name is not found for API version v1beta, or is not supported for countTokens. Call ListModels to see the list of available models and their supported methods.", + "status": "NOT_FOUND", + "details": [ + { + "@type": "type.googleapis.com/google.rpc.DebugInfo", + "detail": "[ORIGINAL ERROR] generic::not_found: models/test-model-name is not found for API version v1beta, or is not supported for countTokens. Call ListModels to see the list of available models and their supported methods. [google.rpc.error_details_ext] { message: \"models/test-model-name is not found for API version v1beta, or is not supported for countTokens. Call ListModels to see the list of available models and their supported methods.\" }" + } + ] + } +} diff --git a/FirebaseVertexAI/Tests/Unit/CountTokenResponses/success-total-tokens.json b/FirebaseVertexAI/Tests/Unit/CountTokenResponses/success-total-tokens.json new file mode 100644 index 00000000000..0bccd9e7c3d --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/CountTokenResponses/success-total-tokens.json @@ -0,0 +1,3 @@ +{ + "totalTokens": 6 +} diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-empty-content.txt b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-empty-content.txt new file mode 100644 index 00000000000..5762b515325 --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-empty-content.txt @@ -0,0 +1 @@ +data: {"candidates": [{"content": {},"index": 0}],"promptFeedback": {"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}} diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-error-mid-stream.txt b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-error-mid-stream.txt new file mode 100644 index 00000000000..aeb4eb09389 --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-error-mid-stream.txt @@ -0,0 +1,17 @@ +data: {"candidates": [{"content": {"parts": [{"text": "First "}]},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}],"promptFeedback": {"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}} + +data: {"candidates": [{"content": {"parts": [{"text": "Second "}]},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}]} + +{ + "error": { + "code": 499, + "message": "The operation was cancelled.", + "status": "CANCELLED", + "details": [ + { + "@type": "type.googleapis.com/google.rpc.DebugInfo", + "detail": "[ORIGINAL ERROR] generic::cancelled: " + } + ] + } +} diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-finish-reason-safety.txt b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-finish-reason-safety.txt new file mode 100644 index 00000000000..b73c75cf505 --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-finish-reason-safety.txt @@ -0,0 +1,2 @@ +data: {"candidates": [{"content": {"parts": [{"text": "No"}]},"finishReason": "SAFETY","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "HIGH"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}],"promptFeedback": {"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}} + diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-invalid-json.txt b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-invalid-json.txt new file mode 100644 index 00000000000..3aea78474cb --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-invalid-json.txt @@ -0,0 +1 @@ +data: {"this": [{"is": {"not": [{"a": "valid"}]}, "response": {}}]} diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-malformed-content.txt b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-malformed-content.txt new file mode 100644 index 00000000000..273b6cdc5ca --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-malformed-content.txt @@ -0,0 +1 @@ +data: {"candidates": [{"content": {"missing-parts": true},"index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_TOXICITY","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_VIOLENCE","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}],"promptFeedback": {"safetyRatings": [{"category": "HARM_CATEGORY_TOXICITY","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_VIOLENCE","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}} diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-prompt-blocked-safety.txt b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-prompt-blocked-safety.txt new file mode 100644 index 00000000000..58c914af08e --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-prompt-blocked-safety.txt @@ -0,0 +1,2 @@ +data: {"promptFeedback": {"blockReason": "SAFETY","safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "HIGH"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}} + diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-recitation-no-content.txt b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-recitation-no-content.txt new file mode 100644 index 00000000000..60ec91d6dfc --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-recitation-no-content.txt @@ -0,0 +1,6 @@ +data: {"candidates": [{"content": {"parts": [{"text": "Some information"}],"role": "model"},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}],"promptFeedback": {"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}} + +data: {"candidates": [{"content": {"parts": [{"text": "Some information cited from an external source"}],"role": "model"},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "LOW"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}],"citationMetadata": {"citationSources": [{"startIndex": 30,"endIndex": 179,"uri": "https://www.example.com/some-citation","license": ""}]}}]} + +data: {"candidates": [{"finishReason": "RECITATION","index": 0}]} + diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-unknown-finish-enum.txt b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-unknown-finish-enum.txt new file mode 100644 index 00000000000..6194abd8957 --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-unknown-finish-enum.txt @@ -0,0 +1,11 @@ +data: {"candidates": [{"content": {"parts": [{"text": "**Cats:**\n\n- **Physical Characteristics:**\n - Size: Cats come"}]},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}],"promptFeedback": {"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}} + +data: {"candidates": [{"content": {"parts": [{"text": " in a wide range of sizes, from small breeds like the Singapura to large breeds like the Maine Coon.\n - Fur: Cats have soft, furry coats"}]},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}]} + +data: {"candidates": [{"content": {"parts": [{"text": " that can vary in length and texture depending on the breed.\n - Eyes: Cats have large, expressive eyes that can be various colors, including green, blue, yellow, and hazel.\n - Ears: Cats have pointed, erect ears that are sensitive to sound.\n - Tail: Cats have long"}]},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}]} + +data: {"candidates": [{"content": {"parts": [{"text": ", flexible tails that they use for balance and communication.\n\n- **Behavior and Personality:**\n - Independent: Cats are often described as independent animals that enjoy spending time alone.\n - Affectionate: Despite their independent nature, cats can be very affectionate and form strong bonds with their owners.\n - Playful: Cats are naturally playful and enjoy engaging in activities such as chasing toys, climbing, and pouncing.\n - Curious: Cats are curious creatures that love to explore their surroundings.\n - Vocal: Cats communicate through a variety of vocalizations, including meows, purrs, hisses, and grow"}]},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}]} + +data: {"candidates": [{"content": {"parts": [{"text": "ls.\n\n- **Health and Care:**\n - Diet: Cats are obligate carnivores, meaning they require animal-based protein for optimal health.\n - Grooming: Cats spend a significant amount of time grooming themselves to keep their fur clean and free of mats.\n - Exercise: Cats need regular exercise to stay healthy and active. This can be achieved through play sessions or access to outdoor space.\n - Veterinary Care: Regular veterinary checkups are essential for maintaining a cat's health and detecting any potential health issues early on.\n\n**Dogs:**\n\n- **Physical Characteristics:**\n - Size: Dogs come in a wide range of sizes, from small breeds like the Chihuahua to giant breeds like the Great Dane.\n - Fur: Dogs have fur coats that can vary in length, texture, and color depending on the breed.\n - Eyes: Dogs have expressive eyes that can be various colors, including brown, blue, green, and hazel.\n - Ears: Dogs have floppy or erect ears that are sensitive to sound.\n - Tail: Dogs have long, wagging tails that they use for communication and expressing emotions.\n\n- **Behavior and Personality:**\n - Loyal: Dogs are known for their loyalty and"}]},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}]} + +data: {"candidates": [{"content": {"parts": [{"text": " devotion to their owners.\n - Friendly: Dogs are generally friendly and outgoing animals that enjoy interacting with people and other animals.\n - Playful: Dogs are playful and energetic creatures that love to engage in activities such as fetching, running, and playing with toys.\n - Trainable: Dogs are highly trainable and can learn a variety of commands and tricks.\n - Vocal: Dogs communicate through a variety of vocalizations, including barking, howling, whining, and growling.\n\n- **Health and Care:**\n - Diet: Dogs are omnivores and can eat a variety of foods, including meat, vegetables, and grains.\n - Grooming: Dogs require regular grooming to keep their fur clean and free of mats. The frequency of grooming depends on the breed and coat type.\n - Exercise: Dogs need regular exercise to stay healthy and active. The amount of exercise required varies depending on the breed and age of the dog.\n - Veterinary Care: Regular veterinary checkups are essential for maintaining a dog's health and detecting any potential health issues early on."}]},"finishReason": "FAKE_ENUM","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT_NEW_ENUM","probability": "NEGLIGIBLE_UNKNOWN_ENUM"}]}]} diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-success-basic-reply-long.txt b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-success-basic-reply-long.txt new file mode 100644 index 00000000000..bca95140490 --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-success-basic-reply-long.txt @@ -0,0 +1,19 @@ +data: {"candidates": [{"content": {"parts": [{"text": "1 "}]},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}],"promptFeedback": {"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}} + +data: {"candidates": [{"content": {"parts": [{"text": "2 "}]},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}]} + +data: {"candidates": [{"content": {"parts": [{"text": "3 "}]},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}]} + +data: {"candidates": [{"content": {"parts": [{"text": "4 "}]},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}]} + +data: {"candidates": [{"content": {"parts": [{"text": "5 "}]},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}]} + +data: {"candidates": [{"content": {"parts": [{"text": "6 "}]},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}]} + +data: {"candidates": [{"content": {"parts": [{"text": "7 "}]},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}]} + +data: {"candidates": [{"content": {"parts": [{"text": "8 "}]},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}]} + +data: {"candidates": [{"content": {"parts": [{"text": "9 "}]},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}]} + +data: {"candidates": [{"content": {"parts": [{"text": "10"}]},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}]} diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-success-basic-reply-short.txt b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-success-basic-reply-short.txt new file mode 100644 index 00000000000..a7f5476954e --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-success-basic-reply-short.txt @@ -0,0 +1,2 @@ +data: {"candidates": [{"content": {"parts": [{"text": "Cheyenne"}]},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}],"promptFeedback": {"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}} + diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-success-citations.txt b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-success-citations.txt new file mode 100644 index 00000000000..4f50be7d32f --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-success-citations.txt @@ -0,0 +1,13 @@ +data: {"candidates": [{"content": {"parts": [{"text": "Some information"}],"role": "model"},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}],"promptFeedback": {"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}} + +data: {"candidates": [{"content": {"parts": [{"text": " More information"}],"role": "model"},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}]} + +data: {"candidates": [{"content": {"parts": [{"text": ", Even more information"}],"role": "model"},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}]} + +data: {"candidates": [{"content": {"parts": [{"text": " Some information cited from an external source"}],"role": "model"},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}],"citationMetadata": {"citationSources": [{"startIndex": 574,"endIndex": 705,"uri": "https://www.example.com/citation-1","license": ""},{"startIndex": 899,"endIndex": 1026,"uri": "https://www.example.com/citation-2","license": ""}]}}]} + +data: {"candidates": [{"content": {"parts": [{"text": "More information cited from an external source"}],"role": "model"},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}],"citationMetadata": {"citationSources": [{"startIndex": 574,"endIndex": 705,"uri": "https://www.example.com/citation-3","license": ""},{"startIndex": 899,"endIndex": 1026,"uri": "https://www.example.com/citation-4","license": ""}]}}]} + +data: {"candidates": [{"content": {"parts": [{"text": "Even more information cited from an external source"}],"role": "model"},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}],"citationMetadata": {"citationSources": [{"startIndex": 574,"endIndex": 705,"uri": "https://www.example.com/citation-5","license": ""},{"startIndex": 899,"endIndex": 1026,"uri": "https://www.example.com/citation-6","license": ""}]}}]} + +data: {"candidates": [{"content": {"parts": [{"text": "Physics (YouTube Channel)"}],"role": "model"},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}],"citationMetadata": {"citationSources": [{"startIndex": 574,"endIndex": 705,"uri": "https://www.google.com","license": ""},{"startIndex": 899,"endIndex": 1026,"uri": "https://www.google.com","license": ""}]}}]} diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-success-citations2.txt b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-success-citations2.txt new file mode 100644 index 00000000000..665e993ad75 --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-success-citations2.txt @@ -0,0 +1,5 @@ +data: {"candidates": [{"content": {"role": "model","parts": [{"text": "In the context of the science fiction comedy novel \"The Hitchhiker's Guide to the Galaxy\" by Douglas Adams, the answer to the \"Ultimate Question of"}]},"safetyRatings": [{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE","probabilityScore": 0.055720285,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.062674366},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE","probabilityScore": 0.03904829,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.03339982},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE","probabilityScore": 0.12220858,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.0540987},{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE","probabilityScore": 0.05781161,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.03149938}]}]} + +data: {"candidates": [{"content": {"role": "model","parts": [{"text": " Life, the Universe, and Everything\" is given as 42. However, in real life, there is no universally accepted meaning of life. The meaning"}]},"safetyRatings": [{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE","probabilityScore": 0.05910154,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.038321976},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE","probabilityScore": 0.049589027,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.019271139},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE","probabilityScore": 0.12787028,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.03986249},{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE","probabilityScore": 0.098946586,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.04177388}],"citationMetadata": {"citations": [{"startIndex": 52,"endIndex": 181,"uri": "https://imcsteakhouse.com.au/faqs"}]}}]} + +data: {"candidates": [{"content": {"role": "model","parts": [{"text": " of life is a philosophical question that has been pondered by humans for centuries, and there is no single answer that is widely agreed upon."}]},"finishReason": "STOP","safetyRatings": [{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE","probabilityScore": 0.044764314,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.028870905},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE","probabilityScore": 0.04240383,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.012576347},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE","probabilityScore": 0.10650458,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.034880884},{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE","probabilityScore": 0.11085559,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.03191915}]}],"usageMetadata": {"promptTokenCount": 9,"candidatesTokenCount": 91,"totalTokenCount": 100}} diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-success-unknown-safety-enum.txt b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-success-unknown-safety-enum.txt new file mode 100644 index 00000000000..8f7a8ff3c1c --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-success-unknown-safety-enum.txt @@ -0,0 +1,2 @@ +data: {"candidates": [{"content": {"parts": [{"text": "Cheyenne"}]},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_SOMETHING_NEW","probability": "NEGLIGIBLE"}]}],"promptFeedback": {"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}} + diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-api-key.json b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-api-key.json new file mode 100644 index 00000000000..ecf6f6b53fa --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-api-key.json @@ -0,0 +1,21 @@ +{ + "error": { + "code": 400, + "message": "API key not valid. Please pass a valid API key.", + "status": "INVALID_ARGUMENT", + "details": [ + { + "@type": "type.googleapis.com/google.rpc.ErrorInfo", + "reason": "API_KEY_INVALID", + "domain": "googleapis.com", + "metadata": { + "service": "generativelanguage.googleapis.com" + } + }, + { + "@type": "type.googleapis.com/google.rpc.DebugInfo", + "detail": "Invalid API key: AIzv00G7VmUCUeC-5OglO3hcXM" + } + ] + } +} diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-empty-content.json b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-empty-content.json new file mode 100644 index 00000000000..4e1889660f2 --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-empty-content.json @@ -0,0 +1,28 @@ +{ + "candidates": [ + { + "content": {}, + "index": 0 + } + ], + "promptFeedback": { + "safetyRatings": [ + { + "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "NEGLIGIBLE" + } + ] + } +} diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-finish-reason-recitation-no-content.json b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-finish-reason-recitation-no-content.json new file mode 100644 index 00000000000..0d99cd71d51 --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-finish-reason-recitation-no-content.json @@ -0,0 +1,46 @@ +{ + "candidates": [ + { + "finishReason": "RECITATION", + "index": 0, + "safetyRatings": [ + { + "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "probability": "HIGH" + }, + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "NEGLIGIBLE" + } + ] + } + ], + "promptFeedback": { + "safetyRatings": [ + { + "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "NEGLIGIBLE" + } + ] + } +} diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-finish-reason-safety-no-content.json b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-finish-reason-safety-no-content.json new file mode 100644 index 00000000000..03958d4afc1 --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-finish-reason-safety-no-content.json @@ -0,0 +1,46 @@ +{ + "candidates": [ + { + "finishReason": "SAFETY", + "index": 0, + "safetyRatings": [ + { + "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "probability": "HIGH" + }, + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "NEGLIGIBLE" + } + ] + } + ], + "promptFeedback": { + "safetyRatings": [ + { + "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "NEGLIGIBLE" + } + ] + } +} diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-finish-reason-safety.json b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-finish-reason-safety.json new file mode 100644 index 00000000000..3249c739cf8 --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-finish-reason-safety.json @@ -0,0 +1,53 @@ +{ + "candidates": [ + { + "content": { + "parts": [ + { + "text": "No" + } + ] + }, + "finishReason": "SAFETY", + "index": 0, + "safetyRatings": [ + { + "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "probability": "HIGH" + }, + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "NEGLIGIBLE" + } + ] + } + ], + "promptFeedback": { + "safetyRatings": [ + { + "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "NEGLIGIBLE" + } + ] + } +} diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-image-rejected.json b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-image-rejected.json new file mode 100644 index 00000000000..9dacdc71e7a --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-image-rejected.json @@ -0,0 +1,13 @@ +{ + "error": { + "code": 400, + "message": "Request contains an invalid argument.", + "status": "INVALID_ARGUMENT", + "details": [ + { + "@type": "type.googleapis.com/google.rpc.DebugInfo", + "detail": "[ORIGINAL ERROR] generic::invalid_argument: invalid status photos.thumbnailer.Status.Code::5: Source image 0 too short" + } + ] + } +} diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-invalid-response.json b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-invalid-response.json new file mode 100644 index 00000000000..49d05e1840b --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-invalid-response.json @@ -0,0 +1,14 @@ +{ + "this": [ + { + "is": { + "not": [ + { + "a": "valid" + } + ] + }, + "response": {} + } + ] +} diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-malformed-content.json b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-malformed-content.json new file mode 100644 index 00000000000..737f2e08548 --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-malformed-content.json @@ -0,0 +1,30 @@ +{ + "candidates": [ + { + "content": { + "invalid-field": true + }, + "index": 0 + } + ], + "promptFeedback": { + "safetyRatings": [ + { + "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "NEGLIGIBLE" + } + ] + } +} diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-prompt-blocked-safety.json b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-prompt-blocked-safety.json new file mode 100644 index 00000000000..9d2abbb23d6 --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-prompt-blocked-safety.json @@ -0,0 +1,23 @@ +{ + "promptFeedback": { + "blockReason": "SAFETY", + "safetyRatings": [ + { + "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "probability": "HIGH" + }, + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "NEGLIGIBLE" + } + ] + } +} diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-unknown-enum-finish-reason.json b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-unknown-enum-finish-reason.json new file mode 100644 index 00000000000..1adbc70cae4 --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-unknown-enum-finish-reason.json @@ -0,0 +1,53 @@ +{ + "candidates": [ + { + "content": { + "parts": [ + { + "text": "Some text" + } + ] + }, + "finishReason": "FAKE_NEW_FINISH_REASON", + "index": 0, + "safetyRatings": [ + { + "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "NEGLIGIBLE" + } + ] + } + ], + "promptFeedback": { + "safetyRatings": [ + { + "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "NEGLIGIBLE" + } + ] + } +} diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-unknown-enum-prompt-blocked.json b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-unknown-enum-prompt-blocked.json new file mode 100644 index 00000000000..a00cbb72cfc --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-unknown-enum-prompt-blocked.json @@ -0,0 +1,23 @@ +{ + "promptFeedback": { + "blockReason": "FAKE_NEW_BLOCK_REASON", + "safetyRatings": [ + { + "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "NEGLIGIBLE" + } + ] + } +} diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-unknown-model.json b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-unknown-model.json new file mode 100644 index 00000000000..60b3f55c978 --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-unknown-model.json @@ -0,0 +1,13 @@ +{ + "error": { + "code": 404, + "message": "models/unknown is not found for API version v1, or is not supported for GenerateContent. Call ListModels to see the list of available models and their supported methods.", + "status": "NOT_FOUND", + "details": [ + { + "@type": "type.googleapis.com/google.rpc.DebugInfo", + "detail": "[ORIGINAL ERROR] generic::not_found: models/unknown is not found for API version v1, or is not supported for GenerateContent. Call ListModels to see the list of available models and their supported methods. [google.rpc.error_details_ext] { message: \"models/unknown is not found for API version v1, or is not supported for GenerateContent. Call ListModels to see the list of available models and their supported methods.\" }" + } + ] + } +} diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-unsupported-user-location.json b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-unsupported-user-location.json new file mode 100644 index 00000000000..c4c2ace4e20 --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-failure-unsupported-user-location.json @@ -0,0 +1,13 @@ +{ + "error": { + "code": 400, + "message": "User location is not supported for the API use.", + "status": "FAILED_PRECONDITION", + "details": [ + { + "@type": "type.googleapis.com/google.rpc.DebugInfo", + "detail": "[ORIGINAL ERROR] generic::failed_precondition: User location is not supported for the API use. [google.rpc.error_details_ext] { message: \"User location is not supported for the API use.\" }" + } + ] + } +} diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-success-basic-reply-long.json b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-success-basic-reply-long.json new file mode 100644 index 00000000000..59b84de92fe --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-success-basic-reply-long.json @@ -0,0 +1,54 @@ +{ + "candidates": [ + { + "content": { + "parts": [ + { + "text": "You can ask me a wide range of questions on various topics. Here are some examples:\n\n1. **General Knowledge:**\n - What is the capital of France?\n - Who painted the Mona Lisa?\n - What is the largest ocean in the world?\n\n2. **Science and Technology:**\n - How does a computer work?\n - What is the difference between a virus and a bacteria?\n - What are the latest advancements in artificial intelligence?\n\n3. **History and Culture:**\n - Who was the first president of the United States?\n - What is the significance of the Great Wall of China?\n - What are some of the most famous works of Shakespeare?\n\n4. **Current Events:**\n - What is the latest news about the COVID-19 pandemic?\n - Who is the current president of Ukraine?\n - What are the major issues being discussed in the upcoming election?\n\n5. **Personal Questions:**\n - What are your hobbies?\n - What is your favorite book or movie?\n - What are your thoughts on the future of technology?\n\n6. **Fun and Games:**\n - Can you tell me a joke?\n - What is the answer to this riddle: \"I have keys but no locks. I have space but no room. You can enter, but can't go outside.\" (Answer: a keyboard)\n - Let's play a game of 20 questions.\n\n7. **Hypothetical Questions:**\n - What would you do if you won the lottery?\n - What would happen if time travel were possible?\n - What is the meaning of life?\n\n8. **Philosophical Questions:**\n - What is the nature of reality?\n - Does free will exist?\n - What is the difference between right and wrong?\n\n9. **Creative Questions:**\n - Write a poem about a sunset.\n - Design a logo for a new company.\n - Compose a song about your favorite season.\n\n10. **Technical Questions:**\n - How can I improve the performance of my computer?\n - What is the best way to troubleshoot a network issue?\n - How do I create a website using HTML and CSS?\n\nRemember, I am still under development and may not be able to answer all questions perfectly. However, I will do my best to provide you with accurate and informative responses." + } + ], + "role": "model" + }, + "finishReason": "STOP", + "index": 0, + "safetyRatings": [ + { + "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "NEGLIGIBLE" + } + ] + } + ], + "promptFeedback": { + "safetyRatings": [ + { + "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "NEGLIGIBLE" + } + ] + } +} diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-success-basic-reply-short.json b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-success-basic-reply-short.json new file mode 100644 index 00000000000..40a9a6da58e --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-success-basic-reply-short.json @@ -0,0 +1,54 @@ +{ + "candidates": [ + { + "content": { + "parts": [ + { + "text": "Mountain View, California, United States" + } + ], + "role": "model" + }, + "finishReason": "STOP", + "index": 0, + "safetyRatings": [ + { + "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "NEGLIGIBLE" + } + ] + } + ], + "promptFeedback": { + "safetyRatings": [ + { + "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "NEGLIGIBLE" + } + ] + } +} diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-success-citations.json b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-success-citations.json new file mode 100644 index 00000000000..ac99fcf0944 --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-success-citations.json @@ -0,0 +1,64 @@ +{ + "candidates": [ + { + "content": { + "parts": [ + { + "text": "Some information cited from an external source" + } + ], + "role": "model" + }, + "finishReason": "STOP", + "index": 0, + "safetyRatings": [ + { + "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "NEGLIGIBLE" + } + ], + "citationMetadata": { + "citationSources": [ + { + "startIndex": 574, + "endIndex": 705, + "uri": "https://www.example.com/some-citation", + "license": "" + } + ] + } + } + ], + "promptFeedback": { + "safetyRatings": [ + { + "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "NEGLIGIBLE" + } + ] + } +} diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-success-missing-safety-ratings.json b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-success-missing-safety-ratings.json new file mode 100644 index 00000000000..645d3e6425e --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-success-missing-safety-ratings.json @@ -0,0 +1,16 @@ +{ + "candidates": [ + { + "content": { + "parts": [ + { + "text": "This is the generated content." + } + ] + }, + "index": 0 + } + ], + "promptFeedback": { + } +} diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-success-quote-reply.json b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-success-quote-reply.json new file mode 100644 index 00000000000..e5bf4f0b497 --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-success-quote-reply.json @@ -0,0 +1,54 @@ +{ + "candidates": [ + { + "content": { + "parts": [ + { + "text": "Google's mission is to \"organize the world's information and make it universally accessible and useful.\"" + } + ], + "role": "model" + }, + "finishReason": "STOP", + "index": 0, + "safetyRatings": [ + { + "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "NEGLIGIBLE" + } + ] + } + ], + "promptFeedback": { + "safetyRatings": [ + { + "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "NEGLIGIBLE" + } + ] + } +} diff --git a/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-success-unknown-enum-safety-ratings.json b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-success-unknown-enum-safety-ratings.json new file mode 100644 index 00000000000..3c47a21322c --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerateContentResponses/unary-success-unknown-enum-safety-ratings.json @@ -0,0 +1,46 @@ +{ + "candidates": [ + { + "content": { + "parts": [ + { + "text": "Some text" + } + ], + "role": "model" + }, + "finishReason": "STOP", + "index": 0, + "safetyRatings": [ + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "MEDIUM" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "FAKE_NEW_HARM_PROBABILITY" + }, + { + "category": "FAKE_NEW_HARM_CATEGORY", + "probability": "HIGH" + } + ] + } + ], + "promptFeedback": { + "safetyRatings": [ + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "MEDIUM" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "FAKE_NEW_HARM_PROBABILITY" + }, + { + "category": "FAKE_NEW_HARM_CATEGORY", + "probability": "HIGH" + } + ] + } +} diff --git a/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift b/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift new file mode 100644 index 00000000000..b835a9ea05f --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift @@ -0,0 +1,930 @@ +// Copyright 2023 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. + +@testable import GoogleGenerativeAI +import XCTest + +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, *) +final class GenerativeModelTests: XCTestCase { + let testPrompt = "What sorts of questions can I ask you?" + let safetyRatingsNegligible: [SafetyRating] = [ + .init(category: .sexuallyExplicit, probability: .negligible), + .init(category: .hateSpeech, probability: .negligible), + .init(category: .harassment, probability: .negligible), + .init(category: .dangerousContent, probability: .negligible), + ] + + var urlSession: URLSession! + var model: GenerativeModel! + + override func setUp() async throws { + let configuration = URLSessionConfiguration.default + configuration.protocolClasses = [MockURLProtocol.self] + urlSession = try XCTUnwrap(URLSession(configuration: configuration)) + model = GenerativeModel(name: "my-model", apiKey: "API_KEY", urlSession: urlSession) + } + + override func tearDown() { + MockURLProtocol.requestHandler = nil + } + + // MARK: - Generate Content + + func testGenerateContent_success_basicReplyLong() async throws { + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "unary-success-basic-reply-long", + withExtension: "json" + ) + + let response = try await model.generateContent(testPrompt) + + XCTAssertEqual(response.candidates.count, 1) + let candidate = try XCTUnwrap(response.candidates.first) + let finishReason = try XCTUnwrap(candidate.finishReason) + XCTAssertEqual(finishReason, .stop) + XCTAssertEqual(candidate.safetyRatings, safetyRatingsNegligible) + XCTAssertEqual(candidate.content.parts.count, 1) + let part = try XCTUnwrap(candidate.content.parts.first) + let partText = try XCTUnwrap(part.text) + XCTAssertTrue(partText.hasPrefix("You can ask me a wide range of questions")) + XCTAssertEqual(response.text, partText) + let promptFeedback = try XCTUnwrap(response.promptFeedback) + XCTAssertNil(promptFeedback.blockReason) + XCTAssertEqual(promptFeedback.safetyRatings, safetyRatingsNegligible) + } + + func testGenerateContent_success_basicReplyShort() async throws { + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "unary-success-basic-reply-short", + withExtension: "json" + ) + + let response = try await model.generateContent(testPrompt) + + XCTAssertEqual(response.candidates.count, 1) + let candidate = try XCTUnwrap(response.candidates.first) + let finishReason = try XCTUnwrap(candidate.finishReason) + XCTAssertEqual(finishReason, .stop) + XCTAssertEqual(candidate.safetyRatings, safetyRatingsNegligible) + XCTAssertEqual(candidate.content.parts.count, 1) + let part = try XCTUnwrap(candidate.content.parts.first) + XCTAssertEqual(part.text, "Mountain View, California, United States") + XCTAssertEqual(response.text, part.text) + let promptFeedback = try XCTUnwrap(response.promptFeedback) + XCTAssertNil(promptFeedback.blockReason) + XCTAssertEqual(promptFeedback.safetyRatings, safetyRatingsNegligible) + } + + func testGenerateContent_success_citations() async throws { + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "unary-success-citations", + withExtension: "json" + ) + + let response = try await model.generateContent(testPrompt) + + XCTAssertEqual(response.candidates.count, 1) + let candidate = try XCTUnwrap(response.candidates.first) + XCTAssertEqual(candidate.content.parts.count, 1) + XCTAssertEqual(response.text, "Some information cited from an external source") + let citationMetadata = try XCTUnwrap(candidate.citationMetadata) + XCTAssertEqual(citationMetadata.citationSources.count, 1) + let citationSource = try XCTUnwrap(citationMetadata.citationSources.first) + XCTAssertEqual(citationSource.uri, "https://www.example.com/some-citation") + XCTAssertEqual(citationSource.startIndex, 574) + XCTAssertEqual(citationSource.endIndex, 705) + XCTAssertEqual(citationSource.license, "") + } + + func testGenerateContent_success_quoteReply() async throws { + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "unary-success-quote-reply", + withExtension: "json" + ) + + let response = try await model.generateContent(testPrompt) + + XCTAssertEqual(response.candidates.count, 1) + let candidate = try XCTUnwrap(response.candidates.first) + let finishReason = try XCTUnwrap(candidate.finishReason) + XCTAssertEqual(finishReason, .stop) + XCTAssertEqual(candidate.safetyRatings, safetyRatingsNegligible) + XCTAssertEqual(candidate.content.parts.count, 1) + let part = try XCTUnwrap(candidate.content.parts.first) + let partText = try XCTUnwrap(part.text) + XCTAssertTrue(partText.hasPrefix("Google")) + XCTAssertEqual(response.text, part.text) + let promptFeedback = try XCTUnwrap(response.promptFeedback) + XCTAssertNil(promptFeedback.blockReason) + XCTAssertEqual(promptFeedback.safetyRatings, safetyRatingsNegligible) + } + + func testGenerateContent_success_unknownEnum_safetyRatings() async throws { + let expectedSafetyRatings = [ + SafetyRating(category: .harassment, probability: .medium), + SafetyRating(category: .dangerousContent, probability: .unknown), + SafetyRating(category: .unknown, probability: .high), + ] + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "unary-success-unknown-enum-safety-ratings", + withExtension: "json" + ) + + let response = try await model.generateContent(testPrompt) + + XCTAssertEqual(response.text, "Some text") + XCTAssertEqual(response.candidates.first?.safetyRatings, expectedSafetyRatings) + XCTAssertEqual(response.promptFeedback?.safetyRatings, expectedSafetyRatings) + } + + func testGenerateContent_success_prefixedModelName() async throws { + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "unary-success-basic-reply-short", + withExtension: "json" + ) + let model = GenerativeModel( + // Model name is prefixed with "models/". + name: "models/test-model", + apiKey: "API_KEY", + urlSession: urlSession + ) + + _ = try await model.generateContent(testPrompt) + } + + func testGenerateContent_failure_invalidAPIKey() async throws { + let expectedStatusCode = 400 + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "unary-failure-api-key", + withExtension: "json", + statusCode: expectedStatusCode + ) + + do { + _ = try await model.generateContent(testPrompt) + XCTFail("Should throw GenerateContentError.internalError; no error thrown.") + } catch let GenerateContentError.invalidAPIKey(message) { + XCTAssertEqual(message, "API key not valid. Please pass a valid API key.") + } catch { + XCTFail("Should throw GenerateContentError.invalidAPIKey; error thrown: \(error)") + } + } + + func testGenerateContent_failure_emptyContent() async throws { + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "unary-failure-empty-content", + withExtension: "json" + ) + + do { + _ = try await model.generateContent(testPrompt) + XCTFail("Should throw GenerateContentError.internalError; no error thrown.") + } catch let GenerateContentError + .internalError(underlying: invalidCandidateError as InvalidCandidateError) { + guard case let .emptyContent(decodingError) = invalidCandidateError else { + XCTFail("Not an InvalidCandidateError.emptyContent error: \(invalidCandidateError)") + return + } + _ = try XCTUnwrap(decodingError as? DecodingError, + "Not a DecodingError: \(decodingError)") + } catch { + XCTFail("Should throw GenerateContentError.internalError; error thrown: \(error)") + } + } + + func testGenerateContent_failure_finishReasonSafety() async throws { + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "unary-failure-finish-reason-safety", + withExtension: "json" + ) + + do { + _ = try await model.generateContent(testPrompt) + XCTFail("Should throw") + } catch let GenerateContentError.responseStoppedEarly(reason, response) { + XCTAssertEqual(reason, .safety) + XCTAssertEqual(response.text, "No") + } catch { + XCTFail("Should throw a responseStoppedEarly") + } + } + + func testGenerateContent_failure_finishReasonSafety_noContent() async throws { + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "unary-failure-finish-reason-safety-no-content", + withExtension: "json" + ) + + do { + _ = try await model.generateContent(testPrompt) + XCTFail("Should throw") + } catch let GenerateContentError.responseStoppedEarly(reason, response) { + XCTAssertEqual(reason, .safety) + XCTAssertNil(response.text) + } catch { + XCTFail("Should throw a responseStoppedEarly") + } + } + + func testGenerateContent_failure_imageRejected() async throws { + let expectedStatusCode = 400 + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "unary-failure-image-rejected", + withExtension: "json", + statusCode: 400 + ) + + do { + _ = try await model.generateContent(testPrompt) + XCTFail("Should throw GenerateContentError.internalError; no error thrown.") + } catch let GenerateContentError.internalError(underlying: rpcError as RPCError) { + XCTAssertEqual(rpcError.status, .invalidArgument) + XCTAssertEqual(rpcError.httpResponseCode, expectedStatusCode) + XCTAssertEqual(rpcError.message, "Request contains an invalid argument.") + } catch { + XCTFail("Should throw GenerateContentError.internalError; error thrown: \(error)") + } + } + + func testGenerateContent_failure_promptBlockedSafety() async throws { + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "unary-failure-prompt-blocked-safety", + withExtension: "json" + ) + + do { + _ = try await model.generateContent(testPrompt) + XCTFail("Should throw") + } catch let GenerateContentError.promptBlocked(response) { + XCTAssertNil(response.text) + } catch { + XCTFail("Should throw a promptBlocked") + } + } + + func testGenerateContent_failure_unknownEnum_finishReason() async throws { + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "unary-failure-unknown-enum-finish-reason", + withExtension: "json" + ) + + do { + _ = try await model.generateContent(testPrompt) + XCTFail("Should throw") + } catch let GenerateContentError.responseStoppedEarly(reason, response) { + XCTAssertEqual(reason, .unknown) + XCTAssertEqual(response.text, "Some text") + } catch { + XCTFail("Should throw a responseStoppedEarly") + } + } + + func testGenerateContent_failure_unknownEnum_promptBlocked() async throws { + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "unary-failure-unknown-enum-prompt-blocked", + withExtension: "json" + ) + + do { + _ = try await model.generateContent(testPrompt) + XCTFail("Should throw") + } catch let GenerateContentError.promptBlocked(response) { + let promptFeedback = try XCTUnwrap(response.promptFeedback) + XCTAssertEqual(promptFeedback.blockReason, .unknown) + } catch { + XCTFail("Should throw a promptBlocked") + } + } + + func testGenerateContent_failure_unknownModel() async throws { + let expectedStatusCode = 404 + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "unary-failure-unknown-model", + withExtension: "json", + statusCode: 404 + ) + + do { + _ = try await model.generateContent(testPrompt) + XCTFail("Should throw GenerateContentError.internalError; no error thrown.") + } catch let GenerateContentError.internalError(underlying: rpcError as RPCError) { + XCTAssertEqual(rpcError.status, .notFound) + XCTAssertEqual(rpcError.httpResponseCode, expectedStatusCode) + XCTAssertTrue(rpcError.message.hasPrefix("models/unknown is not found")) + } catch { + XCTFail("Should throw GenerateContentError.internalError; error thrown: \(error)") + } + } + + func testGenerateContent_failure_unsupportedUserLocation() async throws { + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "unary-failure-unsupported-user-location", + withExtension: "json", + statusCode: 400 + ) + + do { + _ = try await model.generateContent(testPrompt) + XCTFail("Should throw GenerateContentError.unsupportedUserLocation; no error thrown.") + } catch GenerateContentError.unsupportedUserLocation { + return + } + + XCTFail("Expected an unsupported user location error.") + } + + func testGenerateContent_failure_nonHTTPResponse() async throws { + MockURLProtocol.requestHandler = try nonHTTPRequestHandler() + + var responseError: Error? + var content: GenerateContentResponse? + do { + content = try await model.generateContent(testPrompt) + } catch { + responseError = error + } + + XCTAssertNil(content) + XCTAssertNotNil(responseError) + let generateContentError = try XCTUnwrap(responseError as? GenerateContentError) + guard case let .internalError(underlyingError) = generateContentError else { + XCTFail("Not an internal error: \(generateContentError)") + return + } + XCTAssertEqual(underlyingError.localizedDescription, "Response was not an HTTP response.") + } + + func testGenerateContent_failure_invalidResponse() async throws { + MockURLProtocol.requestHandler = try httpRequestHandler( + forResource: "unary-failure-invalid-response", + withExtension: "json" + ) + + var responseError: Error? + var content: GenerateContentResponse? + do { + content = try await model.generateContent(testPrompt) + } catch { + responseError = error + } + + XCTAssertNil(content) + XCTAssertNotNil(responseError) + let generateContentError = try XCTUnwrap(responseError as? GenerateContentError) + guard case let .internalError(underlyingError) = generateContentError else { + XCTFail("Not an internal error: \(generateContentError)") + return + } + let decodingError = try XCTUnwrap(underlyingError as? DecodingError) + guard case let .dataCorrupted(context) = decodingError else { + XCTFail("Not a data corrupted error: \(decodingError)") + return + } + XCTAssert(context.debugDescription.hasPrefix("Failed to decode GenerateContentResponse")) + } + + func testGenerateContent_failure_malformedContent() async throws { + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "unary-failure-malformed-content", + withExtension: "json" + ) + + var responseError: Error? + var content: GenerateContentResponse? + do { + content = try await model.generateContent(testPrompt) + } catch { + responseError = error + } + + XCTAssertNil(content) + XCTAssertNotNil(responseError) + let generateContentError = try XCTUnwrap(responseError as? GenerateContentError) + guard case let .internalError(underlyingError) = generateContentError else { + XCTFail("Not an internal error: \(generateContentError)") + return + } + let invalidCandidateError = try XCTUnwrap(underlyingError as? InvalidCandidateError) + guard case let .malformedContent(malformedContentUnderlyingError) = invalidCandidateError else { + XCTFail("Not a malformed content error: \(invalidCandidateError)") + return + } + _ = try XCTUnwrap( + malformedContentUnderlyingError as? DecodingError, + "Not a decoding error: \(malformedContentUnderlyingError)" + ) + } + + func testGenerateContentMissingSafetyRatings() async throws { + MockURLProtocol.requestHandler = try httpRequestHandler( + forResource: "unary-success-missing-safety-ratings", + withExtension: "json" + ) + + let content = try await model.generateContent(testPrompt) + let promptFeedback = try XCTUnwrap(content.promptFeedback) + XCTAssertEqual(promptFeedback.safetyRatings.count, 0) + XCTAssertEqual(content.text, "This is the generated content.") + } + + func testGenerateContent_requestOptions_customTimeout() async throws { + let expectedTimeout = 150.0 + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "unary-success-basic-reply-short", + withExtension: "json", + timeout: expectedTimeout + ) + let requestOptions = RequestOptions(timeout: expectedTimeout) + model = GenerativeModel( + name: "my-model", + apiKey: "API_KEY", + requestOptions: requestOptions, + urlSession: urlSession + ) + + let response = try await model.generateContent(testPrompt) + + XCTAssertEqual(response.candidates.count, 1) + } + + // MARK: - Generate Content (Streaming) + + func testGenerateContentStream_failureInvalidAPIKey() async throws { + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "unary-failure-api-key", + withExtension: "json" + ) + + do { + let stream = model.generateContentStream("Hi") + for try await _ in stream { + XCTFail("No content is there, this shouldn't happen.") + } + } catch GenerateContentError.invalidAPIKey { + // invalidAPIKey error is as expected, nothing else to check. + return + } + + XCTFail("Should have caught an error.") + } + + func testGenerateContentStream_failureEmptyContent() async throws { + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "streaming-failure-empty-content", + withExtension: "txt" + ) + + do { + let stream = model.generateContentStream("Hi") + for try await _ in stream { + XCTFail("No content is there, this shouldn't happen.") + } + } catch GenerateContentError.internalError(_ as InvalidCandidateError) { + // Underlying error is as expected, nothing else to check. + return + } + + XCTFail("Should have caught an error.") + } + + func testGenerateContentStream_failureFinishReasonSafety() async throws { + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "streaming-failure-finish-reason-safety", + withExtension: "txt" + ) + + do { + let stream = model.generateContentStream("Hi") + for try await _ in stream { + XCTFail("Content shouldn't be shown, this shouldn't happen.") + } + } catch let GenerateContentError.responseStoppedEarly(reason, _) { + XCTAssertEqual(reason, .safety) + return + } + + XCTFail("Should have caught an error.") + } + + func testGenerateContentStream_failurePromptBlockedSafety() async throws { + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "streaming-failure-prompt-blocked-safety", + withExtension: "txt" + ) + + do { + let stream = model.generateContentStream("Hi") + for try await _ in stream { + XCTFail("Content shouldn't be shown, this shouldn't happen.") + } + } catch let GenerateContentError.promptBlocked(response) { + XCTAssertEqual(response.promptFeedback?.blockReason, .safety) + return + } + + XCTFail("Should have caught an error.") + } + + func testGenerateContentStream_failureUnknownFinishEnum() async throws { + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "streaming-failure-unknown-finish-enum", + withExtension: "txt" + ) + + let stream = model.generateContentStream("Hi") + do { + for try await content in stream { + XCTAssertNotNil(content.text) + } + } catch let GenerateContentError.responseStoppedEarly(reason, _) { + XCTAssertEqual(reason, .unknown) + return + } + + XCTFail("Should have caught an error.") + } + + func testGenerateContentStream_successBasicReplyLong() async throws { + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "streaming-success-basic-reply-long", + withExtension: "txt" + ) + + var responses = 0 + let stream = model.generateContentStream("Hi") + for try await content in stream { + XCTAssertNotNil(content.text) + responses += 1 + } + + XCTAssertEqual(responses, 10) + } + + func testGenerateContentStream_successBasicReplyShort() async throws { + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "streaming-success-basic-reply-short", + withExtension: "txt" + ) + + var responses = 0 + let stream = model.generateContentStream("Hi") + for try await content in stream { + XCTAssertNotNil(content.text) + responses += 1 + } + + XCTAssertEqual(responses, 1) + } + + func testGenerateContentStream_successUnknownSafetyEnum() async throws { + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "streaming-success-unknown-safety-enum", + withExtension: "txt" + ) + + var hadUnknown = false + let stream = model.generateContentStream("Hi") + for try await content in stream { + XCTAssertNotNil(content.text) + if let ratings = content.candidates.first?.safetyRatings, + ratings.contains(where: { $0.category == .unknown }) { + hadUnknown = true + } + } + + XCTAssertTrue(hadUnknown) + } + + func testGenerateContentStream_successWithCitations() async throws { + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "streaming-success-citations", + withExtension: "txt" + ) + + let stream = model.generateContentStream("Hi") + var citations: [Citation] = [] + for try await content in stream { + XCTAssertNotNil(content.text) + let candidate = try XCTUnwrap(content.candidates.first) + XCTAssertEqual(candidate.finishReason, .stop) + if let sources = candidate.citationMetadata?.citationSources { + citations.append(contentsOf: sources) + } + } + + XCTAssertEqual(citations.count, 8) + XCTAssertTrue(citations + .contains(where: { $0.startIndex == 574 && $0.endIndex == 705 && !$0.uri.isEmpty })) + XCTAssertTrue(citations + .contains(where: { $0.startIndex == 899 && $0.endIndex == 1026 && !$0.uri.isEmpty })) + } + + func testGenerateContentStream_errorMidStream() async throws { + MockURLProtocol.requestHandler = try httpRequestHandler( + forResource: "streaming-failure-error-mid-stream", + withExtension: "txt" + ) + + var responseCount = 0 + do { + let stream = model.generateContentStream("Hi") + for try await content in stream { + XCTAssertNotNil(content.text) + responseCount += 1 + } + } catch let GenerateContentError.internalError(rpcError as RPCError) { + XCTAssertEqual(rpcError.httpResponseCode, 499) + XCTAssertEqual(rpcError.status, .cancelled) + + // Check the content count is correct. + XCTAssertEqual(responseCount, 2) + return + } + + XCTFail("Expected an internalError with an RPCError.") + } + + func testGenerateContentStream_nonHTTPResponse() async throws { + MockURLProtocol.requestHandler = try nonHTTPRequestHandler() + + let stream = model.generateContentStream("Hi") + do { + for try await content in stream { + XCTFail("Unexpected content in stream: \(content)") + } + } catch let GenerateContentError.internalError(underlying) { + XCTAssertEqual(underlying.localizedDescription, "Response was not an HTTP response.") + return + } + + XCTFail("Expected an internal error.") + } + + func testGenerateContentStream_invalidResponse() async throws { + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "streaming-failure-invalid-json", + withExtension: "txt" + ) + + let stream = model.generateContentStream(testPrompt) + do { + for try await content in stream { + XCTFail("Unexpected content in stream: \(content)") + } + } catch let GenerateContentError.internalError(underlying as DecodingError) { + guard case let .dataCorrupted(context) = underlying else { + XCTFail("Not a data corrupted error: \(underlying)") + return + } + XCTAssert(context.debugDescription.hasPrefix("Failed to decode GenerateContentResponse")) + return + } + + XCTFail("Expected an internal error.") + } + + func testGenerateContentStream_malformedContent() async throws { + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "streaming-failure-malformed-content", + withExtension: "txt" + ) + + let stream = model.generateContentStream(testPrompt) + do { + for try await content in stream { + XCTFail("Unexpected content in stream: \(content)") + } + } catch let GenerateContentError.internalError(underlyingError as InvalidCandidateError) { + guard case let .malformedContent(contentError) = underlyingError else { + XCTFail("Not a malformed content error: \(underlyingError)") + return + } + + XCTAssert(contentError is DecodingError) + return + } + + XCTFail("Expected an internal decoding error.") + } + + func testGenerateContentStream_failure_unsupportedUserLocation() async throws { + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "unary-failure-unsupported-user-location", + withExtension: "json", + statusCode: 400 + ) + + let stream = model.generateContentStream(testPrompt) + do { + for try await content in stream { + XCTFail("Unexpected content in stream: \(content)") + } + } catch GenerateContentError.unsupportedUserLocation { + return + } + + XCTFail("Expected an unsupported user location error.") + } + + func testGenerateContentStream_requestOptions_customTimeout() async throws { + let expectedTimeout = 150.0 + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "streaming-success-basic-reply-short", + withExtension: "txt", + timeout: expectedTimeout + ) + let requestOptions = RequestOptions(timeout: expectedTimeout) + model = GenerativeModel( + name: "my-model", + apiKey: "API_KEY", + requestOptions: requestOptions, + urlSession: urlSession + ) + + var responses = 0 + let stream = model.generateContentStream(testPrompt) + for try await content in stream { + XCTAssertNotNil(content.text) + responses += 1 + } + + XCTAssertEqual(responses, 1) + } + + // MARK: - Count Tokens + + func testCountTokens_succeeds() async throws { + MockURLProtocol.requestHandler = try httpRequestHandler( + forResource: "success-total-tokens", + withExtension: "json" + ) + + let response = try await model.countTokens("Why is the sky blue?") + XCTAssertEqual(response.totalTokens, 6) + } + + func testCountTokens_modelNotFound() async throws { + MockURLProtocol.requestHandler = try httpRequestHandler( + forResource: "failure-model-not-found", withExtension: "json", + statusCode: 404 + ) + + do { + _ = try await model.countTokens("Why is the sky blue?") + XCTFail("Request should not have succeeded.") + } catch let CountTokensError.internalError(rpcError as RPCError) { + XCTAssertEqual(rpcError.httpResponseCode, 404) + XCTAssertEqual(rpcError.status, .notFound) + XCTAssert(rpcError.message.hasPrefix("models/test-model-name is not found")) + return + } + + XCTFail("Expected internal RPCError.") + } + + func testCountTokens_requestOptions_customTimeout() async throws { + let expectedTimeout = 150.0 + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "success-total-tokens", + withExtension: "json", + timeout: expectedTimeout + ) + let requestOptions = RequestOptions(timeout: expectedTimeout) + model = GenerativeModel( + name: "my-model", + apiKey: "API_KEY", + requestOptions: requestOptions, + urlSession: urlSession + ) + + let response = try await model.countTokens(testPrompt) + + XCTAssertEqual(response.totalTokens, 6) + } + + // MARK: - Model Resource Name + + func testModelResourceName_noPrefix() async throws { + let modelName = "my-model" + let modelResourceName = "models/\(modelName)" + + model = GenerativeModel(name: modelName, apiKey: "API_KEY") + + XCTAssertEqual(model.modelResourceName, modelResourceName) + } + + func testModelResourceName_modelsPrefix() async throws { + let modelResourceName = "models/my-model" + + model = GenerativeModel(name: modelResourceName, apiKey: "API_KEY") + + XCTAssertEqual(model.modelResourceName, modelResourceName) + } + + func testModelResourceName_tunedModelsPrefix() async throws { + let tunedModelResourceName = "tunedModels/my-model" + + model = GenerativeModel(name: tunedModelResourceName, apiKey: "API_KEY") + + XCTAssertEqual(model.modelResourceName, tunedModelResourceName) + } + + // MARK: - Helpers + + private func nonHTTPRequestHandler() throws -> ((URLRequest) -> ( + URLResponse, + AsyncLineSequence? + )) { + return { request in + // This is *not* an HTTPURLResponse + let response = URLResponse( + url: request.url!, + mimeType: nil, + expectedContentLength: 0, + textEncodingName: nil + ) + return (response, nil) + } + } + + private func httpRequestHandler(forResource name: String, + withExtension ext: String, + statusCode: Int = 200, + timeout: TimeInterval = URLRequest + .defaultTimeoutInterval()) throws -> ((URLRequest) throws -> ( + URLResponse, + AsyncLineSequence? + )) { + let fileURL = try XCTUnwrap(Bundle.module.url(forResource: name, withExtension: ext)) + return { request in + let requestURL = try XCTUnwrap(request.url) + XCTAssertEqual(requestURL.path.occurrenceCount(of: "models/"), 1) + XCTAssertEqual(request.timeoutInterval, timeout) + let response = try XCTUnwrap(HTTPURLResponse( + url: requestURL, + statusCode: statusCode, + httpVersion: nil, + headerFields: nil + )) + return (response, fileURL.lines) + } + } +} + +private extension String { + /// Returns the number of occurrences of `substring` in the `String`. + func occurrenceCount(of substring: String) -> Int { + return components(separatedBy: substring).count - 1 + } +} + +private extension URLRequest { + /// Returns the default `timeoutInterval` for a `URLRequest`. + static func defaultTimeoutInterval() -> TimeInterval { + let placeholderURL = URL(string: "https://example.com")! + return URLRequest(url: placeholderURL).timeoutInterval + } +} diff --git a/FirebaseVertexAI/Tests/Unit/GoogleAITests.swift b/FirebaseVertexAI/Tests/Unit/GoogleAITests.swift new file mode 100644 index 00000000000..cbc92527a1c --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/GoogleAITests.swift @@ -0,0 +1,175 @@ +// Copyright 2023 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. + +import GoogleGenerativeAI +import XCTest +#if canImport(AppKit) + import AppKit // For NSImage extensions. +#elseif canImport(UIKit) + import UIKit // For UIImage extensions. +#endif + +@available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *) +final class GoogleGenerativeAITests: XCTestCase { + func codeSamples() async throws { + let config = GenerationConfig(temperature: 0.2, + topP: 0.1, + topK: 16, + candidateCount: 4, + maxOutputTokens: 256, + stopSequences: ["..."]) + let filters = [SafetySetting(harmCategory: .dangerousContent, threshold: .blockOnlyHigh)] + + // Permutations without optional arguments. + let _ = GenerativeModel(name: "gemini-1.0-pro", apiKey: "API_KEY") + let _ = GenerativeModel(name: "gemini-1.0-pro", apiKey: "API_KEY", safetySettings: filters) + let _ = GenerativeModel(name: "gemini-1.0-pro", apiKey: "API_KEY", generationConfig: config) + + // All arguments passed. + let genAI = GenerativeModel(name: "gemini-1.0-pro", + apiKey: "API_KEY", + generationConfig: config, // Optional + safetySettings: filters // Optional + ) + // Full Typed Usage + let pngData = Data() // .... + let contents = [ModelContent(role: "user", + parts: [ + .text("Is it a cat?"), + .png(pngData), + ])] + + do { + let response = try await genAI.generateContent(contents) + print(response.text ?? "Couldn't get text... check status") + } catch { + print("Error generating content: \(error)") + } + + // Content input combinations. + let _ = try await genAI.generateContent("Constant String") + let str = "String Variable" + let _ = try await genAI.generateContent(str) + let _ = try await genAI.generateContent([str]) + let _ = try await genAI.generateContent(str, "abc", "def") + #if canImport(UIKit) + _ = try await genAI.generateContent(UIImage()) + _ = try await genAI.generateContent([UIImage()]) + _ = try await genAI + .generateContent([str, UIImage(), ModelContent.Part.text(str)]) + _ = try await genAI.generateContent(str, UIImage(), "def", UIImage()) + _ = try await genAI.generateContent([str, UIImage(), "def", UIImage()]) + _ = try await genAI.generateContent([ModelContent("def", UIImage()), + ModelContent("def", UIImage())]) + #elseif canImport(AppKit) + _ = try await genAI.generateContent(NSImage()) + _ = try await genAI.generateContent([NSImage()]) + _ = try await genAI.generateContent(str, NSImage(), "def", NSImage()) + _ = try await genAI.generateContent([str, NSImage(), "def", NSImage()]) + #endif + + // ThrowingPartsRepresentable combinations. + let _ = ModelContent(parts: [.text(str)]) + let _ = ModelContent(role: "model", parts: [.text(str)]) + let _ = ModelContent(parts: "Constant String") + let _ = ModelContent(parts: str) + // Note: This requires the `try` for some reason. Casting to explicit [PartsRepresentable] also + // doesn't work. + let _ = try ModelContent(parts: [str]) + // Note: without `as [any ThrowingPartsRepresentable]` this will fail to compile with "Cannot + // convert value of type 'String' to expected element type + // 'Array.ArrayLiteralElement'. Not sure if there's a way we can get it to + // work. + let _ = try ModelContent(parts: [str, ModelContent.Part.data( + mimetype: "foo", + Data() + )] as [any ThrowingPartsRepresentable]) + #if canImport(UIKit) + _ = try ModelContent(role: "user", parts: UIImage()) + _ = try ModelContent(role: "user", parts: [UIImage()]) + // Note: without `as [any ThrowingPartsRepresentable]` this will fail to compile with "Cannot + // convert + // value of type `[Any]` to expected type `[any ThrowingPartsRepresentable]`. Not sure if + // there's a + // way we can get it to work. + _ = try ModelContent(parts: [str, UIImage()] as [any ThrowingPartsRepresentable]) + // Alternatively, you can explicitly declare the type in a variable and pass it in. + let representable2: [any ThrowingPartsRepresentable] = [str, UIImage()] + _ = try ModelContent(parts: representable2) + _ = try ModelContent(parts: [str, UIImage(), + ModelContent.Part.text(str)] as [any ThrowingPartsRepresentable]) + #elseif canImport(AppKit) + _ = try ModelContent(role: "user", parts: NSImage()) + _ = try ModelContent(role: "user", parts: [NSImage()]) + // Note: without `as [any ThrowingPartsRepresentable]` this will fail to compile with "Cannot + // convert + // value of type `[Any]` to expected type `[any ThrowingPartsRepresentable]`. Not sure if + // there's a + // way we can get it to work. + _ = try ModelContent(parts: [str, NSImage()] as [any ThrowingPartsRepresentable]) + // Alternatively, you can explicitly declare the type in a variable and pass it in. + let representable2: [any ThrowingPartsRepresentable] = [str, NSImage()] + _ = try ModelContent(parts: representable2) + _ = + try ModelContent(parts: [str, NSImage(), + ModelContent.Part.text(str)] as [any ThrowingPartsRepresentable]) + #endif + + // countTokens API + let _: CountTokensResponse = try await genAI.countTokens("What color is the Sky?") + #if canImport(UIKit) + let _: CountTokensResponse = try await genAI.countTokens("What color is the Sky?", + UIImage()) + let _: CountTokensResponse = try await genAI.countTokens([ + ModelContent("What color is the Sky?", UIImage()), + ModelContent(UIImage(), "What color is the Sky?", UIImage()), + ]) + #endif + + // Chat + _ = genAI.startChat() + _ = genAI.startChat(history: [ModelContent(parts: "abc")]) + } + + // Result builder alternative + + /* + let pngData = Data() // .... + let contents = [GenAIContent(role: "user", + parts: [ + .text("Is it a cat?"), + .png(pngData) + ])] + + // Turns into... + + let contents = GenAIContent { + Role("user") { + Text("Is this a cat?") + Image(png: pngData) + } + } + + GenAIContent { + ForEach(myInput) { input in + Role(input.role) { + input.contents + } + } + } + + // Thoughts: this looks great from a code demo, but since I assume most content will be + // user generated, the result builder may not be the best API. + */ +} diff --git a/FirebaseVertexAI/Tests/Unit/MockURLProtocol.swift b/FirebaseVertexAI/Tests/Unit/MockURLProtocol.swift new file mode 100644 index 00000000000..30528cdf35a --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/MockURLProtocol.swift @@ -0,0 +1,62 @@ +// Copyright 2023 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. + +import Foundation +import XCTest + +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, *) +class MockURLProtocol: URLProtocol { + static var requestHandler: ((URLRequest) throws -> ( + URLResponse, + AsyncLineSequence? + ))? + + override class func canInit(with request: URLRequest) -> Bool { return true } + + override class func canonicalRequest(for request: URLRequest) -> URLRequest { return request } + + override func startLoading() { + guard let requestHandler = MockURLProtocol.requestHandler else { + fatalError("`requestHandler` is nil.") + } + guard let client = client else { + fatalError("`client` is nil.") + } + + Task { + let (response, stream) = try requestHandler(self.request) + client.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed) + if let stream = stream { + do { + for try await line in stream { + guard let data = line.data(using: .utf8) else { + fatalError("Failed to convert \"\(line)\" to UTF8 data.") + } + client.urlProtocol(self, didLoad: data) + // Add a newline character since AsyncLineSequence strips them when reading line by + // line; + // without the following, the whole file is delivered as a single line. + client.urlProtocol(self, didLoad: "\n".data(using: .utf8)!) + } + } catch { + client.urlProtocol(self, didFailWithError: error) + XCTFail("Unexpected failure reading lines from stream: \(error.localizedDescription)") + } + } + client.urlProtocolDidFinishLoading(self) + } + } + + override func stopLoading() {} +} diff --git a/FirebaseVertexAI/Tests/Unit/PartsRepresentableTests.swift b/FirebaseVertexAI/Tests/Unit/PartsRepresentableTests.swift new file mode 100644 index 00000000000..da2f35f727f --- /dev/null +++ b/FirebaseVertexAI/Tests/Unit/PartsRepresentableTests.swift @@ -0,0 +1,138 @@ +// Copyright 2024 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. + +import CoreGraphics +import CoreImage +import GoogleGenerativeAI +import XCTest +#if canImport(UIKit) + import UIKit +#else + import AppKit +#endif + +@available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *) +final class PartsRepresentableTests: XCTestCase { + func testModelContentFromCGImageIsNotEmpty() throws { + // adapted from https://forums.swift.org/t/creating-a-cgimage-from-color-array/18634/2 + var srgbArray = [UInt32](repeating: 0xFFFF_FFFF, count: 8 * 8) + let image = srgbArray.withUnsafeMutableBytes { ptr -> CGImage in + let ctx = CGContext( + data: ptr.baseAddress, + width: 8, + height: 8, + bitsPerComponent: 8, + bytesPerRow: 4 * 8, + space: CGColorSpace(name: CGColorSpace.sRGB)!, + bitmapInfo: CGBitmapInfo.byteOrder32Little.rawValue + + CGImageAlphaInfo.premultipliedFirst.rawValue + )! + return ctx.makeImage()! + } + let modelContent = try image.tryPartsValue() + XCTAssert(modelContent.count > 0, "Expected non-empty model content for CGImage: \(image)") + } + + func testModelContentFromCIImageIsNotEmpty() throws { + let image = CIImage(color: CIColor.red) + .cropped(to: CGRect(origin: CGPointZero, size: CGSize(width: 16, height: 16))) + let modelContent = try image.tryPartsValue() + XCTAssert(modelContent.count > 0, "Expected non-empty model content for CGImage: \(image)") + } + + func testModelContentFromInvalidCIImageThrows() throws { + let image = CIImage.empty() + do { + let _ = try image.tryPartsValue() + } catch { + guard let imageError = (error as? ImageConversionError) else { + XCTFail("Got unexpected error type: \(error)") + return + } + switch imageError { + case let .couldNotConvertToJPEG(source): + // String(describing:) works around a type error. + XCTAssertEqual(String(describing: source), String(describing: image)) + return + case _: + XCTFail("Expected image conversion error, got \(imageError) instead") + return + } + } + XCTFail("Expected model content from invlaid image to error") + } + + #if canImport(UIKit) + func testModelContentFromInvalidUIImageThrows() throws { + let image = UIImage() + do { + _ = try image.tryPartsValue() + } catch { + guard let imageError = (error as? ImageConversionError) else { + XCTFail("Got unexpected error type: \(error)") + return + } + switch imageError { + case let .couldNotConvertToJPEG(source): + // String(describing:) works around a type error. + XCTAssertEqual(String(describing: source), String(describing: image)) + return + case _: + XCTFail("Expected image conversion error, got \(imageError) instead") + return + } + } + XCTFail("Expected model content from invlaid image to error") + } + + func testModelContentFromUIImageIsNotEmpty() throws { + let coreImage = CIImage(color: CIColor.red) + .cropped(to: CGRect(origin: CGPointZero, size: CGSize(width: 16, height: 16))) + let image = UIImage(ciImage: coreImage) + let modelContent = try image.tryPartsValue() + XCTAssert(modelContent.count > 0, "Expected non-empty model content for UIImage: \(image)") + } + #else + func testModelContentFromNSImageIsNotEmpty() throws { + let coreImage = CIImage(color: CIColor.red) + .cropped(to: CGRect(origin: CGPointZero, size: CGSize(width: 16, height: 16))) + let rep = NSCIImageRep(ciImage: coreImage) + let image = NSImage(size: rep.size) + image.addRepresentation(rep) + let modelContent = try image.tryPartsValue() + XCTAssert(modelContent.count > 0, "Expected non-empty model content for NSImage: \(image)") + } + + func testModelContentFromInvalidNSImageThrows() throws { + let image = NSImage() + do { + _ = try image.tryPartsValue() + } catch { + guard let imageError = (error as? ImageConversionError) else { + XCTFail("Got unexpected error type: \(error)") + return + } + switch imageError { + case .invalidUnderlyingImage: + // Pass + return + case _: + XCTFail("Expected image conversion error, got \(imageError) instead") + return + } + } + XCTFail("Expected model content from invlaid image to error") + } + #endif +}