Skip to content

Commit 4686a69

Browse files
brichetbrichet
committed
Allow selecting the tool from the chat input directly
1 parent dab7982 commit 4686a69

File tree

11 files changed

+318
-60
lines changed

11 files changed

+318
-60
lines changed

‎schema/chat.json

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,12 @@
2222
"title": "AI persona name",
2323
"description": "The name of the AI persona",
2424
"default": "Jupyternaut"
25+
},
26+
"UseTool": {
27+
"type": "boolean",
28+
"title": "Use tool",
29+
"description": "Whether to be able to use or not a tool in chat",
30+
"default": false
2531
}
2632
},
2733
"additionalProperties": false

‎schema/provider-registry.json

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,12 +17,6 @@
1717
"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.",
1818
"default": true
1919
},
20-
"UseAgent": {
21-
"type": "boolean",
22-
"title": "Use agent in chat",
23-
"description": "Whether to use or not an agent in chat",
24-
"default": false
25-
},
2620
"AIproviders": {
2721
"type": "object",
2822
"title": "AI providers",

‎src/chat-handler.ts

Lines changed: 93 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -24,10 +24,11 @@ import {
2424
SystemMessage
2525
} from '@langchain/core/messages';
2626
import { UUID } from '@lumino/coreutils';
27+
import { ISignal, Signal } from '@lumino/signaling';
2728

2829
import { DEFAULT_CHAT_SYSTEM_PROMPT } from './default-prompts';
2930
import { jupyternautLiteIcon } from './icons';
30-
import { IAIProviderRegistry } from './tokens';
31+
import { IAIProviderRegistry, IToolRegistry, Tool } from './tokens';
3132
import { AIChatModel } from './types/ai-model';
3233

3334
/**
@@ -58,16 +59,30 @@ export class ChatHandler extends AbstractChatModel {
5859
constructor(options: ChatHandler.IOptions) {
5960
super(options);
6061
this._providerRegistry = options.providerRegistry;
62+
this._toolRegistry = options.toolRegistry;
6163

6264
this._providerRegistry.providerChanged.connect(() => {
6365
this._errorMessage = this._providerRegistry.chatError;
6466
});
6567
}
6668

69+
/**
70+
* Get the tool registry.
71+
*/
72+
get toolRegistry(): IToolRegistry | undefined {
73+
return this._toolRegistry;
74+
}
75+
76+
/**
77+
* Get the agent from the provider registry.
78+
*/
6779
get agent(): AIChatModel | null {
6880
return this._providerRegistry.currentAgent;
6981
}
7082

83+
/**
84+
* Get the chat model from the provider registry.
85+
*/
7186
get chatModel(): AIChatModel | null {
7287
return this._providerRegistry.currentChatModel;
7388
}
@@ -90,12 +105,46 @@ export class ChatHandler extends AbstractChatModel {
90105
}
91106

92107
/**
93-
* Get/set the system prompt for the chat.
108+
* Getter/setter for the use of tools.
109+
*/
110+
get useTool(): boolean {
111+
return this._useTool;
112+
}
113+
set useTool(value: boolean) {
114+
if (this._useTool !== value) {
115+
this._useTool = value;
116+
this._useToolChanged.emit(this._useTool);
117+
}
118+
}
119+
120+
/**
121+
* Get/set a tool, which will build an agent.
122+
*/
123+
get tool(): Tool | null {
124+
return this._tool;
125+
}
126+
set tool(value: Tool | null) {
127+
this._tool = value;
128+
this._providerRegistry.buildAgent(this._tool);
129+
}
130+
131+
/**
132+
* A signal triggered when the setting on tool usage has changed.
133+
*/
134+
get useToolChanged(): ISignal<ChatHandler, boolean> {
135+
return this._useToolChanged;
136+
}
137+
138+
/**
139+
* Get the system prompt for the chat.
94140
*/
95141
get systemPrompt(): string {
96-
return (
97-
this._providerRegistry.chatSystemPrompt ?? DEFAULT_CHAT_SYSTEM_PROMPT
98-
);
142+
let prompt =
143+
this._providerRegistry.chatSystemPrompt ?? DEFAULT_CHAT_SYSTEM_PROMPT;
144+
if (this.useTool && this.agent !== null) {
145+
prompt = prompt.concat('\nPlease use the tool that is provided');
146+
}
147+
return prompt;
99148
}
100149

101150
async sendMessage(message: INewMessage): Promise<boolean> {
@@ -145,7 +194,7 @@ export class ChatHandler extends AbstractChatModel {
145194
const sender = { username: this._personaName, avatar_url: AI_AVATAR };
146195
this.updateWriters([{ user: sender }]);
147196

148-
if (this._providerRegistry.useAgent && this.agent !== null) {
197+
if (this._useTool && this.agent !== null) {
149198
return this._sendAgentMessage(this.agent, messages, sender);
150199
}
151200

@@ -281,12 +330,17 @@ export class ChatHandler extends AbstractChatModel {
281330
this._controller = null;
282331
}
283332
}
333+
284334
private _providerRegistry: IAIProviderRegistry;
285335
private _personaName = 'AI';
286336
private _errorMessage: string = '';
287337
private _history: IChatHistory = { messages: [] };
288338
private _defaultErrorMessage = 'AI provider not configured';
289339
private _controller: AbortController | null = null;
340+
private _useTool: boolean = false;
341+
private _tool: Tool | null = null;
342+
private _toolRegistry?: IToolRegistry;
343+
private _useToolChanged = new Signal<ChatHandler, boolean>(this);
290344
}
291345

292346
export namespace ChatHandler {
@@ -295,13 +349,45 @@ export namespace ChatHandler {
295349
*/
296350
export interface IOptions extends IChatModel.IOptions {
297351
providerRegistry: IAIProviderRegistry;
352+
toolRegistry?: IToolRegistry;
298353
}
299354

300355
/**
301-
* The minimal chat context.
356+
* The chat context.
302357
*/
303358
export class ChatContext extends AbstractChatContext {
304359
users = [];
360+
361+
/**
362+
* The tool registry.
363+
*/
364+
get toolsRegistry(): IToolRegistry | undefined {
365+
return (this._model as ChatHandler).toolRegistry;
366+
}
367+
368+
/**
369+
* Whether to use or not the tool.
370+
*/
371+
get useTool(): boolean {
372+
return (this._model as ChatHandler).useTool;
373+
}
374+
375+
/**
376+
* A signal triggered when the setting on tool usage has changed.
377+
*/
378+
get useToolChanged(): ISignal<ChatHandler, boolean> {
379+
return (this._model as ChatHandler).useToolChanged;
380+
}
381+
382+
/**
383+
* Getter/setter of the tool to use.
384+
*/
385+
get tool(): Tool | null {
386+
return (this._model as ChatHandler).tool;
387+
}
388+
set tool(value: Tool | null) {
389+
(this._model as ChatHandler).tool = value;
390+
}
305391
}
306392

307393
/**

‎src/components/index.ts

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
/*
2+
* Copyright (c) Jupyter Development Team.
3+
* Distributed under the terms of the Modified BSD License.
4+
*/
5+
6+
export * from './stop-button';
7+
export * from './tool-select';

‎src/components/tool-select.tsx

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,150 @@
1+
/*
2+
* Copyright (c) Jupyter Development Team.
3+
* Distributed under the terms of the Modified BSD License.
4+
*/
5+
6+
import { InputToolbarRegistry, TooltippedButton } from '@jupyter/chat';
7+
import { checkIcon } from '@jupyterlab/ui-components';
8+
import BuildIcon from '@mui/icons-material/Build';
9+
import { Menu, MenuItem, Typography } from '@mui/material';
10+
import React, { useCallback, useEffect, useState } from 'react';
11+
12+
import { ChatHandler } from '../chat-handler';
13+
14+
const SELECT_ITEM_CLASS = 'jp-AIToolSelect-item';
15+
16+
/**
17+
* The tool select component.
18+
*/
19+
export function toolSelect(
20+
props: InputToolbarRegistry.IToolbarItemProps
21+
): JSX.Element {
22+
const chatContext = props.model.chatContext as ChatHandler.ChatContext;
23+
const toolRegistry = chatContext.toolsRegistry;
24+
25+
const [useTool, setUseTool] = useState<boolean>(chatContext.useTool);
26+
const [selectedTool, setSelectedTool] = useState<string | null>(null);
27+
const [toolNames, setToolNames] = useState<string[]>(
28+
toolRegistry?.toolNames || []
29+
);
30+
const [menuAnchorEl, setMenuAnchorEl] = useState<HTMLElement | null>(null);
31+
const [menuOpen, setMenuOpen] = useState(false);
32+
33+
const openMenu = useCallback((el: HTMLElement | null) => {
34+
setMenuAnchorEl(el);
35+
setMenuOpen(true);
36+
}, []);
37+
38+
const closeMenu = useCallback(() => {
39+
setMenuOpen(false);
40+
}, []);
41+
42+
const onClick = useCallback(
43+
(tool: string | null) => {
44+
setSelectedTool(tool);
45+
chatContext.tool = toolRegistry?.get(tool) || null;
46+
},
47+
[props.model]
48+
);
49+
50+
useEffect(() => {
51+
const updateTools = () => setToolNames(toolRegistry?.toolNames || []);
52+
toolRegistry?.toolsChanged.connect(updateTools);
53+
return () => {
54+
toolRegistry?.toolsChanged.disconnect(updateTools);
55+
};
56+
}, [toolRegistry]);
57+
58+
useEffect(() => {
59+
const updateUseTool = (_: ChatHandler, value: boolean) => setUseTool(value);
60+
chatContext.useToolChanged.connect(updateUseTool);
61+
return () => {
62+
chatContext.useToolChanged.disconnect(updateUseTool);
63+
};
64+
}, [chatContext]);
65+
66+
return useTool && toolNames.length ? (
67+
<>
68+
<TooltippedButton
69+
onClick={e => {
70+
openMenu(e.currentTarget);
71+
}}
72+
disabled={!toolNames.length}
73+
tooltip="Tool"
74+
buttonProps={{
75+
variant: 'contained',
76+
onKeyDown: e => {
77+
if (e.key !== 'Enter' && e.key !== ' ') {
78+
return;
79+
}
80+
openMenu(e.currentTarget);
81+
// stopping propagation of this event prevents the prompt from being
82+
// sent when the dropdown button is selected and clicked via 'Enter'.
83+
e.stopPropagation();
84+
}
85+
}}
86+
sx={
87+
selectedTool === null
88+
? { backgroundColor: 'var(--jp-layout-color3)' }
89+
: {}
90+
}
91+
>
92+
<BuildIcon />
93+
</TooltippedButton>
94+
<Menu
95+
open={menuOpen}
96+
onClose={closeMenu}
97+
anchorEl={menuAnchorEl}
98+
anchorOrigin={{
99+
vertical: 'top',
100+
horizontal: 'right'
101+
}}
102+
transformOrigin={{
103+
vertical: 'bottom',
104+
horizontal: 'right'
105+
}}
106+
sx={{
107+
'& .MuiMenuItem-root': {
108+
gap: '4px',
109+
padding: '6px'
110+
}
111+
}}
112+
>
113+
<MenuItem
114+
className={SELECT_ITEM_CLASS}
115+
onClick={e => {
116+
onClick(null);
117+
// prevent sending second message with no selection
118+
e.stopPropagation();
119+
}}
120+
>
121+
{selectedTool === null ? (
122+
<checkIcon.react className={'lm-Menu-itemIcon'} />
123+
) : (
124+
<div className={'lm-Menu-itemIcon'} />
125+
)}
126+
<Typography display="block">No tool</Typography>
127+
</MenuItem>
128+
{toolNames.map(tool => (
129+
<MenuItem
130+
className={SELECT_ITEM_CLASS}
131+
onClick={e => {
132+
onClick(tool);
133+
// prevent sending second message with no selection
134+
e.stopPropagation();
135+
}}
136+
>
137+
{selectedTool === tool ? (
138+
<checkIcon.react className={'lm-Menu-itemIcon'} />
139+
) : (
140+
<div className={'lm-Menu-itemIcon'} />
141+
)}
142+
<Typography display="block">{tool}</Typography>
143+
</MenuItem>
144+
))}
145+
</Menu>
146+
</>
147+
) : (
148+
<></>
149+
);
150+
}

0 commit comments

Comments
 (0)