fix(ai): Fix generateContentStream returning wrong inferenceSource. (#9381)

Author: hsubox76

.changeset/gorgeous-rice-carry.md

@@ -0,0 +1,5 @@
+---
+'@firebase/ai': patch
+---
+
+Fix `generateContentStream` returning wrong `inferenceSource`.

packages/ai/src/methods/generate-content.test.ts

@@ -26,6 +26,7 @@ import {
 import * as request from '../requests/request';
 import {
   generateContent,
+  generateContentStream,
   templateGenerateContent,
   templateGenerateContentStream
 } from './generate-content';
@@ -35,6 +36,7 @@ import {
   HarmBlockMethod,
   HarmBlockThreshold,
   HarmCategory,
+  InferenceSource,
   Language,
   Outcome
 } from '../types';
@@ -548,8 +550,7 @@ describe('generateContent()', () => {
       );
     });
   });
-  // TODO: define a similar test for generateContentStream
-  it('on-device', async () => {
+  it('generateContent on-device', async () => {
     const chromeAdapter = fakeChromeAdapter;
     const isAvailableStub = stub(chromeAdapter, 'isAvailable').resolves(true);
     const mockResponse = getMockResponse(
@@ -566,9 +567,35 @@ describe('generateContent()', () => {
       chromeAdapter
     );
     expect(result.response.text()).to.include('Mountain View, California');
+    expect(result.response.inferenceSource).to.equal(InferenceSource.ON_DEVICE);
     expect(isAvailableStub).to.be.called;
     expect(generateContentStub).to.be.calledWith(fakeRequestParams);
   });
+  it('generateContentStream on-device', async () => {
+    const chromeAdapter = fakeChromeAdapter;
+    const isAvailableStub = stub(chromeAdapter, 'isAvailable').resolves(true);
+    const mockResponse = getMockResponseStreaming(
+      'vertexAI',
+      'streaming-success-basic-reply-short.txt'
+    );
+    const generateContentStreamStub = stub(
+      chromeAdapter,
+      'generateContentStream'
+    ).resolves(mockResponse as Response);
+    const result = await generateContentStream(
+      fakeApiSettings,
+      'model',
+      fakeRequestParams,
+      chromeAdapter
+    );
+    const aggregatedResponse = await result.response;
+    expect(aggregatedResponse.text()).to.include('Cheyenne');
+    expect(aggregatedResponse.inferenceSource).to.equal(
+      InferenceSource.ON_DEVICE
+    );
+    expect(isAvailableStub).to.be.called;
+    expect(generateContentStreamStub).to.be.calledWith(fakeRequestParams);
+  });
 });
 
 describe('templateGenerateContent', () => {

packages/ai/src/methods/generate-content.ts

@@ -70,7 +70,11 @@ export async function generateContentStream(
     () =>
       generateContentStreamOnCloud(apiSettings, model, params, requestOptions)
   );
-  return processStream(callResult.response, apiSettings); // TODO: Map streaming responses
+  return processStream(
+    callResult.response,
+    apiSettings,
+    callResult.inferenceSource
+  );
 }
 
 async function generateContentOnCloud(