Fix thread and display tool calls

SamyPesse · SamyPesse · commit d99630ac950b · 2025-06-24T15:03:49.000+02:00
diff --git a/packages/gitbook/src/components/AI/server-actions/AIMessageView.tsx b/packages/gitbook/src/components/AI/server-actions/AIMessageView.tsx
@@ -1,5 +1,7 @@
 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';
 
 /**
@@ -8,9 +10,10 @@ import type { RenderAIMessageOptions } from './types';
 export function AIMessageView(
     props: RenderAIMessageOptions & {
         message: AIMessage;
+        context: GitBookSiteContext;
     }
 ) {
-    const { message } = props;
+    const { message, context } = props;
 
     return (
         <div className="flex flex-col gap-2">
@@ -26,6 +29,9 @@ export function AIMessageView(
                             }}
                             style={['space-y-5']}
                         />
+                        {step.toolCalls && step.toolCalls.length > 0 ? (
+                            <AIToolCallsSummary toolCalls={step.toolCalls} context={context} />
+                        ) : null}
                     </div>
                 );
             })}
diff --git a/packages/gitbook/src/components/AI/server-actions/AIToolCallsSummary.tsx b/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>
+    );
+}
diff --git a/packages/gitbook/src/components/AI/server-actions/api.tsx b/packages/gitbook/src/components/AI/server-actions/api.tsx
@@ -8,6 +8,7 @@ import {
     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';
@@ -85,6 +86,7 @@ export async function streamGenerateAIObject<T>(
  * Stream the generation of a document.
  */
 export async function streamRenderAIMessage(
+    baseContext: GitBookBaseContext,
     rawStream: AsyncIterable<AIStreamResponse>,
     options?: RenderAIMessageOptions
 ) {
@@ -123,10 +125,13 @@ export async function streamRenderAIMessage(
         }
     };
 
+    // Fetch the full-context in the background to avoid blocking the stream.
+    const promiseContext = fetchServerActionSiteContext(baseContext);
+
     return parseResponse<{
         content: React.ReactNode;
         event: AIStreamResponse;
-    }>(rawStream, (event) => {
+    }>(rawStream, async (event) => {
         switch (event.type) {
             /**
              * The agent is processing a tool call in a new message.
@@ -172,7 +177,9 @@ export async function streamRenderAIMessage(
 
         return {
             event,
-            content: <AIMessageView message={message} {...options} />,
+            content: (
+                <AIMessageView message={message} context={await promiseContext} {...options} />
+            ),
         };
     });
 }
@@ -182,7 +189,7 @@ export async function streamRenderAIMessage(
  */
 function parseResponse<T>(
     responseStream: EventIterator<AIStreamResponse>,
-    parse: (response: AIStreamResponse) => T | undefined
+    parse: (response: AIStreamResponse) => T | undefined | Promise<T | undefined>
 ): {
     stream: EventIterator<T>;
     response: Promise<{ responseId: string }>;
@@ -197,14 +204,14 @@ function parseResponse<T>(
             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 });
-                } else {
-                    const parsed = parse(event);
-                    if (parsed !== undefined) {
-                        queue.push(parsed);
-                    }
                 }
             }
 
diff --git a/packages/gitbook/src/components/AI/server-actions/chat.ts b/packages/gitbook/src/components/AI/server-actions/chat.ts
@@ -28,33 +28,29 @@ export async function* streamAIChatResponse({
     previousResponseId?: string;
     options?: RenderAIMessageOptions;
 }) {
-    const { dataFetcher } = await getServerActionBaseContext();
+    const context = await getServerActionBaseContext();
     const siteURLData = await getSiteURLDataFromMiddleware();
 
-    const api = await dataFetcher.api();
-    const rawStream = await api.orgs.streamAiResponseInSite(
-        siteURLData.organization,
-        siteURLData.site,
-        {
-            input: [
-                {
-                    role: AIMessageRole.User,
-                    content: message,
-                },
-            ],
-            output: { type: 'document' },
-            model: AIModel.ReasoningLow,
-            instructions: PROMPT,
-            previousResponseId,
-            tools: {
-                getPageContent: true,
-                getPages: true,
-                search: true,
+    const api = await context.dataFetcher.api();
+    const rawStream = api.orgs.streamAiResponseInSite(siteURLData.organization, siteURLData.site, {
+        input: [
+            {
+                role: AIMessageRole.User,
+                content: message,
             },
-        }
-    );
-
-    const { stream } = await streamRenderAIMessage(rawStream, options);
+        ],
+        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;
diff --git a/packages/gitbook/src/components/AI/server-actions/pages.ts b/packages/gitbook/src/components/AI/server-actions/pages.ts
@@ -32,33 +32,29 @@ export async function* streamGenerateAIPage({
     previousResponseId?: string;
     options?: RenderAIMessageOptions;
 }) {
-    const { dataFetcher } = await getServerActionBaseContext();
+    const context = await getServerActionBaseContext();
     const siteURLData = await getSiteURLDataFromMiddleware();
 
-    const api = await dataFetcher.api();
-    const rawStream = await 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 api = await context.dataFetcher.api();
+    const rawStream = api.orgs.streamAiResponseInSite(siteURLData.organization, siteURLData.site, {
+        input: [
+            {
+                role: AIMessageRole.User,
+                content: query,
             },
-        }
-    );
-
-    const { stream } = await streamRenderAIMessage(rawStream, options);
+        ],
+        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;
diff --git a/packages/gitbook/src/components/AI/server-actions/responses.ts b/packages/gitbook/src/components/AI/server-actions/responses.ts
@@ -14,16 +14,16 @@ export async function* streamAIResponseById({
     responseId: string;
     options?: RenderAIMessageOptions;
 }) {
-    const { dataFetcher } = await getServerActionBaseContext();
+    const context = await getServerActionBaseContext();
     const siteURLData = await getSiteURLDataFromMiddleware();
 
-    const api = await dataFetcher.api();
-    const rawStream = await api.orgs.streamExistingAiResponseInSite(
+    const api = await context.dataFetcher.api();
+    const rawStream = api.orgs.streamExistingAiResponseInSite(
         siteURLData.organization,
         siteURLData.site,
         responseId
     );
-    const { stream } = await streamRenderAIMessage(rawStream, options);
+    const { stream } = await streamRenderAIMessage(context, rawStream, options);
 
     for await (const output of stream) {
         yield output;
