Skip to content

Commit e97cdc3

Browse files
SamyPesseSamyPesse
committed
Add clear and suggestions
1 parent d99630a commit e97cdc3

File tree

3 files changed

+156
-63
lines changed

3 files changed

+156
-63
lines changed

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,2 @@
11
export * from './useAIPage';
2+
export * from './useAIChat';

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

Lines changed: 54 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,8 @@
22
import { AIMessageRole, AIModel } from '@gitbook/api';
33
import { getSiteURLDataFromMiddleware } from '@v2/lib/middleware';
44
import { getServerActionBaseContext } from '@v2/lib/server-actions';
5-
import { streamRenderAIMessage } from './api';
5+
import { z } from 'zod';
6+
import { streamGenerateAIObject, streamRenderAIMessage } from './api';
67
import type { RenderAIMessageOptions } from './types';
78

89
const PROMPT = `
@@ -16,6 +17,22 @@ You analyse the query, and the content of the site, and generate a page that wil
1617
- Always use the provided tools to understand the docs knowledge base, do not make up information.
1718
`;
1819

20+
const FOLLOWUP_PROMPT = `
21+
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.
22+
23+
# Guidelines
24+
25+
- Ensure suggestions are concise and relevant for general chat conversations.
26+
- Limit the length of each suggestion to ensure quick readability and tap selection.
27+
- Suggest at most 3 responses.
28+
- Only suggest responses that are relevant followup to the conversation, otherwise return an empty list.
29+
- When the last message finishes with questions, suggest responses that answer the questions.
30+
31+
# Output Format
32+
33+
Provide the suggestions as a JSON array with each suggestion as a string. Ensure the suggestions are short and suitable for quick tapping.
34+
`;
35+
1936
/**
2037
* Generate a response to a chat message.
2138
*/
@@ -56,3 +73,39 @@ export async function* streamAIChatResponse({
5673
yield output;
5774
}
5875
}
76+
77+
/**
78+
* Stream suggestions of follow-up responses for the user.
79+
*/
80+
export async function* streamAIChatFollowUpResponses({
81+
previousResponseId,
82+
}: {
83+
previousResponseId: string;
84+
}) {
85+
const context = await getServerActionBaseContext();
86+
const siteURLData = await getSiteURLDataFromMiddleware();
87+
88+
const { stream, response } = await streamGenerateAIObject(context, {
89+
organizationId: siteURLData.organization,
90+
siteId: siteURLData.site,
91+
schema: z.object({
92+
suggestions: z.array(z.string()),
93+
}),
94+
previousResponseId,
95+
input: [
96+
{
97+
role: AIMessageRole.User,
98+
content:
99+
'Suggest quick-tap responses the user might want to pick from to continue the previous chat conversation.',
100+
},
101+
],
102+
model: AIModel.Fast,
103+
instructions: FOLLOWUP_PROMPT,
104+
});
105+
106+
for await (const output of stream) {
107+
yield (output.suggestions ?? []).filter((suggestion) => !!suggestion) as string[];
108+
}
109+
110+
console.log('response', { previousResponseId }, await response);
111+
}

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

Lines changed: 101 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ import * as zustand from 'zustand';
44

55
import { AIMessageRole } from '@gitbook/api';
66
import type * as React from 'react';
7-
import { streamAIChatResponse } from './server-actions';
7+
import { streamAIChatFollowUpResponses, streamAIChatResponse } from './server-actions';
88

99
export type AIChatMessage = {
1010
role: AIMessageRole;
@@ -27,6 +27,11 @@ export type AIChatState = {
2727
*/
2828
messages: AIChatMessage[];
2929

30+
/**
31+
* Suggestions for follow-up messages.
32+
*/
33+
followUpSuggestions: string[];
34+
3035
/**
3136
* If true, the session is in progress.
3237
*/
@@ -45,83 +50,117 @@ export type AIChatController = {
4550
/** The message to post to the session. it can be markdown formatted. */
4651
message: string;
4752
}) => void;
53+
54+
/** Clear the conversation */
55+
clear: () => void;
4856
};
4957

