Add a tool registry

Author: brichet

src/index.ts

@@ -23,11 +23,11 @@ import { ISecretsManager, SecretsManager } from 'jupyter-secrets-manager';
 
 import { ChatHandler, welcomeMessage } from './chat-handler';
 import { CompletionProvider } from './completion-provider';
+import { stopItem, toolSelect } from './components';
 import { defaultProviderPlugins } from './default-providers';
 import { AIProviderRegistry } from './provider';
 import { aiSettingsRenderer, textArea } from './settings';
-import { IAIProviderRegistry, PLUGIN_IDS } from './tokens';
-import { stopItem } from './components/stop-button';
+import { IAIProviderRegistry, IToolRegistry, PLUGIN_IDS } from './tokens';
 
 const chatCommandRegistryPlugin: JupyterFrontEndPlugin<IChatCommandRegistry> = {
   id: PLUGIN_IDS.chatCommandRegistry,
@@ -301,12 +301,24 @@ const systemPromptsPlugin: JupyterFrontEndPlugin<void> = {
   }
 };
 
+const toolRegistryPlugin: JupyterFrontEndPlugin<IToolRegistry> = {
+  id: PLUGIN_IDS.toolRegistry,
+  autoStart: true,
+  provides: IToolRegistry,
+  activate: (app: JupyterFrontEnd): IToolRegistry => {
+    const registry = new ToolsRegistry();
+    registry.add(testTool);
+    return registry;
+  }
+};
+
 export default [
   providerRegistryPlugin,
   chatCommandRegistryPlugin,
   chatPlugin,
   completerPlugin,
   systemPromptsPlugin,
+  toolRegistryPlugin,
   ...defaultProviderPlugins
 ];
 

src/tokens.ts

@@ -1,4 +1,5 @@
 import { BaseChatModel } from '@langchain/core/language_models/chat_models';
+import { StructuredToolInterface } from '@langchain/core/tools';
 import { ReadonlyPartialJSONObject, Token } from '@lumino/coreutils';
 import { ISignal } from '@lumino/signaling';
 import { JSONSchema7 } from 'json-schema';
@@ -12,7 +13,8 @@ export const PLUGIN_IDS = {
   completer: '@jupyterlite/ai:completer',
   providerRegistry: '@jupyterlite/ai:provider-registry',
   settingsConnector: '@jupyterlite/ai:settings-connector',
-  systemPrompts: '@jupyterlite/ai:system-prompts'
+  systemPrompts: '@jupyterlite/ai:system-prompts',
+  toolRegistry: '@jupyterlite/ai:tool-registry'
 };
 
 export type ModelRole = 'chat' | 'completer';
@@ -157,10 +159,41 @@ export interface IAIProviderRegistry {
   readonly completerError: string;
 }
 
+/**
+ * The type describing a tool used in langgraph.
+ */
+export type Tool = StructuredToolInterface;
+
+/**
+ * The tool registry interface.
+ */
+export interface IToolRegistry {
+  /**
+   * Get the registered tool names.
+   */
+  readonly toolNames: string[];
+  /**
+   * A signal triggered when the tools has changed;
+   */
+  readonly toolsChanged: ISignal<IToolRegistry, void>;
+  /**
+   * Add a new tool.
+   */
+  add(provider: Tool): void;
+}
+
 /**
  * The provider registry token.
  */
 export const IAIProviderRegistry = new Token<IAIProviderRegistry>(
   '@jupyterlite/ai:provider-registry',
   'Provider for chat and completion LLM provider'
 );
+
+/**
+ * The tool registry token.
+ */
+export const IToolRegistry = new Token<IToolRegistry>(
+  '@jupyterlite/ai:tool-registry',
+  'Tool registry for AI agent'
+);

src/tool-registry.ts

@@ -0,0 +1,23 @@
+import { ISignal, Signal } from '@lumino/signaling';
+import { IToolRegistry, Tool } from './tokens';
+
+export class ToolsRegistry implements IToolRegistry {
+  get toolNames(): string[] {
+    return this._tools.map(tool => tool.name);
+  }
+
+  get toolsChanged(): ISignal<IToolRegistry, void> {
+    return this._toolsChanged;
+  }
+
+  add(tool: Tool): void {
+    const index = this._tools.findIndex(t => t.name === tool.name);
+    if (index === -1) {
+      this._tools.push(tool);
+      this._toolsChanged.emit();
+    }
+  }
+
+  private _tools: Tool[] = [];
+  private _toolsChanged = new Signal<IToolRegistry, void>(this);
+}