[Obs AI Assistant] Boost user prompt in recall (#184933)

Closes: #180995 --------- Co-authored-by: Dario Gieselaar <[email protected]> (cherry picked from commit baa22bb)
elastic · Dec 10, 2024 · 164dc34 · 164dc34
1 parent 4fa7135
commit 164dc34
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 11 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
@@ -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
@@ -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
@@ -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/...ins/observability_solution/observability_ai_assistant_app/server/functions/query/index.ts b/...ins/observability_solution/observability_ai_assistant_app/server/functions/query/index.ts
@@ -5,11 +5,8 @@
  * 2.0.
  */
 
-import {
-  correctCommonEsqlMistakes,
-  isChatCompletionChunkEvent,
-  isOutputEvent,
-} from '@kbn/inference-plugin/common';
+import { ToolDefinition, isChatCompletionChunkEvent, isOutputEvent } from '@kbn/inference-common';
+import { correctCommonEsqlMistakes } from '@kbn/inference-plugin/common';
 import { naturalLanguageToEsql } from '@kbn/inference-plugin/server';
 import {
   FunctionVisibility,
@@ -135,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',
       });