22import { useLanguage } from '@/intl/client' ;
33import { t } from '@/intl/translate' ;
44import { Icon } from '@gitbook/icons' ;
5- import { useEffect , useState } from 'react' ;
5+ import { useEffect } from 'react' ;
66import { create } from 'zustand' ;
7+ import { useShallow } from 'zustand/react/shallow' ;
78import { useVisitedPages } from '../Insights' ;
89import { usePageContext } from '../PageContext' ;
910import { Loading } from '../primitives' ;
1011import { streamLinkPageSummary } from './server-actions/streamLinkPageSummary' ;
1112
12- const useSummaries = create < {
13- cache : Map < string , string > ;
14- setSummary : ( key : string , summary : string ) => void ;
15- } > ( ( set ) => ( {
16- cache : new Map ( ) ,
17- setSummary : ( key , summary ) =>
18- set ( ( state ) => {
19- const newCache = new Map ( state . cache ) ;
20- newCache . set ( key , summary ) ;
21- return { cache : newCache } ;
22- } ) ,
23- } ) ) ;
24-
2513/**
2614 * Get a unique cache key for a page summary
2715 */
2816function getCacheKey ( targetSpaceId : string , targetPageId : string ) : string {
2917 return `${ targetSpaceId } :${ targetPageId } ` ;
3018}
3119
20+ /**
21+ * Global state for the summaries.
22+ */
23+ const useSummaries = create < {
24+ /**
25+ * Cache of all summaries generated so far.
26+ */
27+ cache : Map < string , string > ;
28+
29+ /**
30+ * Get a summary for a page.
31+ */
32+ getSummary : ( params : { targetSpaceId : string ; targetPageId : string } ) => string ;
33+
34+ /**
35+ * Stream the generation of a summary for a page.
36+ */
37+ streamSummary : ( params : {
38+ currentSpaceId : string ;
39+ currentPageId : string ;
40+ currentPageTitle : string ;
41+ targetSpaceId : string ;
42+ targetPageId : string ;
43+ linkPreview ?: string ;
44+ linkTitle ?: string ;
45+ visitedPages : { spaceId : string ; pageId : string } [ ] ;
46+ } ) => Promise < void > ;
47+ } > ( ( set , get ) => ( {
48+ cache : new Map ( ) ,
49+
50+ getSummary : ( {
51+ targetSpaceId,
52+ targetPageId,
53+ } : {
54+ targetSpaceId : string ;
55+ targetPageId : string ;
56+ } ) => {
57+ return get ( ) . cache . get ( getCacheKey ( targetSpaceId , targetPageId ) ) ?? '' ;
58+ } ,
59+
60+ streamSummary : async ( {
61+ currentSpaceId,
62+ currentPageId,
63+ currentPageTitle,
64+ targetSpaceId,
65+ targetPageId,
66+ linkPreview,
67+ linkTitle,
68+ visitedPages,
69+ } ) => {
70+ const cacheKey = getCacheKey ( targetSpaceId , targetPageId ) ;
71+
72+ if ( get ( ) . cache . has ( cacheKey ) ) {
73+ // Already generated or generating
74+ return ;
75+ }
76+
77+ const update = ( summary : string ) => {
78+ set ( ( prev ) => {
79+ const newCache = new Map ( prev . cache ) ;
80+ newCache . set ( cacheKey , summary ) ;
81+ return { cache : newCache } ;
82+ } ) ;
83+ } ;
84+
85+ update ( '' ) ;
86+ const stream = await streamLinkPageSummary ( {
87+ currentSpaceId,
88+ currentPageId,
89+ currentPageTitle,
90+ targetSpaceId,
91+ targetPageId,
92+ linkPreview,
93+ linkTitle,
94+ visitedPages,
95+ } ) ;
96+
97+ let generatedSummary = '' ;
98+ for await ( const highlight of stream ) {
99+ generatedSummary = highlight ?? '' ;
100+ update ( generatedSummary ) ;
101+ }
102+ } ,
103+ } ) ) ;
104+
32105/**
33106 * Summarise a page's content for use in a link preview
34107 */
@@ -44,53 +117,26 @@ export function AIPageLinkSummary(props: {
44117 const currentPage = usePageContext ( ) ;
45118 const language = useLanguage ( ) ;
46119 const visitedPages = useVisitedPages ( ( state ) => state . pages ) ;
47- const [ summary , setSummary ] = useState ( '' ) ;
48- const cacheKey = getCacheKey ( targetSpaceId , targetPageId ) ;
49- const { cachedSummary , setCachedSummary } = useSummaries ( ( state ) => {
50- return {
51- cachedSummary : state . cache . get ( cacheKey ) ?? '' ,
52- setCachedSummary : state . setSummary ,
53- } ;
54- } ) ;
120+ const { summary, streamSummary } = useSummaries (
121+ useShallow ( ( state ) => {
122+ return {
123+ summary : state . getSummary ( { targetSpaceId , targetPageId } ) ,
124+ streamSummary : state . streamSummary ,
125+ } ;
126+ } )
127+ ) ;
55128
56129 useEffect ( ( ) => {
57- let canceled = false ;
58-
59- setSummary ( '' ) ;
60-
61- if ( cachedSummary ) {
62- setSummary ( cachedSummary ) ;
63- return ;
64- }
65-
66- ( async ( ) => {
67- const stream = await streamLinkPageSummary ( {
68- currentSpaceId : currentPage . spaceId ,
69- currentPageId : currentPage . pageId ,
70- currentPageTitle : currentPage . title ,
71- targetSpaceId,
72- targetPageId,
73- linkPreview,
74- linkTitle,
75- visitedPages,
76- } ) ;
77-
78- let generatedSummary = '' ;
79- for await ( const highlight of stream ) {
80- if ( canceled ) return ;
81- generatedSummary = highlight ?? '' ;
82- setSummary ( generatedSummary ) ;
83- }
84-
85- // Cache the complete summary
86- if ( generatedSummary ) {
87- setCachedSummary ( cacheKey , generatedSummary ) ;
88- }
89- } ) ( ) ;
90-
91- return ( ) => {
92- canceled = true ;
93- } ;
130+ streamSummary ( {
131+ currentSpaceId : currentPage . spaceId ,
132+ currentPageId : currentPage . pageId ,
133+ currentPageTitle : currentPage . title ,
134+ targetSpaceId,
135+ targetPageId,
136+ linkPreview,
137+ linkTitle,
138+ visitedPages,
139+ } ) ;
94140 } , [
95141 currentPage . pageId ,
96142 currentPage . spaceId ,
@@ -100,9 +146,7 @@ export function AIPageLinkSummary(props: {
100146 linkPreview ,
101147 linkTitle ,
102148 visitedPages ,
103- cachedSummary ,
104- cacheKey ,
105- setCachedSummary ,
149+ streamSummary ,
106150 ] ) ;
107151
108152 const shimmerBlocks = [
0 commit comments