feat(ui): request the default prompt in search params (#1734)

* feat: default prompt * remove console * update * Update ee/tabby-ui/app/playground/components/chats.tsx Co-authored-by: Meng Zhang <[email protected]> * update --------- Co-authored-by: Meng Zhang <[email protected]>
TabbyML · Mar 27, 2024 · 3c32788 · 3c32788
1 parent 8ecb729
commit 3c32788
Show file tree

Hide file tree

Showing 2 changed files with 82 additions and 33 deletions.
diff --git a/ee/tabby-ui/app/playground/components/chats.tsx b/ee/tabby-ui/app/playground/components/chats.tsx
@@ -3,32 +3,74 @@
 import React from 'react'
 import type { Message } from 'ai'
 
+import useRouterStuff from '@/lib/hooks/use-router-stuff'
 import { useStore } from '@/lib/hooks/use-store'
+import { addChat } from '@/lib/stores/chat-actions'
 import { useChatStore } from '@/lib/stores/chat-store'
 import { getChatById } from '@/lib/stores/utils'
-import { Chat } from '@/components/chat'
+import { truncateText } from '@/lib/utils'
+import { Chat, ChatRef } from '@/components/chat'
+import LoadingWrapper from '@/components/loading-wrapper'
+import { ListSkeleton } from '@/components/skeleton'
 
 import { ChatSessions } from './chat-sessions'
 
 const emptyMessages: Message[] = []
 
 export default function Chats() {
+  const { searchParams, updateSearchParams } = useRouterStuff()
+  const initialMessage = searchParams.get('initialMessage')?.toString()
+  const shouldConsumeInitialMessage = React.useRef(!!initialMessage)
+  const chatRef = React.useRef<ChatRef>(null)
+
   const _hasHydrated = useStore(useChatStore, state => state._hasHydrated)
   const chats = useStore(useChatStore, state => state.chats)
   const activeChatId = useStore(useChatStore, state => state.activeChatId)
+  const chat = getChatById(chats, activeChatId)
+
+  React.useEffect(() => {
+    if (!shouldConsumeInitialMessage.current) return
+    if (!chatRef.current?.append) return
+
+    if (activeChatId && initialMessage) {
+      // request initialMessage
+      chatRef.current
+        .append({
+          role: 'user',
+          content: initialMessage
+        })
+        .then(() => {
+          // Remove the initialMessage params after the request is completed.
+          updateSearchParams({
+            del: 'initialMessage'
+          })
+        })
+      // store as a new chat
+      addChat(activeChatId, truncateText(initialMessage))
 
-  const chatId = activeChatId
-  const chat = getChatById(chats, chatId)
+      shouldConsumeInitialMessage.current = false
+    }
+  }, [chatRef.current?.append])
 
   return (
     <div className="grid flex-1 overflow-hidden lg:grid-cols-[280px_1fr]">
       <ChatSessions className="hidden w-[280px] border-r lg:block" />
-      <Chat
-        loading={!_hasHydrated}
-        id={chatId}
-        key={chatId}
-        initialMessages={chat?.messages ?? emptyMessages}
-      />
+      <LoadingWrapper
+        delay={0}
+        loading={!_hasHydrated || !activeChatId}
+        fallback={
+          <div className="mx-auto w-full max-w-2xl pt-4 md:pt-10">
+            <ListSkeleton />
+          </div>
+        }
+      >
+        <Chat
+          id={activeChatId}
+          key={activeChatId}
+          initialMessages={chat?.messages ?? emptyMessages}
+          ref={chatRef}
+        />
+      </LoadingWrapper>
     </div>
   )
 }
diff --git a/ee/tabby-ui/components/chat.tsx b/ee/tabby-ui/components/chat.tsx
@@ -2,7 +2,7 @@
 
 import React from 'react'
 import { useChat } from 'ai/react'
-import type { Message } from 'ai/react'
+import type { Message, UseChatHelpers } from 'ai/react'
 import { find, findIndex } from 'lodash-es'
 import { toast } from 'sonner'
 
@@ -16,27 +16,22 @@ import { ChatList } from '@/components/chat-list'
 import { ChatPanel } from '@/components/chat-panel'
 import { ChatScrollAnchor } from '@/components/chat-scroll-anchor'
 import { EmptyScreen } from '@/components/empty-screen'
-import { ListSkeleton } from '@/components/skeleton'
 
 export interface ChatProps extends React.ComponentProps<'div'> {
   initialMessages?: Message[]
   id?: string
-  loading?: boolean
 }
 
-export function Chat({ id, initialMessages, loading, className }: ChatProps) {
+export interface ChatRef extends UseChatHelpers {}
+
+function ChatRenderer(
+  { id, initialMessages, className }: ChatProps,
+  ref: React.ForwardedRef<ChatRef>
+) {
   usePatchFetch()
   const chats = useStore(useChatStore, state => state.chats)
-  const {
-    messages,
-    append,
-    reload,
-    stop,
-    isLoading,
-    input,
-    setInput,
-    setMessages
-  } = useChat({
+
+  const useChatHelpers = useChat({
     initialMessages,
     id,
     body: {
@@ -49,6 +44,17 @@ export function Chat({ id, initialMessages, loading, className }: ChatProps) {
     }
   })
 
+  const {
+    messages,
+    append,
+    reload,
+    stop,
+    isLoading,
+    input,
+    setInput,
+    setMessages
+  } = useChatHelpers
+
   const [selectedMessageId, setSelectedMessageId] = React.useState<string>()
 
   const onRegenerateResponse = (messageId: string) => {
@@ -122,20 +128,19 @@ export function Chat({ id, initialMessages, loading, className }: ChatProps) {
     return () => stop()
   }, [])
 
+  React.useImperativeHandle(
+    ref,
+    () => {
+      return useChatHelpers
+    },
+    [useChatHelpers]
+  )
+
   return (
     <div className="flex justify-center overflow-x-hidden">
       <div className="w-full max-w-2xl px-4">
         <div className={cn('pb-[200px] pt-4 md:pt-10', className)}>
-          {loading ? (
-            <div className="group relative mb-4 flex animate-pulse items-start md:-ml-12">
-              <div className="shrink-0">
-                <span className="block h-8 w-8 rounded-md bg-gray-200 dark:bg-gray-700"></span>
-              </div>
-              <div className="ml-4 flex-1 space-y-2 overflow-hidden px-1">
-                <ListSkeleton />
-              </div>
-            </div>
-          ) : messages.length ? (
+          {messages.length ? (
             <>
               <ChatList
                 messages={messages}
@@ -163,3 +168,5 @@ export function Chat({ id, initialMessages, loading, className }: ChatProps) {
     </div>
   )
 }
+
+export const Chat = React.forwardRef<ChatRef, ChatProps>(ChatRenderer)