GenerativeModelGoogleAITests.swift 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568
  1. // Copyright 2025 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. import FirebaseAppCheckInterop
  15. import FirebaseAuthInterop
  16. import FirebaseCore
  17. import XCTest
  18. @testable import FirebaseAI
  19. @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
  20. final class GenerativeModelGoogleAITests: XCTestCase {
  21. let testPrompt = "What sorts of questions can I ask you?"
  22. let safetyRatingsNegligible: [SafetyRating] = [
  23. .init(
  24. category: .sexuallyExplicit,
  25. probability: .negligible,
  26. probabilityScore: 0.0,
  27. severity: SafetyRating.HarmSeverity(rawValue: "HARM_SEVERITY_UNSPECIFIED"),
  28. severityScore: 0.0,
  29. blocked: false
  30. ),
  31. .init(
  32. category: .hateSpeech,
  33. probability: .negligible,
  34. probabilityScore: 0.0,
  35. severity: SafetyRating.HarmSeverity(rawValue: "HARM_SEVERITY_UNSPECIFIED"),
  36. severityScore: 0.0,
  37. blocked: false
  38. ),
  39. .init(
  40. category: .harassment,
  41. probability: .negligible,
  42. probabilityScore: 0.0,
  43. severity: SafetyRating.HarmSeverity(rawValue: "HARM_SEVERITY_UNSPECIFIED"),
  44. severityScore: 0.0,
  45. blocked: false
  46. ),
  47. .init(
  48. category: .dangerousContent,
  49. probability: .negligible,
  50. probabilityScore: 0.0,
  51. severity: SafetyRating.HarmSeverity(rawValue: "HARM_SEVERITY_UNSPECIFIED"),
  52. severityScore: 0.0,
  53. blocked: false
  54. ),
  55. ].sorted()
  56. let testModelName = "test-model"
  57. let testModelResourceName = "projects/test-project-id/models/test-model"
  58. let apiConfig = FirebaseAI.defaultVertexAIAPIConfig
  59. let googleAISubdirectory = "mock-responses/googleai"
  60. var urlSession: URLSession!
  61. var model: GenerativeModel!
  62. override func setUp() async throws {
  63. let configuration = URLSessionConfiguration.default
  64. configuration.protocolClasses = [MockURLProtocol.self]
  65. urlSession = try XCTUnwrap(URLSession(configuration: configuration))
  66. model = GenerativeModel(
  67. modelName: testModelName,
  68. modelResourceName: testModelResourceName,
  69. firebaseInfo: GenerativeModelTestUtil.testFirebaseInfo(),
  70. apiConfig: apiConfig,
  71. tools: nil,
  72. requestOptions: RequestOptions(),
  73. urlSession: urlSession
  74. )
  75. }
  76. override func tearDown() {
  77. MockURLProtocol.requestHandler = nil
  78. }
  79. // MARK: - Generate Content
  80. func testGenerateContent_success_basicReplyLong() async throws {
  81. MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
  82. forResource: "unary-success-basic-reply-long",
  83. withExtension: "json",
  84. subdirectory: googleAISubdirectory
  85. )
  86. let response = try await model.generateContent(testPrompt)
  87. XCTAssertEqual(response.candidates.count, 1)
  88. let candidate = try XCTUnwrap(response.candidates.first)
  89. let finishReason = try XCTUnwrap(candidate.finishReason)
  90. XCTAssertEqual(finishReason, .stop)
  91. XCTAssertEqual(candidate.safetyRatings.count, 4)
  92. XCTAssertEqual(candidate.content.parts.count, 1)
  93. let part = try XCTUnwrap(candidate.content.parts.first)
  94. let partText = try XCTUnwrap(part as? TextPart).text
  95. XCTAssertTrue(partText.hasPrefix("Making professional-quality"))
  96. XCTAssertEqual(response.text, partText)
  97. XCTAssertEqual(response.functionCalls, [])
  98. XCTAssertEqual(response.inlineDataParts, [])
  99. }
  100. func testGenerateContent_success_basicReplyShort() async throws {
  101. MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
  102. forResource: "unary-success-basic-reply-short",
  103. withExtension: "json",
  104. subdirectory: googleAISubdirectory
  105. )
  106. let response = try await model.generateContent(testPrompt)
  107. XCTAssertEqual(response.candidates.count, 1)
  108. let candidate = try XCTUnwrap(response.candidates.first)
  109. let finishReason = try XCTUnwrap(candidate.finishReason)
  110. XCTAssertEqual(finishReason, .stop)
  111. XCTAssertEqual(candidate.safetyRatings.sorted(), safetyRatingsNegligible)
  112. XCTAssertEqual(candidate.content.parts.count, 1)
  113. let part = try XCTUnwrap(candidate.content.parts.first)
  114. let textPart = try XCTUnwrap(part as? TextPart)
  115. XCTAssertTrue(textPart.text.hasPrefix("Google's headquarters"))
  116. XCTAssertEqual(response.text, textPart.text)
  117. XCTAssertEqual(response.functionCalls, [])
  118. XCTAssertEqual(response.inlineDataParts, [])
  119. }
  120. func testGenerateContent_success_citations() async throws {
  121. MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
  122. forResource: "unary-success-citations",
  123. withExtension: "json",
  124. subdirectory: googleAISubdirectory
  125. )
  126. let response = try await model.generateContent(testPrompt)
  127. XCTAssertEqual(response.candidates.count, 1)
  128. let candidate = try XCTUnwrap(response.candidates.first)
  129. XCTAssertEqual(candidate.content.parts.count, 1)
  130. let text = try XCTUnwrap(response.text)
  131. XCTAssertTrue(text.hasPrefix("Okay, let's break down quantum mechanics."))
  132. let citationMetadata = try XCTUnwrap(candidate.citationMetadata)
  133. XCTAssertEqual(citationMetadata.citations.count, 4)
  134. let citationSource1 = try XCTUnwrap(citationMetadata.citations[0])
  135. XCTAssertEqual(citationSource1.uri, "https://www.example.com/some-citation-1")
  136. XCTAssertEqual(citationSource1.startIndex, 548)
  137. XCTAssertEqual(citationSource1.endIndex, 690)
  138. XCTAssertNil(citationSource1.title)
  139. XCTAssertEqual(citationSource1.license, "mit")
  140. XCTAssertNil(citationSource1.publicationDate)
  141. let citationSource2 = try XCTUnwrap(citationMetadata.citations[1])
  142. XCTAssertEqual(citationSource2.uri, "https://www.example.com/some-citation-1")
  143. XCTAssertEqual(citationSource2.startIndex, 1240)
  144. XCTAssertEqual(citationSource2.endIndex, 1407)
  145. XCTAssertNil(citationSource2.title, "some-citation-2")
  146. XCTAssertNil(citationSource2.license)
  147. XCTAssertNil(citationSource2.publicationDate)
  148. let citationSource3 = try XCTUnwrap(citationMetadata.citations[2])
  149. XCTAssertEqual(citationSource3.startIndex, 1942)
  150. XCTAssertEqual(citationSource3.endIndex, 2149)
  151. XCTAssertNil(citationSource3.uri)
  152. XCTAssertNil(citationSource3.license)
  153. XCTAssertNil(citationSource3.title)
  154. XCTAssertNil(citationSource3.publicationDate)
  155. let citationSource4 = try XCTUnwrap(citationMetadata.citations[3])
  156. XCTAssertEqual(citationSource4.startIndex, 2036)
  157. XCTAssertEqual(citationSource4.endIndex, 2175)
  158. XCTAssertNil(citationSource4.uri)
  159. XCTAssertNil(citationSource4.license)
  160. XCTAssertNil(citationSource4.title)
  161. XCTAssertNil(citationSource4.publicationDate)
  162. }
  163. func testGenerateContent_usageMetadata() async throws {
  164. MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
  165. forResource: "unary-success-basic-reply-short",
  166. withExtension: "json",
  167. subdirectory: googleAISubdirectory
  168. )
  169. let response = try await model.generateContent(testPrompt)
  170. let usageMetadata = try XCTUnwrap(response.usageMetadata)
  171. XCTAssertEqual(usageMetadata.promptTokenCount, 7)
  172. XCTAssertEqual(usageMetadata.promptTokensDetails.count, 1)
  173. XCTAssertEqual(usageMetadata.promptTokensDetails[0].modality, .text)
  174. XCTAssertEqual(usageMetadata.promptTokensDetails[0].tokenCount, 7)
  175. XCTAssertEqual(usageMetadata.candidatesTokenCount, 22)
  176. XCTAssertEqual(usageMetadata.candidatesTokensDetails.count, 1)
  177. XCTAssertEqual(usageMetadata.candidatesTokensDetails[0].modality, .text)
  178. XCTAssertEqual(usageMetadata.candidatesTokensDetails[0].tokenCount, 22)
  179. }
  180. func testGenerateContent_groundingMetadata() async throws {
  181. MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
  182. forResource: "unary-success-google-search-grounding",
  183. withExtension: "json",
  184. subdirectory: googleAISubdirectory
  185. )
  186. let response = try await model.generateContent(testPrompt)
  187. XCTAssertEqual(response.candidates.count, 1)
  188. let candidate = try XCTUnwrap(response.candidates.first)
  189. let groundingMetadata = try XCTUnwrap(candidate.groundingMetadata)
  190. XCTAssertEqual(groundingMetadata.webSearchQueries, ["current weather in London"])
  191. let searchEntryPoint = try XCTUnwrap(groundingMetadata.searchEntryPoint)
  192. XCTAssertFalse(searchEntryPoint.renderedContent.isEmpty)
  193. XCTAssertEqual(groundingMetadata.groundingChunks.count, 2)
  194. let firstChunk = try XCTUnwrap(groundingMetadata.groundingChunks.first?.web)
  195. XCTAssertEqual(firstChunk.title, "accuweather.com")
  196. XCTAssertNotNil(firstChunk.uri)
  197. XCTAssertNil(firstChunk.domain) // Domain is not supported by Google AI backend
  198. XCTAssertEqual(groundingMetadata.groundingSupports.count, 3)
  199. let firstSupport = try XCTUnwrap(groundingMetadata.groundingSupports.first)
  200. let segment = try XCTUnwrap(firstSupport.segment)
  201. XCTAssertEqual(segment.text, "The current weather in London, United Kingdom is cloudy.")
  202. XCTAssertEqual(segment.startIndex, 0)
  203. XCTAssertEqual(segment.partIndex, 0)
  204. XCTAssertEqual(segment.endIndex, 56)
  205. XCTAssertEqual(firstSupport.groundingChunkIndices, [0])
  206. }
  207. // This test case can be deleted once https://b.corp.google.com/issues/422779395 (internal) is
  208. // fixed.
  209. func testGenerateContent_groundingMetadata_emptyGroundingChunks() async throws {
  210. MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
  211. forResource: "unary-success-google-search-grounding-empty-grounding-chunks",
  212. withExtension: "json",
  213. subdirectory: googleAISubdirectory
  214. )
  215. let response = try await model.generateContent(testPrompt)
  216. XCTAssertEqual(response.candidates.count, 1)
  217. let candidate = try XCTUnwrap(response.candidates.first)
  218. let groundingMetadata = try XCTUnwrap(candidate.groundingMetadata)
  219. XCTAssertNotNil(groundingMetadata.searchEntryPoint)
  220. XCTAssertEqual(groundingMetadata.webSearchQueries, ["current weather London"])
  221. // Chunks exist, but contain no web information.
  222. XCTAssertEqual(groundingMetadata.groundingChunks.count, 2)
  223. XCTAssertNil(groundingMetadata.groundingChunks[0].web)
  224. XCTAssertNil(groundingMetadata.groundingChunks[1].web)
  225. XCTAssertEqual(groundingMetadata.groundingSupports.count, 1)
  226. let support = try XCTUnwrap(groundingMetadata.groundingSupports.first)
  227. XCTAssertEqual(support.groundingChunkIndices, [0])
  228. XCTAssertEqual(
  229. support.segment.text,
  230. "There is a 0% chance of rain and the humidity is around 41%."
  231. )
  232. }
  233. func testGenerateContent_success_thinking_thoughtSummary() async throws {
  234. MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
  235. forResource: "unary-success-thinking-reply-thought-summary",
  236. withExtension: "json",
  237. subdirectory: googleAISubdirectory
  238. )
  239. let response = try await model.generateContent(testPrompt)
  240. XCTAssertEqual(response.candidates.count, 1)
  241. let candidate = try XCTUnwrap(response.candidates.first)
  242. XCTAssertEqual(candidate.content.parts.count, 2)
  243. let thoughtPart = try XCTUnwrap(candidate.content.parts.first as? TextPart)
  244. XCTAssertTrue(thoughtPart.isThought)
  245. XCTAssertTrue(thoughtPart.text.hasPrefix("**Thinking About Google's Headquarters**"))
  246. XCTAssertEqual(thoughtPart.text, response.thoughtSummary)
  247. let textPart = try XCTUnwrap(candidate.content.parts.last as? TextPart)
  248. XCTAssertFalse(textPart.isThought)
  249. XCTAssertEqual(textPart.text, "Mountain View")
  250. XCTAssertEqual(textPart.text, response.text)
  251. }
  252. func testGenerateContent_success_thinking_functionCall_thoughtSummaryAndSignature() async throws {
  253. MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
  254. forResource: "unary-success-thinking-function-call-thought-summary-signature",
  255. withExtension: "json",
  256. subdirectory: googleAISubdirectory
  257. )
  258. let response = try await model.generateContent(testPrompt)
  259. XCTAssertEqual(response.candidates.count, 1)
  260. let candidate = try XCTUnwrap(response.candidates.first)
  261. XCTAssertEqual(candidate.finishReason, .stop)
  262. XCTAssertEqual(candidate.content.parts.count, 2)
  263. let thoughtPart = try XCTUnwrap(candidate.content.parts.first as? TextPart)
  264. XCTAssertTrue(thoughtPart.isThought)
  265. XCTAssertTrue(thoughtPart.text.hasPrefix("**Thinking Through the New Year's Eve Calculation**"))
  266. let functionCallPart = try XCTUnwrap(candidate.content.parts.last as? FunctionCallPart)
  267. XCTAssertFalse(functionCallPart.isThought)
  268. XCTAssertEqual(functionCallPart.name, "now")
  269. XCTAssertTrue(functionCallPart.args.isEmpty)
  270. let thoughtSignature = try XCTUnwrap(functionCallPart.thoughtSignature)
  271. XCTAssertTrue(thoughtSignature.hasPrefix("CtQOAVSoXO74PmYr9AFu"))
  272. }
  273. func testGenerateContent_failure_invalidAPIKey() async throws {
  274. let expectedStatusCode = 400
  275. MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
  276. forResource: "unary-failure-api-key",
  277. withExtension: "json",
  278. subdirectory: googleAISubdirectory,
  279. statusCode: expectedStatusCode
  280. )
  281. do {
  282. _ = try await model.generateContent(testPrompt)
  283. XCTFail("Should throw GenerateContentError.internalError; no error thrown.")
  284. } catch let GenerateContentError.internalError(error as BackendError) {
  285. XCTAssertEqual(error.httpResponseCode, 400)
  286. XCTAssertEqual(error.status, .invalidArgument)
  287. XCTAssertEqual(error.message, "API key not valid. Please pass a valid API key.")
  288. XCTAssertTrue(error.localizedDescription.contains(error.message))
  289. XCTAssertTrue(error.localizedDescription.contains(error.status.rawValue))
  290. XCTAssertTrue(error.localizedDescription.contains("\(error.httpResponseCode)"))
  291. let nsError = error as NSError
  292. XCTAssertEqual(nsError.domain, "\(Constants.baseErrorDomain).\(BackendError.self)")
  293. XCTAssertEqual(nsError.code, error.httpResponseCode)
  294. return
  295. } catch {
  296. XCTFail("Should throw GenerateContentError.internalError(RPCError); error thrown: \(error)")
  297. }
  298. }
  299. func testGenerateContent_failure_finishReasonSafety() async throws {
  300. MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
  301. forResource: "unary-failure-finish-reason-safety",
  302. withExtension: "json",
  303. subdirectory: googleAISubdirectory
  304. )
  305. do {
  306. _ = try await model.generateContent(testPrompt)
  307. XCTFail("Should throw")
  308. } catch let GenerateContentError.responseStoppedEarly(reason, response) {
  309. XCTAssertEqual(reason, .safety)
  310. XCTAssertEqual(response.text, "Safety error incoming in 5, 4, 3, 2...")
  311. } catch {
  312. XCTFail("Should throw a responseStoppedEarly")
  313. }
  314. }
  315. func testGenerateContent_failure_unknownModel() async throws {
  316. let expectedStatusCode = 404
  317. MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
  318. forResource: "unary-failure-unknown-model",
  319. withExtension: "json",
  320. subdirectory: googleAISubdirectory,
  321. statusCode: 404
  322. )
  323. do {
  324. _ = try await model.generateContent(testPrompt)
  325. XCTFail("Should throw GenerateContentError.internalError; no error thrown.")
  326. } catch let GenerateContentError.internalError(underlying: rpcError as BackendError) {
  327. XCTAssertEqual(rpcError.status, .notFound)
  328. XCTAssertEqual(rpcError.httpResponseCode, expectedStatusCode)
  329. XCTAssertTrue(rpcError.message.hasPrefix("models/gemini-5.0-flash is not found"))
  330. } catch {
  331. XCTFail("Should throw GenerateContentError.internalError; error thrown: \(error)")
  332. }
  333. }
  334. // MARK: - Generate Content (Streaming)
  335. func testGenerateContentStream_successBasicReplyLong() async throws {
  336. MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
  337. forResource: "streaming-success-basic-reply-long",
  338. withExtension: "txt",
  339. subdirectory: googleAISubdirectory
  340. )
  341. var responses = 0
  342. let stream = try model.generateContentStream("Hi")
  343. for try await content in stream {
  344. XCTAssertNotNil(content.text)
  345. responses += 1
  346. }
  347. XCTAssertEqual(responses, 36)
  348. }
  349. func testGenerateContentStream_successBasicReplyShort() async throws {
  350. MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
  351. forResource: "streaming-success-basic-reply-short",
  352. withExtension: "txt",
  353. subdirectory: googleAISubdirectory
  354. )
  355. var responses = 0
  356. let stream = try model.generateContentStream("Hi")
  357. for try await content in stream {
  358. XCTAssertNotNil(content.text)
  359. responses += 1
  360. }
  361. XCTAssertEqual(responses, 3)
  362. }
  363. func testGenerateContentStream_successWithCitations() async throws {
  364. MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
  365. forResource: "streaming-success-citations",
  366. withExtension: "txt",
  367. subdirectory: googleAISubdirectory
  368. )
  369. let stream = try model.generateContentStream("Hi")
  370. var citations = [Citation]()
  371. var responses = [GenerateContentResponse]()
  372. for try await content in stream {
  373. responses.append(content)
  374. XCTAssertNotNil(content.text)
  375. let candidate = try XCTUnwrap(content.candidates.first)
  376. if let sources = candidate.citationMetadata?.citations {
  377. citations.append(contentsOf: sources)
  378. }
  379. }
  380. let lastCandidate = try XCTUnwrap(responses.last?.candidates.first)
  381. XCTAssertEqual(lastCandidate.finishReason, .stop)
  382. XCTAssertEqual(citations.count, 1)
  383. let citation = try XCTUnwrap(citations.first)
  384. XCTAssertEqual(citation.startIndex, 111)
  385. XCTAssertEqual(citation.endIndex, 236)
  386. let citationURI = try XCTUnwrap(citation.uri)
  387. XCTAssertTrue(citationURI.starts(with: "https://www."))
  388. XCTAssertNil(citation.license)
  389. XCTAssertNil(citation.title)
  390. XCTAssertNil(citation.publicationDate)
  391. }
  392. func testGenerateContentStream_successWithThoughtSummary() async throws {
  393. MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
  394. forResource: "streaming-success-thinking-reply-thought-summary",
  395. withExtension: "txt",
  396. subdirectory: googleAISubdirectory
  397. )
  398. var thoughtSummary = ""
  399. var text = ""
  400. let stream = try model.generateContentStream("Hi")
  401. for try await response in stream {
  402. let candidate = try XCTUnwrap(response.candidates.first)
  403. XCTAssertEqual(candidate.content.parts.count, 1)
  404. let textPart = try XCTUnwrap(candidate.content.parts.first as? TextPart)
  405. if textPart.isThought {
  406. let newThought = try XCTUnwrap(response.thoughtSummary)
  407. XCTAssertEqual(textPart.text, newThought)
  408. thoughtSummary.append(newThought)
  409. } else {
  410. let newText = try XCTUnwrap(response.text)
  411. XCTAssertEqual(textPart.text, newText)
  412. text.append(newText)
  413. }
  414. }
  415. XCTAssertTrue(thoughtSummary.hasPrefix("**Exploring Sky Color**"))
  416. XCTAssertTrue(text.hasPrefix("The sky is blue because"))
  417. }
  418. func testGenerateContentStream_success_thinking_functionCall_thoughtSummary_signature() async throws {
  419. MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
  420. forResource: "streaming-success-thinking-function-call-thought-summary-signature",
  421. withExtension: "txt",
  422. subdirectory: googleAISubdirectory
  423. )
  424. var thoughtSummary = ""
  425. var functionCalls: [FunctionCallPart] = []
  426. let stream = try model.generateContentStream("Hi")
  427. for try await response in stream {
  428. let candidate = try XCTUnwrap(response.candidates.first)
  429. XCTAssertEqual(candidate.content.parts.count, 1)
  430. let part = try XCTUnwrap(candidate.content.parts.first)
  431. if part.isThought {
  432. let textPart = try XCTUnwrap(part as? TextPart)
  433. let newThought = try XCTUnwrap(response.thoughtSummary)
  434. XCTAssertEqual(textPart.text, newThought)
  435. thoughtSummary.append(newThought)
  436. } else {
  437. let functionCallPart = try XCTUnwrap(part as? FunctionCallPart)
  438. XCTAssertEqual(response.functionCalls.count, 1)
  439. let newFunctionCall = try XCTUnwrap(response.functionCalls.first)
  440. XCTAssertEqual(functionCallPart, newFunctionCall)
  441. functionCalls.append(newFunctionCall)
  442. }
  443. }
  444. XCTAssertTrue(thoughtSummary.hasPrefix("**Calculating the Days**"))
  445. XCTAssertEqual(functionCalls.count, 1)
  446. let functionCall = try XCTUnwrap(functionCalls.first)
  447. XCTAssertEqual(functionCall.name, "now")
  448. XCTAssertTrue(functionCall.args.isEmpty)
  449. let thoughtSignature = try XCTUnwrap(functionCall.thoughtSignature)
  450. XCTAssertTrue(thoughtSignature.hasPrefix("CiIBVKhc7vB+vaaq6rA"))
  451. }
  452. func testGenerateContentStream_failureInvalidAPIKey() async throws {
  453. MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
  454. forResource: "unary-failure-api-key",
  455. withExtension: "json",
  456. subdirectory: googleAISubdirectory
  457. )
  458. do {
  459. let stream = try model.generateContentStream("Hi")
  460. for try await _ in stream {
  461. XCTFail("No content is there, this shouldn't happen.")
  462. }
  463. } catch let GenerateContentError.internalError(error as BackendError) {
  464. XCTAssertEqual(error.httpResponseCode, 400)
  465. XCTAssertEqual(error.status, .invalidArgument)
  466. XCTAssertEqual(error.message, "API key not valid. Please pass a valid API key.")
  467. XCTAssertTrue(error.localizedDescription.contains(error.message))
  468. XCTAssertTrue(error.localizedDescription.contains(error.status.rawValue))
  469. XCTAssertTrue(error.localizedDescription.contains("\(error.httpResponseCode)"))
  470. let nsError = error as NSError
  471. XCTAssertEqual(nsError.domain, "\(Constants.baseErrorDomain).\(BackendError.self)")
  472. XCTAssertEqual(nsError.code, error.httpResponseCode)
  473. return
  474. }
  475. XCTFail("Should have caught an error.")
  476. }
  477. func testGenerateContentStream_failureFinishRecitation() async throws {
  478. MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
  479. forResource: "streaming-failure-recitation-no-content",
  480. withExtension: "txt",
  481. subdirectory: googleAISubdirectory
  482. )
  483. var responses = [GenerateContentResponse]()
  484. do {
  485. let stream = try model.generateContentStream("Hi")
  486. for try await response in stream {
  487. responses.append(response)
  488. }
  489. XCTFail("Expected a GenerateContentError.responseStoppedEarly error, but got no error.")
  490. } catch let GenerateContentError.responseStoppedEarly(reason, response) {
  491. XCTAssertEqual(reason, .recitation)
  492. let candidate = try XCTUnwrap(response.candidates.first)
  493. XCTAssertEqual(candidate.finishReason, reason)
  494. } catch {
  495. XCTFail("Expected a GenerateContentError.responseStoppedEarly error, but got error: \(error)")
  496. }
  497. XCTAssertEqual(responses.count, 8)
  498. let firstResponse = try XCTUnwrap(responses.first)
  499. XCTAssertEqual(firstResponse.text, "text1")
  500. let lastResponse = try XCTUnwrap(responses.last)
  501. XCTAssertEqual(lastResponse.text, "text8")
  502. }
  503. }