Skip to content

Commit a650b58

Browse files
conico974Nicolas Dorseuil
conico974
and
Nicolas Dorseuil
authored
feat: add tracing to server actions (#3543)
Co-authored-by: Nicolas Dorseuil <[email protected]>
1 parent 388b20d commit a650b58

File tree

8 files changed

+349
-280
lines changed

8 files changed

+349
-280
lines changed

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

Lines changed: 86 additions & 82 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
'use server';
22
import type { GitBookBaseContext } from '@/lib/context';
33
import { fetchServerActionSiteContext } from '@/lib/server-actions';
4+
import { traceErrorOnly } from '@/lib/tracing';
45
import {
56
type AIMessage,
67
AIMessageRole,
@@ -19,97 +20,100 @@ export async function streamRenderAIMessage(
1920
rawStream: AsyncIterable<AIStreamResponse>,
2021
options?: RenderAIMessageOptions
2122
) {
22-
const message: AIMessage = {
23-
id: '',
24-
role: AIMessageRole.Assistant,
25-
steps: [],
26-
};
23+
return traceErrorOnly('AI.streamRenderAIMessage', async () => {
24+
const message: AIMessage = {
25+
id: '',
26+
role: AIMessageRole.Assistant,
27+
steps: [],
28+
};
2729

28-
const updateProcessingMessageStep = (
29-
stepIndex: number,
30-
callback: (step: AIMessageStep) => void
31-
) => {
32-
if (stepIndex > message.steps.length) {
33-
throw new Error(
34-
`Step index out of bounds ${stepIndex} (${message.steps.length} steps)`
35-
);
36-
}
30+
const updateProcessingMessageStep = (
31+
stepIndex: number,
32+
callback: (step: AIMessageStep) => void
33+
) => {
34+
if (stepIndex > message.steps.length) {
35+
throw new Error(
36+
`Step index out of bounds ${stepIndex} (${message.steps.length} steps)`
37+
);
38+
}
3739

38-
if (message.steps[stepIndex]) {
39-
message.steps = [...message.steps];
40-
message.steps[stepIndex] = { ...message.steps[stepIndex] };
41-
callback(message.steps[stepIndex]);
42-
} else {
43-
message.steps = [
44-
...message.steps,
45-
{
46-
content: {
47-
object: 'document',
48-
data: {},
49-
nodes: [],
40+
if (message.steps[stepIndex]) {
41+
message.steps = [...message.steps];
42+
message.steps[stepIndex] = { ...message.steps[stepIndex] };
43+
callback(message.steps[stepIndex]);
44+
} else {
45+
message.steps = [
46+
...message.steps,
47+
{
48+
content: {
49+
object: 'document',
50+
data: {},
51+
nodes: [],
52+
},
5053
},
51-
},
52-
];
53-
callback(message.steps[stepIndex]);
54-
}
55-
};
54+
];
55+
callback(message.steps[stepIndex]);
56+
}
57+
};
5658

57-
// Fetch the full-context in the background to avoid blocking the stream.
58-
const promiseContext = fetchServerActionSiteContext(baseContext);
59+
// Fetch the full-context in the background to avoid blocking the stream.
60+
const promiseContext = fetchServerActionSiteContext(baseContext);
5961

60-
return parseResponse<{
61-
content: React.ReactNode;
62-
event: AIStreamResponse;
63-
}>(rawStream, async (event) => {
64-
switch (event.type) {
65-
/**
66-
* The agent is processing a tool call in a new message.
67-
*/
68-
case 'response_tool_call': {
69-
updateProcessingMessageStep(event.stepIndex, (step) => {
70-
step.toolCalls ??= [];
71-
step.toolCalls.push(event.toolCall);
72-
});
73-
break;
74-
}
62+
return parseResponse<{
63+
content: React.ReactNode;
64+
event: AIStreamResponse;
65+
}>(rawStream, async (event) => {
66+
switch (event.type) {
67+
/**
68+
* The agent is processing a tool call in a new message.
69+
*/
70+
case 'response_tool_call': {
71+
updateProcessingMessageStep(event.stepIndex, (step) => {
72+
step.toolCalls ??= [];
73+
step.toolCalls.push(event.toolCall);
74+
});
75+
break;
76+
}
7577

76-
/**
77-
* The agent is writing the content of a new message.
78-
*/
79-
case 'response_reasoning':
80-
case 'response_document': {
81-
updateProcessingMessageStep(event.stepIndex, (step) => {
82-
const container = event.type === 'response_reasoning' ? 'reasoning' : 'content';
78+
/**
79+
* The agent is writing the content of a new message.
80+
*/
81+
case 'response_reasoning':
82+
case 'response_document': {
83+
updateProcessingMessageStep(event.stepIndex, (step) => {
84+
const container =
85+
event.type === 'response_reasoning' ? 'reasoning' : 'content';
8386

84-
step[container] ??= {
85-
object: 'document',
86-
data: {},
87-
nodes: [],
88-
};
89-
step[container] = {
90-
...step[container],
91-
nodes: [...step[container].nodes],
92-
};
93-
if (event.operation === 'insert') {
94-
step[container].nodes.push(...event.blocks);
95-
} else {
96-
step[container].nodes.splice(
97-
-event.blocks.length,
98-
event.blocks.length,
99-
...event.blocks
100-
);
101-
}
102-
});
103-
break;
87+
step[container] ??= {
88+
object: 'document',
89+
data: {},
90+
nodes: [],
91+
};
92+
step[container] = {
93+
...step[container],
94+
nodes: [...step[container].nodes],
95+
};
96+
if (event.operation === 'insert') {
97+
step[container].nodes.push(...event.blocks);
98+
} else {
99+
step[container].nodes.splice(
100+
-event.blocks.length,
101+
event.blocks.length,
102+
...event.blocks
103+
);
104+
}
105+
});
106+
break;
107+
}
104108
}
105-
}
106109

107-
return {
108-
event,
109-
content: (
110-
<AIMessageView message={message} context={await promiseContext} {...options} />
111-
),
112-
};
110+
return {
111+
event,
112+
content: (
113+
<AIMessageView message={message} context={await promiseContext} {...options} />
114+
),
115+
};
116+
});
113117
});
114118
}
115119

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

Lines changed: 23 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
'use server';
22
import { getSiteURLDataFromMiddleware } from '@/lib/middleware';
33
import { getServerActionBaseContext } from '@/lib/server-actions';
4+
import { traceErrorOnly } from '@/lib/tracing';
45
import { type AIMessageContext, AIMessageRole, AIModel } from '@gitbook/api';
56
import { streamRenderAIMessage } from './api';
67
import type { RenderAIMessageOptions } from './types';
@@ -19,25 +20,31 @@ export async function* streamAIChatResponse({
1920
previousResponseId?: string;
2021
options?: RenderAIMessageOptions;
2122
}) {
22-
const context = await getServerActionBaseContext();
23-
const siteURLData = await getSiteURLDataFromMiddleware();
23+
const { stream } = await traceErrorOnly('AI.streamAIChatResponse', async () => {
24+
const context = await getServerActionBaseContext();
25+
const siteURLData = await getSiteURLDataFromMiddleware();
2426

25-
const api = await context.dataFetcher.api();
26-
const rawStream = api.orgs.streamAiResponseInSite(siteURLData.organization, siteURLData.site, {
27-
mode: 'assistant',
28-
input: [
27+
const api = await context.dataFetcher.api();
28+
const rawStream = api.orgs.streamAiResponseInSite(
29+
siteURLData.organization,
30+
siteURLData.site,
2931
{
30-
role: AIMessageRole.User,
31-
content: message,
32-
context: messageContext,
33-
},
34-
],
35-
output: { type: 'document' },
36-
model: AIModel.ReasoningLow,
37-
previousResponseId,
38-
});
32+
mode: 'assistant',
33+
input: [
34+
{
35+
role: AIMessageRole.User,
36+
content: message,
37+
context: messageContext,
38+
},
39+
],
40+
output: { type: 'document' },
41+
model: AIModel.ReasoningLow,
42+
previousResponseId,
43+
}
44+
);
3945

40-
const { stream } = await streamRenderAIMessage(context, rawStream, options);
46+
return await streamRenderAIMessage(context, rawStream, options);
47+
});
4148

4249
for await (const output of stream) {
4350
yield output;

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

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
'use server';
22
import { getSiteURLDataFromMiddleware } from '@/lib/middleware';
33
import { getServerActionBaseContext } from '@/lib/server-actions';
4+
import { traceErrorOnly } from '@/lib/tracing';
45
import { streamRenderAIMessage } from './api';
56
import type { RenderAIMessageOptions } from './types';
67

@@ -14,16 +15,18 @@ export async function* streamAIResponseById({
1415
responseId: string;
1516
options?: RenderAIMessageOptions;
1617
}) {
17-
const context = await getServerActionBaseContext();
18-
const siteURLData = await getSiteURLDataFromMiddleware();
18+
const { stream } = await traceErrorOnly('AI.streamAIResponseById', async () => {
19+
const context = await getServerActionBaseContext();
20+
const siteURLData = await getSiteURLDataFromMiddleware();
1921

20-
const api = await context.dataFetcher.api();
21-
const rawStream = api.orgs.streamExistingAiResponseInSite(
22-
siteURLData.organization,
23-
siteURLData.site,
24-
responseId
25-
);
26-
const { stream } = await streamRenderAIMessage(context, rawStream, options);
22+
const api = await context.dataFetcher.api();
23+
const rawStream = api.orgs.streamExistingAiResponseInSite(
24+
siteURLData.organization,
25+
siteURLData.site,
26+
responseId
27+
);
28+
return await streamRenderAIMessage(context, rawStream, options);
29+
});
2730

2831
for await (const output of stream) {
2932
yield output;

‎packages/gitbook/src/components/Ads/renderAd.tsx

Lines changed: 37 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import type { SiteInsightsAd, SiteInsightsAdPlacement } from '@gitbook/api';
44
import { headers } from 'next/headers';
55

66
import { getServerActionBaseContext } from '@/lib/server-actions';
7+
import { traceErrorOnly } from '@/lib/tracing';
78
import { AdClassicRendering } from './AdClassicRendering';
89
import { AdCoverRendering } from './AdCoverRendering';
910
import { AdPixels } from './AdPixels';
@@ -40,40 +41,42 @@ interface FetchPlaceholderAdOptions {
4041
* and properly access user-agent and IP.
4142
*/
4243
export async function renderAd(options: FetchAdOptions) {
43-
const [context, result] = await Promise.all([
44-
getServerActionBaseContext(),
45-
options.source === 'live' ? fetchAd(options) : getPlaceholderAd(),
46-
]);
47-
48-
const mode = options.source === 'live' ? options.mode : 'classic';
49-
if (!result || !result.ad.description || !result.ad.statlink) {
50-
return null;
51-
}
52-
53-
const { ad } = result;
54-
55-
const insightsAd: SiteInsightsAd | null =
56-
options.source === 'live'
57-
? {
58-
placement: options.placement,
59-
zoneId: options.zoneId,
60-
domain: 'company' in ad ? ad.company : '',
61-
}
62-
: null;
63-
64-
return {
65-
children: (
66-
<>
67-
{mode === 'classic' || !('callToAction' in ad) ? (
68-
<AdClassicRendering ad={ad} insightsAd={insightsAd} context={context} />
69-
) : (
70-
<AdCoverRendering ad={ad} insightsAd={insightsAd} context={context} />
71-
)}
72-
{ad.pixel ? <AdPixels rawPixel={ad.pixel} /> : null}
73-
</>
74-
),
75-
insightsAd,
76-
};
44+
return traceErrorOnly('Ads.renderAd', async () => {
45+
const [context, result] = await Promise.all([
46+
getServerActionBaseContext(),
47+
options.source === 'live' ? fetchAd(options) : getPlaceholderAd(),
48+
]);
49+
50+
const mode = options.source === 'live' ? options.mode : 'classic';
51+
if (!result || !result.ad.description || !result.ad.statlink) {
52+
return null;
53+
}
54+
55+
const { ad } = result;
56+
57+
const insightsAd: SiteInsightsAd | null =
58+
options.source === 'live'
59+
? {
60+
placement: options.placement,
61+
zoneId: options.zoneId,
62+
domain: 'company' in ad ? ad.company : '',
63+
}
64+
: null;
65+
66+
return {
67+
children: (
68+
<>
69+
{mode === 'classic' || !('callToAction' in ad) ? (
70+
<AdClassicRendering ad={ad} insightsAd={insightsAd} context={context} />
71+
) : (
72+
<AdCoverRendering ad={ad} insightsAd={insightsAd} context={context} />
73+
)}
74+
{ad.pixel ? <AdPixels rawPixel={ad.pixel} /> : null}
75+
</>
76+
),
77+
insightsAd,
78+
};
79+
});
7780
}
7881

7982
async function fetchAd({

0 commit comments

Comments
 (0)