From 3fb6f1cb1d0664fe873ccc586687654630cffef9 Mon Sep 17 00:00:00 2001 From: Joseph McElroy Date: Tue, 27 Feb 2024 14:41:20 +0100 Subject: [PATCH] changes --- .../sources_panel_for_start_chat.tsx | 28 ++++++++-- .../lib/create_query.test.ts | 51 +++++++++++++++++- .../kbn-ai-playground/lib/create_query.ts | 22 ++++++++ .../providers/ai_playground_provider.tsx | 6 ++- .../routes/enterprise_search/ai_playground.ts | 53 ++++++++----------- 5 files changed, 122 insertions(+), 38 deletions(-) diff --git a/packages/kbn-ai-playground/components/sources_panel/sources_panel_for_start_chat.tsx b/packages/kbn-ai-playground/components/sources_panel/sources_panel_for_start_chat.tsx index 0414020954f165..6e6c4aafc4c099 100644 --- a/packages/kbn-ai-playground/components/sources_panel/sources_panel_for_start_chat.tsx +++ b/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 ( { 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'], + }); + }); + }); }); diff --git a/packages/kbn-ai-playground/lib/create_query.ts b/packages/kbn-ai-playground/lib/create_query.ts index 613e5b26a9c731..be6b7e52ca078a 100644 --- a/packages/kbn-ai-playground/lib/create_query.ts +++ b/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( + (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( (acc: IndexFields, index: string) => { diff --git a/packages/kbn-ai-playground/providers/ai_playground_provider.tsx b/packages/kbn-ai-playground/providers/ai_playground_provider.tsx index 708bde9956b791..857f92cc0b1390 100644 --- a/packages/kbn-ai-playground/providers/ai_playground_provider.tsx +++ b/packages/kbn-ai-playground/providers/ai_playground_provider.tsx @@ -24,7 +24,11 @@ export const AIPlaygroundProvider: React.FC = ({ navigateToIndexPage, children, }) => { - const form = useForm(); + const form = useForm({ + defaultValues: { + prompt: 'You are an assistant for question-answering tasks.', + }, + }); return ( diff --git a/x-pack/plugins/enterprise_search/server/routes/enterprise_search/ai_playground.ts b/x-pack/plugins/enterprise_search/server/routes/enterprise_search/ai_playground.ts index c36794fd956687..132bd539a91af3 100644 --- a/x-pack/plugins/enterprise_search/server/routes/enterprise_search/ai_playground.ts +++ b/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); }) ); }