show avatar dialog instead of default (#4033)

show avatar dialog instead of default

- [x] New Feature (non-breaking change which adds functionality)

---------

Co-authored-by: Kevin Hu <[email protected]>

api/apps/canvas_app.py

@@ -23,6 +23,7 @@ from api.utils import get_uuid
 from api.utils.api_utils import get_json_result, server_error_response, validate_request, get_data_error_result
 from agent.canvas import Canvas
 from peewee import MySQLDatabase, PostgresqlDatabase
+from api.db.db_models import APIToken
 
 
 @manager.route('/templates', methods=['GET'])  # noqa: F821
@@ -85,6 +86,20 @@ def get(canvas_id):
         return get_data_error_result(message="canvas not found.")
     return get_json_result(data=c.to_dict())
 
+@manager.route('/getsse/<canvas_id>', methods=['GET'])  # type: ignore # noqa: F821
+def getsse(canvas_id):
+    token = request.headers.get('Authorization').split()
+    if len(token) != 2:
+        return get_data_error_result(message='Authorization is not valid!"')
+    token = token[1]
+    objs = APIToken.query(beta=token)
+    if not objs:
+        return get_data_error_result(message='Token is not valid!"')
+    e, c = UserCanvasService.get_by_id(canvas_id)
+    if not e:
+        return get_data_error_result(message="canvas not found.")
+    return get_json_result(data=c.to_dict())
+
 
 @manager.route('/completion', methods=['POST'])  # noqa: F821
 @validate_request("id")

api/apps/conversation_app.py

@@ -17,6 +17,7 @@ import json
 import re
 import traceback
 from copy import deepcopy
+from api.db.db_models import APIToken
 
 from api.db.services.conversation_service import ConversationService, structure_answer
 from api.db.services.user_service import UserTenantService
@@ -32,7 +33,6 @@ from api.utils.api_utils import get_json_result
 from api.utils.api_utils import server_error_response, get_data_error_result, validate_request
 from graphrag.mind_map_extractor import MindMapExtractor
 
-
 @manager.route('/set', methods=['POST'])  # noqa: F821
 @login_required
 def set_conversation():
@@ -79,12 +79,16 @@ def set_conversation():
 def get():
     conv_id = request.args["conversation_id"]
     try:
+        
         e, conv = ConversationService.get_by_id(conv_id)
         if not e:
             return get_data_error_result(message="Conversation not found!")
         tenants = UserTenantService.query(user_id=current_user.id)
+        avatar =None
         for tenant in tenants:
-            if DialogService.query(tenant_id=tenant.tenant_id, id=conv.dialog_id):
+            dialog = DialogService.query(tenant_id=tenant.tenant_id, id=conv.dialog_id)
+            if dialog and len(dialog)>0:
+                avatar = dialog[0].icon
                 break
         else:
             return get_json_result(
@@ -108,10 +112,31 @@ def get():
             } for ck in ref.get("chunks", [])]
 
         conv = conv.to_dict()
+        conv["avatar"]=avatar
         return get_json_result(data=conv)
     except Exception as e:
         return server_error_response(e)
 
+@manager.route('/getsse/<dialog_id>', methods=['GET'])  # type: ignore # noqa: F821
+def getsse(dialog_id):
+    
+    token = request.headers.get('Authorization').split()
+    if len(token) != 2:
+        return get_data_error_result(message='Authorization is not valid!"')
+    token = token[1]
+    objs = APIToken.query(beta=token)
+    if not objs:
+        return get_data_error_result(message='Token is not valid!"')
+    try:
+        e, conv = DialogService.get_by_id(dialog_id)
+        if not e:
+            return get_data_error_result(message="Dialog not found!")
+        conv = conv.to_dict()
+        conv["avatar"]= conv["icon"]
+        del conv["icon"]
+        return get_json_result(data=conv)
+    except Exception as e:
+        return server_error_response(e)
 
 @manager.route('/rm', methods=['POST'])  # noqa: F821
 @login_required

web/src/components/message-item/index.tsx

