Include langgraph to call tool from the chat prompt

Author: brichet

package.json

@@ -67,7 +67,8 @@
         "@jupyterlab/ui-components": "^4.4.0",
         "@langchain/anthropic": "^0.3.9",
         "@langchain/community": "^0.3.44",
-        "@langchain/core": "^0.3.57",
+        "@langchain/core": "^0.3.61",
+        "@langchain/langgraph": "^0.3.5",
         "@langchain/mistralai": "^0.1.1",
         "@langchain/ollama": "^0.2.0",
         "@langchain/openai": "^0.4.4",

schema/provider-registry.json

@@ -11,6 +11,12 @@
       "description": "Whether to use or not the secrets manager. If not, secrets will be stored in the browser (local storage)",
       "default": true
     },
+    "UseAgent": {
+      "type": "boolean",
+      "title": "Use agent in chat",
+      "description": "Whether to use or not an agent in chat",
+      "default": false
+    },
     "AIprovider": {
       "type": "object",
       "title": "AI provider",

schema/system-prompts.json

@@ -9,7 +9,7 @@
       "type": "string",
       "title": "Chat message system prompt",
       "description": "The system prompt for the chat messages",
-      "default": "You are Jupyternaut, a conversational assistant living in JupyterLab to help users.\nYou are not a language model, but rather an application built on a foundation model from $provider_name$.\nYou are talkative and you provide lots of specific details from the foundation model's context.\nYou may use Markdown to format your response.\nIf your response includes code, they must be enclosed in Markdown fenced code blocks (with triple backticks before and after).\nIf your response includes mathematical notation, they must be expressed in LaTeX markup and enclosed in LaTeX delimiters.\nAll dollar quantities (of USD) must be formatted in LaTeX, with the `$` symbol escaped by a single backslash `\\`.\n- Example prompt: `If I have \\\\$100 and spend \\\\$20, how much money do I have left?`\n- **Correct** response: `You have \\(\\$80\\) remaining.`\n- **Incorrect** response: `You have $80 remaining.`\nIf you do not know the answer to a question, answer truthfully by responding that you do not know.\nThe following is a friendly conversation between you and a human."
+      "default": "You are Jupyternaut, a conversational assistant living in JupyterLab to help users.\nYou are not a language model, but rather an application built on a foundation model from $provider_name$.\nYou are talkative and you provide lots of specific details from the foundation model's context.\nYou may use Markdown to format your response.\nIf your response includes code, they must be enclosed in Markdown fenced code blocks (with triple backticks before and after).\nIf your response includes mathematical notation, they must be expressed in LaTeX markup and enclosed in LaTeX delimiters.\nAll dollar quantities (of USD) must be formatted in LaTeX, with the `$` symbol escaped by a single backslash `\\`.\n- Example prompt: `If I have \\\\$100 and spend \\\\$20, how much money do I have left?`\n- **Correct** response: `You have \\(\\$80\\) remaining.`\n- **Incorrect** response: `You have $80 remaining.`\nIf you do not know the answer to a question, answer truthfully by responding that you do not know.\nAlways use the tool to answer, if available.\nThe following is a friendly conversation between you and a human."
     },
     "completionSystemPrompt": {
       "type": "string",

src/chat-handler.ts

@@ -13,14 +13,17 @@ import {
   IChatMessage,
   IChatModel,
   IInputModel,
-  INewMessage
+  INewMessage,
+  IUser
 } from '@jupyter/chat';
 import {
   AIMessage,
+  BaseMessage,
   HumanMessage,
   mergeMessageRuns,
   SystemMessage
 } from '@langchain/core/messages';
+import { CompiledStateGraph } from '@langchain/langgraph';
 import { UUID } from '@lumino/coreutils';
 
 import { DEFAULT_CHAT_SYSTEM_PROMPT } from './default-prompts';
@@ -62,7 +65,11 @@ export class ChatHandler extends AbstractChatModel {
     });
   }
 
