Skip to content

Commit f0d31ea

Browse files
andrewheardMorgan Chenpaulb777peterfrieseryanwilson
committed
Fork generative-ai-swift tests for Vertex AI (#12584)
Co-authored-by: Morgan Chen <[email protected]> Co-authored-by: Paul Beusterien <[email protected]> Co-authored-by: Peter Friese <[email protected]> Co-authored-by: Ryan Wilson <[email protected]>
1 parent 8af3734 commit f0d31ea

File tree

39 files changed

+2134
-0
lines changed

39 files changed

+2134
-0
lines changed

‎FirebaseVertexAI/Tests/Unit/ChatTests.swift

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
// Copyright 2023 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
import Foundation
16+
@testable import GoogleGenerativeAI
17+
import XCTest
18+
19+
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, *)
20+
final class ChatTests: XCTestCase {
21+
var urlSession: URLSession!
22+
23+
override func setUp() {
24+
let configuration = URLSessionConfiguration.default
25+
configuration.protocolClasses = [MockURLProtocol.self]
26+
urlSession = URLSession(configuration: configuration)
27+
}
28+
29+
override func tearDown() {
30+
MockURLProtocol.requestHandler = nil
31+
}
32+
33+
func testMergingText() async throws {
34+
let fileURL = try XCTUnwrap(Bundle.module.url(
35+
forResource: "streaming-success-basic-reply-long",
36+
withExtension: "txt"
37+
))
38+
39+
MockURLProtocol.requestHandler = { request in
40+
let response = HTTPURLResponse(
41+
url: request.url!,
42+
statusCode: 200,
43+
httpVersion: nil,
44+
headerFields: nil
45+
)!
46+
return (response, fileURL.lines)
47+
}
48+
49+
let model = GenerativeModel(name: "my-model", apiKey: "API_KEY", urlSession: urlSession)
50+
let chat = Chat(model: model, history: [])
51+
let input = "Test input"
52+
let stream = chat.sendMessageStream(input)
53+
54+
// Ensure the values are parsed correctly
55+
for try await value in stream {
56+
XCTAssertNotNil(value.text)
57+
}
58+
59+
XCTAssertEqual(chat.history.count, 2)
60+
XCTAssertEqual(chat.history[0].parts[0].text, input)
61+
62+
let finalText = "1 2 3 4 5 6 7 8 9 10"
63+
let assembledExpectation = ModelContent(role: "model", parts: finalText)
64+
XCTAssertEqual(chat.history[0].parts[0].text, input)
65+
XCTAssertEqual(chat.history[1], assembledExpectation)
66+
}
67+
}

‎FirebaseVertexAI/Tests/Unit/CountTokenResponses/failure-model-not-found.json

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
{
2+
"error": {
3+
"code": 404,
4+
"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.",
5+
"status": "NOT_FOUND",
6+
"details": [
7+
{
8+
"@type": "type.googleapis.com/google.rpc.DebugInfo",
9+
"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.\" }"
10+
}
11+
]
12+
}
13+
}

‎FirebaseVertexAI/Tests/Unit/CountTokenResponses/success-total-tokens.json

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
{
2+
"totalTokens": 6
3+
}

‎FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-empty-content.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
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"}]}}

‎FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-error-mid-stream.txt

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
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"}]}}
2+
3+
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"}]}]}
4+
5+
{
6+
"error": {
7+
"code": 499,
8+
"message": "The operation was cancelled.",
9+
"status": "CANCELLED",
10+
"details": [
11+
{
12+
"@type": "type.googleapis.com/google.rpc.DebugInfo",
13+
"detail": "[ORIGINAL ERROR] generic::cancelled: "
14+
}
15+
]
16+
}
17+
}

‎FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-finish-reason-safety.txt

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
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"}]}}
2+

‎FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-invalid-json.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
data: {"this": [{"is": {"not": [{"a": "valid"}]}, "response": {}}]}

‎FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-malformed-content.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
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"}]}}

‎FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-prompt-blocked-safety.txt

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
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"}]}}
2+

‎FirebaseVertexAI/Tests/Unit/GenerateContentResponses/streaming-failure-recitation-no-content.txt

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
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"}]}}
2+
3+
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": ""}]}}]}
4+
5+
data: {"candidates": [{"finishReason": "RECITATION","index": 0}]}
6+

0 commit comments

Comments
 (0)