From 374456d5816d00b2980f8597500b7bd04e3388ba Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=B8ren=20Louv-Jansen?= <soren.louv@elastic.co>
Date: Mon, 9 Dec 2024 15:00:05 +0100
Subject: [PATCH] [Obs AI Assistant] Fix null pointer in function definition
 (#203344)

Closes #201713

(cherry picked from commit 1d9ca1ebf66f4cb8a367de8a8854d40dd4789ec8)
---
 .../kbn-ai-assistant/src/hooks/use_json_editor_model.ts  | 2 +-
 .../src/utils/create_initialized_object.ts               | 2 +-
 .../observability_ai_assistant/common/functions/types.ts | 4 +++-
 .../server/functions/query/index.ts                      | 9 +++++----
 4 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/x-pack/packages/kbn-ai-assistant/src/hooks/use_json_editor_model.ts b/x-pack/packages/kbn-ai-assistant/src/hooks/use_json_editor_model.ts
index 6412aeb6a870e..5cba3fbb78d75 100644
--- a/x-pack/packages/kbn-ai-assistant/src/hooks/use_json_editor_model.ts
+++ b/x-pack/packages/kbn-ai-assistant/src/hooks/use_json_editor_model.ts
@@ -43,7 +43,7 @@ export const useJsonEditorModel = ({
 
     const initialJsonString = initialJsonValue
       ? JSON.stringify(safeJsonParse(initialJsonValue), null, 4) // prettify the json
-      : functionDefinition.parameters.properties
+      : functionDefinition.parameters?.properties
       ? JSON.stringify(createInitializedObject(functionDefinition.parameters), null, 4)
       : '';
 
diff --git a/x-pack/packages/kbn-ai-assistant/src/utils/create_initialized_object.ts b/x-pack/packages/kbn-ai-assistant/src/utils/create_initialized_object.ts
index e06800aca07a0..22d173f7da67c 100644
--- a/x-pack/packages/kbn-ai-assistant/src/utils/create_initialized_object.ts
+++ b/x-pack/packages/kbn-ai-assistant/src/utils/create_initialized_object.ts
@@ -7,7 +7,7 @@
 
 import type { FunctionDefinition } from '@kbn/observability-ai-assistant-plugin/common';
 
-type Params = FunctionDefinition['parameters'];
+type Params = NonNullable<FunctionDefinition['parameters']>;
 
 export function createInitializedObject(parameters: Params) {
   const emptyObject: Record<string, string | any> = {};
diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/common/functions/types.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/common/functions/types.ts
index bd786e9ba3c75..093aed3aa7120 100644
--- a/x-pack/plugins/observability_solution/observability_ai_assistant/common/functions/types.ts
+++ b/x-pack/plugins/observability_solution/observability_ai_assistant/common/functions/types.ts
@@ -35,7 +35,9 @@ export type FunctionResponse =
     }
   | Observable<ChatCompletionChunkEvent | MessageAddEvent>;
 
-export interface FunctionDefinition<TParameters extends CompatibleJSONSchema = any> {
+export interface FunctionDefinition<
+  TParameters extends CompatibleJSONSchema = CompatibleJSONSchema
+> {
   name: string;
   description: string;
   visibility?: FunctionVisibility;
diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant_app/server/functions/query/index.ts b/x-pack/plugins/observability_solution/observability_ai_assistant_app/server/functions/query/index.ts
index cefec5ae66758..e07a1f140a9ac 100644
--- a/x-pack/plugins/observability_solution/observability_ai_assistant_app/server/functions/query/index.ts
+++ b/x-pack/plugins/observability_solution/observability_ai_assistant_app/server/functions/query/index.ts
@@ -5,7 +5,7 @@
  * 2.0.
  */
 
-import { isChatCompletionChunkEvent, isOutputEvent } from '@kbn/inference-common';
+import { ToolDefinition, isChatCompletionChunkEvent, isOutputEvent } from '@kbn/inference-common';
 import { correctCommonEsqlMistakes } from '@kbn/inference-plugin/common';
 import { naturalLanguageToEsql } from '@kbn/inference-plugin/server';
 import {
@@ -132,9 +132,10 @@ export function registerQueryFunction({
         ),
         logger: resources.logger,
         tools: Object.fromEntries(
-          actions
-            .concat(esqlFunctions)
-            .map((fn) => [fn.name, { description: fn.description, schema: fn.parameters }])
+          [...actions, ...esqlFunctions].map((fn) => [
+            fn.name,
+            { description: fn.description, schema: fn.parameters } as ToolDefinition,
+          ])
         ),
         functionCalling: useSimulatedFunctionCalling ? 'simulated' : 'native',
       });