Merge branch 'main' into eui/v94.2.0

elastic · May 1, 2024 · 875eb6d · 875eb6d
2 parents f071e08 + 09337d4
commit 875eb6d
Show file tree

Hide file tree

Showing 38 changed files with 718 additions and 165 deletions.
diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.test.ts
@@ -99,17 +99,16 @@ const defaultProps = {
   telemetry: mockTelemetry,
   replacements: {},
 };
+const executorMock = initializeAgentExecutorWithOptions as jest.Mock;
 describe('callAgentExecutor', () => {
   beforeEach(() => {
     jest.clearAllMocks();
-    (initializeAgentExecutorWithOptions as jest.Mock).mockImplementation(
-      (_a, _b, { agentType }) => ({
-        // eslint-disable-next-line @typescript-eslint/no-explicit-any
-        call: (props: any, more: any) => mockCall({ ...props, agentType }, more),
-        // eslint-disable-next-line @typescript-eslint/no-explicit-any
-        invoke: (props: any, more: any) => mockInvoke({ ...props, agentType }, more),
-      })
-    );
+    executorMock.mockImplementation((_a, _b, { agentType }) => ({
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      call: (props: any, more: any) => mockCall({ ...props, agentType }, more),
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      invoke: (props: any, more: any) => mockInvoke({ ...props, agentType }, more),
+    }));
   });
 
   describe('callAgentExecutor', () => {
@@ -155,6 +154,41 @@ describe('callAgentExecutor', () => {
       expect(mockCall.mock.calls[0][0].agentType).toEqual('chat-conversational-react-description');
     });
 
+    it('uses the DynamicTool version of ESQLKnowledgeBaseTool', async () => {
+      await callAgentExecutor({
+        ...defaultProps,
+        assistantTools: [
+          {
+            name: 'ESQLKnowledgeBaseTool',
+            id: 'esql-knowledge-base-tool',
+            description: '',
+            sourceRegister: '',
+            isSupported: jest.fn(),
+            getTool: jest.fn().mockReturnValue(() => 'ESQLKnowledgeBaseTool'),
+          },
+          {
+            name: 'ESQLKnowledgeBaseStructuredTool',
+            id: 'esql-knowledge-base-structured-tool',
+            description: '',
+            sourceRegister: '',
+            isSupported: jest.fn(),
+            getTool: jest.fn().mockReturnValue(() => 'ESQLKnowledgeBaseStructuredTool'),
+          },
+          {
+            name: 'UnrelatedTool',
+            id: 'unrelated-tool',
+            description: '',
+            sourceRegister: '',
+            isSupported: jest.fn(),
+            getTool: jest.fn().mockReturnValue(() => 'UnrelatedTool'),
+          },
+        ],
+      });
+
+      expect(executorMock.mock.calls[0][0].length).toEqual(2);
+      expect(executorMock.mock.calls[0][0][0]()).toEqual('ESQLKnowledgeBaseTool');
+    });
+
     it('returns the expected response', async () => {
       const result = await callAgentExecutor(defaultProps);
 
@@ -194,6 +228,42 @@ describe('callAgentExecutor', () => {
       expect(mockInvoke.mock.calls[0][0].agentType).toEqual('openai-functions');
     });
 
+    it('uses the DynamicStructuredTool version of ESQLKnowledgeBaseTool', async () => {
+      await callAgentExecutor({
+        ...defaultProps,
+        isStream: true,
+        assistantTools: [
+          {
+            name: 'ESQLKnowledgeBaseTool',
+            id: 'esql-knowledge-base-tool',
+            description: '',
+            sourceRegister: '',
+            isSupported: jest.fn(),
+            getTool: jest.fn().mockReturnValue(() => 'ESQLKnowledgeBaseTool'),
+          },
+          {
+            name: 'ESQLKnowledgeBaseStructuredTool',
+            id: 'esql-knowledge-base-structured-tool',
+            description: '',
+            sourceRegister: '',
+            isSupported: jest.fn(),
+            getTool: jest.fn().mockReturnValue(() => 'ESQLKnowledgeBaseStructuredTool'),
+          },
+          {
+            name: 'UnrelatedTool',
+            id: 'unrelated-tool',
+            description: '',
+            sourceRegister: '',
+            isSupported: jest.fn(),
+            getTool: jest.fn().mockReturnValue(() => 'UnrelatedTool'),
+          },
+        ],
+      });
+
+      expect(executorMock.mock.calls[0][0].length).toEqual(2);
+      expect(executorMock.mock.calls[0][0][0]()).toEqual('ESQLKnowledgeBaseStructuredTool');
+    });
+
     it('returns the expected response', async () => {
       const result = await callAgentExecutor({ ...defaultProps, isStream: true });
       expect(result.body).toBeInstanceOf(Stream.PassThrough);

diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts
@@ -116,23 +116,31 @@ export const callAgentExecutor: AgentExecutor<true | false> = async ({
     request,
     size,
   };
-  const tools: ToolInterface[] = assistantTools.flatMap(
-    (tool) => tool.getTool(assistantToolParams) ?? []
-  );
+
+  const tools: ToolInterface[] = assistantTools
+    .filter((tool) =>
+      isStream
+        ? tool.id !== 'esql-knowledge-base-tool'
+        : tool.id !== 'esql-knowledge-base-structured-tool'
+    )
+    .flatMap((tool) => tool.getTool(assistantToolParams) ?? []);
 
   logger.debug(`applicable tools: ${JSON.stringify(tools.map((t) => t.name).join(', '), null, 2)}`);
 
+  const executorArgs = {
+    memory,
+    verbose: false,
+    handleParsingErrors: 'Try again, paying close attention to the allowed tool input',
+  };
   // isStream check is not on agentType alone because typescript doesn't like
   const executor = isStream
     ? await initializeAgentExecutorWithOptions(tools, llm, {
         agentType: 'openai-functions',
-        memory,
-        verbose: false,
+        ...executorArgs,
       })
     : await initializeAgentExecutorWithOptions(tools, llm, {
         agentType: 'chat-conversational-react-description',
-        memory,
-        verbose: false,
+        ...executorArgs,
       });
 
   // Sets up tracer for tracing executions to APM. See x-pack/plugins/elastic_assistant/server/lib/langchain/tracers/README.mdx

diff --git a/x-pack/plugins/elastic_assistant/server/types.ts b/x-pack/plugins/elastic_assistant/server/types.ts
@@ -20,7 +20,7 @@ import type {
   SavedObjectsClientContract,
 } from '@kbn/core/server';
 import { type MlPluginSetup } from '@kbn/ml-plugin/server';
-import { Tool } from '@langchain/core/tools';
+import { DynamicStructuredTool, Tool } from '@langchain/core/tools';
 import { SpacesPluginSetup, SpacesPluginStart } from '@kbn/spaces-plugin/server';
 import { TaskManagerSetupContract } from '@kbn/task-manager-plugin/server';
 import { AuthenticatedUser, SecurityPluginStart } from '@kbn/security-plugin/server';
@@ -202,7 +202,7 @@ export interface AssistantTool {
   description: string;
   sourceRegister: string;
   isSupported: (params: AssistantToolParams) => boolean;
-  getTool: (params: AssistantToolParams) => Tool | null;
+  getTool: (params: AssistantToolParams) => Tool | DynamicStructuredTool | null;
 }
 
 export interface AssistantToolParams {

diff --git a/x-pack/plugins/osquery/public/live_queries/form/query_details_flyout.tsx b/x-pack/plugins/osquery/public/live_queries/form/query_details_flyout.tsx
@@ -14,9 +14,9 @@ import {
   EuiFlexItem,
   EuiCodeBlock,
   EuiSpacer,
+  useEuiTheme,
 } from '@elastic/eui';
-
-import React from 'react';
+import React, { useMemo } from 'react';
 import { FormattedMessage } from '@kbn/i18n-react';
 
 interface QueryDetailsFlyoutProps {
@@ -27,46 +27,66 @@ interface QueryDetailsFlyoutProps {
   onClose: () => void;
 }
 
-const QueryDetailsFlyoutComponent: React.FC<QueryDetailsFlyoutProps> = ({ action, onClose }) => (
-  <EuiPortal>
-    <EuiFlyout size="m" ownFocus onClose={onClose} aria-labelledby="flyoutTitle">
-      <EuiFlyoutHeader hasBorder>
-        <EuiTitle size="s">
-          <h2 id="flyoutTitle">
-            <FormattedMessage
-              id="xpack.osquery.liveQueryActions.details.title"
-              defaultMessage="Query Details"
-            />
-          </h2>
-        </EuiTitle>
-      </EuiFlyoutHeader>
-      <EuiFlyoutBody>
-        <EuiFlexItem grow={false}>
-          <strong>
-            <FormattedMessage id="xpack.osquery.liveQueryActions.details.id" defaultMessage="Id" />
-          </strong>
-          <EuiSpacer size="xs" />
-          <EuiCodeBlock fontSize="m" paddingSize="s" isCopyable={true}>
-            {action.id}
-          </EuiCodeBlock>
-        </EuiFlexItem>
-        <EuiSpacer size="m" />
-        <EuiFlexItem grow={false}>
-          <strong>
-            <FormattedMessage
-              id="xpack.osquery.liveQueryActions.details.query"
-              defaultMessage="Query"
-            />
-          </strong>
-          <EuiSpacer size="xs" />
-          <EuiCodeBlock language="sql" fontSize="m" paddingSize="s" isCopyable={true}>
-            {action.query}
-          </EuiCodeBlock>
-        </EuiFlexItem>
-        <EuiSpacer size="m" />
-      </EuiFlyoutBody>
-    </EuiFlyout>
-  </EuiPortal>
-);
+const QueryDetailsFlyoutComponent: React.FC<QueryDetailsFlyoutProps> = ({ action, onClose }) => {
+  const { euiTheme } = useEuiTheme();
+
+  // we need this flyout to be above the timeline flyout (which has a z-index of 1002)
+  const maskProps = useMemo(
+    () => ({ style: `z-index: ${(euiTheme.levels.flyout as number) + 3}` }),
+    [euiTheme.levels.flyout]
+  );
+
+  return (
+    <EuiPortal>
+      <EuiFlyout
+        size="m"
+        ownFocus
+        onClose={onClose}
+        aria-labelledby="flyoutTitle"
+        // EUI TODO: This z-index override of EuiOverlayMask is a workaround, and ideally should be resolved with a cleaner UI/UX flow long-term
+        maskProps={maskProps}
+      >
+        <EuiFlyoutHeader hasBorder>
+          <EuiTitle size="s">
+            <h2 id="flyoutTitle">
+              <FormattedMessage
+                id="xpack.osquery.liveQueryActions.details.title"
+                defaultMessage="Query Details"
+              />
+            </h2>
+          </EuiTitle>
+        </EuiFlyoutHeader>
+        <EuiFlyoutBody>
+          <EuiFlexItem grow={false}>
+            <strong>
+              <FormattedMessage
+                id="xpack.osquery.liveQueryActions.details.id"
+                defaultMessage="Id"
+              />
+            </strong>
+            <EuiSpacer size="xs" />
+            <EuiCodeBlock fontSize="m" paddingSize="s" isCopyable={true}>
+              {action.id}
+            </EuiCodeBlock>
+          </EuiFlexItem>
+          <EuiSpacer size="m" />
+          <EuiFlexItem grow={false}>
+            <strong>
+              <FormattedMessage
+                id="xpack.osquery.liveQueryActions.details.query"
+                defaultMessage="Query"
+              />
+            </strong>
+            <EuiSpacer size="xs" />
+            <EuiCodeBlock language="sql" fontSize="m" paddingSize="s" isCopyable={true}>
+              {action.query}
+            </EuiCodeBlock>
+          </EuiFlexItem>
+          <EuiSpacer size="m" />
+        </EuiFlyoutBody>
+      </EuiFlyout>
+    </EuiPortal>
+  );
+};
 
 export const QueryDetailsFlyout = React.memo(QueryDetailsFlyoutComponent);
diff --git a/x-pack/plugins/search_playground/public/analytics/constants.ts b/x-pack/plugins/search_playground/public/analytics/constants.ts
@@ -0,0 +1,36 @@
+/*
+ * 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.
+ */
+
+export enum AnalyticsEvents {
+  chatPageLoaded = 'chat_page_loaded',
+  chatCleared = 'chat_cleared',
+  chatQuestionSent = 'chat_question_sent',
+  chatRequestStopped = 'chat_request_stopped',
+  chatRegenerateMessages = 'chat_regenerate_messages',
+  citationDetailsExpanded = 'citation_details_expanded',
+  citationDetailsCollapsed = 'citation_details_collapsed',
+  editContextFlyoutOpened = 'edit_context_flyout_opened',
+  editContextFieldToggled = 'edit_context_field_toggled',
+  editContextDocSizeChanged = 'edit_context_doc_size_changed',
+  editContextSaved = 'edit_context_saved',
+  genAiConnectorAdded = 'gen_ai_connector_added',
+  genAiConnectorCreated = 'gen_ai_connector_created',
+  genAiConnectorExists = 'gen_ai_connector_exists',
+  genAiConnectorSetup = 'gen_ai_connector_setup',
+  includeCitations = 'include_citations',
+  instructionsFieldChanged = 'instructions_field_changed',
+  modelSelected = 'model_selected',
+  retrievalDocsFlyoutOpened = 'retrieval_docs_flyout_opened',
+  sourceFieldsLoaded = 'source_fields_loaded',
+  sourceIndexUpdated = 'source_index_updated',
+  startNewChatPageLoaded = 'start_new_chat_page_loaded',
+  viewQueryFlyoutOpened = 'view_query_flyout_opened',
+  viewQueryFieldsUpdated = 'view_query_fields_updated',
+  viewQuerySaved = 'view_query_saved',
+  viewCodeFlyoutOpened = 'view_code_flyout_opened',
+  viewCodeLanguageChange = 'view_code_language_change',
+}
diff --git a/x-pack/plugins/search_playground/public/components/app.tsx b/x-pack/plugins/search_playground/public/components/app.tsx
@@ -5,12 +5,15 @@
  * 2.0.
  */
 
-import React from 'react';
+import React, { useState } from 'react';
 import { KibanaPageTemplate } from '@kbn/shared-ux-page-kibana-template';
 
+import { StartNewChat } from './start_new_chat';
 import { Chat } from './chat';
 
 export const App: React.FC = () => {
+  const [showStartPage, setShowStartPage] = useState(true);
+
   return (
     <KibanaPageTemplate.Section
       alignment="top"
@@ -23,7 +26,7 @@ export const App: React.FC = () => {
       paddingSize="none"
       className="eui-fullHeight"
     >
-      <Chat />
+      {showStartPage ? <StartNewChat onStartClick={() => setShowStartPage(false)} /> : <Chat />}
     </KibanaPageTemplate.Section>
   );
 };