Hooks and APIs for AI v2 (#3385)

Author: SamyPesse

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>
+    );
+}