feat: support maxInputTokens for agent (#4397)

* feat: support maxInputTokens for agent * fix: use formatLocalize * chore: optional maxInputTokens * fix: clear startIndex when dispose * fix: type error * feat: support noTools options
opensumi · Feb 24, 2025 · ebc1030 · ebc1030
1 parent aca7ff8
commit ebc1030
Show file tree

Hide file tree

Showing 9 changed files with 48 additions and 6 deletions.
diff --git a/packages/ai-native/src/browser/chat/chat-proxy.service.ts b/packages/ai-native/src/browser/chat/chat-proxy.service.ts
@@ -125,13 +125,13 @@ export class ChatProxyService extends Disposable {
           } else {
             apiKey = this.preferenceService.get<string>(AINativeSettingSectionsId.AnthropicApiKey, '');
           }
-
+          const MAX_INPUT_TOKENS = 30720;
           const stream = await this.aiBackService.requestStream(
             prompt,
             {
               requestId: request.requestId,
               sessionId: request.sessionId,
-              history: this.aiChatService.getHistoryMessages(),
+              history: this.aiChatService.getHistoryMessages(MAX_INPUT_TOKENS),
               clientId: this.applicationService.clientId,
               apiKey,
               model,

diff --git a/packages/ai-native/src/browser/chat/chat.api.service.ts b/packages/ai-native/src/browser/chat/chat.api.service.ts
@@ -65,8 +65,8 @@ export class ChatService extends Disposable {
     this._onChatMessageListLaunch.fire(list);
   }
 
-  public getHistoryMessages(): IHistoryChatMessage[] {
-    return this.chatInternalService.sessionModel?.history.getMessages() || [];
+  public getHistoryMessages(maxInputTokens?: number): IHistoryChatMessage[] {
+    return this.chatInternalService.sessionModel?.history.getMessages(maxInputTokens) || [];
   }
 
   public scrollToBottom(): void {

diff --git a/packages/ai-native/src/browser/chat/chat.module.less b/packages/ai-native/src/browser/chat/chat.module.less
@@ -279,6 +279,16 @@
   }
 }
 
+
+.chat_tips_container {
+  display: flex;
+  align-items: center;
+  justify-content: center;
+  color: var(--tab-inactiveForeground);
+  font-size: 12px;
+  margin-top: 16px;
+}
+
 .chat_history {
   width: calc(100% - 60px);
   color: var(--design-text-foreground);

diff --git a/packages/ai-native/src/browser/chat/chat.view.tsx b/packages/ai-native/src/browser/chat/chat.view.tsx
@@ -26,6 +26,7 @@ import {
   IChatComponent,
   IChatContent,
   MessageType,
+  formatLocalize,
   localize,
   uuid,
 } from '@opensumi/ide-core-common';
@@ -696,6 +697,13 @@ export const AIChatView = () => {
               dataSource={messageListData}
             />
           </div>
+          {msgHistoryManager.slicedMessageCount ? (
+            <div className={styles.chat_tips_text}>
+              <div className={styles.chat_tips_container}>
+                {formatLocalize('aiNative.chat.ai.assistant.limit.message', msgHistoryManager.slicedMessageCount)}
+              </div>
+            </div>
+          ) : null}
           <div className={styles.chat_input_wrap}>
             <ChatContext />
             <div className={styles.header_operate}>

diff --git a/packages/ai-native/src/browser/model/msg-history-manager.ts b/packages/ai-native/src/browser/model/msg-history-manager.ts
@@ -31,6 +31,7 @@ export class MsgHistoryManager extends Disposable {
   public clearMessages() {
     this.messageMap.clear();
     this.messageAdditionalMap.clear();
+    this.startIndex = 0;
   }
 
   private doAddMessage(message: IExcludeMessage): string {
@@ -50,10 +51,30 @@ export class MsgHistoryManager extends Disposable {
     return id;
   }
 
-  public getMessages(): IHistoryChatMessage[] {
+  private get messageList(): IHistoryChatMessage[] {
     return Array.from(this.messageMap.values()).sort((a, b) => a.order - b.order);
   }
 
+  private startIndex = 0;
+
+  private get totalTokens(): number {
+    const list = this.messageList.slice(this.startIndex);
+    return list.reduce((acc, msg) => acc + (msg.content.length || 0), 0) / 3;
+  }
+
+  public get slicedMessageCount(): number {
+    return this.startIndex;
+  }
+
+  public getMessages(maxTokens?: number): IHistoryChatMessage[] {
+    if (maxTokens && this.totalTokens > maxTokens) {
+      while (this.totalTokens > maxTokens) {
+        this.startIndex++;
+      }
+    }
+    return this.messageList.slice(this.startIndex);
+  }
+
   public addUserMessage(
     message: Required<Pick<IExcludeMessage, 'agentId' | 'agentCommand' | 'content' | 'relationId'>>,
   ): string {

diff --git a/packages/ai-native/src/node/base-language-model.ts b/packages/ai-native/src/node/base-language-model.ts
@@ -45,7 +45,7 @@ export abstract class BaseLanguageModel {
       throw new Error('clientId is required');
     }
     const registry = this.toolInvocationRegistryManager.getRegistry(clientId);
-    const allFunctions = registry.getAllFunctions();
+    const allFunctions = options.noTool ? [] : registry.getAllFunctions();
     return this.handleStreamingRequest(
       provider,
       request,

diff --git a/packages/core-common/src/types/ai-native/index.ts b/packages/core-common/src/types/ai-native/index.ts
@@ -182,6 +182,7 @@ export interface IAIBackServiceOption {
   topP?: number;
   topK?: number;
   providerOptions?: any;
+  noTool?: boolean;
 }
 
 /**

diff --git a/packages/i18n/src/common/en-US.lang.ts b/packages/i18n/src/common/en-US.lang.ts
@@ -1491,6 +1491,7 @@ export const localizationBundle = {
     'aiNative.operate.chatHistory.delete': 'Delete',
 
     'aiNative.chat.welcome.loading.text': 'Initializing...',
+    'aiNative.chat.ai.assistant.limit.message': '{0} earliest messages are dropped due to the input token limit',
 
     'preference.ai.native.inlineChat.title': 'Inline Chat',
     'preference.ai.native.chat.title': 'Chat',

diff --git a/packages/i18n/src/common/zh-CN.lang.ts b/packages/i18n/src/common/zh-CN.lang.ts
@@ -1259,6 +1259,7 @@ export const localizationBundle = {
     'aiNative.operate.chatHistory.delete': '删除',
 
     'aiNative.chat.welcome.loading.text': '初始化中...',
+    'aiNative.chat.ai.assistant.limit.message': '{0} 条最早的消息因输入 Tokens 限制而被丢弃',
 
     'preference.ai.native.inlineChat.title': 'Inline Chat',
     'preference.ai.native.chat.title': 'Chat',