Pass only a subset of the agent

Author: brichet

src/chat-handler.ts

@@ -23,7 +23,6 @@ import {
   mergeMessageRuns,
   SystemMessage
 } from '@langchain/core/messages';
-import { CompiledStateGraph } from '@langchain/langgraph';
 import { UUID } from '@lumino/coreutils';
 
 import { DEFAULT_CHAT_SYSTEM_PROMPT } from './default-prompts';
@@ -65,7 +64,7 @@ export class ChatHandler extends AbstractChatModel {
     });
   }
 
-  get agent(): CompiledStateGraph<any, any> | null {
+  get agent(): AIChatModel | null {
     return this._providerRegistry.currentAgent;
   }
 
@@ -216,7 +215,7 @@ export class ChatHandler extends AbstractChatModel {
   }
 
   private async _sendAgentMessage(
-    agent: CompiledStateGraph<any, any>,
+    agent: AIChatModel,
     messages: BaseMessage[],
     sender: IUser
   ): Promise<boolean> {

src/provider.ts

@@ -159,12 +159,14 @@ export class AIProviderRegistry implements IAIProviderRegistry {
   /**
    * Get the current agent.
    */
-  get currentAgent(): CompiledStateGraph<any, any> | null {
+  get currentAgent(): AIChatModel | null {
     const agent = Private.getAgent();
     if (agent === null) {
       return null;
     }
-    return agent;
+    return {
+      stream: (input: any, options?: any) => agent.stream(input, options)
+    };
   }
 
   /**

src/tokens.ts

@@ -1,4 +1,3 @@
-import { CompiledStateGraph } from '@langchain/langgraph';
 import { BaseChatModel } from '@langchain/core/language_models/chat_models';
 import { ReadonlyPartialJSONObject, Token } from '@lumino/coreutils';
 import { ISignal } from '@lumino/signaling';
@@ -103,7 +102,7 @@ export interface IAIProviderRegistry {
   /**
    * Get the current agent.
    */
-  readonly currentAgent: CompiledStateGraph<any, any> | null;
+  readonly currentAgent: AIChatModel | null;
   /**
    * Getter/setter for the chat system prompt.
    */