[Obs AI Assistant] Gemini prompt improvements

x-pack/platform/packages/shared/kbn-ai-assistant/src/chat/chat_body.test.tsx

@@ -5,8 +5,7 @@
  * 2.0.
  */
 
-import { Message } from '@kbn/observability-ai-assistant-plugin/common';
-import { CONTEXT_FUNCTION_NAME } from '@kbn/observability-ai-assistant-plugin/server/functions/context/context';
+import { CONTEXT_FUNCTION_NAME, Message } from '@kbn/observability-ai-assistant-plugin/common';
 import { reverseToLastUserMessage } from './chat_body';
 
 describe('<ChatBody>', () => {

x-pack/platform/packages/shared/kbn-ai-assistant/src/utils/get_timeline_items_from_conversation.test.tsx

@@ -12,7 +12,7 @@ import { __IntlProvider as IntlProvider } from '@kbn/i18n-react';
 import { ChatState, Message, MessageRole } from '@kbn/observability-ai-assistant-plugin/public';
 import { createMockChatService } from './create_mock_chat_service';
 import { KibanaContextProvider } from '@kbn/triggers-actions-ui-plugin/public/common/lib/kibana';
-import { CONTEXT_FUNCTION_NAME } from '@kbn/observability-ai-assistant-plugin/server/functions/context/context';
+import { CONTEXT_FUNCTION_NAME } from '@kbn/observability-ai-assistant-plugin/common';
 
 const mockChatService = createMockChatService();
 

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

@@ -0,0 +1,44 @@
+/*
+ * 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.
+ */
+
+// Context tools
+export const CONTEXT_FUNCTION_NAME = 'context';
+export const GET_DATA_ON_SCREEN_FUNCTION_NAME = 'get_data_on_screen';
+
+// Alerts tools
+export const GET_ALERTS_DATASET_INFO_FUNCTION_NAME = 'get_alerts_dataset_info';
+export const ALERTS_FUNCTION_NAME = 'alerts';
+
+// Query tools
+export const QUERY_FUNCTION_NAME = 'query';
+export const EXECUTE_QUERY_FUNCTION_NAME = 'execute_query';
+export const VISUALIZE_QUERY_FUNCTION_NAME = 'visualize_query';
+
+// Dataset tools
+export const GET_DATASET_INFO_FUNCTION_NAME = 'get_dataset_info';
+export const SELECT_RELEVANT_FIELDS_NAME = 'select_relevant_fields';
+
+export const ELASTICSEARCH_FUNCTION_NAME = 'elasticsearch';
+
+export const EXECUTE_CONNECTOR_FUNCTION_NAME = 'execute_connector';
+
+export const KIBANA_FUNCTION_NAME = 'kibana';
+
+export const SUMMARIZE_FUNCTION_NAME = 'summarize';
+
+export const CHANGES_FUNCTION_NAME = 'changes';
+
+export const RETRIEVE_ELASTIC_DOC_FUNCTION_NAME = 'retrieve_elastic_doc';
+
+// APM tools
+export const GET_APM_DATASET_INFO_FUNCTION_NAME = 'get_apm_dataset_info';
+export const GET_APM_DOWNSTREAM_DEPENDENCIES_FUNCTION_NAME = 'get_apm_downstream_dependencies';
+export const GET_APM_SERVICES_LIST_FUNCTION_NAME = 'get_apm_services_list';
+
+// Deprecated tools
+export const GET_APM_TIMESERIES_FUNCTION_NAME = 'get_apm_timeseries';
+export const LENS_FUNCTION_NAME = 'lens';

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

@@ -60,3 +60,26 @@ export {
   EIS_PRECONFIGURED_INFERENCE_IDS,
   LEGACY_CUSTOM_INFERENCE_ID,
 } from './preconfigured_inference_ids';
+
+export {
+  GET_ALERTS_DATASET_INFO_FUNCTION_NAME,
+  ALERTS_FUNCTION_NAME,
+  QUERY_FUNCTION_NAME,
+  EXECUTE_QUERY_FUNCTION_NAME,
+  GET_DATA_ON_SCREEN_FUNCTION_NAME,
+  CONTEXT_FUNCTION_NAME,
+  ELASTICSEARCH_FUNCTION_NAME,
+  EXECUTE_CONNECTOR_FUNCTION_NAME,
+  GET_DATASET_INFO_FUNCTION_NAME,
+  SELECT_RELEVANT_FIELDS_NAME,
+  KIBANA_FUNCTION_NAME,
+  SUMMARIZE_FUNCTION_NAME,
+  VISUALIZE_QUERY_FUNCTION_NAME,
+  CHANGES_FUNCTION_NAME,
+  RETRIEVE_ELASTIC_DOC_FUNCTION_NAME,
+  GET_APM_DATASET_INFO_FUNCTION_NAME,
+  GET_APM_DOWNSTREAM_DEPENDENCIES_FUNCTION_NAME,
+  GET_APM_SERVICES_LIST_FUNCTION_NAME,
+  GET_APM_TIMESERIES_FUNCTION_NAME,
+  LENS_FUNCTION_NAME,
+} from './function_names';

x-pack/platform/plugins/shared/observability_ai_assistant/kibana.jsonc

@@ -17,6 +17,7 @@
       "taskManager",
       "dataViews",
       "inference",
+      "llmTasks",
       "productDocBase"
     ],
     "optionalPlugins": ["cloud", "serverless"],

x-pack/platform/plugins/shared/observability_ai_assistant/server/functions/context/context.test.ts

@@ -6,11 +6,9 @@
  */
 
 import { last } from 'lodash';
-import { Message, MessageRole } from '../../../common';
+import { CONTEXT_FUNCTION_NAME, Message, MessageRole } from '../../../common';
 import { removeContextToolRequest } from './context';
 
-const CONTEXT_FUNCTION_NAME = 'context';
-
 describe('removeContextToolRequest', () => {
   const baseMessages: Message[] = [
     {

x-pack/platform/plugins/shared/observability_ai_assistant/server/functions/context/context.ts

@@ -10,15 +10,14 @@ import { encode } from 'gpt-tokenizer';
 import { compact, last } from 'lodash';
 import { Observable } from 'rxjs';
 import { FunctionRegistrationParameters } from '..';
+import { CONTEXT_FUNCTION_NAME } from '../../../common';
 import { MessageAddEvent } from '../../../common/conversation_complete';
 import { Message } from '../../../common/types';
 import { createFunctionResponseMessage } from '../../../common/utils/create_function_response_message';
 import { recallAndScore } from './utils/recall_and_score';
 
 const MAX_TOKEN_COUNT_FOR_DATA_ON_SCREEN = 1000;
 
-export const CONTEXT_FUNCTION_NAME = 'context';
-
 export function registerContextFunction({
   client,
   functions,
@@ -39,8 +38,8 @@ export function registerContextFunction({
         const screenDescription = compact(
           screenContexts.map((context) => context.screenDescription)
         ).join('\n\n');
-        // any data that falls within the token limit, send it automatically
 
+        // any data that falls within the token limit, send it automatically
         const dataWithinTokenLimit = compact(
           screenContexts.flatMap((context) => context.data)
         ).filter(

x-pack/platform/plugins/shared/observability_ai_assistant/server/functions/elasticsearch.ts

@@ -6,25 +6,12 @@
  */
 
 import type { FunctionRegistrationParameters } from '.';
-
-export const ELASTICSEARCH_FUNCTION_NAME = 'elasticsearch';
+import { ELASTICSEARCH_FUNCTION_NAME } from '..';
 
 export function registerElasticsearchFunction({
   functions,
   resources,
 }: FunctionRegistrationParameters) {
-  functions.registerInstruction(({ availableFunctionNames }) => {
-    if (availableFunctionNames.includes(ELASTICSEARCH_FUNCTION_NAME)) {
-      return `You can use the ${ELASTICSEARCH_FUNCTION_NAME} tool to call Elasticsearch APIs on behalf of the user.
-      You are only allowed to perform GET requests (Some examples for GET requests are: Retrieving cluster information, cluster license, cluster health, indices stats, index stats, etc.) and GET/POST requests for the \`/_search\` endpoint (for search operations).
-      If the user asks to perform destructive actions or actions that are not allowed (e.g. PUT, PATCH, DELETE requests or POST requests that are not to the \`/_search\` endpoint), **NEVER** attempt to call the ${ELASTICSEARCH_FUNCTION_NAME} tool.
-      Instead, inform the user that you do not have the capability to perform those actions.
-      If you attempt to call the ${ELASTICSEARCH_FUNCTION_NAME} tool with disallowed methods (PUT, DELETE, PATCH, POST requests that are not to the \`/_search\` endpoint), it will fail.
-      For POST \`/_search\` operations, if a request body is needed, make sure the request body is a valid object.`;
-    }
-    return '';
-  });
-
   functions.registerFunction(
     {
       name: ELASTICSEARCH_FUNCTION_NAME,

x-pack/platform/plugins/shared/observability_ai_assistant/server/functions/execute_connector.ts

@@ -6,8 +6,7 @@
  */
 
 import { FunctionRegistrationParameters } from '.';
-
-export const EXECUTE_CONNECTOR_FUNCTION_NAME = 'execute_connector';
+import { EXECUTE_CONNECTOR_FUNCTION_NAME } from '..';
 
 export function registerExecuteConnectorFunction({
   functions,

x-pack/platform/plugins/shared/observability_ai_assistant/server/functions/get_data_on_screen.ts

@@ -9,8 +9,7 @@ import { compact } from 'lodash';
 import dedent from 'dedent';
 import { ObservabilityAIAssistantScreenContextRequest } from '../../common/types';
 import { ChatFunctionClient } from '../service/chat_function_client';
-
-export const GET_DATA_ON_SCREEN_FUNCTION_NAME = 'get_data_on_screen';
+import { GET_DATA_ON_SCREEN_FUNCTION_NAME, CONTEXT_FUNCTION_NAME } from '../../common';
 
 export function registerGetDataOnScreenFunction(
   functions: ChatFunctionClient,
@@ -49,11 +48,19 @@ export function registerGetDataOnScreenFunction(
     }
   );
 
-  functions.registerInstruction(({ availableFunctionNames }) =>
-    availableFunctionNames.includes(GET_DATA_ON_SCREEN_FUNCTION_NAME)
-      ? `The ${GET_DATA_ON_SCREEN_FUNCTION_NAME} function will retrieve specific content from the user's screen by specifying a data key. Use this tool to provide context-aware responses. Available data: ${dedent(
+  functions.registerInstruction(({ availableFunctionNames }) => {
+    if (availableFunctionNames.includes(GET_DATA_ON_SCREEN_FUNCTION_NAME)) {
+      return `You have access to data on the screen by calling the "${GET_DATA_ON_SCREEN_FUNCTION_NAME}" tool.
+        Use it to help the user understand what they are looking at. 
+
+        This tool will retrieve specific content from the user's screen by specifying a data key. Use this tool to provide context-aware responses. Available data: ${dedent(
           allData.map((data) => `${data.name}: ${data.description}`).join('\n')
-        )}`
-      : undefined
-  );
+        )}
+
+        A short summary of what they are looking at is available in the return of the "${CONTEXT_FUNCTION_NAME}" function.
+        Data that is compact enough automatically gets included in the response for the "${CONTEXT_FUNCTION_NAME}" function.`;
+    }
+
+    return undefined;
+  });
 }

x-pack/platform/plugins/shared/observability_ai_assistant/server/functions/get_dataset_info/get_relevant_field_names.ts

@@ -9,11 +9,15 @@ import type { ElasticsearchClient, SavedObjectsClientContract } from '@kbn/core/
 import type { DataViewsServerPluginStart } from '@kbn/data-views-plugin/server';
 import { castArray, chunk, groupBy, uniq } from 'lodash';
 import { lastValueFrom } from 'rxjs';
-import { MessageRole, ShortIdTable, type Message } from '../../../common';
+import {
+  MessageRole,
+  ShortIdTable,
+  type Message,
+  SELECT_RELEVANT_FIELDS_NAME,
+} from '../../../common';
 import { concatenateChatCompletionChunks } from '../../../common/utils/concatenate_chat_completion_chunks';
 import { FunctionCallChatFunction } from '../../service/types';
 
-export const SELECT_RELEVANT_FIELDS_NAME = 'select_relevant_fields';
 export const GET_RELEVANT_FIELD_NAMES_SYSTEM_MESSAGE = `You are a helpful assistant for Elastic Observability. 
 Your task is to determine which fields are relevant to the conversation by selecting only the field IDs from the provided list. 
 The list in the user message consists of JSON objects that map a human-readable field "name" to its unique "id". 

x-pack/platform/plugins/shared/observability_ai_assistant/server/functions/get_dataset_info/index.ts

@@ -6,13 +6,11 @@
  */
 
 import { IScopedClusterClient, Logger } from '@kbn/core/server';
-import { Message } from '../../../common';
+import { GET_DATASET_INFO_FUNCTION_NAME, Message } from '../../../common';
 import { FunctionRegistrationParameters } from '..';
 import { FunctionCallChatFunction, RespondFunctionResources } from '../../service/types';
 import { getRelevantFieldNames } from './get_relevant_field_names';
 
-export const GET_DATASET_INFO_FUNCTION_NAME = 'get_dataset_info';
-
 export function registerGetDatasetInfoFunction({
   resources,
   functions,