[Obs AI Assistant] Improved error handling for tool response

x-pack/platform/plugins/shared/observability_ai_assistant/common/conversation_complete.ts

@@ -6,6 +6,7 @@
  */
 
 import { i18n } from '@kbn/i18n';
+import type { ErrorObject } from 'ajv';
 import type { ServerSentEventBase } from '@kbn/sse-utils';
 import type { DeanonymizationInput, DeanonymizationOutput } from './types';
 import { type Message } from './types';
@@ -112,8 +113,8 @@ export enum ChatCompletionErrorCode {
   InternalError = 'internalError',
   NotFoundError = 'notFoundError',
   TokenLimitReachedError = 'tokenLimitReachedError',
-  FunctionNotFoundError = 'functionNotFoundError',
   FunctionLimitExceededError = 'functionLimitExceededError',
+  FunctionArgsValidationError = 'functionArgsValidationError',
 }
 
 interface ErrorMetaAttributes {
@@ -123,10 +124,10 @@ interface ErrorMetaAttributes {
     tokenLimit?: number;
     tokenCount?: number;
   };
-  [ChatCompletionErrorCode.FunctionNotFoundError]: {
-    name: string;
-  };
   [ChatCompletionErrorCode.FunctionLimitExceededError]: {};
+  [ChatCompletionErrorCode.FunctionArgsValidationError]: {
+    errors: ErrorObject[];
+  };
 }
 
 export class ChatCompletionError<T extends ChatCompletionErrorCode> extends Error {
@@ -170,13 +171,6 @@ export function createInternalServerError(
   return new ChatCompletionError(ChatCompletionErrorCode.InternalError, originalErrorMessage);
 }
 
-export function createFunctionNotFoundError(name: string) {
-  return new ChatCompletionError(
-    ChatCompletionErrorCode.FunctionNotFoundError,
-    `Function "${name}" called but was not available`
-  );
-}
-
 export function createFunctionLimitExceededError() {
   return new ChatCompletionError(
     ChatCompletionErrorCode.FunctionLimitExceededError,
@@ -193,15 +187,22 @@ export function isTokenLimitReachedError(
   );
 }
 
-export function isFunctionNotFoundError(
+export function isChatCompletionError(error: Error): error is ChatCompletionError<any> {
+  return error instanceof ChatCompletionError;
+}
+export function createFunctionArgsValidationError(errors: ErrorObject[]) {
+  return new ChatCompletionError(
+    ChatCompletionErrorCode.FunctionArgsValidationError,
+    'Function arguments are invalid',
+    { errors }
+  );
+}
+
+export function isFunctionArgsValidationError(
   error: Error
-): error is ChatCompletionError<ChatCompletionErrorCode.FunctionNotFoundError> {
+): error is ChatCompletionError<ChatCompletionErrorCode.FunctionArgsValidationError> {
   return (
     error instanceof ChatCompletionError &&
-    error.code === ChatCompletionErrorCode.FunctionNotFoundError
+    error.code === ChatCompletionErrorCode.FunctionArgsValidationError
   );
 }
-
-export function isChatCompletionError(error: Error): error is ChatCompletionError<any> {
-  return error instanceof ChatCompletionError;
-}

x-pack/platform/plugins/shared/observability_ai_assistant/common/index.ts

@@ -39,7 +39,6 @@ export {
   createInternalServerError,
   isTokenLimitReachedError,
   isChatCompletionError,
-  createFunctionNotFoundError,
 } from './conversation_complete';
 
 export {

x-pack/platform/plugins/shared/observability_ai_assistant/server/service/chat_function_client/index.ts

@@ -4,11 +4,11 @@
  * 2.0; you may not use this file except in compliance with the Elastic License
  * 2.0.
  */
-/* eslint-disable max-classes-per-file*/
 
-import Ajv, { type ErrorObject, type ValidateFunction } from 'ajv';
+import Ajv, { type ValidateFunction } from 'ajv';
 import { compact, keyBy } from 'lodash';
 import type { Logger } from '@kbn/logging';
+import { createFunctionArgsValidationError } from '../../../common/conversation_complete';
 import { type FunctionResponse } from '../../../common/functions/types';
 import type { Message, ObservabilityAIAssistantScreenContextRequest } from '../../../common/types';
 import { filterFunctionDefinitions } from '../../../common/utils/filter_function_definitions';
@@ -22,12 +22,6 @@ import type {
 } from '../types';
 import { registerGetDataOnScreenFunction } from '../../functions/get_data_on_screen';
 
-export class FunctionArgsValidationError extends Error {
-  constructor(public readonly errors: ErrorObject[]) {
-    super('Function arguments are invalid');
-  }
-}
-
 const ajv = new Ajv({
   strict: false,
 });
@@ -71,7 +65,7 @@ export class ChatFunctionClient {
 
     const result = validator(parameters);
     if (!result) {
-      throw new FunctionArgsValidationError(validator.errors!);
+      throw createFunctionArgsValidationError(validator.errors!);
     }
   }
 

x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.ts

@@ -67,7 +67,6 @@ import type { ChatFunctionClient } from '../chat_function_client';
 import type { KnowledgeBaseService, RecalledEntry } from '../knowledge_base_service';
 import { getAccessQuery } from '../util/get_access_query';
 import { getSystemMessageFromInstructions } from '../util/get_system_message_from_instructions';
-import { failOnNonExistingFunctionCall } from './operators/fail_on_non_existing_function_call';
 import { getContextFunctionRequestIfNeeded } from './get_context_function_request_if_needed';
 import { continueConversation } from './operators/continue_conversation';
 import { convertInferenceEventsToStreamingEvents } from './operators/convert_inference_events_to_streaming_events';
@@ -578,7 +577,6 @@ export class ObservabilityAIAssistantClient {
         })
       ).pipe(
         convertInferenceEventsToStreamingEvents(),
-        failOnNonExistingFunctionCall({ functions }),
         tap((event) => {
           if (event.type === StreamingChatResponseEventType.ChatCompletionChunk) {
             this.dependencies.logger.trace(

x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/operators/catch_function_not_found_error.ts

@@ -8,15 +8,14 @@
 import type { OperatorFunction } from 'rxjs';
 import { catchError, filter, of, share, throwError } from 'rxjs';
 import { i18n } from '@kbn/i18n';
+import { isToolNotFoundError } from '@kbn/inference-common';
 import { MessageRole } from '../../../../common';
 import type {
   ChatCompletionChunkEvent,
+  MessageAddEvent,
   MessageOrChatEvent,
 } from '../../../../common/conversation_complete';
-import {
-  isFunctionNotFoundError,
-  StreamingChatResponseEventType,
-} from '../../../../common/conversation_complete';
+import { StreamingChatResponseEventType } from '../../../../common/conversation_complete';
 import { emitWithConcatenatedMessage } from '../../../../common/utils/emit_with_concatenated_message';
 
 function appendFunctionLimitExceededErrorMessageToAssistantResponse(): OperatorFunction<
@@ -75,13 +74,30 @@ export function catchFunctionNotFoundError(
 
     return shared$.pipe(
       catchError((error) => {
-        if (isFunctionNotFoundError(error)) {
+        if (isToolNotFoundError(error)) {
           if (functionLimitExceeded) {
             return chunksWithoutErrors$.pipe(
               appendFunctionLimitExceededErrorMessageToAssistantResponse()
             );
           }
-          return chunksWithoutErrors$.pipe(emitWithConcatenatedMessage());
+          // Instead of throwing error, return a message with the function name, to be handled by the function client
+          const simpleMessage: MessageAddEvent = {
+            type: StreamingChatResponseEventType.MessageAdd as const,
+            id: 'error_recovery_message',
+            message: {
+              '@timestamp': new Date().toISOString(),
+              message: {
+                content: '',
+                role: MessageRole.Assistant,
+                function_call: {
+                  name: error.meta.name,
+                  arguments: '',
+                  trigger: MessageRole.Assistant,
+                },
+              },
+            },
+          };
+          return of(simpleMessage);
         }
         return throwError(() => error);
       })

x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/operators/continue_conversation.ts

@@ -23,6 +23,7 @@ import {
   throwError,
 } from 'rxjs';
 import { withExecuteToolSpan } from '@kbn/inference-tracing';
+import { createToolNotFoundError } from '@kbn/inference-plugin/common/chat_complete/errors';
 import type { AnalyticsServiceStart } from '@kbn/core/server';
 import type { Connector } from '@kbn/actions-plugin/server';
 import type { AssistantScope } from '@kbn/ai-assistant-common';
@@ -32,12 +33,14 @@ import { toolCallEventType } from '../../../analytics/tool_call';
 import type { Message, CompatibleJSONSchema, MessageAddEvent } from '../../../../common';
 import {
   CONTEXT_FUNCTION_NAME,
-  createFunctionNotFoundError,
   MessageRole,
   StreamingChatResponseEventType,
 } from '../../../../common';
 import type { MessageOrChatEvent } from '../../../../common/conversation_complete';
-import { createFunctionLimitExceededError } from '../../../../common/conversation_complete';
+import {
+  createFunctionLimitExceededError,
+  isFunctionArgsValidationError,
+} from '../../../../common/conversation_complete';
 import type { Instruction } from '../../../../common/types';
 import { createFunctionResponseMessage } from '../../../../common/utils/create_function_response_message';
 import { emitWithConcatenatedMessage } from '../../../../common/utils/emit_with_concatenated_message';
@@ -111,6 +114,15 @@ export function executeFunctionAndCatchError({
       catchError((error) => {
         span?.recordException(error);
         logger.error(`Encountered error running function ${name}: ${JSON.stringify(error)}`);
+
+        if (isFunctionArgsValidationError(error)) {
+          return of(
+            createFunctionResponseMessage({
+              name,
+              content: { message: error.message, errors: error.meta.errors },
+            })
+          );
+        }
         // We want to catch the error only when a promise occurs
         // if it occurs in the Observable, we cannot easily recover
         // from it because the function may have already emitted
@@ -321,7 +333,7 @@ export function continueConversation({
       return of(
         createServerSideFunctionResponseError({
           name: functionCallName,
-          error: createFunctionNotFoundError(functionCallName),
+          error: createToolNotFoundError(functionCallName),
         })
       );
     }

x-pack/platform/plugins/shared/observability_ai_assistant/tsconfig.json

@@ -62,7 +62,7 @@
     "@kbn/product-doc-base-plugin",
     "@kbn/inference-endpoint-plugin",
     "@kbn/spaces-utils",
-    "@kbn/usage-collection-plugin"
+    "@kbn/usage-collection-plugin",
   ],
   "exclude": ["target/**/*"]
 }