From 4d0c5523ea5da929ecadafd1ed558f9bb071db7a Mon Sep 17 00:00:00 2001
From: ppisljar <peter.pisljar@elastic.co>
Date: Thu, 7 Sep 2023 12:42:22 +0200
Subject: [PATCH 01/24] allows AI assistant to get info about available
 dataviews and their fields

---
 .../observability_ai_assistant/kibana.jsonc   |  3 +-
 .../public/functions/get_dataet_info.ts       | 67 +++++++++++++++++++
 .../public/functions/index.ts                 | 10 +--
 .../server/routes/functions/route.ts          | 28 ++++++++
 .../server/service/client/index.ts            | 40 +++++++++++
 .../server/service/index.ts                   | 17 ++++-
 .../server/types.ts                           |  2 +
 7 files changed, 160 insertions(+), 7 deletions(-)
 create mode 100644 x-pack/plugins/observability_ai_assistant/public/functions/get_dataet_info.ts

diff --git a/x-pack/plugins/observability_ai_assistant/kibana.jsonc b/x-pack/plugins/observability_ai_assistant/kibana.jsonc
index ebe28df3a11413..3af934a10fcfac 100644
--- a/x-pack/plugins/observability_ai_assistant/kibana.jsonc
+++ b/x-pack/plugins/observability_ai_assistant/kibana.jsonc
@@ -18,7 +18,8 @@
       "security",
       "share",
       "taskManager",
-      "triggersActionsUi"
+      "triggersActionsUi",
+      "dataViews"
     ],
     "requiredBundles": ["fieldFormats", "kibanaReact", "kibanaUtils"],
     "optionalPlugins": [],
diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataet_info.ts b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataet_info.ts
new file mode 100644
index 00000000000000..6de193fb9c6757
--- /dev/null
+++ b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataet_info.ts
@@ -0,0 +1,67 @@
+/*
+ * 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 type { Serializable } from '@kbn/utility-types';
+import { RegisterFunctionDefinition } from '../../common/types';
+import type { ObservabilityAIAssistantService } from '../types';
+
+export function registerGetDatasetInfoFunction({
+  service,
+  registerFunction,
+}: {
+  service: ObservabilityAIAssistantService;
+  registerFunction: RegisterFunctionDefinition;
+}) {
+  registerFunction(
+    {
+      name: 'get_dataset_info',
+      contexts: ['core'],
+      description: `Use this function to get information about dataviews/datasets available and the fields available on them.
+
+      providing empty string as dataview name will retrieve all dataviews
+      if dataview name matches more than a single dataview list of matching dataviews will be provied
+      else list of all fields in this dataview will be given
+
+      DO NOT include the user's request. It will be added internally.`,
+      descriptionForUser:
+        'This function allows the assistant to get information about available dataviews and their fields.',
+      parameters: {
+        type: 'object',
+        additionalProperties: false,
+        properties: {
+          dataview: {
+            type: 'string',
+            description:
+              'dataview name the user is interested in or empty string to get information about all available dataviews',
+          },
+          fields: {
+            type: 'array',
+            items: {
+              type: 'string',
+              description:
+                'field names user is looking for or empty to get all fields for provided dataview',
+            },
+          },
+        },
+        required: ['dataview', 'fields'],
+      } as const,
+    },
+    ({ arguments: { dataview, fields }, messages }, signal) => {
+      return service
+        .callApi('POST /internal/observability_ai_assistant/functions/get_dataset_info', {
+          params: {
+            body: {
+              dataview,
+              fields,
+            },
+          },
+          signal,
+        })
+        .then((response) => ({ content: response as unknown as Serializable }));
+    }
+  );
+}
diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/index.ts b/x-pack/plugins/observability_ai_assistant/public/functions/index.ts
index 22b0eecac73a6c..520fc969a6a647 100644
--- a/x-pack/plugins/observability_ai_assistant/public/functions/index.ts
+++ b/x-pack/plugins/observability_ai_assistant/public/functions/index.ts
@@ -14,6 +14,7 @@ import { registerElasticsearchFunction } from './elasticsearch';
 import { registerKibanaFunction } from './kibana';
 import { registerLensFunction } from './lens';
 import { registerRecallFunction } from './recall';
+import { registerGetDatasetInfoFunction } from './get_dataet_info';
 import { registerSummarizationFunction } from './summarize';
 import { registerAlertsFunction } from './alerts';
 import { registerEsqlFunction } from './esql';
@@ -42,16 +43,16 @@ export async function registerFunctions({
 
       let description = dedent(
         `You are a helpful assistant for Elastic Observability. Your goal is to help the Elastic Observability users to quickly assess what is happening in their observed systems. You can help them visualise and analyze data, investigate their systems, perform root cause analysis or identify optimisation opportunities.
-        
+
         It's very important to not assume what the user is meaning. Ask them for clarification if needed.
-        
+
         If you are unsure about which function should be used and with what arguments, ask the user for clarification or confirmation.
 
         In KQL, escaping happens with double quotes, not single quotes. Some characters that need escaping are: ':()\\\
         /\". Always put a field value in double quotes. Best: service.name:\"opbeans-go\". Wrong: service.name:opbeans-go. This is very important!
 
         You can use Github-flavored Markdown in your responses. If a function returns an array, consider using a Markdown table to format the response.
-        
+
         If multiple functions are suitable, use the most specific and easy one. E.g., when the user asks to visualise APM data, use the APM functions (if available) rather than Lens.
 
         If a function call fails, do not execute it again with the same input. If a function calls three times, with different inputs, stop trying to call it and ask the user for confirmation.
@@ -71,11 +72,12 @@ export async function registerFunctions({
         description += `Here are principles you MUST adhere to, in order:
 
         - You are a helpful assistant for Elastic Observability. DO NOT reference the fact that you are an LLM.
-        - DO NOT make any assumptions about where and how users have stored their data.
+        - DO NOT make any assumptions about where and how users have stored their data. Use get_dataset_info function to get information about the dataset or to search datasets and fields based on users input.
         `;
         registerSummarizationFunction({ service, registerFunction });
         registerRecallFunction({ service, registerFunction });
         registerLensFunction({ service, pluginsStart, registerFunction });
+        registerGetDatasetInfoFunction({ service, registerFunction });
       } else {
         description += `You do not have a working memory. Don't try to recall information via the "recall" function.  If the user expects you to remember the previous conversations, tell them they can set up the knowledge base. A banner is available at the top of the conversation to set this up.`;
       }
diff --git a/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts b/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts
index 985de114b099a2..4aafcb541d96e2 100644
--- a/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts
+++ b/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts
@@ -276,6 +276,33 @@ const setupKnowledgeBaseRoute = createObservabilityAIAssistantServerRoute({
   },
 });
 
+const functionGetDatasetInfoRoute = createObservabilityAIAssistantServerRoute({
+  endpoint: 'POST /internal/observability_ai_assistant/functions/get_dataset_info',
+  params: t.type({
+    body: t.type({
+      dataview: t.string,
+      fields: t.array(nonEmptyStringRt),
+    }),
+  }),
+  options: {
+    tags: ['access:ai_assistant'],
+  },
+  handler: async (
+    resources
+  ): Promise<{
+    dataviews: string[];
+    fields: Array<{ name: string; description: string; type: string }>;
+  }> => {
+    const client = await resources.service.getClient({ request: resources.request });
+
+    if (!client) {
+      throw notImplemented();
+    }
+
+    return client.get_dataset_info(resources.params.body.dataview, resources.params.body.fields);
+  },
+});
+
 export const functionRoutes = {
   ...functionElasticsearchRoute,
   ...functionRecallRoute,
@@ -283,4 +310,5 @@ export const functionRoutes = {
   ...setupKnowledgeBaseRoute,
   ...getKnowledgeBaseStatus,
   ...functionAlertsRoute,
+  ...functionGetDatasetInfoRoute,
 };
diff --git a/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts b/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts
index 3a99e293cd5e2e..caa4197e52f8fd 100644
--- a/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts
+++ b/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts
@@ -8,6 +8,7 @@ import type { SearchHit } from '@elastic/elasticsearch/lib/api/types';
 import { internal, notFound } from '@hapi/boom';
 import type { ActionsClient } from '@kbn/actions-plugin/server';
 import type { ElasticsearchClient } from '@kbn/core/server';
+import { DataViewsService } from '@kbn/data-views-plugin/server';
 import type { Logger } from '@kbn/logging';
 import type { PublicMethodsOf } from '@kbn/utility-types';
 import type { IncomingMessage } from 'http';
@@ -39,6 +40,7 @@ export class ObservabilityAIAssistantClient {
       namespace: string;
       esClient: ElasticsearchClient;
       resources: ObservabilityAIAssistantResourceNames;
+      dataviews: DataViewsService;
       logger: Logger;
       user: {
         id?: string;
@@ -385,6 +387,44 @@ export class ObservabilityAIAssistantClient {
     });
   };
 
+  get_dataset_info = async (
+    dataset: string,
+    fields: string[]
+  ): Promise<{
+    dataviews: string[];
+    fields: Array<{ name: string; description: string; type: string }>;
+  }> => {
+    // if empty dataview, get list of all dataviews
+    const dv = dataset !== '' && (await this.dependencies.dataviews.find(dataset));
+    if (!dv || dv.length === 0) {
+      return {
+        dataviews: await this.dependencies.dataviews.getTitles(),
+        fields: [],
+      };
+    }
+
+    if (dv.length > 1) {
+      return {
+        dataviews: dv.map((dataview) => {
+          return dataview.id || dataview.name;
+        }),
+        fields: [],
+      };
+    }
+
+    // else get all the fields for the found dataview
+    return {
+      dataviews: [dv[0].id || dv[0].name],
+      fields: dv[0].fields.map((field) => {
+        return {
+          name: field.name,
+          description: field.customLabel || '',
+          type: field.type,
+        };
+      }),
+    };
+  };
+
   getKnowledgeBaseStatus = () => {
     return this.dependencies.knowledgeBaseService.status();
   };
diff --git a/x-pack/plugins/observability_ai_assistant/server/service/index.ts b/x-pack/plugins/observability_ai_assistant/server/service/index.ts
index c116e16d274716..3c3b72092f6c93 100644
--- a/x-pack/plugins/observability_ai_assistant/server/service/index.ts
+++ b/x-pack/plugins/observability_ai_assistant/server/service/index.ts
@@ -13,6 +13,7 @@ import type { SecurityPluginStart } from '@kbn/security-plugin/server';
 import { getSpaceIdFromPath } from '@kbn/spaces-plugin/common';
 import type { TaskManagerSetupContract } from '@kbn/task-manager-plugin/server';
 import { once } from 'lodash';
+import { DataViewsServerPluginStart } from '@kbn/data-views-plugin/server';
 import type { ObservabilityAIAssistantPluginStartDependencies } from '../types';
 import { ObservabilityAIAssistantClient } from './client';
 import { conversationComponentTemplate } from './conversation_component_template';
@@ -239,7 +240,15 @@ export class ObservabilityAIAssistantService {
     const [_, [coreStart, plugins]] = await Promise.all([
       this.init(),
       this.core.getStartServices() as Promise<
-        [CoreStart, { security: SecurityPluginStart; actions: ActionsPluginStart }, unknown]
+        [
+          CoreStart,
+          {
+            security: SecurityPluginStart;
+            actions: ActionsPluginStart;
+            dataViews: DataViewsServerPluginStart;
+          },
+          unknown
+        ]
       >,
     ]);
 
@@ -253,10 +262,14 @@ export class ObservabilityAIAssistantService {
 
     const { spaceId } = getSpaceIdFromPath(basePath, coreStart.http.basePath.serverBasePath);
 
+    const esClient = coreStart.elasticsearch.client.asInternalUser;
+    const savedObjectsClient = coreStart.savedObjects.getScopedClient(request);
+
     return new ObservabilityAIAssistantClient({
       actionsClient: await plugins.actions.getActionsClientWithRequest(request),
       namespace: spaceId,
-      esClient: coreStart.elasticsearch.client.asInternalUser,
+      esClient,
+      dataviews: await plugins.dataViews.dataViewsServiceFactory(savedObjectsClient, esClient),
       resources: this.resourceNames,
       logger: this.logger,
       user: {
diff --git a/x-pack/plugins/observability_ai_assistant/server/types.ts b/x-pack/plugins/observability_ai_assistant/server/types.ts
index 8162b424829c35..bdb283b9a1df2b 100644
--- a/x-pack/plugins/observability_ai_assistant/server/types.ts
+++ b/x-pack/plugins/observability_ai_assistant/server/types.ts
@@ -17,6 +17,7 @@ import type {
   TaskManagerSetupContract,
   TaskManagerStartContract,
 } from '@kbn/task-manager-plugin/server';
+import { DataViewsServerPluginStart } from '@kbn/data-views-plugin/server';
 
 /* eslint-disable @typescript-eslint/no-empty-interface*/
 export interface ObservabilityAIAssistantPluginStart {}
