Skip to content

Commit

Permalink
telemetry(amazonq): measure e2e latency for amazon q requests
Browse files Browse the repository at this point in the history
Problem:
- We can't measure e2e how long a chat request takes

Solution:
- Take an initial approach, see what works and what doesn't

Original approach:
- The original approach was to "listen" to all events coming out of telemetry. I think this is something we probably want to do eventually because it means we won't need to manually instrument flows ourselves but I think the integration needed to be more refined

Current approach:
- Plumb the trace id through to the webview -> vscode -> vscode. This allows us to know exactly what is happening to a message and how long everything took to get there
  • Loading branch information
jpinkney-aws committed Sep 12, 2024
1 parent 5eea5f9 commit 22db33e
Show file tree
Hide file tree
Showing 15 changed files with 232 additions and 13 deletions.
88 changes: 88 additions & 0 deletions packages/core/src/amazonq/messages/chatMessageDuration.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,88 @@
/*!
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
* SPDX-License-Identifier: Apache-2.0
*/

import { globals } from '../../shared'
import { telemetry } from '../../shared/telemetry'
import { uiEventRecorder } from '../util/eventRecorder'

export class AmazonQChatMessageDuration {
/**
* Record the initial requests in the chat message flow
*/
static startListening(msg: { traceId: string; startTime: number; trigger?: string }) {
const { traceId, startTime, trigger } = msg

uiEventRecorder.set(traceId, {
events: {
chatMessageSent: startTime,
},
})
uiEventRecorder.set(traceId, {
events: {
editorReceivedMessage: globals.clock.Date.now(),
},
})
if (trigger) {
uiEventRecorder.set(traceId, {
trigger,
})
}
}

/**
* Stop listening to all incoming events and emit what we've found
*/
static stopListening(msg: { traceId: string }) {
const { traceId } = msg

// We can't figure out what trace this event was associated with
if (!traceId) {
return
}

uiEventRecorder.set(traceId, {
events: {
messageDisplayed: globals.clock.Date.now(),
},
})

const metrics = uiEventRecorder.get(traceId)

// get events sorted by the time they were created
const events = Object.entries(metrics.events)
.map((x) => ({
event: x[0],
duration: x[1],
}))
.sort((a, b) => {
return a.duration - b.duration
})

// Get the total duration by subtracting the when the message was displayed and when the chat message was first sent
const totalDuration = events[events.length - 1].duration - events[0].duration

/**
* Find the time it took to get between two metric events
*/
const timings = new Map<string, number>()
for (let i = 1; i < events.length; i++) {
const currentEvent = events[i]
const previousEvent = events[i - 1]

const timeDifference = currentEvent.duration - previousEvent.duration

timings.set(currentEvent.event, timeDifference)
}

telemetry.event_flow.emit({
name: metrics.trigger,
duration: totalDuration,
result: 'Succeeded',
...Object.fromEntries(timings),
} as any)

uiEventRecorder.delete(traceId)
}
}
22 changes: 22 additions & 0 deletions packages/core/src/amazonq/util/eventRecorder.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
/*!
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
* SPDX-License-Identifier: Apache-2.0
*/

import { MetricName } from '../../shared/telemetry/telemetry.gen'
import { RecordMap } from '../../shared/utilities/map'

type Event =
| 'chatMessageSent' // initial on chat prompt event in the ui
| 'editorReceivedMessage' // message gets from the chat prompt to VSCode
| 'featureReceivedMessage' // message gets redirected from VSCode -> Partner team features implementation
| MetricName // any telemetry event emitted between when the partner teams code gets the original message to when its sent to the UI
| 'messageDisplayed' // message gets received in the UI

/**
* For a given traceID, map an event to a time
*/
export const uiEventRecorder = new RecordMap<{
trigger: string
events: Partial<Record<Event, number>>
}>()
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import { TabType } from '../ui/storages/tabsStorage'
import { getLogger } from '../../../shared/logger'
import { amazonqMark } from '../../../shared/performance/marks'
import { telemetry } from '../../../shared/telemetry'
import { AmazonQChatMessageDuration } from '../../messages/chatMessageDuration'

