Merge pull request elastic#22 from joemcelroy/add-view-code-fields-re…

…design

Streaming support + setDefaultQuery

packages/kbn-ai-playground/components/sources_panel/sources_panel_for_start_chat.tsx

@@ -1,23 +1,43 @@
 /*
  * 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.
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
  */
 
 import { EuiFlexGroup, EuiFlexItem, EuiLoadingSpinner } from '@elastic/eui';
-import React from 'react';
+import React, { useEffect } from 'react';
+import { i18n } from '@kbn/i18n';
+import { useController } from 'react-hook-form';
 import { AddIndicesField } from './add_indices_field';
 import { IndicesTable } from './indices_table';
 import { StartChatPanel } from '../start_chat_panel';
 import { CreateIndexCallout } from './create_index_callout';
 import { useSourceIndicesField } from '../../hooks/useSourceIndicesField';
 import { useQueryIndices } from '../../hooks/useQueryIndices';
-import { i18n } from '@kbn/i18n';
+import { ChatFormFields } from '../../types';
+import { useIndicesFields } from '../../hooks/useIndicesFields';
+import { createQuery, getDefaultQueryFields } from '../../lib/create_query';
 
 export const SourcesPanelForStartChat: React.FC = () => {
   const { selectedIndices, removeIndex, addIndex } = useSourceIndicesField();
   const { indices, isLoading } = useQueryIndices();
+  const { fields } = useIndicesFields(selectedIndices || []);
+
+  const {
+    field: { onChange: elasticsearchQueryOnChange },
+  } = useController({
+    name: ChatFormFields.elasticsearchQuery,
+    defaultValue: {},
+  });
+
+  useEffect(() => {
+    if (fields) {
+      const defaultFields = getDefaultQueryFields(fields);
+      elasticsearchQueryOnChange(createQuery(defaultFields, fields));
+    }
+  }, [selectedIndices, fields, elasticsearchQueryOnChange]);
 
   return (
     <StartChatPanel

packages/kbn-ai-playground/lib/create_query.test.ts

@@ -7,7 +7,7 @@
  */
 
 import { IndicesQuerySourceFields } from '../types';
-import { createQuery, getDefaultQueryFields } from './create_query';
+import { createQuery, getDefaultQueryFields, getDefaultSourceFields } from './create_query';
 
 describe('create_query', () => {
   describe('createQuery', () => {
@@ -379,4 +379,53 @@ describe('create_query', () => {
       });
     });
   });
+
+  describe('getDefaultSourceFields', () => {
+    it('should return default source fields', () => {
+      const fieldDescriptors: IndicesQuerySourceFields = {
+        'search-search-labs': {
+          elser_query_fields: [],
+          dense_vector_query_fields: [],
+          bm25_query_fields: [
+            'additional_urls',
+            'title',
+            'links',
+            'id',
+            'url_host',
+            'url_path',
+            'url_path_dir3',
+            'body_content',
+            'domains',
+            'url',
+            'url_scheme',
+            'meta_description',
+            'headings',
+            'url_path_dir2',
+            'url_path_dir1',
+          ],
+          source_fields: [
+            'additional_urls',
+            'title',
+            'links',
+            'id',
+            'url_host',
+            'url_path',
+            'url_path_dir3',
+            'body_content',
+            'domains',
+            'url',
+            'url_scheme',
+            'meta_description',
+            'headings',
+            'url_path_dir2',
+            'url_path_dir1',
+          ],
+        },
+      };
+
+      expect(getDefaultSourceFields(fieldDescriptors)).toEqual({
+        'search-search-labs': ['body_content'],
+      });
+    });
+  });
 });

packages/kbn-ai-playground/lib/create_query.ts

@@ -21,6 +21,8 @@ const SUGGESTED_BM25_FIELDS = ['title', 'body_content', 'text', 'content'];
 
 const SUGGESTED_DENSE_VECTOR_FIELDS = ['content_vector.tokens'];
 
+const SUGGESTED_SOURCE_FIELDS = ['body_content', 'content', 'text'];
+
 interface Matches {
   queryMatches: any[];
   knnMatches: any[];
@@ -127,6 +129,26 @@ export function createQuery(fields: IndexFields, fieldDescriptors: IndicesQueryS
   };
 }
 
