[Obs AI Assistant] Remove TokenCountEvent (elastic#209549)

Closes elastic#205479 This filters out the `ChatCompletionTokenCountEvent` from the inference plugin. This greatly simplifies handling ChatCompletion events in the Obs AI Assistant. (cherry picked from commit c4826bd)
kibanamachine · Feb 20, 2025 · 29a3e77 · 29a3e77
1 parent 09125b6
commit 29a3e77
Show file tree

Hide file tree

Showing 26 changed files with 100 additions and 545 deletions.
diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/common/conversation_complete.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/common/conversation_complete.ts
@@ -6,7 +6,7 @@
  */
 
 import { i18n } from '@kbn/i18n';
-import { TokenCount as TokenCountType, type Message } from './types';
+import { type Message } from './types';
 
 export enum StreamingChatResponseEventType {
   ChatCompletionChunk = 'chatCompletionChunk',
@@ -16,7 +16,6 @@ export enum StreamingChatResponseEventType {
   MessageAdd = 'messageAdd',
   ChatCompletionError = 'chatCompletionError',
   BufferFlush = 'bufferFlush',
-  TokenCount = 'tokenCount',
 }
 
 type StreamingChatResponseEventBase<
@@ -54,7 +53,6 @@ export type ConversationCreateEvent = StreamingChatResponseEventBase<
       id: string;
       title: string;
       last_updated: string;
-      token_count?: TokenCountType;
     };
   }
 >;
@@ -66,7 +64,6 @@ export type ConversationUpdateEvent = StreamingChatResponseEventBase<
       id: string;
       title: string;
       last_updated: string;
-      token_count?: TokenCountType;
     };
   }
 >;
@@ -95,33 +92,21 @@ export type BufferFlushEvent = StreamingChatResponseEventBase<
   }
 >;
 
-export type TokenCountEvent = StreamingChatResponseEventBase<
-  StreamingChatResponseEventType.TokenCount,
-  {
-    tokens: {
-      completion: number;
-      prompt: number;
-      total: number;
-    };
-  }
->;
-
 export type StreamingChatResponseEvent =
   | ChatCompletionChunkEvent
   | ChatCompletionMessageEvent
   | ConversationCreateEvent
   | ConversationUpdateEvent
   | MessageAddEvent
   | ChatCompletionErrorEvent
-  | TokenCountEvent
   | BufferFlushEvent;
 
 export type StreamingChatResponseEventWithoutError = Exclude<
   StreamingChatResponseEvent,
   ChatCompletionErrorEvent
 >;
 
