wip

dcSpark · Feb 8, 2024 · b08cfd0 · b08cfd0
1 parent c8da734
commit b08cfd0
Show file tree

Hide file tree

Showing 3 changed files with 86 additions and 40 deletions.
diff --git a/apps/shinkai-visor/src/components/inbox/inbox.tsx b/apps/shinkai-visor/src/components/inbox/inbox.tsx
@@ -126,30 +126,6 @@ export const Inbox = () => {
     }
   };
 
-  const retrySendMessage = (message: ChatConversationMessage) => {
-    console.log('retrySendMessage');
-    if (!auth) return;
-    const jobId = extractJobIdFromInbox(decodedInboxId);
-    sendMessageToJob({
-      nodeAddress: auth.node_address,
-      jobId,
-      // message: 'what is html in 1 one line',
-      message: 'what is css in 3 words',
-      files_inbox: '',
-      parent:
-        '57b61385ac57ccdbc3447a8175b7442d7a2fe9aafcbf5203dc9e719411f6db7d',
-      // parent:
-      // '530f2a9a6107e9918a88f74e69aab8ed12e81aa07cc92569bf2fdac4e463f047',
-      shinkaiIdentity: auth.shinkai_identity,
-      profile: auth.profile,
-      my_device_encryption_sk: auth.my_device_encryption_sk,
-      my_device_identity_sk: auth.my_device_identity_sk,
-      node_encryption_pk: auth.node_encryption_pk,
-      profile_encryption_sk: auth.profile_encryption_sk,
-      profile_identity_sk: auth.profile_identity_sk,
-    });
-  };
-
   useEffect(() => {
     const chatContainerElement = chatContainerRef.current;
     if (!chatContainerElement) return;
@@ -260,7 +236,7 @@ export const Inbox = () => {
                           </span>
                         </div>
                         <div className="flex flex-col gap-4">
-                          <ChatMessages messages={messages} />
+                          <ChatMessages initialMessages={messages} />
                         </div>
                       </div>
                     );
@@ -295,16 +271,70 @@ export const Inbox = () => {
   );
 };
 
-function ChatMessages({ messages }: { messages: ChatConversationMessage[] }) {
-  console.log(messages, 'messages');
-  const [mainThread, setMainThread] = useState<ChatConversationMessage[]>([]);
+function ChatMessages({
+  initialMessages,
+}: {
+  initialMessages: ChatConversationMessage[];
+}) {
+  const auth = useAuth((state) => state.auth);
+  const { inboxId } = useParams<{ inboxId: string }>();
+
+  const { mutateAsync: sendMessageToJob, isPending: isSendingMessageToJob } =
+    useSendMessageToJob();
+
+  console.log(initialMessages, 'messages');
+  const [activeThread, setActiveThread] = useState<ChatConversationMessage[]>(
+    [],
+  );
+  const messagesRef = useRef<ChatConversationMessage[]>([]);
+  console.log(messagesRef, 'messagesRef');
+  const buildActiveThread = useCallback(
+    (messages: ChatConversationMessage[]) => {
+      const thread = [];
+      let currentMessage = messages.find((m) => m.parentHash === 'root'); // Start from root
+
+      while (currentMessage) {
+        thread.push(currentMessage);
+        currentMessage = messages.find(
+          (message) => message.parentHash === currentMessage?.hash,
+        );
+      }
+
+      return thread;
+    },
+    [],
+  );
 
   useEffect(() => {
-    // TODO:  get main thread
-    setMainThread(messages);
-  }, [messages]);
+    const thread = buildActiveThread(initialMessages);
+    setActiveThread(thread);
+  }, [initialMessages, buildActiveThread]);
+
+  const regenerateLastMessage = () => {
+    const content = activeThread.at(-2)?.content;
+    const parentHash = activeThread.at(-3)?.hash;
+    console.log(content, parentHash, 'activeThread');
+    console.log('retrySendMessage');
+    if (!auth) return;
+    const decodedInboxId = decodeURIComponent(inboxId);
+    const jobId = extractJobIdFromInbox(decodedInboxId);
+    sendMessageToJob({
+      nodeAddress: auth.node_address,
+      jobId,
+      message: content!,
+      files_inbox: '',
+      parent: parentHash!,
+      shinkaiIdentity: auth.shinkai_identity,
+      profile: auth.profile,
+      my_device_encryption_sk: auth.my_device_encryption_sk,
+      my_device_identity_sk: auth.my_device_identity_sk,
+      node_encryption_pk: auth.node_encryption_pk,
+      profile_encryption_sk: auth.profile_encryption_sk,
+      profile_identity_sk: auth.profile_identity_sk,
+    });
+  };
 
-  return mainThread.map((message, index) => {
+  return activeThread.map((message, index) => {
     return (
       <div
         className={cn('relative pl-2')}
@@ -313,13 +343,11 @@ function ChatMessages({ messages }: { messages: ChatConversationMessage[] }) {
         }`}
         key={`${index}-${message.scheduledTime}`}
       >
-        <Message message={message} />
-        {message.isLocal && (
-          <button className="absolute left-2 top-1/2 -translate-y-1/2 transform rounded-full bg-gray-300">
-            <RotateCcw className="h-4 w-4" />
-            <span className="sr-only">Retry</span>
-          </button>
-        )}
+        <Message
+          isLastMessage={index === activeThread.length - 1}
+          message={message}
+          regenerateLastMessage={regenerateLastMessage}
+        />
       </div>
     );
   });

diff --git a/apps/shinkai-visor/src/components/message/message.tsx b/apps/shinkai-visor/src/components/message/message.tsx
@@ -3,10 +3,12 @@ import {
   Avatar,
   AvatarFallback,
   AvatarImage,
+  buttonVariants,
   CopyToClipboardIcon,
   MarkdownPreview,
 } from '@shinkai_network/shinkai-ui';
 import { cn } from '@shinkai_network/shinkai-ui/utils';
+import { RotateCcw } from 'lucide-react';
 import React from 'react';
 
 import shinkaiMiniLogo from '../../assets/icons/shinkai-min.svg';
@@ -17,6 +19,8 @@ import { FileList } from '../file-list/file-list';
 
 type MessageProps = {
   message: ChatConversationMessage;
+  isLastMessage?: boolean;
+  regenerateLastMessage: () => void;
 };
 const copyToClipboard = (content: string) => {
   sendMessage({
@@ -25,7 +29,11 @@ const copyToClipboard = (content: string) => {
   });
 };
 
-export const Message = ({ message }: MessageProps) => {
+export const Message = ({
+  message,
+  isLastMessage,
+  regenerateLastMessage,
+}: MessageProps) => {
   const openMarkdownLink = (
     event: React.MouseEvent<HTMLAnchorElement, MouseEvent>,
     url?: string,
@@ -81,6 +89,15 @@ export const Message = ({ message }: MessageProps) => {
           }}
           source={message.content}
         />
+        {isLastMessage && (
+          <button
+            className={cn('mt-2 h-7 w-7 rounded-full bg-gray-500 p-2')}
+            onClick={regenerateLastMessage}
+          >
+            <RotateCcw className="h-full w-full" />
+            <span className="sr-only">Regenerate</span>
+          </button>
+        )}
         {!!message.fileInbox?.files?.length && (
           <FileList
             actions={[]}

diff --git a/libs/shinkai-node-state/src/lib/queries/getChatConversation/types.ts b/libs/shinkai-node-state/src/lib/queries/getChatConversation/types.ts
@@ -25,6 +25,7 @@ export type ChatConversationMessage = {
       name: string;
     }[];
   };
+  parentHash?: string;
 };
 
 export type GetChatConversationOutput = ChatConversationMessage[];