Add Chat Agent Pinning (#14716)

Co-authored-by: mihaela21k <[email protected]> Co-authored-by: Fuyao Tang <[email protected]> Co-authored-by: Zidong Wang <[email protected]>
eclipse-theia · Feb 17, 2025 · b119b20 · b119b20
1 parent 0419c0a
commit b119b20
Show file tree

Hide file tree

Showing 8 changed files with 143 additions and 22 deletions.
diff --git a/packages/ai-chat-ui/src/browser/ai-chat-ui-frontend-module.ts b/packages/ai-chat-ui/src/browser/ai-chat-ui-frontend-module.ts
@@ -61,7 +61,8 @@ export default new ContainerModule((bind, _unbind, _isBound, rebind) => {
 
     bind(AIChatInputWidget).toSelf();
     bind(AIChatInputConfiguration).toConstantValue({
-        showContext: false
+        showContext: false,
+        showPinnedAgent: true
     });
     bind(WidgetFactory).toDynamicValue(({ container }) => ({
         id: AIChatInputWidget.ID,

diff --git a/packages/ai-chat-ui/src/browser/chat-input-widget.tsx b/packages/ai-chat-ui/src/browser/chat-input-widget.tsx
@@ -13,7 +13,7 @@
 //
 // SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
 // *****************************************************************************
-import { ChangeSet, ChangeSetElement, ChatChangeEvent, ChatModel, ChatRequestModel } from '@theia/ai-chat';
+import { ChangeSet, ChangeSetElement, ChatAgent, ChatChangeEvent, ChatModel, ChatRequestModel } from '@theia/ai-chat';
 import { Disposable, UntitledResourceResolver } from '@theia/core';
 import { ContextMenuRenderer, LabelProvider, Message, ReactWidget } from '@theia/core/lib/browser';
 import { Deferred } from '@theia/core/lib/common/promise-util';
@@ -25,13 +25,15 @@ import { MonacoEditorProvider } from '@theia/monaco/lib/browser/monaco-editor-pr
 import { CHAT_VIEW_LANGUAGE_EXTENSION } from './chat-view-language-contribution';
 
 type Query = (query: string) => Promise<void>;
+type Unpin = () => void;
 type Cancel = (requestModel: ChatRequestModel) => void;
 type DeleteChangeSet = (requestModel: ChatRequestModel) => void;
 type DeleteChangeSetElement = (requestModel: ChatRequestModel, index: number) => void;
 
 export const AIChatInputConfiguration = Symbol('AIChatInputConfiguration');
 export interface AIChatInputConfiguration {
     showContext?: boolean;
+    showPinnedAgent?: boolean;
 }
 
 @injectable()
@@ -63,6 +65,10 @@ export class AIChatInputWidget extends ReactWidget {
     set onQuery(query: Query) {
         this._onQuery = query;
     }
+    private _onUnpin: Unpin;
+    set onUnpin(unpin: Unpin) {
+        this._onUnpin = unpin;
+    }
     private _onCancel: Cancel;
     set onCancel(cancel: Cancel) {
         this._onCancel = cancel;
@@ -80,6 +86,11 @@ export class AIChatInputWidget extends ReactWidget {
         this._chatModel = chatModel;
         this.update();
     }
+    private _pinnedAgent: ChatAgent | undefined;
+    set pinnedAgent(pinnedAgent: ChatAgent | undefined) {
+        this._pinnedAgent = pinnedAgent;
+        this.update();
+    }
 
     @postConstruct()
     protected init(): void {
@@ -101,10 +112,12 @@ export class AIChatInputWidget extends ReactWidget {
         return (
             <ChatInput
                 onQuery={this._onQuery.bind(this)}
+                onUnpin={this._onUnpin.bind(this)}
                 onCancel={this._onCancel.bind(this)}
                 onDeleteChangeSet={this._onDeleteChangeSet.bind(this)}
                 onDeleteChangeSetElement={this._onDeleteChangeSetElement.bind(this)}
                 chatModel={this._chatModel}
+                pinnedAgent={this._pinnedAgent}
                 editorProvider={this.editorProvider}
                 untitledResourceResolver={this.untitledResourceResolver}
                 contextMenuCallback={this.handleContextMenu.bind(this)}
@@ -114,6 +127,7 @@ export class AIChatInputWidget extends ReactWidget {
                     this.editorReady.resolve();
                 }}
                 showContext={this.configuration?.showContext}
+                showPinnedAgent={this.configuration?.showPinnedAgent}
                 labelProvider={this.labelProvider}
             />
         );
@@ -137,15 +151,18 @@ export class AIChatInputWidget extends ReactWidget {
 interface ChatInputProperties {
     onCancel: (requestModel: ChatRequestModel) => void;
     onQuery: (query: string) => void;
+    onUnpin: () => void;
     onDeleteChangeSet: (sessionId: string) => void;
     onDeleteChangeSetElement: (sessionId: string, index: number) => void;
     isEnabled?: boolean;
     chatModel: ChatModel;
+    pinnedAgent?: ChatAgent;
     editorProvider: MonacoEditorProvider;
     untitledResourceResolver: UntitledResourceResolver;
     contextMenuCallback: (event: IMouseEvent) => void;
     setEditorRef: (editor: MonacoEditor | undefined) => void;
     showContext?: boolean;
+    showPinnedAgent?: boolean;
     labelProvider: LabelProvider;
 }
 
@@ -319,11 +336,42 @@ const ChatInput: React.FunctionComponent<ChatInputProperties> = (props: ChatInpu
         }
     };
 
-    const leftOptions = props.showContext ? [{
-        title: 'Attach elements to context',
-        handler: () => { /* TODO */ },
-        className: 'codicon-add'
-    }] : [];
+    const handlePin = () => {
+        if (editorRef.current) {
+            editorRef.current.getControl().getModel()?.applyEdits([{
+                range: {
+                    startLineNumber: 1,
+                    startColumn: 1,
+                    endLineNumber: 1,
+                    endColumn: 1
+                },
+                text: '@ ',
+            }]);
+            editorRef.current.getControl().setPosition({ lineNumber: 1, column: 2 });
+            editorRef.current.getControl().getAction('editor.action.triggerSuggest')?.run();
+        }
+    };
+
+    const leftOptions = [
+        ...(props.showContext
+            ? [{
+                title: 'Attach elements to context',
+                handler: () => { /* TODO */ },
+                className: 'codicon-add'
+            }]
+            : []),
+        ...(props.showPinnedAgent
+            ? [{
+                title: props.pinnedAgent ? 'Unpin Agent' : 'Pin Agent',
+                handler: props.pinnedAgent ? props.onUnpin : handlePin,
+                className: 'at-icon',
+                text: {
+                    align: 'right',
+                    content: props.pinnedAgent && props.pinnedAgent.name
+                },
+            }]
+            : []),
+    ] as Option[];
 
     const rightOptions = inProgress
         ? [{
@@ -456,6 +504,10 @@ interface Option {
     handler: () => void;
     className: string;
     disabled?: boolean;
+    text?: {
+        align?: 'left' | 'right';
+        content: string;
+    };
 }
 
 const ChatInputOptions: React.FunctionComponent<ChatInputOptionsProps> = ({ leftOptions, rightOptions }) => (
@@ -464,20 +516,26 @@ const ChatInputOptions: React.FunctionComponent<ChatInputOptionsProps> = ({ left
             {leftOptions.map((option, index) => (
                 <span
                     key={index}
-                    className={`codicon ${option.className} option ${option.disabled ? 'disabled' : ''}`}
+                    className={`option ${option.disabled ? 'disabled' : ''} ${option.text?.align === 'right' ? 'reverse' : ''}`}
                     title={option.title}
                     onClick={option.handler}
-                />
+                >
+                    <span>{option.text?.content}</span>
+                    <span className={`codicon ${option.className}`} />
+                </span>
             ))}
         </div>
         <div className="theia-ChatInputOptions-right">
             {rightOptions.map((option, index) => (
                 <span
                     key={index}
-                    className={`codicon ${option.className} option ${option.disabled ? 'disabled' : ''}`}
+                    className={`option ${option.disabled ? 'disabled' : ''} ${option.text?.align === 'right' ? 'reverse' : ''}`}
                     title={option.title}
                     onClick={option.handler}
-                />
+                >
+                    <span>{option.text?.content}</span>
+                    <span className={`codicon ${option.className}`}/>
+                </span>
             ))}
         </div>
     </div>

diff --git a/packages/ai-chat-ui/src/browser/chat-view-widget.tsx b/packages/ai-chat-ui/src/browser/chat-view-widget.tsx
@@ -91,8 +91,10 @@ export class ChatViewWidget extends BaseWidget implements ExtractableWidget, Sta
         this.chatSession = this.chatService.createSession();
 
         this.inputWidget.onQuery = this.onQuery.bind(this);
+        this.inputWidget.onUnpin = this.onUnpin.bind(this);
         this.inputWidget.onCancel = this.onCancel.bind(this);
         this.inputWidget.chatModel = this.chatSession.model;
+        this.inputWidget.pinnedAgent = this.chatSession.pinnedAgent;
         this.inputWidget.onDeleteChangeSet = this.onDeleteChangeSet.bind(this);
         this.inputWidget.onDeleteChangeSetElement = this.onDeleteChangeSetElement.bind(this);
         this.treeWidget.trackChatModel(this.chatSession.model);
@@ -117,6 +119,7 @@ export class ChatViewWidget extends BaseWidget implements ExtractableWidget, Sta
                     this.chatSession = session;
                     this.treeWidget.trackChatModel(this.chatSession.model);
                     this.inputWidget.chatModel = this.chatSession.model;
+                    this.inputWidget.pinnedAgent = this.chatSession.pinnedAgent;
                     if (event.focus) {
                         this.show();
                     }
@@ -169,6 +172,8 @@ export class ChatViewWidget extends BaseWidget implements ExtractableWidget, Sta
             if (responseModel.isError) {
                 this.messageService.error(responseModel.errorObject?.message ?? 'An error occurred during chat service invocation.');
             }
+        }).finally(() => {
+            this.inputWidget.pinnedAgent = this.chatSession.pinnedAgent;
         });
         if (!requestProgress) {
             this.messageService.error(`Was not able to send request "${chatRequest.text}" to session ${this.chatSession.id}`);
@@ -177,6 +182,11 @@ export class ChatViewWidget extends BaseWidget implements ExtractableWidget, Sta
         // Tree Widget currently tracks the ChatModel itself. Therefore no notification necessary.
     }
 
+    protected onUnpin(): void {
+        this.chatSession.pinnedAgent = undefined;
+        this.inputWidget.pinnedAgent = this.chatSession.pinnedAgent;
+    }
+
     protected onCancel(requestModel: ChatRequestModel): void {
         this.chatService.cancelRequest(requestModel.session.id, requestModel.id);
     }

diff --git a/packages/ai-chat-ui/src/browser/style/index.css b/packages/ai-chat-ui/src/browser/style/index.css
@@ -392,10 +392,13 @@ div:last-child>.theia-ChatNode {
 }
 
 .theia-ChatInputOptions .option {
-  width: 21px;
+  min-width: 21px;
   height: 21px;
   padding: 2px;
-  display: inline-block;
+  display: flex;
+  justify-content: space-between;
+  align-items: center;
+  gap: 2px;
   box-sizing: border-box;
   user-select: none;
   background-repeat: no-repeat;
@@ -416,6 +419,10 @@ div:last-child>.theia-ChatNode {
   background-color: var(--theia-toolbar-hoverBackground);
 }
 
+.theia-ChatInputOptions .reverse {
+    flex-direction: row-reverse;
+}
+
 .theia-CodePartRenderer-root {
   display: flex;
   flex-direction: column;

diff --git a/packages/ai-chat/src/browser/ai-chat-frontend-module.ts b/packages/ai-chat/src/browser/ai-chat-frontend-module.ts
@@ -26,6 +26,7 @@ import {
     ChatRequestParserImpl,
     ChatService,
     ToolCallChatResponseContentFactory,
+    PinChatAgent
 } from '../common';
 import { ChatAgentsVariableContribution } from '../common/chat-agents-variable-contribution';
 import { CustomChatAgent } from '../common/custom-chat-agent';
@@ -47,6 +48,7 @@ export default new ContainerModule(bind => {
 
     bind(ChatAgentServiceImpl).toSelf().inSingletonScope();
     bind(ChatAgentService).toService(ChatAgentServiceImpl);
+    bind(PinChatAgent).toConstantValue(true);
 
     bindContributionProvider(bind, ResponseContentMatcherProvider);
     bind(DefaultResponseContentMatcherProvider).toSelf().inSingletonScope();

diff --git a/packages/ai-chat/src/browser/ai-chat-preferences.ts b/packages/ai-chat/src/browser/ai-chat-preferences.ts
@@ -18,6 +18,7 @@ import { AI_CORE_PREFERENCES_TITLE } from '@theia/ai-core/lib/browser/ai-core-pr
 import { PreferenceSchema } from '@theia/core/lib/browser/preferences/preference-contribution';
 
 export const DEFAULT_CHAT_AGENT_PREF = 'ai-features.chat.defaultChatAgent';
+export const PIN_CHAT_AGENT_PREF = 'ai-features.chat.pinChatAgent';
 
 export const aiChatPreferences: PreferenceSchema = {
     type: 'object',
@@ -27,6 +28,14 @@ export const aiChatPreferences: PreferenceSchema = {
             description: 'Optional: <agent-name> of the Chat Agent that shall be invoked, if no agent is explicitly mentioned with @<agent-name> in the user query.\
             If no Default Agent is configured, Theia´s defaults will be applied.',
             title: AI_CORE_PREFERENCES_TITLE,
+        },
+        [PIN_CHAT_AGENT_PREF]: {
+            type: 'boolean',
+            description: 'Enable agent pinning to automatically keep a mentioned chat agent active across prompts, reducing the need for repeated mentions.\
+            \n\
+            You can manually unpin or switch agents anytime.',
+            default: true,
+            title: AI_CORE_PREFERENCES_TITLE,
         }
     }
 };
diff --git a/packages/ai-chat/src/browser/frontend-chat-service.ts b/packages/ai-chat/src/browser/frontend-chat-service.ts
@@ -15,9 +15,9 @@
 // *****************************************************************************
 
 import { inject, injectable } from '@theia/core/shared/inversify';
-import { ChatAgent, ChatServiceImpl, ParsedChatRequest } from '../common';
+import { ChatAgent, ChatServiceImpl, ChatSession, ParsedChatRequest } from '../common';
 import { PreferenceService } from '@theia/core/lib/browser';
-import { DEFAULT_CHAT_AGENT_PREF } from './ai-chat-preferences';
+import { DEFAULT_CHAT_AGENT_PREF, PIN_CHAT_AGENT_PREF } from './ai-chat-preferences';
 
 /**
  * Customizes the ChatServiceImpl to consider preference based default chat agent
@@ -28,15 +28,28 @@ export class FrontendChatServiceImpl extends ChatServiceImpl {
     @inject(PreferenceService)
     protected preferenceService: PreferenceService;
 
-    protected override getAgent(parsedRequest: ParsedChatRequest): ChatAgent | undefined {
+    protected override getAgent(parsedRequest: ParsedChatRequest, session: ChatSession): ChatAgent | undefined {
+        let agent = this.initialAgentSelection(parsedRequest);
+        if (!this.preferenceService.get<boolean>(PIN_CHAT_AGENT_PREF)) {
+            return agent;
+        }
+        if (!session.pinnedAgent && agent && agent.id !== this.defaultChatAgentId?.id) {
+            session.pinnedAgent = agent;
+        } else if (session.pinnedAgent && this.getMentionedAgent(parsedRequest) === undefined) {
+            agent = session.pinnedAgent;
+        }
+        return agent;
+    }
+
+    protected override initialAgentSelection(parsedRequest: ParsedChatRequest): ChatAgent | undefined {
         const agentPart = this.getMentionedAgent(parsedRequest);
         if (!agentPart) {
             const configuredDefaultChatAgent = this.getConfiguredDefaultChatAgent();
             if (configuredDefaultChatAgent) {
                 return configuredDefaultChatAgent;
             }
         }
-        return super.getAgent(parsedRequest);
+        return super.initialAgentSelection(parsedRequest);
     }
 
     protected getConfiguredDefaultChatAgent(): ChatAgent | undefined {