[Agent Builder] Expose thinking complete event and update round to include time to first and last token (#241681)

Author: chrisbmar

x-pack/platform/packages/shared/onechat/onechat-common/chat/conversation.ts

@@ -151,17 +151,33 @@ export interface ConversationRound {
   steps: ConversationRoundStep[];
   /** The final response from the assistant */
   response: AssistantResponse;
+  /** when the round was started */
+  started_at: string;
+  /** time it took to first token, in ms */
+  time_to_first_token: number;
+  /** time it took to last token, in ms */
+  time_to_last_token: number;
   /** when tracing is enabled, contains the traceId associated with this round */
   trace_id?: string;
 }
 
+/**
+ * Main structure representing a conversation with an agent.
+ */
 export interface Conversation {
+  /** unique id for this conversation */
   id: string;
+  /** id of the agent this conversation is bound to */
   agent_id: string;
+  /** info of the owner of the discussion */
   user: UserIdAndName;
+  /** title of the conversation */
   title: string;
+  /** creation date */
   created_at: string;
+  /** update date */
   updated_at: string;
+  /** list of round for this conversation */
   rounds: ConversationRound[];
 }
 

x-pack/platform/packages/shared/onechat/onechat-common/chat/events.ts

@@ -16,6 +16,7 @@ export enum ChatEventType {
   reasoning = 'reasoning',
   messageChunk = 'message_chunk',
   messageComplete = 'message_complete',
+  thinkingComplete = 'thinking_complete',
   roundComplete = 'round_complete',
   conversationCreated = 'conversation_created',
   conversationUpdated = 'conversation_updated',
@@ -122,6 +123,24 @@ export const isMessageCompleteEvent = (
   return event.type === ChatEventType.messageComplete;
 };
 
+// Thinking complete
+
+export interface ThinkingCompleteEventData {
+  /** time elapsed from round start to first token arrival, in ms */
+  time_to_first_token: number;
+}
+
+export type ThinkingCompleteEvent = ChatEventBase<
+  ChatEventType.thinkingComplete,
+  ThinkingCompleteEventData
+>;
+
+export const isThinkingCompleteEvent = (
+  event: OnechatEvent<string, any>
+): event is ThinkingCompleteEvent => {
+  return event.type === ChatEventType.thinkingComplete;
+};
+
 // Round complete
 
 export interface RoundCompleteEventData {
@@ -200,6 +219,7 @@ export type ChatAgentEvent =
   | ReasoningEvent
   | MessageChunkEvent
   | MessageCompleteEvent
+  | ThinkingCompleteEvent
   | RoundCompleteEvent;
 
 /**

x-pack/platform/packages/shared/onechat/onechat-common/chat/index.ts

@@ -46,6 +46,8 @@ export {
   type MessageChunkEvent,
   type MessageCompleteEventData,
   type MessageCompleteEvent,
+  type ThinkingCompleteEventData,
+  type ThinkingCompleteEvent,
   type RoundCompleteEventData,
   type RoundCompleteEvent,
   isToolCallEvent,
@@ -54,6 +56,7 @@ export {
   isReasoningEvent,
   isMessageChunkEvent,
   isMessageCompleteEvent,
+  isThinkingCompleteEvent,
   isRoundCompleteEvent,
   isConversationCreatedEvent,
   isConversationUpdatedEvent,

x-pack/platform/packages/shared/onechat/onechat-common/index.ts

@@ -123,6 +123,8 @@ export {
   type MessageChunkEvent,
   type MessageCompleteEventData,
   type MessageCompleteEvent,
+  type ThinkingCompleteEventData,
+  type ThinkingCompleteEvent,
   type RoundCompleteEventData,
   type RoundCompleteEvent,
   type ToolCallProgress,
@@ -131,6 +133,7 @@ export {
   isReasoningEvent,
   isMessageChunkEvent,
   isMessageCompleteEvent,
+  isThinkingCompleteEvent,
   isRoundCompleteEvent,
   isConversationCreatedEvent,
   isConversationUpdatedEvent,

x-pack/platform/packages/shared/onechat/onechat-genai-utils/langchain/graph_events.ts

@@ -10,6 +10,7 @@ import type {
   MessageChunkEvent,
   MessageCompleteEvent,
   ReasoningEvent,
+  ThinkingCompleteEvent,
   ToolCallEvent,
   ToolResultEvent,
 } from '@kbn/onechat-common';
@@ -108,3 +109,12 @@ export const createReasoningEvent = (
     },
   };
 };
+
+export const createThinkingCompleteEvent = (timeToFirstToken: number): ThinkingCompleteEvent => {
+  return {
+    type: ChatEventType.thinkingComplete,
+    data: {
+      time_to_first_token: timeToFirstToken,
+    },
+  };
+};

x-pack/platform/packages/shared/onechat/onechat-genai-utils/langchain/index.ts

@@ -17,6 +17,7 @@ export {
   createReasoningEvent,
   createToolCallEvent,
   createToolResultEvent,
+  createThinkingCompleteEvent,
 } from './graph_events';
 export {
   extractTextContent,

x-pack/platform/plugins/shared/onechat/public/application/components/conversations/conversation_rounds/round_thinking/round_thinking.tsx

@@ -13,6 +13,7 @@ import React, { useState } from 'react';
 import { useSendMessage } from '../../../../context/send_message/send_message_context';
 import { RoundFlyout } from '../round_flyout';
 import { RoundSteps } from './steps/round_steps';
+import { ThinkingTimeDisplay } from './thinking_time_display';
 
 interface RoundThinkingProps {
   rawRound: ConversationRound;
@@ -91,6 +92,7 @@ export const RoundThinking: React.FC<RoundThinkingProps> = ({ steps, isLoading,
             {rawResponseButtonLabel}
           </EuiButton>
         )}
+        <ThinkingTimeDisplay timeToFirstToken={rawRound.time_to_first_token} />
       </EuiPanel>
       <RoundFlyout isOpen={showFlyout} onClose={toggleFlyout} rawRound={rawRound} />
     </EuiAccordion>

x-pack/platform/plugins/shared/onechat/public/application/components/conversations/conversation_rounds/round_thinking/thinking_time_display.tsx

@@ -0,0 +1,37 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import React from 'react';
+import { EuiSpacer, EuiText } from '@elastic/eui';
+import { FormattedMessage } from '@kbn/i18n-react';
+import { css } from '@emotion/react';
+
+export const ThinkingTimeDisplay = ({ timeToFirstToken }: { timeToFirstToken: number }) => {
+  if (timeToFirstToken === 0) {
+    return null;
+  }
+  return (
+    <>
+      <EuiSpacer size="m" />
+      <EuiText
+        size="s"
+        color="subdued"
+        css={css`
+          font-style: italic;
+        `}
+      >
+        <FormattedMessage
+          id="xpack.onechat.conversation.thinking.timeToFirstToken"
+          defaultMessage="Thought for {timeToFirstToken} seconds."
+          values={{
+            timeToFirstToken: Math.round(timeToFirstToken / 1000),
+          }}
+        />
+      </EuiText>
+    </>
+  );
+};

x-pack/platform/plugins/shared/onechat/public/application/context/conversation/use_conversation_actions.ts

@@ -49,6 +49,7 @@ export interface ConversationActions {
   }) => void;
   setAssistantMessage: ({ assistantMessage }: { assistantMessage: string }) => void;
   addAssistantMessageChunk: ({ messageChunk }: { messageChunk: string }) => void;
+  setTimeToFirstToken: ({ timeToFirstToken }: { timeToFirstToken: number }) => void;
   onConversationCreated: ({
     conversationId,
     title,
@@ -110,6 +111,9 @@ const createConversationActions = ({
             input: { message: userMessage },
             response: { message: '' },
             steps: [],
+            started_at: new Date().toISOString(),
+            time_to_first_token: 0,
+            time_to_last_token: 0,
           };
 
           if (!draft) {
@@ -129,7 +133,6 @@ const createConversationActions = ({
         })
       );
     },
-
     setAgentId: (agentId: string) => {
       // We allow to change agent only at the start of the conversation
       if (conversationId) {
@@ -148,7 +151,6 @@ const createConversationActions = ({
       );
       setAgentIdStorage(agentId);
     },
-
     addReasoningStep: ({ step }: { step: ReasoningStep }) => {
       setCurrentRound((round) => {
         round.steps.push(step);
@@ -194,7 +196,11 @@ const createConversationActions = ({
         round.response.message += messageChunk;
       });
     },
-
+    setTimeToFirstToken: ({ timeToFirstToken }: { timeToFirstToken: number }) => {
+      setCurrentRound((round) => {
+        round.time_to_first_token = timeToFirstToken;
+      });
+    },
     onConversationCreated: ({
       conversationId: id,
       title,
@@ -224,7 +230,6 @@ const createConversationActions = ({
         onConversationCreated({ conversationId: id, title });
       }
     },
-
     deleteConversation: async (id: string) => {
       await conversationsService.delete({ conversationId: id });
 

x-pack/platform/plugins/shared/onechat/public/application/context/send_message/use_subscribe_to_chat_events.ts

@@ -15,6 +15,7 @@ import {
   isToolCallEvent,
   isToolProgressEvent,
   isToolResultEvent,
+  isThinkingCompleteEvent,
 } from '@kbn/onechat-common';
 import { createReasoningStep, createToolCallStep } from '@kbn/onechat-common/chat/conversation';
 import type { Observable } from 'rxjs';
@@ -77,6 +78,10 @@ export const useSubscribeToChatEvents = ({
           } else if (isConversationCreatedEvent(event)) {
             const { conversation_id: id, title } = event.data;
             conversationActions.onConversationCreated({ conversationId: id, title });
+          } else if (isThinkingCompleteEvent(event)) {
+            conversationActions.setTimeToFirstToken({
+              timeToFirstToken: event.data.time_to_first_token,
+            });
           }
         },
         complete: () => {