+export function getDefaultSourceFields(fieldDescriptors: IndicesQuerySourceFields): IndexFields {
+  const indexFields = Object.keys(fieldDescriptors).reduce<IndexFields>(
+    (acc: IndexFields, index: string) => {
+      const indexFieldDescriptors = fieldDescriptors[index];
+
+      const suggested = indexFieldDescriptors.source_fields.filter((x) =>
+        SUGGESTED_SOURCE_FIELDS.includes(x)
+      );
+
+      return {
+        ...acc,
+        [index]: suggested,
+      };
+    },
+    {}
+  );
+
+  return indexFields;
+}
+
 export function getDefaultQueryFields(fieldDescriptors: IndicesQuerySourceFields): IndexFields {
   const indexFields = Object.keys(fieldDescriptors).reduce<IndexFields>(
     (acc: IndexFields, index: string) => {

packages/kbn-ai-playground/providers/ai_playground_provider.tsx

@@ -24,7 +24,11 @@ export const AIPlaygroundProvider: React.FC<AIPlaygroundProviderProps> = ({
   navigateToIndexPage,
   children,
 }) => {
-  const form = useForm<ChatForm>();
+  const form = useForm<ChatForm>({
+    defaultValues: {
+      prompt: 'You are an assistant for question-answering tasks.',
+    },
+  });
 
   return (
     <ChatContext.Provider value={{ navigateToIndexPage }}>

x-pack/plugins/enterprise_search/server/routes/enterprise_search/ai_playground.ts

@@ -5,14 +5,13 @@
  * 2.0.
  */
 
-import Stream from 'stream';
-
 import Assist, { ConversationalChain } from '@elastic/ai-assist';
 
 import { Prompt } from '@elastic/ai-assist';
 import { ChatOpenAI } from '@elastic/ai-assist/models';
 import { fetchFields } from '@kbn/ai-playground/lib/fetch_query_source_fields';
 import { schema } from '@kbn/config-schema';
+import { streamFactory } from '@kbn/ml-response-stream/server';
 
 import { RouteDependencies } from '../../plugin';
 import { elasticsearchErrorHandler } from '../../utils/elasticsearch_error_handler';
@@ -68,14 +67,12 @@ export function registerAIPlaygroundRoutes({ log, router }: RouteDependencies) {
         rag: {
           index: data.indices,
           retriever: (question: string) => {
-            return {
-              text_expansion: {
-                'vector.tokens': {
-                  model_id: '.elser_model_2',
-                  model_text: question,
-                },
-              },
-            };
+            try {
+              const query = JSON.parse(data.elasticsearchQuery.replace(/{query}/g, question));
+              return query.query;
+            } catch (e) {
+              log.error('Failed to parse the Elasticsearch query', e);
+            }
           },
         },
         prompt: Prompt(data.prompt, {
@@ -87,33 +84,25 @@ export function registerAIPlaygroundRoutes({ log, router }: RouteDependencies) {
 
       const stream = await chain.stream(aiClient, messages);
 
+      const { end, push, responseWithHeaders } = streamFactory(request.headers, log);
+
       const reader = (stream as ReadableStream).getReader();
+      const textDecoder = new TextDecoder();
 
-      class UIStream extends Stream.Readable {
-        _read() {
-          const that = this;
-
-          function read() {
-            reader.read().then(({ done, value }: { done: boolean; value?: string }) => {
-              if (done) {
-                that.push(null);
-                return;
-              }
-              that.push(value);
-              read();
-            });
+      async function pushStreamUpdate() {
+        reader.read().then(({ done, value }: { done: boolean; value?: Uint8Array }) => {
+          if (done) {
+            end();
+            return;
           }
-          read();
-        }
+          push(textDecoder.decode(value));
+          pushStreamUpdate();
+        });
       }
 
-      return response.custom({
-        body: new UIStream(),
-        statusCode: 200,
-        headers: {
-          'Content-Type': 'text/plain; charset=utf-8',
-        },
-      });
+      pushStreamUpdate();
+
+      return response.ok(responseWithHeaders);
     })
   );
 }