-export type ChatEvent = ChatCompletionChunkEvent | TokenCountEvent | ChatCompletionMessageEvent;
+export type ChatEvent = ChatCompletionChunkEvent | ChatCompletionMessageEvent;
 export type MessageOrChatEvent = ChatEvent | MessageAddEvent;
 
 export enum ChatCompletionErrorCode {

diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/common/index.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/common/index.ts
@@ -18,7 +18,6 @@ export {
 export type {
   ChatCompletionChunkEvent,
   ChatCompletionMessageEvent,
-  TokenCountEvent,
   ConversationCreateEvent,
   ConversationUpdateEvent,
   MessageAddEvent,

diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/common/types.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/common/types.ts
@@ -45,12 +45,6 @@ export interface Message {
   };
 }
 
-export interface TokenCount {
-  prompt: number;
-  completion: number;
-  total: number;
-}
-
 export interface Conversation {
   '@timestamp': string;
   user?: {
@@ -61,7 +55,6 @@ export interface Conversation {
     id: string;
     title: string;
     last_updated: string;
-    token_count?: TokenCount;
   };
   systemMessage?: string;
   messages: Message[];
@@ -71,8 +64,8 @@ export interface Conversation {
   public: boolean;
 }
 
-export type ConversationRequestBase = Omit<Conversation, 'user' | 'conversation' | 'namespace'> & {
-  conversation: { title: string; token_count?: TokenCount; id?: string };
+type ConversationRequestBase = Omit<Conversation, 'user' | 'conversation' | 'namespace'> & {
+  conversation: { title: string; id?: string };
 };
 
 export type ConversationCreateRequest = ConversationRequestBase;

diff --git a/.../plugins/shared/observability_ai_assistant/common/utils/emit_with_concatenated_message.ts b/.../plugins/shared/observability_ai_assistant/common/utils/emit_with_concatenated_message.ts
@@ -16,7 +16,6 @@ import {
   withLatestFrom,
   filter,
 } from 'rxjs';
-import { withoutTokenCountEvents } from './without_token_count_events';
 import {
   type ChatCompletionChunkEvent,
   ChatEvent,
@@ -69,15 +68,12 @@ export function emitWithConcatenatedMessage<T extends ChatEvent>(
   return (source$) => {
     const shared = source$.pipe(shareReplay());
 
-    const withoutTokenCount$ = shared.pipe(filterChunkEvents());
-
     const response$ = concat(
       shared,
       shared.pipe(
-        withoutTokenCountEvents(),
         concatenateChatCompletionChunks(),
         last(),
-        withLatestFrom(withoutTokenCount$),
+        withLatestFrom(shared.pipe(filterChunkEvents())),
         mergeMap(([message, chunkEvent]) => {
           return mergeWithEditedMessage(message, chunkEvent, callback);
         })

diff --git a/...form/plugins/shared/observability_ai_assistant/common/utils/without_token_count_events.ts b/...form/plugins/shared/observability_ai_assistant/common/utils/without_token_count_events.ts
diff --git a/...tform/plugins/shared/observability_ai_assistant/public/analytics/schemas/chat_feedback.ts b/...tform/plugins/shared/observability_ai_assistant/public/analytics/schemas/chat_feedback.ts
@@ -65,28 +65,6 @@ export const chatFeedbackEventSchema: EventTypeOpts<ChatFeedback> = {
                 description: 'The timestamp of the last message in the conversation.',
               },
             },
-            token_count: {
-              properties: {
-                completion: {
-                  type: 'long',
-                  _meta: {
-                    description: 'The number of tokens in the completion.',
-                  },
-                },
-                prompt: {
-                  type: 'long',
-                  _meta: {
-                    description: 'The number of tokens in the prompt.',
-                  },
-                },
-                total: {
-                  type: 'long',
-                  _meta: {
-                    description: 'The total number of tokens in the conversation.',
-                  },
-                },
-              },
-            },
           },
         },
         labels: {

diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/routes/chat/route.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/routes/chat/route.ts
@@ -13,7 +13,6 @@ import { Readable } from 'stream';
 import { AssistantScope } from '@kbn/ai-assistant-common';
 import { aiAssistantSimulatedFunctionCalling } from '../..';
 import { createFunctionResponseMessage } from '../../../common/utils/create_function_response_message';
-import { withoutTokenCountEvents } from '../../../common/utils/without_token_count_events';
 import { LangTracer } from '../../service/client/instrumentation/lang_tracer';
 import { flushBuffer } from '../../service/util/flush_buffer';
 import { observableIntoOpenAIStream } from '../../service/util/observable_into_openai_stream';
@@ -203,16 +202,14 @@ const chatRecallRoute = createObservabilityAIAssistantServerRoute({
       recallAndScore({
         analytics: (await resources.plugins.core.start()).analytics,
         chat: (name, params) =>
-          client
-            .chat(name, {
-              ...params,
-              stream: true,
-              connectorId,
-              simulateFunctionCalling,
-              signal,
-              tracer: new LangTracer(otelContext.active()),
-            })
-            .pipe(withoutTokenCountEvents()),
+          client.chat(name, {
+            ...params,
+            stream: true,
+            connectorId,
+            simulateFunctionCalling,
+            signal,
+            tracer: new LangTracer(otelContext.active()),
+          }),
         context,
         logger: resources.logger,
         messages: [],

diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/server/routes/runtime_types.ts b/x-pack/platform/plugins/shared/observability_ai_assistant/server/routes/runtime_types.ts
@@ -7,9 +7,7 @@
 import * as t from 'io-ts';
 import { toBooleanRt } from '@kbn/io-ts-utils';
 import {
-  type Conversation,
   type ConversationCreateRequest,
-  type ConversationRequestBase,
   type ConversationUpdateRequest,
   type Message,
   MessageRole,
@@ -57,17 +55,12 @@ const tokenCountRt = t.type({
   total: t.number,
 });
 
-export const baseConversationRt: t.Type<ConversationRequestBase> = t.intersection([
+export const conversationCreateRt: t.Type<ConversationCreateRequest> = t.intersection([
   t.type({
     '@timestamp': t.string,
-    conversation: t.intersection([
-      t.type({
-        title: t.string,
-      }),
-      t.partial({
-        token_count: tokenCountRt,
-      }),
-    ]),
+    conversation: t.type({
+      title: t.string,
+    }),
     messages: t.array(messageRt),
     labels: t.record(t.string, t.string),
     numeric_labels: t.record(t.string, t.number),
@@ -84,51 +77,21 @@ export const assistantScopeType = t.union([
   t.literal('all'),
 ]);
 
-export const conversationCreateRt: t.Type<ConversationCreateRequest> = t.intersection([
-  baseConversationRt,
-  t.type({
-    conversation: t.type({
-      title: t.string,
-    }),
-  }),
-]);
-
 export const conversationUpdateRt: t.Type<ConversationUpdateRequest> = t.intersection([
-  baseConversationRt,
+  conversationCreateRt,
   t.type({
     conversation: t.intersection([
       t.type({
         id: t.string,
         title: t.string,
       }),
       t.partial({
-        token_count: tokenCountRt,
+        token_count: tokenCountRt, // deprecated, but kept for backwards compatibility
       }),
     ]),
   }),
 ]);
 
-export const conversationRt: t.Type<Conversation> = t.intersection([
-  baseConversationRt,
-  t.intersection([
-    t.type({
-      namespace: t.string,
-      conversation: t.intersection([
-        t.type({
-          id: t.string,
-          last_updated: t.string,
-        }),
-        t.partial({
-          token_count: tokenCountRt,
-        }),
-      ]),
-    }),
-    t.partial({
-      user: t.intersection([t.type({ name: t.string }), t.partial({ id: t.string })]),
-    }),
-  ]),
-]);
-
 export const functionRt = t.intersection([
   t.type({
     name: t.string,