feat: Migrate Chrome AI integration to a separate module

src/models/ChatChromeAi.ts

@@ -10,15 +10,9 @@ import {
 import { BaseMessage, AIMessageChunk } from "@langchain/core/messages"
 import { ChatGenerationChunk } from "@langchain/core/outputs"
 import { IterableReadableStream } from "@langchain/core/utils/stream"
+import { AITextSession, checkChromeAIAvailability, createAITextSession } from "./utils/chrome"
 
 
-export interface AITextSession {
-  prompt(input: string): Promise<string>
-  promptStreaming(input: string): ReadableStream
-  destroy(): void
-  clone(): AITextSession
-}
-
 export interface AITextSessionOptions {
   topK: number
   temperature: number
@@ -45,10 +39,6 @@ function formatPrompt(messages: BaseMessage[]): string {
   return messages
     .map((message) => {
       if (typeof message.content !== "string") {
-        // console.log(message.content)
-        // throw new Error(
-        //   "ChatChromeAI does not support non-string message content."
-        // )
         if (message.content.length > 0) {
           //@ts-ignore
           return message.content[0]?.text || ""
@@ -61,25 +51,6 @@ function formatPrompt(messages: BaseMessage[]): string {
     .join("\n")
 }
 
-/**
- * To use this model you need to have the `Built-in AI Early Preview Program`
- * for Chrome. You can find more information about the program here:
- * @link https://developer.chrome.com/docs/ai/built-in
- *
- * @example
- * ```typescript
- * // Initialize the ChatChromeAI model.
- * const model = new ChatChromeAI({
- *   temperature: 0.5, // Optional. Default is 0.5.
- *   topK: 40, // Optional. Default is 40.
- * });
- *
- * // Call the model with a message and await the response.
- * const response = await model.invoke([
- *   new HumanMessage({ content: "My name is John." }),
- * ]);
- * ```
- */
 export class ChatChromeAI extends SimpleChatModel<ChromeAICallOptions> {
   session?: AITextSession
 
@@ -115,18 +86,16 @@ export class ChatChromeAI extends SimpleChatModel<ChromeAICallOptions> {
       throw new Error("ChatChromeAI can only be used in the browser.")
     }
 
-    const { ai } = window as any
+    const canCreateTextSession = await checkChromeAIAvailability()
-    const capabilities = await ai?.assistant?.capabilities()
-    const canCreateTextSession = capabilities?.available
     if (canCreateTextSession === AIModelAvailability.No) {
       throw new Error("The AI model is not available.")
     } else if (canCreateTextSession === AIModelAvailability.AfterDownload) {
       throw new Error("The AI model is not yet downloaded.")
     }
 
-    this.session = await ai?.assistant?.create({
+    this.session = await createAITextSession({
-      temperature: this.temperature,
+      topK: this.topK,
-      topK: this.topK
+      temperature: this.temperature
     })
   }
 

src/models/utils/chrome.ts

@@ -0,0 +1,54 @@
+
+export const checkChromeAIAvailability = async (): Promise<"readily" | "no" | "after-download"> => {
+    try {
+        const ai = (window as any).ai;
+
+        // upcoming version change 
+        if (ai?.assistant?.capabilities) {
+            const capabilities = await ai.assistant.capabilities();
+            return capabilities?.available ?? "no";
+        }
+
+        // old version 
+        if (ai?.canCreateTextSession) {
+            const available = await ai.canCreateTextSession();
+            return available ?? "no";
+        }
+
+        return "no";
+    } catch (e) {
+        console.error("Error checking Chrome AI availability:", e);
+        return "no";
+    }
+}
+
+export interface AITextSession {
+    prompt(input: string): Promise<string>
+    promptStreaming(input: string): ReadableStream
+    destroy(): void
+    clone(): AITextSession
+}
+
+
+export const createAITextSession = async (data: any): Promise<AITextSession> => {
+    const ai = (window as any).ai;
+
+    // upcoming version change 
+    if (ai?.assistant?.create) {
+        const session = await ai.assistant.create({
+            ...data
+        })
+        return session
+    }
+
+    // old version
+    if (ai.createTextSession) {
+        const session = await ai.createTextSession({
+            ...data
+        })
+
+        return session
+    }
+
+    throw new Error("Chrome AI is not available.")
+}

src/utils/chrome.ts

@@ -1,3 +1,5 @@
+import { checkChromeAIAvailability } from "@/models/utils/chrome"
+
 export const getChromeAISupported = async () => {
   try {
     let browserInfo = navigator.userAgent.match(/Chrom(e|ium)\/([0-9]+)\./)
@@ -11,9 +13,8 @@ export const getChromeAISupported = async () => {
       return "ai_not_supported"
     }
 
-    //@ts-ignore
+    const capabilities = await checkChromeAIAvailability()
-    const capabilities = await ai?.assistant?.capabilities()
+    if (capabilities !== "readily") {
-    if (capabilities?.available !== "readily") {
       return "ai_not_ready"
     }