-  get provider(): AIChatModel | null {
+  get agent(): CompiledStateGraph<any, any> | null {
+    return this._providerRegistry.currentAgent;
+  }
+
+  get chatModel(): AIChatModel | null {
     return this._providerRegistry.currentChatModel;
   }
 
@@ -110,7 +117,9 @@ export class ChatHandler extends AbstractChatModel {
     };
     this.messageAdded(msg);
 
-    if (this._providerRegistry.currentChatModel === null) {
+    const chatModel = this.chatModel;
+
+    if (chatModel === null) {
       const errorMsg: IChatMessage = {
         id: UUID.uuid4(),
         body: `**${this._errorMessage ? this._errorMessage : this._defaultErrorMessage}**`,
@@ -137,23 +146,52 @@ export class ChatHandler extends AbstractChatModel {
     const sender = { username: this._personaName, avatar_url: AI_AVATAR };
     this.updateWriters([{ user: sender }]);
 
-    // create an empty message to be filled by the AI provider
+    if (this._providerRegistry.useAgent && this.agent !== null) {
+      return this._sendAgentMessage(this.agent, messages, sender);
+    }
+
+    return this._sentChatMessage(chatModel, messages, sender);
+  }
+
+  async getHistory(): Promise<IChatHistory> {
+    return this._history;
+  }
+
+  dispose(): void {
+    super.dispose();
+  }
+
+  messageAdded(message: IChatMessage): void {
+    super.messageAdded(message);
+  }
+
+  stopStreaming(): void {
+    this._controller?.abort();
+  }
+
+  createChatContext(): IChatContext {
+    return new ChatHandler.ChatContext({ model: this });
+  }
+
+  private async _sentChatMessage(
+    chatModel: AIChatModel,
+    messages: BaseMessage[],
+    sender: IUser
+  ): Promise<boolean> {
+    // Create an empty message to be filled by the AI provider
     const botMsg: IChatMessage = {
       id: UUID.uuid4(),
       body: '',
       sender,
       time: Private.getTimestampMs(),
       type: 'msg'
     };
-
     let content = '';
-
     this._controller = new AbortController();
     try {
-      for await (const chunk of await this._providerRegistry.currentChatModel.stream(
-        messages,
-        { signal: this._controller.signal }
-      )) {
+      for await (const chunk of await chatModel.stream(messages, {
+        signal: this._controller.signal
+      })) {
         content += chunk.content ?? chunk;
         botMsg.body = content;
         this.messageAdded(botMsg);
@@ -177,26 +215,73 @@ export class ChatHandler extends AbstractChatModel {
     }
   }
 
-  async getHistory(): Promise<IChatHistory> {
-    return this._history;
-  }
-
-  dispose(): void {
-    super.dispose();
-  }
-
-  messageAdded(message: IChatMessage): void {
-    super.messageAdded(message);
-  }
-
-  stopStreaming(): void {
-    this._controller?.abort();
-  }
-
-  createChatContext(): IChatContext {
-    return new ChatHandler.ChatContext({ model: this });
+  private async _sendAgentMessage(
+    agent: CompiledStateGraph<any, any>,
+    messages: BaseMessage[],
+    sender: IUser
+  ): Promise<boolean> {
+    this._controller = new AbortController();
+    try {
+      for await (const chunk of await agent.stream(
+        { messages },
+        {
+          streamMode: 'updates',
+          signal: this._controller.signal
+        }
+      )) {
+        if ((chunk as any).agent) {
+          messages = (chunk as any).agent.messages;
+          messages.forEach(message => {
+            const contents: string[] = [];
+            if (typeof message.content === 'string') {
+              contents.push(message.content);
+            } else if (Array.isArray(message.content)) {
+              message.content.forEach(content => {
+                if (content.type === 'text') {
+                  contents.push(content.text);
+                }
+              });
+            }
+            contents.forEach(content => {
+              this.messageAdded({
+                id: UUID.uuid4(),
+                body: content,
+                sender,
+                time: Private.getTimestampMs(),
+                type: 'msg'
+              });
+            });
+          });
+        } else if ((chunk as any).tools) {
+          messages = (chunk as any).tools.messages;
+          messages.forEach(message => {
+            this.messageAdded({
+              id: UUID.uuid4(),
+              body: message.content as string,
+              sender: { username: `Tool "${message.name}"` },
+              time: Private.getTimestampMs(),
+              type: 'msg'
+            });
+          });
+        }
+      }
+      return true;
+    } catch (reason) {
+      const error = this._providerRegistry.formatErrorMessage(reason);
+      const errorMsg: IChatMessage = {
+        id: UUID.uuid4(),
+        body: `**${error}**`,
+        sender: { username: 'ERROR' },
+        time: Private.getTimestampMs(),
+        type: 'msg'
+      };
+      this.messageAdded(errorMsg);
+      return false;
+    } finally {
+      this.updateWriters([]);
+      this._controller = null;
+    }
   }
-
   private _providerRegistry: IAIProviderRegistry;
   private _personaName = 'AI';
   private _errorMessage: string = '';

src/index.ts

@@ -203,7 +203,13 @@ const providerRegistryPlugin: JupyterFrontEndPlugin<IAIProviderRegistry> =
           if (!secretsManager) {
             delete settings.schema.properties?.['UseSecretsManager'];
           }
+
           const updateProvider = () => {
+            // Update agent usage if necessary.
+            const useAgent =
+              (settings.get('UseAgent').composite as boolean) ?? false;
+            providerRegistry.useAgent = useAgent;
+
             // Update the settings to the AI providers.
             const providerSettings = (settings.get('AIprovider').composite ?? {
               provider: 'None'

src/provider.ts

@@ -3,6 +3,8 @@ import {
   IInlineCompletionContext
 } from '@jupyterlab/completer';
 import { BaseLanguageModel } from '@langchain/core/language_models/base';
+import { CompiledStateGraph } from '@langchain/langgraph';
+import { createReactAgent } from '@langchain/langgraph/prebuilt';
 import { BaseChatModel } from '@langchain/core/language_models/chat_models';
 import { ISignal, Signal } from '@lumino/signaling';
 import { ReadonlyPartialJSONObject } from '@lumino/coreutils';
@@ -18,6 +20,7 @@ import {
   ISetProviderOptions,
   PLUGIN_IDS
 } from './tokens';
+import { testTool } from './tools/test_tool';
 import { AIChatModel, AICompleter } from './types/ai-model';
 
 const SECRETS_NAMESPACE = PLUGIN_IDS.providerRegistry;
@@ -149,6 +152,17 @@ export class AIProviderRegistry implements IAIProviderRegistry {
     };
   }
 
+  /**
+   * Get the current agent.
+   */
+  get currentAgent(): CompiledStateGraph<any, any> | null {
+    const agent = Private.getAgent();
+    if (agent === null) {
+      return null;
+    }
+    return agent;
+  }
+
   /**
    * Getter/setter for the chat system prompt.
    */
@@ -159,6 +173,20 @@ export class AIProviderRegistry implements IAIProviderRegistry {
     this._chatPrompt = value;
   }
 
+  /**
+   * Getter/setter for the use of agent in chat.
+   */
+  get useAgent(): boolean {
+    return this._useAgentInChat;
+  }
+  set useAgent(value: boolean) {
+    if (value === this._useAgentInChat) {
+      return;
+    }
+    this._useAgentInChat = value;
+    this._buildAgent();
+  }
+
   /**
    * Get the settings schema of a given provider.
    */
@@ -286,6 +314,9 @@ export class AIProviderRegistry implements IAIProviderRegistry {
             ...fullSettings
           })
         );
+        if (this._useAgentInChat) {
+          this._buildAgent();
+        }
         this._chatError = '';
       } catch (e: any) {
         this._chatError = e.message;
@@ -305,13 +336,37 @@ export class AIProviderRegistry implements IAIProviderRegistry {
     return this._providerChanged;
   }
 
+  /**
+   * Build an agent.
+   */
+  private _buildAgent() {
+    if (this._useAgentInChat) {
+      const chatModel = Private.getChatModel();
+      if (chatModel === null) {
+        Private.setAgent(null);
+        return;
+      }
+      chatModel.bindTools?.([testTool]);
+      Private.setChatModel(chatModel);
+      Private.setAgent(
+        createReactAgent({
+          llm: chatModel,
+          tools: [testTool]
+        })
+      );
+    } else {
+      Private.setAgent(null);
+    }
+  }
+
   private _secretsManager: ISecretsManager | null;
   private _providerChanged = new Signal<IAIProviderRegistry, void>(this);
   private _chatError: string = '';
   private _completerError: string = '';
   private _deferredProvider: ISetProviderOptions | null = null;
   private _chatPrompt: string = '';
   private _completerPrompt: string = '';
+  private _useAgentInChat = false;
 }
 
 export namespace AIProviderRegistry {
@@ -433,4 +488,15 @@ namespace Private {
   export function getCompleter(): IBaseCompleter | null {
     return completer;
   }
+
+  /**
+   * The agent getter and setter.
+   */
+  let agent: CompiledStateGraph<any, any> | null = null;
+  export function setAgent(value: CompiledStateGraph<any, any> | null): void {
+    agent = value;
+  }
+  export function getAgent(): CompiledStateGraph<any, any> | null {
+    return agent;
+  }
 }