Hooks and APIs for AI v2 (#3385)

packages/gitbook-v2/src/lib/data/api.ts

@@ -169,10 +169,6 @@ export function createDataFetcher(
         getUserById(userId) {
             return trace('getUserById', () => getUserById(input, { userId }));
         },
-
-        streamAIResponse(params) {
-            return streamAIResponse(input, params);
-        },
     };
 }
 
@@ -657,29 +653,6 @@ const renderIntegrationUi = cache(
     }
 );
 
-async function* streamAIResponse(
-    input: DataFetcherInput,
-    params: Parameters<GitBookDataFetcher['streamAIResponse']>[0]
-) {
-    const api = apiClient(input);
-    const res = await api.orgs.streamAiResponseInSite(
-        params.organizationId,
-        params.siteId,
-        {
-            input: params.input,
-            output: params.output,
-            model: params.model,
-        },
-        {
-            ...noCacheFetchOptions,
-        }
-    );
-
-    for await (const event of res) {
-        yield event;
-    }
-}
-
 /**
  * Create a new API client.
  */

packages/gitbook-v2/src/lib/data/types.ts

@@ -160,15 +160,4 @@ export interface GitBookDataFetcher {
         integrationName: string;
         request: api.RenderIntegrationUI;
     }): Promise<DataFetcherResponse<api.ContentKitRenderOutput>>;
-
-    /**
-     * Stream an AI response.
-     */
-    streamAIResponse(params: {
-        organizationId: string;
-        siteId: string;
-        input: api.AIMessageInput[];
-        output: api.AIOutputFormat;
-        model: api.AIModel;
-    }): AsyncGenerator<api.AIStreamResponse, void, unknown>;
 }

packages/gitbook/src/components/AI/index.ts

@@ -0,0 +1,2 @@
+export * from './useAIPage';
+export * from './useAIChat';

packages/gitbook/src/components/AI/server-actions/AIMessageView.tsx

@@ -0,0 +1,40 @@
+import type { AIMessage } from '@gitbook/api';
+import type { GitBookSiteContext } from '@v2/lib/context';
+import { DocumentView } from '../../DocumentView';
+import { AIToolCallsSummary } from './AIToolCallsSummary';
+import type { RenderAIMessageOptions } from './types';
+
+/**
+ * Render a message from the API backend.
+ */
+export function AIMessageView(
+    props: RenderAIMessageOptions & {
+        message: AIMessage;
+        context: GitBookSiteContext;
+    }
+) {
+    const { message, context, renderToolCalls = true } = props;
+
+    return (
+        <div className="flex flex-col gap-2">
+            {message.steps.map((step, index) => {
+                return (
+                    <div key={index} className="flex flex-col gap-2">
+                        <DocumentView
+                            document={step.content}
+                            context={{
+                                mode: 'default',
+                                contentContext: undefined,
+                                wrapBlocksInSuspense: false,
+                            }}
+                            style={['space-y-5']}
+                        />
+                        {renderToolCalls && step.toolCalls && step.toolCalls.length > 0 ? (
+                            <AIToolCallsSummary toolCalls={step.toolCalls} context={context} />
+                        ) : null}
+                    </div>
+                );
+            })}
+        </div>
+    );
+}

packages/gitbook/src/components/AI/server-actions/AIToolCallsSummary.tsx

@@ -0,0 +1,142 @@
+import { Link } from '@/components/primitives';
+import { resolveContentRef } from '@/lib/references';
+import type { AIToolCall, ContentRef } from '@gitbook/api';
+import { Icon, type IconName } from '@gitbook/icons';
+import type { GitBookSiteContext } from '@v2/lib/context';
+import type * as React from 'react';
+
+/**
+ * Display the tool calls in a message or step.
+ */
+export function AIToolCallsSummary(props: {
+    toolCalls: AIToolCall[];
+    context: GitBookSiteContext;
+}) {
+    const { toolCalls, context } = props;
+
+    return (
+        <div className="flex flex-col gap-1">
+            {toolCalls.map((toolCall, index) => (
+                <ToolCallSummary key={index} toolCall={toolCall} context={context} />
+            ))}
+        </div>
+    );
+}
+
+function ToolCallSummary(props: {
+    toolCall: AIToolCall;
+    context: GitBookSiteContext;
+}) {
+    const { toolCall, context } = props;
+
+    return (
+        <p className="text-slate-700 text-sm">
+            <Icon
+                icon={getIconForToolCall(toolCall)}
+                className="mr-1 inline-block size-3 text-slate-300"
+            />
+            {getDescriptionForToolCall(toolCall, context)}
+        </p>
+    );
+}
+
+function getDescriptionForToolCall(
+    toolCall: AIToolCall,
+    context: GitBookSiteContext
+): React.ReactNode {
+    switch (toolCall.tool) {
+        case 'getPageContent':
+            return (
+                <>
+                    Read page{' '}
+                    <ContentRefLink
+                        contentRef={{
+                            kind: 'page',
+                            page: toolCall.page.id,
+                            space: toolCall.spaceId,
+                        }}
+                        context={context}
+                        fallback={toolCall.page.title}
+                    />
+                    <OtherSpaceLink spaceId={toolCall.spaceId} context={context} />
+                </>
+            );
+        case 'search':
+            // TODO: Show in a popover the results using the list `toolCall.results`.
+            return (
+                <>
+                    Searched <strong>{toolCall.query}</strong>
+                </>
+            );
+        case 'getPages':
+            return (
+                <>
+                    Listed the pages
+                    <OtherSpaceLink spaceId={toolCall.spaceId} context={context} />
+                </>
+            );
+        default:
+            return <>{toolCall.tool}</>;
+    }
+}
+
+function getIconForToolCall(toolCall: AIToolCall): IconName {
+    switch (toolCall.tool) {
+        case 'getPageContent':
+            return 'memo';
+        case 'search':
+            return 'magnifying-glass';
+        case 'getPages':
+            return 'files';
+        default:
+            return 'hammer';
+    }
+}
+
+/**
+ * Link to a space that is not the current space.
+ */
+function OtherSpaceLink(props: {
+    spaceId: string;
+    context: GitBookSiteContext;
+    prefix?: React.ReactNode;
+}) {
+    const { spaceId, prefix = ' in ', context } = props;
+
+    if (context.space.id === spaceId) {
+        return null;
+    }
+
+    return (
+        <>
+            {prefix}
+            <ContentRefLink
+                contentRef={{
+                    kind: 'space',
+                    space: spaceId,
+                }}
+                context={context}
+            />
+        </>
+    );
+}
+
+async function ContentRefLink(props: {
+    contentRef: ContentRef;
+    context: GitBookSiteContext;
+    fallback?: React.ReactNode;
+}) {
+    const { contentRef, context, fallback } = props;
+
+    const resolved = await resolveContentRef(contentRef, context);
+
+    if (!resolved) {
+        return <span>{fallback}</span>;
+    }
+
+    return (
+        <Link href={resolved.href} className="text-inherit underline decoration-dashed">
+            {resolved.text}
+        </Link>
+    );
+}

packages/gitbook/src/components/AI/server-actions/api.tsx

