fix ts errors

Author: hsubox76

packages/ai/src/methods/count-tokens.ts

@@ -28,6 +28,7 @@ import { ApiSettings } from '../types/internal';
 import * as GoogleAIMapper from '../googleai-mappers';
 import { BackendType } from '../public-types';
 import { ChromeAdapter } from '../types/chrome-adapter';
+import { ChromeAdapterImpl } from './chrome-adapter';
 
 export async function countTokensOnCloud(
   apiSettings: ApiSettings,
@@ -60,7 +61,7 @@ export async function countTokens(
   chromeAdapter?: ChromeAdapter,
   requestOptions?: RequestOptions
 ): Promise<CountTokensResponse> {
-  if (chromeAdapter?.mode === InferenceMode.ONLY_ON_DEVICE) {
+  if ((chromeAdapter as ChromeAdapterImpl)?.mode === InferenceMode.ONLY_ON_DEVICE) {
     throw new AIError(
       AIErrorCode.UNSUPPORTED,
       'countTokens() is not supported for on-device models.'

packages/ai/src/methods/helpers.test.ts

@@ -19,25 +19,26 @@ import { use, expect } from 'chai';
 import { SinonStub, SinonStubbedInstance, restore, stub } from 'sinon';
 import { callCloudOrDevice } from './helpers';
 import {
-  ChromeAdapter,
   GenerateContentRequest,
   InferenceMode,
   AIErrorCode
 } from '../types';
 import { AIError } from '../errors';
 import sinonChai from 'sinon-chai';
 import chaiAsPromised from 'chai-as-promised';
+import { ChromeAdapterImpl } from './chrome-adapter';
 
 use(sinonChai);
 use(chaiAsPromised);
 
 describe('callCloudOrDevice', () => {
-  let chromeAdapter: SinonStubbedInstance<ChromeAdapter>;
+  let chromeAdapter: SinonStubbedInstance<ChromeAdapterImpl>;
   let onDeviceCall: SinonStub;
   let inCloudCall: SinonStub;
   let request: GenerateContentRequest;
 
   beforeEach(() => {
+    // @ts-ignore
     chromeAdapter = {
       mode: InferenceMode.PREFER_ON_DEVICE,
       isAvailable: stub(),

packages/ai/src/methods/helpers.ts

@@ -16,7 +16,12 @@
  */
 
 import { AIError } from '../errors';
-import { GenerateContentRequest, InferenceMode, AIErrorCode } from '../types';
+import {
+  GenerateContentRequest,
+  InferenceMode,
+  AIErrorCode,
+  ChromeAdapter
+} from '../types';
 import { ChromeAdapterImpl } from './chrome-adapter';
 
 /**
@@ -31,14 +36,14 @@ import { ChromeAdapterImpl } from './chrome-adapter';
  */
 export async function callCloudOrDevice<Response>(
   request: GenerateContentRequest,
-  chromeAdapter: ChromeAdapterImpl | undefined,
+  chromeAdapter: ChromeAdapter | undefined,
   onDeviceCall: () => Promise<Response>,
   inCloudCall: () => Promise<Response>
 ): Promise<Response> {
   if (!chromeAdapter) {
     return inCloudCall();
   }
-  switch (chromeAdapter.mode) {
+  switch ((chromeAdapter as ChromeAdapterImpl).mode) {
     case InferenceMode.ONLY_ON_DEVICE:
       if (await chromeAdapter.isAvailable(request)) {
         return onDeviceCall();

packages/ai/src/models/generative-model.test.ts

@@ -30,7 +30,6 @@ import {
 } from '../../test-utils/mock-response';
 import sinonChai from 'sinon-chai';
 import { VertexAIBackend } from '../backend';
-import { ChromeAdapter } from '../types/chrome-adapter';
 import { ChromeAdapterImpl } from '../methods/chrome-adapter';
 import { AIError } from '../errors';
 import chaiAsPromised from 'chai-as-promised';
@@ -421,7 +420,7 @@ describe('GenerativeModel', () => {
 
 describe('GenerativeModel dispatch logic', () => {
   let makeRequestStub: SinonStub;
-  let mockChromeAdapter: ChromeAdapter;
+  let mockChromeAdapter: ChromeAdapterImpl;
 
   function stubMakeRequest(stream?: boolean): void {
     if (stream) {
@@ -442,6 +441,7 @@ describe('GenerativeModel dispatch logic', () => {
   }
 
   beforeEach(() => {
+    // @ts-ignore
     mockChromeAdapter = {
       isAvailable: stub(),
       generateContent: stub().resolves(new Response(JSON.stringify({}))),