From 9a3572ca5cae82330c0d0e9b6dfff30631313bf7 Mon Sep 17 00:00:00 2001
From: Nakul <nakulverma.py@gmail.com>
Date: Sat, 16 Aug 2025 19:30:36 +0530
Subject: [PATCH 1/2] fixing chatprovider error

---
 src/default-providers/index.ts | 2 --
 src/tokens.ts                  | 3 ++-
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/src/default-providers/index.ts b/src/default-providers/index.ts
index b953648..7e71290 100644
--- a/src/default-providers/index.ts
+++ b/src/default-providers/index.ts
@@ -56,8 +56,6 @@ const AIProviders: IAIProvider[] = [
   },
   {
     name: 'ChromeAI',
-    // TODO: fix
-    // @ts-expect-error: missing properties
     chat: ChromeAI,
     completer: ChromeCompleter,
     instructions: ChromeAIInstructions,
diff --git a/src/tokens.ts b/src/tokens.ts
index 9d02124..e69ec1c 100644
--- a/src/tokens.ts
+++ b/src/tokens.ts
@@ -3,6 +3,7 @@ import { StructuredToolInterface } from '@langchain/core/tools';
 import { ReadonlyPartialJSONObject, Token } from '@lumino/coreutils';
 import { ISignal } from '@lumino/signaling';
 import { JSONSchema7 } from 'json-schema';
+import { BaseLLM } from '@langchain/core/language_models/llms';
 
 import { IBaseCompleter } from './base-completer';
 import { AIChatModel, AICompleter } from './types/ai-model';
@@ -38,7 +39,7 @@ export interface IAIProvider {
   /**
    * The chat model class to use.
    */
-  chat?: IType<BaseChatModel>;
+  chat?: IType<BaseChatModel> | IType<BaseLLM>;
   /**
    * The completer class to use.
    */

From a9a812ae3c89422543320bee9d9093426ffae4c1 Mon Sep 17 00:00:00 2001
From: Nakul <nakulverma.py@gmail.com>
Date: Sat, 16 Aug 2025 20:04:14 +0530
Subject: [PATCH 2/2] Adding a wrapper to ChromeAI

---
 src/default-providers/ChromeAI/chat.ts | 48 ++++++++++++++++++++++++++
 src/default-providers/index.ts         |  4 +--
 src/tokens.ts                          |  3 +-
 3 files changed, 51 insertions(+), 4 deletions(-)
 create mode 100644 src/default-providers/ChromeAI/chat.ts

diff --git a/src/default-providers/ChromeAI/chat.ts b/src/default-providers/ChromeAI/chat.ts
new file mode 100644
index 0000000..77e4775
--- /dev/null
+++ b/src/default-providers/ChromeAI/chat.ts
@@ -0,0 +1,48 @@
+import {
+  BaseChatModel,
+  BaseChatModelCallOptions
+} from '@langchain/core/language_models/chat_models';
+import {
+  AIMessageChunk,
+  BaseMessage,
+  AIMessage
+} from '@langchain/core/messages';
+import { ChromeAI as ChromeLLM } from '@langchain/community/experimental/llms/chrome_ai';
+import { CallbackManagerForLLMRun } from '@langchain/core/callbacks/manager';
+import { ChatResult, ChatGeneration } from '@langchain/core/outputs';
+
+export interface IChromeChatCallOptions extends BaseChatModelCallOptions {}
+
+export class ChromeChatModel extends BaseChatModel<
+  IChromeChatCallOptions,
+  AIMessageChunk
+> {
+  private llm: ChromeLLM;
+
+  constructor(fields?: ConstructorParameters<typeof ChromeLLM>[0]) {
+    super(fields ?? {});
+    this.llm = new ChromeLLM(fields ?? {});
+  }
+
+  _llmType() {
+    return 'chrome-chat';
+  }
+
+  async _generate(
+    messages: BaseMessage[],
+    options: IChromeChatCallOptions,
+    runManager?: CallbackManagerForLLMRun
+  ): Promise<ChatResult> {
+    const text = messages.map(m => m.content).join('\n');
+    const completion = await this.llm.invoke(text, options);
+
+    const generations: ChatGeneration[] = [
+      {
+        text: completion,
+        message: new AIMessage(completion)
+      }
+    ];
+
+    return { generations };
+  }
+}
diff --git a/src/default-providers/index.ts b/src/default-providers/index.ts
index 7e71290..3561970 100644
--- a/src/default-providers/index.ts
+++ b/src/default-providers/index.ts
@@ -6,7 +6,7 @@ import { Notification } from '@jupyterlab/apputils';
 
 import { ChatAnthropic } from '@langchain/anthropic';
 import { ChatWebLLM } from '@langchain/community/chat_models/webllm';
-import { ChromeAI } from '@langchain/community/experimental/llms/chrome_ai';
+import { ChromeChatModel } from '../default-providers/ChromeAI/chat';
 import { ChatGoogleGenerativeAI } from '@langchain/google-genai';
 import { ChatMistralAI } from '@langchain/mistralai';
 import { ChatOllama } from '@langchain/ollama';
@@ -56,7 +56,7 @@ const AIProviders: IAIProvider[] = [
   },
   {
     name: 'ChromeAI',
-    chat: ChromeAI,
+    chat: ChromeChatModel,
     completer: ChromeCompleter,
     instructions: ChromeAIInstructions,
     settingsSchema: ChromeAISettings,
diff --git a/src/tokens.ts b/src/tokens.ts
index e69ec1c..9d02124 100644
--- a/src/tokens.ts
+++ b/src/tokens.ts
@@ -3,7 +3,6 @@ import { StructuredToolInterface } from '@langchain/core/tools';
 import { ReadonlyPartialJSONObject, Token } from '@lumino/coreutils';
 import { ISignal } from '@lumino/signaling';
 import { JSONSchema7 } from 'json-schema';
-import { BaseLLM } from '@langchain/core/language_models/llms';
 
 import { IBaseCompleter } from './base-completer';
 import { AIChatModel, AICompleter } from './types/ai-model';
@@ -39,7 +38,7 @@ export interface IAIProvider {
   /**
    * The chat model class to use.
    */
-  chat?: IType<BaseChatModel> | IType<BaseLLM>;
+  chat?: IType<BaseChatModel>;
   /**
    * The completer class to use.
    */