@@ -0,0 +1,232 @@
+'use server';
+import {
+    type AIMessage,
+    type AIMessageInput,
+    AIMessageRole,
+    type AIMessageStep,
+    type AIModel,
+    type AIStreamResponse,
+} from '@gitbook/api';
+import type { GitBookBaseContext } from '@v2/lib/context';
+import { fetchServerActionSiteContext } from '@v2/lib/server-actions';
+import { EventIterator } from 'event-iterator';
+import type { MaybePromise } from 'p-map';
+import * as partialJson from 'partial-json';
+import type { DeepPartial } from 'ts-essentials';
+import type { z } from 'zod';
+import { zodToJsonSchema } from 'zod-to-json-schema';
+import { AIMessageView } from './AIMessageView';
+import type { RenderAIMessageOptions } from './types';
+
+type StreamGenerateInput = {
+    organizationId: string;
+    siteId: string;
+    instructions?: string;
+    previousResponseId?: string;
+    input: AIMessageInput[];
+    model: AIModel;
+};
+
+/**
+ * Get the latest value from a stream and the response id.
+ */
+export async function generate<T>(
+    promise: MaybePromise<{
+        stream: EventIterator<T>;
+        response: Promise<{ responseId: string }>;
+    }>
+) {
+    const input = await promise;
+    let value: T | undefined;
+
+    for await (const event of input.stream) {
+        value = event;
+    }
+
+    const { responseId } = await input.response;
+    return {
+        responseId,
+        value,
+    };
+}
+
+/**
+ * Stream the generation of an object using the AI.
+ */
+export async function streamGenerateAIObject<T>(
+    context: GitBookBaseContext,
+    {
+        schema,
+        ...input
+    }: StreamGenerateInput & {
+        schema: z.ZodSchema<T>;
+    }
+) {
+    const api = await context.dataFetcher.api();
+    const rawStream = await api.orgs.streamAiResponseInSite(input.organizationId, input.siteId, {
+        input: input.input,
+        output: { type: 'object', schema: zodToJsonSchema(schema) },
+        model: input.model,
+        instructions: input.instructions,
+        previousResponseId: input.previousResponseId,
+    });
+
+    let json = '';
+    return parseResponse<DeepPartial<T>>(rawStream, (event) => {
+        if (event.type === 'response_object') {
+            json += event.jsonChunk;
+
+            const parsed = partialJson.parse(json, partialJson.ALL);
+            return parsed;
+        }
+    });
+}
+
+/**
+ * Stream the generation of a document.
+ */
+export async function streamRenderAIMessage(
+    baseContext: GitBookBaseContext,
+    rawStream: AsyncIterable<AIStreamResponse>,
+    options?: RenderAIMessageOptions
+) {
+    const message: AIMessage = {
+        id: '',
+        role: AIMessageRole.Assistant,
+        steps: [],
+    };
+
+    const updateProcessingMessageStep = (
+        stepIndex: number,
+        callback: (step: AIMessageStep) => void
+    ) => {
+        if (stepIndex > message.steps.length) {
+            throw new Error(
+                `Step index out of bounds ${stepIndex} (${message.steps.length} steps)`
+            );
+        }
+
+        if (message.steps[stepIndex]) {
+            message.steps = [...message.steps];
+            message.steps[stepIndex] = { ...message.steps[stepIndex] };
+            callback(message.steps[stepIndex]);
+        } else {
+            message.steps = [
+                ...message.steps,
+                {
+                    content: {
+                        object: 'document',
+                        data: {},
+                        nodes: [],
+                    },
+                },
+            ];
+            callback(message.steps[stepIndex]);
+        }
+    };
+
+    // Fetch the full-context in the background to avoid blocking the stream.
+    const promiseContext = fetchServerActionSiteContext(baseContext);
+
+    return parseResponse<{
+        content: React.ReactNode;
+        event: AIStreamResponse;
+    }>(rawStream, async (event) => {
+        switch (event.type) {
+            /**
+             * The agent is processing a tool call in a new message.
+             */
+            case 'response_tool_call': {
+                updateProcessingMessageStep(event.stepIndex, (step) => {
+                    step.toolCalls ??= [];
+                    step.toolCalls.push(event.toolCall);
+                });
+                break;
+            }
+
+            /**
+             * The agent is writing the content of a new message.
+             */
+            case 'response_reasoning':
+            case 'response_document': {
+                updateProcessingMessageStep(event.stepIndex, (step) => {
+                    const container = event.type === 'response_reasoning' ? 'reasoning' : 'content';
+
+                    step[container] ??= {
+                        object: 'document',
+                        data: {},
+                        nodes: [],
+                    };
+                    step[container] = {
+                        ...step[container],
+                        nodes: [...step[container].nodes],
+                    };
+                    if (event.operation === 'insert') {
+                        step[container].nodes.push(...event.blocks);
+                    } else {
+                        step[container].nodes.splice(
+                            -event.blocks.length,
+                            event.blocks.length,
+                            ...event.blocks
+                        );
+                    }
+                });
+                break;
+            }
+        }
+
+        return {
+            event,
+            content: (
+                <AIMessageView message={message} context={await promiseContext} {...options} />
+            ),
+        };
+    });
+}
+
+/**
+ * Parse a stream from the API to extract the responseId.
+ */
+function parseResponse<T>(
+    responseStream: EventIterator<AIStreamResponse>,
+    parse: (response: AIStreamResponse) => T | undefined | Promise<T | undefined>
+): {
+    stream: EventIterator<T>;
+    response: Promise<{ responseId: string }>;
+} {
+    let resolveResponse: (value: { responseId: string }) => void;
+    const response = new Promise<{ responseId: string }>((resolve) => {
+        resolveResponse = resolve;
+    });
+
+    const stream = new EventIterator<T>((queue) => {
+        (async () => {
+            let foundResponse = false;
+
+            for await (const event of responseStream) {
+                const parsed = await parse(event);
+                if (parsed !== undefined) {
+                    queue.push(parsed);
+                }
+
+                if (event.type === 'response_finish') {
+                    foundResponse = true;
+                    resolveResponse({ responseId: event.responseId });
+                }
+            }
+
+            if (!foundResponse) {
+                throw new Error('No response found');
+            }
+        })().then(
+            () => {
+                queue.stop();
+            },
+            (error) => {
+                queue.fail(error);
+            }
+        );
+    });
+
+    return { stream, response };
+}

packages/gitbook/src/components/AI/server-actions/chat.ts

@@ -0,0 +1,117 @@
+'use server';
+import { type AIMessageContext, AIMessageRole, AIModel } from '@gitbook/api';
+import { getSiteURLDataFromMiddleware } from '@v2/lib/middleware';
+import { getServerActionBaseContext } from '@v2/lib/server-actions';
+import { z } from 'zod';
+import { streamGenerateAIObject, streamRenderAIMessage } from './api';
+import { MARKDOWN_LINKS_PROMPT } from './prompts';
+import type { RenderAIMessageOptions } from './types';
+
+const PROMPT = `
+You are GitBook Docs Assistant, a helpful docs assistant that answers questions from the user about a documentation site.
+
+You analyse the query, and the content of the site, and generate a short, concise answer that will help the user.
+
+# Instructions
+
+- Generate a response formatted in markdown
+- Always use the provided tools to understand the docs knowledge base, do not make up information.
+
+${MARKDOWN_LINKS_PROMPT}
+`;
+
+const FOLLOWUP_PROMPT = `
+Generate a short JSON list with message suggestions for a user to post in a chat. The suggestions will be displayed next to the text input, allowing the user to quickly tap and pick one.
+
+# Guidelines
+
+- Ensure suggestions are concise and relevant for general chat conversations.
+- Limit the length of each suggestion to ensure quick readability and tap selection.
+- Suggest at most 3 responses.
+- Only suggest responses that are relevant followup to the conversation, otherwise return an empty list.
+- When the last message finishes with questions, suggest responses that answer the questions.
+
+# Output Format
+
+Provide the suggestions as a JSON array with each suggestion as a string. Ensure the suggestions are short and suitable for quick tapping.
+`;
+
+/**
+ * Generate a response to a chat message.
+ */
+export async function* streamAIChatResponse({
+    message,
+    messageContext,
+    previousResponseId,
+    options,
+}: {
+    message: string;
+    messageContext: AIMessageContext;
+    previousResponseId?: string;
+    options?: RenderAIMessageOptions;
+}) {
+    const context = await getServerActionBaseContext();
+    const siteURLData = await getSiteURLDataFromMiddleware();
+
+    const api = await context.dataFetcher.api();
+    const rawStream = api.orgs.streamAiResponseInSite(siteURLData.organization, siteURLData.site, {
+        input: [
+            {
+                role: AIMessageRole.User,
+                content: message,
+                context: messageContext,
+            },
+        ],
+        output: { type: 'document' },
+        model: AIModel.ReasoningLow,
+        instructions: PROMPT,
+        previousResponseId,
+        tools: {
+            getPageContent: true,
+            getPages: true,
+            search: true,
+        },
+    });
+
+    const { stream } = await streamRenderAIMessage(context, rawStream, options);
+
+    for await (const output of stream) {
+        yield output;
+    }
+}
+
+/**
+ * Stream suggestions of follow-up responses for the user.
+ */
+export async function* streamAIChatFollowUpResponses({
+    previousResponseId,
+}: {
+    previousResponseId: string;
+}) {
+    const context = await getServerActionBaseContext();
+    const siteURLData = await getSiteURLDataFromMiddleware();
+
+    const { stream, response } = await streamGenerateAIObject(context, {
+        organizationId: siteURLData.organization,
+        siteId: siteURLData.site,
+        schema: z.object({
+            suggestions: z.array(z.string()),
+        }),
+        previousResponseId,
+        input: [
+            {
+                role: AIMessageRole.User,
+                content:
+                    'Suggest quick-tap responses the user might want to pick from to continue the previous chat conversation.',
+            },
+        ],
+        model: AIModel.Fast,
+        instructions: FOLLOWUP_PROMPT,
+    });
+
+    for await (const output of stream) {
+        yield (output.suggestions ?? []).filter((suggestion) => !!suggestion) as string[];
+    }
+
+    console.log('response', { previousResponseId }, await response);
+}