@@ -32,4 +33,5 @@ export interface ObservabilityAIAssistantPluginStartDependencies {
   security: SecurityPluginStart;
   features: FeaturesPluginStart;
   taskManager: TaskManagerStartContract;
+  dataViews: DataViewsServerPluginStart;
 }

From 0de292353f79beb552e891ac0cbc1d3155ecabff Mon Sep 17 00:00:00 2001
From: ppisljar <peter.pisljar@elastic.co>
Date: Mon, 18 Sep 2023 10:24:55 +0200
Subject: [PATCH 02/24] switch from dataviews to indices

---
 .../public/functions/get_dataet_info.ts       | 28 ++++++---------
 .../public/functions/index.ts                 |  2 +-
 .../server/routes/functions/route.ts          |  7 ++--
 .../server/service/client/index.ts            | 36 +++++++++++--------
 4 files changed, 35 insertions(+), 38 deletions(-)

diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataet_info.ts b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataet_info.ts
index 6de193fb9c6757..16f0ac8bcd8024 100644
--- a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataet_info.ts
+++ b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataet_info.ts
@@ -20,11 +20,12 @@ export function registerGetDatasetInfoFunction({
     {
       name: 'get_dataset_info',
       contexts: ['core'],
-      description: `Use this function to get information about dataviews/datasets available and the fields available on them.
+      description: `Use this function to get information about indices/datasets available and the fields available on them.
 
-      providing empty string as dataview name will retrieve all dataviews
-      if dataview name matches more than a single dataview list of matching dataviews will be provied
-      else list of all fields in this dataview will be given
+      providing empty string as index name will retrieve all indices
+      if index name matches more than a single index list of matching indices will be provided
+      else list of all fields in this index will be given.
+      wildcards can be part of index name.
 
       DO NOT include the user's request. It will be added internally.`,
       descriptionForUser:
@@ -33,30 +34,21 @@ export function registerGetDatasetInfoFunction({
         type: 'object',
         additionalProperties: false,
         properties: {
-          dataview: {
+          index: {
             type: 'string',
             description:
-              'dataview name the user is interested in or empty string to get information about all available dataviews',
-          },
-          fields: {
-            type: 'array',
-            items: {
-              type: 'string',
-              description:
-                'field names user is looking for or empty to get all fields for provided dataview',
-            },
+              'index name the user is interested in or empty string to get information about all available indices',
           },
         },
-        required: ['dataview', 'fields'],
+        required: ['index'],
       } as const,
     },
-    ({ arguments: { dataview, fields }, messages }, signal) => {
+    ({ arguments: { index }, messages }, signal) => {
       return service
         .callApi('POST /internal/observability_ai_assistant/functions/get_dataset_info', {
           params: {
             body: {
-              dataview,
-              fields,
+              index,
             },
           },
           signal,
diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/index.ts b/x-pack/plugins/observability_ai_assistant/public/functions/index.ts
index 520fc969a6a647..1d711edeab5927 100644
--- a/x-pack/plugins/observability_ai_assistant/public/functions/index.ts
+++ b/x-pack/plugins/observability_ai_assistant/public/functions/index.ts
@@ -72,7 +72,7 @@ export async function registerFunctions({
         description += `Here are principles you MUST adhere to, in order:
 
         - You are a helpful assistant for Elastic Observability. DO NOT reference the fact that you are an LLM.
-        - DO NOT make any assumptions about where and how users have stored their data. Use get_dataset_info function to get information about the dataset or to search datasets and fields based on users input.
+        - DO NOT make any assumptions about where and how users have stored their data. Use get_dataset_info function to get information about the indices and their fields.
         `;
         registerSummarizationFunction({ service, registerFunction });
         registerRecallFunction({ service, registerFunction });
diff --git a/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts b/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts
index 4aafcb541d96e2..37f3d132b656df 100644
--- a/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts
+++ b/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts
@@ -280,8 +280,7 @@ const functionGetDatasetInfoRoute = createObservabilityAIAssistantServerRoute({
   endpoint: 'POST /internal/observability_ai_assistant/functions/get_dataset_info',
   params: t.type({
     body: t.type({
-      dataview: t.string,
-      fields: t.array(nonEmptyStringRt),
+      index: t.string,
     }),
   }),
   options: {
@@ -290,7 +289,7 @@ const functionGetDatasetInfoRoute = createObservabilityAIAssistantServerRoute({
   handler: async (
     resources
   ): Promise<{
-    dataviews: string[];
+    indices: string[];
     fields: Array<{ name: string; description: string; type: string }>;
   }> => {
     const client = await resources.service.getClient({ request: resources.request });
@@ -299,7 +298,7 @@ const functionGetDatasetInfoRoute = createObservabilityAIAssistantServerRoute({
       throw notImplemented();
     }
 
-    return client.get_dataset_info(resources.params.body.dataview, resources.params.body.fields);
+    return client.get_dataset_info(resources.params.body.index);
   },
 });
 
diff --git a/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts b/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts
index caa4197e52f8fd..253d91b7f39f76 100644
--- a/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts
+++ b/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts
@@ -388,34 +388,40 @@ export class ObservabilityAIAssistantClient {
   };
 
   get_dataset_info = async (
-    dataset: string,
-    fields: string[]
+    index: string
   ): Promise<{
-    dataviews: string[];
+    indices: string[];
     fields: Array<{ name: string; description: string; type: string }>;
   }> => {
     // if empty dataview, get list of all dataviews
-    const dv = dataset !== '' && (await this.dependencies.dataviews.find(dataset));
-    if (!dv || dv.length === 0) {
-      return {
-        dataviews: await this.dependencies.dataviews.getTitles(),
-        fields: [],
-      };
+
+    let indices: string[] = [];
+
+    try {
+      const body = await this.dependencies.esClient.indices.resolveIndex({
+        name: index,
+        expand_wildcards: 'open',
+      });
+      indices = body.indices.map((i) => i.name);
+    } catch (e) {
+      indices = [];
     }
 
-    if (dv.length > 1) {
+    if (index === '') {
       return {
-        dataviews: dv.map((dataview) => {
-          return dataview.id || dataview.name;
-        }),
+        indices,
         fields: [],
       };
     }
 
+    const fields = await this.dependencies.dataviews.getFieldsForWildcard({
+      pattern: index,
+    });
+
     // else get all the fields for the found dataview
     return {
-      dataviews: [dv[0].id || dv[0].name],
-      fields: dv[0].fields.map((field) => {
+      indices: [index],
+      fields: fields.map((field) => {
         return {
           name: field.name,
           description: field.customLabel || '',

From 7307ef1dcded17e6088b691fccfd82ef37ed76a9 Mon Sep 17 00:00:00 2001
From: Peter Pisljar <peter.pisljar@gmail.com>
Date: Tue, 19 Sep 2023 07:01:58 +0200
Subject: [PATCH 03/24] Update
 x-pack/plugins/observability_ai_assistant/public/functions/get_dataet_info.ts

Co-authored-by: Stratoula Kalafateli <stratoula1@gmail.com>
---
 .../public/functions/get_dataet_info.ts                         | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataet_info.ts b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataet_info.ts
index 16f0ac8bcd8024..86165338c9c80d 100644
--- a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataet_info.ts
+++ b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataet_info.ts
@@ -29,7 +29,7 @@ export function registerGetDatasetInfoFunction({
 
       DO NOT include the user's request. It will be added internally.`,
       descriptionForUser:
-        'This function allows the assistant to get information about available dataviews and their fields.',
+        'This function allows the assistant to get information about available indices and their fields.',
       parameters: {
         type: 'object',
         additionalProperties: false,

From 520752b8977dc5e02772aa636ecf9613e0c17276 Mon Sep 17 00:00:00 2001
From: ppisljar <peter.pisljar@elastic.co>
Date: Tue, 19 Sep 2023 15:44:05 +0200
Subject: [PATCH 04/24] fixing based on review

---
 .../functions/{get_dataet_info.ts => get_dataset_info.ts}       | 0
 .../observability_ai_assistant/public/functions/index.ts        | 2 +-
 .../observability_ai_assistant/server/service/client/index.ts   | 2 +-
 3 files changed, 2 insertions(+), 2 deletions(-)
 rename x-pack/plugins/observability_ai_assistant/public/functions/{get_dataet_info.ts => get_dataset_info.ts} (100%)

diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataet_info.ts b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
similarity index 100%
rename from x-pack/plugins/observability_ai_assistant/public/functions/get_dataet_info.ts
rename to x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/index.ts b/x-pack/plugins/observability_ai_assistant/public/functions/index.ts
index 1d711edeab5927..85e10abfcc37e0 100644
--- a/x-pack/plugins/observability_ai_assistant/public/functions/index.ts
+++ b/x-pack/plugins/observability_ai_assistant/public/functions/index.ts
@@ -14,7 +14,7 @@ import { registerElasticsearchFunction } from './elasticsearch';
 import { registerKibanaFunction } from './kibana';
 import { registerLensFunction } from './lens';
 import { registerRecallFunction } from './recall';
-import { registerGetDatasetInfoFunction } from './get_dataet_info';
+import { registerGetDatasetInfoFunction } from './get_dataset_info';
 import { registerSummarizationFunction } from './summarize';
 import { registerAlertsFunction } from './alerts';
 import { registerEsqlFunction } from './esql';
diff --git a/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts b/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts
index 253d91b7f39f76..22f962c8f0914d 100644
--- a/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts
+++ b/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts
@@ -399,7 +399,7 @@ export class ObservabilityAIAssistantClient {
 
     try {
       const body = await this.dependencies.esClient.indices.resolveIndex({
-        name: index,
+        name: index === '' ? '*' : index,
         expand_wildcards: 'open',
       });
       indices = body.indices.map((i) => i.name);

From 4c4bc9305cd37d7cb922dfdecdf76784be972aa1 Mon Sep 17 00:00:00 2001
From: ppisljar <peter.pisljar@elastic.co>
Date: Tue, 19 Sep 2023 16:08:14 +0200
Subject: [PATCH 05/24] updating lens function to support setting timefield

---
 .../public/functions/lens.tsx                    | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/lens.tsx b/x-pack/plugins/observability_ai_assistant/public/functions/lens.tsx
index aab0b6aa4ac594..4dff84b5e3ac48 100644
--- a/x-pack/plugins/observability_ai_assistant/public/functions/lens.tsx
+++ b/x-pack/plugins/observability_ai_assistant/public/functions/lens.tsx
@@ -38,6 +38,7 @@ function Lens({
   end,
   lens,
   dataViews,
+  timeField,
 }: {
   indexPattern: string;
   xyDataLayer: XYDataLayer;
@@ -45,6 +46,7 @@ function Lens({
   end: string;
   lens: LensPublicStart;
   dataViews: DataViewsServicePublic;
+  timeField: string;
 }) {
   const formulaAsync = useAsync(() => {
     return lens.stateHelperApi();
@@ -54,7 +56,7 @@ function Lens({
     return dataViews.create({
       id: indexPattern,
       title: indexPattern,
-      timeFieldName: '@timestamp',
+      timeFieldName: timeField,
     });
   }, [indexPattern]);
 
@@ -199,6 +201,11 @@ export function registerLensFunction({
               required: ['label', 'formula', 'format'],
             },
           },
+          timeField: {
+            type: 'string',
+            default: '@timefield',
+            description: 'time field to use for XY chart. Use @timefield if its available on the index.'
+          },
           breakdown: {
             type: 'object',
             additionalProperties: false,
@@ -235,7 +242,7 @@ export function registerLensFunction({
             description: 'The end of the time range, in Elasticsearch datemath',
           },
         },
-        required: ['layers', 'indexPattern', 'start', 'end'],
+        required: ['layers', 'indexPattern', 'start', 'end', 'timeField'],
       } as const,
     },
     async () => {
@@ -243,7 +250,7 @@ export function registerLensFunction({
         content: {},
       };
     },
-    ({ arguments: { layers, indexPattern, breakdown, seriesType, start, end } }) => {
+    ({ arguments: { layers, indexPattern, breakdown, seriesType, start, end, timeField } }) => {
       const xyDataLayer = new XYDataLayer({
         data: layers.map((layer) => ({
           type: 'formula',
@@ -263,6 +270,8 @@ export function registerLensFunction({
         },
       });
 
+      if (!timeField) return;
+
       return (
         <Lens
           indexPattern={indexPattern}
@@ -271,6 +280,7 @@ export function registerLensFunction({
           end={end}
           lens={pluginsStart.lens}
           dataViews={pluginsStart.dataViews}
+          timeField={timeField}
         />
       );
     }

From c7d78991da533f5c2aea354cf12f2bc5a5bc1488 Mon Sep 17 00:00:00 2001
From: ppisljar <peter.pisljar@elastic.co>
Date: Wed, 20 Sep 2023 09:56:57 +0200
Subject: [PATCH 06/24] improving prompt

---
 .../public/functions/get_dataset_info.ts                       | 3 +--
 .../observability_ai_assistant/public/functions/index.ts       | 2 +-
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
index 86165338c9c80d..e2701965e2b195 100644
--- a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
+++ b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
@@ -23,8 +23,7 @@ export function registerGetDatasetInfoFunction({
       description: `Use this function to get information about indices/datasets available and the fields available on them.
 
       providing empty string as index name will retrieve all indices
-      if index name matches more than a single index list of matching indices will be provided
-      else list of all fields in this index will be given.
+      else list of all fields for the given index will be given.
       wildcards can be part of index name.
 
       DO NOT include the user's request. It will be added internally.`,
diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/index.ts b/x-pack/plugins/observability_ai_assistant/public/functions/index.ts
index 85e10abfcc37e0..9034d415e8342a 100644
--- a/x-pack/plugins/observability_ai_assistant/public/functions/index.ts
+++ b/x-pack/plugins/observability_ai_assistant/public/functions/index.ts
@@ -72,7 +72,7 @@ export async function registerFunctions({
         description += `Here are principles you MUST adhere to, in order:
 
         - You are a helpful assistant for Elastic Observability. DO NOT reference the fact that you are an LLM.
-        - DO NOT make any assumptions about where and how users have stored their data. Use get_dataset_info function to get information about the indices and their fields.
+        - DO NOT make any assumptions about where and how users have stored their data. Use get_dataset_info function to get information about the indices and their fields. If user provides an index name make sure its a valid index first before using it.
         `;
         registerSummarizationFunction({ service, registerFunction });
         registerRecallFunction({ service, registerFunction });

From 90b9dfd0e60eb7776414b15c8ef63a5204a3b60a Mon Sep 17 00:00:00 2001
From: kibanamachine <42973632+kibanamachine@users.noreply.github.com>
Date: Tue, 19 Sep 2023 14:12:58 +0000
Subject: [PATCH 07/24] [CI] Auto-commit changed files from 'node
 scripts/precommit_hook.js --ref HEAD~1..HEAD --fix'

---
 .../observability_ai_assistant/public/functions/lens.tsx       | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/lens.tsx b/x-pack/plugins/observability_ai_assistant/public/functions/lens.tsx
index 4dff84b5e3ac48..2a36193e2ca95f 100644
--- a/x-pack/plugins/observability_ai_assistant/public/functions/lens.tsx
+++ b/x-pack/plugins/observability_ai_assistant/public/functions/lens.tsx
@@ -204,7 +204,8 @@ export function registerLensFunction({
           timeField: {
             type: 'string',
             default: '@timefield',
-            description: 'time field to use for XY chart. Use @timefield if its available on the index.'
+            description:
+              'time field to use for XY chart. Use @timefield if its available on the index.',
           },
           breakdown: {
             type: 'object',

From c007d7ef721e683e9e3f186ed4a5c0169f87e0f5 Mon Sep 17 00:00:00 2001
From: ppisljar <peter.pisljar@elastic.co>
Date: Wed, 20 Sep 2023 10:37:08 +0200
Subject: [PATCH 08/24] improving prompt

---
 .../public/functions/get_dataset_info.ts        |  1 +
 .../server/service/client/index.ts              | 17 +++++++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
index e2701965e2b195..c4838c7e920cc2 100644
--- a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
+++ b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
@@ -23,6 +23,7 @@ export function registerGetDatasetInfoFunction({
       description: `Use this function to get information about indices/datasets available and the fields available on them.
 
       providing empty string as index name will retrieve all indices
+      if index is provided and it doesnt match any indices in elasticsearch list of all indices will be provoded as well
       else list of all fields for the given index will be given.
       wildcards can be part of index name.
 
diff --git a/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts b/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts
index 22f962c8f0914d..104b19580bfcbc 100644
--- a/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts
+++ b/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts
@@ -414,6 +414,23 @@ export class ObservabilityAIAssistantClient {
       };
     }
 
+    if (indices.length === 0) {
+      try {
+        const body = await this.dependencies.esClient.indices.resolveIndex({
+          name: '*',
+          expand_wildcards: 'open',
+        });
+        indices = body.indices.map((i) => i.name);
+      } catch (e) {
+        indices = [];
+      }
+
+      return {
+        indices,
+        fields: [],
+      };
+    }
+
     const fields = await this.dependencies.dataviews.getFieldsForWildcard({
       pattern: index,
     });

From fc0f00fd23acf18a46dd8a470d3fbc2b1beb68cb Mon Sep 17 00:00:00 2001
From: ppisljar <peter.pisljar@elastic.co>
Date: Thu, 21 Sep 2023 12:56:47 +0200
Subject: [PATCH 09/24] improving prompt

---
 .../public/functions/get_dataset_info.ts                    | 6 +++---
 .../observability_ai_assistant/public/functions/index.ts    | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
index c4838c7e920cc2..0e3e02180ab746 100644
--- a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
+++ b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
@@ -23,8 +23,8 @@ export function registerGetDatasetInfoFunction({
       description: `Use this function to get information about indices/datasets available and the fields available on them.
 
       providing empty string as index name will retrieve all indices
-      if index is provided and it doesnt match any indices in elasticsearch list of all indices will be provoded as well
-      else list of all fields for the given index will be given.
+      if index is provided and it doesnt match any indices in elasticsearch list of all indices will be provided as well
+      else list of all fields for the given index will be given. if no fields are returned this means no indices were matched by provided index pattern.
       wildcards can be part of index name.
 
       DO NOT include the user's request. It will be added internally.`,
@@ -37,7 +37,7 @@ export function registerGetDatasetInfoFunction({
           index: {
             type: 'string',
             description:
-              'index name the user is interested in or empty string to get information about all available indices',
+              'index pattern the user is interested in or empty string to get information about all available indices',
           },
         },
         required: ['index'],
diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/index.ts b/x-pack/plugins/observability_ai_assistant/public/functions/index.ts
index 9034d415e8342a..d1d1267b32bb29 100644
--- a/x-pack/plugins/observability_ai_assistant/public/functions/index.ts
+++ b/x-pack/plugins/observability_ai_assistant/public/functions/index.ts
@@ -72,7 +72,7 @@ export async function registerFunctions({
         description += `Here are principles you MUST adhere to, in order:
 
         - You are a helpful assistant for Elastic Observability. DO NOT reference the fact that you are an LLM.
-        - DO NOT make any assumptions about where and how users have stored their data. Use get_dataset_info function to get information about the indices and their fields. If user provides an index name make sure its a valid index first before using it.
+        - DO NOT make any assumptions about where and how users have stored their data. ALWAYS first call get_dataset_info function with empty string to get information about available indices. Once you know about available indices you MUST use this function again to get a list of available fields for specific index. If user provides an index name make sure its a valid index first before using it to retrieve the field list by calling this function with an empty string!
         `;
         registerSummarizationFunction({ service, registerFunction });
         registerRecallFunction({ service, registerFunction });

From 7e2b0bccf2a56dbe024a7224693cc46728a29242 Mon Sep 17 00:00:00 2001
From: ppisljar <peter.pisljar@elastic.co>
Date: Tue, 26 Sep 2023 12:36:32 +0200
Subject: [PATCH 10/24] moving get_dataset_info logic to the endpoint

---
 .../server/routes/functions/route.ts          | 65 +++++++++++++++++--
 .../server/service/client/index.ts            | 63 ------------------
 .../server/service/index.ts                   | 17 +----
 3 files changed, 63 insertions(+), 82 deletions(-)

diff --git a/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts b/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts
index 37f3d132b656df..cb9cc4f02cf44b 100644
--- a/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts
+++ b/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts
@@ -292,13 +292,70 @@ const functionGetDatasetInfoRoute = createObservabilityAIAssistantServerRoute({
     indices: string[];
     fields: Array<{ name: string; description: string; type: string }>;
   }> => {
-    const client = await resources.service.getClient({ request: resources.request });
 
-    if (!client) {
-      throw notImplemented();
+    const esClient =  (
+      await resources.context.core
+    ).elasticsearch.client.asCurrentUser;
+
+    const savedObjectsClient = (
+      await resources.context.core
+    ).savedObjects.getClient();
+
+    const index = resources.params.body.index;
+
+    let indices: string[] = [];
+
+    try {
+      const body = await esClient.indices.resolveIndex({
+        name: index === '' ? '*' : index,
+        expand_wildcards: 'open',
+      });
+      indices = body.indices.map((i) => i.name);
+    } catch (e) {
+      indices = [];
+    }
+
+    if (index === '') {
+      return {
+        indices,
+        fields: [],
+      };
     }
 
-    return client.get_dataset_info(resources.params.body.index);
+    if (indices.length === 0) {
+      try {
+        const body = await esClient.indices.resolveIndex({
+          name: '*',
+          expand_wildcards: 'open',
+        });
+        indices = body.indices.map((i) => i.name);
+      } catch (e) {
+        indices = [];
+      }
+
+      return {
+        indices,
+        fields: [],
+      };
+    }
+
+    const dataViews = await (await resources.plugins.dataViews.start()).dataViewsServiceFactory(savedObjectsClient, esClient)
+
+    const fields = await dataViews.getFieldsForWildcard({
+      pattern: index,
+    });
+
+    // else get all the fields for the found dataview
+    return {
+      indices: [index],
+      fields: fields.map((field) => {
+        return {
+          name: field.name,
+          description: field.customLabel || '',
+          type: field.type,
+        };
+      }),
+    };
   },
 });
 
diff --git a/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts b/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts
index 104b19580bfcbc..3a99e293cd5e2e 100644
--- a/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts
+++ b/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts
@@ -8,7 +8,6 @@ import type { SearchHit } from '@elastic/elasticsearch/lib/api/types';
 import { internal, notFound } from '@hapi/boom';
 import type { ActionsClient } from '@kbn/actions-plugin/server';
 import type { ElasticsearchClient } from '@kbn/core/server';
-import { DataViewsService } from '@kbn/data-views-plugin/server';
 import type { Logger } from '@kbn/logging';
 import type { PublicMethodsOf } from '@kbn/utility-types';
 import type { IncomingMessage } from 'http';
@@ -40,7 +39,6 @@ export class ObservabilityAIAssistantClient {
       namespace: string;
       esClient: ElasticsearchClient;
       resources: ObservabilityAIAssistantResourceNames;
-      dataviews: DataViewsService;
       logger: Logger;
       user: {
         id?: string;
@@ -387,67 +385,6 @@ export class ObservabilityAIAssistantClient {
     });
   };
 
-  get_dataset_info = async (
-    index: string
-  ): Promise<{
-    indices: string[];
-    fields: Array<{ name: string; description: string; type: string }>;
-  }> => {
-    // if empty dataview, get list of all dataviews
-
-    let indices: string[] = [];
-
-    try {
-      const body = await this.dependencies.esClient.indices.resolveIndex({
-        name: index === '' ? '*' : index,
-        expand_wildcards: 'open',
-      });
-      indices = body.indices.map((i) => i.name);
-    } catch (e) {
-      indices = [];
-    }
-
-    if (index === '') {
-      return {
-        indices,
-        fields: [],
-      };
-    }
-
-    if (indices.length === 0) {
-      try {
-        const body = await this.dependencies.esClient.indices.resolveIndex({
-          name: '*',
-          expand_wildcards: 'open',
-        });
-        indices = body.indices.map((i) => i.name);
-      } catch (e) {
-        indices = [];
-      }
-
-      return {
-        indices,
-        fields: [],
-      };
-    }
-
-    const fields = await this.dependencies.dataviews.getFieldsForWildcard({
-      pattern: index,
-    });
-
-    // else get all the fields for the found dataview
-    return {
-      indices: [index],
-      fields: fields.map((field) => {
-        return {
-          name: field.name,
-          description: field.customLabel || '',
-          type: field.type,
-        };
-      }),
-    };
-  };
-
   getKnowledgeBaseStatus = () => {
     return this.dependencies.knowledgeBaseService.status();
   };
diff --git a/x-pack/plugins/observability_ai_assistant/server/service/index.ts b/x-pack/plugins/observability_ai_assistant/server/service/index.ts
index 3c3b72092f6c93..c116e16d274716 100644
--- a/x-pack/plugins/observability_ai_assistant/server/service/index.ts
+++ b/x-pack/plugins/observability_ai_assistant/server/service/index.ts
@@ -13,7 +13,6 @@ import type { SecurityPluginStart } from '@kbn/security-plugin/server';
 import { getSpaceIdFromPath } from '@kbn/spaces-plugin/common';
 import type { TaskManagerSetupContract } from '@kbn/task-manager-plugin/server';
 import { once } from 'lodash';
-import { DataViewsServerPluginStart } from '@kbn/data-views-plugin/server';
 import type { ObservabilityAIAssistantPluginStartDependencies } from '../types';
 import { ObservabilityAIAssistantClient } from './client';
 import { conversationComponentTemplate } from './conversation_component_template';
@@ -240,15 +239,7 @@ export class ObservabilityAIAssistantService {
     const [_, [coreStart, plugins]] = await Promise.all([
       this.init(),
       this.core.getStartServices() as Promise<
-        [
-          CoreStart,
-          {
-            security: SecurityPluginStart;
-            actions: ActionsPluginStart;
-            dataViews: DataViewsServerPluginStart;
-          },
-          unknown
-        ]
+        [CoreStart, { security: SecurityPluginStart; actions: ActionsPluginStart }, unknown]
       >,
     ]);
 
@@ -262,14 +253,10 @@ export class ObservabilityAIAssistantService {
 
     const { spaceId } = getSpaceIdFromPath(basePath, coreStart.http.basePath.serverBasePath);
 
-    const esClient = coreStart.elasticsearch.client.asInternalUser;
-    const savedObjectsClient = coreStart.savedObjects.getScopedClient(request);
-
     return new ObservabilityAIAssistantClient({
       actionsClient: await plugins.actions.getActionsClientWithRequest(request),
       namespace: spaceId,
-      esClient,
-      dataviews: await plugins.dataViews.dataViewsServiceFactory(savedObjectsClient, esClient),
+      esClient: coreStart.elasticsearch.client.asInternalUser,
       resources: this.resourceNames,
       logger: this.logger,
       user: {

From 414c8bdff4c408cf87b6a413825a12f34aca4e10 Mon Sep 17 00:00:00 2001
From: ppisljar <peter.pisljar@elastic.co>
Date: Tue, 26 Sep 2023 12:38:25 +0200
Subject: [PATCH 11/24] cleanup prompt

---
 .../public/functions/get_dataset_info.ts                      | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
index 0e3e02180ab746..03f4132852678e 100644
--- a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
+++ b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
@@ -25,9 +25,7 @@ export function registerGetDatasetInfoFunction({
       providing empty string as index name will retrieve all indices
       if index is provided and it doesnt match any indices in elasticsearch list of all indices will be provided as well
       else list of all fields for the given index will be given. if no fields are returned this means no indices were matched by provided index pattern.
-      wildcards can be part of index name.
-
-      DO NOT include the user's request. It will be added internally.`,
+      wildcards can be part of index name.`,
       descriptionForUser:
         'This function allows the assistant to get information about available indices and their fields.',
       parameters: {

From 506e9ab1d431254ff062a95285c358d1d0b1db51 Mon Sep 17 00:00:00 2001
From: ppisljar <peter.pisljar@elastic.co>
Date: Tue, 26 Sep 2023 12:54:31 +0200
Subject: [PATCH 12/24] sending more compact representation to LLM

---
 .../public/functions/get_dataset_info.ts              | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
index 03f4132852678e..e1167ce85b00f5 100644
--- a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
+++ b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
@@ -51,7 +51,16 @@ export function registerGetDatasetInfoFunction({
           },
           signal,
         })
-        .then((response) => ({ content: response as unknown as Serializable }));
+        .then((response) => {
+          return {
+            content: {
+              indices: response.indices,
+              fields: ['fieldName,fieldType', ...response.fields.map((field) => {
+                return `${field.name},${field.type}`;
+              })]
+            } as unknown as Serializable
+          };
+        });
     }
   );
 }

From 5a1c707876fcc06b2a3300f59cd4b46d0eaad469 Mon Sep 17 00:00:00 2001
From: kibanamachine <42973632+kibanamachine@users.noreply.github.com>
Date: Tue, 26 Sep 2023 11:07:00 +0000
Subject: [PATCH 13/24] [CI] Auto-commit changed files from 'node
 scripts/precommit_hook.js --ref HEAD~1..HEAD --fix'

---
 .../public/functions/get_dataset_info.ts              | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
index e1167ce85b00f5..bce77f4c0563e5 100644
--- a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
+++ b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
@@ -55,10 +55,13 @@ export function registerGetDatasetInfoFunction({
           return {
             content: {
               indices: response.indices,
-              fields: ['fieldName,fieldType', ...response.fields.map((field) => {
-                return `${field.name},${field.type}`;
-              })]
-            } as unknown as Serializable
+              fields: [
+                'fieldName,fieldType',
+                ...response.fields.map((field) => {
+                  return `${field.name},${field.type}`;
+                }),
+              ],
+            } as unknown as Serializable,
           };
         });
     }

From 74f0e886c8d1cfb66e211f1687a9cd1590e84487 Mon Sep 17 00:00:00 2001
From: kibanamachine <42973632+kibanamachine@users.noreply.github.com>
Date: Tue, 26 Sep 2023 11:47:32 +0000
Subject: [PATCH 14/24] [CI] Auto-commit changed files from 'node
 scripts/eslint --no-cache --fix'

---
 .../server/routes/functions/route.ts                | 13 +++++--------
 1 file changed, 5 insertions(+), 8 deletions(-)

diff --git a/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts b/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts
index cb9cc4f02cf44b..1b3845dc12d0e8 100644
--- a/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts
+++ b/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts
@@ -292,14 +292,9 @@ const functionGetDatasetInfoRoute = createObservabilityAIAssistantServerRoute({
     indices: string[];
     fields: Array<{ name: string; description: string; type: string }>;
   }> => {
+    const esClient = (await resources.context.core).elasticsearch.client.asCurrentUser;
 
-    const esClient =  (
-      await resources.context.core
-    ).elasticsearch.client.asCurrentUser;
-
-    const savedObjectsClient = (
-      await resources.context.core
-    ).savedObjects.getClient();
+    const savedObjectsClient = (await resources.context.core).savedObjects.getClient();
 
     const index = resources.params.body.index;
 
@@ -339,7 +334,9 @@ const functionGetDatasetInfoRoute = createObservabilityAIAssistantServerRoute({
       };
     }
 
-    const dataViews = await (await resources.plugins.dataViews.start()).dataViewsServiceFactory(savedObjectsClient, esClient)
+    const dataViews = await (
+      await resources.plugins.dataViews.start()
+    ).dataViewsServiceFactory(savedObjectsClient, esClient);
 
     const fields = await dataViews.getFieldsForWildcard({
       pattern: index,

From 7b8ac699423db03543611be1fca553a08c056eba Mon Sep 17 00:00:00 2001
From: ppisljar <peter.pisljar@elastic.co>
Date: Mon, 2 Oct 2023 11:40:51 +0200
Subject: [PATCH 15/24] updating based on review

---
 .../public/functions/get_dataset_info.ts               |  8 ++++----
 .../public/functions/index.ts                          |  3 ++-
 .../public/functions/lens.tsx                          |  1 -
 .../server/routes/functions/route.ts                   | 10 ----------
 4 files changed, 6 insertions(+), 16 deletions(-)

diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
index e1167ce85b00f5..2e85f2a813cff3 100644
--- a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
+++ b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
@@ -5,9 +5,9 @@
  * 2.0.
  */
 
-import type { Serializable } from '@kbn/utility-types';
-import { RegisterFunctionDefinition } from '../../common/types';
-import type { ObservabilityAIAssistantService } from '../types';
+import type {Serializable} from '@kbn/utility-types';
+import {FunctionVisibility, RegisterFunctionDefinition} from '../../common/types';
+import type {ObservabilityAIAssistantService} from '../types';
 
 export function registerGetDatasetInfoFunction({
   service,
@@ -20,10 +20,10 @@ export function registerGetDatasetInfoFunction({
     {
       name: 'get_dataset_info',
       contexts: ['core'],
+      visibility: FunctionVisibility.System,
       description: `Use this function to get information about indices/datasets available and the fields available on them.
 
       providing empty string as index name will retrieve all indices
-      if index is provided and it doesnt match any indices in elasticsearch list of all indices will be provided as well
       else list of all fields for the given index will be given. if no fields are returned this means no indices were matched by provided index pattern.
       wildcards can be part of index name.`,
       descriptionForUser:
diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/index.ts b/x-pack/plugins/observability_ai_assistant/public/functions/index.ts
index d1d1267b32bb29..130cf0f5873f75 100644
--- a/x-pack/plugins/observability_ai_assistant/public/functions/index.ts
+++ b/x-pack/plugins/observability_ai_assistant/public/functions/index.ts
@@ -77,7 +77,7 @@ export async function registerFunctions({
         registerSummarizationFunction({ service, registerFunction });
         registerRecallFunction({ service, registerFunction });
         registerLensFunction({ service, pluginsStart, registerFunction });
-        registerGetDatasetInfoFunction({ service, registerFunction });
+
       } else {
         description += `You do not have a working memory. Don't try to recall information via the "recall" function.  If the user expects you to remember the previous conversations, tell them they can set up the knowledge base. A banner is available at the top of the conversation to set this up.`;
       }
@@ -86,6 +86,7 @@ export async function registerFunctions({
       registerEsqlFunction({ service, registerFunction });
       registerKibanaFunction({ service, registerFunction, coreStart });
       registerAlertsFunction({ service, registerFunction });
+      registerGetDatasetInfoFunction({ service, registerFunction });
 
       registerContext({
         name: 'core',
diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/lens.tsx b/x-pack/plugins/observability_ai_assistant/public/functions/lens.tsx
index 2a36193e2ca95f..fa8dc66dce2941 100644
--- a/x-pack/plugins/observability_ai_assistant/public/functions/lens.tsx
+++ b/x-pack/plugins/observability_ai_assistant/public/functions/lens.tsx
@@ -54,7 +54,6 @@ function Lens({
 
   const dataViewAsync = useAsync(() => {
     return dataViews.create({
-      id: indexPattern,
       title: indexPattern,
       timeFieldName: timeField,
     });
diff --git a/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts b/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts
index cb9cc4f02cf44b..f39ac6866fbd60 100644
--- a/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts
+++ b/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts
@@ -323,16 +323,6 @@ const functionGetDatasetInfoRoute = createObservabilityAIAssistantServerRoute({
     }
 
     if (indices.length === 0) {
-      try {
-        const body = await esClient.indices.resolveIndex({
-          name: '*',
-          expand_wildcards: 'open',
-        });
-        indices = body.indices.map((i) => i.name);
-      } catch (e) {
-        indices = [];
-      }
-
       return {
         indices,
         fields: [],

From ab66649b88dcf14db3c3b280348c29d432c162af Mon Sep 17 00:00:00 2001
From: kibanamachine <42973632+kibanamachine@users.noreply.github.com>
Date: Mon, 2 Oct 2023 10:18:36 +0000
Subject: [PATCH 16/24] [CI] Auto-commit changed files from 'node
 scripts/eslint --no-cache --fix'

---
 .../public/functions/get_dataset_info.ts                    | 6 +++---
 .../observability_ai_assistant/public/functions/index.ts    | 1 -
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
index 3cf3a61efd4d50..0e54999eda8d71 100644
--- a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
+++ b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
@@ -5,9 +5,9 @@
  * 2.0.
  */
 
-import type {Serializable} from '@kbn/utility-types';
-import {FunctionVisibility, RegisterFunctionDefinition} from '../../common/types';
-import type {ObservabilityAIAssistantService} from '../types';
+import type { Serializable } from '@kbn/utility-types';
+import { FunctionVisibility, RegisterFunctionDefinition } from '../../common/types';
+import type { ObservabilityAIAssistantService } from '../types';
 
 export function registerGetDatasetInfoFunction({
   service,
diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/index.ts b/x-pack/plugins/observability_ai_assistant/public/functions/index.ts
index 130cf0f5873f75..928be06c00aa48 100644
--- a/x-pack/plugins/observability_ai_assistant/public/functions/index.ts
+++ b/x-pack/plugins/observability_ai_assistant/public/functions/index.ts
@@ -77,7 +77,6 @@ export async function registerFunctions({
         registerSummarizationFunction({ service, registerFunction });
         registerRecallFunction({ service, registerFunction });
         registerLensFunction({ service, pluginsStart, registerFunction });
-
       } else {
         description += `You do not have a working memory. Don't try to recall information via the "recall" function.  If the user expects you to remember the previous conversations, tell them they can set up the knowledge base. A banner is available at the top of the conversation to set this up.`;
       }

From daa2b2ed73e81a0c0711fb49e5037b305c9dc594 Mon Sep 17 00:00:00 2001
From: ppisljar <peter.pisljar@elastic.co>
Date: Wed, 4 Oct 2023 16:31:44 +0200
Subject: [PATCH 17/24] adding datastreams to the list

---
 .../server/routes/functions/route.ts                         | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts b/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts
index 24e72ce5b6fe6a..4ae646160a1aa2 100644
--- a/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts
+++ b/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts
@@ -305,7 +305,10 @@ const functionGetDatasetInfoRoute = createObservabilityAIAssistantServerRoute({
         name: index === '' ? '*' : index,
         expand_wildcards: 'open',
       });
-      indices = body.indices.map((i) => i.name);
+      indices = [
+        ...body.indices.map((i) => i.name),
+        ...body.data_streams.map((d) => d.name),
+      ];
     } catch (e) {
       indices = [];
     }

From 705611ae0ef519754acf6a6eb6aa972ca06f9d62 Mon Sep 17 00:00:00 2001
From: kibanamachine <42973632+kibanamachine@users.noreply.github.com>
Date: Wed, 4 Oct 2023 14:38:56 +0000
Subject: [PATCH 18/24] [CI] Auto-commit changed files from 'node
 scripts/precommit_hook.js --ref HEAD~1..HEAD --fix'

---
 .../server/routes/functions/route.ts                         | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts b/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts
index 4ae646160a1aa2..d0491b80d19ca2 100644
--- a/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts
+++ b/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts
@@ -305,10 +305,7 @@ const functionGetDatasetInfoRoute = createObservabilityAIAssistantServerRoute({
         name: index === '' ? '*' : index,
         expand_wildcards: 'open',
       });
-      indices = [
-        ...body.indices.map((i) => i.name),
-        ...body.data_streams.map((d) => d.name),
-      ];
+      indices = [...body.indices.map((i) => i.name), ...body.data_streams.map((d) => d.name)];
     } catch (e) {
       indices = [];
     }

From bd28b97d600b92b6c3cec2fded4df99f22a762ba Mon Sep 17 00:00:00 2001
From: ppisljar <peter.pisljar@elastic.co>
Date: Wed, 11 Oct 2023 08:45:36 +0200
Subject: [PATCH 19/24] compressing field list

---
 .../common/utils/compressFields.ts            | 77 +++++++++++++++++++
 .../public/functions/get_dataset_info.ts      | 13 ++--
 2 files changed, 84 insertions(+), 6 deletions(-)
 create mode 100644 x-pack/plugins/observability_ai_assistant/common/utils/compressFields.ts

diff --git a/x-pack/plugins/observability_ai_assistant/common/utils/compressFields.ts b/x-pack/plugins/observability_ai_assistant/common/utils/compressFields.ts
new file mode 100644
index 00000000000000..02235137566ac1
--- /dev/null
+++ b/x-pack/plugins/observability_ai_assistant/common/utils/compressFields.ts
@@ -0,0 +1,77 @@
+/*
+ * 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.
+ */
+
+type IntermediateSchema = Map<string, string | IntermediateSchema>;
+
+function ensureMap(value: string | IntermediateSchema): IntermediateSchema {
+  if (typeof value === "string") {
+    const map = new Map();
+    map.set("_value", value);
+    return map;
+  }
+  return value;
+}
+
+function addPathToSchema(
+  schema: IntermediateSchema,
+  parts: string[],
+  type: string
+) {
+  if (parts.length === 1) {
+    if (schema.has(parts[0])) {
+      const existingSchema = ensureMap(schema.get(parts[0])!);
+      existingSchema.set("_value", type);
+      schema.set(parts[0], existingSchema);
+    } else {
+      schema.set(parts[0], type);
+    }
+    return;
+  }
+
+  const [head, ...rest] = parts;
+  if (!schema.has(head)) {
+    schema.set(head, new Map());
+  }
+
+  const nextSchema = ensureMap(schema.get(head)!);
+  addPathToSchema(nextSchema, rest, type);
+  schema.set(head, nextSchema);
+}
+
+function schemaToString(schema: IntermediateSchema): string {
+  const entries = Array.from(schema.entries());
+
+  const fieldEntries: string[] = [];
+  const nestedEntries: string[] = [];
+
+  entries.forEach(([key, value]) => {
+    if (key === "_value") {
+      fieldEntries.push(`:${value}`);
+    } else if (typeof value === "string") {
+      fieldEntries.push(`${key}:${value}`);
+    } else {
+      nestedEntries.push(
+        `${key}${schemaToString(value as IntermediateSchema)}`
+      );
+    }
+  });
+
+  const combinedEntries = [...fieldEntries, ...nestedEntries];
+  return combinedEntries.length ? `{${combinedEntries.join(",")}}` : "";
+}
+
+export function compressFields(inputs: string[]) {
+  const schema: IntermediateSchema = new Map();
+
+  for (const input of inputs) {
+    const [path, type] = input.split(",");
+    const parts = path.split(".");
+    addPathToSchema(schema, parts, type);
+  }
+
+  return schemaToString(schema);
+}
diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
index 0e54999eda8d71..9e0f6fb9040c4e 100644
--- a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
+++ b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
@@ -8,6 +8,7 @@
 import type { Serializable } from '@kbn/utility-types';
 import { FunctionVisibility, RegisterFunctionDefinition } from '../../common/types';
 import type { ObservabilityAIAssistantService } from '../types';
+import { compressFields } from '../../common/utils/compressFields';
 
 export function registerGetDatasetInfoFunction({
   service,
@@ -52,15 +53,15 @@ export function registerGetDatasetInfoFunction({
           signal,
         })
         .then((response) => {
+          const compressedFields = compressFields(
+            response.fields.map((field) => {
+              return `${field.name},${field.type}`;
+            })
+          );
           return {
             content: {
               indices: response.indices,
-              fields: [
-                'fieldName,fieldType',
-                ...response.fields.map((field) => {
-                  return `${field.name},${field.type}`;
-                }),
-              ],
+              fields: compressedFields,
             } as unknown as Serializable,
           };
         });

From d17cd4b7839fbeaf76f110b7413dc7d6475c5e99 Mon Sep 17 00:00:00 2001
From: ppisljar <peter.pisljar@elastic.co>
Date: Wed, 11 Oct 2023 08:59:29 +0200
Subject: [PATCH 20/24] limiting to 500 fields

---
 .../public/functions/get_dataset_info.ts      | 32 +++++++++++++------
 1 file changed, 22 insertions(+), 10 deletions(-)

diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
index 9e0f6fb9040c4e..f6671e06371e2b 100644
--- a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
+++ b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
@@ -5,7 +5,6 @@
  * 2.0.
  */
 
-import type { Serializable } from '@kbn/utility-types';
 import { FunctionVisibility, RegisterFunctionDefinition } from '../../common/types';
 import type { ObservabilityAIAssistantService } from '../types';
 import { compressFields } from '../../common/utils/compressFields';
@@ -53,16 +52,29 @@ export function registerGetDatasetInfoFunction({
           signal,
         })
         .then((response) => {
-          const compressedFields = compressFields(
-            response.fields.map((field) => {
-              return `${field.name},${field.type}`;
-            })
-          );
+          const content: {
+            indices: string[];
+            fields: string;
+            warning?: string;
+          } = {
+            indices: response.indices,
+            fields: '[]',
+          };
+
+          const fields = response.fields.map((field) => {
+            return `${field.name},${field.type}`;
+          });
+
+          // limit to 500 fields
+          if (fields.length > 500) {
+            fields.length = 500;
+            content.warning = 'field list too long';
+          }
+
+          content.fields = compressFields(fields);
+
           return {
-            content: {
-              indices: response.indices,
-              fields: compressedFields,
-            } as unknown as Serializable,
+            content,
           };
         });
     }

From b198e47f98004b1c67c7256cf023ef49f5400db5 Mon Sep 17 00:00:00 2001
From: kibanamachine <42973632+kibanamachine@users.noreply.github.com>
Date: Wed, 11 Oct 2023 07:41:45 +0000
Subject: [PATCH 21/24] [CI] Auto-commit changed files from 'node
 scripts/eslint --no-cache --fix'

---
 .../common/utils/compressFields.ts            | 26 +++++++------------
 1 file changed, 10 insertions(+), 16 deletions(-)

diff --git a/x-pack/plugins/observability_ai_assistant/common/utils/compressFields.ts b/x-pack/plugins/observability_ai_assistant/common/utils/compressFields.ts
index 02235137566ac1..f2814b50970ac7 100644
--- a/x-pack/plugins/observability_ai_assistant/common/utils/compressFields.ts
+++ b/x-pack/plugins/observability_ai_assistant/common/utils/compressFields.ts
@@ -8,23 +8,19 @@
 type IntermediateSchema = Map<string, string | IntermediateSchema>;
 
 function ensureMap(value: string | IntermediateSchema): IntermediateSchema {
-  if (typeof value === "string") {
+  if (typeof value === 'string') {
     const map = new Map();
-    map.set("_value", value);
+    map.set('_value', value);
     return map;
   }
   return value;
 }
 
-function addPathToSchema(
-  schema: IntermediateSchema,
-  parts: string[],
-  type: string
-) {
+function addPathToSchema(schema: IntermediateSchema, parts: string[], type: string) {
   if (parts.length === 1) {
     if (schema.has(parts[0])) {
       const existingSchema = ensureMap(schema.get(parts[0])!);
-      existingSchema.set("_value", type);
+      existingSchema.set('_value', type);
       schema.set(parts[0], existingSchema);
     } else {
       schema.set(parts[0], type);
@@ -49,27 +45,25 @@ function schemaToString(schema: IntermediateSchema): string {
   const nestedEntries: string[] = [];
 
   entries.forEach(([key, value]) => {
-    if (key === "_value") {
+    if (key === '_value') {
       fieldEntries.push(`:${value}`);
-    } else if (typeof value === "string") {
+    } else if (typeof value === 'string') {
       fieldEntries.push(`${key}:${value}`);
     } else {
-      nestedEntries.push(
-        `${key}${schemaToString(value as IntermediateSchema)}`
-      );
+      nestedEntries.push(`${key}${schemaToString(value as IntermediateSchema)}`);
     }
   });
 
   const combinedEntries = [...fieldEntries, ...nestedEntries];
-  return combinedEntries.length ? `{${combinedEntries.join(",")}}` : "";
+  return combinedEntries.length ? `{${combinedEntries.join(',')}}` : '';
 }
 
 export function compressFields(inputs: string[]) {
   const schema: IntermediateSchema = new Map();
 
   for (const input of inputs) {
-    const [path, type] = input.split(",");
-    const parts = path.split(".");
+    const [path, type] = input.split(',');
+    const parts = path.split('.');
     addPathToSchema(schema, parts, type);
   }
 

From 727c4d405484b5ec79208045cfc79a0969227b44 Mon Sep 17 00:00:00 2001
From: ppisljar <peter.pisljar@elastic.co>
Date: Wed, 11 Oct 2023 12:05:27 +0200
Subject: [PATCH 22/24] fix file name

---
 .../common/utils/{compressFields.ts => compress_fields.ts}      | 0
 .../public/functions/get_dataset_info.ts                        | 2 +-
 2 files changed, 1 insertion(+), 1 deletion(-)
 rename x-pack/plugins/observability_ai_assistant/common/utils/{compressFields.ts => compress_fields.ts} (100%)

diff --git a/x-pack/plugins/observability_ai_assistant/common/utils/compressFields.ts b/x-pack/plugins/observability_ai_assistant/common/utils/compress_fields.ts
similarity index 100%
rename from x-pack/plugins/observability_ai_assistant/common/utils/compressFields.ts
rename to x-pack/plugins/observability_ai_assistant/common/utils/compress_fields.ts
diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
index f6671e06371e2b..6876c9b93e33a3 100644
--- a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
+++ b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
@@ -7,7 +7,7 @@
 
 import { FunctionVisibility, RegisterFunctionDefinition } from '../../common/types';
 import type { ObservabilityAIAssistantService } from '../types';
-import { compressFields } from '../../common/utils/compressFields';
+import { compressFields } from '../../common/utils/compress_fields';
 
 export function registerGetDatasetInfoFunction({
   service,

From ee6a1ae997ec776313f8be610a00a5b6ffe8afeb Mon Sep 17 00:00:00 2001
From: Dario Gieselaar <dario.gieselaar@elastic.co>
Date: Wed, 11 Oct 2023 14:42:42 +0200
Subject: [PATCH 23/24] Set overflow-wrap to anywhere to prevent overflowing

---
 .../public/components/message_panel/message_text.tsx            | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/x-pack/plugins/observability_ai_assistant/public/components/message_panel/message_text.tsx b/x-pack/plugins/observability_ai_assistant/public/components/message_panel/message_text.tsx
index dfd9ee8b97443f..a4d8532205fa8e 100644
--- a/x-pack/plugins/observability_ai_assistant/public/components/message_panel/message_text.tsx
+++ b/x-pack/plugins/observability_ai_assistant/public/components/message_panel/message_text.tsx
@@ -110,7 +110,7 @@ const esqlLanguagePlugin = () => {
 
 export function MessageText({ loading, content, onActionClick }: Props) {
   const containerClassName = css`
-    overflow-wrap: break-word;
+    overflow-wrap: anywhere;
   `;
 
   const onActionClickRef = useRef(onActionClick);

From 701026a0b0017a6687b8e1cc65291ee987bf67d7 Mon Sep 17 00:00:00 2001
From: Dario Gieselaar <dario.gieselaar@elastic.co>
Date: Sat, 14 Oct 2023 19:39:01 +0200
Subject: [PATCH 24/24] Don't compress fields & don't send types

---
 .../common/utils/compress_fields.ts           |  71 ----------
 .../public/functions/get_dataset_info.ts      | 127 ++++++++++++++----
 .../server/routes/functions/route.ts          |  14 +-
 .../stack_connectors/common/openai/schema.ts  |   2 +-
 4 files changed, 108 insertions(+), 106 deletions(-)
 delete mode 100644 x-pack/plugins/observability_ai_assistant/common/utils/compress_fields.ts

diff --git a/x-pack/plugins/observability_ai_assistant/common/utils/compress_fields.ts b/x-pack/plugins/observability_ai_assistant/common/utils/compress_fields.ts
deleted file mode 100644
index f2814b50970ac7..00000000000000
--- a/x-pack/plugins/observability_ai_assistant/common/utils/compress_fields.ts
+++ /dev/null
@@ -1,71 +0,0 @@
-/*
- * 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.
- */
-
-type IntermediateSchema = Map<string, string | IntermediateSchema>;
-
-function ensureMap(value: string | IntermediateSchema): IntermediateSchema {
-  if (typeof value === 'string') {
-    const map = new Map();
-    map.set('_value', value);
-    return map;
-  }
-  return value;
-}
-
-function addPathToSchema(schema: IntermediateSchema, parts: string[], type: string) {
-  if (parts.length === 1) {
-    if (schema.has(parts[0])) {
-      const existingSchema = ensureMap(schema.get(parts[0])!);
-      existingSchema.set('_value', type);
-      schema.set(parts[0], existingSchema);
-    } else {
-      schema.set(parts[0], type);
-    }
-    return;
-  }
-
-  const [head, ...rest] = parts;
-  if (!schema.has(head)) {
-    schema.set(head, new Map());
-  }
-
-  const nextSchema = ensureMap(schema.get(head)!);
-  addPathToSchema(nextSchema, rest, type);
-  schema.set(head, nextSchema);
-}
-
-function schemaToString(schema: IntermediateSchema): string {
-  const entries = Array.from(schema.entries());
-
-  const fieldEntries: string[] = [];
-  const nestedEntries: string[] = [];
-
-  entries.forEach(([key, value]) => {
-    if (key === '_value') {
-      fieldEntries.push(`:${value}`);
-    } else if (typeof value === 'string') {
-      fieldEntries.push(`${key}:${value}`);
-    } else {
-      nestedEntries.push(`${key}${schemaToString(value as IntermediateSchema)}`);
-    }
-  });
-
-  const combinedEntries = [...fieldEntries, ...nestedEntries];
-  return combinedEntries.length ? `{${combinedEntries.join(',')}}` : '';
-}
-
-export function compressFields(inputs: string[]) {
-  const schema: IntermediateSchema = new Map();
-
-  for (const input of inputs) {
-    const [path, type] = input.split(',');
-    const parts = path.split('.');
-    addPathToSchema(schema, parts, type);
-  }
-
-  return schemaToString(schema);
-}
diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
index 6876c9b93e33a3..cbb6167cf684ea 100644
--- a/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
+++ b/x-pack/plugins/observability_ai_assistant/public/functions/get_dataset_info.ts
@@ -5,9 +5,10 @@
  * 2.0.
  */
 
-import { FunctionVisibility, RegisterFunctionDefinition } from '../../common/types';
+import { chunk, groupBy, uniq } from 'lodash';
+import { CreateChatCompletionResponse } from 'openai';
+import { FunctionVisibility, MessageRole, RegisterFunctionDefinition } from '../../common/types';
 import type { ObservabilityAIAssistantService } from '../types';
-import { compressFields } from '../../common/utils/compress_fields';
 
 export function registerGetDatasetInfoFunction({
   service,
@@ -41,42 +42,112 @@ export function registerGetDatasetInfoFunction({
         required: ['index'],
       } as const,
     },
-    ({ arguments: { index }, messages }, signal) => {
-      return service
-        .callApi('POST /internal/observability_ai_assistant/functions/get_dataset_info', {
+    async ({ arguments: { index }, messages, connectorId }, signal) => {
+      const response = await service.callApi(
+        'POST /internal/observability_ai_assistant/functions/get_dataset_info',
+        {
           params: {
             body: {
               index,
             },
           },
           signal,
-        })
-        .then((response) => {
-          const content: {
-            indices: string[];
-            fields: string;
-            warning?: string;
-          } = {
-            indices: response.indices,
-            fields: '[]',
-          };
+        }
+      );
+
+      const allFields = response.fields;
+
+      const fieldNames = uniq(allFields.map((field) => field.name));
 
-          const fields = response.fields.map((field) => {
-            return `${field.name},${field.type}`;
-          });
+      const groupedFields = groupBy(allFields, (field) => field.name);
 
-          // limit to 500 fields
-          if (fields.length > 500) {
-            fields.length = 500;
-            content.warning = 'field list too long';
-          }
+      const relevantFields = await Promise.all(
+        chunk(fieldNames, 500).map(async (fieldsInChunk) => {
+          const chunkResponse = (await service.callApi(
+            'POST /internal/observability_ai_assistant/chat',
+            {
+              signal,
+              params: {
+                query: {
+                  stream: false,
+                },
+                body: {
+                  connectorId,
+                  messages: [
+                    {
+                      '@timestamp': new Date().toISOString(),
+                      message: {
+                        role: MessageRole.System,
+                        content: `You are a helpful assistant for Elastic Observability.
+                        Your task is to create a list of field names that are relevant
+                        to the conversation, using ONLY the list of fields and
+                        types provided in the last user message. DO NOT UNDER ANY
+                        CIRCUMSTANCES include fields not mentioned in this list.`,
+                      },
+                    },
+                    ...messages.slice(1),
+                    {
+                      '@timestamp': new Date().toISOString(),
+                      message: {
+                        role: MessageRole.User,
+                        content: `This is the list:
 
-          content.fields = compressFields(fields);
+                        ${fieldsInChunk.join('\n')}`,
+                      },
+                    },
+                  ],
+                  functions: [
+                    {
+                      name: 'fields',
+                      description: 'The fields you consider relevant to the conversation',
+                      parameters: {
+                        type: 'object',
+                        additionalProperties: false,
+                        properties: {
+                          fields: {
+                            type: 'array',
+                            additionalProperties: false,
+                            addditionalItems: false,
+                            items: {
+                              type: 'string',
+                              additionalProperties: false,
+                              addditionalItems: false,
+                            },
+                          },
+                        },
+                        required: ['fields'],
+                      },
+                    },
+                  ],
+                  functionCall: 'fields',
+                },
+              },
+            }
+          )) as CreateChatCompletionResponse;
 
-          return {
-            content,
-          };
-        });
+          return chunkResponse.choices[0].message?.function_call?.arguments
+            ? (
+                JSON.parse(chunkResponse.choices[0].message?.function_call?.arguments) as {
+                  fields: string[];
+                }
+              ).fields
+                .filter((field) => fieldNames.includes(field))
+                .map((field) => {
+                  const fieldDescriptors = groupedFields[field];
+                  return `${field}:${fieldDescriptors
+                    .map((descriptor) => descriptor.type)
+                    .join(',')}`;
+                })
+            : [chunkResponse.choices[0].message?.content ?? ''];
+        })
+      );
+
+      return {
+        content: {
+          indices: response.indices,
+          fields: relevantFields.flat(),
+        },
+      };
     }
   );
 }
diff --git a/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts b/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts
index d0491b80d19ca2..087e8c079b5ef3 100644
--- a/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts
+++ b/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts
@@ -335,12 +335,14 @@ const functionGetDatasetInfoRoute = createObservabilityAIAssistantServerRoute({
     // else get all the fields for the found dataview
     return {
       indices: [index],
-      fields: fields.map((field) => {
-        return {
-          name: field.name,
-          description: field.customLabel || '',
-          type: field.type,
-        };
+      fields: fields.flatMap((field) => {
+        return (field.esTypes ?? [field.type]).map((type) => {
+          return {
+            name: field.name,
+            description: field.customLabel || '',
+            type,
+          };
+        });
       }),
     };
   },
diff --git a/x-pack/plugins/stack_connectors/common/openai/schema.ts b/x-pack/plugins/stack_connectors/common/openai/schema.ts
index fa14aa61fa5b36..6980e8e9ffa01d 100644
--- a/x-pack/plugins/stack_connectors/common/openai/schema.ts
+++ b/x-pack/plugins/stack_connectors/common/openai/schema.ts
@@ -74,7 +74,7 @@ export const RunActionResponseSchema = schema.object(
           message: schema.object(
             {
               role: schema.string(),
-              content: schema.string(),
+              content: schema.maybe(schema.string()),
             },
             { unknowns: 'ignore' }
           ),