VertexAI: migrate test cases to use updated set of golden files

firebase · Jun 18, 2024 · 9f726e6 · 9f726e6
1 parent a90255a
commit 9f726e6
Show file tree

Hide file tree

Showing 13 changed files with 149 additions and 138 deletions.
diff --git a/.changeset/tame-lions-raise.md b/.changeset/tame-lions-raise.md
@@ -0,0 +1,2 @@
+---
+---
diff --git a/packages/vertexai/src/methods/generate-content.test.ts b/packages/vertexai/src/methods/generate-content.test.ts
@@ -70,7 +70,7 @@ describe('generateContent()', () => {
       'model',
       fakeRequestParams
     );
-    expect(result.response.text()).to.include('Helena');
+    expect(result.response.text()).to.include('Mountain View, California');
     expect(makeRequestStub).to.be.calledWith(
       'model',
       Task.GENERATE_CONTENT,
@@ -112,10 +112,12 @@ describe('generateContent()', () => {
       'model',
       fakeRequestParams
     );
-    expect(result.response.text()).to.include('Quantum mechanics is');
+    expect(result.response.text()).to.include(
+      'Some information cited from an external source'
+    );
     expect(
       result.response.candidates?.[0].citationMetadata?.citations.length
-    ).to.equal(1);
+    ).to.equal(3);
     expect(makeRequestStub).to.be.calledWith(
       'model',
       Task.GENERATE_CONTENT,
@@ -186,7 +188,7 @@ describe('generateContent()', () => {
     );
   });
   it('unknown enum - should ignore', async () => {
-    const mockResponse = getMockResponse('unary-unknown-enum.json');
+    const mockResponse = getMockResponse('unary-success-unknown-enum.json');
     const makeRequestStub = stub(request, 'makeRequest').resolves(
       mockResponse as Response
     );

diff --git a/packages/vertexai/src/requests/stream-reader.test.ts b/packages/vertexai/src/requests/stream-reader.test.ts
@@ -173,7 +173,9 @@ describe('processStream', () => {
     }
   });
   it('unknown enum - should ignore', async () => {
-    const fakeResponse = getMockResponseStreaming('streaming-unknown-enum.txt');
+    const fakeResponse = getMockResponseStreaming(
+      'streaming-success-unknown-enum.txt'
+    );
     const result = processStream(fakeResponse as Response);
     const aggregatedResponse = await result.response;
     expect(aggregatedResponse.text()).to.include('Cats');

diff --git a/packages/vertexai/test-utils/mock-responses/streaming-failure-finish-reason-safety.txt b/packages/vertexai/test-utils/mock-responses/streaming-failure-finish-reason-safety.txt
@@ -1,2 +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"}]}}
+data: {"candidates": [{"content": {"parts": [{"text": "<redacted>"}],"role": "model"},"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/packages/vertexai/test-utils/mock-responses/streaming-success-basic-reply-long.txt b/packages/vertexai/test-utils/mock-responses/streaming-success-basic-reply-long.txt
@@ -1,8 +1,8 @@
-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": "**Cats:**\n\n- **Physical Characteristics:**\n  - Size: Cats come"}],"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": " 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": " 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"}],"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": " 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": " 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"}],"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": ", 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"}]}]}
 

diff --git a/packages/vertexai/test-utils/mock-responses/streaming-success-basic-reply-short.txt b/packages/vertexai/test-utils/mock-responses/streaming-success-basic-reply-short.txt
@@ -1,2 +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"}]}}
+data: {"candidates": [{"content": {"role": "model","parts": [{"text": "Cheyenne"}]},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE","probabilityScore": 0.02854415,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.052424565},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE","probabilityScore": 0.24926445,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.0996453},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE","probabilityScore": 0.087096825,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.043123372},{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE","probabilityScore": 0.14402841,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.086169556}]}],"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"}]},"usageMetadata": {"promptTokenCount": 6,"candidatesTokenCount": 4,"totalTokenCount": 10}}
 
diff --git a/packages/vertexai/test-utils/mock-responses/streaming-success-citations.txt b/packages/vertexai/test-utils/mock-responses/streaming-success-citations.txt
@@ -1,4 +1,4 @@
-data: {"candidates": [{"content": {"parts": [{"text": "1. **Definition:**\nQuantum mechanics is a fundamental theory in physics that provides"}],"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": "1. **Definition:**\nQuantum mechanics is a fundamental theory in physics that provides"}],"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": " the foundation for understanding the physical properties of nature at the scale of atoms and subatomic particles. It is based on the idea that energy, momentum, angular momentum"}],"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"}]}]}