packages/gitbook/src/components/AI/server-actions/index.ts

@@ -0,0 +1,4 @@
+export * from './pages';
+export * from './types';
+export * from './responses';
+export * from './chat';

packages/gitbook/src/components/AI/server-actions/pages.ts

@@ -0,0 +1,62 @@
+'use server';
+import { AIMessageRole, AIModel } from '@gitbook/api';
+import { getSiteURLDataFromMiddleware } from '@v2/lib/middleware';
+import { getServerActionBaseContext } from '@v2/lib/server-actions';
+import { streamRenderAIMessage } from './api';
+import { MARKDOWN_SYNTAX_PROMPT } from './prompts';
+import type { RenderAIMessageOptions } from './types';
+
+const PROMPT = `
+You are GitBook AI, a helpful docs assistant that can generate an optimized page for a given query.
+
+You analyse the query, and the content of the site, and generate a page that will help the user understand the content of the site.
+
+# Instructions
+
+- Generate a complete page formatted in markdown
+- Always start the page with a markdown heading 1 (\`# Title of the page\`)
+- Use the provided tools to understand the site content.
+
+${MARKDOWN_SYNTAX_PROMPT}
+`;
+
+/**
+ * Generate a page using AI.
+ */
+export async function* streamGenerateAIPage({
+    query,
+    previousResponseId,
+    options,
+}: {
+    query: string;
+    previousResponseId?: string;
+    options?: RenderAIMessageOptions;
+}) {
+    const context = await getServerActionBaseContext();
+    const siteURLData = await getSiteURLDataFromMiddleware();
+
+    const api = await context.dataFetcher.api();
+    const rawStream = api.orgs.streamAiResponseInSite(siteURLData.organization, siteURLData.site, {
+        input: [
+            {
+                role: AIMessageRole.User,
+                content: query,
+            },
+        ],
+        output: { type: 'document' },
+        model: AIModel.ReasoningLow,
+        instructions: PROMPT,
+        previousResponseId,
+        tools: {
+            getPageContent: true,
+            getPages: true,
+            search: true,
+        },
+    });
+
+    const { stream } = await streamRenderAIMessage(context, rawStream, options);
+
+    for await (const output of stream) {
+        yield output;
+    }
+}

packages/gitbook/src/components/AI/server-actions/prompts.ts

@@ -0,0 +1,71 @@
+/*
+ * Set of common prompts used to generate AI responses.
+ * We'll move this to GBX once we have finished experimenting.
+ */
+
+/**
+ * Prompt to explain the markdown syntax supported by GitBook.
+ */
+export const MARKDOWN_SYNTAX_PROMPT = `
+## Markdown syntax
+
+You can use all the markdown syntax supported by GitHub Flavored Markdown (headings, paragraphs, code blocks, lists, tables, etc).
+
+And you also can use advanced blocks using Liquid syntax, the supported advanced blocks are:
+
+#### Tabs
+
+The tabs block can be used to represent alternatives of content (programming languages, operating systems, etc).
+
+Syntax example:
+
+\`\`\`
+{% tabs %}
+{% tab title="Foo" %}
+First tab content.
+{% endtab %}
+
+{% tab title="Bar" %}
+Second tab content.
+{% endtab %}
+{% endtabs %}
+\`\`\`
+
+#### Stepper
+
+The stepper block can be used to represent a multi-steps process to the user.
+
+Syntax example:
+
+\`\`\`
+{% stepper %}
+{% step %}
+## First step
+
+First step content.
+{% endstep %}
+
+{% step %}
+## Second step
+
+Second step content.
+{% endstep %}
+{% endstepper %}
+\`\`\`
+
+`;
+
+/**
+ * Prompts to indicate how to format links to pages.
+ */
+export const MARKDOWN_LINKS_PROMPT = `
+## Instructions for referring to pages
+
+You MUST use the following format when referring to pages: markdown links with the following format:
+
+\`\`\`
+[Page Title](/spaces/:spaceId/pages/:pageId)
+\`\`\`
+
+Always refer to pages using links and their titles. NEVER refer to pages using their IDs or as "the page".
+`;

packages/gitbook/src/components/AI/server-actions/responses.ts

@@ -0,0 +1,31 @@
+'use server';
+import { getSiteURLDataFromMiddleware } from '@v2/lib/middleware';
+import { getServerActionBaseContext } from '@v2/lib/server-actions';
+import { streamRenderAIMessage } from './api';
+import type { RenderAIMessageOptions } from './types';
+
+/**
+ * Stream an existing AI responses.
+ */
+export async function* streamAIResponseById({
+    responseId,
+    options,
+}: {
+    responseId: string;
+    options?: RenderAIMessageOptions;
+}) {
+    const context = await getServerActionBaseContext();
+    const siteURLData = await getSiteURLDataFromMiddleware();
+
+    const api = await context.dataFetcher.api();
+    const rawStream = api.orgs.streamExistingAiResponseInSite(
+        siteURLData.organization,
+        siteURLData.site,
+        responseId
+    );
+    const { stream } = await streamRenderAIMessage(context, rawStream, options);
+
+    for await (const output of stream) {
+        yield output;
+    }
+}

packages/gitbook/src/components/AI/server-actions/types.ts

@@ -0,0 +1,20 @@
+import type { AIStreamResponse } from '@gitbook/api';
+
+/**
+ * Stream when rendering an AI message.
+ */
+export type AIMessageRenderStream = AsyncIterable<{
+    content: React.ReactNode;
+    event: AIStreamResponse;
+}>;
+
+/**
+ * Options that can be passed when generating a AI message stream.
+ */
+export type RenderAIMessageOptions = {
+    /**
+     * Render the tool calls.
+     * @default true
+     */
+    renderToolCalls?: boolean;
+};

packages/gitbook/src/components/AI/useAIChat.tsx

