diff --git a/src/components/workspace-shell.tsx b/src/components/workspace-shell.tsx index 204f79f..4e2af57 100644 --- a/src/components/workspace-shell.tsx +++ b/src/components/workspace-shell.tsx @@ -17,6 +17,7 @@ import type { ChatMessageView, ChatThreadView, } from "@/domains/chat/types" +import type { ParsedChunkView } from "@/domains/chunks/types" import type { SourceView } from "@/domains/sources/types" export type { PanelId } from "@/components/workspace-shell-layout" @@ -41,6 +42,7 @@ export type WorkspaceShellProps = { activeChatThreadId?: string | null chatMessages?: ChatMessageView[] dashboardUrl?: string + initialPrefetchedChunksBySourceId?: Record isGuest?: boolean loginUrl?: string } @@ -68,6 +70,7 @@ function WorkspaceShellContent({ activeChatThreadId, chatMessages: initialChatMessages, dashboardUrl, + initialPrefetchedChunksBySourceId, isGuest = false, loginUrl, }: WorkspaceShellProps): ReactElement { @@ -80,6 +83,8 @@ function WorkspaceShellContent({ }) const citationFocus = useWorkspaceCitationFocus({ fetchChunks: workspaceClient.fetchChunks, + initialPrefetchedChunksBySourceId: + initialPrefetchedChunksBySourceId ?? undefined, onSelectSource: sourceWorkflow.setSelectedSourceId, selectedSourceId: sourceWorkflow.selectedSourceId, sources: sourceWorkflow.sources, diff --git a/src/domains/workspace/initial-state.ts b/src/domains/workspace/initial-state.ts index 39f2fc5..d11e1fb 100644 --- a/src/domains/workspace/initial-state.ts +++ b/src/domains/workspace/initial-state.ts @@ -1,8 +1,10 @@ import "server-only" +import { unstable_cache } from "next/cache" import { Effect } from "effect" import type { ChatMessageView } from "@/domains/chat/types" +import type { ParsedChunkView } from "@/domains/chunks/types" import { demoView } from "@/domains/demo/view" import { getMaterializedDemoSourceViewOptionsBySourceId, @@ -33,6 +35,7 @@ type WorkspaceShellInitialState = { readonly chatMessages?: ChatMessageView[] readonly chatThreads?: ReturnType[] readonly dashboardUrl?: string + readonly initialPrefetchedChunksBySourceId?: Record readonly isGuest?: boolean readonly loginUrl?: string readonly sources?: SourceView[] @@ -47,6 +50,38 @@ type WorkspaceShellInitialState = { } } +const getCachedDemoChunksForSource = (demoSourceId: string) => + unstable_cache( + async (): Promise => { + const chunkPage = await knowhereDemoApi.fetchChunkPage({ + demoSourceId, + page: 1, + pageSize: 100, + }) + const sourceView = demoView.toSourceView({ + demoSourceId: chunkPage.demoSourceId, + canonicalDocumentId: chunkPage.canonicalDocumentId, + title: chunkPage.title, + mimeType: chunkPage.mimeType, + sizeBytes: 0, + status: "ready", + chunkCount: chunkPage.pagination.total, + originalFile: { + url: "", + mimeType: "", + sizeBytes: 0, + canDownload: false, + }, + examples: [], + }) + return chunkPage.chunks.map((chunk) => + demoView.toParsedChunkView(sourceView, chunk), + ) + }, + ["demo-chunks", demoSourceId], + { revalidate: false }, + )() + type WorkspaceShellInitialStateClient = Parameters[1] @@ -117,11 +152,32 @@ export const loadWorkspaceShellInitialStateEffect = ( deps.fetchDemoCatalog(), ) const guestContext = yield* Effect.tryPromise(() => deps.getGuest()) + + const firstDemoSource = demoCatalog.sources[0] + let initialPrefetchedChunksBySourceId: Record< + string, + ParsedChunkView[] + > = {} + if (firstDemoSource) { + const chunks = yield* Effect.catchAll( + Effect.tryPromise(() => + getCachedDemoChunksForSource(firstDemoSource.demoSourceId), + ), + () => Effect.succeed([] as ParsedChunkView[]), + ) + if (chunks.length > 0) { + initialPrefetchedChunksBySourceId = { + [firstDemoSource.demoSourceId]: chunks, + } + } + } + return { isGuest: true, sources: demoCatalog.sources.map(demoView.toSourceView), chatMessages: demoView.toChatMessages(demoCatalog), dashboardUrl: resolveDashboardUrl(), + initialPrefetchedChunksBySourceId, loginUrl: guestContext.loginUrl, } }