refactor: implement external stream protocol

backend/app/rag/chat.py

@@ -508,12 +508,22 @@ def _external_chat(self) -> Generator[ChatEvent | str, None, None]:
         )
 
         stream_chat_api_url = self.chat_engine_config.external_engine_config.stream_chat_api_url
-        logger.debug(f"Chatting with external chat engine (api_url: {stream_chat_api_url}) to answer for user question: {self.user_question}")
+        stream_chat_type = self.chat_engine_config.external_engine_config.type
+        logger.debug(f"Chatting with external chat engine (api_url: {stream_chat_api_url}, type: {stream_chat_type}) to answer for user question: {self.user_question}")
         chat_params = {
             "goal": self.user_question
         }
         res = requests.post(stream_chat_api_url, json=chat_params, stream=True)
 
+        yield ChatEvent(
+            event_type=ChatEventType.MESSAGE_ANNOTATIONS_PART,
+            payload=ChatStreamMessagePayload(
+                state=ChatMessageSate.EXTERNAL_STREAM_START,
+                display="Using external engine",
+                context={ 'type': stream_chat_type },
+            ),
+        )
+
         # Notice: External type chat engine doesn't support non-streaming mode for now.
         response_text = ""
         for line in res.iter_lines():
@@ -527,6 +537,13 @@ def _external_chat(self) -> Generator[ChatEvent | str, None, None]:
 
             yield line + b'\n'
 
+        yield ChatEvent(
+            event_type=ChatEventType.MESSAGE_ANNOTATIONS_PART,
+            payload=ChatStreamMessagePayload(
+                state=ChatMessageSate.FINISHED,
+            ),
+        )
+
         db_assistant_message.content = response_text
         db_assistant_message.updated_at = datetime.now(UTC)
         db_assistant_message.finished_at = datetime.now(UTC)

backend/app/rag/chat_config.py

@@ -75,6 +75,7 @@ class KnowledgeGraphOption(BaseModel):
 
 class ExternalChatEngine(BaseModel):
     stream_chat_api_url: str = None
+    type: str = 'StackVM'
 
 
 class ChatEngineConfig(BaseModel):

backend/app/rag/types.py

@@ -59,3 +59,5 @@ class ChatMessageSate(int, enum.Enum):
     SEARCH_RELATED_DOCUMENTS = 4
     GENERATE_ANSWER = 5
     FINISHED = 9
+    # See https://github.com/pingcap/tidb.ai/issues/345
+    EXTERNAL_STREAM_START = 10

frontend/app/src/components/chat/chat-controller.ts

@@ -1,5 +1,5 @@
 import { chat, type Chat, type ChatMessage, type PostChatParams } from '@/api/chats';
-import { BaseChatMessageController, ChatMessageController, LegacyChatMessageController, type OngoingState, StackVMChatMessageController } from '@/components/chat/chat-message-controller';
+import { BaseChatMessageController, ChatMessageController, ExternalChatMessageController, LegacyChatMessageController, type OngoingState } from '@/components/chat/chat-message-controller';
 import { AppChatStreamState, type BaseAnnotation, chatDataPartSchema, fixChatInitialData, type StackVMState } from '@/components/chat/chat-stream-state';
 import type { GtagFn } from '@/components/gtag-provider';
 import { getErrorMessage } from '@/lib/errors';
@@ -34,7 +34,7 @@ export interface ChatControllerEventsMap<State = AppChatStreamState, Annotation
 export class ChatController<State extends AppChatStreamState = AppChatStreamState, Annotation extends BaseAnnotation<State> = BaseAnnotation<State>> extends EventEmitter<ChatControllerEventsMap<State, Annotation>> {
   public chat: Chat | undefined;
 
-  private _messages: Map<number, ChatMessageController | StackVMChatMessageController> = new Map();
+  private _messages: Map<number, ChatMessageController | ExternalChatMessageController> = new Map();
 
   private _postParams: Omit<PostChatParams, 'chat_id'> | undefined = undefined;
   private _postError: unknown = undefined;
@@ -140,7 +140,7 @@ export class ChatController<State extends AppChatStreamState = AppChatStreamStat
     }
   }
 