@@ -0,0 +1,172 @@
+'use client';
+
+import * as zustand from 'zustand';
+
+import { AIMessageRole } from '@gitbook/api';
+import * as React from 'react';
+import { streamAIChatFollowUpResponses, streamAIChatResponse } from './server-actions';
+import { useAIMessageContextRef } from './useAIMessageContext';
+
+export type AIChatMessage = {
+    role: AIMessageRole;
+    content: React.ReactNode;
+};
+
+export type AIChatState = {
+    /**
+     * If true, the chat is open.
+     */
+    opened: boolean;
+
+    /**
+     * ID of the latest AI response.
+     */
+    responseId: string | null;
+
+    /**
+     * Messages in the session.
+     */
+    messages: AIChatMessage[];
+
+    /**
+     * Suggestions for follow-up messages.
+     */
+    followUpSuggestions: string[];
+
+    /**
+     * If true, the session is in progress.
+     */
+    loading: boolean;
+};
+
+export type AIChatController = {
+    /** Open the dialog */
+    open: () => void;
+
+    /** Close the dialog */
+    close: () => void;
+
+    /** Post a message to the session */
+    postMessage: (input: {
+        /** The message to post to the session. it can be markdown formatted. */
+        message: string;
+    }) => void;
+
+    /** Clear the conversation */
+    clear: () => void;
+};
+
+const globalState = zustand.create<{
+    state: AIChatState;
+    setState: (fn: (state: AIChatState) => Partial<AIChatState>) => void;
+}>((set) => {
+    return {
+        state: {
+            opened: false,
+            responseId: null,
+            messages: [],
+            followUpSuggestions: [],
+            loading: false,
+        },
+        setState: (fn) => set((state) => ({ state: { ...state.state, ...fn(state.state) } })),
+    };
+});
+
+/**
+ * Get the current state of the AI chat.
+ */
+export function useAIChatState(): AIChatState {
+    const state = zustand.useStore(globalState, (state) => state.state);
+    return state;
+}
+
+/**
+ * Get the controller to interact with the AI chat.
+ */
+export function useAIChatController(): AIChatController {
+    const messageContextRef = useAIMessageContextRef();
+    const setState = zustand.useStore(globalState, (state) => state.setState);
+
+    return React.useMemo(() => {
+        /**
+         * Refresh the follow-up suggestions.
+         */
+        const fetchFollowUpSuggestions = async (previousResponseId: string) => {
+            const stream = await streamAIChatFollowUpResponses({
+                previousResponseId,
+            });
+
+            for await (const suggestions of stream) {
+                setState((state) => ({ ...state, followUpSuggestions: suggestions }));
+            }
+        };
+
+        return {
+            open: () => setState((state) => ({ ...state, opened: true })),
+            close: () => setState((state) => ({ ...state, opened: false })),
+            clear: () =>
+                setState((state) => ({
+                    opened: state.opened,
+                    loading: false,
+                    messages: [],
+                    followUpSuggestions: [],
+                    responseId: null,
+                })),
+            postMessage: async (input: { message: string }) => {
+                setState((state) => {
+                    return {
+                        ...state,
+                        messages: [
+                            ...state.messages,
+                            {
+                                // TODO: how to handle markdown here?
+                                // to avoid rendering as plain text
+                                role: AIMessageRole.User,
+                                content: input.message,
+                            },
+                            {
+                                role: AIMessageRole.Assistant,
+                                content: null,
+                            },
+                        ],
+                        followUpSuggestions: [],
+                        loading: true,
+                    };
+                });
+
+                const stream = await streamAIChatResponse({
+                    message: input.message,
+                    messageContext: messageContextRef.current,
+                    previousResponseId: globalState.getState().state.responseId ?? undefined,
+                });
+
+                for await (const data of stream) {
+                    if (!data) continue;
+
+                    const event = data.event;
+                    if (event.type === 'response_finish') {
+                        setState((state) => ({ ...state, responseId: event.responseId }));
+
+                        fetchFollowUpSuggestions(event.responseId);
+                    }
+
+                    setState((state) => ({
+                        ...state,
+                        messages: [
+                            ...state.messages.slice(0, -1),
+                            {
+                                role: AIMessageRole.Assistant,
+                                content: data.content,
+                            },
+                        ],
+                    }));
+                }
+
+                setState((state) => ({
+                    ...state,
+                    loading: false,
+                }));
+            },
+        };
+    }, [messageContextRef, setState]);
+}

packages/gitbook/src/components/AI/useAIMessageContext.ts

@@ -0,0 +1,35 @@
+import type { AIMessageContext } from '@gitbook/api';
+import React from 'react';
+import { useCurrentPage } from '../hooks';
+
+/**
+ * Return the context for the AI message.
+ */
+export function useAIMessageContext(): AIMessageContext {
+    const currentPage = useCurrentPage();
+
+    return React.useMemo(() => {
+        return {
+            location: currentPage
+                ? {
+                      spaceId: currentPage.spaceId,
+                      pageId: currentPage.pageId,
+                  }
+                : undefined,
+        };
+    }, [currentPage]);
+}
+
+/**
+ * Return the context for the AI message as a mutable React ref
+ */
+export function useAIMessageContextRef(): React.MutableRefObject<AIMessageContext> {
+    const context = useAIMessageContext();
+    const ref = React.useRef(context);
+
+    React.useEffect(() => {
+        ref.current = context;
+    }, [context]);
+
+    return ref;
+}

packages/gitbook/src/components/AI/useAIPage.tsx

@@ -0,0 +1,130 @@
+'use client';
+
+import React from 'react';
+import {
+    type AIMessageRenderStream,
+    streamAIResponseById,
+    streamGenerateAIPage,
+} from './server-actions';
+
+export type AIPageState = {
+    /**
+     * The body of the page.
+     */
+    body: React.ReactNode;
+
+    /**
+     * The ID of the latest AI response.
+     */
+    responseId: string | null;
+};
+
+export type AIPageController = {
+    /**
+     * Generate a new page for a query.
+     */
+    generate: (query: string) => void;
+};
+
+/**
+ * Hook to generate a page using AI.
+ */
+export function useAIPage(
+    props: {
+        initialResponseId?: string;
+    } = {}
+): [AIPageState, AIPageController] {
+    const { initialResponseId } = props;
+    const [responseId, setResponseId] = React.useState<string | null>(null);
+    const [body, setBody] = React.useState<React.ReactNode>('');
+    const currentStreamRef = React.useRef<AIMessageRenderStream | null>(null);
+    const lastResponseIdRef = React.useRef<string | undefined>(props.initialResponseId);
+
+    /**
+     * Update the page body with the content of the stream.
+     */
+    const generateFromStream = React.useCallback(
+        async (rawStream: AIMessageRenderStream | Promise<AIMessageRenderStream>) => {
+            currentStreamRef.current = null;
+            const stream = await rawStream;
+            if (currentStreamRef.current) {
+                // If there's already a stream, we don't want to process this one.
+                return;
+            }
+            currentStreamRef.current = stream;
+
+            try {
+                for await (const data of stream) {
+                    if (currentStreamRef.current !== stream) {
+                        // If the stream has changed, we don't want to process this one.
+                        return;
+                    }
+                    if (!data) continue;
+
+                    setBody(data.content);
+
+                    switch (data.event.type) {
+                        case 'response_finish':
+                            lastResponseIdRef.current = data.event.responseId;
+                            setResponseId(data.event.responseId);
+                            break;
+                    }
+                }
+            } catch (error) {
+                console.error('Error in summary stream:', error);
+            }
+        },
+        []
+    );
+
+    /**
+     * Initialize the page with the initial response id
+     */
+    React.useEffect(() => {
+        if (initialResponseId) {
+            generateFromStream(
+                streamAIResponseById({
+                    responseId: initialResponseId,
+                    options: {
+                        renderToolCalls: false,
+                    },
+                })
+            );
+        }
+    }, [generateFromStream, initialResponseId]);
+
+    /**
+     * Generate a new page for a query.
+     */
+    const generate = React.useCallback(
+        async (query: string) => {
+            generateFromStream(
+                streamGenerateAIPage({
+                    query,
+                    previousResponseId: lastResponseIdRef.current,
+                    options: {
+                        renderToolCalls: false,
+                    },
+                })
+            );
+        },
+        [generateFromStream]
+    );
+
+    const state = React.useMemo(
+        () => ({
+            body,
+            responseId,
+        }),
+        [body, responseId]
+    );
+
+    const controller = React.useMemo(
+        () => ({
+            generate,
+        }),
+        [generate]
+    );
+
+    return [state, controller];
+}

packages/gitbook/src/components/AIPageLinkSummary/AIPageLinkSummary.tsx

@@ -5,8 +5,8 @@ import { Icon } from '@gitbook/icons';
 import { useEffect } from 'react';
 import { create } from 'zustand';
 import { useShallow } from 'zustand/react/shallow';
-import { useVisitedPages } from '../Insights';
 import { usePageContext } from '../PageContext';
+import { useVisitedPages } from '../hooks';
 import { Loading } from '../primitives';
 import { streamLinkPageSummary } from './server-actions/streamLinkPageSummary';
 
