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

@@ -17,6 +17,12 @@
       "description": "Whether to use only one provider for the chat and the completer.\nThis will overwrite all the settings for the completer, and copy the ones from the chat.",
       "default": true
     },
+    "UseAgent": {
+      "type": "boolean",
+      "title": "Use agent in chat",
+      "description": "Whether to use or not an agent in chat",
+      "default": false
+    },
     "AIproviders": {
       "type": "object",
       "title": "AI providers",

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;
+
             // Get the Ai provider settings.
             const providerSettings = settings.get('AIproviders')
               .composite as ReadonlyPartialJSONObject;

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 {
   ModelRole,
   PLUGIN_IDS
 } from './tokens';
+import { testTool } from './tools/test_tool';
 import { AIChatModel, AICompleter } from './types/ai-model';
 
 const SECRETS_NAMESPACE = PLUGIN_IDS.providerRegistry;
@@ -153,6 +156,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.
    */
@@ -166,6 +180,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.
    */
@@ -329,6 +357,9 @@ export class AIProviderRegistry implements IAIProviderRegistry {
             ...fullSettings
           })
         );
+        if (this._useAgentInChat) {
+          this._buildAgent();
+        }
       } catch (e: any) {
         this._chatError = e.message;
         Private.setChatModel(null);
@@ -374,6 +405,29 @@ export class AIProviderRegistry implements IAIProviderRegistry {
     return fullSettings;
   }
 
+  /**
+   * 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, ModelRole>(this);
   private _chatError: string = '';
@@ -386,6 +440,7 @@ export class AIProviderRegistry implements IAIProviderRegistry {
   };
   private _chatPrompt: string = '';
   private _completerPrompt: string = '';
+  private _useAgentInChat = false;
 }
 
 export namespace AIProviderRegistry {
@@ -510,4 +565,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;
+  }
 }

src/tokens.ts

@@ -1,3 +1,4 @@
+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';
@@ -90,19 +91,27 @@ export interface IAIProviderRegistry {
   /**
    * Get the current completer of the completion provider.
    */
-  currentCompleter: AICompleter | null;
+  readonly currentCompleter: AICompleter | null;
   /**
    * Getter/setter for the completer system prompt.
    */
   completerSystemPrompt: string;
   /**
    * Get the current llm chat model.
    */
-  currentChatModel: AIChatModel | null;
+  readonly currentChatModel: AIChatModel | null;
+  /**
+   * Get the current agent.
+   */
+  readonly currentAgent: CompiledStateGraph<any, any> | null;
   /**
    * Getter/setter for the chat system prompt.
    */
   chatSystemPrompt: string;
+  /**
+   * Getter/setter for the use of agent in chat.
+   */
+  useAgent: boolean;
   /**
    * Get the settings schema of a given provider.
    */