From 0502292c03131d4de204511febd8ee44824615d2 Mon Sep 17 00:00:00 2001 From: pgayvallet Date: Thu, 19 Sep 2024 20:27:47 +0200 Subject: [PATCH 01/21] Work in progress --- .../adapters/gemini/process_vertex_stream.ts | 4 + .../nl_to_esql/actions/correct_generation.ts | 97 ++++++++++ .../tasks/nl_to_esql/actions/generate_esql.ts | 169 ++++++------------ .../server/tasks/nl_to_esql/actions/index.ts | 3 +- .../actions/request_documentation.ts | 93 +++++----- .../nl_to_esql/actions/review_generation.ts | 55 ++++++ .../server/tasks/nl_to_esql/actions/shared.ts | 1 + .../actions/summarize_discussion.ts | 141 +++++++++++++++ .../server/tasks/nl_to_esql/actions/types.ts | 28 +++ .../nl_to_esql/esql_docs/esql-bucket.txt | 19 +- .../nl_to_esql/esql_docs/esql-date_trunc.txt | 8 + .../nl_to_esql/esql_docs/esql-syntax.txt | 21 ++- .../inference/server/tasks/nl_to_esql/task.ts | 130 ++++++++++---- .../server/tasks/nl_to_esql/types.ts | 4 +- .../nl_to_esql/utils/extract_queries.test.ts | 75 ++++++++ .../tasks/nl_to_esql/utils/extract_queries.ts | 36 ++++ .../server/tasks/nl_to_esql/utils/index.ts | 9 + .../utils/validate_query_ast.test.ts | 31 ++++ .../nl_to_esql/utils/validate_query_ast.ts | 63 +++++++ .../tasks/nl_to_esql/validate_esql_query.ts | 72 -------- 20 files changed, 782 insertions(+), 277 deletions(-) create mode 100644 x-pack/plugins/inference/server/tasks/nl_to_esql/actions/correct_generation.ts create mode 100644 x-pack/plugins/inference/server/tasks/nl_to_esql/actions/review_generation.ts create mode 100644 x-pack/plugins/inference/server/tasks/nl_to_esql/actions/summarize_discussion.ts create mode 100644 x-pack/plugins/inference/server/tasks/nl_to_esql/actions/types.ts create mode 100644 x-pack/plugins/inference/server/tasks/nl_to_esql/utils/extract_queries.test.ts create mode 100644 x-pack/plugins/inference/server/tasks/nl_to_esql/utils/extract_queries.ts create mode 100644 x-pack/plugins/inference/server/tasks/nl_to_esql/utils/index.ts create mode 100644 x-pack/plugins/inference/server/tasks/nl_to_esql/utils/validate_query_ast.test.ts create mode 100644 x-pack/plugins/inference/server/tasks/nl_to_esql/utils/validate_query_ast.ts delete mode 100644 x-pack/plugins/inference/server/tasks/nl_to_esql/validate_esql_query.ts diff --git a/x-pack/plugins/inference/server/chat_complete/adapters/gemini/process_vertex_stream.ts b/x-pack/plugins/inference/server/chat_complete/adapters/gemini/process_vertex_stream.ts index e2a6c74a0447f..25e58a3c477ce 100644 --- a/x-pack/plugins/inference/server/chat_complete/adapters/gemini/process_vertex_stream.ts +++ b/x-pack/plugins/inference/server/chat_complete/adapters/gemini/process_vertex_stream.ts @@ -30,6 +30,10 @@ export function processVertexStream() { }); } + if(!value.candidates?.[0].content) { + console.log('*** WTF', JSON.stringify(value)); + } + const contentPart = value.candidates?.[0].content.parts[0]; const completion = contentPart?.text; const toolCall = contentPart?.functionCall; diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/correct_generation.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/correct_generation.ts new file mode 100644 index 0000000000000..fc22d5407a3a2 --- /dev/null +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/correct_generation.ts @@ -0,0 +1,97 @@ +/* + * 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 { lastValueFrom, tap } from 'rxjs'; +import { ToolOptions } from '../../../../common/chat_complete/tools'; +import { Message, MessageRole, ToolSchema } from '../../../../common'; +import type { ActionsOptionsBase } from './types'; +import { withoutTokenCountEvents } from '../../../../common/chat_complete/without_token_count_events'; + +const correctGenerationSchema = { + type: 'object', + properties: { + correction: { + type: 'string', + description: 'The corrected answer', + }, + }, + required: ['correction'] as const, +} satisfies ToolSchema; + +export const correctEsqlGenerationFn = ({ + client, + connectorId, + logger, + output$, + functionCalling, +}: ActionsOptionsBase) => { + return async function correctEsqlGeneration({ + messages, + generatedQuery, + documentation, + review, + }: { + messages: Message[]; + generatedQuery: string; + documentation: Record; + review: string; + }) { + const result = await lastValueFrom( + client + .chatComplete({ + connectorId, + functionCalling, + system: ` + You are an helpful Elastic assistant in charge of improving an answer related to an ES|QL question. + + You will be given: + - the original user question + - the answer another assistant provided + - a review of this answer + + Given those, please improve the answer according to the review. + You must correct and return the whole answer. If there is text content in the answer, return it too. + If the review is valid, just don't do any modifications and return the answer as it was provided as input. + + Here is the ES|QL documentation that was used to generate the query. Please use it if necessary to improve the query: + + \`\`\`json + ${JSON.stringify(documentation, undefined, 2)} + \`\`\` + + `, + messages: [ + ...messages, + { + role: MessageRole.User, + content: ` + # User question + + (See previous messages) + + # Provided answer + + ${generatedQuery} + + # Review + + ${JSON.stringify(review)} + `, + }, + ], + }) + .pipe( + withoutTokenCountEvents(), + tap((event) => { + output$.next(event); + }) + ) + ); + + return result.content; + }; +}; diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/generate_esql.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/generate_esql.ts index d31952e2f5252..d25c576067d32 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/generate_esql.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/generate_esql.ts @@ -5,7 +5,7 @@ * 2.0. */ -import { Observable, map, merge, of, switchMap } from 'rxjs'; +import { Observable, map, merge, of, switchMap, tap, lastValueFrom } from 'rxjs'; import type { Logger } from '@kbn/logging'; import { ToolCall, ToolOptions } from '../../../../common/chat_complete/tools'; import { @@ -15,67 +15,43 @@ import { Message, MessageRole, } from '../../../../common'; -import { InferenceClient, withoutTokenCountEvents } from '../../..'; +import { InferenceClient, withoutTokenCountEvents, withoutChunkEvents } from '../../..'; import { OutputCompleteEvent, OutputEventType } from '../../../../common/output'; import { INLINE_ESQL_QUERY_REGEX } from '../../../../common/tasks/nl_to_esql/constants'; import { EsqlDocumentBase } from '../doc_base'; import { requestDocumentationSchema } from './shared'; import type { NlToEsqlTaskEvent } from '../types'; -import type { FunctionCallingMode } from '../../../../common/chat_complete'; +import type { ActionsOptionsBase } from './types'; export const generateEsqlTask = ({ - chatCompleteApi, + client, connectorId, systemMessage, - messages, - toolOptions: { tools, toolChoice }, - docBase, functionCalling, logger, -}: { - connectorId: string; + output$, +}: ActionsOptionsBase & { systemMessage: string; - messages: Message[]; - toolOptions: ToolOptions; - chatCompleteApi: InferenceClient['chatComplete']; - docBase: EsqlDocumentBase; - functionCalling?: FunctionCallingMode; - logger: Pick; }) => { - return function askLlmToRespond({ - documentationRequest: { commands, functions }, + return async function generateEsql({ + messages, + documentationRequest: { keywords, requestedDocumentation }, }: { - documentationRequest: { commands?: string[]; functions?: string[] }; - }): Observable> { - const keywords = [...(commands ?? []), ...(functions ?? [])]; - const requestedDocumentation = docBase.getDocumentation(keywords); - const fakeRequestDocsToolCall = createFakeTooCall(commands, functions); + messages: Message[]; + documentationRequest: { keywords: string[]; requestedDocumentation: Record }; + }) { + const fakeRequestDocsToolCall = createFakeTooCall(keywords); - return merge( - of< - OutputCompleteEvent< - 'request_documentation', - { keywords: string[]; requestedDocumentation: Record } - > - >({ - type: OutputEventType.OutputComplete, - id: 'request_documentation', - output: { - keywords, - requestedDocumentation, - }, - content: '', - }), - chatCompleteApi({ - connectorId, - functionCalling, - system: `${systemMessage} + const result = await lastValueFrom( + client + .chatComplete({ + connectorId, + functionCalling, + system: `${systemMessage} # Current task - Your current task is to respond to the user's question. If there is a tool - suitable for answering the user's question, use that tool, preferably - with a natural language reply included. + Your current task is to respond to the user's question. Format any ES|QL query as follows: \`\`\`esql @@ -95,74 +71,47 @@ export const generateEsqlTask = ({ When converting queries from one language to ES|QL, make sure that the functions are available and documented in ES|QL. E.g., for SPL's LEN, use LENGTH. For IF, use CASE. `, - messages: [ - ...messages, - { - role: MessageRole.Assistant, - content: null, - toolCalls: [fakeRequestDocsToolCall], - }, - { - role: MessageRole.Tool, - response: { - documentation: requestedDocumentation, + messages: [ + ...messages, + { + role: MessageRole.Assistant, + content: null, + toolCalls: [fakeRequestDocsToolCall], }, - toolCallId: fakeRequestDocsToolCall.toolCallId, - }, - ], - toolChoice, - tools: { - ...tools, - request_documentation: { - description: 'Request additional ES|QL documentation if needed', - schema: requestDocumentationSchema, - }, - }, - }).pipe( - withoutTokenCountEvents(), - map((generateEvent) => { - if (isChatCompletionMessageEvent(generateEvent)) { - return { - ...generateEvent, - content: generateEvent.content - ? correctEsqlMistakes({ content: generateEvent.content, logger }) - : generateEvent.content, - }; - } - - return generateEvent; - }), - switchMap((generateEvent) => { - if (isChatCompletionMessageEvent(generateEvent)) { - const onlyToolCall = - generateEvent.toolCalls.length === 1 ? generateEvent.toolCalls[0] : undefined; - - if (onlyToolCall?.function.name === 'request_documentation') { - const args = onlyToolCall.function.arguments; - - return askLlmToRespond({ - documentationRequest: { - commands: args.commands, - functions: args.functions, - }, - }); - } - } - - return of(generateEvent); + { + role: MessageRole.Tool, + response: { + documentation: requestedDocumentation, + }, + toolCallId: fakeRequestDocsToolCall.toolCallId, + }, + ], }) - ) + .pipe( + withoutTokenCountEvents(), + map((event) => { + if (isChatCompletionMessageEvent(event)) { + return { + ...event, + content: event.content + ? correctEsqlMistakes({ content: event.content, logger }) + : event.content, + }; + } + return event; + }), + tap((event) => { + output$.next(event); + }), + withoutChunkEvents() + ) ); + + return { content: result.content }; }; }; -const correctEsqlMistakes = ({ - content, - logger, -}: { - content: string; - logger: Pick; -}) => { +const correctEsqlMistakes = ({ content, logger }: { content: string; logger: Logger }) => { return content.replaceAll(INLINE_ESQL_QUERY_REGEX, (_match, query) => { const correction = correctCommonEsqlMistakes(query); if (correction.isCorrection) { @@ -172,16 +121,12 @@ const correctEsqlMistakes = ({ }); }; -const createFakeTooCall = ( - commands: string[] | undefined, - functions: string[] | undefined -): ToolCall => { +const createFakeTooCall = (keywords: string[]): ToolCall => { return { function: { name: 'request_documentation', arguments: { - commands, - functions, + keywords, }, }, toolCallId: generateFakeToolCallId(), diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/index.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/index.ts index ec1d54dd8a26b..c7291f46203da 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/index.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/index.ts @@ -5,5 +5,6 @@ * 2.0. */ -export { requestDocumentation } from './request_documentation'; +export { requestDocumentationFn } from './request_documentation'; export { generateEsqlTask } from './generate_esql'; +export type { ActionsOptionsBase } from './types'; diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/request_documentation.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/request_documentation.ts index d4eb3060f59bb..3d9019c8ff493 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/request_documentation.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/request_documentation.ts @@ -5,59 +5,68 @@ * 2.0. */ -import { isEmpty } from 'lodash'; +import { lastValueFrom, tap, map } from 'rxjs'; import { InferenceClient, withoutOutputUpdateEvents } from '../../..'; import { Message } from '../../../../common'; -import { ToolChoiceType, ToolOptions } from '../../../../common/chat_complete/tools'; -import { requestDocumentationSchema } from './shared'; +import { EsqlDocumentBase } from '../doc_base'; import type { FunctionCallingMode } from '../../../../common/chat_complete'; +import { requestDocumentationSchema } from './shared'; +import type { ActionsOptionsBase } from './types'; -export const requestDocumentation = ({ - outputApi, +export const requestDocumentationFn = ({ + client, system, - messages, connectorId, functionCalling, - toolOptions: { tools, toolChoice }, -}: { - outputApi: InferenceClient['output']; + docBase, + logger, + output$, +}: ActionsOptionsBase & { + docBase: EsqlDocumentBase; system: string; - messages: Message[]; - connectorId: string; - functionCalling?: FunctionCallingMode; - toolOptions: ToolOptions; }) => { - const hasTools = !isEmpty(tools) && toolChoice !== ToolChoiceType.none; + return async ({ messages }: { messages: Message[] }) => { + const result = await lastValueFrom( + client + .output('request_documentation', { + connectorId, + functionCalling, + system, + previousMessages: messages, + input: `Based on the previous conversation, request documentation + for commands or functions listed in the ES|QL handbook to help you + get the right information needed to generate a query. - return outputApi('request_documentation', { - connectorId, - functionCalling, - system, - previousMessages: messages, - input: `Based on the previous conversation, request documentation - from the ES|QL handbook to help you get the right information - needed to generate a query. + Make sure to request documentation for any command or function you think you may use, + even if you end up not using all of them. - Examples for functions and commands: - - Do you need to group data? Request \`STATS\`. - - Extract data? Request \`DISSECT\` AND \`GROK\`. + Example: if you need to... + - Group or aggregate data? Request \`STATS\`. + - Extract data? Request \`DISSECT\` and \`GROK\`. - Convert a column based on a set of conditionals? Request \`EVAL\` and \`CASE\`. - - ${ - hasTools - ? `### Tools - - The following tools will be available to be called in the step after this. - - \`\`\`json - ${JSON.stringify({ - tools, - toolChoice, - })} - \`\`\`` - : '' - } + - Group data by time intervals? Request \`BUCKET\` `, - schema: requestDocumentationSchema, - }).pipe(withoutOutputUpdateEvents()); + schema: requestDocumentationSchema, + }) + .pipe( + withoutOutputUpdateEvents(), + map((event) => { + const keywords = [...(event.output.commands ?? []), ...(event.output.functions ?? [])]; + const requestedDocumentation = docBase.getDocumentation(keywords); + return { + ...event, + output: { + keywords, + requestedDocumentation, + }, + }; + }), + tap((event) => { + output$.next(event); + }) + ) + ); + + return result.output; + }; }; diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/review_generation.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/review_generation.ts new file mode 100644 index 0000000000000..a2ad727245b50 --- /dev/null +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/review_generation.ts @@ -0,0 +1,55 @@ +/* + * 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 { Observable, map, merge, of, switchMap, lastValueFrom, tap } from 'rxjs'; +import type { Logger } from '@kbn/logging'; +import { ToolCall, ToolOptions } from '../../../../common/chat_complete/tools'; +import { + correctCommonEsqlMistakes, + generateFakeToolCallId, + isChatCompletionMessageEvent, + Message, + MessageRole, + ToolSchema, +} from '../../../../common'; +import { InferenceClient, withoutOutputUpdateEvents, withoutTokenCountEvents } from '../../..'; +import { OutputCompleteEvent, OutputEventType } from '../../../../common/output'; +import { INLINE_ESQL_QUERY_REGEX } from '../../../../common/tasks/nl_to_esql/constants'; +import { EsqlDocumentBase } from '../doc_base'; +import type { NlToEsqlTaskEvent } from '../types'; +import type { FunctionCallingMode } from '../../../../common/chat_complete'; +import type { ActionsOptionsBase } from './types'; +import { extractQueries, validateQueryAst } from '../utils'; + +export const reviewEsqlGenerationFn = ({ + logger, + output$, +}: ActionsOptionsBase) => { + return async function reviewEsqlGeneration({ messageWithQuery }: { messageWithQuery: string }) { + const queries = extractQueries(messageWithQuery); + + const queriesWithErrors = await Promise.all( + queries.map(async (query) => { + const astErrors = await validateQueryAst(query.query); + return { + ...query, + astErrors, + }; + }) + ); + + const hasErrors = queriesWithErrors.some((query) => query.astErrors.length > 0); + + // TODO: simulate a tool call event + // output$.next(event); + + return { + hasErrors, + queries: queriesWithErrors, + }; + }; +}; diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/shared.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/shared.ts index f0fc796173b23..cd34962367897 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/shared.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/shared.ts @@ -26,4 +26,5 @@ export const requestDocumentationSchema = { description: 'ES|QL functions you want to analyze before generating the query.', }, }, + required: ['commands', 'functions'] as const, } satisfies ToolSchema; diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/summarize_discussion.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/summarize_discussion.ts new file mode 100644 index 0000000000000..b7ef217303441 --- /dev/null +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/summarize_discussion.ts @@ -0,0 +1,141 @@ +/* + * 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 { lastValueFrom, tap } from 'rxjs'; +import { withoutOutputUpdateEvents } from '../../..'; +import { Message } from '../../../../common'; +import { ToolSchema } from '../../../../common'; +import type { ActionsOptionsBase } from './types'; + +const summarizeDiscussionSchema = { + type: 'object', + properties: { + summary: { + type: 'string', + description: 'The requested summary of the conversation', + }, + }, + required: ['summary'] as const, +} satisfies ToolSchema; + +export const summarizeDiscussionFn = ({ + client, + connectorId, + functionCalling, + logger, + output$, +}: ActionsOptionsBase) => { + return async function summarizeDiscussion({ messages }: { messages: Message[] }) { + const result = await lastValueFrom( + client + .output('summarize_discussion', { + connectorId, + functionCalling, + system: ` + You are an helpful Elastic assistant in charge of helping the user generating + ES|QL queries. Specifically, you current role is to extract from a discussion + all the information that could be useful to answer the user's question and generate + an ES|QL query. + + Make sure to include ALL information that can be useful to generate, including: + - which index or data view to perform the query against + - the known fields, their type and their description + - the question that needs to be answered with an ES|QL query + - Any previous query that may be mentioned + + Your role is only to retrieve and extract the info that are present in the conversation, not to invent + any. E.g. if there are no description attached to the index's schema, don't invent any. + + Here are a few examples: + ${examples} + `, + previousMessages: messages, + input: `Please extract the requested summary from the current discussion, as specified in your system + instructions. + `, + schema: summarizeDiscussionSchema, + }) + .pipe( + withoutOutputUpdateEvents(), + tap((event) => { + output$.next(event); + }) + ) + ); + const summary = result.output.summary; + return { summary }; + }; +}; + +const examples = ` + + #example 1 + + ## message list: + + [ + { + "role": "user", + "content" "Show me a query to display the amount of order per month over the past year" + }, + { + "content": { + "indices": [ + "kibana_sample_data_ecommerce" + ], + "fields": [ + "order_date:date", + "order_id:keyword" + ] + } + } + ] + + ## expected summary + + The user wants to generate an ES|QL query to display the amount of orders per month over the past year. + The query should target the index 'kibana_sample_data_ecommerce'. + The relevant fields for this query are: + - order_date of type date + - order_id of type keyword + + #example 2 + + ## message list: + + [ + { + "role": "user", + "content" "Show me a query to display the amount of order per month over the past year" + } + ] + + ## expected summary + + The user wants to generate an ES|QL query to display the amount of orders per month over the past year. + There are no indications about which index should be targeted. + There are no indications about the index's schema. + + #example 3 + + ## message list: + + [ + { + "role": "user", + "content" "Can you convert this SPL query to ESQL? \`index=network_firewall "SYN Timeout" | stats count by dest\`" + } + ] + + ## expected summary + + The user wants to convert a SPL query to ESQL. + The SPL query is: \`index=network_firewall "SYN Timeout" | stats count by dest\`. + The query should target the index 'network_firewall'. + The query should re-use the fields from the SPL query. +\` +`; diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/types.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/types.ts new file mode 100644 index 0000000000000..dad0c836212b5 --- /dev/null +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/types.ts @@ -0,0 +1,28 @@ +/* + * 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 { Logger } from '@kbn/logging'; +import { Subject } from 'rxjs'; +import type { InferenceClient } from '../../../types'; +import type { FunctionCallingMode } from '../../../../common/chat_complete'; +import { NlToEsqlTaskEvent } from '../types'; + +export interface ActionsOptionsBase { + client: Pick; + connectorId: string; + logger: Logger; + functionCalling?: FunctionCallingMode; + output$: Subject>; +} + +/* + client, + connectorId, + functionCalling, + logger, + subject, + */ diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/esql_docs/esql-bucket.txt b/x-pack/plugins/inference/server/tasks/nl_to_esql/esql_docs/esql-bucket.txt index 585a0321ef818..617952666542a 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/esql_docs/esql-bucket.txt +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/esql_docs/esql-bucket.txt @@ -4,7 +4,7 @@ The BUCKET function allows you to create groups of values, known as buckets, fro ## Syntax -`BUCKET(field, buckets, from, to)` +`BUCKET(field, buckets [, from, to])` ### Parameters @@ -18,15 +18,24 @@ The target number of buckets, or the desired bucket size if `from` and `to` para #### from -The start of the range. This can be a number, a date, or a date expressed as a string. +(optional) The start of the range. This can be a number, a date, or a date expressed as a string. #### to -The end of the range. This can be a number, a date, or a date expressed as a string. +(optional) The end of the range. This can be a number, a date, or a date expressed as a string. -## Examples +## Important notes: + +BUCKET can operate in two modes: +- one where the bucket size is computed based on a bucket count recommendation and a range, +- and another where the bucket size is provided directly. -BUCKET can operate in two modes: one where the bucket size is computed based on a bucket count recommendation and a range, and another where the bucket size is provided directly. +When the bucket size is provided directly for time interval, +it is expressed as a *timespan literal*, e.g. +- GOOD: `BUCKET(@timestamp, 1 month)` +- BAD: `BUCKET(@timestamp, "month")` + +## Examples For instance, asking for at most 20 buckets over a year results in monthly buckets: diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/esql_docs/esql-date_trunc.txt b/x-pack/plugins/inference/server/tasks/nl_to_esql/esql_docs/esql-date_trunc.txt index bd1d4b68043b1..7df5249933821 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/esql_docs/esql-date_trunc.txt +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/esql_docs/esql-date_trunc.txt @@ -16,6 +16,14 @@ This is the interval to which the date will be rounded down. It is expressed usi This is the date expression that will be rounded down. +## Important notes + +The *interval* parameter of DATE_TRUNC is a timespan literal, NOT a string. + - GOOD: `DATE_TRUNC(1 year, date)` + - BAD: `DATE_TRUNC("year", date)` + +When grouping data by time interval, it is recommended to use BUCKET instead of DATE_TRUNC. + ## Examples The following example rounds down the hire_date to the nearest year: diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/esql_docs/esql-syntax.txt b/x-pack/plugins/inference/server/tasks/nl_to_esql/esql_docs/esql-syntax.txt index 8259b8d8d5546..85df775422801 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/esql_docs/esql-syntax.txt +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/esql_docs/esql-syntax.txt @@ -43,7 +43,7 @@ FROM index ## Literals -ES|QL currently supports numeric and string literals. +ES|QL currently supports numeric, string and timespan literals. ### String Literals @@ -93,11 +93,11 @@ These qualifiers are supported: - `quarter`/`quarters`/`q` - `year`/`years`/`yr`/`y` -Timespan literals are not whitespace sensitive. These expressions are all valid: - -- 1day -- 1 day -- 1 day +Timespan literals are not whitespace sensitive, and should not be wrapped with quotes: +- GOOD: 1day +- GOOD: 1 day +- BAD: "day" +- BAD: "2 days" ## Example Queries with Timespan Literals @@ -137,6 +137,15 @@ FROM sales | SORT week ``` +4. The same example with BUCKET instead of DATE_TRUNC: + +```esql +FROM sales +| WHERE @timestamp > NOW() - 1 quarter +| STATS weekly_sales = SUM(sales_amount) BY week = BUCKET(@timestamp, 1 week) +| SORT week +``` + 5. Retrieve error logs from the last 15 minutes and group by error type: ```esql diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/task.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/task.ts index e0c5a838ea148..8bee60e15da3d 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/task.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/task.ts @@ -6,12 +6,18 @@ */ import { once } from 'lodash'; -import { Observable, from, switchMap } from 'rxjs'; +import { from, Observable, Subject, switchMap } from 'rxjs'; import { Message, MessageRole } from '../../../common/chat_complete'; import type { ToolOptions } from '../../../common/chat_complete/tools'; import { EsqlDocumentBase } from './doc_base'; -import { requestDocumentation, generateEsqlTask } from './actions'; -import { NlToEsqlTaskParams, NlToEsqlTaskEvent } from './types'; +import { type ActionsOptionsBase, generateEsqlTask, requestDocumentationFn } from './actions'; +import { NlToEsqlTaskEvent, NlToEsqlTaskParams } from './types'; +import { summarizeDiscussionFn } from './actions/summarize_discussion'; +import { reviewEsqlGenerationFn } from './actions/review_generation'; +import { correctEsqlGenerationFn } from './actions/correct_generation'; +import { withoutChunkEvents } from '../../../common/chat_complete/without_chunk_events'; +import { withoutTokenCountEvents } from '../../../common/chat_complete/without_token_count_events'; +import { INLINE_ESQL_QUERY_REGEX } from '/../../../common/tasks/nl_to_esql/constants'; const loadDocBase = once(() => EsqlDocumentBase.load()); @@ -24,46 +30,94 @@ export function naturalLanguageToEsql({ functionCalling, ...rest }: NlToEsqlTaskParams): Observable> { - return from(loadDocBase()).pipe( - switchMap((docBase) => { - const systemMessage = docBase.getSystemMessage(); + const output$ = new Subject>(); + + const baseOptions: ActionsOptionsBase = { + client, + connectorId, + functionCalling, + logger, + output$, + }; + + loadDocBase() + .then(async (docBase) => { const messages: Message[] = 'input' in rest ? [{ role: MessageRole.User, content: rest.input }] : rest.messages; - const askLlmToRespond = generateEsqlTask({ - connectorId, - chatCompleteApi: client.chatComplete, - messages, + const summarizeDiscussion = summarizeDiscussionFn({ + ...baseOptions, + }); + + const requestDocumentation = requestDocumentationFn({ + ...baseOptions, docBase, - logger, - systemMessage, - functionCalling, - toolOptions: { - tools, - toolChoice, - }, + system: docBase.getSystemMessage(), + }); + + const generateEsql = generateEsqlTask({ + ...baseOptions, + systemMessage: docBase.getSystemMessage(), + }); + + const reviewEsqlGeneration = reviewEsqlGenerationFn({ + ...baseOptions, + }); + + const correctEsqlGeneration = correctEsqlGenerationFn({ + ...baseOptions, + }); + + //// actual workflow + + const discussionSummary = await summarizeDiscussion({ messages }); + + console.log('discussion summary:', discussionSummary); + + const documentationRequest = await requestDocumentation({ + messages: discussionFromSummary(discussionSummary.summary), + }); + + console.log('**** requested keywords:', documentationRequest.keywords); + + const generated = await generateEsql({ + documentationRequest, + messages: discussionFromSummary(discussionSummary.summary), + }); + + console.log('**** generated:', generated); + + const review = await reviewEsqlGeneration({ + messageWithQuery: generated.content, }); - return requestDocumentation({ - connectorId, - functionCalling, - outputApi: client.output, - messages, - system: systemMessage, - toolOptions: { - tools, - toolChoice, - }, - }).pipe( - switchMap((documentationEvent) => { - return askLlmToRespond({ - documentationRequest: { - commands: documentationEvent.output.commands, - functions: documentationEvent.output.functions, - }, - }); - }) - ); + if (review.hasErrors) { + console.log('**** review:', JSON.stringify(review)); + + const correctedAnswer = await correctEsqlGeneration({ + documentation: documentationRequest.requestedDocumentation, + messages: discussionFromSummary(discussionSummary.summary), + generatedQuery: generated.content, + review: JSON.stringify(review.queries), // TODO: fix + }); + + console.log('**** corrected answer:', correctedAnswer); + } + + // TODO + + // TODO: correctEsqlMistakes + + output$.complete(); }) - ); + .catch((err) => { + // TODO: throw inference error + output$.error(err); + }); + + return output$.asObservable(); } + +const discussionFromSummary = (summary: string): Message[] => { + return [{ role: MessageRole.User, content: summary }]; +}; diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/types.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/types.ts index a0bcd635081ea..6ec469ca4eb88 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/types.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/types.ts @@ -21,13 +21,15 @@ export type NlToEsqlTaskEvent = 'request_documentation', { keywords: string[]; requestedDocumentation: Record } > + | OutputCompleteEvent<'summarize_discussion', { summary: string }> + | OutputCompleteEvent<'review_generation', { valid: boolean; review: string }> | ChatCompletionChunkEvent | ChatCompletionMessageEvent; export type NlToEsqlTaskParams = { client: Pick; connectorId: string; - logger: Pick; + logger: Logger; functionCalling?: FunctionCallingMode; } & TToolOptions & ({ input: string } | { messages: Message[] }); diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/extract_queries.test.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/extract_queries.test.ts new file mode 100644 index 0000000000000..78b17f0902991 --- /dev/null +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/extract_queries.test.ts @@ -0,0 +1,75 @@ +/* + * 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 { extractQueries } from './extract_queries'; + +describe('extractQueries', () => { + it('finds a single query', () => { + const input = ` + Some text + + \`\`\`esql + FROM my-index + | WHERE foo == bar + | SORT foo + \`\`\` + `; + + const queries = extractQueries(input); + + expect(queries).toEqual([ + { + query: `FROM my-index + | WHERE foo == bar + | SORT foo`, + startPos: 20, + endPos: 91, + }, + ]); + }); + + it('finds multiple queries', () => { + const input = ` + Some text + + \`\`\`esql + FROM my-index + | WHERE foo == bar + | SORT foo + \`\`\` + + Another block of text + + \`\`\`esql + FROM kibana_sample_data_logs + | WHERE @timestamp > NOW() - 10 days + | STATS volume = COUNT(*) BY BUCKET(@timestamp, 1 day) + \`\`\` + + Some final block of text + `; + + const queries = extractQueries(input); + + expect(queries).toEqual([ + { + query: `FROM my-index + | WHERE foo == bar + | SORT foo`, + startPos: 20, + endPos: 91, + }, + { + query: `FROM kibana_sample_data_logs + | WHERE @timestamp > NOW() - 10 days + | STATS volume = COUNT(*) BY BUCKET(@timestamp, 1 day)`, + startPos: 124, + endPos: 272, + }, + ]); + }); +}); diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/extract_queries.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/extract_queries.ts new file mode 100644 index 0000000000000..1b6a715091a28 --- /dev/null +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/extract_queries.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. + */ + +import { INLINE_ESQL_QUERY_REGEX } from '../../../../common/tasks/nl_to_esql/constants'; + +export interface QueryWithPosition { + /** the query text, without the ```esql``` */ + query: string; + /** the query start position in the text - INCLUDING the ```esql``` */ + startPos: number; + /** the query end position in the text - INCLUDING the ```esql``` */ + endPos: number; +} + +/** + * Extract the esql queries from a given string content. + * + * Note: it will only find queries wrapped with ```esql [query] ``` + */ +export const extractQueries = (text: string): QueryWithPosition[] => { + return Array.from(text.matchAll(INLINE_ESQL_QUERY_REGEX)) + .filter((match) => { + return match[1] !== undefined && match.index !== undefined; + }) + .map((match) => { + return { + query: match[1], + startPos: match.index!, + endPos: match.index! + match[0].length, + }; + }); +}; diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/index.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/index.ts new file mode 100644 index 0000000000000..1d2ba2d78ccc2 --- /dev/null +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/index.ts @@ -0,0 +1,9 @@ +/* + * 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 { extractQueries, type QueryWithPosition } from './extract_queries'; +export { validateQueryAst } from './validate_query_ast'; diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/validate_query_ast.test.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/validate_query_ast.test.ts new file mode 100644 index 0000000000000..79bcda6f08e2e --- /dev/null +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/validate_query_ast.test.ts @@ -0,0 +1,31 @@ +/* + * 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 { validateQueryAst } from './validate_query_ast'; + +describe('validateQueryAst', () => { + it('returns AST validation errors', async () => { + const query = `FROM kibana_sample_data_logs + | KOUKOU foobar + | WHERE`; + + const errors = await validateQueryAst(query); + + expect(errors).toEqual([ + { + startPos: 36, + endPos: 42, + message: expect.stringContaining(`mismatched input 'KOUKOU'`), + }, + { + startPos: 61, + endPos: 61, + message: expect.stringContaining(`mismatched input ''`), + }, + ]); + }); +}); diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/validate_query_ast.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/validate_query_ast.ts new file mode 100644 index 0000000000000..7e70b14aa595c --- /dev/null +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/validate_query_ast.ts @@ -0,0 +1,63 @@ +/* + * 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 { validateQuery } from '@kbn/esql-validation-autocomplete'; +import { getAstAndSyntaxErrors } from '@kbn/esql-ast'; +import { splitIntoCommands } from '../../../../common/tasks/nl_to_esql/correct_common_esql_mistakes'; + +interface QuerySyntaxError { + message: string; + startPos: number; + endPos: number; + code?: string; +} + +/** + * Run the query through the kibana + * @param query + */ +export const validateQueryAst = async (query: string): Promise => { + const { errors: astErrors } = await validateQuery(query, getAstAndSyntaxErrors, { + // setting this to true, we don't want to validate the index / fields existence + ignoreOnMissingCallbacks: true, + }); + + const asCommands = splitIntoCommands(query); + + const errors = (astErrors ?? []).map((error) => { + if ('location' in error) { + return { + message: error.text, + code: error.code, + startPos: error.location.min, + endPos: error.location.max, + }; + } else { + return { + message: error.message, + code: error.code, + startPos: getPosition(asCommands, error.startLineNumber, error.startColumn), + endPos: getPosition(asCommands, error.endLineNumber, error.endColumn), + }; + } + }); + + return errors; +}; + +const getPosition = ( + commands: ReturnType, + line: number, + column: number +): number => { + const previousCommands = commands.slice(0, line - 1); + return ( + previousCommands.reduce((count, command) => { + return count + command.command.length; + }, 0) + column + ); +}; diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/validate_esql_query.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/validate_esql_query.ts deleted file mode 100644 index 823344f52a891..0000000000000 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/validate_esql_query.ts +++ /dev/null @@ -1,72 +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. - */ - -import { validateQuery } from '@kbn/esql-validation-autocomplete'; -import { getAstAndSyntaxErrors } from '@kbn/esql-ast'; -import type { ElasticsearchClient } from '@kbn/core/server'; -import { ESQLSearchResponse, ESQLRow } from '@kbn/es-types'; -import { esFieldTypeToKibanaFieldType } from '@kbn/field-types'; -import { DatatableColumn, DatatableColumnType } from '@kbn/expressions-plugin/common'; -import { splitIntoCommands } from '../../../common/tasks/nl_to_esql/correct_common_esql_mistakes'; - -export async function runAndValidateEsqlQuery({ - query, - client, -}: { - query: string; - client: ElasticsearchClient; -}): Promise<{ - columns?: DatatableColumn[]; - rows?: ESQLRow[]; - error?: Error; - errorMessages?: string[]; -}> { - const { errors } = await validateQuery(query, getAstAndSyntaxErrors, { - // setting this to true, we don't want to validate the index / fields existence - ignoreOnMissingCallbacks: true, - }); - - const asCommands = splitIntoCommands(query); - - const errorMessages = errors?.map((error) => { - if ('location' in error) { - const commandsUntilEndOfError = splitIntoCommands(query.substring(0, error.location.max)); - const lastCompleteCommand = asCommands[commandsUntilEndOfError.length - 1]; - if (lastCompleteCommand) { - return `Error in ${lastCompleteCommand.command}\n: ${error.text}`; - } - } - return 'text' in error ? error.text : error.message; - }); - - return client.transport - .request({ - method: 'POST', - path: '_query', - body: { - query, - }, - }) - .then((res) => { - const esqlResponse = res as ESQLSearchResponse; - - const columns = - esqlResponse.columns?.map(({ name, type }) => ({ - id: name, - name, - meta: { type: esFieldTypeToKibanaFieldType(type) as DatatableColumnType }, - })) ?? []; - - return { columns, rows: esqlResponse.values }; - }) - .catch((error) => { - return { - error, - ...(errorMessages.length ? { errorMessages } : {}), - }; - }); -} From 71dfeb2ee1d87756b6e4d4f27ec9df422035a9b6 Mon Sep 17 00:00:00 2001 From: pgayvallet Date: Tue, 29 Oct 2024 13:13:29 +0100 Subject: [PATCH 02/21] remove additional things --- .../adapters/gemini/process_vertex_stream.ts | 4 - .../nl_to_esql/actions/correct_generation.ts | 97 ------------------- .../tasks/nl_to_esql/actions/generate_esql.ts | 3 +- .../server/tasks/nl_to_esql/actions/index.ts | 3 +- .../actions/request_documentation.ts | 29 +++++- .../nl_to_esql/actions/review_generation.ts | 55 ----------- .../server/tasks/nl_to_esql/actions/shared.ts | 30 ------ .../inference/server/tasks/nl_to_esql/task.ts | 52 ++-------- 8 files changed, 36 insertions(+), 237 deletions(-) delete mode 100644 x-pack/plugins/inference/server/tasks/nl_to_esql/actions/correct_generation.ts delete mode 100644 x-pack/plugins/inference/server/tasks/nl_to_esql/actions/review_generation.ts delete mode 100644 x-pack/plugins/inference/server/tasks/nl_to_esql/actions/shared.ts diff --git a/x-pack/plugins/inference/server/chat_complete/adapters/gemini/process_vertex_stream.ts b/x-pack/plugins/inference/server/chat_complete/adapters/gemini/process_vertex_stream.ts index 25e58a3c477ce..e2a6c74a0447f 100644 --- a/x-pack/plugins/inference/server/chat_complete/adapters/gemini/process_vertex_stream.ts +++ b/x-pack/plugins/inference/server/chat_complete/adapters/gemini/process_vertex_stream.ts @@ -30,10 +30,6 @@ export function processVertexStream() { }); } - if(!value.candidates?.[0].content) { - console.log('*** WTF', JSON.stringify(value)); - } - const contentPart = value.candidates?.[0].content.parts[0]; const completion = contentPart?.text; const toolCall = contentPart?.functionCall; diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/correct_generation.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/correct_generation.ts deleted file mode 100644 index fc22d5407a3a2..0000000000000 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/correct_generation.ts +++ /dev/null @@ -1,97 +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. - */ - -import { lastValueFrom, tap } from 'rxjs'; -import { ToolOptions } from '../../../../common/chat_complete/tools'; -import { Message, MessageRole, ToolSchema } from '../../../../common'; -import type { ActionsOptionsBase } from './types'; -import { withoutTokenCountEvents } from '../../../../common/chat_complete/without_token_count_events'; - -const correctGenerationSchema = { - type: 'object', - properties: { - correction: { - type: 'string', - description: 'The corrected answer', - }, - }, - required: ['correction'] as const, -} satisfies ToolSchema; - -export const correctEsqlGenerationFn = ({ - client, - connectorId, - logger, - output$, - functionCalling, -}: ActionsOptionsBase) => { - return async function correctEsqlGeneration({ - messages, - generatedQuery, - documentation, - review, - }: { - messages: Message[]; - generatedQuery: string; - documentation: Record; - review: string; - }) { - const result = await lastValueFrom( - client - .chatComplete({ - connectorId, - functionCalling, - system: ` - You are an helpful Elastic assistant in charge of improving an answer related to an ES|QL question. - - You will be given: - - the original user question - - the answer another assistant provided - - a review of this answer - - Given those, please improve the answer according to the review. - You must correct and return the whole answer. If there is text content in the answer, return it too. - If the review is valid, just don't do any modifications and return the answer as it was provided as input. - - Here is the ES|QL documentation that was used to generate the query. Please use it if necessary to improve the query: - - \`\`\`json - ${JSON.stringify(documentation, undefined, 2)} - \`\`\` - - `, - messages: [ - ...messages, - { - role: MessageRole.User, - content: ` - # User question - - (See previous messages) - - # Provided answer - - ${generatedQuery} - - # Review - - ${JSON.stringify(review)} - `, - }, - ], - }) - .pipe( - withoutTokenCountEvents(), - tap((event) => { - output$.next(event); - }) - ) - ); - - return result.content; - }; -}; diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/generate_esql.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/generate_esql.ts index d25c576067d32..47d40a3ce4e32 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/generate_esql.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/generate_esql.ts @@ -19,11 +19,10 @@ import { InferenceClient, withoutTokenCountEvents, withoutChunkEvents } from '.. import { OutputCompleteEvent, OutputEventType } from '../../../../common/output'; import { INLINE_ESQL_QUERY_REGEX } from '../../../../common/tasks/nl_to_esql/constants'; import { EsqlDocumentBase } from '../doc_base'; -import { requestDocumentationSchema } from './shared'; import type { NlToEsqlTaskEvent } from '../types'; import type { ActionsOptionsBase } from './types'; -export const generateEsqlTask = ({ +export const generateEsqlTaskFn = ({ client, connectorId, systemMessage, diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/index.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/index.ts index c7291f46203da..563a1d2099157 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/index.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/index.ts @@ -6,5 +6,6 @@ */ export { requestDocumentationFn } from './request_documentation'; -export { generateEsqlTask } from './generate_esql'; +export { generateEsqlTaskFn } from './generate_esql'; +export { summarizeDiscussionFn } from './summarize_discussion'; export type { ActionsOptionsBase } from './types'; diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/request_documentation.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/request_documentation.ts index 3d9019c8ff493..b98bfb8f3fae6 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/request_documentation.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/request_documentation.ts @@ -6,20 +6,39 @@ */ import { lastValueFrom, tap, map } from 'rxjs'; -import { InferenceClient, withoutOutputUpdateEvents } from '../../..'; -import { Message } from '../../../../common'; +import { withoutOutputUpdateEvents } from '../../..'; +import { Message, ToolSchema } from '../../../../common'; import { EsqlDocumentBase } from '../doc_base'; -import type { FunctionCallingMode } from '../../../../common/chat_complete'; -import { requestDocumentationSchema } from './shared'; import type { ActionsOptionsBase } from './types'; +export const requestDocumentationSchema = { + type: 'object', + properties: { + commands: { + type: 'array', + items: { + type: 'string', + }, + description: + 'ES|QL source and processing commands you want to analyze before generating the query.', + }, + functions: { + type: 'array', + items: { + type: 'string', + }, + description: 'ES|QL functions you want to analyze before generating the query.', + }, + }, + required: ['commands', 'functions'] as const, +} satisfies ToolSchema; + export const requestDocumentationFn = ({ client, system, connectorId, functionCalling, docBase, - logger, output$, }: ActionsOptionsBase & { docBase: EsqlDocumentBase; diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/review_generation.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/review_generation.ts deleted file mode 100644 index a2ad727245b50..0000000000000 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/review_generation.ts +++ /dev/null @@ -1,55 +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. - */ - -import { Observable, map, merge, of, switchMap, lastValueFrom, tap } from 'rxjs'; -import type { Logger } from '@kbn/logging'; -import { ToolCall, ToolOptions } from '../../../../common/chat_complete/tools'; -import { - correctCommonEsqlMistakes, - generateFakeToolCallId, - isChatCompletionMessageEvent, - Message, - MessageRole, - ToolSchema, -} from '../../../../common'; -import { InferenceClient, withoutOutputUpdateEvents, withoutTokenCountEvents } from '../../..'; -import { OutputCompleteEvent, OutputEventType } from '../../../../common/output'; -import { INLINE_ESQL_QUERY_REGEX } from '../../../../common/tasks/nl_to_esql/constants'; -import { EsqlDocumentBase } from '../doc_base'; -import type { NlToEsqlTaskEvent } from '../types'; -import type { FunctionCallingMode } from '../../../../common/chat_complete'; -import type { ActionsOptionsBase } from './types'; -import { extractQueries, validateQueryAst } from '../utils'; - -export const reviewEsqlGenerationFn = ({ - logger, - output$, -}: ActionsOptionsBase) => { - return async function reviewEsqlGeneration({ messageWithQuery }: { messageWithQuery: string }) { - const queries = extractQueries(messageWithQuery); - - const queriesWithErrors = await Promise.all( - queries.map(async (query) => { - const astErrors = await validateQueryAst(query.query); - return { - ...query, - astErrors, - }; - }) - ); - - const hasErrors = queriesWithErrors.some((query) => query.astErrors.length > 0); - - // TODO: simulate a tool call event - // output$.next(event); - - return { - hasErrors, - queries: queriesWithErrors, - }; - }; -}; diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/shared.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/shared.ts deleted file mode 100644 index cd34962367897..0000000000000 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/shared.ts +++ /dev/null @@ -1,30 +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. - */ - -import { ToolSchema } from '../../../../common'; - -export const requestDocumentationSchema = { - type: 'object', - properties: { - commands: { - type: 'array', - items: { - type: 'string', - }, - description: - 'ES|QL source and processing commands you want to analyze before generating the query.', - }, - functions: { - type: 'array', - items: { - type: 'string', - }, - description: 'ES|QL functions you want to analyze before generating the query.', - }, - }, - required: ['commands', 'functions'] as const, -} satisfies ToolSchema; diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/task.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/task.ts index 8bee60e15da3d..ebd7626949f07 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/task.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/task.ts @@ -6,18 +6,17 @@ */ import { once } from 'lodash'; -import { from, Observable, Subject, switchMap } from 'rxjs'; +import { Observable, Subject } from 'rxjs'; import { Message, MessageRole } from '../../../common/chat_complete'; import type { ToolOptions } from '../../../common/chat_complete/tools'; import { EsqlDocumentBase } from './doc_base'; -import { type ActionsOptionsBase, generateEsqlTask, requestDocumentationFn } from './actions'; +import { + type ActionsOptionsBase, + generateEsqlTaskFn, + requestDocumentationFn, + summarizeDiscussionFn, +} from './actions'; import { NlToEsqlTaskEvent, NlToEsqlTaskParams } from './types'; -import { summarizeDiscussionFn } from './actions/summarize_discussion'; -import { reviewEsqlGenerationFn } from './actions/review_generation'; -import { correctEsqlGenerationFn } from './actions/correct_generation'; -import { withoutChunkEvents } from '../../../common/chat_complete/without_chunk_events'; -import { withoutTokenCountEvents } from '../../../common/chat_complete/without_token_count_events'; -import { INLINE_ESQL_QUERY_REGEX } from '/../../../common/tasks/nl_to_esql/constants'; const loadDocBase = once(() => EsqlDocumentBase.load()); @@ -48,26 +47,16 @@ export function naturalLanguageToEsql({ const summarizeDiscussion = summarizeDiscussionFn({ ...baseOptions, }); - const requestDocumentation = requestDocumentationFn({ ...baseOptions, docBase, system: docBase.getSystemMessage(), }); - - const generateEsql = generateEsqlTask({ + const generateEsql = generateEsqlTaskFn({ ...baseOptions, systemMessage: docBase.getSystemMessage(), }); - const reviewEsqlGeneration = reviewEsqlGenerationFn({ - ...baseOptions, - }); - - const correctEsqlGeneration = correctEsqlGenerationFn({ - ...baseOptions, - }); - //// actual workflow const discussionSummary = await summarizeDiscussion({ messages }); @@ -80,34 +69,11 @@ export function naturalLanguageToEsql({ console.log('**** requested keywords:', documentationRequest.keywords); - const generated = await generateEsql({ + await generateEsql({ documentationRequest, messages: discussionFromSummary(discussionSummary.summary), }); - console.log('**** generated:', generated); - - const review = await reviewEsqlGeneration({ - messageWithQuery: generated.content, - }); - - if (review.hasErrors) { - console.log('**** review:', JSON.stringify(review)); - - const correctedAnswer = await correctEsqlGeneration({ - documentation: documentationRequest.requestedDocumentation, - messages: discussionFromSummary(discussionSummary.summary), - generatedQuery: generated.content, - review: JSON.stringify(review.queries), // TODO: fix - }); - - console.log('**** corrected answer:', correctedAnswer); - } - - // TODO - - // TODO: correctEsqlMistakes - output$.complete(); }) .catch((err) => { From 1748dce0275ade8ed2ef750b9a239804b0f67ceb Mon Sep 17 00:00:00 2001 From: pgayvallet Date: Tue, 29 Oct 2024 15:53:07 +0100 Subject: [PATCH 03/21] prepare for type extraction --- .../inference/common/chat_complete/api.ts | 35 ++++++ .../common/chat_complete/event_utils.ts | 57 +++++++++ .../inference/common/chat_complete/events.ts | 49 ++++++++ .../inference/common/chat_complete/index.ts | 118 ++++-------------- .../is_chat_completion_chunk_event.ts | 14 --- .../chat_complete/is_chat_completion_event.ts | 17 --- .../is_chat_completion_message_event.ts | 15 --- .../common/chat_complete/messages.ts | 33 +++++ .../inference/common/chat_complete/request.ts | 3 +- .../chat_complete/without_chunk_events.ts | 19 --- .../without_token_count_events.ts | 19 --- x-pack/plugins/inference/common/index.ts | 19 ++- x-pack/plugins/inference/common/output/api.ts | 36 ++++++ .../common/output/create_output_api.ts | 2 +- .../inference/common/output/event_utils.ts | 37 ++++++ .../plugins/inference/common/output/events.ts | 39 ++++++ .../plugins/inference/common/output/index.ts | 88 +++---------- .../common/output/is_output_complete_event.ts | 14 --- .../common/output/is_output_event.ts | 15 --- .../common/output/is_output_update_event.ts | 14 --- .../output/without_output_update_events.ts | 18 --- .../generate_fake_tool_call_id.ts | 0 .../common/{util => utils}/truncate_list.ts | 0 .../actions/summarize_discussion.ts | 1 - .../inference/server/tasks/nl_to_esql/task.ts | 19 ++- .../server/tasks/nl_to_esql/types.ts | 1 + 26 files changed, 347 insertions(+), 335 deletions(-) create mode 100644 x-pack/plugins/inference/common/chat_complete/api.ts create mode 100644 x-pack/plugins/inference/common/chat_complete/event_utils.ts create mode 100644 x-pack/plugins/inference/common/chat_complete/events.ts delete mode 100644 x-pack/plugins/inference/common/chat_complete/is_chat_completion_chunk_event.ts delete mode 100644 x-pack/plugins/inference/common/chat_complete/is_chat_completion_event.ts delete mode 100644 x-pack/plugins/inference/common/chat_complete/is_chat_completion_message_event.ts create mode 100644 x-pack/plugins/inference/common/chat_complete/messages.ts delete mode 100644 x-pack/plugins/inference/common/chat_complete/without_chunk_events.ts delete mode 100644 x-pack/plugins/inference/common/chat_complete/without_token_count_events.ts create mode 100644 x-pack/plugins/inference/common/output/api.ts create mode 100644 x-pack/plugins/inference/common/output/event_utils.ts create mode 100644 x-pack/plugins/inference/common/output/events.ts delete mode 100644 x-pack/plugins/inference/common/output/is_output_complete_event.ts delete mode 100644 x-pack/plugins/inference/common/output/is_output_event.ts delete mode 100644 x-pack/plugins/inference/common/output/is_output_update_event.ts delete mode 100644 x-pack/plugins/inference/common/output/without_output_update_events.ts rename x-pack/plugins/inference/common/{chat_complete => utils}/generate_fake_tool_call_id.ts (100%) rename x-pack/plugins/inference/common/{util => utils}/truncate_list.ts (100%) diff --git a/x-pack/plugins/inference/common/chat_complete/api.ts b/x-pack/plugins/inference/common/chat_complete/api.ts new file mode 100644 index 0000000000000..5e8d6908a10cc --- /dev/null +++ b/x-pack/plugins/inference/common/chat_complete/api.ts @@ -0,0 +1,35 @@ +/* + * 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 { Observable } from 'rxjs'; +import type { ToolOptions } from './tools'; +import type { Message } from './messages'; +import type { ChatCompletionEvent } from './events'; + +/** + * Request a completion from the LLM based on a prompt or conversation. + * + * @param {string} options.connectorId The ID of the connector to use + * @param {string} [options.system] A system message that defines the behavior of the LLM. + * @param {Message[]} options.message A list of messages that make up the conversation to be completed. + * @param {ToolChoice} [options.toolChoice] Force the LLM to call a (specific) tool, or no tool + * @param {Record} [options.tools] A map of tools that can be called by the LLM + */ +export type ChatCompleteAPI = ( + options: { + connectorId: string; + system?: string; + messages: Message[]; + functionCalling?: FunctionCallingMode; + } & TToolOptions +) => ChatCompletionResponse; + +export type ChatCompletionResponse = Observable< + ChatCompletionEvent +>; + +export type FunctionCallingMode = 'native' | 'simulated'; diff --git a/x-pack/plugins/inference/common/chat_complete/event_utils.ts b/x-pack/plugins/inference/common/chat_complete/event_utils.ts new file mode 100644 index 0000000000000..bbe812c89d4df --- /dev/null +++ b/x-pack/plugins/inference/common/chat_complete/event_utils.ts @@ -0,0 +1,57 @@ +/* + * 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 { filter, OperatorFunction } from 'rxjs'; +import { InferenceTaskEvent } from '../inference_task'; +import { + ChatCompletionEventType, + ChatCompletionEvent, + ChatCompletionChunkEvent, + ChatCompletionMessageEvent, + ChatCompletionTokenCountEvent, +} from './events'; +import type { ToolOptions } from './tools'; + +export function isChatCompletionChunkEvent( + event: ChatCompletionEvent +): event is ChatCompletionChunkEvent { + return event.type === ChatCompletionEventType.ChatCompletionChunk; +} + +export function isChatCompletionMessageEvent>( + event: ChatCompletionEvent +): event is ChatCompletionMessageEvent { + return event.type === ChatCompletionEventType.ChatCompletionMessage; +} + +export function isChatCompletionEvent(event: InferenceTaskEvent): event is ChatCompletionEvent { + return ( + event.type === ChatCompletionEventType.ChatCompletionChunk || + event.type === ChatCompletionEventType.ChatCompletionMessage || + event.type === ChatCompletionEventType.ChatCompletionTokenCount + ); +} + +export function withoutChunkEvents(): OperatorFunction< + T, + Exclude +> { + return filter( + (event): event is Exclude => + event.type !== ChatCompletionEventType.ChatCompletionChunk + ); +} + +export function withoutTokenCountEvents(): OperatorFunction< + T, + Exclude +> { + return filter( + (event): event is Exclude => + event.type !== ChatCompletionEventType.ChatCompletionTokenCount + ); +} diff --git a/x-pack/plugins/inference/common/chat_complete/events.ts b/x-pack/plugins/inference/common/chat_complete/events.ts new file mode 100644 index 0000000000000..924b9631239b7 --- /dev/null +++ b/x-pack/plugins/inference/common/chat_complete/events.ts @@ -0,0 +1,49 @@ +/* + * 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 { InferenceTaskEventBase } from '../inference_task'; +import type { ToolCallsOf, ToolOptions } from './tools'; + +export enum ChatCompletionEventType { + ChatCompletionChunk = 'chatCompletionChunk', + ChatCompletionTokenCount = 'chatCompletionTokenCount', + ChatCompletionMessage = 'chatCompletionMessage', +} + +export type ChatCompletionMessageEvent = + InferenceTaskEventBase & { + content: string; + } & { toolCalls: ToolCallsOf['toolCalls'] }; + +export interface ChatCompletionChunkToolCall { + index: number; + toolCallId: string; + function: { + name: string; + arguments: string; + }; +} + +export type ChatCompletionChunkEvent = + InferenceTaskEventBase & { + content: string; + tool_calls: ChatCompletionChunkToolCall[]; + }; + +export type ChatCompletionTokenCountEvent = + InferenceTaskEventBase & { + tokens: { + prompt: number; + completion: number; + total: number; + }; + }; + +export type ChatCompletionEvent = + | ChatCompletionChunkEvent + | ChatCompletionTokenCountEvent + | ChatCompletionMessageEvent; diff --git a/x-pack/plugins/inference/common/chat_complete/index.ts b/x-pack/plugins/inference/common/chat_complete/index.ts index aef9de12ba7a9..3f4bcfb13015b 100644 --- a/x-pack/plugins/inference/common/chat_complete/index.ts +++ b/x-pack/plugins/inference/common/chat_complete/index.ts @@ -5,95 +5,29 @@ * 2.0. */ -import type { Observable } from 'rxjs'; -import type { InferenceTaskEventBase } from '../inference_task'; -import type { ToolCall, ToolCallsOf, ToolOptions } from './tools'; - -export enum MessageRole { - User = 'user', - Assistant = 'assistant', - Tool = 'tool', -} - -interface MessageBase { - role: TRole; -} - -export type UserMessage = MessageBase & { content: string }; - -export type AssistantMessage = MessageBase & { - content: string | null; - toolCalls?: Array | undefined>>; -}; - -export type ToolMessage | unknown> = - MessageBase & { - toolCallId: string; - response: TToolResponse; - }; - -export type Message = UserMessage | AssistantMessage | ToolMessage; - -export type ChatCompletionMessageEvent = - InferenceTaskEventBase & { - content: string; - } & { toolCalls: ToolCallsOf['toolCalls'] }; - -export type ChatCompletionResponse = Observable< - ChatCompletionEvent ->; - -export enum ChatCompletionEventType { - ChatCompletionChunk = 'chatCompletionChunk', - ChatCompletionTokenCount = 'chatCompletionTokenCount', - ChatCompletionMessage = 'chatCompletionMessage', -} - -export interface ChatCompletionChunkToolCall { - index: number; - toolCallId: string; - function: { - name: string; - arguments: string; - }; -} - -export type ChatCompletionChunkEvent = - InferenceTaskEventBase & { - content: string; - tool_calls: ChatCompletionChunkToolCall[]; - }; - -export type ChatCompletionTokenCountEvent = - InferenceTaskEventBase & { - tokens: { - prompt: number; - completion: number; - total: number; - }; - }; - -export type ChatCompletionEvent = - | ChatCompletionChunkEvent - | ChatCompletionTokenCountEvent - | ChatCompletionMessageEvent; - -export type FunctionCallingMode = 'native' | 'simulated'; - -/** - * Request a completion from the LLM based on a prompt or conversation. - * - * @param {string} options.connectorId The ID of the connector to use - * @param {string} [options.system] A system message that defines the behavior of the LLM. - * @param {Message[]} options.message A list of messages that make up the conversation to be completed. - * @param {ToolChoice} [options.toolChoice] Force the LLM to call a (specific) tool, or no tool - * @param {Record} [options.tools] A map of tools that can be called by the LLM - */ -export type ChatCompleteAPI = ( - options: { - connectorId: string; - system?: string; - messages: Message[]; - functionCalling?: FunctionCallingMode; - } & TToolOptions -) => ChatCompletionResponse; +export type { ChatCompletionResponse, ChatCompleteAPI, FunctionCallingMode } from './api'; +export { + ChatCompletionEventType, + type ChatCompletionMessageEvent, + type ChatCompletionChunkEvent, + type ChatCompletionEvent, + type ChatCompletionChunkToolCall, + type ChatCompletionTokenCountEvent, +} from './events'; +export { + MessageRole, + type Message, + type AssistantMessage, + type UserMessage, + type ToolMessage, +} from './messages'; +export { type ToolSchema, type ToolSchemaType, type FromToolSchema } from './tool_schema'; +export { + ToolChoiceType, + type ToolOptions, + type ToolDefinition, + type ToolCall, + type ToolCallsOf, + type UnvalidatedToolCall, + type ToolChoice, +} from './tools'; diff --git a/x-pack/plugins/inference/common/chat_complete/is_chat_completion_chunk_event.ts b/x-pack/plugins/inference/common/chat_complete/is_chat_completion_chunk_event.ts deleted file mode 100644 index 1630d765ab81e..0000000000000 --- a/x-pack/plugins/inference/common/chat_complete/is_chat_completion_chunk_event.ts +++ /dev/null @@ -1,14 +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. - */ - -import { ChatCompletionChunkEvent, ChatCompletionEvent, ChatCompletionEventType } from '.'; - -export function isChatCompletionChunkEvent( - event: ChatCompletionEvent -): event is ChatCompletionChunkEvent { - return event.type === ChatCompletionEventType.ChatCompletionChunk; -} diff --git a/x-pack/plugins/inference/common/chat_complete/is_chat_completion_event.ts b/x-pack/plugins/inference/common/chat_complete/is_chat_completion_event.ts deleted file mode 100644 index d4d9305cac94b..0000000000000 --- a/x-pack/plugins/inference/common/chat_complete/is_chat_completion_event.ts +++ /dev/null @@ -1,17 +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. - */ - -import { ChatCompletionEvent, ChatCompletionEventType } from '.'; -import { InferenceTaskEvent } from '../inference_task'; - -export function isChatCompletionEvent(event: InferenceTaskEvent): event is ChatCompletionEvent { - return ( - event.type === ChatCompletionEventType.ChatCompletionChunk || - event.type === ChatCompletionEventType.ChatCompletionMessage || - event.type === ChatCompletionEventType.ChatCompletionTokenCount - ); -} diff --git a/x-pack/plugins/inference/common/chat_complete/is_chat_completion_message_event.ts b/x-pack/plugins/inference/common/chat_complete/is_chat_completion_message_event.ts deleted file mode 100644 index 172e55df9e4b4..0000000000000 --- a/x-pack/plugins/inference/common/chat_complete/is_chat_completion_message_event.ts +++ /dev/null @@ -1,15 +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. - */ - -import { ChatCompletionEvent, ChatCompletionEventType, ChatCompletionMessageEvent } from '.'; -import type { ToolOptions } from './tools'; - -export function isChatCompletionMessageEvent>( - event: ChatCompletionEvent -): event is ChatCompletionMessageEvent { - return event.type === ChatCompletionEventType.ChatCompletionMessage; -} diff --git a/x-pack/plugins/inference/common/chat_complete/messages.ts b/x-pack/plugins/inference/common/chat_complete/messages.ts new file mode 100644 index 0000000000000..df29d9580700b --- /dev/null +++ b/x-pack/plugins/inference/common/chat_complete/messages.ts @@ -0,0 +1,33 @@ +/* + * 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 { ToolCall } from './tools'; + +export enum MessageRole { + User = 'user', + Assistant = 'assistant', + Tool = 'tool', +} + +interface MessageBase { + role: TRole; +} + +export type UserMessage = MessageBase & { content: string }; + +export type AssistantMessage = MessageBase & { + content: string | null; + toolCalls?: Array | undefined>>; +}; + +export type ToolMessage | unknown> = + MessageBase & { + toolCallId: string; + response: TToolResponse; + }; + +export type Message = UserMessage | AssistantMessage | ToolMessage; diff --git a/x-pack/plugins/inference/common/chat_complete/request.ts b/x-pack/plugins/inference/common/chat_complete/request.ts index 1038e481a6260..60341bfdaffa5 100644 --- a/x-pack/plugins/inference/common/chat_complete/request.ts +++ b/x-pack/plugins/inference/common/chat_complete/request.ts @@ -5,7 +5,8 @@ * 2.0. */ -import type { Message, FunctionCallingMode } from '.'; +import type { FunctionCallingMode } from './api'; +import type { Message } from './messages'; import type { ToolOptions } from './tools'; export type ChatCompleteRequestBody = { diff --git a/x-pack/plugins/inference/common/chat_complete/without_chunk_events.ts b/x-pack/plugins/inference/common/chat_complete/without_chunk_events.ts deleted file mode 100644 index 58e72e2c90903..0000000000000 --- a/x-pack/plugins/inference/common/chat_complete/without_chunk_events.ts +++ /dev/null @@ -1,19 +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. - */ - -import { filter, OperatorFunction } from 'rxjs'; -import { ChatCompletionChunkEvent, ChatCompletionEvent, ChatCompletionEventType } from '.'; - -export function withoutChunkEvents(): OperatorFunction< - T, - Exclude -> { - return filter( - (event): event is Exclude => - event.type !== ChatCompletionEventType.ChatCompletionChunk - ); -} diff --git a/x-pack/plugins/inference/common/chat_complete/without_token_count_events.ts b/x-pack/plugins/inference/common/chat_complete/without_token_count_events.ts deleted file mode 100644 index 1b7dbdb9c1372..0000000000000 --- a/x-pack/plugins/inference/common/chat_complete/without_token_count_events.ts +++ /dev/null @@ -1,19 +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. - */ - -import { filter, OperatorFunction } from 'rxjs'; -import { ChatCompletionEvent, ChatCompletionEventType, ChatCompletionTokenCountEvent } from '.'; - -export function withoutTokenCountEvents(): OperatorFunction< - T, - Exclude -> { - return filter( - (event): event is Exclude => - event.type !== ChatCompletionEventType.ChatCompletionTokenCount - ); -} diff --git a/x-pack/plugins/inference/common/index.ts b/x-pack/plugins/inference/common/index.ts index 58c84a47c1804..b84e03e2ad771 100644 --- a/x-pack/plugins/inference/common/index.ts +++ b/x-pack/plugins/inference/common/index.ts @@ -10,22 +10,21 @@ export { splitIntoCommands, } from './tasks/nl_to_esql/correct_common_esql_mistakes'; -export { isChatCompletionChunkEvent } from './chat_complete/is_chat_completion_chunk_event'; -export { isChatCompletionMessageEvent } from './chat_complete/is_chat_completion_message_event'; -export { isChatCompletionEvent } from './chat_complete/is_chat_completion_event'; - -export { isOutputUpdateEvent } from './output/is_output_update_event'; -export { isOutputCompleteEvent } from './output/is_output_complete_event'; -export { isOutputEvent } from './output/is_output_event'; +export { + isChatCompletionChunkEvent, + isChatCompletionMessageEvent, + isChatCompletionEvent, +} from './chat_complete/event_utils'; -export type { ToolSchema } from './chat_complete/tool_schema'; +export { isOutputUpdateEvent, isOutputCompleteEvent, isOutputEvent } from './output/event_utils'; export { - type Message, MessageRole, + type Message, type ToolMessage, type AssistantMessage, type UserMessage, + type ToolSchema, } from './chat_complete'; -export { generateFakeToolCallId } from './chat_complete/generate_fake_tool_call_id'; +export { generateFakeToolCallId } from './utils/generate_fake_tool_call_id'; diff --git a/x-pack/plugins/inference/common/output/api.ts b/x-pack/plugins/inference/common/output/api.ts new file mode 100644 index 0000000000000..f6bb530185d41 --- /dev/null +++ b/x-pack/plugins/inference/common/output/api.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. + */ + +import type { Observable } from 'rxjs'; +import type { Message, FunctionCallingMode, FromToolSchema, ToolSchema } from '../chat_complete'; +import type { OutputEvent } from './events'; + +/** + * Generate a response with the LLM for a prompt, optionally based on a schema. + * + * @param {string} id The id of the operation + * @param {string} options.connectorId The ID of the connector that is to be used. + * @param {string} options.input The prompt for the LLM. + * @param {string} options.messages Previous messages in a conversation. + * @param {ToolSchema} [options.schema] The schema the response from the LLM should adhere to. + */ +export type OutputAPI = < + TId extends string = string, + TOutputSchema extends ToolSchema | undefined = ToolSchema | undefined +>( + id: TId, + options: { + connectorId: string; + system?: string; + input: string; + schema?: TOutputSchema; + previousMessages?: Message[]; + functionCalling?: FunctionCallingMode; + } +) => Observable< + OutputEvent : undefined> +>; diff --git a/x-pack/plugins/inference/common/output/create_output_api.ts b/x-pack/plugins/inference/common/output/create_output_api.ts index 848135beefb0f..d232d15843794 100644 --- a/x-pack/plugins/inference/common/output/create_output_api.ts +++ b/x-pack/plugins/inference/common/output/create_output_api.ts @@ -7,7 +7,7 @@ import { map } from 'rxjs'; import { ChatCompleteAPI, ChatCompletionEventType, MessageRole } from '../chat_complete'; -import { withoutTokenCountEvents } from '../chat_complete/without_token_count_events'; +import { withoutTokenCountEvents } from '../chat_complete/event_utils'; import { OutputAPI, OutputEvent, OutputEventType } from '.'; import { ensureMultiTurn } from '../ensure_multi_turn'; diff --git a/x-pack/plugins/inference/common/output/event_utils.ts b/x-pack/plugins/inference/common/output/event_utils.ts new file mode 100644 index 0000000000000..b005b67d8a6dc --- /dev/null +++ b/x-pack/plugins/inference/common/output/event_utils.ts @@ -0,0 +1,37 @@ +/* + * 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 { filter, OperatorFunction } from 'rxjs'; +import { OutputEvent, OutputEventType, OutputUpdateEvent } from '.'; +import type { InferenceTaskEvent } from '../inference_task'; + +export function isOutputCompleteEvent( + event: TOutputEvent +): event is Exclude { + return event.type === OutputEventType.OutputComplete; +} + +export function isOutputEvent(event: InferenceTaskEvent): event is OutputEvent { + return ( + event.type === OutputEventType.OutputComplete || event.type === OutputEventType.OutputUpdate + ); +} + +export function isOutputUpdateEvent( + event: OutputEvent +): event is OutputUpdateEvent { + return event.type === OutputEventType.OutputComplete; +} + +export function withoutOutputUpdateEvents(): OperatorFunction< + T, + Exclude +> { + return filter( + (event): event is Exclude => event.type !== OutputEventType.OutputUpdate + ); +} diff --git a/x-pack/plugins/inference/common/output/events.ts b/x-pack/plugins/inference/common/output/events.ts new file mode 100644 index 0000000000000..d95b821bec6de --- /dev/null +++ b/x-pack/plugins/inference/common/output/events.ts @@ -0,0 +1,39 @@ +/* + * 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 { ServerSentEventBase } from '@kbn/sse-utils'; + +export enum OutputEventType { + OutputUpdate = 'output', + OutputComplete = 'complete', +} + +export type Output = Record | undefined | unknown; + +export type OutputUpdateEvent = ServerSentEventBase< + OutputEventType.OutputUpdate, + { + id: TId; + content: string; + } +>; + +export type OutputCompleteEvent< + TId extends string = string, + TOutput extends Output = Output +> = ServerSentEventBase< + OutputEventType.OutputComplete, + { + id: TId; + output: TOutput; + content: string; + } +>; + +export type OutputEvent = + | OutputUpdateEvent + | OutputCompleteEvent; diff --git a/x-pack/plugins/inference/common/output/index.ts b/x-pack/plugins/inference/common/output/index.ts index 0f7655f8f1cd4..7b294af455dd2 100644 --- a/x-pack/plugins/inference/common/output/index.ts +++ b/x-pack/plugins/inference/common/output/index.ts @@ -5,77 +5,17 @@ * 2.0. */ -import { Observable } from 'rxjs'; -import { ServerSentEventBase } from '@kbn/sse-utils'; -import { FromToolSchema, ToolSchema } from '../chat_complete/tool_schema'; -import type { Message, FunctionCallingMode } from '../chat_complete'; - -export enum OutputEventType { - OutputUpdate = 'output', - OutputComplete = 'complete', -} - -type Output = Record | undefined | unknown; - -export type OutputUpdateEvent = ServerSentEventBase< - OutputEventType.OutputUpdate, - { - id: TId; - content: string; - } ->; - -export type OutputCompleteEvent< - TId extends string = string, - TOutput extends Output = Output -> = ServerSentEventBase< - OutputEventType.OutputComplete, - { - id: TId; - output: TOutput; - content: string; - } ->; - -export type OutputEvent = - | OutputUpdateEvent - | OutputCompleteEvent; - -/** - * Generate a response with the LLM for a prompt, optionally based on a schema. - * - * @param {string} id The id of the operation - * @param {string} options.connectorId The ID of the connector that is to be used. - * @param {string} options.input The prompt for the LLM. - * @param {string} options.messages Previous messages in a conversation. - * @param {ToolSchema} [options.schema] The schema the response from the LLM should adhere to. - */ -export type OutputAPI = < - TId extends string = string, - TOutputSchema extends ToolSchema | undefined = ToolSchema | undefined ->( - id: TId, - options: { - connectorId: string; - system?: string; - input: string; - schema?: TOutputSchema; - previousMessages?: Message[]; - functionCalling?: FunctionCallingMode; - } -) => Observable< - OutputEvent : undefined> ->; - -export function createOutputCompleteEvent( - id: TId, - output: TOutput, - content?: string -): OutputCompleteEvent { - return { - type: OutputEventType.OutputComplete, - id, - output, - content: content ?? '', - }; -} +export type { OutputAPI } from './api'; +export { + OutputEventType, + type OutputCompleteEvent, + type OutputUpdateEvent, + type Output, + type OutputEvent, +} from './events'; +export { + isOutputCompleteEvent, + isOutputUpdateEvent, + isOutputEvent, + withoutOutputUpdateEvents, +} from './event_utils'; diff --git a/x-pack/plugins/inference/common/output/is_output_complete_event.ts b/x-pack/plugins/inference/common/output/is_output_complete_event.ts deleted file mode 100644 index bac3443b8258c..0000000000000 --- a/x-pack/plugins/inference/common/output/is_output_complete_event.ts +++ /dev/null @@ -1,14 +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. - */ - -import { OutputEvent, OutputEventType, OutputUpdateEvent } from '.'; - -export function isOutputCompleteEvent( - event: TOutputEvent -): event is Exclude { - return event.type === OutputEventType.OutputComplete; -} diff --git a/x-pack/plugins/inference/common/output/is_output_event.ts b/x-pack/plugins/inference/common/output/is_output_event.ts deleted file mode 100644 index dad2b0967a6ac..0000000000000 --- a/x-pack/plugins/inference/common/output/is_output_event.ts +++ /dev/null @@ -1,15 +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. - */ - -import { OutputEvent, OutputEventType } from '.'; -import type { InferenceTaskEvent } from '../inference_task'; - -export function isOutputEvent(event: InferenceTaskEvent): event is OutputEvent { - return ( - event.type === OutputEventType.OutputComplete || event.type === OutputEventType.OutputUpdate - ); -} diff --git a/x-pack/plugins/inference/common/output/is_output_update_event.ts b/x-pack/plugins/inference/common/output/is_output_update_event.ts deleted file mode 100644 index 459436e64014e..0000000000000 --- a/x-pack/plugins/inference/common/output/is_output_update_event.ts +++ /dev/null @@ -1,14 +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. - */ - -import { OutputEvent, OutputEventType, OutputUpdateEvent } from '.'; - -export function isOutputUpdateEvent( - event: OutputEvent -): event is OutputUpdateEvent { - return event.type === OutputEventType.OutputComplete; -} diff --git a/x-pack/plugins/inference/common/output/without_output_update_events.ts b/x-pack/plugins/inference/common/output/without_output_update_events.ts deleted file mode 100644 index 38f26c8c8ece1..0000000000000 --- a/x-pack/plugins/inference/common/output/without_output_update_events.ts +++ /dev/null @@ -1,18 +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. - */ - -import { filter, OperatorFunction } from 'rxjs'; -import { OutputEvent, OutputEventType, OutputUpdateEvent } from '.'; - -export function withoutOutputUpdateEvents(): OperatorFunction< - T, - Exclude -> { - return filter( - (event): event is Exclude => event.type !== OutputEventType.OutputUpdate - ); -} diff --git a/x-pack/plugins/inference/common/chat_complete/generate_fake_tool_call_id.ts b/x-pack/plugins/inference/common/utils/generate_fake_tool_call_id.ts similarity index 100% rename from x-pack/plugins/inference/common/chat_complete/generate_fake_tool_call_id.ts rename to x-pack/plugins/inference/common/utils/generate_fake_tool_call_id.ts diff --git a/x-pack/plugins/inference/common/util/truncate_list.ts b/x-pack/plugins/inference/common/utils/truncate_list.ts similarity index 100% rename from x-pack/plugins/inference/common/util/truncate_list.ts rename to x-pack/plugins/inference/common/utils/truncate_list.ts diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/summarize_discussion.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/summarize_discussion.ts index b7ef217303441..bc89e3beeedb4 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/summarize_discussion.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/summarize_discussion.ts @@ -26,7 +26,6 @@ export const summarizeDiscussionFn = ({ client, connectorId, functionCalling, - logger, output$, }: ActionsOptionsBase) => { return async function summarizeDiscussion({ messages }: { messages: Message[] }) { diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/task.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/task.ts index ebd7626949f07..f9f21e9825230 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/task.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/task.ts @@ -27,6 +27,7 @@ export function naturalLanguageToEsql({ toolChoice, logger, functionCalling, + summarizeInput = false, ...rest }: NlToEsqlTaskParams): Observable> { const output$ = new Subject>(); @@ -41,7 +42,7 @@ export function naturalLanguageToEsql({ loadDocBase() .then(async (docBase) => { - const messages: Message[] = + let messages: Message[] = 'input' in rest ? [{ role: MessageRole.User, content: rest.input }] : rest.messages; const summarizeDiscussion = summarizeDiscussionFn({ @@ -57,27 +58,23 @@ export function naturalLanguageToEsql({ systemMessage: docBase.getSystemMessage(), }); - //// actual workflow - - const discussionSummary = await summarizeDiscussion({ messages }); - - console.log('discussion summary:', discussionSummary); + if (summarizeInput) { + const discussionSummary = await summarizeDiscussion({ messages }); + messages = discussionFromSummary(discussionSummary.summary); + } const documentationRequest = await requestDocumentation({ - messages: discussionFromSummary(discussionSummary.summary), + messages, }); - console.log('**** requested keywords:', documentationRequest.keywords); - await generateEsql({ documentationRequest, - messages: discussionFromSummary(discussionSummary.summary), + messages, }); output$.complete(); }) .catch((err) => { - // TODO: throw inference error output$.error(err); }); diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/types.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/types.ts index 6ec469ca4eb88..1fa300daaba00 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/types.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/types.ts @@ -31,5 +31,6 @@ export type NlToEsqlTaskParams = { connectorId: string; logger: Logger; functionCalling?: FunctionCallingMode; + summarizeInput?: boolean; } & TToolOptions & ({ input: string } | { messages: Message[] }); From ffd9db577c0d5e7405120af3a9bbd9c7a1a69c70 Mon Sep 17 00:00:00 2001 From: pgayvallet Date: Wed, 30 Oct 2024 08:22:43 +0100 Subject: [PATCH 04/21] move types to @kbn/inference-common --- package.json | 1 + tsconfig.base.json | 2 + .../ai-infra/inference-common/README.md | 7 +++ .../ai-infra/inference-common/index.ts | 55 +++++++++++++++++++ .../ai-infra/inference-common/jest.config.js | 12 ++++ .../ai-infra/inference-common/kibana.jsonc | 5 ++ .../ai-infra/inference-common/package.json | 7 +++ .../src}/chat_complete/api.ts | 0 .../src}/chat_complete/errors.ts | 0 .../src}/chat_complete/event_utils.ts | 0 .../src}/chat_complete/events.ts | 0 .../src}/chat_complete/index.ts | 7 +++ .../src}/chat_complete/messages.ts | 0 .../src}/chat_complete/request.ts | 0 .../src}/chat_complete/tool_schema.ts | 0 .../src}/chat_complete/tools.ts | 0 .../inference-common/src}/inference_task.ts | 0 .../inference-common/src}/output/api.ts | 0 .../src}/output/event_utils.ts | 0 .../inference-common/src}/output/events.ts | 0 .../inference-common/src}/output/index.ts | 0 .../ai-infra/inference-common/tsconfig.json | 19 +++++++ .../common/{output => }/create_output_api.ts | 0 x-pack/plugins/inference/common/index.ts | 17 ------ .../nl_to_esql/correct_query_with_actions.ts | 1 + .../common/{ => utils}/ensure_multi_turn.ts | 2 +- yarn.lock | 4 ++ 27 files changed, 121 insertions(+), 18 deletions(-) create mode 100644 x-pack/packages/ai-infra/inference-common/README.md create mode 100644 x-pack/packages/ai-infra/inference-common/index.ts create mode 100644 x-pack/packages/ai-infra/inference-common/jest.config.js create mode 100644 x-pack/packages/ai-infra/inference-common/kibana.jsonc create mode 100644 x-pack/packages/ai-infra/inference-common/package.json rename x-pack/{plugins/inference/common => packages/ai-infra/inference-common/src}/chat_complete/api.ts (100%) rename x-pack/{plugins/inference/common => packages/ai-infra/inference-common/src}/chat_complete/errors.ts (100%) rename x-pack/{plugins/inference/common => packages/ai-infra/inference-common/src}/chat_complete/event_utils.ts (100%) rename x-pack/{plugins/inference/common => packages/ai-infra/inference-common/src}/chat_complete/events.ts (100%) rename x-pack/{plugins/inference/common => packages/ai-infra/inference-common/src}/chat_complete/index.ts (85%) rename x-pack/{plugins/inference/common => packages/ai-infra/inference-common/src}/chat_complete/messages.ts (100%) rename x-pack/{plugins/inference/common => packages/ai-infra/inference-common/src}/chat_complete/request.ts (100%) rename x-pack/{plugins/inference/common => packages/ai-infra/inference-common/src}/chat_complete/tool_schema.ts (100%) rename x-pack/{plugins/inference/common => packages/ai-infra/inference-common/src}/chat_complete/tools.ts (100%) rename x-pack/{plugins/inference/common => packages/ai-infra/inference-common/src}/inference_task.ts (100%) rename x-pack/{plugins/inference/common => packages/ai-infra/inference-common/src}/output/api.ts (100%) rename x-pack/{plugins/inference/common => packages/ai-infra/inference-common/src}/output/event_utils.ts (100%) rename x-pack/{plugins/inference/common => packages/ai-infra/inference-common/src}/output/events.ts (100%) rename x-pack/{plugins/inference/common => packages/ai-infra/inference-common/src}/output/index.ts (100%) create mode 100644 x-pack/packages/ai-infra/inference-common/tsconfig.json rename x-pack/plugins/inference/common/{output => }/create_output_api.ts (100%) rename x-pack/plugins/inference/common/{ => utils}/ensure_multi_turn.ts (92%) diff --git a/package.json b/package.json index 09e19ab86b1ca..16e83f75b1269 100644 --- a/package.json +++ b/package.json @@ -571,6 +571,7 @@ "@kbn/index-management-plugin": "link:x-pack/plugins/index_management", "@kbn/index-management-shared-types": "link:x-pack/packages/index-management/index_management_shared_types", "@kbn/index-patterns-test-plugin": "link:test/plugin_functional/plugins/index_patterns", + "@kbn/inference-common": "link:x-pack/packages/ai-infra/inference-common", "@kbn/inference-plugin": "link:x-pack/plugins/inference", "@kbn/inference_integration_flyout": "link:x-pack/packages/ml/inference_integration_flyout", "@kbn/infra-forge": "link:x-pack/packages/kbn-infra-forge", diff --git a/tsconfig.base.json b/tsconfig.base.json index 4471cb1bc6754..727cb930bc606 100644 --- a/tsconfig.base.json +++ b/tsconfig.base.json @@ -1044,6 +1044,8 @@ "@kbn/index-patterns-test-plugin/*": ["test/plugin_functional/plugins/index_patterns/*"], "@kbn/inference_integration_flyout": ["x-pack/packages/ml/inference_integration_flyout"], "@kbn/inference_integration_flyout/*": ["x-pack/packages/ml/inference_integration_flyout/*"], + "@kbn/inference-common": ["x-pack/packages/ai-infra/inference-common"], + "@kbn/inference-common/*": ["x-pack/packages/ai-infra/inference-common/*"], "@kbn/inference-plugin": ["x-pack/plugins/inference"], "@kbn/inference-plugin/*": ["x-pack/plugins/inference/*"], "@kbn/infra-forge": ["x-pack/packages/kbn-infra-forge"], diff --git a/x-pack/packages/ai-infra/inference-common/README.md b/x-pack/packages/ai-infra/inference-common/README.md new file mode 100644 index 0000000000000..f16f1ce9cea75 --- /dev/null +++ b/x-pack/packages/ai-infra/inference-common/README.md @@ -0,0 +1,7 @@ +# @kbn/inference-common + +Common types and utilities for the inference APIs and features. + +The main purpose of the package is to have a clean line between the inference plugin's +implementation and the underlying types, so that other packages or plugins can leverage the +types without directly depending on the plugin. diff --git a/x-pack/packages/ai-infra/inference-common/index.ts b/x-pack/packages/ai-infra/inference-common/index.ts new file mode 100644 index 0000000000000..2f043c850c92f --- /dev/null +++ b/x-pack/packages/ai-infra/inference-common/index.ts @@ -0,0 +1,55 @@ +/* + * 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 { + MessageRole, + ChatCompletionEventType, + ToolChoiceType, + type Message, + type AssistantMessage, + type ToolMessage, + type UserMessage, + type ToolSchemaType, + type FromToolSchema, + type ToolSchema, + type UnvalidatedToolCall, + type ToolCallsOf, + type ToolCall, + type ToolDefinition, + type ToolOptions, + type FunctionCallingMode, + type ToolChoice, + type ChatCompleteAPI, + type ChatCompletionResponse, + type ChatCompletionTokenCountEvent, + type ChatCompletionEvent, + type ChatCompletionChunkEvent, + type ChatCompletionChunkToolCall, + type ChatCompletionMessageEvent, + withoutTokenCountEvents, + withoutChunkEvents, + isChatCompletionMessageEvent, + isChatCompletionEvent, + isChatCompletionChunkEvent, +} from './src/chat_complete'; +export { + OutputEventType, + type OutputAPI, + type OutputCompleteEvent, + type OutputUpdateEvent, + type Output, + type OutputEvent, + isOutputCompleteEvent, + isOutputUpdateEvent, + isOutputEvent, + withoutOutputUpdateEvents, +} from './src/output'; +export { + InferenceTaskEventType, + type InferenceTaskEvent, + type InferenceTaskEventBase, +} from './src/inference_task'; diff --git a/x-pack/packages/ai-infra/inference-common/jest.config.js b/x-pack/packages/ai-infra/inference-common/jest.config.js new file mode 100644 index 0000000000000..faa0d30b40233 --- /dev/null +++ b/x-pack/packages/ai-infra/inference-common/jest.config.js @@ -0,0 +1,12 @@ +/* + * 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. + */ + +module.exports = { + preset: '@kbn/test', + rootDir: '../../../..', + roots: ['/x-pack/packages/ai-infra/inference-common'], +}; diff --git a/x-pack/packages/ai-infra/inference-common/kibana.jsonc b/x-pack/packages/ai-infra/inference-common/kibana.jsonc new file mode 100644 index 0000000000000..568755d303c3b --- /dev/null +++ b/x-pack/packages/ai-infra/inference-common/kibana.jsonc @@ -0,0 +1,5 @@ +{ + "type": "shared-common", + "id": "@kbn/inference-common", + "owner": "@elastic/appex-ai-infra" +} diff --git a/x-pack/packages/ai-infra/inference-common/package.json b/x-pack/packages/ai-infra/inference-common/package.json new file mode 100644 index 0000000000000..8deaf0a5504f0 --- /dev/null +++ b/x-pack/packages/ai-infra/inference-common/package.json @@ -0,0 +1,7 @@ +{ + "name": "@kbn/inference-common", + "private": true, + "version": "1.0.0", + "license": "Elastic License 2.0 OR AGPL-3.0-only OR SSPL-1.0", + "sideEffects": false +} diff --git a/x-pack/plugins/inference/common/chat_complete/api.ts b/x-pack/packages/ai-infra/inference-common/src/chat_complete/api.ts similarity index 100% rename from x-pack/plugins/inference/common/chat_complete/api.ts rename to x-pack/packages/ai-infra/inference-common/src/chat_complete/api.ts diff --git a/x-pack/plugins/inference/common/chat_complete/errors.ts b/x-pack/packages/ai-infra/inference-common/src/chat_complete/errors.ts similarity index 100% rename from x-pack/plugins/inference/common/chat_complete/errors.ts rename to x-pack/packages/ai-infra/inference-common/src/chat_complete/errors.ts diff --git a/x-pack/plugins/inference/common/chat_complete/event_utils.ts b/x-pack/packages/ai-infra/inference-common/src/chat_complete/event_utils.ts similarity index 100% rename from x-pack/plugins/inference/common/chat_complete/event_utils.ts rename to x-pack/packages/ai-infra/inference-common/src/chat_complete/event_utils.ts diff --git a/x-pack/plugins/inference/common/chat_complete/events.ts b/x-pack/packages/ai-infra/inference-common/src/chat_complete/events.ts similarity index 100% rename from x-pack/plugins/inference/common/chat_complete/events.ts rename to x-pack/packages/ai-infra/inference-common/src/chat_complete/events.ts diff --git a/x-pack/plugins/inference/common/chat_complete/index.ts b/x-pack/packages/ai-infra/inference-common/src/chat_complete/index.ts similarity index 85% rename from x-pack/plugins/inference/common/chat_complete/index.ts rename to x-pack/packages/ai-infra/inference-common/src/chat_complete/index.ts index 3f4bcfb13015b..1a66a535c1ebc 100644 --- a/x-pack/plugins/inference/common/chat_complete/index.ts +++ b/x-pack/packages/ai-infra/inference-common/src/chat_complete/index.ts @@ -31,3 +31,10 @@ export { type UnvalidatedToolCall, type ToolChoice, } from './tools'; +export { + isChatCompletionChunkEvent, + isChatCompletionEvent, + isChatCompletionMessageEvent, + withoutChunkEvents, + withoutTokenCountEvents, +} from './event_utils'; diff --git a/x-pack/plugins/inference/common/chat_complete/messages.ts b/x-pack/packages/ai-infra/inference-common/src/chat_complete/messages.ts similarity index 100% rename from x-pack/plugins/inference/common/chat_complete/messages.ts rename to x-pack/packages/ai-infra/inference-common/src/chat_complete/messages.ts diff --git a/x-pack/plugins/inference/common/chat_complete/request.ts b/x-pack/packages/ai-infra/inference-common/src/chat_complete/request.ts similarity index 100% rename from x-pack/plugins/inference/common/chat_complete/request.ts rename to x-pack/packages/ai-infra/inference-common/src/chat_complete/request.ts diff --git a/x-pack/plugins/inference/common/chat_complete/tool_schema.ts b/x-pack/packages/ai-infra/inference-common/src/chat_complete/tool_schema.ts similarity index 100% rename from x-pack/plugins/inference/common/chat_complete/tool_schema.ts rename to x-pack/packages/ai-infra/inference-common/src/chat_complete/tool_schema.ts diff --git a/x-pack/plugins/inference/common/chat_complete/tools.ts b/x-pack/packages/ai-infra/inference-common/src/chat_complete/tools.ts similarity index 100% rename from x-pack/plugins/inference/common/chat_complete/tools.ts rename to x-pack/packages/ai-infra/inference-common/src/chat_complete/tools.ts diff --git a/x-pack/plugins/inference/common/inference_task.ts b/x-pack/packages/ai-infra/inference-common/src/inference_task.ts similarity index 100% rename from x-pack/plugins/inference/common/inference_task.ts rename to x-pack/packages/ai-infra/inference-common/src/inference_task.ts diff --git a/x-pack/plugins/inference/common/output/api.ts b/x-pack/packages/ai-infra/inference-common/src/output/api.ts similarity index 100% rename from x-pack/plugins/inference/common/output/api.ts rename to x-pack/packages/ai-infra/inference-common/src/output/api.ts diff --git a/x-pack/plugins/inference/common/output/event_utils.ts b/x-pack/packages/ai-infra/inference-common/src/output/event_utils.ts similarity index 100% rename from x-pack/plugins/inference/common/output/event_utils.ts rename to x-pack/packages/ai-infra/inference-common/src/output/event_utils.ts diff --git a/x-pack/plugins/inference/common/output/events.ts b/x-pack/packages/ai-infra/inference-common/src/output/events.ts similarity index 100% rename from x-pack/plugins/inference/common/output/events.ts rename to x-pack/packages/ai-infra/inference-common/src/output/events.ts diff --git a/x-pack/plugins/inference/common/output/index.ts b/x-pack/packages/ai-infra/inference-common/src/output/index.ts similarity index 100% rename from x-pack/plugins/inference/common/output/index.ts rename to x-pack/packages/ai-infra/inference-common/src/output/index.ts diff --git a/x-pack/packages/ai-infra/inference-common/tsconfig.json b/x-pack/packages/ai-infra/inference-common/tsconfig.json new file mode 100644 index 0000000000000..b05325b824a67 --- /dev/null +++ b/x-pack/packages/ai-infra/inference-common/tsconfig.json @@ -0,0 +1,19 @@ +{ + "extends": "../../../../tsconfig.base.json", + "compilerOptions": { + "outDir": "target/types", + "types": [ + "jest", + "node", + "react" + ] + }, + "include": [ + "**/*.ts", + "**/*.tsx", + ], + "exclude": [ + "target/**/*" + ], + "kbn_references": [] +} diff --git a/x-pack/plugins/inference/common/output/create_output_api.ts b/x-pack/plugins/inference/common/create_output_api.ts similarity index 100% rename from x-pack/plugins/inference/common/output/create_output_api.ts rename to x-pack/plugins/inference/common/create_output_api.ts diff --git a/x-pack/plugins/inference/common/index.ts b/x-pack/plugins/inference/common/index.ts index b84e03e2ad771..f317311d433b1 100644 --- a/x-pack/plugins/inference/common/index.ts +++ b/x-pack/plugins/inference/common/index.ts @@ -10,21 +10,4 @@ export { splitIntoCommands, } from './tasks/nl_to_esql/correct_common_esql_mistakes'; -export { - isChatCompletionChunkEvent, - isChatCompletionMessageEvent, - isChatCompletionEvent, -} from './chat_complete/event_utils'; - -export { isOutputUpdateEvent, isOutputCompleteEvent, isOutputEvent } from './output/event_utils'; - -export { - MessageRole, - type Message, - type ToolMessage, - type AssistantMessage, - type UserMessage, - type ToolSchema, -} from './chat_complete'; - export { generateFakeToolCallId } from './utils/generate_fake_tool_call_id'; diff --git a/x-pack/plugins/inference/common/tasks/nl_to_esql/correct_query_with_actions.ts b/x-pack/plugins/inference/common/tasks/nl_to_esql/correct_query_with_actions.ts index 15b050c3a3897..30e2c11adb6de 100644 --- a/x-pack/plugins/inference/common/tasks/nl_to_esql/correct_query_with_actions.ts +++ b/x-pack/plugins/inference/common/tasks/nl_to_esql/correct_query_with_actions.ts @@ -4,6 +4,7 @@ * 2.0; you may not use this file except in compliance with the Elastic License * 2.0. */ + import { validateQuery, getActions } from '@kbn/esql-validation-autocomplete'; import { getAstAndSyntaxErrors } from '@kbn/esql-ast'; diff --git a/x-pack/plugins/inference/common/ensure_multi_turn.ts b/x-pack/plugins/inference/common/utils/ensure_multi_turn.ts similarity index 92% rename from x-pack/plugins/inference/common/ensure_multi_turn.ts rename to x-pack/plugins/inference/common/utils/ensure_multi_turn.ts index 8d222564f3e72..476ecec108e94 100644 --- a/x-pack/plugins/inference/common/ensure_multi_turn.ts +++ b/x-pack/plugins/inference/common/utils/ensure_multi_turn.ts @@ -5,7 +5,7 @@ * 2.0. */ -import { Message, MessageRole } from './chat_complete'; +import { Message, MessageRole } from '@kbn/inference-common'; function isUserMessage(message: Message): boolean { return message.role !== MessageRole.Assistant; diff --git a/yarn.lock b/yarn.lock index ce86d0b1e6d77..623b0f39bf3a3 100644 --- a/yarn.lock +++ b/yarn.lock @@ -5340,6 +5340,10 @@ version "0.0.0" uid "" +"@kbn/inference-common@link:x-pack/packages/ai-infra/inference-common": + version "0.0.0" + uid "" + "@kbn/inference-plugin@link:x-pack/plugins/inference": version "0.0.0" uid "" From 3e2173eecffe51aec75a134ceedb54edfb59528a Mon Sep 17 00:00:00 2001 From: pgayvallet Date: Wed, 30 Oct 2024 09:18:37 +0100 Subject: [PATCH 05/21] fix imports, first batch --- .../ai-infra/inference-common/index.ts | 12 +++++ .../src/chat_complete/errors.ts | 37 -------------- .../ai-infra/inference-common/src}/errors.ts | 6 +-- x-pack/plugins/inference/common/connectors.ts | 4 -- .../inference/common/create_output_api.ts | 14 +++-- .../inference/common/http_apis.ts} | 9 ++-- x-pack/plugins/inference/common/index.ts | 4 ++ .../index.ts => chat_complete.ts} | 6 +-- x-pack/plugins/inference/public/index.ts | 4 +- x-pack/plugins/inference/public/plugin.tsx | 7 +-- .../create_observable_from_http_response.ts | 3 +- .../http_response_into_observable.test.ts | 8 +-- .../util/http_response_into_observable.ts | 5 +- .../scripts/evaluation/evaluation_client.ts | 3 +- .../evaluation/scenarios/esql/index.spec.ts | 5 +- .../load_esql_docs/utils/output_executor.ts | 2 +- .../inference/scripts/util/cli_options.ts | 1 + .../inference/scripts/util/kibana_client.ts | 17 ++++--- .../bedrock/bedrock_claude_adapter.test.ts | 3 +- .../bedrock/bedrock_claude_adapter.ts | 12 +++-- .../bedrock/process_completion_chunks.ts | 2 +- .../serde_eventstream_into_observable.ts | 2 +- .../adapters/gemini/gemini_adapter.test.ts | 3 +- .../adapters/gemini/gemini_adapter.ts | 11 ++-- .../gemini/process_vertex_stream.test.ts | 2 +- .../adapters/gemini/process_vertex_stream.ts | 2 +- .../adapters/openai/openai_adapter.test.ts | 2 +- .../adapters/openai/openai_adapter.ts | 8 +-- .../inference/server/chat_complete/api.ts | 7 ++- .../inference/server/chat_complete/errors.ts | 51 +++++++++++++++++++ .../get_system_instructions.ts | 2 +- .../parse_inline_function_calls.ts | 4 +- .../wrap_with_simulated_function_calling.ts | 13 +++-- .../inference/server/chat_complete/types.ts | 4 +- .../utils/chunks_into_message.test.ts | 5 +- .../utils/chunks_into_message.ts | 7 +-- x-pack/plugins/inference/server/index.ts | 11 ++-- .../server/inference_client/index.ts | 2 +- .../inference/server/routes/chat_complete.ts | 5 +- .../tasks/nl_to_esql/actions/generate_esql.ts | 16 +++--- .../actions/request_documentation.ts | 3 +- .../actions/summarize_discussion.ts | 4 +- .../server/tasks/nl_to_esql/actions/types.ts | 10 +--- .../inference/server/tasks/nl_to_esql/task.ts | 3 +- .../server/tasks/nl_to_esql/types.ts | 6 +-- x-pack/plugins/inference/server/types.ts | 3 +- .../server/util/get_connector_by_id.ts | 2 +- ...bservable_into_event_source_stream.test.ts | 2 +- .../observable_into_event_source_stream.ts | 4 +- .../server/util/validate_tool_calls.test.ts | 4 +- .../server/util/validate_tool_calls.ts | 7 +-- 51 files changed, 203 insertions(+), 166 deletions(-) rename x-pack/{plugins/inference/common => packages/ai-infra/inference-common/src}/errors.ts (93%) rename x-pack/{packages/ai-infra/inference-common/src/chat_complete/request.ts => plugins/inference/common/http_apis.ts} (66%) rename x-pack/plugins/inference/public/{chat_complete/index.ts => chat_complete.ts} (79%) create mode 100644 x-pack/plugins/inference/server/chat_complete/errors.ts diff --git a/x-pack/packages/ai-infra/inference-common/index.ts b/x-pack/packages/ai-infra/inference-common/index.ts index 2f043c850c92f..8ceaab2e4b78f 100644 --- a/x-pack/packages/ai-infra/inference-common/index.ts +++ b/x-pack/packages/ai-infra/inference-common/index.ts @@ -53,3 +53,15 @@ export { type InferenceTaskEvent, type InferenceTaskEventBase, } from './src/inference_task'; +export { + InferenceTaskError, + InferenceTaskErrorCode, + type InferenceTaskErrorEvent, + type InferenceTaskInternalError, + type InferenceTaskRequestError, + createInferenceInternalError, + createInferenceRequestError, + isInferenceError, + isInferenceInternalError, + isInferenceRequestError, +} from './src/errors'; diff --git a/x-pack/packages/ai-infra/inference-common/src/chat_complete/errors.ts b/x-pack/packages/ai-infra/inference-common/src/chat_complete/errors.ts index 8497350d7b49b..b4c2d6f0a0890 100644 --- a/x-pack/packages/ai-infra/inference-common/src/chat_complete/errors.ts +++ b/x-pack/packages/ai-infra/inference-common/src/chat_complete/errors.ts @@ -5,7 +5,6 @@ * 2.0. */ -import { i18n } from '@kbn/i18n'; import { InferenceTaskError } from '../errors'; import type { UnvalidatedToolCall } from './tools'; @@ -40,42 +39,6 @@ export type ChatCompletionToolValidationError = InferenceTaskError< } >; -export function createTokenLimitReachedError( - tokenLimit?: number, - tokenCount?: number -): ChatCompletionTokenLimitReachedError { - return new InferenceTaskError( - ChatCompletionErrorCode.TokenLimitReachedError, - i18n.translate('xpack.inference.chatCompletionError.tokenLimitReachedError', { - defaultMessage: `Token limit reached. Token limit is {tokenLimit}, but the current conversation has {tokenCount} tokens.`, - values: { tokenLimit, tokenCount }, - }), - { tokenLimit, tokenCount } - ); -} - -export function createToolNotFoundError(name: string): ChatCompletionToolNotFoundError { - return new InferenceTaskError( - ChatCompletionErrorCode.ToolNotFoundError, - `Tool ${name} called but was not available`, - { - name, - } - ); -} - -export function createToolValidationError( - message: string, - meta: { - name?: string; - arguments?: string; - errorsText?: string; - toolCalls?: UnvalidatedToolCall[]; - } -): ChatCompletionToolValidationError { - return new InferenceTaskError(ChatCompletionErrorCode.ToolValidationError, message, meta); -} - export function isToolValidationError(error?: Error): error is ChatCompletionToolValidationError { return ( error instanceof InferenceTaskError && diff --git a/x-pack/plugins/inference/common/errors.ts b/x-pack/packages/ai-infra/inference-common/src/errors.ts similarity index 93% rename from x-pack/plugins/inference/common/errors.ts rename to x-pack/packages/ai-infra/inference-common/src/errors.ts index e8bcd4cf60aaf..acf8d34b4275a 100644 --- a/x-pack/plugins/inference/common/errors.ts +++ b/x-pack/packages/ai-infra/inference-common/src/errors.ts @@ -4,7 +4,7 @@ * 2.0; you may not use this file except in compliance with the Elastic License * 2.0. */ -import { i18n } from '@kbn/i18n'; + import { InferenceTaskEventBase, InferenceTaskEventType } from './inference_task'; export enum InferenceTaskErrorCode { @@ -51,9 +51,7 @@ export type InferenceTaskRequestError = InferenceTaskError< >; export function createInferenceInternalError( - message: string = i18n.translate('xpack.inference.internalError', { - defaultMessage: 'An internal error occurred', - }), + message = 'An internal error occurred', meta?: Record ): InferenceTaskInternalError { return new InferenceTaskError(InferenceTaskErrorCode.internalError, message, meta ?? {}); diff --git a/x-pack/plugins/inference/common/connectors.ts b/x-pack/plugins/inference/common/connectors.ts index f7ad616741d79..ee628f520feff 100644 --- a/x-pack/plugins/inference/common/connectors.ts +++ b/x-pack/plugins/inference/common/connectors.ts @@ -22,7 +22,3 @@ export interface InferenceConnector { export function isSupportedConnectorType(id: string): id is InferenceConnectorType { return allSupportedConnectorTypes.includes(id as InferenceConnectorType); } - -export interface GetConnectorsResponseBody { - connectors: InferenceConnector[]; -} diff --git a/x-pack/plugins/inference/common/create_output_api.ts b/x-pack/plugins/inference/common/create_output_api.ts index d232d15843794..450114c892cba 100644 --- a/x-pack/plugins/inference/common/create_output_api.ts +++ b/x-pack/plugins/inference/common/create_output_api.ts @@ -6,10 +6,16 @@ */ import { map } from 'rxjs'; -import { ChatCompleteAPI, ChatCompletionEventType, MessageRole } from '../chat_complete'; -import { withoutTokenCountEvents } from '../chat_complete/event_utils'; -import { OutputAPI, OutputEvent, OutputEventType } from '.'; -import { ensureMultiTurn } from '../ensure_multi_turn'; +import { + OutputAPI, + OutputEvent, + OutputEventType, + ChatCompleteAPI, + ChatCompletionEventType, + MessageRole, + withoutTokenCountEvents, +} from '@kbn/inference-common'; +import { ensureMultiTurn } from './utils/ensure_multi_turn'; export function createOutputApi(chatCompleteApi: ChatCompleteAPI): OutputAPI { return (id, { connectorId, input, schema, system, previousMessages, functionCalling }) => { diff --git a/x-pack/packages/ai-infra/inference-common/src/chat_complete/request.ts b/x-pack/plugins/inference/common/http_apis.ts similarity index 66% rename from x-pack/packages/ai-infra/inference-common/src/chat_complete/request.ts rename to x-pack/plugins/inference/common/http_apis.ts index 60341bfdaffa5..c07fcd29b2211 100644 --- a/x-pack/packages/ai-infra/inference-common/src/chat_complete/request.ts +++ b/x-pack/plugins/inference/common/http_apis.ts @@ -5,9 +5,8 @@ * 2.0. */ -import type { FunctionCallingMode } from './api'; -import type { Message } from './messages'; -import type { ToolOptions } from './tools'; +import type { FunctionCallingMode, Message, ToolOptions } from '@kbn/inference-common'; +import { InferenceConnector } from './connectors'; export type ChatCompleteRequestBody = { connectorId: string; @@ -16,3 +15,7 @@ export type ChatCompleteRequestBody = { messages: Message[]; functionCalling?: FunctionCallingMode; } & ToolOptions; + +export interface GetConnectorsResponseBody { + connectors: InferenceConnector[]; +} diff --git a/x-pack/plugins/inference/common/index.ts b/x-pack/plugins/inference/common/index.ts index f317311d433b1..19b24d53a389a 100644 --- a/x-pack/plugins/inference/common/index.ts +++ b/x-pack/plugins/inference/common/index.ts @@ -11,3 +11,7 @@ export { } from './tasks/nl_to_esql/correct_common_esql_mistakes'; export { generateFakeToolCallId } from './utils/generate_fake_tool_call_id'; + +export { createOutputApi } from './create_output_api'; + +export type { ChatCompleteRequestBody, GetConnectorsResponseBody } from './http_apis'; diff --git a/x-pack/plugins/inference/public/chat_complete/index.ts b/x-pack/plugins/inference/public/chat_complete.ts similarity index 79% rename from x-pack/plugins/inference/public/chat_complete/index.ts rename to x-pack/plugins/inference/public/chat_complete.ts index e229f6c8f8eae..5319f7c31c381 100644 --- a/x-pack/plugins/inference/public/chat_complete/index.ts +++ b/x-pack/plugins/inference/public/chat_complete.ts @@ -7,9 +7,9 @@ import { from } from 'rxjs'; import type { HttpStart } from '@kbn/core/public'; -import type { ChatCompleteAPI } from '../../common/chat_complete'; -import type { ChatCompleteRequestBody } from '../../common/chat_complete/request'; -import { httpResponseIntoObservable } from '../util/http_response_into_observable'; +import type { ChatCompleteAPI } from '@kbn/inference-common'; +import type { ChatCompleteRequestBody } from '../common/http_apis'; +import { httpResponseIntoObservable } from './util/http_response_into_observable'; export function createChatCompleteApi({ http }: { http: HttpStart }): ChatCompleteAPI { return ({ connectorId, messages, system, toolChoice, tools, functionCalling }) => { diff --git a/x-pack/plugins/inference/public/index.ts b/x-pack/plugins/inference/public/index.ts index 82d36a7abe82d..4928242879b3b 100644 --- a/x-pack/plugins/inference/public/index.ts +++ b/x-pack/plugins/inference/public/index.ts @@ -4,9 +4,8 @@ * 2.0; you may not use this file except in compliance with the Elastic License * 2.0. */ -import type { PluginInitializer, PluginInitializerContext } from '@kbn/core/public'; -import { InferencePlugin } from './plugin'; +import type { PluginInitializer, PluginInitializerContext } from '@kbn/core/public'; import type { InferencePublicSetup, InferencePublicStart, @@ -14,6 +13,7 @@ import type { InferenceStartDependencies, ConfigSchema, } from './types'; +import { InferencePlugin } from './plugin'; export { httpResponseIntoObservable } from './util/http_response_into_observable'; diff --git a/x-pack/plugins/inference/public/plugin.tsx b/x-pack/plugins/inference/public/plugin.tsx index 13ef4a0373845..f1023bc9c2546 100644 --- a/x-pack/plugins/inference/public/plugin.tsx +++ b/x-pack/plugins/inference/public/plugin.tsx @@ -7,8 +7,8 @@ import type { CoreSetup, CoreStart, Plugin, PluginInitializerContext } from '@kbn/core/public'; import type { Logger } from '@kbn/logging'; -import { createOutputApi } from '../common/output/create_output_api'; -import type { GetConnectorsResponseBody } from '../common/connectors'; +import { createOutputApi } from '../common/create_output_api'; +import type { GetConnectorsResponseBody } from '../common/http_apis'; import { createChatCompleteApi } from './chat_complete'; import type { ConfigSchema, @@ -41,10 +41,11 @@ export class InferencePlugin start(coreStart: CoreStart, pluginsStart: InferenceStartDependencies): InferencePublicStart { const chatComplete = createChatCompleteApi({ http: coreStart.http }); + const output = createOutputApi(chatComplete); return { chatComplete, - output: createOutputApi(chatComplete), + output, getConnectors: async () => { const res = await coreStart.http.get( '/internal/inference/connectors' diff --git a/x-pack/plugins/inference/public/util/create_observable_from_http_response.ts b/x-pack/plugins/inference/public/util/create_observable_from_http_response.ts index 09e9b9b2d5f5e..862986ce1c73a 100644 --- a/x-pack/plugins/inference/public/util/create_observable_from_http_response.ts +++ b/x-pack/plugins/inference/public/util/create_observable_from_http_response.ts @@ -4,9 +4,10 @@ * 2.0; you may not use this file except in compliance with the Elastic License * 2.0. */ + import { createParser } from 'eventsource-parser'; import { Observable, throwError } from 'rxjs'; -import { createInferenceInternalError } from '../../common/errors'; +import { createInferenceInternalError } from '@kbn/inference-common'; export interface StreamedHttpResponse { response?: { body: ReadableStream | null | undefined }; diff --git a/x-pack/plugins/inference/public/util/http_response_into_observable.test.ts b/x-pack/plugins/inference/public/util/http_response_into_observable.test.ts index 2b99b6f1db6f5..a0964da025af8 100644 --- a/x-pack/plugins/inference/public/util/http_response_into_observable.test.ts +++ b/x-pack/plugins/inference/public/util/http_response_into_observable.test.ts @@ -7,10 +7,12 @@ import { lastValueFrom, of, toArray } from 'rxjs'; import { httpResponseIntoObservable } from './http_response_into_observable'; +import { + ChatCompletionEventType, + InferenceTaskEventType, + InferenceTaskErrorCode, +} from '@kbn/inference-common'; import type { StreamedHttpResponse } from './create_observable_from_http_response'; -import { ChatCompletionEventType } from '../../common/chat_complete'; -import { InferenceTaskEventType } from '../../common/inference_task'; -import { InferenceTaskErrorCode } from '../../common/errors'; function toSse(...events: Array>) { return events.map((event) => new TextEncoder().encode(`data: ${JSON.stringify(event)}\n\n`)); diff --git a/x-pack/plugins/inference/public/util/http_response_into_observable.ts b/x-pack/plugins/inference/public/util/http_response_into_observable.ts index c63a7bcb3cd15..0aab09cdebe0c 100644 --- a/x-pack/plugins/inference/public/util/http_response_into_observable.ts +++ b/x-pack/plugins/inference/public/util/http_response_into_observable.ts @@ -10,8 +10,9 @@ import { createInferenceInternalError, InferenceTaskError, InferenceTaskErrorEvent, -} from '../../common/errors'; -import { InferenceTaskEvent, InferenceTaskEventType } from '../../common/inference_task'; + InferenceTaskEvent, + InferenceTaskEventType, +} from '@kbn/inference-common'; import { createObservableFromHttpResponse, StreamedHttpResponse, diff --git a/x-pack/plugins/inference/scripts/evaluation/evaluation_client.ts b/x-pack/plugins/inference/scripts/evaluation/evaluation_client.ts index acf2fece1d0ff..d35c214542255 100644 --- a/x-pack/plugins/inference/scripts/evaluation/evaluation_client.ts +++ b/x-pack/plugins/inference/scripts/evaluation/evaluation_client.ts @@ -7,8 +7,7 @@ import { remove } from 'lodash'; import { lastValueFrom } from 'rxjs'; -import type { OutputAPI } from '../../common/output'; -import { withoutOutputUpdateEvents } from '../../common/output/without_output_update_events'; +import { type OutputAPI, withoutOutputUpdateEvents } from '@kbn/inference-common'; import type { EvaluationResult } from './types'; export interface InferenceEvaluationClient { diff --git a/x-pack/plugins/inference/scripts/evaluation/scenarios/esql/index.spec.ts b/x-pack/plugins/inference/scripts/evaluation/scenarios/esql/index.spec.ts index 3aeca67030366..d9071b3f0ae3f 100644 --- a/x-pack/plugins/inference/scripts/evaluation/scenarios/esql/index.spec.ts +++ b/x-pack/plugins/inference/scripts/evaluation/scenarios/esql/index.spec.ts @@ -8,11 +8,12 @@ /// import expect from '@kbn/expect'; +import type { Logger } from '@kbn/logging'; import { firstValueFrom, lastValueFrom, filter } from 'rxjs'; +import { isOutputCompleteEvent } from '@kbn/inference-common'; import { naturalLanguageToEsql } from '../../../../server/tasks/nl_to_esql'; import { chatClient, evaluationClient, logger } from '../../services'; import { EsqlDocumentBase } from '../../../../server/tasks/nl_to_esql/doc_base'; -import { isOutputCompleteEvent } from '../../../../common'; interface TestCase { title: string; @@ -40,7 +41,7 @@ const callNaturalLanguageToEsql = async (question: string) => { debug: (source) => { logger.debug(typeof source === 'function' ? source() : source); }, - }, + } as Logger, }) ); }; diff --git a/x-pack/plugins/inference/scripts/load_esql_docs/utils/output_executor.ts b/x-pack/plugins/inference/scripts/load_esql_docs/utils/output_executor.ts index 6697446f93cec..62cfd8f877e3f 100644 --- a/x-pack/plugins/inference/scripts/load_esql_docs/utils/output_executor.ts +++ b/x-pack/plugins/inference/scripts/load_esql_docs/utils/output_executor.ts @@ -6,7 +6,7 @@ */ import { lastValueFrom } from 'rxjs'; -import type { OutputAPI } from '../../../common/output'; +import type { OutputAPI } from '@kbn/inference-common'; export interface Prompt { system?: string; diff --git a/x-pack/plugins/inference/scripts/util/cli_options.ts b/x-pack/plugins/inference/scripts/util/cli_options.ts index 13bac131922ff..8bbb6dabe406e 100644 --- a/x-pack/plugins/inference/scripts/util/cli_options.ts +++ b/x-pack/plugins/inference/scripts/util/cli_options.ts @@ -4,6 +4,7 @@ * 2.0; you may not use this file except in compliance with the Elastic License * 2.0. */ + import { format, parse } from 'url'; import { readKibanaConfig } from './read_kibana_config'; diff --git a/x-pack/plugins/inference/scripts/util/kibana_client.ts b/x-pack/plugins/inference/scripts/util/kibana_client.ts index ca26ef76b2c72..b599ab81a4af4 100644 --- a/x-pack/plugins/inference/scripts/util/kibana_client.ts +++ b/x-pack/plugins/inference/scripts/util/kibana_client.ts @@ -13,18 +13,19 @@ import { from, map, switchMap, throwError } from 'rxjs'; import { UrlObject, format, parse } from 'url'; import { inspect } from 'util'; import { isReadable } from 'stream'; -import type { ChatCompleteAPI, ChatCompletionEvent } from '../../common/chat_complete'; -import { ChatCompleteRequestBody } from '../../common/chat_complete/request'; -import type { InferenceConnector } from '../../common/connectors'; import { + ChatCompleteAPI, + OutputAPI, + ChatCompletionEvent, InferenceTaskError, InferenceTaskErrorEvent, + InferenceTaskEventType, createInferenceInternalError, -} from '../../common/errors'; -import { InferenceTaskEventType } from '../../common/inference_task'; -import type { OutputAPI } from '../../common/output'; -import { createOutputApi } from '../../common/output/create_output_api'; -import { withoutOutputUpdateEvents } from '../../common/output/without_output_update_events'; + withoutOutputUpdateEvents, +} from '@kbn/inference-common'; +import type { ChatCompleteRequestBody } from '../../common/http_apis'; +import type { InferenceConnector } from '../../common/connectors'; +import { createOutputApi } from '../../common/create_output_api'; import { eventSourceStreamIntoObservable } from '../../server/util/event_source_stream_into_observable'; // eslint-disable-next-line spaced-comment diff --git a/x-pack/plugins/inference/server/chat_complete/adapters/bedrock/bedrock_claude_adapter.test.ts b/x-pack/plugins/inference/server/chat_complete/adapters/bedrock/bedrock_claude_adapter.test.ts index d34b8693cb85f..ca6f60dd45a55 100644 --- a/x-pack/plugins/inference/server/chat_complete/adapters/bedrock/bedrock_claude_adapter.test.ts +++ b/x-pack/plugins/inference/server/chat_complete/adapters/bedrock/bedrock_claude_adapter.test.ts @@ -8,8 +8,7 @@ import { PassThrough } from 'stream'; import { loggerMock } from '@kbn/logging-mocks'; import type { InferenceExecutor } from '../../utils/inference_executor'; -import { MessageRole } from '../../../../common/chat_complete'; -import { ToolChoiceType } from '../../../../common/chat_complete/tools'; +import { MessageRole, ToolChoiceType } from '@kbn/inference-common'; import { bedrockClaudeAdapter } from './bedrock_claude_adapter'; import { addNoToolUsageDirective } from './prompts'; diff --git a/x-pack/plugins/inference/server/chat_complete/adapters/bedrock/bedrock_claude_adapter.ts b/x-pack/plugins/inference/server/chat_complete/adapters/bedrock/bedrock_claude_adapter.ts index a0b48e6fc8631..e73d9c9344c98 100644 --- a/x-pack/plugins/inference/server/chat_complete/adapters/bedrock/bedrock_claude_adapter.ts +++ b/x-pack/plugins/inference/server/chat_complete/adapters/bedrock/bedrock_claude_adapter.ts @@ -7,10 +7,15 @@ import { filter, from, map, switchMap, tap } from 'rxjs'; import { Readable } from 'stream'; +import { + Message, + MessageRole, + createInferenceInternalError, + ToolChoiceType, + ToolSchemaType, + type ToolOptions, +} from '@kbn/inference-common'; import { parseSerdeChunkMessage } from './serde_utils'; -import { Message, MessageRole } from '../../../../common/chat_complete'; -import { createInferenceInternalError } from '../../../../common/errors'; -import { ToolChoiceType, type ToolOptions } from '../../../../common/chat_complete/tools'; import { InferenceConnectorAdapter } from '../../types'; import type { BedRockMessage, BedrockToolChoice } from './types'; import { @@ -19,7 +24,6 @@ import { } from './serde_eventstream_into_observable'; import { processCompletionChunks } from './process_completion_chunks'; import { addNoToolUsageDirective } from './prompts'; -import { ToolSchemaType } from '../../../../common/chat_complete/tool_schema'; export const bedrockClaudeAdapter: InferenceConnectorAdapter = { chatComplete: ({ executor, system, messages, toolChoice, tools }) => { diff --git a/x-pack/plugins/inference/server/chat_complete/adapters/bedrock/process_completion_chunks.ts b/x-pack/plugins/inference/server/chat_complete/adapters/bedrock/process_completion_chunks.ts index 5513cc9028ac9..8a5c9805ddf63 100644 --- a/x-pack/plugins/inference/server/chat_complete/adapters/bedrock/process_completion_chunks.ts +++ b/x-pack/plugins/inference/server/chat_complete/adapters/bedrock/process_completion_chunks.ts @@ -11,7 +11,7 @@ import { ChatCompletionTokenCountEvent, ChatCompletionChunkToolCall, ChatCompletionEventType, -} from '../../../../common/chat_complete'; +} from '@kbn/inference-common'; import type { CompletionChunk, MessageStopChunk } from './types'; export function processCompletionChunks() { diff --git a/x-pack/plugins/inference/server/chat_complete/adapters/bedrock/serde_eventstream_into_observable.ts b/x-pack/plugins/inference/server/chat_complete/adapters/bedrock/serde_eventstream_into_observable.ts index 24a245ab2efcc..5ab264750e5a9 100644 --- a/x-pack/plugins/inference/server/chat_complete/adapters/bedrock/serde_eventstream_into_observable.ts +++ b/x-pack/plugins/inference/server/chat_complete/adapters/bedrock/serde_eventstream_into_observable.ts @@ -11,7 +11,7 @@ import { identity } from 'lodash'; import { Observable } from 'rxjs'; import { Readable } from 'stream'; import { Message } from '@smithy/types'; -import { createInferenceInternalError } from '../../../../common/errors'; +import { createInferenceInternalError } from '@kbn/inference-common'; interface ModelStreamErrorException { name: 'ModelStreamErrorException'; diff --git a/x-pack/plugins/inference/server/chat_complete/adapters/gemini/gemini_adapter.test.ts b/x-pack/plugins/inference/server/chat_complete/adapters/gemini/gemini_adapter.test.ts index a9f4305a3c532..c3410b2af3623 100644 --- a/x-pack/plugins/inference/server/chat_complete/adapters/gemini/gemini_adapter.test.ts +++ b/x-pack/plugins/inference/server/chat_complete/adapters/gemini/gemini_adapter.test.ts @@ -11,8 +11,7 @@ import { noop, tap, lastValueFrom, toArray, Subject } from 'rxjs'; import { loggerMock } from '@kbn/logging-mocks'; import type { InferenceExecutor } from '../../utils/inference_executor'; import { observableIntoEventSourceStream } from '../../../util/observable_into_event_source_stream'; -import { MessageRole } from '../../../../common/chat_complete'; -import { ToolChoiceType } from '../../../../common/chat_complete/tools'; +import { MessageRole, ToolChoiceType } from '@kbn/inference-common'; import { geminiAdapter } from './gemini_adapter'; describe('geminiAdapter', () => { diff --git a/x-pack/plugins/inference/server/chat_complete/adapters/gemini/gemini_adapter.ts b/x-pack/plugins/inference/server/chat_complete/adapters/gemini/gemini_adapter.ts index 2e86adcc82a85..80d0439449066 100644 --- a/x-pack/plugins/inference/server/chat_complete/adapters/gemini/gemini_adapter.ts +++ b/x-pack/plugins/inference/server/chat_complete/adapters/gemini/gemini_adapter.ts @@ -8,10 +8,15 @@ import * as Gemini from '@google/generative-ai'; import { from, map, switchMap } from 'rxjs'; import { Readable } from 'stream'; +import { + Message, + MessageRole, + ToolChoiceType, + ToolOptions, + ToolSchema, + ToolSchemaType, +} from '@kbn/inference-common'; import type { InferenceConnectorAdapter } from '../../types'; -import { Message, MessageRole } from '../../../../common/chat_complete'; -import { ToolChoiceType, ToolOptions } from '../../../../common/chat_complete/tools'; -import type { ToolSchema, ToolSchemaType } from '../../../../common/chat_complete/tool_schema'; import { eventSourceStreamIntoObservable } from '../../../util/event_source_stream_into_observable'; import { processVertexStream } from './process_vertex_stream'; import type { GenerateContentResponseChunk, GeminiMessage, GeminiToolConfig } from './types'; diff --git a/x-pack/plugins/inference/server/chat_complete/adapters/gemini/process_vertex_stream.test.ts b/x-pack/plugins/inference/server/chat_complete/adapters/gemini/process_vertex_stream.test.ts index 78e0da0a384b8..8613799846e3b 100644 --- a/x-pack/plugins/inference/server/chat_complete/adapters/gemini/process_vertex_stream.test.ts +++ b/x-pack/plugins/inference/server/chat_complete/adapters/gemini/process_vertex_stream.test.ts @@ -6,7 +6,7 @@ */ import { TestScheduler } from 'rxjs/testing'; -import { ChatCompletionEventType } from '../../../../common/chat_complete'; +import { ChatCompletionEventType } from '@kbn/inference-common'; import { processVertexStream } from './process_vertex_stream'; import type { GenerateContentResponseChunk } from './types'; diff --git a/x-pack/plugins/inference/server/chat_complete/adapters/gemini/process_vertex_stream.ts b/x-pack/plugins/inference/server/chat_complete/adapters/gemini/process_vertex_stream.ts index e2a6c74a0447f..3081317882c65 100644 --- a/x-pack/plugins/inference/server/chat_complete/adapters/gemini/process_vertex_stream.ts +++ b/x-pack/plugins/inference/server/chat_complete/adapters/gemini/process_vertex_stream.ts @@ -10,7 +10,7 @@ import { ChatCompletionChunkEvent, ChatCompletionTokenCountEvent, ChatCompletionEventType, -} from '../../../../common/chat_complete'; +} from '@kbn/inference-common'; import { generateFakeToolCallId } from '../../../../common'; import type { GenerateContentResponseChunk } from './types'; diff --git a/x-pack/plugins/inference/server/chat_complete/adapters/openai/openai_adapter.test.ts b/x-pack/plugins/inference/server/chat_complete/adapters/openai/openai_adapter.test.ts index 813e88760de8c..ff1bbc71a876d 100644 --- a/x-pack/plugins/inference/server/chat_complete/adapters/openai/openai_adapter.test.ts +++ b/x-pack/plugins/inference/server/chat_complete/adapters/openai/openai_adapter.test.ts @@ -12,7 +12,7 @@ import { pick } from 'lodash'; import { lastValueFrom, Subject, toArray } from 'rxjs'; import type { Logger } from '@kbn/logging'; import { loggerMock } from '@kbn/logging-mocks'; -import { ChatCompletionEventType, MessageRole } from '../../../../common/chat_complete'; +import { ChatCompletionEventType, MessageRole } from '@kbn/inference-common'; import { observableIntoEventSourceStream } from '../../../util/observable_into_event_source_stream'; import { InferenceExecutor } from '../../utils/inference_executor'; import { openAIAdapter } from '.'; diff --git a/x-pack/plugins/inference/server/chat_complete/adapters/openai/openai_adapter.ts b/x-pack/plugins/inference/server/chat_complete/adapters/openai/openai_adapter.ts index f1821be4d4d57..121ba96ab115a 100644 --- a/x-pack/plugins/inference/server/chat_complete/adapters/openai/openai_adapter.ts +++ b/x-pack/plugins/inference/server/chat_complete/adapters/openai/openai_adapter.ts @@ -20,10 +20,10 @@ import { ChatCompletionEventType, Message, MessageRole, -} from '../../../../common/chat_complete'; -import type { ToolOptions } from '../../../../common/chat_complete/tools'; -import { createTokenLimitReachedError } from '../../../../common/chat_complete/errors'; -import { createInferenceInternalError } from '../../../../common/errors'; + ToolOptions, + createInferenceInternalError, +} from '@kbn/inference-common'; +import { createTokenLimitReachedError } from '../../errors'; import { eventSourceStreamIntoObservable } from '../../../util/event_source_stream_into_observable'; import type { InferenceConnectorAdapter } from '../../types'; import { diff --git a/x-pack/plugins/inference/server/chat_complete/api.ts b/x-pack/plugins/inference/server/chat_complete/api.ts index ca9e61ff3627f..62a1ea8b26146 100644 --- a/x-pack/plugins/inference/server/chat_complete/api.ts +++ b/x-pack/plugins/inference/server/chat_complete/api.ts @@ -9,8 +9,11 @@ import { last } from 'lodash'; import { defer, switchMap, throwError } from 'rxjs'; import type { Logger } from '@kbn/logging'; import type { KibanaRequest } from '@kbn/core-http-server'; -import type { ChatCompleteAPI, ChatCompletionResponse } from '../../common/chat_complete'; -import { createInferenceRequestError } from '../../common/errors'; +import { + type ChatCompleteAPI, + type ChatCompletionResponse, + createInferenceRequestError, +} from '@kbn/inference-common'; import type { InferenceStartDependencies } from '../types'; import { getConnectorById } from '../util/get_connector_by_id'; import { getInferenceAdapter } from './adapters'; diff --git a/x-pack/plugins/inference/server/chat_complete/errors.ts b/x-pack/plugins/inference/server/chat_complete/errors.ts new file mode 100644 index 0000000000000..a830f57fec559 --- /dev/null +++ b/x-pack/plugins/inference/server/chat_complete/errors.ts @@ -0,0 +1,51 @@ +/* + * 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 { InferenceTaskError, type UnvalidatedToolCall } from '@kbn/inference-common'; +import { i18n } from '@kbn/i18n'; +import { + ChatCompletionErrorCode, + ChatCompletionTokenLimitReachedError, + ChatCompletionToolNotFoundError, + ChatCompletionToolValidationError, +} from '@kbn/inference-common/src/chat_complete/errors'; + +export function createTokenLimitReachedError( + tokenLimit?: number, + tokenCount?: number +): ChatCompletionTokenLimitReachedError { + return new InferenceTaskError( + ChatCompletionErrorCode.TokenLimitReachedError, + i18n.translate('xpack.inference.chatCompletionError.tokenLimitReachedError', { + defaultMessage: `Token limit reached. Token limit is {tokenLimit}, but the current conversation has {tokenCount} tokens.`, + values: { tokenLimit, tokenCount }, + }), + { tokenLimit, tokenCount } + ); +} + +export function createToolNotFoundError(name: string): ChatCompletionToolNotFoundError { + return new InferenceTaskError( + ChatCompletionErrorCode.ToolNotFoundError, + `Tool ${name} called but was not available`, + { + name, + } + ); +} + +export function createToolValidationError( + message: string, + meta: { + name?: string; + arguments?: string; + errorsText?: string; + toolCalls?: UnvalidatedToolCall[]; + } +): ChatCompletionToolValidationError { + return new InferenceTaskError(ChatCompletionErrorCode.ToolValidationError, message, meta); +} diff --git a/x-pack/plugins/inference/server/chat_complete/simulated_function_calling/get_system_instructions.ts b/x-pack/plugins/inference/server/chat_complete/simulated_function_calling/get_system_instructions.ts index abfc48dfa2ef2..c4adfae7e3f19 100644 --- a/x-pack/plugins/inference/server/chat_complete/simulated_function_calling/get_system_instructions.ts +++ b/x-pack/plugins/inference/server/chat_complete/simulated_function_calling/get_system_instructions.ts @@ -5,8 +5,8 @@ * 2.0. */ +import { ToolDefinition } from '@kbn/inference-common'; import { TOOL_USE_END, TOOL_USE_START } from './constants'; -import { ToolDefinition } from '../../../common/chat_complete/tools'; export function getSystemMessageInstructions({ tools, diff --git a/x-pack/plugins/inference/server/chat_complete/simulated_function_calling/parse_inline_function_calls.ts b/x-pack/plugins/inference/server/chat_complete/simulated_function_calling/parse_inline_function_calls.ts index 3436d7a7edac5..73d03ee2f00af 100644 --- a/x-pack/plugins/inference/server/chat_complete/simulated_function_calling/parse_inline_function_calls.ts +++ b/x-pack/plugins/inference/server/chat_complete/simulated_function_calling/parse_inline_function_calls.ts @@ -8,11 +8,11 @@ import { Observable } from 'rxjs'; import { Logger } from '@kbn/logging'; import { + createInferenceInternalError, ChatCompletionChunkEvent, ChatCompletionTokenCountEvent, ChatCompletionEventType, -} from '../../../common/chat_complete'; -import { createInferenceInternalError } from '../../../common/errors'; +} from '@kbn/inference-common'; import { TOOL_USE_END, TOOL_USE_START } from './constants'; function matchOnSignalStart(buffer: string) { diff --git a/x-pack/plugins/inference/server/chat_complete/simulated_function_calling/wrap_with_simulated_function_calling.ts b/x-pack/plugins/inference/server/chat_complete/simulated_function_calling/wrap_with_simulated_function_calling.ts index d8cfc373b66cc..4eb6cfd8d50e1 100644 --- a/x-pack/plugins/inference/server/chat_complete/simulated_function_calling/wrap_with_simulated_function_calling.ts +++ b/x-pack/plugins/inference/server/chat_complete/simulated_function_calling/wrap_with_simulated_function_calling.ts @@ -5,9 +5,16 @@ * 2.0. */ -import { AssistantMessage, Message, ToolMessage, UserMessage } from '../../../common'; -import { MessageRole } from '../../../common/chat_complete'; -import { ToolChoice, ToolChoiceType, ToolDefinition } from '../../../common/chat_complete/tools'; +import { + MessageRole, + AssistantMessage, + Message, + ToolMessage, + UserMessage, + ToolChoice, + ToolChoiceType, + ToolDefinition, +} from '@kbn/inference-common'; import { TOOL_USE_END, TOOL_USE_START } from './constants'; import { getSystemMessageInstructions } from './get_system_instructions'; diff --git a/x-pack/plugins/inference/server/chat_complete/types.ts b/x-pack/plugins/inference/server/chat_complete/types.ts index 394fe370240ef..64cc542ff6119 100644 --- a/x-pack/plugins/inference/server/chat_complete/types.ts +++ b/x-pack/plugins/inference/server/chat_complete/types.ts @@ -12,8 +12,8 @@ import type { ChatCompletionTokenCountEvent, FunctionCallingMode, Message, -} from '../../common/chat_complete'; -import type { ToolOptions } from '../../common/chat_complete/tools'; + ToolOptions, +} from '@kbn/inference-common'; import type { InferenceExecutor } from './utils'; /** diff --git a/x-pack/plugins/inference/server/chat_complete/utils/chunks_into_message.test.ts b/x-pack/plugins/inference/server/chat_complete/utils/chunks_into_message.test.ts index 0c5552a0113b8..c6e5b032120a3 100644 --- a/x-pack/plugins/inference/server/chat_complete/utils/chunks_into_message.test.ts +++ b/x-pack/plugins/inference/server/chat_complete/utils/chunks_into_message.test.ts @@ -4,13 +4,14 @@ * 2.0; you may not use this file except in compliance with the Elastic License * 2.0. */ + import { lastValueFrom, of } from 'rxjs'; import { + ToolChoiceType, ChatCompletionChunkEvent, ChatCompletionEventType, ChatCompletionTokenCountEvent, -} from '../../../common/chat_complete'; -import { ToolChoiceType } from '../../../common/chat_complete/tools'; +} from '@kbn/inference-common'; import { chunksIntoMessage } from './chunks_into_message'; import type { Logger } from '@kbn/logging'; diff --git a/x-pack/plugins/inference/server/chat_complete/utils/chunks_into_message.ts b/x-pack/plugins/inference/server/chat_complete/utils/chunks_into_message.ts index 902289182a37a..fe9b745f442fc 100644 --- a/x-pack/plugins/inference/server/chat_complete/utils/chunks_into_message.ts +++ b/x-pack/plugins/inference/server/chat_complete/utils/chunks_into_message.ts @@ -7,14 +7,15 @@ import { last, map, merge, OperatorFunction, scan, share } from 'rxjs'; import type { Logger } from '@kbn/logging'; -import type { UnvalidatedToolCall, ToolOptions } from '../../../common/chat_complete/tools'; import { + UnvalidatedToolCall, + ToolOptions, ChatCompletionChunkEvent, ChatCompletionEventType, ChatCompletionMessageEvent, ChatCompletionTokenCountEvent, -} from '../../../common/chat_complete'; -import { withoutTokenCountEvents } from '../../../common/chat_complete/without_token_count_events'; + withoutTokenCountEvents, +} from '@kbn/inference-common'; import { validateToolCalls } from '../../util/validate_tool_calls'; export function chunksIntoMessage({ diff --git a/x-pack/plugins/inference/server/index.ts b/x-pack/plugins/inference/server/index.ts index d02dfec733941..60ce870020feb 100644 --- a/x-pack/plugins/inference/server/index.ts +++ b/x-pack/plugins/inference/server/index.ts @@ -4,25 +4,22 @@ * 2.0; you may not use this file except in compliance with the Elastic License * 2.0. */ + import type { PluginInitializer, PluginInitializerContext } from '@kbn/core/server'; import type { InferenceConfig } from './config'; -import { InferencePlugin } from './plugin'; import type { InferenceServerSetup, InferenceServerStart, InferenceSetupDependencies, InferenceStartDependencies, } from './types'; - -export { withoutTokenCountEvents } from '../common/chat_complete/without_token_count_events'; -export { withoutChunkEvents } from '../common/chat_complete/without_chunk_events'; -export { withoutOutputUpdateEvents } from '../common/output/without_output_update_events'; +import { InferencePlugin } from './plugin'; export type { InferenceClient } from './types'; -export { naturalLanguageToEsql } from './tasks/nl_to_esql'; - export type { InferenceServerSetup, InferenceServerStart }; +export { naturalLanguageToEsql } from './tasks/nl_to_esql'; + export const plugin: PluginInitializer< InferenceServerSetup, InferenceServerStart, diff --git a/x-pack/plugins/inference/server/inference_client/index.ts b/x-pack/plugins/inference/server/inference_client/index.ts index 25208bebc54bb..03da0e3da200f 100644 --- a/x-pack/plugins/inference/server/inference_client/index.ts +++ b/x-pack/plugins/inference/server/inference_client/index.ts @@ -9,7 +9,7 @@ import type { Logger } from '@kbn/logging'; import type { KibanaRequest } from '@kbn/core-http-server'; import type { InferenceClient, InferenceStartDependencies } from '../types'; import { createChatCompleteApi } from '../chat_complete'; -import { createOutputApi } from '../../common/output/create_output_api'; +import { createOutputApi } from '../../common/create_output_api'; import { getConnectorById } from '../util/get_connector_by_id'; export function createInferenceClient({ diff --git a/x-pack/plugins/inference/server/routes/chat_complete.ts b/x-pack/plugins/inference/server/routes/chat_complete.ts index fdf33fbf0af82..d4d0d012a78cd 100644 --- a/x-pack/plugins/inference/server/routes/chat_complete.ts +++ b/x-pack/plugins/inference/server/routes/chat_complete.ts @@ -7,9 +7,8 @@ import { schema, Type } from '@kbn/config-schema'; import type { CoreSetup, IRouter, Logger, RequestHandlerContext } from '@kbn/core/server'; -import { MessageRole } from '../../common/chat_complete'; -import type { ChatCompleteRequestBody } from '../../common/chat_complete/request'; -import { ToolCall, ToolChoiceType } from '../../common/chat_complete/tools'; +import { MessageRole, ToolCall, ToolChoiceType } from '@kbn/inference-common'; +import type { ChatCompleteRequestBody } from '../../common/http_apis'; import { createInferenceClient } from '../inference_client'; import { InferenceServerStart, InferenceStartDependencies } from '../types'; import { observableIntoEventSourceStream } from '../util/observable_into_event_source_stream'; diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/generate_esql.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/generate_esql.ts index 47d40a3ce4e32..3ee4297a13c41 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/generate_esql.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/generate_esql.ts @@ -5,21 +5,19 @@ * 2.0. */ -import { Observable, map, merge, of, switchMap, tap, lastValueFrom } from 'rxjs'; +import { map, tap, lastValueFrom } from 'rxjs'; import type { Logger } from '@kbn/logging'; -import { ToolCall, ToolOptions } from '../../../../common/chat_complete/tools'; import { - correctCommonEsqlMistakes, - generateFakeToolCallId, + ToolCall, + ToolOptions, isChatCompletionMessageEvent, + withoutTokenCountEvents, + withoutChunkEvents, Message, MessageRole, -} from '../../../../common'; -import { InferenceClient, withoutTokenCountEvents, withoutChunkEvents } from '../../..'; -import { OutputCompleteEvent, OutputEventType } from '../../../../common/output'; +} from '@kbn/inference-common'; +import { correctCommonEsqlMistakes, generateFakeToolCallId } from '../../../../common'; import { INLINE_ESQL_QUERY_REGEX } from '../../../../common/tasks/nl_to_esql/constants'; -import { EsqlDocumentBase } from '../doc_base'; -import type { NlToEsqlTaskEvent } from '../types'; import type { ActionsOptionsBase } from './types'; export const generateEsqlTaskFn = ({ diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/request_documentation.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/request_documentation.ts index b98bfb8f3fae6..c8eaed176da73 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/request_documentation.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/request_documentation.ts @@ -6,8 +6,7 @@ */ import { lastValueFrom, tap, map } from 'rxjs'; -import { withoutOutputUpdateEvents } from '../../..'; -import { Message, ToolSchema } from '../../../../common'; +import { withoutOutputUpdateEvents, Message, ToolSchema } from '@kbn/inference-common'; import { EsqlDocumentBase } from '../doc_base'; import type { ActionsOptionsBase } from './types'; diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/summarize_discussion.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/summarize_discussion.ts index bc89e3beeedb4..a5625c4423232 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/summarize_discussion.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/summarize_discussion.ts @@ -6,9 +6,7 @@ */ import { lastValueFrom, tap } from 'rxjs'; -import { withoutOutputUpdateEvents } from '../../..'; -import { Message } from '../../../../common'; -import { ToolSchema } from '../../../../common'; +import { Message, ToolSchema, withoutOutputUpdateEvents } from '@kbn/inference-common'; import type { ActionsOptionsBase } from './types'; const summarizeDiscussionSchema = { diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/types.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/types.ts index dad0c836212b5..dc198d6b44686 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/types.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/types.ts @@ -7,8 +7,8 @@ import type { Logger } from '@kbn/logging'; import { Subject } from 'rxjs'; +import type { FunctionCallingMode } from '@kbn/inference-common'; import type { InferenceClient } from '../../../types'; -import type { FunctionCallingMode } from '../../../../common/chat_complete'; import { NlToEsqlTaskEvent } from '../types'; export interface ActionsOptionsBase { @@ -18,11 +18,3 @@ export interface ActionsOptionsBase { functionCalling?: FunctionCallingMode; output$: Subject>; } - -/* - client, - connectorId, - functionCalling, - logger, - subject, - */ diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/task.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/task.ts index f9f21e9825230..c0ce44d77971c 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/task.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/task.ts @@ -7,8 +7,7 @@ import { once } from 'lodash'; import { Observable, Subject } from 'rxjs'; -import { Message, MessageRole } from '../../../common/chat_complete'; -import type { ToolOptions } from '../../../common/chat_complete/tools'; +import { MessageRole, type Message, type ToolOptions } from '@kbn/inference-common'; import { EsqlDocumentBase } from './doc_base'; import { type ActionsOptionsBase, diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/types.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/types.ts index 1fa300daaba00..ca3b44fd18c4f 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/types.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/types.ts @@ -11,9 +11,9 @@ import type { ChatCompletionMessageEvent, FunctionCallingMode, Message, -} from '../../../common/chat_complete'; -import type { ToolOptions } from '../../../common/chat_complete/tools'; -import type { OutputCompleteEvent } from '../../../common/output'; + ToolOptions, + OutputCompleteEvent, +} from '@kbn/inference-common'; import type { InferenceClient } from '../../types'; export type NlToEsqlTaskEvent = diff --git a/x-pack/plugins/inference/server/types.ts b/x-pack/plugins/inference/server/types.ts index 20679ffd4cedf..f538448372e36 100644 --- a/x-pack/plugins/inference/server/types.ts +++ b/x-pack/plugins/inference/server/types.ts @@ -10,9 +10,8 @@ import type { PluginSetupContract as ActionsPluginSetup, } from '@kbn/actions-plugin/server'; import type { KibanaRequest } from '@kbn/core-http-server'; -import { ChatCompleteAPI } from '../common/chat_complete'; +import { ChatCompleteAPI, OutputAPI } from '@kbn/inference-common'; import { InferenceConnector } from '../common/connectors'; -import { OutputAPI } from '../common/output'; /* eslint-disable @typescript-eslint/no-empty-interface*/ diff --git a/x-pack/plugins/inference/server/util/get_connector_by_id.ts b/x-pack/plugins/inference/server/util/get_connector_by_id.ts index 3fd77630ad3d1..1dbf9a6f0d75e 100644 --- a/x-pack/plugins/inference/server/util/get_connector_by_id.ts +++ b/x-pack/plugins/inference/server/util/get_connector_by_id.ts @@ -6,8 +6,8 @@ */ import type { ActionsClient, ActionResult as ActionConnector } from '@kbn/actions-plugin/server'; +import { createInferenceRequestError } from '@kbn/inference-common'; import { isSupportedConnectorType, type InferenceConnector } from '../../common/connectors'; -import { createInferenceRequestError } from '../../common/errors'; /** * Retrieves a connector given the provided `connectorId` and asserts it's an inference connector diff --git a/x-pack/plugins/inference/server/util/observable_into_event_source_stream.test.ts b/x-pack/plugins/inference/server/util/observable_into_event_source_stream.test.ts index ed5466ba1e027..8ece214c27599 100644 --- a/x-pack/plugins/inference/server/util/observable_into_event_source_stream.test.ts +++ b/x-pack/plugins/inference/server/util/observable_into_event_source_stream.test.ts @@ -8,7 +8,7 @@ import { createParser } from 'eventsource-parser'; import { partition } from 'lodash'; import { merge, of, throwError } from 'rxjs'; -import type { InferenceTaskEvent } from '../../common/inference_task'; +import type { InferenceTaskEvent } from '@kbn/inference-common'; import { observableIntoEventSourceStream } from './observable_into_event_source_stream'; import type { Logger } from '@kbn/logging'; diff --git a/x-pack/plugins/inference/server/util/observable_into_event_source_stream.ts b/x-pack/plugins/inference/server/util/observable_into_event_source_stream.ts index bcd1ef60ce1da..62eae6609441f 100644 --- a/x-pack/plugins/inference/server/util/observable_into_event_source_stream.ts +++ b/x-pack/plugins/inference/server/util/observable_into_event_source_stream.ts @@ -9,11 +9,11 @@ import { catchError, map, Observable, of } from 'rxjs'; import { PassThrough } from 'stream'; import type { Logger } from '@kbn/logging'; import { + InferenceTaskEventType, InferenceTaskErrorCode, InferenceTaskErrorEvent, isInferenceError, -} from '../../common/errors'; -import { InferenceTaskEventType } from '../../common/inference_task'; +} from '@kbn/inference-common'; export function observableIntoEventSourceStream( source$: Observable, diff --git a/x-pack/plugins/inference/server/util/validate_tool_calls.test.ts b/x-pack/plugins/inference/server/util/validate_tool_calls.test.ts index 96bf202fa236b..2a6aed608dced 100644 --- a/x-pack/plugins/inference/server/util/validate_tool_calls.test.ts +++ b/x-pack/plugins/inference/server/util/validate_tool_calls.test.ts @@ -5,8 +5,8 @@ * 2.0. */ -import { isToolValidationError } from '../../common/chat_complete/errors'; -import { ToolChoiceType } from '../../common/chat_complete/tools'; +import { ToolChoiceType } from '@kbn/inference-common'; +import { isToolValidationError } from '@kbn/inference-common/src/chat_complete/errors'; import { validateToolCalls } from './validate_tool_calls'; describe('validateToolCalls', () => { diff --git a/x-pack/plugins/inference/server/util/validate_tool_calls.ts b/x-pack/plugins/inference/server/util/validate_tool_calls.ts index 5d1e659bc36f5..ffc2482774b23 100644 --- a/x-pack/plugins/inference/server/util/validate_tool_calls.ts +++ b/x-pack/plugins/inference/server/util/validate_tool_calls.ts @@ -5,16 +5,13 @@ * 2.0. */ import Ajv from 'ajv'; -import { - createToolNotFoundError, - createToolValidationError, -} from '../../common/chat_complete/errors'; import { ToolCallsOf, ToolChoiceType, ToolOptions, UnvalidatedToolCall, -} from '../../common/chat_complete/tools'; +} from '@kbn/inference-common'; +import { createToolNotFoundError, createToolValidationError } from '../chat_complete/errors'; export function validateToolCalls({ toolCalls, From 4607c49844cf6c26b4abb3c94cd242d56c4c242b Mon Sep 17 00:00:00 2001 From: pgayvallet Date: Wed, 30 Oct 2024 09:24:27 +0100 Subject: [PATCH 06/21] add kbn reference --- x-pack/plugins/inference/tsconfig.json | 1 + 1 file changed, 1 insertion(+) diff --git a/x-pack/plugins/inference/tsconfig.json b/x-pack/plugins/inference/tsconfig.json index cc81eec1da96c..4e4a91d15e57b 100644 --- a/x-pack/plugins/inference/tsconfig.json +++ b/x-pack/plugins/inference/tsconfig.json @@ -36,5 +36,6 @@ "@kbn/es-types", "@kbn/field-types", "@kbn/expressions-plugin", + "@kbn/inference-common", ] } From 68f3e72e7b8c500f527e99aa909457262a15f409 Mon Sep 17 00:00:00 2001 From: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Date: Wed, 30 Oct 2024 08:35:11 +0000 Subject: [PATCH 07/21] [CI] Auto-commit changed files from 'node scripts/generate codeowners' --- .github/CODEOWNERS | 1 + x-pack/packages/ai-infra/inference-common/package.json | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 1451c647f658e..c27a0ef0776f5 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -525,6 +525,7 @@ x-pack/plugins/index_management @elastic/kibana-management x-pack/packages/index-management/index_management_shared_types @elastic/kibana-management test/plugin_functional/plugins/index_patterns @elastic/kibana-data-discovery x-pack/packages/ml/inference_integration_flyout @elastic/ml-ui +x-pack/packages/ai-infra/inference-common @elastic/appex-ai-infra x-pack/plugins/inference @elastic/appex-ai-infra x-pack/packages/kbn-infra-forge @elastic/obs-ux-management-team x-pack/plugins/observability_solution/infra @elastic/obs-ux-logs-team @elastic/obs-ux-infra_services-team diff --git a/x-pack/packages/ai-infra/inference-common/package.json b/x-pack/packages/ai-infra/inference-common/package.json index 8deaf0a5504f0..0c67ca7815f16 100644 --- a/x-pack/packages/ai-infra/inference-common/package.json +++ b/x-pack/packages/ai-infra/inference-common/package.json @@ -2,6 +2,6 @@ "name": "@kbn/inference-common", "private": true, "version": "1.0.0", - "license": "Elastic License 2.0 OR AGPL-3.0-only OR SSPL-1.0", + "license": "Elastic License 2.0", "sideEffects": false } From 7d5ec2571148a44e10d4820cdca6ae1d1509368d Mon Sep 17 00:00:00 2001 From: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Date: Wed, 30 Oct 2024 08:53:01 +0000 Subject: [PATCH 08/21] [CI] Auto-commit changed files from 'node scripts/notice' --- x-pack/packages/ai-infra/inference-common/tsconfig.json | 4 +++- x-pack/plugins/inference/tsconfig.json | 4 ---- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/x-pack/packages/ai-infra/inference-common/tsconfig.json b/x-pack/packages/ai-infra/inference-common/tsconfig.json index b05325b824a67..7f261ba1dcb82 100644 --- a/x-pack/packages/ai-infra/inference-common/tsconfig.json +++ b/x-pack/packages/ai-infra/inference-common/tsconfig.json @@ -15,5 +15,7 @@ "exclude": [ "target/**/*" ], - "kbn_references": [] + "kbn_references": [ + "@kbn/sse-utils", + ] } diff --git a/x-pack/plugins/inference/tsconfig.json b/x-pack/plugins/inference/tsconfig.json index 4e4a91d15e57b..9bf696c45ccb7 100644 --- a/x-pack/plugins/inference/tsconfig.json +++ b/x-pack/plugins/inference/tsconfig.json @@ -19,7 +19,6 @@ ], "kbn_references": [ "@kbn/i18n", - "@kbn/sse-utils", "@kbn/esql-ast", "@kbn/esql-validation-autocomplete", "@kbn/core", @@ -33,9 +32,6 @@ "@kbn/core-http-server", "@kbn/actions-plugin", "@kbn/config-schema", - "@kbn/es-types", - "@kbn/field-types", - "@kbn/expressions-plugin", "@kbn/inference-common", ] } From f354cfefc3669e85cddef6a09f82577fb31cf64e Mon Sep 17 00:00:00 2001 From: pgayvallet Date: Wed, 30 Oct 2024 10:31:16 +0100 Subject: [PATCH 09/21] why are you even importing that from there? --- .../investigate_app/server/lib/get_sample_documents.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/x-pack/plugins/observability_solution/investigate_app/server/lib/get_sample_documents.ts b/x-pack/plugins/observability_solution/investigate_app/server/lib/get_sample_documents.ts index 75e21526a6506..9621a5ff1a4ed 100644 --- a/x-pack/plugins/observability_solution/investigate_app/server/lib/get_sample_documents.ts +++ b/x-pack/plugins/observability_solution/investigate_app/server/lib/get_sample_documents.ts @@ -7,7 +7,7 @@ import pLimit from 'p-limit'; import { estypes } from '@elastic/elasticsearch'; import { castArray, sortBy, uniq, partition, shuffle } from 'lodash'; -import { truncateList } from '@kbn/inference-plugin/common/util/truncate_list'; +import { truncateList } from '@kbn/inference-plugin/common/utils/truncate_list'; import { QueryDslQueryContainer } from '@kbn/data-views-plugin/common/types'; import type { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; import { rangeQuery, excludeFrozenQuery } from './queries'; From 35ec20471e60c758d1f432e332d6cc92b6e606d3 Mon Sep 17 00:00:00 2001 From: pgayvallet Date: Wed, 30 Oct 2024 12:08:37 +0100 Subject: [PATCH 10/21] more import fixes --- x-pack/plugins/inference/public/types.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/x-pack/plugins/inference/public/types.ts b/x-pack/plugins/inference/public/types.ts index df80256679ab4..735abfb5459a0 100644 --- a/x-pack/plugins/inference/public/types.ts +++ b/x-pack/plugins/inference/public/types.ts @@ -4,9 +4,9 @@ * 2.0; you may not use this file except in compliance with the Elastic License * 2.0. */ -import type { ChatCompleteAPI } from '../common/chat_complete'; + +import type { ChatCompleteAPI, OutputAPI } from '@kbn/inference-common'; import type { InferenceConnector } from '../common/connectors'; -import type { OutputAPI } from '../common/output'; /* eslint-disable @typescript-eslint/no-empty-interface*/ From 512ec1dc3efe0712e93f26fcda2717ed7e543051 Mon Sep 17 00:00:00 2001 From: pgayvallet Date: Wed, 30 Oct 2024 12:56:06 +0100 Subject: [PATCH 11/21] more type fix --- .../observability_ai_assistant/common/types.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/common/types.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/common/types.ts index 7595c42e4dc93..51ae37b39d90f 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/common/types.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/common/types.ts @@ -5,7 +5,7 @@ * 2.0. */ import { IconType } from '@elastic/eui'; -import type { ToolSchema } from '@kbn/inference-plugin/common'; +import type { ToolSchema } from '@kbn/inference-common'; import type { AssistantScope } from '@kbn/ai-assistant-common'; import type { ObservabilityAIAssistantChatService } from '../public'; import type { FunctionResponse } from './functions/types'; From ac2aedcb1e2cae707c2e3aa93806e7a2ce4d42a7 Mon Sep 17 00:00:00 2001 From: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Date: Wed, 30 Oct 2024 12:07:58 +0000 Subject: [PATCH 12/21] [CI] Auto-commit changed files from 'node scripts/lint_ts_projects --fix' --- .../observability_ai_assistant/tsconfig.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/tsconfig.json b/x-pack/plugins/observability_solution/observability_ai_assistant/tsconfig.json index 7c2f2212ee946..fac579124d888 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/tsconfig.json +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/tsconfig.json @@ -14,7 +14,6 @@ ], "kbn_references": [ "@kbn/i18n", - "@kbn/inference-plugin", "@kbn/logging", "@kbn/kibana-utils-plugin", "@kbn/core-analytics-browser", @@ -47,6 +46,7 @@ "@kbn/inference-plugin", "@kbn/management-settings-ids", "@kbn/ai-assistant-common", + "@kbn/inference-common", ], "exclude": ["target/**/*"] } From 7387e47eb3ab27941a0b6c34672debddc6f64aad Mon Sep 17 00:00:00 2001 From: pgayvallet Date: Wed, 30 Oct 2024 13:15:16 +0100 Subject: [PATCH 13/21] last type fixes --- .../common/convert_messages_for_inference.ts | 4 ++-- .../server/functions/query/index.ts | 7 ++----- .../server/assistant/tools/esql/nl_to_esql_tool.ts | 6 +----- 3 files changed, 5 insertions(+), 12 deletions(-) diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant_app/common/convert_messages_for_inference.ts b/x-pack/plugins/observability_solution/observability_ai_assistant_app/common/convert_messages_for_inference.ts index 1dc8638626d0b..7ab9516440988 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant_app/common/convert_messages_for_inference.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant_app/common/convert_messages_for_inference.ts @@ -10,8 +10,8 @@ import { AssistantMessage, Message as InferenceMessage, MessageRole as InferenceMessageRole, - generateFakeToolCallId, -} from '@kbn/inference-plugin/common'; +} from '@kbn/inference-common'; +import { generateFakeToolCallId } from '@kbn/inference-plugin/common'; export function convertMessagesForInference(messages: Message[]): InferenceMessage[] { const inferenceMessages: InferenceMessage[] = []; diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant_app/server/functions/query/index.ts b/x-pack/plugins/observability_solution/observability_ai_assistant_app/server/functions/query/index.ts index 3643c54365248..210dee20339af 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant_app/server/functions/query/index.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant_app/server/functions/query/index.ts @@ -5,11 +5,8 @@ * 2.0. */ -import { - correctCommonEsqlMistakes, - isChatCompletionChunkEvent, - isOutputEvent, -} from '@kbn/inference-plugin/common'; +import { isChatCompletionChunkEvent, isOutputEvent } from '@kbn/inference-common'; +import { correctCommonEsqlMistakes } from '@kbn/inference-plugin/common'; import { naturalLanguageToEsql } from '@kbn/inference-plugin/server'; import { FunctionVisibility, diff --git a/x-pack/plugins/security_solution/server/assistant/tools/esql/nl_to_esql_tool.ts b/x-pack/plugins/security_solution/server/assistant/tools/esql/nl_to_esql_tool.ts index 1205fb03b0458..11ca8ffd94edd 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/esql/nl_to_esql_tool.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/esql/nl_to_esql_tool.ts @@ -49,11 +49,7 @@ export const NL_TO_ESQL_TOOL: AssistantTool = { connectorId, input: question, ...(isOssModel ? { functionCalling: 'simulated' } : {}), - logger: { - debug: (source) => { - logger.debug(typeof source === 'function' ? source() : source); - }, - }, + logger, }) ); }; From 71c0d75d76a6e4bfbc4cb5026e26982ce0a258ba Mon Sep 17 00:00:00 2001 From: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Date: Wed, 30 Oct 2024 12:28:17 +0000 Subject: [PATCH 14/21] [CI] Auto-commit changed files from 'node scripts/lint_ts_projects --fix' --- .../observability_ai_assistant/tsconfig.json | 1 - .../observability_ai_assistant_app/tsconfig.json | 1 + 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/tsconfig.json b/x-pack/plugins/observability_solution/observability_ai_assistant/tsconfig.json index fac579124d888..750bf69477653 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/tsconfig.json +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/tsconfig.json @@ -43,7 +43,6 @@ "@kbn/serverless", "@kbn/core-elasticsearch-server", "@kbn/core-ui-settings-server", - "@kbn/inference-plugin", "@kbn/management-settings-ids", "@kbn/ai-assistant-common", "@kbn/inference-common", diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant_app/tsconfig.json b/x-pack/plugins/observability_solution/observability_ai_assistant_app/tsconfig.json index af04a677f5e94..6608799caaf61 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant_app/tsconfig.json +++ b/x-pack/plugins/observability_solution/observability_ai_assistant_app/tsconfig.json @@ -69,6 +69,7 @@ "@kbn/cloud-plugin", "@kbn/logs-data-access-plugin", "@kbn/ai-assistant-common", + "@kbn/inference-common", ], "exclude": [ "target/**/*" From 4e28e144e899191feabbca798fa491780692da07 Mon Sep 17 00:00:00 2001 From: pgayvallet Date: Wed, 30 Oct 2024 14:20:41 +0100 Subject: [PATCH 15/21] extract option type --- x-pack/packages/ai-infra/inference-common/index.ts | 1 + .../inference-common/src/chat_complete/api.ts | 14 ++++++++------ .../inference-common/src/chat_complete/index.ts | 7 ++++++- 3 files changed, 15 insertions(+), 7 deletions(-) diff --git a/x-pack/packages/ai-infra/inference-common/index.ts b/x-pack/packages/ai-infra/inference-common/index.ts index 8ceaab2e4b78f..eacf31afea5e9 100644 --- a/x-pack/packages/ai-infra/inference-common/index.ts +++ b/x-pack/packages/ai-infra/inference-common/index.ts @@ -24,6 +24,7 @@ export { type FunctionCallingMode, type ToolChoice, type ChatCompleteAPI, + type ChatCompleteOptions, type ChatCompletionResponse, type ChatCompletionTokenCountEvent, type ChatCompletionEvent, diff --git a/x-pack/packages/ai-infra/inference-common/src/chat_complete/api.ts b/x-pack/packages/ai-infra/inference-common/src/chat_complete/api.ts index 5e8d6908a10cc..681727f9051f2 100644 --- a/x-pack/packages/ai-infra/inference-common/src/chat_complete/api.ts +++ b/x-pack/packages/ai-infra/inference-common/src/chat_complete/api.ts @@ -10,6 +10,13 @@ import type { ToolOptions } from './tools'; import type { Message } from './messages'; import type { ChatCompletionEvent } from './events'; +export type ChatCompleteOptions = { + connectorId: string; + system?: string; + messages: Message[]; + functionCalling?: FunctionCallingMode; +} & TToolOptions; + /** * Request a completion from the LLM based on a prompt or conversation. * @@ -20,12 +27,7 @@ import type { ChatCompletionEvent } from './events'; * @param {Record} [options.tools] A map of tools that can be called by the LLM */ export type ChatCompleteAPI = ( - options: { - connectorId: string; - system?: string; - messages: Message[]; - functionCalling?: FunctionCallingMode; - } & TToolOptions + options: ChatCompleteOptions ) => ChatCompletionResponse; export type ChatCompletionResponse = Observable< diff --git a/x-pack/packages/ai-infra/inference-common/src/chat_complete/index.ts b/x-pack/packages/ai-infra/inference-common/src/chat_complete/index.ts index 1a66a535c1ebc..7183ff6361e5b 100644 --- a/x-pack/packages/ai-infra/inference-common/src/chat_complete/index.ts +++ b/x-pack/packages/ai-infra/inference-common/src/chat_complete/index.ts @@ -5,7 +5,12 @@ * 2.0. */ -export type { ChatCompletionResponse, ChatCompleteAPI, FunctionCallingMode } from './api'; +export type { + ChatCompletionResponse, + ChatCompleteAPI, + ChatCompleteOptions, + FunctionCallingMode, +} from './api'; export { ChatCompletionEventType, type ChatCompletionMessageEvent, From 7904a14c8758c3454295ad04daf574c28bcfb605 Mon Sep 17 00:00:00 2001 From: pgayvallet Date: Thu, 31 Oct 2024 08:33:46 +0100 Subject: [PATCH 16/21] export chat complete errors from entry point --- x-pack/packages/ai-infra/inference-common/index.ts | 7 +++++++ .../ai-infra/inference-common/src/chat_complete/index.ts | 9 +++++++++ .../inference/server/util/validate_tool_calls.test.ts | 3 +-- 3 files changed, 17 insertions(+), 2 deletions(-) diff --git a/x-pack/packages/ai-infra/inference-common/index.ts b/x-pack/packages/ai-infra/inference-common/index.ts index eacf31afea5e9..283cac2594c29 100644 --- a/x-pack/packages/ai-infra/inference-common/index.ts +++ b/x-pack/packages/ai-infra/inference-common/index.ts @@ -36,6 +36,13 @@ export { isChatCompletionMessageEvent, isChatCompletionEvent, isChatCompletionChunkEvent, + ChatCompletionErrorCode, + type ChatCompletionToolNotFoundError, + type ChatCompletionToolValidationError, + type ChatCompletionTokenLimitReachedError, + isToolValidationError, + isTokenLimitReachedError, + isToolNotFoundError, } from './src/chat_complete'; export { OutputEventType, diff --git a/x-pack/packages/ai-infra/inference-common/src/chat_complete/index.ts b/x-pack/packages/ai-infra/inference-common/src/chat_complete/index.ts index 7183ff6361e5b..6bf73f605c49e 100644 --- a/x-pack/packages/ai-infra/inference-common/src/chat_complete/index.ts +++ b/x-pack/packages/ai-infra/inference-common/src/chat_complete/index.ts @@ -43,3 +43,12 @@ export { withoutChunkEvents, withoutTokenCountEvents, } from './event_utils'; +export { + ChatCompletionErrorCode, + type ChatCompletionToolNotFoundError, + type ChatCompletionToolValidationError, + type ChatCompletionTokenLimitReachedError, + isToolValidationError, + isTokenLimitReachedError, + isToolNotFoundError, +} from './errors'; diff --git a/x-pack/plugins/inference/server/util/validate_tool_calls.test.ts b/x-pack/plugins/inference/server/util/validate_tool_calls.test.ts index 2a6aed608dced..57b030771c6c0 100644 --- a/x-pack/plugins/inference/server/util/validate_tool_calls.test.ts +++ b/x-pack/plugins/inference/server/util/validate_tool_calls.test.ts @@ -5,8 +5,7 @@ * 2.0. */ -import { ToolChoiceType } from '@kbn/inference-common'; -import { isToolValidationError } from '@kbn/inference-common/src/chat_complete/errors'; +import { ToolChoiceType, isToolValidationError } from '@kbn/inference-common'; import { validateToolCalls } from './validate_tool_calls'; describe('validateToolCalls', () => { From c14b4f64d2e358970a1d5aefeba8ca809547ba43 Mon Sep 17 00:00:00 2001 From: pgayvallet Date: Thu, 31 Oct 2024 13:21:15 +0100 Subject: [PATCH 17/21] revert changes to ESQL task --- .../inference-common/src/chat_complete/api.ts | 63 +++++- .../src/chat_complete/events.ts | 14 ++ .../tasks/nl_to_esql/actions/generate_esql.ts | 184 ++++++++++++------ .../server/tasks/nl_to_esql/actions/index.ts | 6 +- .../actions/request_documentation.ts | 117 +++++------ .../server/tasks/nl_to_esql/actions/shared.ts | 29 +++ .../actions/summarize_discussion.ts | 138 ------------- .../server/tasks/nl_to_esql/actions/types.ts | 20 -- .../inference/server/tasks/nl_to_esql/task.ts | 96 ++++----- .../server/tasks/nl_to_esql/types.ts | 11 +- .../nl_to_esql/utils/extract_queries.test.ts | 75 ------- .../tasks/nl_to_esql/utils/extract_queries.ts | 36 ---- .../server/tasks/nl_to_esql/utils/index.ts | 9 - .../utils/validate_query_ast.test.ts | 31 --- .../nl_to_esql/utils/validate_query_ast.ts | 63 ------ .../tasks/nl_to_esql/validate_esql_query.ts | 72 +++++++ 16 files changed, 380 insertions(+), 584 deletions(-) create mode 100644 x-pack/plugins/inference/server/tasks/nl_to_esql/actions/shared.ts delete mode 100644 x-pack/plugins/inference/server/tasks/nl_to_esql/actions/summarize_discussion.ts delete mode 100644 x-pack/plugins/inference/server/tasks/nl_to_esql/actions/types.ts delete mode 100644 x-pack/plugins/inference/server/tasks/nl_to_esql/utils/extract_queries.test.ts delete mode 100644 x-pack/plugins/inference/server/tasks/nl_to_esql/utils/extract_queries.ts delete mode 100644 x-pack/plugins/inference/server/tasks/nl_to_esql/utils/index.ts delete mode 100644 x-pack/plugins/inference/server/tasks/nl_to_esql/utils/validate_query_ast.test.ts delete mode 100644 x-pack/plugins/inference/server/tasks/nl_to_esql/utils/validate_query_ast.ts create mode 100644 x-pack/plugins/inference/server/tasks/nl_to_esql/validate_esql_query.ts diff --git a/x-pack/packages/ai-infra/inference-common/src/chat_complete/api.ts b/x-pack/packages/ai-infra/inference-common/src/chat_complete/api.ts index 681727f9051f2..a872b3f950527 100644 --- a/x-pack/packages/ai-infra/inference-common/src/chat_complete/api.ts +++ b/x-pack/packages/ai-infra/inference-common/src/chat_complete/api.ts @@ -10,28 +10,71 @@ import type { ToolOptions } from './tools'; import type { Message } from './messages'; import type { ChatCompletionEvent } from './events'; +/** + * Request a completion from the LLM based on a prompt or conversation. + * + * @example using the API to get an event observable. + * ```ts + * const events$ = chatComplete({ + * connectorId: 'my-connector', + * system: "You are a helpful assistant", + * messages: [ + * { role: MessageRole.User, content: "First question?"}, + * { role: MessageRole.Assistant, content: "Some answer"}, + * { role: MessageRole.User, content: "Another question?"}, + * ] + * }); + * + * ``` + * @example directly accessing the LLM message without using the observable API + * ```ts + * const message = await chatComplete({ + * connectorId: 'my-connector', + * system: "You are a helpful assistant", + * messages: [ + * { role: MessageRole.User, content: "Some question?"}, + * ] + * }).getMessage(); + * ``` + */ +export type ChatCompleteAPI = ( + options: ChatCompleteOptions +) => ChatCompletionResponse; + +/** + * Options used to call the {@link ChatCompleteAPI} + */ export type ChatCompleteOptions = { + /** + * The ID of the connector to use + */ connectorId: string; + /** + * Optional system message for the LLM. + */ system?: string; + /** + * The list of messages for the current conversation + */ messages: Message[]; + /** + * Function calling mode, defaults to "native". + */ functionCalling?: FunctionCallingMode; } & TToolOptions; /** - * Request a completion from the LLM based on a prompt or conversation. + * Response from the {@link ChatCompleteAPI}. * - * @param {string} options.connectorId The ID of the connector to use - * @param {string} [options.system] A system message that defines the behavior of the LLM. - * @param {Message[]} options.message A list of messages that make up the conversation to be completed. - * @param {ToolChoice} [options.toolChoice] Force the LLM to call a (specific) tool, or no tool - * @param {Record} [options.tools] A map of tools that can be called by the LLM + * Observable of {@link ChatCompletionEvent} */ -export type ChatCompleteAPI = ( - options: ChatCompleteOptions -) => ChatCompletionResponse; - export type ChatCompletionResponse = Observable< ChatCompletionEvent >; +/** + * Define the function calling mode when using inference APIs. + * - native will use the LLM's native function calling (requires the LLM to have native support) + * - simulated: will emulate function calling with function calling instructions + */ export type FunctionCallingMode = 'native' | 'simulated'; diff --git a/x-pack/packages/ai-infra/inference-common/src/chat_complete/events.ts b/x-pack/packages/ai-infra/inference-common/src/chat_complete/events.ts index 924b9631239b7..c6f5766405c40 100644 --- a/x-pack/packages/ai-infra/inference-common/src/chat_complete/events.ts +++ b/x-pack/packages/ai-infra/inference-common/src/chat_complete/events.ts @@ -43,6 +43,20 @@ export type ChatCompletionTokenCountEvent = }; }; +/** + * Events emitted from the {@link ChatCompletionResponse} observable + * returned from the {@link ChatCompleteAPI}. + * + * The chatComplete API returns 3 type of events: + * - {@link ChatCompletionChunkEvent}: message chunk events + * - {@link ChatCompletionTokenCountEvent}: token count event + * - {@link ChatCompletionMessageEvent}: message event + * + * Note that chunk events can be emitted any amount of times, but token count will be emitted + * at most once (could not be emitted depending on the underlying connector), and message + * event will be emitted ex + * + */ export type ChatCompletionEvent = | ChatCompletionChunkEvent | ChatCompletionTokenCountEvent diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/generate_esql.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/generate_esql.ts index 3ee4297a13c41..d31952e2f5252 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/generate_esql.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/generate_esql.ts @@ -5,50 +5,77 @@ * 2.0. */ -import { map, tap, lastValueFrom } from 'rxjs'; +import { Observable, map, merge, of, switchMap } from 'rxjs'; import type { Logger } from '@kbn/logging'; +import { ToolCall, ToolOptions } from '../../../../common/chat_complete/tools'; import { - ToolCall, - ToolOptions, + correctCommonEsqlMistakes, + generateFakeToolCallId, isChatCompletionMessageEvent, - withoutTokenCountEvents, - withoutChunkEvents, Message, MessageRole, -} from '@kbn/inference-common'; -import { correctCommonEsqlMistakes, generateFakeToolCallId } from '../../../../common'; +} from '../../../../common'; +import { InferenceClient, withoutTokenCountEvents } from '../../..'; +import { OutputCompleteEvent, OutputEventType } from '../../../../common/output'; import { INLINE_ESQL_QUERY_REGEX } from '../../../../common/tasks/nl_to_esql/constants'; -import type { ActionsOptionsBase } from './types'; +import { EsqlDocumentBase } from '../doc_base'; +import { requestDocumentationSchema } from './shared'; +import type { NlToEsqlTaskEvent } from '../types'; +import type { FunctionCallingMode } from '../../../../common/chat_complete'; -export const generateEsqlTaskFn = ({ - client, +export const generateEsqlTask = ({ + chatCompleteApi, connectorId, systemMessage, + messages, + toolOptions: { tools, toolChoice }, + docBase, functionCalling, logger, - output$, -}: ActionsOptionsBase & { +}: { + connectorId: string; systemMessage: string; + messages: Message[]; + toolOptions: ToolOptions; + chatCompleteApi: InferenceClient['chatComplete']; + docBase: EsqlDocumentBase; + functionCalling?: FunctionCallingMode; + logger: Pick; }) => { - return async function generateEsql({ - messages, - documentationRequest: { keywords, requestedDocumentation }, + return function askLlmToRespond({ + documentationRequest: { commands, functions }, }: { - messages: Message[]; - documentationRequest: { keywords: string[]; requestedDocumentation: Record }; - }) { - const fakeRequestDocsToolCall = createFakeTooCall(keywords); + documentationRequest: { commands?: string[]; functions?: string[] }; + }): Observable> { + const keywords = [...(commands ?? []), ...(functions ?? [])]; + const requestedDocumentation = docBase.getDocumentation(keywords); + const fakeRequestDocsToolCall = createFakeTooCall(commands, functions); - const result = await lastValueFrom( - client - .chatComplete({ - connectorId, - functionCalling, - system: `${systemMessage} + return merge( + of< + OutputCompleteEvent< + 'request_documentation', + { keywords: string[]; requestedDocumentation: Record } + > + >({ + type: OutputEventType.OutputComplete, + id: 'request_documentation', + output: { + keywords, + requestedDocumentation, + }, + content: '', + }), + chatCompleteApi({ + connectorId, + functionCalling, + system: `${systemMessage} # Current task - Your current task is to respond to the user's question. + Your current task is to respond to the user's question. If there is a tool + suitable for answering the user's question, use that tool, preferably + with a natural language reply included. Format any ES|QL query as follows: \`\`\`esql @@ -68,47 +95,74 @@ export const generateEsqlTaskFn = ({ When converting queries from one language to ES|QL, make sure that the functions are available and documented in ES|QL. E.g., for SPL's LEN, use LENGTH. For IF, use CASE. `, - messages: [ - ...messages, - { - role: MessageRole.Assistant, - content: null, - toolCalls: [fakeRequestDocsToolCall], + messages: [ + ...messages, + { + role: MessageRole.Assistant, + content: null, + toolCalls: [fakeRequestDocsToolCall], + }, + { + role: MessageRole.Tool, + response: { + documentation: requestedDocumentation, }, - { - role: MessageRole.Tool, - response: { - documentation: requestedDocumentation, - }, - toolCallId: fakeRequestDocsToolCall.toolCallId, - }, - ], - }) - .pipe( - withoutTokenCountEvents(), - map((event) => { - if (isChatCompletionMessageEvent(event)) { - return { - ...event, - content: event.content - ? correctEsqlMistakes({ content: event.content, logger }) - : event.content, - }; + toolCallId: fakeRequestDocsToolCall.toolCallId, + }, + ], + toolChoice, + tools: { + ...tools, + request_documentation: { + description: 'Request additional ES|QL documentation if needed', + schema: requestDocumentationSchema, + }, + }, + }).pipe( + withoutTokenCountEvents(), + map((generateEvent) => { + if (isChatCompletionMessageEvent(generateEvent)) { + return { + ...generateEvent, + content: generateEvent.content + ? correctEsqlMistakes({ content: generateEvent.content, logger }) + : generateEvent.content, + }; + } + + return generateEvent; + }), + switchMap((generateEvent) => { + if (isChatCompletionMessageEvent(generateEvent)) { + const onlyToolCall = + generateEvent.toolCalls.length === 1 ? generateEvent.toolCalls[0] : undefined; + + if (onlyToolCall?.function.name === 'request_documentation') { + const args = onlyToolCall.function.arguments; + + return askLlmToRespond({ + documentationRequest: { + commands: args.commands, + functions: args.functions, + }, + }); } - return event; - }), - tap((event) => { - output$.next(event); - }), - withoutChunkEvents() - ) - ); + } - return { content: result.content }; + return of(generateEvent); + }) + ) + ); }; }; -const correctEsqlMistakes = ({ content, logger }: { content: string; logger: Logger }) => { +const correctEsqlMistakes = ({ + content, + logger, +}: { + content: string; + logger: Pick; +}) => { return content.replaceAll(INLINE_ESQL_QUERY_REGEX, (_match, query) => { const correction = correctCommonEsqlMistakes(query); if (correction.isCorrection) { @@ -118,12 +172,16 @@ const correctEsqlMistakes = ({ content, logger }: { content: string; logger: Log }); }; -const createFakeTooCall = (keywords: string[]): ToolCall => { +const createFakeTooCall = ( + commands: string[] | undefined, + functions: string[] | undefined +): ToolCall => { return { function: { name: 'request_documentation', arguments: { - keywords, + commands, + functions, }, }, toolCallId: generateFakeToolCallId(), diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/index.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/index.ts index 563a1d2099157..ec1d54dd8a26b 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/index.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/index.ts @@ -5,7 +5,5 @@ * 2.0. */ -export { requestDocumentationFn } from './request_documentation'; -export { generateEsqlTaskFn } from './generate_esql'; -export { summarizeDiscussionFn } from './summarize_discussion'; -export type { ActionsOptionsBase } from './types'; +export { requestDocumentation } from './request_documentation'; +export { generateEsqlTask } from './generate_esql'; diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/request_documentation.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/request_documentation.ts index c8eaed176da73..d4eb3060f59bb 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/request_documentation.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/request_documentation.ts @@ -5,86 +5,59 @@ * 2.0. */ -import { lastValueFrom, tap, map } from 'rxjs'; -import { withoutOutputUpdateEvents, Message, ToolSchema } from '@kbn/inference-common'; -import { EsqlDocumentBase } from '../doc_base'; -import type { ActionsOptionsBase } from './types'; +import { isEmpty } from 'lodash'; +import { InferenceClient, withoutOutputUpdateEvents } from '../../..'; +import { Message } from '../../../../common'; +import { ToolChoiceType, ToolOptions } from '../../../../common/chat_complete/tools'; +import { requestDocumentationSchema } from './shared'; +import type { FunctionCallingMode } from '../../../../common/chat_complete'; -export const requestDocumentationSchema = { - type: 'object', - properties: { - commands: { - type: 'array', - items: { - type: 'string', - }, - description: - 'ES|QL source and processing commands you want to analyze before generating the query.', - }, - functions: { - type: 'array', - items: { - type: 'string', - }, - description: 'ES|QL functions you want to analyze before generating the query.', - }, - }, - required: ['commands', 'functions'] as const, -} satisfies ToolSchema; - -export const requestDocumentationFn = ({ - client, +export const requestDocumentation = ({ + outputApi, system, + messages, connectorId, functionCalling, - docBase, - output$, -}: ActionsOptionsBase & { - docBase: EsqlDocumentBase; + toolOptions: { tools, toolChoice }, +}: { + outputApi: InferenceClient['output']; system: string; + messages: Message[]; + connectorId: string; + functionCalling?: FunctionCallingMode; + toolOptions: ToolOptions; }) => { - return async ({ messages }: { messages: Message[] }) => { - const result = await lastValueFrom( - client - .output('request_documentation', { - connectorId, - functionCalling, - system, - previousMessages: messages, - input: `Based on the previous conversation, request documentation - for commands or functions listed in the ES|QL handbook to help you - get the right information needed to generate a query. + const hasTools = !isEmpty(tools) && toolChoice !== ToolChoiceType.none; - Make sure to request documentation for any command or function you think you may use, - even if you end up not using all of them. + return outputApi('request_documentation', { + connectorId, + functionCalling, + system, + previousMessages: messages, + input: `Based on the previous conversation, request documentation + from the ES|QL handbook to help you get the right information + needed to generate a query. - Example: if you need to... - - Group or aggregate data? Request \`STATS\`. - - Extract data? Request \`DISSECT\` and \`GROK\`. + Examples for functions and commands: + - Do you need to group data? Request \`STATS\`. + - Extract data? Request \`DISSECT\` AND \`GROK\`. - Convert a column based on a set of conditionals? Request \`EVAL\` and \`CASE\`. - - Group data by time intervals? Request \`BUCKET\` - `, - schema: requestDocumentationSchema, - }) - .pipe( - withoutOutputUpdateEvents(), - map((event) => { - const keywords = [...(event.output.commands ?? []), ...(event.output.functions ?? [])]; - const requestedDocumentation = docBase.getDocumentation(keywords); - return { - ...event, - output: { - keywords, - requestedDocumentation, - }, - }; - }), - tap((event) => { - output$.next(event); - }) - ) - ); - return result.output; - }; + ${ + hasTools + ? `### Tools + + The following tools will be available to be called in the step after this. + + \`\`\`json + ${JSON.stringify({ + tools, + toolChoice, + })} + \`\`\`` + : '' + } + `, + schema: requestDocumentationSchema, + }).pipe(withoutOutputUpdateEvents()); }; diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/shared.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/shared.ts new file mode 100644 index 0000000000000..f0fc796173b23 --- /dev/null +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/shared.ts @@ -0,0 +1,29 @@ +/* + * 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 { ToolSchema } from '../../../../common'; + +export const requestDocumentationSchema = { + type: 'object', + properties: { + commands: { + type: 'array', + items: { + type: 'string', + }, + description: + 'ES|QL source and processing commands you want to analyze before generating the query.', + }, + functions: { + type: 'array', + items: { + type: 'string', + }, + description: 'ES|QL functions you want to analyze before generating the query.', + }, + }, +} satisfies ToolSchema; diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/summarize_discussion.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/summarize_discussion.ts deleted file mode 100644 index a5625c4423232..0000000000000 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/summarize_discussion.ts +++ /dev/null @@ -1,138 +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. - */ - -import { lastValueFrom, tap } from 'rxjs'; -import { Message, ToolSchema, withoutOutputUpdateEvents } from '@kbn/inference-common'; -import type { ActionsOptionsBase } from './types'; - -const summarizeDiscussionSchema = { - type: 'object', - properties: { - summary: { - type: 'string', - description: 'The requested summary of the conversation', - }, - }, - required: ['summary'] as const, -} satisfies ToolSchema; - -export const summarizeDiscussionFn = ({ - client, - connectorId, - functionCalling, - output$, -}: ActionsOptionsBase) => { - return async function summarizeDiscussion({ messages }: { messages: Message[] }) { - const result = await lastValueFrom( - client - .output('summarize_discussion', { - connectorId, - functionCalling, - system: ` - You are an helpful Elastic assistant in charge of helping the user generating - ES|QL queries. Specifically, you current role is to extract from a discussion - all the information that could be useful to answer the user's question and generate - an ES|QL query. - - Make sure to include ALL information that can be useful to generate, including: - - which index or data view to perform the query against - - the known fields, their type and their description - - the question that needs to be answered with an ES|QL query - - Any previous query that may be mentioned - - Your role is only to retrieve and extract the info that are present in the conversation, not to invent - any. E.g. if there are no description attached to the index's schema, don't invent any. - - Here are a few examples: - ${examples} - `, - previousMessages: messages, - input: `Please extract the requested summary from the current discussion, as specified in your system - instructions. - `, - schema: summarizeDiscussionSchema, - }) - .pipe( - withoutOutputUpdateEvents(), - tap((event) => { - output$.next(event); - }) - ) - ); - const summary = result.output.summary; - return { summary }; - }; -}; - -const examples = ` - - #example 1 - - ## message list: - - [ - { - "role": "user", - "content" "Show me a query to display the amount of order per month over the past year" - }, - { - "content": { - "indices": [ - "kibana_sample_data_ecommerce" - ], - "fields": [ - "order_date:date", - "order_id:keyword" - ] - } - } - ] - - ## expected summary - - The user wants to generate an ES|QL query to display the amount of orders per month over the past year. - The query should target the index 'kibana_sample_data_ecommerce'. - The relevant fields for this query are: - - order_date of type date - - order_id of type keyword - - #example 2 - - ## message list: - - [ - { - "role": "user", - "content" "Show me a query to display the amount of order per month over the past year" - } - ] - - ## expected summary - - The user wants to generate an ES|QL query to display the amount of orders per month over the past year. - There are no indications about which index should be targeted. - There are no indications about the index's schema. - - #example 3 - - ## message list: - - [ - { - "role": "user", - "content" "Can you convert this SPL query to ESQL? \`index=network_firewall "SYN Timeout" | stats count by dest\`" - } - ] - - ## expected summary - - The user wants to convert a SPL query to ESQL. - The SPL query is: \`index=network_firewall "SYN Timeout" | stats count by dest\`. - The query should target the index 'network_firewall'. - The query should re-use the fields from the SPL query. -\` -`; diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/types.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/types.ts deleted file mode 100644 index dc198d6b44686..0000000000000 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/types.ts +++ /dev/null @@ -1,20 +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. - */ - -import type { Logger } from '@kbn/logging'; -import { Subject } from 'rxjs'; -import type { FunctionCallingMode } from '@kbn/inference-common'; -import type { InferenceClient } from '../../../types'; -import { NlToEsqlTaskEvent } from '../types'; - -export interface ActionsOptionsBase { - client: Pick; - connectorId: string; - logger: Logger; - functionCalling?: FunctionCallingMode; - output$: Subject>; -} diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/task.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/task.ts index c0ce44d77971c..e0c5a838ea148 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/task.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/task.ts @@ -6,16 +6,12 @@ */ import { once } from 'lodash'; -import { Observable, Subject } from 'rxjs'; -import { MessageRole, type Message, type ToolOptions } from '@kbn/inference-common'; +import { Observable, from, switchMap } from 'rxjs'; +import { Message, MessageRole } from '../../../common/chat_complete'; +import type { ToolOptions } from '../../../common/chat_complete/tools'; import { EsqlDocumentBase } from './doc_base'; -import { - type ActionsOptionsBase, - generateEsqlTaskFn, - requestDocumentationFn, - summarizeDiscussionFn, -} from './actions'; -import { NlToEsqlTaskEvent, NlToEsqlTaskParams } from './types'; +import { requestDocumentation, generateEsqlTask } from './actions'; +import { NlToEsqlTaskParams, NlToEsqlTaskEvent } from './types'; const loadDocBase = once(() => EsqlDocumentBase.load()); @@ -26,60 +22,48 @@ export function naturalLanguageToEsql({ toolChoice, logger, functionCalling, - summarizeInput = false, ...rest }: NlToEsqlTaskParams): Observable> { - const output$ = new Subject>(); - - const baseOptions: ActionsOptionsBase = { - client, - connectorId, - functionCalling, - logger, - output$, - }; - - loadDocBase() - .then(async (docBase) => { - let messages: Message[] = + return from(loadDocBase()).pipe( + switchMap((docBase) => { + const systemMessage = docBase.getSystemMessage(); + const messages: Message[] = 'input' in rest ? [{ role: MessageRole.User, content: rest.input }] : rest.messages; - const summarizeDiscussion = summarizeDiscussionFn({ - ...baseOptions, - }); - const requestDocumentation = requestDocumentationFn({ - ...baseOptions, - docBase, - system: docBase.getSystemMessage(), - }); - const generateEsql = generateEsqlTaskFn({ - ...baseOptions, - systemMessage: docBase.getSystemMessage(), - }); - - if (summarizeInput) { - const discussionSummary = await summarizeDiscussion({ messages }); - messages = discussionFromSummary(discussionSummary.summary); - } - - const documentationRequest = await requestDocumentation({ + const askLlmToRespond = generateEsqlTask({ + connectorId, + chatCompleteApi: client.chatComplete, messages, + docBase, + logger, + systemMessage, + functionCalling, + toolOptions: { + tools, + toolChoice, + }, }); - await generateEsql({ - documentationRequest, + return requestDocumentation({ + connectorId, + functionCalling, + outputApi: client.output, messages, - }); - - output$.complete(); + system: systemMessage, + toolOptions: { + tools, + toolChoice, + }, + }).pipe( + switchMap((documentationEvent) => { + return askLlmToRespond({ + documentationRequest: { + commands: documentationEvent.output.commands, + functions: documentationEvent.output.functions, + }, + }); + }) + ); }) - .catch((err) => { - output$.error(err); - }); - - return output$.asObservable(); + ); } - -const discussionFromSummary = (summary: string): Message[] => { - return [{ role: MessageRole.User, content: summary }]; -}; diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/types.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/types.ts index ca3b44fd18c4f..a0bcd635081ea 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/types.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/types.ts @@ -11,9 +11,9 @@ import type { ChatCompletionMessageEvent, FunctionCallingMode, Message, - ToolOptions, - OutputCompleteEvent, -} from '@kbn/inference-common'; +} from '../../../common/chat_complete'; +import type { ToolOptions } from '../../../common/chat_complete/tools'; +import type { OutputCompleteEvent } from '../../../common/output'; import type { InferenceClient } from '../../types'; export type NlToEsqlTaskEvent = @@ -21,16 +21,13 @@ export type NlToEsqlTaskEvent = 'request_documentation', { keywords: string[]; requestedDocumentation: Record } > - | OutputCompleteEvent<'summarize_discussion', { summary: string }> - | OutputCompleteEvent<'review_generation', { valid: boolean; review: string }> | ChatCompletionChunkEvent | ChatCompletionMessageEvent; export type NlToEsqlTaskParams = { client: Pick; connectorId: string; - logger: Logger; + logger: Pick; functionCalling?: FunctionCallingMode; - summarizeInput?: boolean; } & TToolOptions & ({ input: string } | { messages: Message[] }); diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/extract_queries.test.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/extract_queries.test.ts deleted file mode 100644 index 78b17f0902991..0000000000000 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/extract_queries.test.ts +++ /dev/null @@ -1,75 +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. - */ - -import { extractQueries } from './extract_queries'; - -describe('extractQueries', () => { - it('finds a single query', () => { - const input = ` - Some text - - \`\`\`esql - FROM my-index - | WHERE foo == bar - | SORT foo - \`\`\` - `; - - const queries = extractQueries(input); - - expect(queries).toEqual([ - { - query: `FROM my-index - | WHERE foo == bar - | SORT foo`, - startPos: 20, - endPos: 91, - }, - ]); - }); - - it('finds multiple queries', () => { - const input = ` - Some text - - \`\`\`esql - FROM my-index - | WHERE foo == bar - | SORT foo - \`\`\` - - Another block of text - - \`\`\`esql - FROM kibana_sample_data_logs - | WHERE @timestamp > NOW() - 10 days - | STATS volume = COUNT(*) BY BUCKET(@timestamp, 1 day) - \`\`\` - - Some final block of text - `; - - const queries = extractQueries(input); - - expect(queries).toEqual([ - { - query: `FROM my-index - | WHERE foo == bar - | SORT foo`, - startPos: 20, - endPos: 91, - }, - { - query: `FROM kibana_sample_data_logs - | WHERE @timestamp > NOW() - 10 days - | STATS volume = COUNT(*) BY BUCKET(@timestamp, 1 day)`, - startPos: 124, - endPos: 272, - }, - ]); - }); -}); diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/extract_queries.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/extract_queries.ts deleted file mode 100644 index 1b6a715091a28..0000000000000 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/extract_queries.ts +++ /dev/null @@ -1,36 +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. - */ - -import { INLINE_ESQL_QUERY_REGEX } from '../../../../common/tasks/nl_to_esql/constants'; - -export interface QueryWithPosition { - /** the query text, without the ```esql``` */ - query: string; - /** the query start position in the text - INCLUDING the ```esql``` */ - startPos: number; - /** the query end position in the text - INCLUDING the ```esql``` */ - endPos: number; -} - -/** - * Extract the esql queries from a given string content. - * - * Note: it will only find queries wrapped with ```esql [query] ``` - */ -export const extractQueries = (text: string): QueryWithPosition[] => { - return Array.from(text.matchAll(INLINE_ESQL_QUERY_REGEX)) - .filter((match) => { - return match[1] !== undefined && match.index !== undefined; - }) - .map((match) => { - return { - query: match[1], - startPos: match.index!, - endPos: match.index! + match[0].length, - }; - }); -}; diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/index.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/index.ts deleted file mode 100644 index 1d2ba2d78ccc2..0000000000000 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/index.ts +++ /dev/null @@ -1,9 +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. - */ - -export { extractQueries, type QueryWithPosition } from './extract_queries'; -export { validateQueryAst } from './validate_query_ast'; diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/validate_query_ast.test.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/validate_query_ast.test.ts deleted file mode 100644 index 79bcda6f08e2e..0000000000000 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/validate_query_ast.test.ts +++ /dev/null @@ -1,31 +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. - */ - -import { validateQueryAst } from './validate_query_ast'; - -describe('validateQueryAst', () => { - it('returns AST validation errors', async () => { - const query = `FROM kibana_sample_data_logs - | KOUKOU foobar - | WHERE`; - - const errors = await validateQueryAst(query); - - expect(errors).toEqual([ - { - startPos: 36, - endPos: 42, - message: expect.stringContaining(`mismatched input 'KOUKOU'`), - }, - { - startPos: 61, - endPos: 61, - message: expect.stringContaining(`mismatched input ''`), - }, - ]); - }); -}); diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/validate_query_ast.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/validate_query_ast.ts deleted file mode 100644 index 7e70b14aa595c..0000000000000 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/utils/validate_query_ast.ts +++ /dev/null @@ -1,63 +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. - */ - -import { validateQuery } from '@kbn/esql-validation-autocomplete'; -import { getAstAndSyntaxErrors } from '@kbn/esql-ast'; -import { splitIntoCommands } from '../../../../common/tasks/nl_to_esql/correct_common_esql_mistakes'; - -interface QuerySyntaxError { - message: string; - startPos: number; - endPos: number; - code?: string; -} - -/** - * Run the query through the kibana - * @param query - */ -export const validateQueryAst = async (query: string): Promise => { - const { errors: astErrors } = await validateQuery(query, getAstAndSyntaxErrors, { - // setting this to true, we don't want to validate the index / fields existence - ignoreOnMissingCallbacks: true, - }); - - const asCommands = splitIntoCommands(query); - - const errors = (astErrors ?? []).map((error) => { - if ('location' in error) { - return { - message: error.text, - code: error.code, - startPos: error.location.min, - endPos: error.location.max, - }; - } else { - return { - message: error.message, - code: error.code, - startPos: getPosition(asCommands, error.startLineNumber, error.startColumn), - endPos: getPosition(asCommands, error.endLineNumber, error.endColumn), - }; - } - }); - - return errors; -}; - -const getPosition = ( - commands: ReturnType, - line: number, - column: number -): number => { - const previousCommands = commands.slice(0, line - 1); - return ( - previousCommands.reduce((count, command) => { - return count + command.command.length; - }, 0) + column - ); -}; diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/validate_esql_query.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/validate_esql_query.ts new file mode 100644 index 0000000000000..823344f52a891 --- /dev/null +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/validate_esql_query.ts @@ -0,0 +1,72 @@ +/* + * 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 { validateQuery } from '@kbn/esql-validation-autocomplete'; +import { getAstAndSyntaxErrors } from '@kbn/esql-ast'; +import type { ElasticsearchClient } from '@kbn/core/server'; +import { ESQLSearchResponse, ESQLRow } from '@kbn/es-types'; +import { esFieldTypeToKibanaFieldType } from '@kbn/field-types'; +import { DatatableColumn, DatatableColumnType } from '@kbn/expressions-plugin/common'; +import { splitIntoCommands } from '../../../common/tasks/nl_to_esql/correct_common_esql_mistakes'; + +export async function runAndValidateEsqlQuery({ + query, + client, +}: { + query: string; + client: ElasticsearchClient; +}): Promise<{ + columns?: DatatableColumn[]; + rows?: ESQLRow[]; + error?: Error; + errorMessages?: string[]; +}> { + const { errors } = await validateQuery(query, getAstAndSyntaxErrors, { + // setting this to true, we don't want to validate the index / fields existence + ignoreOnMissingCallbacks: true, + }); + + const asCommands = splitIntoCommands(query); + + const errorMessages = errors?.map((error) => { + if ('location' in error) { + const commandsUntilEndOfError = splitIntoCommands(query.substring(0, error.location.max)); + const lastCompleteCommand = asCommands[commandsUntilEndOfError.length - 1]; + if (lastCompleteCommand) { + return `Error in ${lastCompleteCommand.command}\n: ${error.text}`; + } + } + return 'text' in error ? error.text : error.message; + }); + + return client.transport + .request({ + method: 'POST', + path: '_query', + body: { + query, + }, + }) + .then((res) => { + const esqlResponse = res as ESQLSearchResponse; + + const columns = + esqlResponse.columns?.map(({ name, type }) => ({ + id: name, + name, + meta: { type: esFieldTypeToKibanaFieldType(type) as DatatableColumnType }, + })) ?? []; + + return { columns, rows: esqlResponse.values }; + }) + .catch((error) => { + return { + error, + ...(errorMessages.length ? { errorMessages } : {}), + }; + }); +} From f194638951f240f966332f5c4d46938aeee65686 Mon Sep 17 00:00:00 2001 From: pgayvallet Date: Thu, 31 Oct 2024 13:31:29 +0100 Subject: [PATCH 18/21] fix imports again due to revert --- .../tasks/nl_to_esql/actions/generate_esql.ts | 16 +++++++++------- .../nl_to_esql/actions/request_documentation.ts | 12 ++++++++---- .../server/tasks/nl_to_esql/actions/shared.ts | 2 +- .../inference/server/tasks/nl_to_esql/task.ts | 3 +-- .../inference/server/tasks/nl_to_esql/types.ts | 6 +++--- 5 files changed, 22 insertions(+), 17 deletions(-) diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/generate_esql.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/generate_esql.ts index d31952e2f5252..26a8fb63ce013 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/generate_esql.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/generate_esql.ts @@ -7,21 +7,23 @@ import { Observable, map, merge, of, switchMap } from 'rxjs'; import type { Logger } from '@kbn/logging'; -import { ToolCall, ToolOptions } from '../../../../common/chat_complete/tools'; import { - correctCommonEsqlMistakes, - generateFakeToolCallId, + ToolCall, + ToolOptions, + withoutTokenCountEvents, isChatCompletionMessageEvent, Message, MessageRole, -} from '../../../../common'; -import { InferenceClient, withoutTokenCountEvents } from '../../..'; -import { OutputCompleteEvent, OutputEventType } from '../../../../common/output'; + OutputCompleteEvent, + OutputEventType, + FunctionCallingMode, +} from '@kbn/inference-common'; +import { correctCommonEsqlMistakes, generateFakeToolCallId } from '../../../../common'; +import { InferenceClient } from '../../..'; import { INLINE_ESQL_QUERY_REGEX } from '../../../../common/tasks/nl_to_esql/constants'; import { EsqlDocumentBase } from '../doc_base'; import { requestDocumentationSchema } from './shared'; import type { NlToEsqlTaskEvent } from '../types'; -import type { FunctionCallingMode } from '../../../../common/chat_complete'; export const generateEsqlTask = ({ chatCompleteApi, diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/request_documentation.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/request_documentation.ts index d4eb3060f59bb..aea428208be1d 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/request_documentation.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/request_documentation.ts @@ -6,11 +6,15 @@ */ import { isEmpty } from 'lodash'; -import { InferenceClient, withoutOutputUpdateEvents } from '../../..'; -import { Message } from '../../../../common'; -import { ToolChoiceType, ToolOptions } from '../../../../common/chat_complete/tools'; +import { + ToolChoiceType, + ToolOptions, + Message, + withoutOutputUpdateEvents, + FunctionCallingMode, +} from '@kbn/inference-common'; +import { InferenceClient } from '../../..'; import { requestDocumentationSchema } from './shared'; -import type { FunctionCallingMode } from '../../../../common/chat_complete'; export const requestDocumentation = ({ outputApi, diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/shared.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/shared.ts index f0fc796173b23..60114188ea37f 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/shared.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/actions/shared.ts @@ -5,7 +5,7 @@ * 2.0. */ -import { ToolSchema } from '../../../../common'; +import { ToolSchema } from '@kbn/inference-common'; export const requestDocumentationSchema = { type: 'object', diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/task.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/task.ts index e0c5a838ea148..56c48b73f4994 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/task.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/task.ts @@ -7,8 +7,7 @@ import { once } from 'lodash'; import { Observable, from, switchMap } from 'rxjs'; -import { Message, MessageRole } from '../../../common/chat_complete'; -import type { ToolOptions } from '../../../common/chat_complete/tools'; +import { Message, MessageRole, ToolOptions } from '@kbn/inference-common'; import { EsqlDocumentBase } from './doc_base'; import { requestDocumentation, generateEsqlTask } from './actions'; import { NlToEsqlTaskParams, NlToEsqlTaskEvent } from './types'; diff --git a/x-pack/plugins/inference/server/tasks/nl_to_esql/types.ts b/x-pack/plugins/inference/server/tasks/nl_to_esql/types.ts index a0bcd635081ea..ce45d9a15e4b3 100644 --- a/x-pack/plugins/inference/server/tasks/nl_to_esql/types.ts +++ b/x-pack/plugins/inference/server/tasks/nl_to_esql/types.ts @@ -11,9 +11,9 @@ import type { ChatCompletionMessageEvent, FunctionCallingMode, Message, -} from '../../../common/chat_complete'; -import type { ToolOptions } from '../../../common/chat_complete/tools'; -import type { OutputCompleteEvent } from '../../../common/output'; + ToolOptions, + OutputCompleteEvent, +} from '@kbn/inference-common'; import type { InferenceClient } from '../../types'; export type NlToEsqlTaskEvent = From 08d1e30d329a620d708529b5c5922a4b9becd6ea Mon Sep 17 00:00:00 2001 From: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Date: Thu, 31 Oct 2024 12:41:49 +0000 Subject: [PATCH 19/21] [CI] Auto-commit changed files from 'node scripts/yarn_deduplicate' --- x-pack/plugins/inference/tsconfig.json | 3 +++ 1 file changed, 3 insertions(+) diff --git a/x-pack/plugins/inference/tsconfig.json b/x-pack/plugins/inference/tsconfig.json index 9bf696c45ccb7..92327007829a9 100644 --- a/x-pack/plugins/inference/tsconfig.json +++ b/x-pack/plugins/inference/tsconfig.json @@ -33,5 +33,8 @@ "@kbn/actions-plugin", "@kbn/config-schema", "@kbn/inference-common", + "@kbn/es-types", + "@kbn/field-types", + "@kbn/expressions-plugin", ] } From 9d858a25a5b1f32b82afb028a6ded80b00797dd6 Mon Sep 17 00:00:00 2001 From: pgayvallet Date: Thu, 31 Oct 2024 15:19:37 +0100 Subject: [PATCH 20/21] Add TSdoc for most public types --- .../ai-infra/inference-common/index.ts | 2 + .../inference-common/src/chat_complete/api.ts | 15 +--- .../src/chat_complete/errors.ts | 27 ++++++ .../src/chat_complete/event_utils.ts | 26 +++++- .../src/chat_complete/events.ts | 57 +++++++++++- .../src/chat_complete/index.ts | 1 + .../src/chat_complete/messages.ts | 46 +++++++++- .../src/chat_complete/tool_schema.ts | 6 ++ .../src/chat_complete/tools.ts | 87 ++++++++++++++++--- .../ai-infra/inference-common/src/errors.ts | 6 ++ .../inference-common/src/inference_task.ts | 6 ++ .../inference-common/src/output/api.ts | 12 ++- .../src/output/event_utils.ts | 16 +++- .../inference-common/src/output/events.ts | 54 +++++++++--- .../inference-common/src/output/index.ts | 2 +- 15 files changed, 318 insertions(+), 45 deletions(-) diff --git a/x-pack/packages/ai-infra/inference-common/index.ts b/x-pack/packages/ai-infra/inference-common/index.ts index 283cac2594c29..6de7ce3bb8008 100644 --- a/x-pack/packages/ai-infra/inference-common/index.ts +++ b/x-pack/packages/ai-infra/inference-common/index.ts @@ -36,6 +36,7 @@ export { isChatCompletionMessageEvent, isChatCompletionEvent, isChatCompletionChunkEvent, + isChatCompletionTokenCountEvent, ChatCompletionErrorCode, type ChatCompletionToolNotFoundError, type ChatCompletionToolValidationError, @@ -47,6 +48,7 @@ export { export { OutputEventType, type OutputAPI, + type OutputResponse, type OutputCompleteEvent, type OutputUpdateEvent, type Output, diff --git a/x-pack/packages/ai-infra/inference-common/src/chat_complete/api.ts b/x-pack/packages/ai-infra/inference-common/src/chat_complete/api.ts index a872b3f950527..c6ffa9d4c8d5d 100644 --- a/x-pack/packages/ai-infra/inference-common/src/chat_complete/api.ts +++ b/x-pack/packages/ai-infra/inference-common/src/chat_complete/api.ts @@ -24,18 +24,6 @@ import type { ChatCompletionEvent } from './events'; * { role: MessageRole.User, content: "Another question?"}, * ] * }); - * - * ``` - * @example directly accessing the LLM message without using the observable API - * ```ts - * const message = await chatComplete({ - * connectorId: 'my-connector', - * system: "You are a helpful assistant", - * messages: [ - * { role: MessageRole.User, content: "Some question?"}, - * ] - * }).getMessage(); - * ``` */ export type ChatCompleteAPI = ( options: ChatCompleteOptions @@ -46,7 +34,8 @@ export type ChatCompleteAPI = ( */ export type ChatCompleteOptions = { /** - * The ID of the connector to use + * The ID of the connector to use. + * Must be a genAI compatible connector, or an error will be thrown. */ connectorId: string; /** diff --git a/x-pack/packages/ai-infra/inference-common/src/chat_complete/errors.ts b/x-pack/packages/ai-infra/inference-common/src/chat_complete/errors.ts index b4c2d6f0a0890..b9d859a666761 100644 --- a/x-pack/packages/ai-infra/inference-common/src/chat_complete/errors.ts +++ b/x-pack/packages/ai-infra/inference-common/src/chat_complete/errors.ts @@ -8,12 +8,19 @@ import { InferenceTaskError } from '../errors'; import type { UnvalidatedToolCall } from './tools'; +/** + * List of code of error that are specific to the {@link ChatCompleteAPI} + */ export enum ChatCompletionErrorCode { TokenLimitReachedError = 'tokenLimitReachedError', ToolNotFoundError = 'toolNotFoundError', ToolValidationError = 'toolValidationError', } +/** + * Error thrown if the completion call fails because of a token limit + * error, e.g. when the context window is higher than the limit + */ export type ChatCompletionTokenLimitReachedError = InferenceTaskError< ChatCompletionErrorCode.TokenLimitReachedError, { @@ -22,13 +29,24 @@ export type ChatCompletionTokenLimitReachedError = InferenceTaskError< } >; +/** + * Error thrown if the LLM called a tool that was not provided + * in the list of available tools. + */ export type ChatCompletionToolNotFoundError = InferenceTaskError< ChatCompletionErrorCode.ToolNotFoundError, { + /** The name of the tool that got called */ name: string; } >; +/** + * Error thrown when the LLM called a tool with parameters that + * don't match the tool's schema. + * + * The level of details on the error vary depending on the underlying LLM. + */ export type ChatCompletionToolValidationError = InferenceTaskError< ChatCompletionErrorCode.ToolValidationError, { @@ -39,6 +57,9 @@ export type ChatCompletionToolValidationError = InferenceTaskError< } >; +/** + * Check if an error is a {@link ChatCompletionToolValidationError} + */ export function isToolValidationError(error?: Error): error is ChatCompletionToolValidationError { return ( error instanceof InferenceTaskError && @@ -46,6 +67,9 @@ export function isToolValidationError(error?: Error): error is ChatCompletionToo ); } +/** + * Check if an error is a {@link ChatCompletionTokenLimitReachedError} + */ export function isTokenLimitReachedError( error: Error ): error is ChatCompletionTokenLimitReachedError { @@ -55,6 +79,9 @@ export function isTokenLimitReachedError( ); } +/** + * Check if an error is a {@link ChatCompletionToolNotFoundError} + */ export function isToolNotFoundError(error: Error): error is ChatCompletionToolNotFoundError { return ( error instanceof InferenceTaskError && error.code === ChatCompletionErrorCode.ToolNotFoundError diff --git a/x-pack/packages/ai-infra/inference-common/src/chat_complete/event_utils.ts b/x-pack/packages/ai-infra/inference-common/src/chat_complete/event_utils.ts index bbe812c89d4df..4749673264aff 100644 --- a/x-pack/packages/ai-infra/inference-common/src/chat_complete/event_utils.ts +++ b/x-pack/packages/ai-infra/inference-common/src/chat_complete/event_utils.ts @@ -16,18 +16,36 @@ import { } from './events'; import type { ToolOptions } from './tools'; +/** + * Check if the provided {@link ChatCompletionEvent} is a {@link ChatCompletionChunkEvent} + */ export function isChatCompletionChunkEvent( event: ChatCompletionEvent ): event is ChatCompletionChunkEvent { return event.type === ChatCompletionEventType.ChatCompletionChunk; } -export function isChatCompletionMessageEvent>( +/** + * Check if the provided {@link ChatCompletionEvent} is a {@link ChatCompletionMessageEvent} + */ +export function isChatCompletionMessageEvent( event: ChatCompletionEvent ): event is ChatCompletionMessageEvent { return event.type === ChatCompletionEventType.ChatCompletionMessage; } +/** + * Check if the provided {@link ChatCompletionEvent} is a {@link ChatCompletionMessageEvent} + */ +export function isChatCompletionTokenCountEvent( + event: ChatCompletionEvent +): event is ChatCompletionTokenCountEvent { + return event.type === ChatCompletionEventType.ChatCompletionTokenCount; +} + +/** + * Check if the provided {@link InferenceTaskEvent} is a {@link ChatCompletionEvent} + */ export function isChatCompletionEvent(event: InferenceTaskEvent): event is ChatCompletionEvent { return ( event.type === ChatCompletionEventType.ChatCompletionChunk || @@ -36,6 +54,9 @@ export function isChatCompletionEvent(event: InferenceTaskEvent): event is ChatC ); } +/** + * Operator filtering out the chunk events from the provided observable. + */ export function withoutChunkEvents(): OperatorFunction< T, Exclude @@ -46,6 +67,9 @@ export function withoutChunkEvents(): OperatorFun ); } +/** + * Operator filtering out the token count events from the provided observable. + */ export function withoutTokenCountEvents(): OperatorFunction< T, Exclude diff --git a/x-pack/packages/ai-infra/inference-common/src/chat_complete/events.ts b/x-pack/packages/ai-infra/inference-common/src/chat_complete/events.ts index c6f5766405c40..92c49e6ee7fc0 100644 --- a/x-pack/packages/ai-infra/inference-common/src/chat_complete/events.ts +++ b/x-pack/packages/ai-infra/inference-common/src/chat_complete/events.ts @@ -8,37 +8,92 @@ import type { InferenceTaskEventBase } from '../inference_task'; import type { ToolCallsOf, ToolOptions } from './tools'; +/** + * List possible values of {@link ChatCompletionEvent} types. + */ export enum ChatCompletionEventType { ChatCompletionChunk = 'chatCompletionChunk', ChatCompletionTokenCount = 'chatCompletionTokenCount', ChatCompletionMessage = 'chatCompletionMessage', } +/** + * Message event, sent only once, after all the chunks were emitted, and containing + * the whole text content and potential tool calls of the response. + */ export type ChatCompletionMessageEvent = InferenceTaskEventBase & { + /** + * The text content of the LLM response. + */ content: string; - } & { toolCalls: ToolCallsOf['toolCalls'] }; + /** + * The eventual tool calls performed by the LLM. + */ + toolCalls: ToolCallsOf['toolCalls']; + }; +/** + * Represent a partial tool call present in a chunk event. + * + * Note that all properties of the structure, except from the index, + * are partial and must be aggregated. + */ export interface ChatCompletionChunkToolCall { + /** + * The tool call index (position in the tool call array). + */ index: number; + /** + * chunk of tool call id. + */ toolCallId: string; function: { + /** + * chunk of tool name. + */ name: string; + /** + * chunk of tool call arguments. + */ arguments: string; }; } +/** + * Chunk event, containing a fragment of the total content, + * and potentially chunks of tool calls. + */ export type ChatCompletionChunkEvent = InferenceTaskEventBase & { + /** + * The content chunk + */ content: string; + /** + * The tool call chunks + */ tool_calls: ChatCompletionChunkToolCall[]; }; +/** + * Token count event, send only once, usually (but not necessarily) + * before the message event + */ export type ChatCompletionTokenCountEvent = InferenceTaskEventBase & { tokens: { + /** + * Input token count + */ prompt: number; + /** + * Output token count + */ completion: number; + /** + * Total token count + */ total: number; }; }; diff --git a/x-pack/packages/ai-infra/inference-common/src/chat_complete/index.ts b/x-pack/packages/ai-infra/inference-common/src/chat_complete/index.ts index 6bf73f605c49e..8199af4cf068b 100644 --- a/x-pack/packages/ai-infra/inference-common/src/chat_complete/index.ts +++ b/x-pack/packages/ai-infra/inference-common/src/chat_complete/index.ts @@ -40,6 +40,7 @@ export { isChatCompletionChunkEvent, isChatCompletionEvent, isChatCompletionMessageEvent, + isChatCompletionTokenCountEvent, withoutChunkEvents, withoutTokenCountEvents, } from './event_utils'; diff --git a/x-pack/packages/ai-infra/inference-common/src/chat_complete/messages.ts b/x-pack/packages/ai-infra/inference-common/src/chat_complete/messages.ts index df29d9580700b..ca74b094e0a3b 100644 --- a/x-pack/packages/ai-infra/inference-common/src/chat_complete/messages.ts +++ b/x-pack/packages/ai-infra/inference-common/src/chat_complete/messages.ts @@ -7,27 +7,69 @@ import type { ToolCall } from './tools'; +/** + * Enum for all possible {@link Message} roles. + */ export enum MessageRole { User = 'user', Assistant = 'assistant', Tool = 'tool', } +/** + * Base type for all subtypes of {@link Message}. + */ interface MessageBase { role: TRole; } -export type UserMessage = MessageBase & { content: string }; +/** + * Represents a message from the user. + */ +export type UserMessage = MessageBase & { + /** + * The text content of the user message + */ + content: string; +}; +/** + * Represents a message from the LLM. + */ export type AssistantMessage = MessageBase & { + /** + * The text content of the message. + * Can be null if the LLM called a tool. + */ content: string | null; - toolCalls?: Array | undefined>>; + /** + * A potential list of {@ToolCall} the LLM asked to execute. + * Note that LLM with parallel tool invocation can potentially call multiple tools at the same time. + */ + toolCalls?: ToolCall[]; }; +/** + * Represents a tool invocation result, following a request from the LLM to execute a tool. + */ export type ToolMessage | unknown> = MessageBase & { + /** + * The call id matching the {@link ToolCall} this tool message is for. + */ toolCallId: string; + /** + * The response from the tool invocation. + */ response: TToolResponse; }; +/** + * Mixin composed of all the possible types of messages in a chatComplete discussion. + * + * Message can be of three types: + * - {@link UserMessage} + * - {@link AssistantMessage} + * - {@link ToolMessage} + */ export type Message = UserMessage | AssistantMessage | ToolMessage; diff --git a/x-pack/packages/ai-infra/inference-common/src/chat_complete/tool_schema.ts b/x-pack/packages/ai-infra/inference-common/src/chat_complete/tool_schema.ts index 2a2c61f8e9b70..bb4742c6b74d9 100644 --- a/x-pack/packages/ai-infra/inference-common/src/chat_complete/tool_schema.ts +++ b/x-pack/packages/ai-infra/inference-common/src/chat_complete/tool_schema.ts @@ -72,8 +72,14 @@ type FromToolSchemaString = ? ValuesType : string; +/** + * Defines the schema for a {@link ToolDefinition} + */ export type ToolSchema = ToolSchemaTypeObject; +/** + * Utility type to infer the shape of a tool call from its schema. + */ export type FromToolSchema = TToolSchema extends ToolSchemaTypeObject ? FromToolSchemaObject diff --git a/x-pack/packages/ai-infra/inference-common/src/chat_complete/tools.ts b/x-pack/packages/ai-infra/inference-common/src/chat_complete/tools.ts index a5db86c7c996d..0c7d5c6755f31 100644 --- a/x-pack/packages/ai-infra/inference-common/src/chat_complete/tools.ts +++ b/x-pack/packages/ai-infra/inference-common/src/chat_complete/tools.ts @@ -4,15 +4,12 @@ * 2.0; you may not use this file except in compliance with the Elastic License * 2.0. */ + import type { ValuesType } from 'utility-types'; import { FromToolSchema, ToolSchema } from './tool_schema'; type Assert = TValue extends TType ? TValue & TType : never; -interface CustomToolChoice { - function: TName; -} - type ToolsOfChoice = TToolOptions['toolChoice'] extends { function: infer TToolName; } @@ -21,6 +18,9 @@ type ToolsOfChoice = TToolOptions['toolChoice' : TToolOptions['tools'] : TToolOptions['tools']; +/** + * Utility type to infer the tool calls response shape. + */ type ToolResponsesOf | undefined> = TTools extends Record ? Array< @@ -30,18 +30,64 @@ type ToolResponsesOf | undefined> > : never[]; +/** + * Utility type to infer the tool call response shape. + */ type ToolResponseOf = ToolCall< TName, TToolDefinition extends { schema: ToolSchema } ? FromToolSchema : {} >; +/** + * Tool invocation choice type. + * + * Refer to {@link ToolChoice} for more details. + */ +export enum ToolChoiceType { + none = 'none', + auto = 'auto', + required = 'required', +} + +/** + * Represent a tool choice where the LLM is forced to call a specific tool. + * + * Refer to {@link ToolChoice} for more details. + */ +interface CustomToolChoice { + function: TName; +} + +/** + * Defines the tool invocation for {@link ToolOptions}, either a {@link ToolChoiceType} or {@link CustomToolChoice}. + * - {@link ToolChoiceType.none}: the LLM will never call a tool + * - {@link ToolChoiceType.auto}: the LLM will decide if it should call a tool or provide a text response + * - {@link ToolChoiceType.required}: the LLM will always call a tool, but will decide with one to call + * - {@link CustomToolChoice}: the LLM will always call the specified tool + */ export type ToolChoice = ToolChoiceType | CustomToolChoice; +/** + * The definition of a tool that will be provided to the LLM for it to eventually call. + */ export interface ToolDefinition { + /** + * A description of what the tool does. Note that this will be exposed to the LLM, + * so the description should be explicit about what the tool does and when to call it. + */ description: string; + /** + * The input schema for the tool, representing the shape of the tool's parameters + * + * Even if optional, it is highly recommended to define a schema for all tool definitions, unless + * the tool is supposed to be called without parameters. + */ schema?: ToolSchema; } +/** + * Utility type to infer the toolCall type of {@link ChatCompletionMessageEvent}. + */ export type ToolCallsOf = TToolOptions extends { tools?: Record; } @@ -52,12 +98,11 @@ export type ToolCallsOf = TToolOptions extends } : { toolCalls: never }; -export enum ToolChoiceType { - none = 'none', - auto = 'auto', - required = 'required', -} - +/** + * Represents a tool call from the LLM before correctly converted to the schema type. + * + * Only publicly exposed because referenced by {@link ChatCompletionToolValidationError} + */ export interface UnvalidatedToolCall { toolCallId: string; function: { @@ -66,17 +111,39 @@ export interface UnvalidatedToolCall { }; } +/** + * Represents a tool call performed by the LLM. + */ export interface ToolCall< TName extends string = string, TArguments extends Record | undefined = Record | undefined > { + /** + * The id of the tool call, that must be re-used when providing the tool call response + */ toolCallId: string; function: { + /** + * The name of the tool that was called + */ name: TName; } & (TArguments extends Record ? { arguments: TArguments } : {}); } +/** + * Tool-related parameters of {@link ChatCompleteAPI} + */ export interface ToolOptions { + /** + * The choice of tool execution. + * + * Refer to {@link ToolChoice} + */ toolChoice?: ToolChoice; + /** + * The list of tool definitions that will be exposed to the LLM. + * + * Refer to {@link ToolDefinition}. + */ tools?: Record; } diff --git a/x-pack/packages/ai-infra/inference-common/src/errors.ts b/x-pack/packages/ai-infra/inference-common/src/errors.ts index acf8d34b4275a..5a99adc4321d9 100644 --- a/x-pack/packages/ai-infra/inference-common/src/errors.ts +++ b/x-pack/packages/ai-infra/inference-common/src/errors.ts @@ -7,11 +7,17 @@ import { InferenceTaskEventBase, InferenceTaskEventType } from './inference_task'; +/** + * Enum for generic inference error codes. + */ export enum InferenceTaskErrorCode { internalError = 'internalError', requestError = 'requestError', } +/** + * Base class for all inference API errors. + */ export class InferenceTaskError< TCode extends string, TMeta extends Record | undefined diff --git a/x-pack/packages/ai-infra/inference-common/src/inference_task.ts b/x-pack/packages/ai-infra/inference-common/src/inference_task.ts index 7b8f65b7af2c9..15449e1275a5b 100644 --- a/x-pack/packages/ai-infra/inference-common/src/inference_task.ts +++ b/x-pack/packages/ai-infra/inference-common/src/inference_task.ts @@ -5,7 +5,13 @@ * 2.0. */ +/** + * Base interface for all inference events. + */ export interface InferenceTaskEventBase { + /** + * Unique identifier of the event type. + */ type: TEventType; } diff --git a/x-pack/packages/ai-infra/inference-common/src/output/api.ts b/x-pack/packages/ai-infra/inference-common/src/output/api.ts index f6bb530185d41..677d2f7015c2a 100644 --- a/x-pack/packages/ai-infra/inference-common/src/output/api.ts +++ b/x-pack/packages/ai-infra/inference-common/src/output/api.ts @@ -31,6 +31,16 @@ export type OutputAPI = < previousMessages?: Message[]; functionCalling?: FunctionCallingMode; } -) => Observable< +) => OutputResponse; + +/** + * Response from the {@link OutputAPI}. + * + * Observable of {@link OutputEvent} + */ +export type OutputResponse< + TId extends string = string, + TOutputSchema extends ToolSchema | undefined = ToolSchema | undefined +> = Observable< OutputEvent : undefined> >; diff --git a/x-pack/packages/ai-infra/inference-common/src/output/event_utils.ts b/x-pack/packages/ai-infra/inference-common/src/output/event_utils.ts index b005b67d8a6dc..1139bac92c610 100644 --- a/x-pack/packages/ai-infra/inference-common/src/output/event_utils.ts +++ b/x-pack/packages/ai-infra/inference-common/src/output/event_utils.ts @@ -6,27 +6,39 @@ */ import { filter, OperatorFunction } from 'rxjs'; -import { OutputEvent, OutputEventType, OutputUpdateEvent } from '.'; +import { OutputCompleteEvent, OutputEvent, OutputEventType, OutputUpdateEvent } from '.'; import type { InferenceTaskEvent } from '../inference_task'; +/** + * Check if the provided {@link ChatCompletionEvent} is a {@link ChatCompletionChunkEvent} + */ export function isOutputCompleteEvent( event: TOutputEvent -): event is Exclude { +): event is Extract { return event.type === OutputEventType.OutputComplete; } +/** + * Check if the provided {@link InferenceTaskEvent} is a {@link OutputEvent} + */ export function isOutputEvent(event: InferenceTaskEvent): event is OutputEvent { return ( event.type === OutputEventType.OutputComplete || event.type === OutputEventType.OutputUpdate ); } +/** + * Check if the provided {@link OutputEvent} is a {@link OutputUpdateEvent} + */ export function isOutputUpdateEvent( event: OutputEvent ): event is OutputUpdateEvent { return event.type === OutputEventType.OutputComplete; } +/** + * Operator filtering out the update events from the provided observable. + */ export function withoutOutputUpdateEvents(): OperatorFunction< T, Exclude diff --git a/x-pack/packages/ai-infra/inference-common/src/output/events.ts b/x-pack/packages/ai-infra/inference-common/src/output/events.ts index d95b821bec6de..794f58bd7db79 100644 --- a/x-pack/packages/ai-infra/inference-common/src/output/events.ts +++ b/x-pack/packages/ai-infra/inference-common/src/output/events.ts @@ -5,35 +5,61 @@ * 2.0. */ -import { ServerSentEventBase } from '@kbn/sse-utils'; +import { InferenceTaskEventBase } from '../inference_task'; +/** + * List possible values of {@link OutputEvent} types. + */ export enum OutputEventType { OutputUpdate = 'output', OutputComplete = 'complete', } +/** + * Task output of a {@link OutputCompleteEvent} + */ export type Output = Record | undefined | unknown; -export type OutputUpdateEvent = ServerSentEventBase< - OutputEventType.OutputUpdate, - { +/** + * Update (chunk) event for the {@link OutputAPI} + */ +export type OutputUpdateEvent = + InferenceTaskEventBase & { + /** + * The id of the operation, as provided as input + */ id: TId; + /** + * The text content of the chunk + */ content: string; - } ->; + }; +/** + * Completion (complete message) event for the {@link OutputAPI} + */ export type OutputCompleteEvent< TId extends string = string, TOutput extends Output = Output -> = ServerSentEventBase< - OutputEventType.OutputComplete, - { - id: TId; - output: TOutput; - content: string; - } ->; +> = InferenceTaskEventBase & { + /** + * The id of the operation, as provided as input + */ + id: TId; + /** + * The task output, following the schema specified as input + */ + output: TOutput; + /** + * Potential text content provided by the LLM, + * if it was provided in addition to the tool call + */ + content: string; +}; +/** + * Events emitted from the {@link OutputEvent}. + */ export type OutputEvent = | OutputUpdateEvent | OutputCompleteEvent; diff --git a/x-pack/packages/ai-infra/inference-common/src/output/index.ts b/x-pack/packages/ai-infra/inference-common/src/output/index.ts index 7b294af455dd2..ceac178f47faa 100644 --- a/x-pack/packages/ai-infra/inference-common/src/output/index.ts +++ b/x-pack/packages/ai-infra/inference-common/src/output/index.ts @@ -5,7 +5,7 @@ * 2.0. */ -export type { OutputAPI } from './api'; +export type { OutputAPI, OutputResponse } from './api'; export { OutputEventType, type OutputCompleteEvent, From 798a60bf2e36ab8799a27cf3ba877d851db7607d Mon Sep 17 00:00:00 2001 From: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Date: Thu, 31 Oct 2024 14:31:04 +0000 Subject: [PATCH 21/21] [CI] Auto-commit changed files from 'node scripts/notice' --- x-pack/packages/ai-infra/inference-common/tsconfig.json | 1 - 1 file changed, 1 deletion(-) diff --git a/x-pack/packages/ai-infra/inference-common/tsconfig.json b/x-pack/packages/ai-infra/inference-common/tsconfig.json index 7f261ba1dcb82..86d57b8d692f7 100644 --- a/x-pack/packages/ai-infra/inference-common/tsconfig.json +++ b/x-pack/packages/ai-infra/inference-common/tsconfig.json @@ -16,6 +16,5 @@ "target/**/*" ], "kbn_references": [ - "@kbn/sse-utils", ] }