export function dispatchWebViewMessagesToApps(
webview: Webview,
Expand All @@ -33,6 +34,14 @@ export function dispatchWebViewMessagesToApps(
return
}

if (msg.type === 'startListening') {
AmazonQChatMessageDuration.startListening(msg)
return
} else if (msg.type === 'stopListening') {
AmazonQChatMessageDuration.stopListening(msg)
return
}

if (msg.type === 'error') {
const event = msg.event === 'webview_load' ? telemetry.webview_load : telemetry.webview_error
event.emit({
Expand Down
9 changes: 7 additions & 2 deletions packages/core/src/amazonq/webview/ui/apps/cwChatConnector.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,9 +8,11 @@ import { ExtensionMessage } from '../commands'
import { CodeReference } from './amazonqCommonsConnector'
import { TabOpenType, TabsStorage } from '../storages/tabsStorage'
import { FollowUpGenerator } from '../followUps/generator'
import { TracedChatItem } from '../connector'

interface ChatPayload {
chatMessage: string
traceId?: string
chatCommand?: string
}

Expand Down Expand Up @@ -188,6 +190,7 @@ export class Connector {
command: 'chat-prompt',
chatMessage: payload.chatMessage,
chatCommand: payload.chatCommand,
traceId: payload.traceId,
tabType: 'cwc',
})
})
Expand Down Expand Up @@ -258,13 +261,14 @@ export class Connector {
}
: undefined

const answer: ChatItem = {
const answer: TracedChatItem = {
type: messageData.messageType,
messageId: messageData.messageID ?? messageData.triggerID,
body: messageData.message,
followUp: followUps,
canBeVoted: true,
codeReference: messageData.codeReference,
traceId: messageData.traceId,
}

// If it is not there we will not set it
Expand All @@ -291,7 +295,7 @@ export class Connector {
return
}
if (messageData.messageType === ChatItemType.ANSWER) {
const answer: ChatItem = {
const answer: TracedChatItem = {
type: messageData.messageType,
body: undefined,
relatedContent: undefined,
Expand All @@ -304,6 +308,7 @@ export class Connector {
options: messageData.followUps,
}
: undefined,
traceId: messageData.traceId,
}
this.onChatAnswerReceived(messageData.tabID, answer)

Expand Down
9 changes: 7 additions & 2 deletions packages/core/src/amazonq/webview/ui/connector.ts
Original file line number Diff line number Diff line change
Expand Up @@ -26,18 +26,23 @@ export interface CodeReference {

export interface ChatPayload {
chatMessage: string
traceId?: string // TODO: instrumented for cwc, not for gumby/featuredev. Remove the ? once those changes are merged
chatCommand?: string
}

export interface TracedChatItem extends ChatItem {
traceId?: string
}

export interface ConnectorProps {
sendMessageToExtension: (message: ExtensionMessage) => void
onMessageReceived?: (tabID: string, messageData: any, needToShowAPIDocsTab: boolean) => void
onChatAnswerUpdated?: (tabID: string, message: ChatItem) => void
onChatAnswerReceived?: (tabID: string, message: ChatItem) => void
onChatAnswerReceived?: (tabID: string, message: TracedChatItem) => void
onWelcomeFollowUpClicked: (tabID: string, welcomeFollowUpType: WelcomeFollowupType) => void
onAsyncEventProgress: (tabID: string, inProgress: boolean, message: string | undefined) => void
onQuickHandlerCommand: (tabID: string, command?: string, eventId?: string) => void
onCWCContextCommandMessage: (message: ChatItem, command?: string) => string | undefined
onCWCContextCommandMessage: (message: TracedChatItem, command?: string) => string | undefined
onOpenSettingsMessage: (tabID: string) => void
onError: (tabID: string, message: string, title: string) => void
onWarning: (tabID: string, message: string, title: string) => void
Expand Down
27 changes: 24 additions & 3 deletions packages/core/src/amazonq/webview/ui/main.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
* SPDX-License-Identifier: Apache-2.0
*/
import { Connector } from './connector'
import { Connector, TracedChatItem } from './connector'
import { ChatItem, ChatItemType, MynahIcons, MynahUI, MynahUIDataModel, NotificationType } from '@aws/mynah-ui'
import { ChatPrompt } from '@aws/mynah-ui/dist/static'
import { TabsStorage, TabType } from './storages/tabsStorage'
Expand Down Expand Up @@ -189,7 +189,7 @@ export const createMynahUI = (ideApi: any, amazonQEnabled: boolean) => {
} as ChatItem)
}
},
onChatAnswerReceived: (tabID: string, item: ChatItem) => {
onChatAnswerReceived: (tabID: string, item: TracedChatItem) => {
if (item.type === ChatItemType.ANSWER_PART || item.type === ChatItemType.CODE_RESULT) {
mynahUI.updateLastChatAnswer(tabID, {
...(item.messageId !== undefined ? { messageId: item.messageId } : {}),
Expand Down Expand Up @@ -234,6 +234,15 @@ export const createMynahUI = (ideApi: any, amazonQEnabled: boolean) => {
promptInputDisabledState: tabsStorage.isTabDead(tabID),
})
tabsStorage.updateTabStatus(tabID, 'free')

if (item.traceId) {
ideApi.postMessage({
type: 'stopListening',
tabID,
traceId: item.traceId,
tabType: tabsStorage.getTab(tabID)?.type,
})
}
}
},
onMessageReceived: (tabID: string, messageData: MynahUIDataModel) => {
Expand Down Expand Up @@ -362,6 +371,7 @@ export const createMynahUI = (ideApi: any, amazonQEnabled: boolean) => {
},
onTabRemove: connector.onTabRemove,
onTabChange: connector.onTabChange,
// TODO: update mynah-ui this type doesn't seem correct https://github.com/aws/mynah-ui/blob/3777a39eb534a91fd6b99d6cf421ce78ee5c7526/src/main.ts#L372
onChatPrompt: (tabID: string, prompt: ChatPrompt, eventId: string | undefined) => {
if ((prompt.prompt ?? '') === '' && (prompt.command ?? '') === '') {
return
Expand All @@ -374,6 +384,7 @@ export const createMynahUI = (ideApi: any, amazonQEnabled: boolean) => {
} else if (tabsStorage.getTab(tabID)?.type === 'gumby') {
connector.requestAnswer(tabID, {
chatMessage: prompt.prompt ?? '',
traceId: eventId as string,
})
return
}
Expand All @@ -383,7 +394,17 @@ export const createMynahUI = (ideApi: any, amazonQEnabled: boolean) => {
return
}

textMessageHandler.handle(prompt, tabID)
// When a user presses "enter"
ideApi.postMessage({
type: 'startListening',
trigger: 'onChatPrompt',
tabID,
traceId: eventId,
tabType: tabsStorage.getTab(tabID)?.type,
startTime: Date.now(),
})

textMessageHandler.handle(prompt, tabID, eventId as string)
},
onVote: connector.onChatItemVoted,
onInBodyButtonClicked: (tabId, messageId, action, eventId) => {
Expand Down
3 changes: 2 additions & 1 deletion packages/core/src/amazonq/webview/ui/messages/handler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ export class TextMessageHandler {
this.tabsStorage = props.tabsStorage
}

public handle(chatPrompt: ChatPrompt, tabID: string) {
public handle(chatPrompt: ChatPrompt, tabID: string, eventID: string) {
this.tabsStorage.updateTabTypeFromUnknown(tabID, 'cwc')
this.tabsStorage.resetTabTimer(tabID)
this.connector.onUpdateTabType(tabID)
Expand All @@ -44,6 +44,7 @@ export class TextMessageHandler {
.requestGenerativeAIAnswer(tabID, {
chatMessage: chatPrompt.prompt ?? '',
chatCommand: chatPrompt.command,
traceId: eventID,
})
.then(() => {})
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,8 @@ import { CodeWhispererSettings } from '../../../codewhisperer/util/codewhisperer
import { getSelectedCustomization } from '../../../codewhisperer/util/customizationUtil'
import { FeatureConfigProvider } from '../../../shared/featureConfig'
import { getHttpStatusCode, AwsClientResponseError } from '../../../shared/errors'
import { uiEventRecorder } from '../../../amazonq/util/eventRecorder'
import { globals } from '../../../shared'

export interface ChatControllerMessagePublishers {
readonly processPromptChatMessage: MessagePublisher<PromptMessage>
Expand Down Expand Up @@ -122,6 +124,13 @@ export class ChatController {
})

this.chatControllerMessageListeners.processPromptChatMessage.onMessage((data) => {
if (data.traceId) {
uiEventRecorder.set(data.traceId, {
events: {
featureReceivedMessage: globals.clock.Date.now(),
},
})
}
return this.processPromptChatMessage(data)
})

Expand Down Expand Up @@ -485,6 +494,7 @@ export class ChatController {
codeQuery: context?.focusAreaContext?.names,
userIntent: this.userIntentRecognizer.getFromPromptChatMessage(message),
customization: getSelectedCustomization(),
traceId: message.traceId,
},
triggerID
)
Expand Down Expand Up @@ -604,7 +614,7 @@ export class ChatController {
try {
this.messenger.sendInitalStream(tabID, triggerID)
response = await session.chat(request)
this.telemetryHelper.recordEnterFocusConversation(triggerEvent.tabID)
this.telemetryHelper.recordEnterFocusConversation(triggerEvent.tabID, triggerPayload.traceId)
this.telemetryHelper.recordStartConversation(triggerEvent, triggerPayload)

getLogger().info(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -293,6 +293,7 @@ export class Messenger {
relatedSuggestions: undefined,
triggerID,
messageID,
traceId: triggerPayload.traceId,
},
tabID
)
Expand Down
2 changes: 2 additions & 0 deletions packages/core/src/codewhispererChat/controllers/chat/model.ts
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,7 @@ export type ChatPromptCommandType =
export interface PromptMessage {
message: string | undefined
messageId: string
traceId: string
command: ChatPromptCommandType | undefined
userIntent: UserIntent | undefined
tabID: string
Expand Down Expand Up @@ -143,6 +144,7 @@ export interface TriggerPayload {
readonly customization: Customization
relevantTextDocuments?: RelevantTextDocument[]
useRelevantDocuments?: boolean
traceId?: string
}

export interface InsertedCode {
Expand Down
Loading

0 comments on commit 22db33e

Please sign in to comment.