-  get messages (): (ChatMessageController | StackVMChatMessageController)[] {
+  get messages (): (ChatMessageController | ExternalChatMessageController)[] {
     return Array.from(this._messages.values()).sort((a, b) => a.message.ordinal - b.message.ordinal);
   }
 
@@ -163,7 +163,7 @@ export class ChatController<State extends AppChatStreamState = AppChatStreamStat
     this._postInitialized = false;
     this.emit('post', params);
 
-    let ongoingMessageController: ChatMessageController | StackVMChatMessageController | undefined = undefined;
+    let ongoingMessageController: ChatMessageController | ExternalChatMessageController | undefined = undefined;
 
     try {
       const streamParts = chat({
@@ -218,7 +218,7 @@ export class ChatController<State extends AppChatStreamState = AppChatStreamStat
     }
   }
 
-  _processPart (ongoingMessageController: ChatMessageController | StackVMChatMessageController | undefined, part: ReturnType<StreamPart<any, any, any>['parse']>) {
+  _processPart (ongoingMessageController: ChatMessageController | ExternalChatMessageController | undefined, part: ReturnType<StreamPart<any, any, any>['parse']>) {
     switch (part.type) {
       case 'data':
         // Data part contains chat and chat_message info from server. will be sent twice (beginning and finished).
@@ -240,19 +240,20 @@ export class ChatController<State extends AppChatStreamState = AppChatStreamStat
         break;
       case 'tool_result':
         this._processToolResultPart(ongoingMessageController, part);
-        break
+        break;
       case 'finish_step':
-
+        // We don't need this yet
+        break;
       case 'finish_message':
-        // TODO
+        this._processMessageFinishPart(ongoingMessageController, part);
         break;
       default:
         console.warn('unsupported stream part', part);
     }
     return ongoingMessageController;
   }
 
-  private _processDataPart (ongoingMessageController: ChatMessageController | StackVMChatMessageController | undefined, part: ReturnType<StreamPart<any, 'data', JSONValue[]>['parse']>): ChatMessageController | StackVMChatMessageController {
+  private _processDataPart (ongoingMessageController: ChatMessageController | ExternalChatMessageController | undefined, part: ReturnType<StreamPart<any, 'data', JSONValue[]>['parse']>): ChatMessageController | ExternalChatMessageController {
     const { chat, user_message, assistant_message } = chatDataPartSchema.parse(fixChatInitialData(part.value[0]));
     this.updateChat(chat);
     this.upsertMessage(user_message);
@@ -267,34 +268,41 @@ export class ChatController<State extends AppChatStreamState = AppChatStreamStat
     return ongoingMessageController;
   }
 
-  private _processMessageAnnotationPart (ongoingMessageController: ChatMessageController | StackVMChatMessageController | undefined, part: ReturnType<StreamPart<any, 'message_annotations', JSONValue[]>['parse']>) {
+  private _processMessageAnnotationPart (ongoingMessageController: ChatMessageController | ExternalChatMessageController | undefined, part: ReturnType<StreamPart<any, 'message_annotations', JSONValue[]>['parse']>) {
     assertNonNull(ongoingMessageController, 'Cannot handle chat stream part: no ongoingMessageController', part);
     const annotation = ongoingMessageController.parseAnnotation(part.value[0]);
-    ongoingMessageController.applyStreamAnnotation(annotation as never);
+    if (annotation !== null) {
+      ongoingMessageController.applyStreamAnnotation(annotation as never);
+    }
   }
 
-  private _processTextPart (ongoingMessageController: ChatMessageController | StackVMChatMessageController | undefined, part: ReturnType<StreamPart<any, 'text', string>['parse']>) {
+  private _processTextPart (ongoingMessageController: ChatMessageController | ExternalChatMessageController | undefined, part: ReturnType<StreamPart<any, 'text', string>['parse']>) {
     if (part.value) { // ignore leading empty chunks.
       assertNonNull(ongoingMessageController, 'Cannot handle chat stream part: no ongoingMessageController', part);
       ongoingMessageController.applyDelta(part.value);
     }
   }
 
-  private _processErrorPart (ongoingMessageController: ChatMessageController | StackVMChatMessageController | undefined, part: ReturnType<StreamPart<any, 'error', string>['parse']>) {
+  private _processErrorPart (ongoingMessageController: ChatMessageController | ExternalChatMessageController | undefined, part: ReturnType<StreamPart<any, 'error', string>['parse']>) {
     assertNonNull(ongoingMessageController, 'Cannot handle chat stream part: no ongoingMessageController', part);
     ongoingMessageController.applyError(part.value);
   }
 
-  private _processToolCallPart (ongoingMessageController: ChatMessageController | StackVMChatMessageController | undefined, part: ReturnType<StreamPart<any, 'tool_call', { toolCallId: string, toolName: string, args: any }>['parse']>) {
+  private _processToolCallPart (ongoingMessageController: ChatMessageController | ExternalChatMessageController | undefined, part: ReturnType<StreamPart<any, 'tool_call', { toolCallId: string, toolName: string, args: any }>['parse']>) {
     assertNonNull(ongoingMessageController, 'Cannot handle chat stream part: no ongoingMessageController', part);
     ongoingMessageController.applyToolCall(part.value);
   }
 
-  private _processToolResultPart (ongoingMessageController: ChatMessageController | StackVMChatMessageController | undefined, part: ReturnType<StreamPart<any, 'tool_call', { toolCallId: string, result: any }>['parse']>) {
+  private _processToolResultPart (ongoingMessageController: ChatMessageController | ExternalChatMessageController | undefined, part: ReturnType<StreamPart<any, 'tool_call', { toolCallId: string, result: any }>['parse']>) {
     assertNonNull(ongoingMessageController, 'Cannot handle chat stream part: no ongoingMessageController', part);
     ongoingMessageController.applyToolResult(part.value);
   }
 
+  private _processMessageFinishPart (ongoingMessageController: ChatMessageController | ExternalChatMessageController | undefined, part: ReturnType<StreamPart<any, 'finish_message', unknown>['parse']>) {
+    assertNonNull(ongoingMessageController, 'Cannot handle chat stream part: no ongoingMessageController', part);
+    ongoingMessageController.applyFinishMessage(part.value);
+  }
+
   private createMessage (message: ChatMessage, initialOngoingState?: true) {
     if (!this.chat?.engine_options) {
       throw new Error('Unable to decide which chat engine used.');
@@ -315,7 +323,7 @@ export class ChatController<State extends AppChatStreamState = AppChatStreamStat
   }
 
   private createStackVMMessage (message: ChatMessage, initialOngoingState?: true | OngoingState<StackVMState>) {
-    const controller = new StackVMChatMessageController(message, initialOngoingState);
+    const controller = new ExternalChatMessageController(message, initialOngoingState);
     this._messages.set(message.id, controller);
     this.emit('message-loaded', controller as any);
     return controller;

frontend/app/src/components/chat/chat-message-controller.ts

@@ -24,15 +24,17 @@ export interface ChatMessageControllerEventsMap<State = AppChatStreamState> {
 
   'stream-tool-call': [id: string, name: string, args: any];
   'stream-tool-result': [id: string, result: any];
+
+  'stream-state-change': [];
 }
 
 export abstract class BaseChatMessageController<
   State,
-  Annotation extends BaseAnnotation<State>
+  Annotation extends BaseAnnotation<any>,
 > extends EventEmitter<ChatMessageControllerEventsMap<State>> {
   protected _message: ChatMessage;
   protected _ongoing: OngoingState<State> | undefined;
-  protected _ongoingHistory: OngoingStateHistoryItem<State>[] | undefined;
+  protected _ongoingHistory: OngoingStateHistoryItem<Annotation['state']>[] | undefined;
   public readonly role: ChatMessageRole;
   public readonly id: number;
 
@@ -137,6 +139,9 @@ export abstract class BaseChatMessageController<
     this.emit('stream-tool-result', toolCallId, result);
   }
 
+  applyFinishMessage (_: unknown) {
+  }
+
   finish () {
     this._ongoing = undefined;
     this.emit('stream-finished', this._message);
@@ -155,7 +160,7 @@ export abstract class BaseChatMessageController<
     return this._ongoingHistory;
   }
 
-  abstract parseAnnotation (raw: unknown): Annotation;
+  abstract parseAnnotation (raw: unknown): Annotation | null;
 
   abstract createInitialOngoingState (): OngoingState<State>;
 
@@ -164,7 +169,7 @@ export abstract class BaseChatMessageController<
   protected abstract _polishMessage (message: ChatMessage, ongoing: OngoingState<State>, annotation: Annotation): ChatMessage
 }
 
-export type ChatMessageController = LegacyChatMessageController | StackVMChatMessageController;
+export type ChatMessageController = LegacyChatMessageController | ExternalChatMessageController;
 export type ChatMessageControllerAnnotationState<C extends ChatMessageController> = C extends BaseChatMessageController<infer State, any> ? State : never;
 
 export class LegacyChatMessageController extends BaseChatMessageController<AppChatStreamState, ChatMessageAnnotation> {
@@ -209,8 +214,9 @@ export class LegacyChatMessageController extends BaseChatMessageController<AppCh
   }
 }
 
-export class StackVMChatMessageController extends BaseChatMessageController<StackVMState, StackVMStateAnnotation> {
+export class ExternalChatMessageController extends BaseChatMessageController<StackVMState, StackVMStateAnnotation> {
   readonly version = 'StackVM';
+  private externalStreamStart = false;
 
   applyToolCall (payload: { toolCallId: string; toolName: string; args: any }) {
     super.applyToolCall(payload);
@@ -243,14 +249,34 @@ export class StackVMChatMessageController extends BaseChatMessageController<Stac
     }
   }
 
-  parseAnnotation (raw: unknown): StackVMStateAnnotation {
-    const { state: rawState, plan_id } = raw as { state: StackVM.State, plan_id: string };
-    const state = StackVM.model.parseState(rawState);
+  applyFinishMessage (_: unknown) {
+    this.externalStreamStart = false;
+    this.emit('stream-state-change');
+  }
 
-    return {
-      state: { plan_id, state, toolCalls: [] },
-      display: '[deprecated]',
-    };
+  parseAnnotation (raw: unknown): StackVMStateAnnotation | null {
+    if (this.externalStreamStart) {
+      const { state: rawState, plan_id } = raw as { state: StackVM.State, plan_id: string };
+      const state = StackVM.model.parseState(rawState);
+
+      return {
+        state: { plan_id, state, toolCalls: [] },
+        display: '[deprecated]',
+      };
+    } else {
+      const { state, display } = raw as ChatMessageAnnotation;
+
+      if (state === AppChatStreamState.EXTERNAL_STREAM_START) {
+        this.externalStreamStart = true;
+        this.emit('stream-state-change');
+        return null;
+      } else {
+        return {
+          state: null,
+          display: display ?? state,
+        };
+      }
+    }
   }
 
   createInitialOngoingState (): OngoingState<StackVMState> {
@@ -306,8 +332,8 @@ export class StackVMChatMessageController extends BaseChatMessageController<Stac
   _polishMessage (message: ChatMessage, ongoing: OngoingState<StackVMState>, annotation: StackVMStateAnnotation): ChatMessage {
     // FIX Initial state
     // First step reasoning finished with PC = 1, we need to insert a PC = 0 state first.
-    if (annotation.state.state.program_counter === 1) {
-      if (!this._ongoingHistory?.find(item => item.state.state.state.program_counter === 0)) {
+    if (annotation.state && annotation.state.state.program_counter === 1) {
+      if (!this._ongoingHistory?.find(item => item.state.state?.state.program_counter === 0)) {
         const lastState = {
           state: {
             state: {
@@ -329,4 +355,4 @@ export class StackVMChatMessageController extends BaseChatMessageController<Stac
 
     return message;
   }
-}
+}

frontend/app/src/components/chat/chat-stream-state.ts

@@ -6,6 +6,7 @@
  *     SEARCH_RELATED_DOCUMENTS = 4
  *     GENERATE_ANSWER = 5
  *     FINISHED = 9
+ *     EXTERNAL_STREAM_START = 10
  */
 import { Chat, ChatMessage, chatMessageSchema, ChatMessageSource, chatSchema } from '@/api/chats';
 import { StackVM } from '@/lib/stackvm';
@@ -26,6 +27,7 @@ export const enum AppChatStreamState {
   RERANKING = 'RERANKING',
   GENERATE_ANSWER = 'GENERATE_ANSWER',
   FINISHED = 'FINISHED',
+  EXTERNAL_STREAM_START = 'EXTERNAL_STREAM_START',
   FAILED = 'FAILED',
   UNKNOWN = 'UNKNOWN',
 }
@@ -55,13 +57,18 @@ export interface RefineQuestionAnnotation extends BaseAnnotation<AppChatStreamSt
   message?: string;
 }
 
+export interface ExternalStreamStartAnnotation extends BaseAnnotation<AppChatStreamState.EXTERNAL_STREAM_START> {
+  context: { type: 'StackVM' }
+}
+
 export type ChatMessageAnnotation =
-  BaseAnnotation<Exclude<AppChatStreamState, AppChatStreamState.TRACE | AppChatStreamState.SOURCE_NODES | AppChatStreamState.REFINE_QUESTION>>
+  BaseAnnotation<Exclude<AppChatStreamState, AppChatStreamState.TRACE | AppChatStreamState.SOURCE_NODES | AppChatStreamState.REFINE_QUESTION | AppChatStreamState.EXTERNAL_STREAM_START>>
   | TraceAnnotation
   | SourceNodesAnnotation
-  | RefineQuestionAnnotation;
+  | RefineQuestionAnnotation
+  | ExternalStreamStartAnnotation;
 
-export interface StackVMStateAnnotation extends BaseAnnotation<StackVMState> {
+export interface StackVMStateAnnotation extends BaseAnnotation<StackVMState | null> {
 }
 
 export type ChatInitialData = {