Skip to content
This repository was archived by the owner on Sep 29, 2025. It is now read-only.

Commit 06b42d2

Browse files
mongodbenBen Perlmutter
mongodben
and
Ben Perlmutter
authored
(EAI-1112): Segment tracing fix (#780)
* add custom request custom data even if not in generate content func * more elegantly handle missing segment data * fix test bug * Test fix * test fix * elegently handle missing segment * add segment check --------- Co-authored-by: Ben Perlmutter <mongodben@mongodb.com>
1 parent a22d0ea commit 06b42d2

File tree

10 files changed

+241
-91
lines changed

10 files changed

+241
-91
lines changed

‎packages/chatbot-server-mongodb-public/src/config.ts‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -324,12 +324,17 @@ export const config: AppConfig = {
324324
req,
325325
res
326326
);
327-
return {
327+
const customData = {
328328
...defaultCustomData,
329-
segmentUserId: res.locals.customData.segmentUserId ?? undefined,
330-
segmentAnonymousId:
331-
res.locals.customData.segmentAnonymousId ?? undefined,
332329
};
330+
if (res.locals.customData.segmentUserId) {
331+
customData.segmentUserId = res.locals.customData.segmentUserId;
332+
}
333+
if (res.locals.customData.segmentAnonymousId) {
334+
customData.segmentAnonymousId =
335+
res.locals.customData.segmentAnonymousId;
336+
}
337+
return customData;
333338
},
334339
addMessageToConversationUpdateTrace:
335340
makeAddMessageToConversationUpdateTrace({

‎packages/chatbot-server-mongodb-public/src/tracing/extractTracingData.test.ts‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -128,4 +128,110 @@ describe("extractTracingData", () => {
128128
expect(tracingData.userMessageIndex).toBe(0);
129129
expect(tracingData.assistantMessageIndex).toBe(1);
130130
});
131+
132+
test("should extract origin from customData", () => {
133+
const messagesWithOrigin: Message[] = [
134+
{
135+
...baseUserMessage,
136+
customData: {
137+
origin: "https://example.com/chat",
138+
},
139+
},
140+
baseAssistantMessage,
141+
];
142+
const tracingData = extractTracingData(
143+
messagesWithOrigin,
144+
msgId,
145+
conversationId
146+
);
147+
expect(tracingData.origin).toBe("https://example.com/chat");
148+
});
149+
150+
test("should handle missing origin", () => {
151+
const messagesNoOrigin: Message[] = [
152+
{
153+
...baseUserMessage,
154+
customData: {},
155+
},
156+
baseAssistantMessage,
157+
];
158+
const tracingData = extractTracingData(
159+
messagesNoOrigin,
160+
msgId,
161+
conversationId
162+
);
163+
expect(tracingData.origin).toBeUndefined();
164+
});
165+
166+
test("should handle non-string origin", () => {
167+
const messagesInvalidOrigin: Message[] = [
168+
{
169+
...baseUserMessage,
170+
customData: {
171+
origin: 123, // non-string value
172+
},
173+
},
174+
baseAssistantMessage,
175+
];
176+
const tracingData = extractTracingData(
177+
messagesInvalidOrigin,
178+
msgId,
179+
conversationId
180+
);
181+
expect(tracingData.origin).toBeUndefined();
182+
});
183+
184+
test("should extract rejectionReason from customData", () => {
185+
const messagesWithRejection: Message[] = [
186+
{
187+
...baseUserMessage,
188+
customData: {
189+
rejectionReason: "Query contains inappropriate content",
190+
},
191+
},
192+
baseAssistantMessage,
193+
];
194+
const tracingData = extractTracingData(
195+
messagesWithRejection,
196+
msgId,
197+
conversationId
198+
);
199+
expect(tracingData.rejectionReason).toBe(
200+
"Query contains inappropriate content"
201+
);
202+
});
203+
204+
test("should use default rejection reason when missing", () => {
205+
const messagesNoRejection: Message[] = [
206+
{
207+
...baseUserMessage,
208+
customData: {},
209+
},
210+
baseAssistantMessage,
211+
];
212+
const tracingData = extractTracingData(
213+
messagesNoRejection,
214+
msgId,
215+
conversationId
216+
);
217+
expect(tracingData.rejectionReason).toBe("Unknown rejection reason");
218+
});
219+
220+
test("should use default rejection reason for non-string value", () => {
221+
const messagesInvalidRejection: Message[] = [
222+
{
223+
...baseUserMessage,
224+
customData: {
225+
rejectionReason: { reason: "complex object" }, // non-string value
226+
},
227+
},
228+
baseAssistantMessage,
229+
];
230+
const tracingData = extractTracingData(
231+
messagesInvalidRejection,
232+
msgId,
233+
conversationId
234+
);
235+
expect(tracingData.rejectionReason).toBe("Unknown rejection reason");
236+
});
131237
});

‎packages/chatbot-server-mongodb-public/src/tracing/extractTracingData.ts‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,10 +76,20 @@ export function extractTracingData(
7676
const rating = evalAssistantMessage.rating;
7777
const comment = evalAssistantMessage.userComment;
7878

79+
const maybeOrigin = previousUserMessage.customData?.origin;
80+
const origin = typeof maybeOrigin === "string" ? maybeOrigin : undefined;
81+
82+
const maybeRejectionReason = previousUserMessage.customData?.rejectionReason;
83+
const rejectionReason =
84+
typeof maybeRejectionReason === "string"
85+
? maybeRejectionReason
86+
: "Unknown rejection reason";
87+
7988
return {
8089
conversationId: conversationId,
8190
tags,
8291
rejectQuery,
92+
rejectionReason,
8393
isVerifiedAnswer,
8494
llmDoesNotKnow,
8595
numRetrievedChunks,
@@ -88,6 +98,7 @@ export function extractTracingData(
8898
userMessageIndex: previousUserMessageIdx,
8999
assistantMessage: evalAssistantMessage,
90100
assistantMessageIndex: evalAssistantMessageIdx,
101+
origin,
91102
rating,
92103
comment,
93104
};

‎packages/chatbot-server-mongodb-public/src/tracing/getLlmAsAJudgeScores.test.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,8 @@ describe("getLlmAsAJudgeScores", () => {
5555
conversationId: new ObjectId(),
5656
userMessageIndex: 0,
5757
assistantMessageIndex: 1,
58+
rejectionReason: "no rejection",
59+
origin: "unknown",
5860
} satisfies Parameters<typeof getLlmAsAJudgeScores>[1];
5961

6062
it("shouldn't judge verified answer", async () => {

‎packages/chatbot-server-mongodb-public/src/tracing/routesUpdateTraceHandlers.ts‎

Lines changed: 90 additions & 82 deletions
Original file line numberDiff line numberDiff line change
@@ -128,52 +128,57 @@ ${tracingData.assistantMessage?.content}
128128

129129
// Send Segment events
130130
try {
131-
if (segmentTrackUserSentMessage) {
131+
const userMessage = tracingData.userMessage;
132+
const { userId, anonymousId } = getSegmentIds(userMessage);
133+
const hasSegmentId = !!userId || !!anonymousId;
134+
if (segmentTrackUserSentMessage && hasSegmentId) {
132135
logRequest({
133136
reqId,
134137
message: `Sending addMessageToConversation event to Segment for conversation ${conversation._id}`,
135138
});
136-
}
137-
const userMessage = tracingData.userMessage;
138-
const { userId, anonymousId } = getSegmentIds(userMessage);
139-
if (userMessage) {
140-
segmentTrackUserSentMessage?.({
141-
userId,
142-
anonymousId,
143-
conversationId: conversation._id,
144-
origin: userMessage.customData?.origin as string,
145-
createdAt: userMessage.createdAt,
146-
tags: tracingData.tags,
147-
});
148-
} else {
149-
throw new Error(
150-
`Missing required data ${JSON.stringify({
151-
userMessage,
152-
})}`
153-
);
154-
}
155139

156-
const assistantMessage = tracingData.assistantMessage;
157-
if (userMessage && assistantMessage) {
158-
segmentTrackAssistantResponded?.({
159-
userId,
160-
anonymousId,
161-
conversationId: conversation._id,
162-
origin: userMessage.customData?.origin as string,
163-
createdAt: assistantMessage.createdAt,
164-
isVerifiedAnswer: tracingData.isVerifiedAnswer ?? false,
165-
rejectionReason: tracingData.rejectQuery
166-
? (userMessage.customData?.rejectionReason as string | undefined) ??
167-
"Unknown rejection reason"
168-
: undefined,
169-
});
140+
if (userMessage) {
141+
segmentTrackUserSentMessage?.({
142+
userId,
143+
anonymousId,
144+
conversationId: conversation._id,
145+
origin: tracingData.origin,
146+
createdAt: userMessage.createdAt,
147+
tags: tracingData.tags,
148+
});
149+
} else {
150+
throw new Error(
151+
`Missing required data ${JSON.stringify({
152+
userMessage,
153+
})}`
154+
);
155+
}
156+
157+
const assistantMessage = tracingData.assistantMessage;
158+
if (userMessage && assistantMessage) {
159+
segmentTrackAssistantResponded?.({
160+
userId,
161+
anonymousId,
162+
conversationId: conversation._id,
163+
origin: tracingData.origin,
164+
createdAt: assistantMessage.createdAt,
165+
isVerifiedAnswer: tracingData.isVerifiedAnswer ?? false,
166+
rejectionReason: tracingData.rejectionReason,
167+
});
168+
} else {
169+
throw new Error(
170+
`Missing required data ${JSON.stringify({
171+
userMessage,
172+
assistantMessage,
173+
})}`
174+
);
175+
}
170176
} else {
171-
throw new Error(
172-
`Missing required data ${JSON.stringify({
173-
userMessage,
174-
assistantMessage,
175-
})}`
176-
);
177+
logRequest({
178+
reqId,
179+
message: `Not sending events to Segement because no Segment configuration or ID present.`,
180+
type: "info",
181+
});
177182
}
178183
} catch (error) {
179184
logRequest({
@@ -300,29 +305,30 @@ export function makeRateMessageUpdateTrace({
300305
}
301306

302307
try {
303-
if (segmentTrackUserRatedMessage) {
308+
const hasSegmentId = !!userId || !!anonymousId;
309+
if (segmentTrackUserRatedMessage && hasSegmentId) {
304310
logRequest({
305311
reqId: traceId,
306312
message: `Sending rateMessage event to Segment for conversation ${conversation._id}`,
307313
});
308-
}
309-
if (userMessage && assistantMessage && rating !== undefined) {
310-
segmentTrackUserRatedMessage?.({
311-
userId,
312-
anonymousId,
313-
conversationId: conversation._id,
314-
origin: userMessage.customData?.origin as string,
315-
createdAt: new Date(),
316-
rating,
317-
});
318-
} else {
319-
throw new Error(
320-
`Missing required data ${JSON.stringify({
321-
userMessage,
322-
assistantMessage,
314+
if (userMessage && assistantMessage && rating !== undefined) {
315+
segmentTrackUserRatedMessage?.({
316+
userId,
317+
anonymousId,
318+
conversationId: conversation._id,
319+
origin: tracingData.origin,
320+
createdAt: new Date(),
323321
rating,
324-
})}`
325-
);
322+
});
323+
} else {
324+
throw new Error(
325+
`Missing required data ${JSON.stringify({
326+
userMessage,
327+
assistantMessage,
328+
rating,
329+
})}`
330+
);
331+
}
326332
}
327333
} catch (error) {
328334
logRequest({
@@ -431,36 +437,38 @@ export function makeCommentMessageUpdateTrace({
431437
}
432438

433439
try {
434-
if (segmentTrackUserCommentedMessage) {
440+
const hasSegmentId = !!userId || !!anonymousId;
441+
if (segmentTrackUserCommentedMessage && hasSegmentId) {
435442
logRequest({
436443
reqId: traceId,
437444
message: `Sending commentMessage event to Segment for conversation ${conversation._id}`,
438445
});
439-
}
440-
if (
441-
userMessage &&
442-
assistantMessage &&
443-
rating !== undefined &&
444-
comment !== undefined
445-
) {
446-
segmentTrackUserCommentedMessage?.({
447-
userId,
448-
anonymousId,
449-
conversationId: conversation._id,
450-
origin: userMessage.customData?.origin as string,
451-
createdAt: new Date(),
452-
rating,
453-
comment,
454-
});
455-
} else {
456-
throw new Error(
457-
`Missing required data ${JSON.stringify({
458-
userMessage,
459-
assistantMessage,
446+
447+
if (
448+
userMessage &&
449+
assistantMessage &&
450+
rating !== undefined &&
451+
comment !== undefined
452+
) {
453+
segmentTrackUserCommentedMessage?.({
454+
userId,
455+
anonymousId,
456+
conversationId: conversation._id,
457+
origin: tracingData.origin,
458+
createdAt: new Date(),
460459
rating,
461460
comment,
462-
})}`
463-
);
461+
});
462+
} else {
463+
throw new Error(
464+
`Missing required data ${JSON.stringify({
465+
userMessage,
466+
assistantMessage,
467+
rating,
468+
comment,
469+
})}`
470+
);
471+
}
464472
}
465473
} catch (error) {
466474
logRequest({

‎packages/chatbot-server-mongodb-public/src/tracing/scrubbedMessages/makeScrubbedMessagesFromTracingData.test.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,8 @@ describe("makeScrubbedMessagesFromTracingData", () => {
3333
},
3434
contextContent: [],
3535
assistantMessageIndex: 1,
36+
rejectionReason: "no rejection",
37+
origin: "unknown",
3638
} as ReturnType<typeof extractTracingData>;
3739

3840
it("should create scrubbed messages from tracing data", async () => {

‎packages/chatbot-server-mongodb-public/src/tracing/segment.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,7 @@ const BaseTrackEventParamsSchema = z.object({
7878
userId: z.string().optional(),
7979
anonymousId: z.string().optional(),
8080
conversationId: z.instanceof(ObjectId),
81-
origin: z.string(),
81+
origin: z.string().optional(),
8282
createdAt: z.date(),
8383
});
8484

0 commit comments

Comments
 (0)