refactor(ui): event handling refactor (#805)

Author: dazy-ds

packages/ragbits-chat/CHANGELOG.md

@@ -1,6 +1,7 @@
 # CHANGELOG
 
 ## Unreleased
+- Refactor chat handlers in the UI to use registry (#805)
 - Add auth token storage and automatic logout on 401 (#802)
 - Improve user settings storage when history is disabled (#799)
 - Remove redundant test for `/api/config` endpoint (#795)

scripts/generate_typescript_from_json_schema.py

@@ -203,18 +203,11 @@ def _generate_chat_response_union_type() -> str:
         ("MessageUsageChatResponse", "usage", "Record<string, MessageUsage>"),
     ]
 
-    exported_response_interfaces = [
+    internal_response_interfaces = [
         ("ChunkedChatResponse", "chunked_content", "ChunkedContent"),
     ]
 
-    for interface_name, response_type, content_type in response_interfaces:
-        lines.append(f"interface {interface_name} {{")
-        lines.append(f"    type: '{response_type}'")
-        lines.append(f"    content: {content_type}")
-        lines.append("}")
-        lines.append("")
-
-    for interface_name, response_type, content_type in exported_response_interfaces:
+    for interface_name, response_type, content_type in [*response_interfaces, *internal_response_interfaces]:
         lines.append(f"export interface {interface_name} {{")
         lines.append(f"    type: '{response_type}'")
         lines.append(f"    content: {content_type}")

typescript/@ragbits/api-client/src/autogen.types.ts

@@ -414,52 +414,52 @@ export interface User {
 /**
  * Specific chat response types
  */
-interface TextChatResponse {
+export interface TextChatResponse {
     type: 'text'
     content: string
 }
 
-interface ReferenceChatResponse {
+export interface ReferenceChatResponse {
     type: 'reference'
     content: Reference
 }
 
-interface MessageIdChatResponse {
+export interface MessageIdChatResponse {
     type: 'message_id'
     content: string
 }
 
-interface ConversationIdChatResponse {
+export interface ConversationIdChatResponse {
     type: 'conversation_id'
     content: string
 }
 
-interface StateUpdateChatResponse {
+export interface StateUpdateChatResponse {
     type: 'state_update'
     content: ServerState
 }
 
-interface LiveUpdateChatResponse {
+export interface LiveUpdateChatResponse {
     type: 'live_update'
     content: LiveUpdate
 }
 
-interface FollowupMessagesChatResponse {
+export interface FollowupMessagesChatResponse {
     type: 'followup_messages'
     content: string[]
 }
 
-interface ImageChatResponse {
+export interface ImageChatResponse {
     type: 'image'
     content: Image
 }
 
-interface ClearMessageResponse {
+export interface ClearMessageResponse {
     type: 'clear_message'
     content: never
 }
 
-interface MessageUsageChatResponse {
+export interface MessageUsageChatResponse {
     type: 'usage'
     content: Record<string, MessageUsage>
 }

typescript/ui/src/core/stores/HistoryStore/eventHandlers/eventHandlerRegistry.ts

@@ -0,0 +1,103 @@
+import { ChatResponse, ChatResponseType } from "@ragbits/api-client-react";
+import { Conversation, HistoryStore } from "../../../../types/history";
+import {
+  handleAfterConversationId,
+  handleConversationId,
+  handleFollowupMessages,
+  handleStateUpdate,
+} from "./nonMessageHandlers";
+import {
+  handleClearMessage,
+  handleImage,
+  handleLiveUpdate,
+  handleMessageId,
+  handleReference,
+  handleText,
+  handleUsage,
+} from "./messageHandlers";
+
+type BaseHandlerContext = {
+  conversationIdRef: { current: string };
+  messageId: string;
+};
+
+export type PrimaryHandler<
+  T extends ChatResponse,
+  Ctx extends object = object,
+> = (response: T, draft: Conversation, ctx: BaseHandlerContext) => Ctx | void;
+
+export type AfterHandler<
+  T extends ChatResponse,
+  Ctx extends object = object,
+> = (response: T, draft: HistoryStore, ctx: BaseHandlerContext & Ctx) => void;
+
+export type HandlerEntry<
+  T extends ChatResponse,
+  Ctx extends object = object,
+> = {
+  handle: PrimaryHandler<T, Ctx>;
+  after?: AfterHandler<T, Ctx>;
+};
+
+class ChatResponseHandlerRegistry {
+  private handlers = new Map<ChatResponseType, unknown>();
+
+  register<
+    K extends ChatResponseType,
+    R extends Extract<ChatResponse, { type: K }>,
+    Ctx extends object = object,
+  >(type: K, entry: HandlerEntry<R, Ctx>) {
+    if (this.handlers.has(type)) {
+      console.warn(
+        `Handler for ${String(type)} already registered - overwriting.`,
+      );
+    }
+    this.handlers.set(type, entry);
+  }
+
+  get<
+    K extends ChatResponseType,
+    R extends Extract<ChatResponse, { type: K }>,
+    Ctx extends object = object,
+  >(type: K): HandlerEntry<R, Ctx> {
+    const raw = this.handlers.get(type);
+    if (!raw) {
+      throw new Error(`No handler registered for type: ${String(type)}`);
+    }
+    return raw as HandlerEntry<R, Ctx>;
+  }
+}
+
+export const ChatHandlerRegistry = new ChatResponseHandlerRegistry();
+ChatHandlerRegistry.register(ChatResponseType.StateUpdate, {
+  handle: handleStateUpdate,
+});
+ChatHandlerRegistry.register(ChatResponseType.ConversationId, {
+  handle: handleConversationId,
+  after: handleAfterConversationId,
+});
+ChatHandlerRegistry.register(ChatResponseType.FollowupMessages, {
+  handle: handleFollowupMessages,
+});
+
+ChatHandlerRegistry.register(ChatResponseType.Text, {
+  handle: handleText,
+});
+ChatHandlerRegistry.register(ChatResponseType.Reference, {
+  handle: handleReference,
+});
+ChatHandlerRegistry.register(ChatResponseType.MessageId, {
+  handle: handleMessageId,
+});
+ChatHandlerRegistry.register(ChatResponseType.LiveUpdate, {
+  handle: handleLiveUpdate,
+});
+ChatHandlerRegistry.register(ChatResponseType.Image, {
+  handle: handleImage,
+});
+ChatHandlerRegistry.register(ChatResponseType.ClearMessage, {
+  handle: handleClearMessage,
+});
+ChatHandlerRegistry.register(ChatResponseType.Usage, {
+  handle: handleUsage,
+});

typescript/ui/src/core/stores/HistoryStore/eventHandlers/messageHandlers.ts

@@ -0,0 +1,99 @@
+import {
+  ClearMessageResponse,
+  ImageChatResponse,
+  LiveUpdateChatResponse,
+  LiveUpdateType,
+  MessageIdChatResponse,
+  MessageUsageChatResponse,
+  ReferenceChatResponse,
+  TextChatResponse,
+} from "@ragbits/api-client-react";
+import { PrimaryHandler } from "./eventHandlerRegistry";
+import { produce } from "immer";
+
+export const handleText: PrimaryHandler<TextChatResponse> = (
+  response,
+  draft,
+  ctx,
+) => {
+  const message = draft.history[ctx.messageId];
+  message.content += response.content;
+};
+
+export const handleReference: PrimaryHandler<ReferenceChatResponse> = (
+  response,
+  draft,
+  ctx,
+) => {
+  const message = draft.history[ctx.messageId];
+  message.references = [...(message.references ?? []), response.content];
+};
+
+export const handleMessageId: PrimaryHandler<MessageIdChatResponse> = (
+  response,
+  draft,
+  ctx,
+) => {
+  const message = draft.history[ctx.messageId];
+  message.serverId = response.content;
+};
+
+export const handleLiveUpdate: PrimaryHandler<LiveUpdateChatResponse> = (
+  response,
+  draft,
+  ctx,
+) => {
+  const message = draft.history[ctx.messageId];
+
+  const { update_id, content, type } = response.content;
+  const liveUpdates = produce(message.liveUpdates ?? {}, (draft) => {
+    if (type === LiveUpdateType.Start && update_id in draft) {
+      console.error(
+        `Got duplicate start event for update_id: ${update_id}. Ignoring the event.`,
+      );
+    }
+
+    draft[update_id] = content;
+  });
+  message.liveUpdates = liveUpdates;
+};
+
+export const handleImage: PrimaryHandler<ImageChatResponse> = (
+  response,
+  draft,
+  ctx,
+) => {
+  const message = draft.history[ctx.messageId];
+  const image = response.content;
+  message.images = produce(message.images ?? {}, (draft) => {
+    if (draft[image.id]) {
+      console.error(
+        `Got duplicate image event for image_id: ${image.id}. Ignoring the event.`,
+      );
+    }
+
+    draft[image.id] = image.url;
+  });
+};
+
+export const handleClearMessage: PrimaryHandler<ClearMessageResponse> = (
+  _,
+  draft,
+  ctx,
+) => {
+  const message = draft.history[ctx.messageId];
+  draft.history[ctx.messageId] = {
+    id: message.id,
+    role: message.role,
+    content: "",
+  };
+};
+
+export const handleUsage: PrimaryHandler<MessageUsageChatResponse> = (
+  response,
+  draft,
+  ctx,
+) => {
+  const message = draft.history[ctx.messageId];
+  message.usage = response.content;
+};

typescript/ui/src/core/stores/HistoryStore/eventHandlers/nonMessageHandlers.ts

@@ -0,0 +1,47 @@
+import {
+  ConversationIdChatResponse,
+  FollowupMessagesChatResponse,
+  StateUpdateChatResponse,
+} from "@ragbits/api-client-react";
+import { AfterHandler, PrimaryHandler } from "./eventHandlerRegistry";
+
+export const handleStateUpdate: PrimaryHandler<StateUpdateChatResponse> = (
+  { content },
+  draft,
+) => {
+  draft.serverState = content;
+};
+
+export const handleConversationId: PrimaryHandler<
+  ConversationIdChatResponse,
+  { originalConversationId: string }
+> = ({ content }, draft, ctx) => {
+  const originalConversationId = ctx.conversationIdRef.current;
+  draft.conversationId = content;
+  // Update the ref to propagate the change
+  ctx.conversationIdRef.current = content;
+  return { originalConversationId };
+};
+
+export const handleAfterConversationId: AfterHandler<
+  ConversationIdChatResponse,
+  { originalConversationId: string }
+> = (_, draft, { originalConversationId, conversationIdRef }) => {
+  const oldConversation = draft.conversations[originalConversationId];
+
+  if (!oldConversation) {
+    throw new Error("Received events for non-existent conversation");
+  }
+
+  draft.conversations[conversationIdRef.current] = oldConversation;
+  if (draft.currentConversation === originalConversationId) {
+    draft.currentConversation = conversationIdRef.current;
+  }
+  delete draft.conversations[originalConversationId];
+};
+
+export const handleFollowupMessages: PrimaryHandler<
+  FollowupMessagesChatResponse
+> = ({ content }, draft) => {
+  draft.followupMessages = content;
+};