From 51f59063ac4573bbb8288fcc7aee51cf5a68e097 Mon Sep 17 00:00:00 2001 From: Zeno Kapitein Date: Wed, 7 May 2025 15:05:31 +0200 Subject: [PATCH] Add chat --- packages/gitbook-v2/src/lib/data/api.ts | 1 + packages/gitbook-v2/src/lib/data/types.ts | 1 + .../src/components/Adaptive/AIPageSummary.tsx | 140 +++++++++++++++++- .../src/components/Adaptive/AdaptivePane.tsx | 2 +- .../components/Adaptive/server-actions/api.ts | 2 + .../server-actions/streamPageQuestion.ts | 77 ++++++++++ .../server-actions/streamPageSummary.ts | 22 ++- .../src/components/primitives/Button.tsx | 16 +- 8 files changed, 250 insertions(+), 11 deletions(-) create mode 100644 packages/gitbook/src/components/Adaptive/server-actions/streamPageQuestion.ts diff --git a/packages/gitbook-v2/src/lib/data/api.ts b/packages/gitbook-v2/src/lib/data/api.ts index 0af1e9e36..13061c3af 100644 --- a/packages/gitbook-v2/src/lib/data/api.ts +++ b/packages/gitbook-v2/src/lib/data/api.ts @@ -1395,6 +1395,7 @@ async function* streamAIResponse( output: params.output, model: params.model, tools: params.tools, + previousResponseId: params.previousResponseId, }); for await (const event of res) { diff --git a/packages/gitbook-v2/src/lib/data/types.ts b/packages/gitbook-v2/src/lib/data/types.ts index a2b151938..0b6920523 100644 --- a/packages/gitbook-v2/src/lib/data/types.ts +++ b/packages/gitbook-v2/src/lib/data/types.ts @@ -190,5 +190,6 @@ export interface GitBookDataFetcher { output: api.AIOutputFormat; model: api.AIModel; tools?: api.AIToolCapabilities; + previousResponseId?: string; }): AsyncGenerator; } diff --git a/packages/gitbook/src/components/Adaptive/AIPageSummary.tsx b/packages/gitbook/src/components/Adaptive/AIPageSummary.tsx index a30d3fbdd..d3a3d74e1 100644 --- a/packages/gitbook/src/components/Adaptive/AIPageSummary.tsx +++ b/packages/gitbook/src/components/Adaptive/AIPageSummary.tsx @@ -2,9 +2,18 @@ import { useEffect, useRef, useState } from 'react'; import { useVisitedPages } from '../Insights'; import { usePageContext } from '../PageContext'; +import { Button } from '../primitives/Button'; import { useAdaptiveContext } from './AdaptiveContext'; +import { streamPageQuestion } from './server-actions/streamPageQuestion'; import { streamPageSummary } from './server-actions/streamPageSummary'; +interface ChatMessage { + type: 'question' | 'answer'; + content: string; +} + +type StreamData = { answer: string } | { newResponseId: string } | { toolUsage: boolean }; + export function AIPageSummary() { const { toggle, setLoading, setToggle } = useAdaptiveContext(); @@ -17,8 +26,66 @@ export function AIPageSummary() { bigPicture?: string; }>({}); + const [question, setQuestion] = useState(''); + const [chatHistory, setChatHistory] = useState([]); + const [isAsking, setIsAsking] = useState(false); + const [responseId, setResponseId] = useState(null); + const [showTypingIndicator, setShowTypingIndicator] = useState(false); + + const handleSubmit = async () => { + if (!question.trim() || isAsking) return; + + const currentQuestion = question; + setQuestion(''); + setIsAsking(true); + setShowTypingIndicator(true); + + // Add question to chat history + setChatHistory((prev) => [...prev, { type: 'question', content: currentQuestion }]); + + try { + const stream = await streamPageQuestion(currentQuestion, responseId ?? ''); + let currentAnswer = ''; + + for await (const data of stream as AsyncIterableIterator) { + if ('answer' in data && data.answer) { + currentAnswer = data.answer; + setShowTypingIndicator(false); + // Update the last message in chat history with the streaming answer + setChatHistory((prev) => { + const newHistory = [...prev]; + const lastMessage = newHistory[newHistory.length - 1]; + if (lastMessage?.type === 'answer') { + lastMessage.content = currentAnswer; + } else { + newHistory.push({ type: 'answer', content: currentAnswer }); + } + return newHistory; + }); + } else if ('newResponseId' in data && data.newResponseId) { + setResponseId(data.newResponseId); + } else if ('toolUsage' in data) { + // Show typing indicator when tools are being used + setShowTypingIndicator(true); + } + } + } finally { + setIsAsking(false); + setShowTypingIndicator(false); + } + }; + + const handleKeyPress = (e: React.KeyboardEvent) => { + if (e.key === 'Enter') { + handleSubmit(); + } + }; + useEffect(() => { if (!summary.keyFacts) setLoading(true); + }, [summary.keyFacts, setLoading]); + + useEffect(() => { if (!visitedPages?.length) return; // Skip if the visited pages haven't changed @@ -41,10 +108,17 @@ export function AIPageSummary() { visitedPages: visitedPages, }); - for await (const summary of stream) { + for await (const data of stream) { if (canceled) return; - setSummary(summary); + if ('responseId' in data && data.responseId !== undefined) { + setResponseId(data.responseId); + } + + setSummary((prev) => ({ + keyFacts: data.keyFacts ?? prev.keyFacts, + bigPicture: data.bigPicture ?? prev.bigPicture, + })); } })().finally(() => { setLoading(false); @@ -53,7 +127,7 @@ export function AIPageSummary() { return () => { canceled = true; }; - }, [currentPage, visitedPages, toggle, setLoading, setToggle]); + }, [currentPage, visitedPages, setLoading]); const shimmerBlocks = [20, 35, 25, 10, 45, 30, 30, 35, 25, 10, 40, 30]; // Widths in percentages @@ -91,6 +165,66 @@ export function AIPageSummary() { {summary?.bigPicture} ) : null} + + {chatHistory.length > 0 && ( +
+ {chatHistory.map((message) => ( +
+
+ {message.content} +
+
+ ))} + + {showTypingIndicator && ( +
+
+ + + • + + + • + +
+
+ )} +
+ )} + +
+ setQuestion(e.target.value)} + onKeyPress={handleKeyPress} + disabled={isAsking || !responseId} + /> +
) ); diff --git a/packages/gitbook/src/components/Adaptive/AdaptivePane.tsx b/packages/gitbook/src/components/Adaptive/AdaptivePane.tsx index f0e6f1a08..a719d3b21 100644 --- a/packages/gitbook/src/components/Adaptive/AdaptivePane.tsx +++ b/packages/gitbook/src/components/Adaptive/AdaptivePane.tsx @@ -10,7 +10,7 @@ export function AdaptivePane() { return (
diff --git a/packages/gitbook/src/components/Adaptive/server-actions/api.ts b/packages/gitbook/src/components/Adaptive/server-actions/api.ts index fdc642c15..74d54295e 100644 --- a/packages/gitbook/src/components/Adaptive/server-actions/api.ts +++ b/packages/gitbook/src/components/Adaptive/server-actions/api.ts @@ -51,6 +51,7 @@ export async function streamGenerateObject( { schema, messages, + previousResponseId, model = AIModel.Fast, tools = {}, }: { @@ -64,6 +65,7 @@ export async function streamGenerateObject( const rawStream = context.dataFetcher.streamAIResponse({ organizationId, siteId, + previousResponseId, input: messages, output: { type: 'object', diff --git a/packages/gitbook/src/components/Adaptive/server-actions/streamPageQuestion.ts b/packages/gitbook/src/components/Adaptive/server-actions/streamPageQuestion.ts new file mode 100644 index 000000000..3790abe4f --- /dev/null +++ b/packages/gitbook/src/components/Adaptive/server-actions/streamPageQuestion.ts @@ -0,0 +1,77 @@ +'use server'; +import { getV1BaseContext } from '@/lib/v1'; +import { isV2 } from '@/lib/v2'; +import { AIMessageRole } from '@gitbook/api'; +import { getSiteURLDataFromMiddleware } from '@v2/lib/middleware'; +import { fetchServerActionSiteContext, getServerActionBaseContext } from '@v2/lib/server-actions'; +import { z } from 'zod'; +import { streamGenerateObject } from './api'; + +/** + * Get a summary of a page, in the context of another page + */ +export async function* streamPageQuestion(question: string, responseId: string) { + const baseContext = isV2() ? await getServerActionBaseContext() : await getV1BaseContext(); + const siteURLData = await getSiteURLDataFromMiddleware(); + + const [{ stream, response }] = await Promise.all([ + streamGenerateObject( + baseContext, + { + organizationId: siteURLData.organization, + siteId: siteURLData.site, + }, + { + schema: z.object({ + answer: z.string().describe('The answer to the question'), + }), + previousResponseId: responseId, + tools: { + search: true, + getPageContent: true, + getPages: true, + }, + messages: [ + { + role: AIMessageRole.Developer, + content: + 'The user is asking a question about the page. Use your knowledge of the page and the context to answer the question. Be succinct in your answers, do not repeat information already in the key facts or big picture.', + }, + { + role: AIMessageRole.Developer, + content: + 'Use the tools available to you to find the answers (read page content, etc).', + }, + { + role: AIMessageRole.User, + content: question, + }, + ], + } + ), + fetchServerActionSiteContext(baseContext), + ]); + + // Get the responseId asynchronously in the background + let newResponseId: string | null = null; + const responseIdPromise = response + .then((r) => { + newResponseId = r.responseId; + }) + .catch((error) => { + console.error('Error getting responseId:', error); + }); + + // Start processing the stream immediately + for await (const value of stream) { + if (!value.answer) continue; + + yield { + answer: value.answer, + }; + } + + // Wait for the responseId to be available and yield one final time + await responseIdPromise; + yield { newResponseId }; +} diff --git a/packages/gitbook/src/components/Adaptive/server-actions/streamPageSummary.ts b/packages/gitbook/src/components/Adaptive/server-actions/streamPageSummary.ts index 3f64bb00e..6e9699209 100644 --- a/packages/gitbook/src/components/Adaptive/server-actions/streamPageSummary.ts +++ b/packages/gitbook/src/components/Adaptive/server-actions/streamPageSummary.ts @@ -21,7 +21,6 @@ export async function* streamPageSummary({ }; currentSpace: { id: string; - // title: string; }; visitedPages: { pageId: string; @@ -31,7 +30,7 @@ export async function* streamPageSummary({ const baseContext = isV2() ? await getServerActionBaseContext() : await getV1BaseContext(); const siteURLData = await getSiteURLDataFromMiddleware(); - const [{ stream }] = await Promise.all([ + const [{ stream, response }] = await Promise.all([ streamGenerateObject( baseContext, { @@ -54,10 +53,6 @@ export async function* streamPageSummary({ ) : z.undefined(), }), - tools: { - // getPages: true, - // getPageContent: true, - }, messages: [ { role: AIMessageRole.Developer, @@ -200,6 +195,17 @@ export async function* streamPageSummary({ fetchServerActionSiteContext(baseContext), ]); + // Get the responseId asynchronously in the background + let responseId: string | null = null; + const responseIdPromise = response + .then((r) => { + responseId = r.responseId; + }) + .catch((error) => { + console.error('Error getting responseId:', error); + }); + + // Start processing the stream immediately for await (const value of stream) { const keyFacts = value.keyFacts; const bigPicture = value.bigPicture; @@ -211,4 +217,8 @@ export async function* streamPageSummary({ bigPicture, }; } + + // Wait for the responseId to be available and yield one final time + await responseIdPromise; + yield { responseId }; } diff --git a/packages/gitbook/src/components/primitives/Button.tsx b/packages/gitbook/src/components/primitives/Button.tsx index b7edac7ab..74befc41b 100644 --- a/packages/gitbook/src/components/primitives/Button.tsx +++ b/packages/gitbook/src/components/primitives/Button.tsx @@ -14,6 +14,7 @@ type ButtonProps = { iconOnly?: boolean; size?: 'default' | 'medium' | 'small'; className?: ClassValue; + disabled?: boolean; label?: string; } & LinkInsightsProps & HTMLAttributes; @@ -33,10 +34,12 @@ const variantClasses = { 'ring-0', 'shadow-none', 'hover:bg-primary-hover', + 'disabled:hover:bg-transparent', 'hover:text-primary', 'hover:scale-1', 'hover:shadow-none', 'contrast-more:bg-tint-subtle', + 'disabled:hover:shadow-none', ], secondary: [ 'bg-tint', @@ -57,6 +60,7 @@ export function Button({ label, icon, iconOnly = false, + disabled = false, ...rest }: ButtonProps & { target?: HTMLAttributeAnchorTarget }) { const sizes = { @@ -94,6 +98,10 @@ export function Button({ 'active:scale-100', 'transition-all', + 'disabled:opacity-5', + 'disabled:cursor-not-allowed', + 'disabled:hover:shadow-none', + 'grow-0', 'shrink-0', 'truncate', @@ -119,7 +127,13 @@ export function Button({ } return ( -