@@ -116,7 +116,7 @@ export function AIPageLinkSummary(props: {
 
     const currentPage = usePageContext();
     const language = useLanguage();
-    const visitedPages = useVisitedPages((state) => state.pages);
+    const visitedPages = useVisitedPages();
     const { summary, streamSummary } = useSummaries(
         useShallow((state) => {
             return {

packages/gitbook/src/components/AIPageLinkSummary/server-actions/streamLinkPageSummary.ts

@@ -2,11 +2,11 @@
 import { filterOutNullable } from '@/lib/typescript';
 import { getV1BaseContext } from '@/lib/v1';
 import { isV2 } from '@/lib/v2';
-import { AIMessageRole } from '@gitbook/api';
+import { AIMessageRole, AIModel } from '@gitbook/api';
 import { getSiteURLDataFromMiddleware } from '@v2/lib/middleware';
 import { getServerActionBaseContext } from '@v2/lib/server-actions';
 import { z } from 'zod';
-import { streamGenerateObject } from './api';
+import { streamGenerateAIObject } from '../../AI/server-actions/api';
 
 /**
  * Get a summary of a page, in the context of another page
@@ -32,23 +32,17 @@ export async function* streamLinkPageSummary({
     const baseContext = isV2() ? await getServerActionBaseContext() : await getV1BaseContext();
     const siteURLData = await getSiteURLDataFromMiddleware();
 
-    const { stream } = await streamGenerateObject(
-        baseContext,
-        {
-            organizationId: siteURLData.organization,
-            siteId: siteURLData.site,
-        },
-        {
-            schema: z.object({
-                highlight: z
-                    .string()
-                    .describe('The reason why the user should read the target page.'),
-                // questions: z.array(z.string().describe('The questions to sea')).max(3),
-            }),
-            messages: [
-                {
-                    role: AIMessageRole.Developer,
-                    content: `# 1. Role
+    const { stream } = await streamGenerateAIObject(baseContext, {
+        organizationId: siteURLData.organization,
+        siteId: siteURLData.site,
+        model: AIModel.Fast,
+        schema: z.object({
+            highlight: z.string().describe('The reason why the user should read the target page.'),
+        }),
+        input: [
+            {
+                role: AIMessageRole.Developer,
+                content: `# 1. Role
 You are a contextual fact extractor. Your job is to find the exact fact from the linked page that directly answers the implied question in the current paragraph.
 
 # 2. Task
@@ -64,60 +58,60 @@ Extract a contextually-relevant fact that:
 3. Ensure the fact relates directly to the context of the paragraph containing the link
 4. Avoid ALL instructional language including words like "use", "click", "select", "create"
 5. Keep it under 30 words, factual and declarative about what EXISTS or IS TRUE`,
-                },
-                {
-                    role: AIMessageRole.Developer,
-                    content: `# 4. Current page
+            },
+            {
+                role: AIMessageRole.Developer,
+                content: `# 4. Current page
 The content of the current page is:`,
-                    attachments: [
-                        {
-                            type: 'page' as const,
-                            spaceId: currentSpaceId,
-                            pageId: currentPageId,
-                        },
-                    ],
-                },
-                ...(visitedPages
-                    ? [
-                          {
-                              role: AIMessageRole.Developer,
-                              content: '# 5. Previous pages',
-                          },
-                          ...visitedPages.map(({ spaceId, pageId }) => ({
-                              role: AIMessageRole.Developer,
-                              content: `## Page ${pageId}`,
-                              attachments: [
-                                  {
-                                      type: 'page' as const,
-                                      spaceId,
-                                      pageId,
-                                  },
-                              ],
-                          })),
-                      ]
-                    : []),
-                {
-                    role: AIMessageRole.Developer,
-                    content: `# 6. Target page
+                attachments: [
+                    {
+                        type: 'page' as const,
+                        spaceId: currentSpaceId,
+                        pageId: currentPageId,
+                    },
+                ],
+            },
+            ...(visitedPages
+                ? [
+                      {
+                          role: AIMessageRole.Developer,
+                          content: '# 5. Previous pages',
+                      },
+                      ...visitedPages.map(({ spaceId, pageId }) => ({
+                          role: AIMessageRole.Developer,
+                          content: `## Page ${pageId}`,
+                          attachments: [
+                              {
+                                  type: 'page' as const,
+                                  spaceId,
+                                  pageId,
+                              },
+                          ],
+                      })),
+                  ]
+                : []),
+            {
+                role: AIMessageRole.Developer,
+                content: `# 6. Target page
 The content of the target page is:`,
-                    attachments: [
-                        {
-                            type: 'page' as const,
-                            spaceId: targetSpaceId,
-                            pageId: targetPageId,
-                        },
-                    ],
-                },
-                {
-                    role: AIMessageRole.Developer,
-                    content: `# 7. Link preview
+                attachments: [
+                    {
+                        type: 'page' as const,
+                        spaceId: targetSpaceId,
+                        pageId: targetPageId,
+                    },
+                ],
+            },
+            {
+                role: AIMessageRole.Developer,
+                content: `# 7. Link preview
 The content of the link preview is:
 > ${linkPreview}
 > Page ID: ${targetPageId}`,
-                },
-                {
-                    role: AIMessageRole.Developer,
-                    content: `# 8. Guidelines & Examples
+            },
+            {
+                role: AIMessageRole.Developer,
+                content: `# 8. Guidelines & Examples
 ALWAYS:
 - ALWAYS choose facts that directly fulfill the contextual need where the link appears
 - ALWAYS connect target page information specifically to the current paragraph context
@@ -146,14 +140,13 @@ Current paragraph: "Your team mentioned issues with conflicting edits. Need to c
 Preview: "Live Edit: Real-time collaborative editing."
 ✓ "Teams with GitHub repositories (like yours) cannot use this feature due to sync limitations."
 ✗ "Incompatible with GitHub/GitLab sync and requires specific visibility settings."`,
-                },
-                {
-                    role: AIMessageRole.User,
-                    content: `I'm considering reading the link titled "${linkTitle}" pointing to page ${targetPageId}. Why should I read it? Relate it to the paragraph I'm currently reading.`,
-                },
-            ].filter(filterOutNullable),
-        }
-    );
+            },
+            {
+                role: AIMessageRole.User,
+                content: `I'm considering reading the link titled "${linkTitle}" pointing to page ${targetPageId}. Why should I read it? Relate it to the paragraph I'm currently reading.`,
+            },
+        ].filter(filterOutNullable),
+    });
 
     for await (const value of stream) {
         const highlight = value.highlight;

packages/gitbook/src/components/Adaptive/server-actions/api.ts

@@ -1,124 +0,0 @@
-'use server';
-import { type AIMessageInput, AIModel, type AIStreamResponse } from '@gitbook/api';
-import type { GitBookBaseContext } from '@v2/lib/context';
-import { EventIterator } from 'event-iterator';
-import type { MaybePromise } from 'p-map';
-import * as partialJson from 'partial-json';
-import type { DeepPartial } from 'ts-essentials';
-import type { z } from 'zod';
-import { zodToJsonSchema } from 'zod-to-json-schema';
-
-/**
- * Get the latest value from a stream and the response id.
- */
-export async function generate<T>(
-    promise: MaybePromise<{
-        stream: EventIterator<T>;
-        response: Promise<{ responseId: string }>;
-    }>
-) {
-    const input = await promise;
-    let value: T | undefined;
-
-    for await (const event of input.stream) {
-        value = event;
-    }
-
-    const { responseId } = await input.response;
-    return {
-        responseId,
-        value,
-    };
-}
-
-/**
- * Stream the generation of an object using the AI.
- */
-export async function streamGenerateObject<T>(
-    context: GitBookBaseContext,
-    {
-        organizationId,
-        siteId,
-    }: {
-        organizationId: string;
-        siteId: string;
-    },
-    {
-        schema,
-        messages,
-        model = AIModel.Fast,
-    }: {
-        schema: z.ZodSchema<T>;
-        messages: AIMessageInput[];
-        model?: AIModel;
-        previousResponseId?: string;
-    }
-) {
-    const rawStream = context.dataFetcher.streamAIResponse({
-        organizationId,
-        siteId,
-        input: messages,
-        output: {
-            type: 'object',
-            schema: zodToJsonSchema(schema),
-        },
-        model,
-    });
-
-    let json = '';
-    return parseResponse<DeepPartial<T>>(rawStream, (event) => {
-        if (event.type === 'response_object') {
-            json += event.jsonChunk;
-
-            const parsed = partialJson.parse(json, partialJson.ALL);
-            return parsed;
-        }
-    });
-}
-
-/**
- * Parse a stream from the API to extract the responseId.
- */
-function parseResponse<T>(
-    responseStream: EventIterator<AIStreamResponse>,
-    parse: (response: AIStreamResponse) => T | undefined
-): {
-    stream: EventIterator<T>;
-    response: Promise<{ responseId: string }>;
-} {
-    let resolveResponse: (value: { responseId: string }) => void;
-    const response = new Promise<{ responseId: string }>((resolve) => {
-        resolveResponse = resolve;
-    });
-
-    const stream = new EventIterator<T>((queue) => {
-        (async () => {
-            let foundResponse = false;
-
-            for await (const event of responseStream) {
-                if (event.type === 'response_finish') {
-                    foundResponse = true;
-                    resolveResponse({ responseId: event.responseId });
-                } else {
-                    const parsed = parse(event);
-                    if (parsed !== undefined) {
-                        queue.push(parsed);
-                    }
-                }
-            }
-
-            if (!foundResponse) {
-                throw new Error('No response found');
-            }
-        })().then(
-            () => {
-                queue.stop();
-            },
-            (error) => {
-                queue.fail(error);
-            }
-        );
-    });
-
-    return { stream, response };
-}

packages/gitbook/src/components/DocumentView/InlineLink/InlineLinkTooltipImpl.tsx

@@ -3,7 +3,7 @@ import { tcls } from '@/lib/tailwind';
 import { Icon } from '@gitbook/icons';
 import * as Tooltip from '@radix-ui/react-tooltip';
 import { Fragment } from 'react';
-import { AIPageLinkSummary } from '../../Adaptive';
+import { AIPageLinkSummary } from '../../AIPageLinkSummary';
 import { Button, StyledLink } from '../../primitives';
 
 export function InlineLinkTooltipImpl(props: {

packages/gitbook/src/components/Insights/InsightsProvider.tsx

@@ -5,28 +5,13 @@ import { OpenAPIOperationContextProvider } from '@gitbook/react-openapi';
 import * as React from 'react';
 import { useDebounceCallback, useEventCallback } from 'usehooks-ts';
 
-import type { VisitorAuthClaims } from '@/lib/adaptive';
 import { getAllBrowserCookiesMap } from '@/lib/browser-cookies';
+import { type CurrentContentContext, useCurrentContent } from '../hooks';
 import { getSession } from './sessions';
-import { useVisitedPages } from './useVisitedPages';
 import { getVisitorId } from './visitorId';
 
 export type InsightsEventName = api.SiteInsightsEvent['type'];
 
-/**
- * Global context for all events in the session.
- */
-type InsightsEventContext = {
-    organizationId: string;
-    siteId: string;
-    siteSectionId: string | null;
-    siteSpaceId: string | null;
-    siteShareKey: string | null;
-    spaceId: string;
-    revisionId: string;
-    visitorAuthClaims: VisitorAuthClaims;
-};
-
 /**
  * Context for an event on a page.
  */
@@ -63,7 +48,7 @@ type TrackEventCallback = <EventName extends InsightsEventName>(
 
 const InsightsContext = React.createContext<TrackEventCallback>(() => {});
 
-interface InsightsProviderProps extends InsightsEventContext {
+interface InsightsProviderProps {
     /** If true, the events will be sent to the server. */
     enabled: boolean;
 
@@ -84,16 +69,16 @@ interface InsightsProviderProps extends InsightsEventContext {
  * Wrap the content of the app with the InsightsProvider to track events.
  */
 export function InsightsProvider(props: InsightsProviderProps) {
-    const { enabled, appURL, apiHost, children, visitorCookieTrackingEnabled, ...context } = props;
+    const { enabled, appURL, apiHost, children, visitorCookieTrackingEnabled } = props;
 
-    const addVisitedPage = useVisitedPages((state) => state.addPage);
+    const currentContent = useCurrentContent();
     const visitorIdRef = React.useRef<string | null>(null);
     const eventsRef = React.useRef<{
         [pathname: string]:
             | {
                   url: string;
                   events: TrackEventInput<InsightsEventName>[];
-                  context: InsightsEventContext;
+                  context: CurrentContentContext;
                   pageContext?: InsightsEventPageContext;
               }
             | undefined;
@@ -124,7 +109,7 @@ export function InsightsProvider(props: InsightsProviderProps) {
                 ...transformEvents({
                     url: eventsForPathname.url,
                     events: eventsForPathname.events,
-                    context,
+                    context: currentContent,
                     pageContext: eventsForPathname.pageContext,
                     visitorId,
                     sessionId: session.id,
@@ -136,22 +121,14 @@ export function InsightsProvider(props: InsightsProviderProps) {
                 ...eventsForPathname,
                 events: [],
             };
-
-            // Mark the page as visited in our local state
-            if (eventsForPathname.pageContext.pageId) {
-                addVisitedPage({
-                    spaceId: context.spaceId,
-                    pageId: eventsForPathname.pageContext.pageId,
-                });
-            }
         }
 
         if (allEvents.length > 0) {
             if (enabled) {
                 sendEvents({
                     apiHost,
-                    organizationId: context.organizationId,
-                    siteId: context.siteId,
+                    organizationId: currentContent.organizationId,
+                    siteId: currentContent.siteId,
                     events: allEvents,
                 });
             } else {
@@ -185,7 +162,7 @@ export function InsightsProvider(props: InsightsProviderProps) {
                         timestamp: new Date().toISOString(),
                     },
                 ],
-                context,
+                context: currentContent,
             };
 
             if (eventsRef.current[pathname].pageContext !== undefined) {
@@ -222,7 +199,7 @@ export function InsightsProvider(props: InsightsProviderProps) {
                     trackEvent({ type: 'api_client_open', operation });
                 }}
             >
-                {props.children}
+                {children}
             </OpenAPIOperationContextProvider>
         </InsightsContext.Provider>
     );
@@ -269,7 +246,7 @@ function sendEvents(args: {
 function transformEvents(input: {
     url: string;
     events: TrackEventInput<InsightsEventName>[];
-    context: InsightsEventContext;
+    context: CurrentContentContext;
     pageContext: InsightsEventPageContext;
     visitorId: string;
     sessionId: string;

packages/gitbook/src/components/Insights/TrackPageViewEvent.tsx

@@ -2,13 +2,14 @@
 
 import * as React from 'react';
 
-import { type InsightsEventPageContext, useTrackEvent } from './InsightsProvider';
+import { useCurrentPage } from '../hooks';
+import { useTrackEvent } from './InsightsProvider';
 
 /**
  * Track a page view event.
  */
-export function TrackPageViewEvent(props: InsightsEventPageContext) {
-    const { pageId } = props;
+export function TrackPageViewEvent() {
+    const page = useCurrentPage();
     const trackEvent = useTrackEvent();
 
     React.useEffect(() => {
@@ -17,10 +18,10 @@ export function TrackPageViewEvent(props: InsightsEventPageContext) {
                 type: 'page_view',
             },
             {
-                pageId,
+                pageId: page?.pageId ?? null,
             }
         );
-    }, [pageId, trackEvent]);
+    }, [page, trackEvent]);
 
     return null;
 }

packages/gitbook/src/components/Insights/index.ts

@@ -2,4 +2,3 @@ export * from './InsightsProvider';
 export * from './visitorId';
 export * from './cookies';
 export * from './TrackPageViewEvent';
-export * from './useVisitedPages';

packages/gitbook/src/components/Insights/useVisitedPages.tsx

@@ -1,25 +0,0 @@
-import { create } from 'zustand';
-
-type VisitedPage = {
-    spaceId: string;
-    pageId: string;
-};
-
-/**
- * A store for the pages that have been visited in the current session.
- */
-export const useVisitedPages = create<{
-    pages: VisitedPage[];
-    addPage: (page: VisitedPage) => void;
-}>((set) => ({
-    pages: [],
-    addPage: (page) =>
-        set((state) => {
-            const lastPage = state.pages[state.pages.length - 1];
-            if (lastPage && lastPage.spaceId === page.spaceId && lastPage.pageId === page.pageId) {
-                return { pages: state.pages };
-            }
-
-            return { pages: [...state.pages, page] };
-        }),
-}));

packages/gitbook/src/components/PageBody/PageBody.tsx

@@ -10,6 +10,7 @@ import { tcls } from '@/lib/tailwind';
 import { DocumentView, DocumentViewSkeleton } from '../DocumentView';
 import { TrackPageViewEvent } from '../Insights';
 import { PageFeedbackForm } from '../PageFeedback';
+import { CurrentPageProvider } from '../hooks/useCurrentPage';
 import { DateRelative } from '../primitives';
 import { PageBodyBlankslate } from './PageBodyBlankslate';
 import { PageCover } from './PageCover';
@@ -45,7 +46,7 @@ export function PageBody(props: {
     const updatedAt = page.updatedAt ?? page.createdAt;
 
     return (
-        <>
+        <CurrentPageProvider page={{ spaceId: context.space.id, pageId: page.id }}>
             <main
                 className={tcls(
                     'relative min-w-0 flex-1',
@@ -106,7 +107,7 @@ export function PageBody(props: {
                 </div>
             </main>
 
-            <TrackPageViewEvent pageId={page.id} />
-        </>
+            <TrackPageViewEvent />
+        </CurrentPageProvider>
     );
 }

packages/gitbook/src/components/SitePage/SitePageNotFound.tsx

@@ -6,6 +6,7 @@ import { tcls } from '@/lib/tailwind';
 import { useRouter, useSearchParams } from 'next/navigation';
 import { useEffect } from 'react';
 import { useSpaceBasePath } from '../SpaceLayout/SpaceLayoutContext';
+import { CurrentPageProvider } from '../hooks';
 
 /**
  * Component that displays a "page not found" message.
@@ -25,26 +26,28 @@ export function SitePageNotFound() {
     }, [basePath, fallback, router]);
 
     return (
-        <div
-            className={tcls(
-                'flex-1',
-                'flex',
-                'flex-row',
-                'items-center',
-                'justify-center',
-                'py-9',
-                'min-h-[calc(100vh-64px)] lg:min-h-fit'
-            )}
-        >
-            <div className={tcls('max-w-80')}>
-                <h2 className={tcls('text-2xl', 'font-semibold', 'mb-2')}>
-                    {t(language, 'notfound_title')}
-                </h2>
-                <p className={tcls('text-base', 'mb-4')}>{t(language, 'notfound')}</p>
-            </div>
+        <CurrentPageProvider page={null}>
+            <div
+                className={tcls(
+                    'flex-1',
+                    'flex',
+                    'flex-row',
+                    'items-center',
+                    'justify-center',
+                    'py-9',
+                    'min-h-[calc(100vh-64px)] lg:min-h-fit'
+                )}
+            >
+                <div className={tcls('max-w-80')}>
+                    <h2 className={tcls('text-2xl', 'font-semibold', 'mb-2')}>
+                        {t(language, 'notfound_title')}
+                    </h2>
+                    <p className={tcls('text-base', 'mb-4')}>{t(language, 'notfound')}</p>
+                </div>
 
-            {/* Track the page not found as a page view */}
-            <TrackPageViewEvent pageId={null} />
-        </div>
+                {/* Track the page not found as a page view */}
+                <TrackPageViewEvent />
+            </div>
+        </CurrentPageProvider>
     );
 }

packages/gitbook/src/components/SpaceLayout/SpaceLayout.tsx

@@ -17,6 +17,7 @@ import { Announcement } from '../Announcement';
 import { SpacesDropdown } from '../Header/SpacesDropdown';
 import { InsightsProvider } from '../Insights';
 import { SiteSectionList, encodeClientSiteSections } from '../SiteSections';
+import { CurrentContentProvider } from '../hooks';
 import { SpaceLayoutContextProvider } from './SpaceLayoutContext';
 
 /**
@@ -50,110 +51,116 @@ export function SpaceLayout(props: {
 
     return (
         <SpaceLayoutContextProvider basePath={context.linker.toPathInSpace('')}>
-            <InsightsProvider
-                enabled={withTracking}
-                appURL={GITBOOK_APP_URL}
-                apiHost={GITBOOK_API_PUBLIC_URL}
+            <CurrentContentProvider
                 organizationId={context.organizationId}
                 siteId={context.site.id}
                 siteSectionId={context.sections?.current?.id ?? null}
                 siteSpaceId={context.siteSpace.id}
                 siteShareKey={context.shareKey ?? null}
-                revisionId={context.revisionId}
                 spaceId={context.space.id}
+                revisionId={context.revisionId}
                 visitorAuthClaims={visitorAuthClaims}
-                visitorCookieTrackingEnabled={context.customization.insights?.trackingCookie}
             >
-                <Announcement context={context} />
-                <Header withTopHeader={withTopHeader} context={context} />
-                <div className="scroll-nojump">
-                    <div
-                        className={tcls(
-                            'flex',
-                            'flex-col',
-                            'lg:flex-row',
-                            CONTAINER_STYLE,
-                            'site-full-width:max-w-full',
+                <InsightsProvider
+                    enabled={withTracking}
+                    appURL={GITBOOK_APP_URL}
+                    apiHost={GITBOOK_API_PUBLIC_URL}
+                    visitorCookieTrackingEnabled={context.customization.insights?.trackingCookie}
+                >
+                    <Announcement context={context} />
+                    <Header withTopHeader={withTopHeader} context={context} />
+                    <div className="scroll-nojump">
+                        <div
+                            className={tcls(
+                                'flex',
+                                'flex-col',
+                                'lg:flex-row',
+                                CONTAINER_STYLE,
+                                'site-full-width:max-w-full',
 
-                            // Ensure the footer is display below the viewport even if the content is not enough
-                            withFooter && 'min-h-[calc(100vh-64px)]',
-                            withTopHeader ? null : 'lg:min-h-screen'
-                        )}
-                    >
-                        <TableOfContents
-                            context={context}
-                            header={
-                                withTopHeader ? null : (
-                                    <div
-                                        className={tcls(
-                                            'hidden',
-                                            'pr-4',
-                                            'lg:flex',
-                                            'grow-0',
-                                            'flex-wrap',
-                                            'dark:shadow-light/1',
-                                            'text-base/tight'
-                                        )}
-                                    >
-                                        <HeaderLogo context={context} />
-                                    </div>
-                                )
-                            }
-                            innerHeader={
-                                // displays the search button and/or the space dropdown in the ToC according to the header/variant settings. E.g if there is no header, the search button will be displayed in the ToC.
-                                <>
-                                    {!withTopHeader && (
-                                        <div className={tcls('hidden', 'lg:block')}>
-                                            <React.Suspense fallback={null}>
-                                                <SearchButton>
-                                                    <span className={tcls('flex-1')}>
-                                                        {t(
-                                                            getSpaceLanguage(customization),
-                                                            customization.aiSearch.enabled
-                                                                ? 'search_or_ask'
-                                                                : 'search'
-                                                        )}
-                                                        ...
-                                                    </span>
-                                                </SearchButton>
-                                            </React.Suspense>
-                                        </div>
-                                    )}
-                                    {!withTopHeader && withSections && sections && (
-                                        <SiteSectionList
-                                            className={tcls('hidden', 'lg:block')}
-                                            sections={encodeClientSiteSections(context, sections)}
-                                        />
-                                    )}
-                                    {isMultiVariants && !sections && (
-                                        <SpacesDropdown
-                                            context={context}
-                                            siteSpace={siteSpace}
-                                            siteSpaces={siteSpaces}
+                                // Ensure the footer is display below the viewport even if the content is not enough
+                                withFooter && 'min-h-[calc(100vh-64px)]',
+                                withTopHeader ? null : 'lg:min-h-screen'
+                            )}
+                        >
+                            <TableOfContents
+                                context={context}
+                                header={
+                                    withTopHeader ? null : (
+                                        <div
                                             className={tcls(
-                                                'w-full',
-                                                'page-no-toc:hidden',
-                                                'site-header-none:page-no-toc:flex'
+                                                'hidden',
+                                                'pr-4',
+                                                'lg:flex',
+                                                'grow-0',
+                                                'flex-wrap',
+                                                'dark:shadow-light/1',
+                                                'text-base/tight'
                                             )}
-                                        />
-                                    )}
-                                </>
-                            }
-                        />
-                        <div className="flex min-w-0 flex-1 flex-col">{children}</div>
+                                        >
+                                            <HeaderLogo context={context} />
+                                        </div>
+                                    )
+                                }
+                                innerHeader={
+                                    // displays the search button and/or the space dropdown in the ToC according to the header/variant settings. E.g if there is no header, the search button will be displayed in the ToC.
+                                    <>
+                                        {!withTopHeader && (
+                                            <div className={tcls('hidden', 'lg:block')}>
+                                                <React.Suspense fallback={null}>
+                                                    <SearchButton>
+                                                        <span className={tcls('flex-1')}>
+                                                            {t(
+                                                                getSpaceLanguage(customization),
+                                                                customization.aiSearch.enabled
+                                                                    ? 'search_or_ask'
+                                                                    : 'search'
+                                                            )}
+                                                            ...
+                                                        </span>
+                                                    </SearchButton>
+                                                </React.Suspense>
+                                            </div>
+                                        )}
+                                        {!withTopHeader && withSections && sections && (
+                                            <SiteSectionList
+                                                className={tcls('hidden', 'lg:block')}
+                                                sections={encodeClientSiteSections(
+                                                    context,
+                                                    sections
+                                                )}
+                                            />
+                                        )}
+                                        {isMultiVariants && !sections && (
+                                            <SpacesDropdown
+                                                context={context}
+                                                siteSpace={siteSpace}
+                                                siteSpaces={siteSpaces}
+                                                className={tcls(
+                                                    'w-full',
+                                                    'page-no-toc:hidden',
+                                                    'site-header-none:page-no-toc:flex'
+                                                )}
+                                            />
+                                        )}
+                                    </>
+                                }
+                            />
+                            <div className="flex min-w-0 flex-1 flex-col">{children}</div>
+                        </div>
                     </div>
-                </div>
 
-                {withFooter ? <Footer context={context} /> : null}
+                    {withFooter ? <Footer context={context} /> : null}
 
-                <React.Suspense fallback={null}>
-                    <SearchModal
-                        spaceTitle={siteSpace.title}
-                        withAsk={customization.aiSearch.enabled}
-                        isMultiVariants={isMultiVariants}
-                    />
-                </React.Suspense>
-            </InsightsProvider>
+                    <React.Suspense fallback={null}>
+                        <SearchModal
+                            spaceTitle={siteSpace.title}
+                            withAsk={customization.aiSearch.enabled}
+                            isMultiVariants={isMultiVariants}
+                        />
+                    </React.Suspense>
+                </InsightsProvider>
+            </CurrentContentProvider>
         </SpaceLayoutContextProvider>
     );
 }

packages/gitbook/src/components/hooks/index.ts

@@ -4,3 +4,5 @@ export * from './useHash';
 export * from './useIsMounted';
 export * from './useToggleAnimation';
 export * from './useCurrentPagePath';
+export * from './useCurrentContent';
+export * from './useCurrentPage';

packages/gitbook/src/components/hooks/useCurrentContent.tsx

@@ -0,0 +1,65 @@
+'use client';
+
+import type { VisitorAuthClaims } from '@/lib/adaptive';
+import * as React from 'react';
+
+/**
+ * Global context for the current content.
+ */
+export type CurrentContentContext = {
+    organizationId: string;
+    siteId: string;
+    siteSectionId: string | null;
+    siteSpaceId: string | null;
+    siteShareKey: string | null;
+    spaceId: string;
+    revisionId: string;
+    visitorAuthClaims: VisitorAuthClaims;
+};
+
+const ReactCurrentContentContext = React.createContext<CurrentContentContext | null>(null);
+
+/**
+ * Hook to get the current content.
+ */
+export function useCurrentContent(): CurrentContentContext {
+    const context = React.useContext(ReactCurrentContentContext);
+    if (!context) {
+        throw new Error('useCurrentContent must be used within a CurrentContentProvider');
+    }
+
+    return context;
+}
+
+/**
+ * Provider for the current content.
+ */
+export function CurrentContentProvider(props: React.PropsWithChildren<CurrentContentContext>) {
+    const contextValue = React.useMemo(() => {
+        return {
+            organizationId: props.organizationId,
+            siteId: props.siteId,
+            siteSectionId: props.siteSectionId,
+            siteSpaceId: props.siteSpaceId,
+            siteShareKey: props.siteShareKey,
+            spaceId: props.spaceId,
+            revisionId: props.revisionId,
+            visitorAuthClaims: props.visitorAuthClaims,
+        };
+    }, [
+        props.organizationId,
+        props.siteId,
+        props.siteSectionId,
+        props.siteSpaceId,
+        props.siteShareKey,
+        props.spaceId,
+        props.revisionId,
+        props.visitorAuthClaims,
+    ]);
+
+    return (
+        <ReactCurrentContentContext.Provider value={contextValue}>
+            {props.children}
+        </ReactCurrentContentContext.Provider>
+    );
+}

packages/gitbook/src/components/hooks/useCurrentPage.tsx

@@ -0,0 +1,90 @@
+'use client';
+
+import * as React from 'react';
+import * as zustand from 'zustand';
+
+export type PagePointer = {
+    spaceId: string;
+    pageId: string;
+};
+
+const ReactCurrentPageContext = React.createContext<PagePointer | null | undefined>(undefined);
+
+/**
+ * A store for the pages that have been visited in the current session.
+ */
+const visitedPagesStore = zustand.create<{
+    pages: PagePointer[];
+    addPage: (page: PagePointer) => void;
+}>((set) => ({
+    pages: [],
+    addPage: (page) =>
+        set((state) => {
+            const lastPage = state.pages[state.pages.length - 1];
+            if (lastPage && lastPage.spaceId === page.spaceId && lastPage.pageId === page.pageId) {
+                return { pages: state.pages };
+            }
+
+            return { pages: [...state.pages, page] };
+        }),
+}));
+
+/**
+ * Provider for the current page.
+ */
+export function CurrentPageProvider(
+    props: React.PropsWithChildren<{
+        page: PagePointer | null;
+    }>
+) {
+    const { page, children } = props;
+    const addPage = zustand.useStore(visitedPagesStore, (state) => state.addPage);
+
+    const contextValue = React.useMemo(() => {
+        return page
+            ? {
+                  spaceId: page.spaceId,
+                  pageId: page.pageId,
+              }
+            : null;
+    }, [page]);
+
+    React.useEffect(() => {
+        if (contextValue) {
+            addPage(contextValue);
+        }
+    }, [contextValue, addPage]);
+
+    return (
+        <ReactCurrentPageContext.Provider value={contextValue}>
+            {children}
+        </ReactCurrentPageContext.Provider>
+    );
+}
+
+/**
+ * Get the pointer to the current page.
+ */
+export function useCurrentPage() {
+    const currentPage = React.useContext(ReactCurrentPageContext);
+    const lastVisitedPage = zustand.useStore(
+        visitedPagesStore,
+        (state) => state.pages[state.pages.length - 1]
+    );
+
+    if (currentPage !== undefined) {
+        return currentPage;
+    }
+
+    // If the context is undefined, we are in a layout component (outside the page component)
+    // We use the "deferred" value using the visited pages.
+    return lastVisitedPage ?? null;
+}
+
+/**
+ * Return the list of recently visited pages.
+ */
+export function useVisitedPages() {
+    const pages = zustand.useStore(visitedPagesStore, (state) => state.pages);
+    return pages;
+}

packages/gitbook/src/lib/v1.ts

@@ -306,10 +306,6 @@ function getDataFetcherV1(apiTokenOverride?: string): GitBookDataFetcher {
                 })
             );
         },
-
-        streamAIResponse() {
-            throw new Error('Not implemented in v1');
-        },
     };
 
     return dataFetcher;