5058
const globalState = zustand.create<{
5159
state: AIChatState;
5260
controller: AIChatController;
53-
}>((set, get) => ({
54-
state: {
55-
opened: false,
56-
responseId: null,
57-
messages: [],
58-
loading: false,
59-
},
60-
controller: {
61-
open: () => set((state) => ({ state: { ...state.state, opened: true } })),
62-
close: () => set((state) => ({ state: { ...state.state, opened: false } })),
63-
postMessage: async (input: { message: string }) => {
64-
set((previous) => {
65-
return {
66-
...previous,
61+
}>((set, get) => {
62+
/**
63+
* Refresh the follow-up suggestions.
64+
*/
65+
const fetchFollowUpSuggestions = async (previousResponseId: string) => {
66+
const stream = await streamAIChatFollowUpResponses({
67+
previousResponseId,
68+
});
69+
70+
for await (const suggestions of stream) {
71+
set((state) => ({ state: { ...state.state, followUpSuggestions: suggestions } }));
72+
}
73+
};
74+
75+
return {
76+
state: {
77+
opened: false,
78+
responseId: null,
79+
messages: [],
80+
followUpSuggestions: [],
81+
loading: false,
82+
},
83+
controller: {
84+
open: () => set((state) => ({ state: { ...state.state, opened: true } })),
85+
close: () => set((state) => ({ state: { ...state.state, opened: false } })),
86+
clear: () =>
87+
set(() => ({
6788
state: {
68-
...previous.state,
69-
messages: [
70-
...previous.state.messages,
71-
{
72-
// TODO: how to handle markdown here?
73-
// to avoid rendering as plain text
74-
role: AIMessageRole.User,
75-
content: input.message,
76-
},
77-
{
78-
role: AIMessageRole.Assistant,
79-
content: null,
80-
},
81-
],
82-
loading: true,
89+
opened: false,
90+
loading: false,
91+
messages: [],
92+
followUpSuggestions: [],
93+
responseId: null,
8394
},
84-
};
85-
});
86-
87-
const stream = await streamAIChatResponse({
88-
message: input.message,
89-
previousResponseId: get().state.responseId ?? undefined,
90-
});
91-
92-
for await (const data of stream) {
93-
if (!data) continue;
94-
95-
const event = data.event;
96-
if (event.type === 'response_finish') {
97-
set((state) => ({ state: { ...state.state, responseId: event.responseId } }));
95+
})),
96+
postMessage: async (input: { message: string }) => {
97+
set((previous) => {
98+
return {
99+
...previous,
100+
state: {
101+
...previous.state,
102+
messages: [
103+
...previous.state.messages,
104+
{
105+
// TODO: how to handle markdown here?
106+
// to avoid rendering as plain text
107+
role: AIMessageRole.User,
108+
content: input.message,
109+
},
110+
{
111+
role: AIMessageRole.Assistant,
112+
content: null,
113+
},
114+
],
115+
followUpSuggestions: [],
116+
loading: true,
117+
},
118+
};
119+
});
120+
121+
const stream = await streamAIChatResponse({
122+
message: input.message,
123+
previousResponseId: get().state.responseId ?? undefined,
124+
});
125+
126+
for await (const data of stream) {
127+
if (!data) continue;
128+
129+
const event = data.event;
130+
if (event.type === 'response_finish') {
131+
set((state) => ({
132+
state: { ...state.state, responseId: event.responseId },
133+
}));
134+
135+
fetchFollowUpSuggestions(event.responseId);
136+
}
137+
138+
set((previous) => ({
139+
...previous,
140+
state: {
141+
...previous.state,
142+
messages: [
143+
...previous.state.messages.slice(0, -1),
144+
{
145+
role: AIMessageRole.Assistant,
146+
content: data.content,
147+
},
148+
],
149+
},
150+
}));
98151
}
99152

100153
set((previous) => ({
101154
...previous,
102155
state: {
103156
...previous.state,
104-
messages: [
105-
...previous.state.messages.slice(0, -1),
106-
{
107-
role: AIMessageRole.Assistant,
108-
content: data.content,
109-
},
110-
],
157+
loading: false,
111158
},
112159
}));
113-
}
114-
115-
set((previous) => ({
116-
...previous,
117-
state: {
118-
...previous.state,
119-
loading: false,
120-
},
121-
}));
160+
},
122161
},
123-
},
124-
}));
162+
};
163+
});
125164

126165
/**
127166
* Get the current state of the AI chat.

0 commit comments

Comments
 (0)