@@ -30,6 +30,7 @@ interface IProps extends Partial<IRemoveMessageById>, IRegenerateMessage {
   sendLoading?: boolean;
   nickname?: string;
   avatar?: string;
+  avatardialog?: string | null;
   clickDocumentButton?: (documentId: string, chunk: IReferenceChunk) => void;
   index: number;
   showLikeButton?: boolean;
@@ -40,6 +41,7 @@ const MessageItem = ({
   reference,
   loading = false,
   avatar,
+  avatardialog,
   sendLoading = false,
   clickDocumentButton,
   index,
@@ -103,8 +105,10 @@ const MessageItem = ({
         >
           {item.role === MessageType.User ? (
             <Avatar size={40} src={avatar ?? '/logo.svg'} />
+          ) : avatardialog ? (
+            <Avatar size={40} src={avatardialog} />
           ) : (
-            <AssistantIcon></AssistantIcon>
+            <AssistantIcon />
           )}
           <Flex vertical gap={8} flex={1}>
             <Space>

web/src/hooks/chat-hooks.ts

@@ -11,6 +11,7 @@ import {
 } from '@/interfaces/request/chat';
 import i18n from '@/locales/config';
 import { IClientConversation } from '@/pages/chat/interface';
+import { useGetSharedChatSearchParams } from '@/pages/chat/shared-hooks';
 import chatService from '@/services/chat-service';
 import {
   buildMessageListWithUuid,
@@ -27,6 +28,7 @@ import { history, useSearchParams } from 'umi';
 //#region logic
 
 export const useClickDialogCard = () => {
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
   const [_, setSearchParams] = useSearchParams();
 
   const newQueryParameters: URLSearchParams = useMemo(() => {
@@ -243,6 +245,7 @@ export const useFetchNextConversationList = () => {
 
 export const useFetchNextConversation = () => {
   const { isNew, conversationId } = useGetChatSearchParams();
+  const { sharedId } = useGetSharedChatSearchParams();
   const {
     data,
     isFetching: loading,
@@ -254,8 +257,13 @@ export const useFetchNextConversation = () => {
     gcTime: 0,
     refetchOnWindowFocus: false,
     queryFn: async () => {
-      if (isNew !== 'true' && isConversationIdExist(conversationId)) {
-        const { data } = await chatService.getConversation({ conversationId });
+      if (
+        isNew !== 'true' &&
+        isConversationIdExist(sharedId || conversationId)
+      ) {
+        const { data } = await chatService.getConversation({
+          conversationId: conversationId || sharedId,
+        });
 
         const conversation = data?.data ?? {};
 
@@ -270,6 +278,33 @@ export const useFetchNextConversation = () => {
   return { data, loading, refetch };
 };
 
+export const useFetchNextConversationSSE = () => {
+  const { isNew } = useGetChatSearchParams();
+  const { sharedId } = useGetSharedChatSearchParams();
+  const {
+    data,
+    isFetching: loading,
+    refetch,
+  } = useQuery<IClientConversation>({
+    queryKey: ['fetchConversationSSE', sharedId],
+    initialData: {} as IClientConversation,
+    gcTime: 0,
+    refetchOnWindowFocus: false,
+    queryFn: async () => {
+      if (isNew !== 'true' && isConversationIdExist(sharedId || '')) {
+        if (!sharedId) return {};
+        const { data } = await chatService.getConversationSSE({}, sharedId);
+        const conversation = data?.data ?? {};
+        const messageList = buildMessageListWithUuid(conversation?.message);
+        return { ...conversation, message: messageList };
+      }
+      return { message: [] };
+    },
+  });
+
+  return { data, loading, refetch };
+};
+
 export const useFetchManualConversation = () => {
   const {
     data,
@@ -547,7 +582,7 @@ export const useFetchMindMap = () => {
       try {
         const ret = await chatService.getMindMap(params);
         return ret?.data?.data ?? {};
-      } catch (error) {
+      } catch (error: any) {
         if (has(error, 'message')) {
           message.error(error.message);
         }

web/src/hooks/flow-hooks.ts

@@ -2,6 +2,7 @@ import { ResponseType } from '@/interfaces/database/base';
 import { DSL, IFlow, IFlowTemplate } from '@/interfaces/database/flow';
 import { IDebugSingleRequestBody } from '@/interfaces/request/flow';
 import i18n from '@/locales/config';
+import { useGetSharedChatSearchParams } from '@/pages/chat/shared-hooks';
 import flowService from '@/services/flow-service';
 import { buildMessageListWithUuid } from '@/utils/chat';
 import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query';
@@ -91,6 +92,8 @@ export const useFetchFlow = (): {
   refetch: () => void;
 } => {
   const { id } = useParams();
+  const { sharedId } = useGetSharedChatSearchParams();
+
   const {
     data,
     isFetching: loading,
@@ -103,7 +106,41 @@ export const useFetchFlow = (): {
     refetchOnWindowFocus: false,
     gcTime: 0,
     queryFn: async () => {
-      const { data } = await flowService.getCanvas({}, id);
+      const { data } = await flowService.getCanvas({}, sharedId || id);
+
+      const messageList = buildMessageListWithUuid(
+        get(data, 'data.dsl.messages', []),
+      );
+      set(data, 'data.dsl.messages', messageList);
+
+      return data?.data ?? {};
+    },
+  });
+
+  return { data, loading, refetch };
+};
+
+export const useFetchFlowSSE = (): {
+  data: IFlow;
+  loading: boolean;
+  refetch: () => void;
+} => {
+  const { sharedId } = useGetSharedChatSearchParams();
+
+  const {
+    data,
+    isFetching: loading,
+    refetch,
+  } = useQuery({
+    queryKey: ['flowDetailSSE'],
+    initialData: {} as IFlow,
+    refetchOnReconnect: false,
+    refetchOnMount: false,
+    refetchOnWindowFocus: false,
+    gcTime: 0,
+    queryFn: async () => {
+      if (!sharedId) return {};
+      const { data } = await flowService.getCanvasSSE({}, sharedId);
 
       const messageList = buildMessageListWithUuid(
         get(data, 'data.dsl.messages', []),

web/src/interfaces/database/chat.ts

@@ -57,6 +57,7 @@ export interface IConversation {
   create_time: number;
   dialog_id: string;
   id: string;
+  avatar: string;
   message: Message[];
   reference: IReference[];
   name: string;

web/src/interfaces/database/flow.ts

@@ -29,8 +29,8 @@ export interface IGraph {
   edges: Edge[];
 }
 
-export interface IFlow {
-  avatar: null;
+export declare interface IFlow {
+  avatar?: null | string;
   canvas_type: null;
   create_date: string;
   create_time: number;

web/src/pages/chat/chat-container/index.tsx

@@ -68,6 +68,7 @@ const ChatContainer = ({ controller }: IProps) => {
                     item={message}
                     nickname={userInfo.nickname}
                     avatar={userInfo.avatar}
+                    avatardialog={conversation.avatar}
                     reference={buildMessageItemReference(
                       {
                         message: derivedMessages,

web/src/pages/chat/share/large.tsx

@@ -4,7 +4,7 @@ import { useClickDrawer } from '@/components/pdf-drawer/hooks';
 import { MessageType } from '@/constants/chat';
 import { useSendButtonDisabled } from '@/pages/chat/hooks';
 import { Flex, Spin } from 'antd';
-import { forwardRef } from 'react';
+import { forwardRef, useMemo } from 'react';
 import {
   useGetSharedChatSearchParams,
   useSendSharedMessage,
@@ -12,6 +12,8 @@ import {
 import { buildMessageItemReference } from '../utils';
 
 import PdfDrawer from '@/components/pdf-drawer';
+import { useFetchNextConversationSSE } from '@/hooks/chat-hooks';
+import { useFetchFlowSSE } from '@/hooks/flow-hooks';
 import styles from './index.less';
 
 const ChatContainer = () => {
@@ -30,6 +32,14 @@ const ChatContainer = () => {
     hasError,
   } = useSendSharedMessage();
   const sendDisabled = useSendButtonDisabled(value);
+  const useData = (from: SharedFrom) =>
+    useMemo(() => {
+      return from === SharedFrom.Agent
+        ? useFetchFlowSSE
+        : useFetchNextConversationSSE;
+    }, [from]);
+
+  const { data: InforForm } = useData(from)();
 
   if (!conversationId) {
     return <div>empty</div>;
@@ -45,6 +55,7 @@ const ChatContainer = () => {
                 return (
                   <MessageItem
                     key={message.id}
+                    avatardialog={InforForm?.avatar}
                     item={message}
                     nickname="You"
                     reference={buildMessageItemReference(

web/src/pages/flow/chat/box.tsx

@@ -9,6 +9,7 @@ import { useSendNextMessage } from './hooks';
 
 import PdfDrawer from '@/components/pdf-drawer';
 import { useClickDrawer } from '@/components/pdf-drawer/hooks';
+import { useFetchFlow } from '@/hooks/flow-hooks';
 import { useFetchUserInfo } from '@/hooks/user-setting-hooks';
 import styles from './index.less';
 
@@ -29,6 +30,7 @@ const FlowChatBox = () => {
   useGetFileIcon();
   const { t } = useTranslate('chat');
   const { data: userInfo } = useFetchUserInfo();
+  const { data: cavasInfo } = useFetchFlow();
 
   return (
     <>
@@ -47,6 +49,7 @@ const FlowChatBox = () => {
                     key={message.id}
                     nickname={userInfo.nickname}
                     avatar={userInfo.avatar}
+                    avatardialog={cavasInfo.avatar}
                     item={message}
                     reference={buildMessageItemReference(
                       { message: derivedMessages, reference },

web/src/pages/knowledge/index.tsx

@@ -89,9 +89,12 @@ const KnowledgeList = () => {
             className={styles.knowledgeCardContainer}
           >
             {nextList?.length > 0 ? (
-              nextList.map((item: any) => {
+              nextList.map((item: any, index: number) => {
                 return (
-                  <KnowledgeCard item={item} key={item.name}></KnowledgeCard>
+                  <KnowledgeCard
+                    item={item}
+                    key={`${item?.name}-${index}`}
+                  ></KnowledgeCard>
                 );
               })
             ) : (

web/src/services/chat-service.ts

@@ -8,6 +8,7 @@ const {
   listDialog,
   removeDialog,
   getConversation,
+  getConversationSSE,
   setConversation,
   completeConversation,
   listConversation,
@@ -53,6 +54,10 @@ const methods = {
     url: getConversation,
     method: 'get',
   },
+  getConversationSSE: {
+    url: getConversationSSE,
+    method: 'get',
+  },
   setConversation: {
     url: setConversation,
     method: 'post',

web/src/services/flow-service.ts

@@ -4,6 +4,7 @@ import request from '@/utils/request';
 
 const {
   getCanvas,
+  getCanvasSSE,
   setCanvas,
   listCanvas,
   resetCanvas,
@@ -20,6 +21,10 @@ const methods = {
     url: getCanvas,
     method: 'get',
   },
+  getCanvasSSE: {
+    url: getCanvasSSE,
+    method: 'get',
+  },
   setCanvas: {
     url: setCanvas,
     method: 'post',

web/src/utils/api.ts

@@ -71,6 +71,7 @@ export default {
   listDialog: `${api_host}/dialog/list`,
   setConversation: `${api_host}/conversation/set`,
   getConversation: `${api_host}/conversation/get`,
+  getConversationSSE: `${api_host}/conversation/getsse`,
   listConversation: `${api_host}/conversation/list`,
   removeConversation: `${api_host}/conversation/rm`,
   completeConversation: `${api_host}/conversation/completion`,
@@ -113,6 +114,7 @@ export default {
   listTemplates: `${api_host}/canvas/templates`,
   listCanvas: `${api_host}/canvas/list`,
   getCanvas: `${api_host}/canvas/get`,
+  getCanvasSSE: `${api_host}/canvas/getsse`,
   removeCanvas: `${api_host}/canvas/rm`,
   setCanvas: `${api_host}/canvas/set`,
   resetCanvas: `${api_host}/canvas/reset`,