diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/common/conversation_complete.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/common/conversation_complete.ts index 3c4e2cd609f8b..29818937b74a4 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/common/conversation_complete.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/common/conversation_complete.ts @@ -6,6 +6,7 @@ */ import { i18n } from '@kbn/i18n'; +import type { ChatCompletionChunkEvent as InferenceChatCompletionChunkEvent } from '@kbn/inference-common'; import { TokenCount as TokenCountType, type Message } from './types'; export enum StreamingChatResponseEventType { @@ -105,14 +106,19 @@ export type StreamingChatResponseEvent = | MessageAddEvent | ChatCompletionErrorEvent | TokenCountEvent - | BufferFlushEvent; + | BufferFlushEvent + | InferenceChatCompletionChunkEvent; export type StreamingChatResponseEventWithoutError = Exclude< StreamingChatResponseEvent, ChatCompletionErrorEvent >; -export type ChatEvent = ChatCompletionChunkEvent | TokenCountEvent; +export type ChatEvent = + | InferenceChatCompletionChunkEvent + | ChatCompletionChunkEvent + | TokenCountEvent; + export type MessageOrChatEvent = ChatEvent | MessageAddEvent; export enum ChatCompletionErrorCode { 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/common/convert_messages_for_inference.ts similarity index 96% rename from x-pack/plugins/observability_solution/observability_ai_assistant_app/common/convert_messages_for_inference.ts rename to x-pack/plugins/observability_solution/observability_ai_assistant/common/convert_messages_for_inference.ts index 7ab9516440988..974b002ea93c6 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/common/convert_messages_for_inference.ts @@ -5,13 +5,13 @@ * 2.0. */ -import { Message, MessageRole } from '@kbn/observability-ai-assistant-plugin/common'; import { AssistantMessage, Message as InferenceMessage, MessageRole as InferenceMessageRole, } from '@kbn/inference-common'; import { generateFakeToolCallId } from '@kbn/inference-plugin/common'; +import { Message, MessageRole } from '.'; export function convertMessagesForInference(messages: Message[]): InferenceMessage[] { const inferenceMessages: InferenceMessage[] = []; diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/common/utils/concatenate_chat_completion_chunks.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/common/utils/concatenate_chat_completion_chunks.ts index bead0974b91a3..692c4d727285a 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/common/utils/concatenate_chat_completion_chunks.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/common/utils/concatenate_chat_completion_chunks.ts @@ -6,10 +6,11 @@ */ import { cloneDeep } from 'lodash'; -import { type Observable, scan } from 'rxjs'; +import { type Observable, scan, filter } from 'rxjs'; +import type { ChatCompletionChunkEvent as InferenceChatCompletionChunkEvent } from '@kbn/inference-common'; import type { ChatCompletionChunkEvent } from '../conversation_complete'; +import { StreamingChatResponseEventType } from '../conversation_complete'; import { MessageRole } from '../types'; - export interface ConcatenatedMessage { message: { content: string; @@ -24,14 +25,21 @@ export interface ConcatenatedMessage { export const concatenateChatCompletionChunks = () => - (source: Observable): Observable => + ( + source: Observable + ): Observable => source.pipe( + filter( + (event): event is InferenceChatCompletionChunkEvent => + event.type === StreamingChatResponseEventType.ChatCompletionChunk + ), scan( - (acc, { message }) => { - acc.message.content += message.content ?? ''; - acc.message.function_call.name += message.function_call?.name ?? ''; - acc.message.function_call.arguments += message.function_call?.arguments ?? ''; - + (acc, event) => { + acc.message.content += event.content ?? ''; + if (event.tool_calls.length > 0) { + acc.message.function_call.name += event.tool_calls[0].function.name ?? ''; + acc.message.function_call.arguments += event.tool_calls[0].function.arguments ?? ''; + } return cloneDeep(acc); }, { diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/common/utils/emit_with_concatenated_message.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/common/utils/emit_with_concatenated_message.ts index 47370cc48cf00..86bc13f9d397b 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/common/utils/emit_with_concatenated_message.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/common/utils/emit_with_concatenated_message.ts @@ -15,6 +15,7 @@ import { shareReplay, withLatestFrom, } from 'rxjs'; +import { ChatCompletionChunkEvent as InferenceChatCompletionChunkEvent } from '@kbn/inference-common'; import { withoutTokenCountEvents } from './without_token_count_events'; import { ChatCompletionChunkEvent, @@ -33,14 +34,14 @@ type ConcatenateMessageCallback = ( function mergeWithEditedMessage( originalMessage: ConcatenatedMessage, - chunkEvent: ChatCompletionChunkEvent, + chunkEvent: ChatCompletionChunkEvent | InferenceChatCompletionChunkEvent, callback?: ConcatenateMessageCallback ): Observable { return from( (callback ? callback(originalMessage) : Promise.resolve(originalMessage)).then((message) => { const next: MessageAddEvent = { type: StreamingChatResponseEventType.MessageAdd as const, - id: chunkEvent.id, + id: 'id' in chunkEvent ? chunkEvent.id : '', message: { '@timestamp': new Date().toISOString(), ...message, diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/kibana.jsonc b/x-pack/plugins/observability_solution/observability_ai_assistant/kibana.jsonc index e7a6a905a8bd2..ed106c9b6a791 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/kibana.jsonc +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/kibana.jsonc @@ -1,38 +1,26 @@ { "type": "plugin", "id": "@kbn/observability-ai-assistant-plugin", - "owner": [ - "@elastic/obs-ai-assistant" - ], + "owner": ["@elastic/obs-ai-assistant"], "group": "platform", "visibility": "shared", "plugin": { "id": "observabilityAIAssistant", "browser": true, "server": true, - "configPath": [ - "xpack", - "observabilityAIAssistant" - ], + "configPath": ["xpack", "observabilityAIAssistant"], "requiredPlugins": [ "actions", "features", "licensing", "security", "taskManager", - "dataViews" - ], - "optionalPlugins": [ - "cloud", - "serverless" - ], - "requiredBundles": [ - "kibanaReact", - "kibanaUtils" - ], - "runtimePluginDependencies": [ - "ml" + "dataViews", + "inference" ], + "optionalPlugins": ["cloud", "serverless"], + "requiredBundles": ["kibanaReact", "kibanaUtils"], + "runtimePluginDependencies": ["ml"], "extraPublicDirs": [] } } diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/public/hooks/use_chat.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/public/hooks/use_chat.ts index 86aeb8f519e87..c281c04099c1d 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/public/hooks/use_chat.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/public/hooks/use_chat.ts @@ -11,6 +11,7 @@ import { useCallback, useEffect, useMemo, useRef, useState } from 'react'; import { AbortError } from '@kbn/kibana-utils-plugin/common'; import type { NotificationsStart } from '@kbn/core/public'; import type { AssistantScope } from '@kbn/ai-assistant-common'; +import { ChatCompletionEventType } from '@kbn/inference-common'; import { MessageRole, type Message, @@ -199,24 +200,22 @@ function useChatWithoutContext({ const subscription = next$.subscribe({ next: (event) => { switch (event.type) { - case StreamingChatResponseEventType.ChatCompletionChunk: + case ChatCompletionEventType.ChatCompletionChunk: if (!pendingMessage) { pendingMessage = { '@timestamp': new Date().toISOString(), message: { - content: event.message.content || '', + content: event.content || '', function_call: { - name: event.message.function_call?.name || '', - arguments: event.message.function_call?.arguments || '', + name: '', + arguments: '', }, }, }; } else { - pendingMessage.message.content += event.message.content || ''; - pendingMessage.message.function_call.name += - event.message.function_call?.name || ''; - pendingMessage.message.function_call.arguments += - event.message.function_call?.arguments || ''; + pendingMessage.message.content += event.content || ''; + pendingMessage.message.function_call.name += ''; + pendingMessage.message.function_call.arguments += ''; } break; diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/client/index.test.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/client/index.test.ts index 8da2a0d843b11..8546ca43297bb 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/client/index.test.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/client/index.test.ts @@ -13,6 +13,7 @@ import type OpenAI from 'openai'; import { Subject } from 'rxjs'; import { EventEmitter, PassThrough, type Readable } from 'stream'; import { finished } from 'stream/promises'; +import type { InferenceClient } from '@kbn/inference-plugin/server'; import { ObservabilityAIAssistantClient } from '.'; import { MessageRole, type Message } from '../../../common'; import { ObservabilityAIAssistantConnectorType } from '../../../common/connectors'; @@ -94,6 +95,10 @@ describe('Observability AI Assistant client', () => { get: jest.fn(), } as any; + const inferenceClientMock: DeeplyMockedKeys = { + chatComplete: jest.fn(), + } as any; + const uiSettingsClientMock: DeeplyMockedKeys = { get: jest.fn(), } as any; @@ -183,6 +188,7 @@ describe('Observability AI Assistant client', () => { asInternalUser: internalUserEsClientMock, asCurrentUser: currentUserEsClientMock, }, + inferenceClient: inferenceClientMock, knowledgeBaseService: knowledgeBaseServiceMock, logger: loggerMock, namespace: 'default', diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/client/index.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/client/index.ts index 048bbd2d362c2..aeadb5de5103b 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/client/index.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/client/index.ts @@ -10,7 +10,7 @@ import type { ActionsClient } from '@kbn/actions-plugin/server'; import type { ElasticsearchClient, IUiSettingsClient } from '@kbn/core/server'; import type { Logger } from '@kbn/logging'; import type { PublicMethodsOf } from '@kbn/utility-types'; -import { SpanKind, context } from '@opentelemetry/api'; +import { context } from '@opentelemetry/api'; import { last, merge, omit } from 'lodash'; import { catchError, @@ -28,23 +28,28 @@ import { tap, throwError, } from 'rxjs'; -import { Readable } from 'stream'; import { v4 } from 'uuid'; import type { AssistantScope } from '@kbn/ai-assistant-common'; +import type { InferenceClient } from '@kbn/inference-plugin/server'; +import { + isChatCompletionChunkEvent, + withoutTokenCountEvents as inferenceWithoutTokenCountEvents, + ToolChoiceType, + type ChatCompletionChunkEvent as InferenceChatCompletionChunkEvent, +} from '@kbn/inference-common'; + import { resourceNames } from '..'; -import { ObservabilityAIAssistantConnectorType } from '../../../common/connectors'; import { - ChatCompletionChunkEvent, + ChatEvent, ChatCompletionErrorEvent, ConversationCreateEvent, ConversationUpdateEvent, createConversationNotFoundError, - createInternalServerError, - createTokenLimitReachedError, StreamingChatResponseEventType, TokenCountEvent, type StreamingChatResponseEvent, } from '../../../common/conversation_complete'; +import { convertMessagesForInference } from '../../../common/convert_messages_for_inference'; import { CompatibleJSONSchema } from '../../../common/functions/types'; import { type AdHocInstruction, @@ -55,6 +60,7 @@ import { type Message, KnowledgeBaseType, KnowledgeBaseEntryRole, + MessageRole, } from '../../../common/types'; import { withoutTokenCountEvents } from '../../../common/utils/without_token_count_events'; import { CONTEXT_FUNCTION_NAME } from '../../functions/context'; @@ -63,23 +69,16 @@ import { KnowledgeBaseService, RecalledEntry } from '../knowledge_base_service'; import { getAccessQuery } from '../util/get_access_query'; import { getSystemMessageFromInstructions } from '../util/get_system_message_from_instructions'; import { replaceSystemMessage } from '../util/replace_system_message'; -import { withAssistantSpan } from '../util/with_assistant_span'; -import { createBedrockClaudeAdapter } from './adapters/bedrock/bedrock_claude_adapter'; -import { failOnNonExistingFunctionCall } from './adapters/fail_on_non_existing_function_call'; -import { createGeminiAdapter } from './adapters/gemini/gemini_adapter'; -import { createOpenAiAdapter } from './adapters/openai_adapter'; -import { LlmApiAdapter } from './adapters/types'; import { getContextFunctionRequestIfNeeded } from './get_context_function_request_if_needed'; import { LangTracer } from './instrumentation/lang_tracer'; import { continueConversation } from './operators/continue_conversation'; import { extractMessages } from './operators/extract_messages'; import { extractTokenCount } from './operators/extract_token_count'; import { getGeneratedTitle } from './operators/get_generated_title'; -import { instrumentAndCountTokens } from './operators/instrument_and_count_tokens'; import { - LangtraceServiceProvider, - withLangtraceChatCompleteSpan, -} from './operators/with_langtrace_chat_complete_span'; + instrumentAndCountTokens, + instrumentStream, +} from './operators/instrument_and_count_tokens'; const MAX_FUNCTION_CALLS = 8; @@ -93,6 +92,7 @@ export class ObservabilityAIAssistantClient { asInternalUser: ElasticsearchClient; asCurrentUser: ElasticsearchClient; }; + inferenceClient: InferenceClient; logger: Logger; user?: { id?: string; @@ -477,126 +477,42 @@ export class ObservabilityAIAssistantClient { simulateFunctionCalling?: boolean; tracer: LangTracer; } - ): Observable => { - return defer(() => - from( - withAssistantSpan('get_connector', () => - this.dependencies.actionsClient.get({ id: connectorId, throwIfSystemAction: true }) - ) - ) - ).pipe( - switchMap((connector) => { - this.dependencies.logger.debug(`Creating "${connector.actionTypeId}" adapter`); - - let adapter: LlmApiAdapter; - - switch (connector.actionTypeId) { - case ObservabilityAIAssistantConnectorType.OpenAI: - adapter = createOpenAiAdapter({ - messages, - functions, - functionCall, - logger: this.dependencies.logger, - simulateFunctionCalling, - }); - break; - - case ObservabilityAIAssistantConnectorType.Bedrock: - adapter = createBedrockClaudeAdapter({ - messages, - functions, - functionCall, - logger: this.dependencies.logger, - }); - break; - - case ObservabilityAIAssistantConnectorType.Gemini: - adapter = createGeminiAdapter({ - messages, - functions, - functionCall, - logger: this.dependencies.logger, - }); - break; - - default: - throw new Error(`Connector type is not supported: ${connector.actionTypeId}`); - } - - const subAction = adapter.getSubAction(); - - if (this.dependencies.logger.isLevelEnabled('trace')) { - this.dependencies.logger.trace(JSON.stringify(subAction.subActionParams, null, 2)); - } - - return from( - withAssistantSpan('get_execute_result', () => - this.dependencies.actionsClient.execute({ - actionId: connectorId, - params: subAction, - }) - ) - ).pipe( - switchMap((executeResult) => { - if (executeResult.status === 'error' && executeResult?.serviceMessage) { - const tokenLimitRegex = - /This model's maximum context length is (\d+) tokens\. However, your messages resulted in (\d+) tokens/g; - const tokenLimitRegexResult = tokenLimitRegex.exec(executeResult.serviceMessage); - - if (tokenLimitRegexResult) { - const [, tokenLimit, tokenCount] = tokenLimitRegexResult; - throw createTokenLimitReachedError( - parseInt(tokenLimit, 10), - parseInt(tokenCount, 10) - ); - } - } - - if (executeResult.status === 'error') { - throw createInternalServerError( - `${executeResult?.message} - ${executeResult?.serviceMessage}` - ); + ): Observable => { + const tools = functions?.reduce((acc, fn) => { + acc[fn.name] = { + description: fn.description, + schema: fn.parameters, + }; + return acc; + }, {} as Record); + + const chatComplete$ = defer(() => + this.dependencies.inferenceClient.chatComplete({ + connectorId, + stream: true, + messages: convertMessagesForInference( + messages.filter((message) => message.message.role !== MessageRole.System) + ), + functionCalling: simulateFunctionCalling ? 'simulated' : 'native', + toolChoice: functionCall + ? { + function: functionCall, } - - const response = executeResult.data as Readable; - - signal.addEventListener('abort', () => response.destroy()); - - return tracer.startActiveSpan( - '/chat/completions', - { - kind: SpanKind.CLIENT, - }, - ({ span }) => { - return adapter.streamIntoObservable(response).pipe( - withLangtraceChatCompleteSpan({ - span, - messages, - functions, - model: connector.name, - serviceProvider: - connector.actionTypeId === ObservabilityAIAssistantConnectorType.OpenAI - ? LangtraceServiceProvider.OpenAI - : LangtraceServiceProvider.Anthropic, - }) - ); - } - ); - }) - ); - }), - instrumentAndCountTokens(name), - failOnNonExistingFunctionCall({ functions }), + : ToolChoiceType.auto, + tools, + }) + ).pipe( + inferenceWithoutTokenCountEvents(), tap((event) => { - if ( - event.type === StreamingChatResponseEventType.ChatCompletionChunk && - this.dependencies.logger.isLevelEnabled('trace') - ) { - this.dependencies.logger.trace(`Received chunk: ${JSON.stringify(event.message)}`); + if (isChatCompletionChunkEvent(event) && this.dependencies.logger.isLevelEnabled('trace')) { + this.dependencies.logger.trace(`Received chunk: ${JSON.stringify(event)}`); } }), + instrumentStream(name), shareReplay() ); + + return chatComplete$ as Observable; }; find = async (options?: { query?: string }): Promise<{ conversations: Conversation[] }> => { diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/client/operators/extract_token_count.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/client/operators/extract_token_count.ts index 0d11db24732f3..5372988d2405a 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/client/operators/extract_token_count.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/client/operators/extract_token_count.ts @@ -5,7 +5,7 @@ * 2.0. */ -import { filter, OperatorFunction, scan } from 'rxjs'; +import { filter, OperatorFunction, scan, endWith } from 'rxjs'; import { StreamingChatResponseEvent, StreamingChatResponseEventType, @@ -24,13 +24,15 @@ export function extractTokenCount(): OperatorFunction< ), scan( (acc, event) => { - acc.completion += event.tokens.completion; - acc.prompt += event.tokens.prompt; - acc.total += event.tokens.total; + acc.completion += event?.tokens?.completion ?? 0; + acc.prompt += event?.tokens?.prompt ?? 0; + acc.total += event?.tokens?.total ?? 0; return acc; }, { completion: 0, prompt: 0, total: 0 } - ) + ), + // Emit a default value at the end if no events were processed + endWith({ completion: 0, prompt: 0, total: 0 }) ); }; } diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/client/operators/instrument_and_count_tokens.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/client/operators/instrument_and_count_tokens.ts index 094b2606ae533..34d8da61636eb 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/client/operators/instrument_and_count_tokens.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/client/operators/instrument_and_count_tokens.ts @@ -20,6 +20,36 @@ import { import type { StreamingChatResponseEvent } from '../../../../common/conversation_complete'; import { extractTokenCount } from './extract_token_count'; +export function instrumentStream(name: string): OperatorFunction { + return (source$) => { + const span = apm.startSpan(name); + + if (!span) { + return source$; + } + + span.addLabels({ + plugin: 'observability_ai_assistant', + }); + + return source$.pipe( + shareReplay(), + tap({ + complete: () => { + span.setOutcome('success'); + }, + }), + catchError((error) => { + span.setOutcome('failure'); + return throwError(() => error); + }), + finalize(() => { + span.end(); + }) + ); + }; +} + export function instrumentAndCountTokens( name: string ): OperatorFunction { diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/client/operators/with_langtrace_chat_complete_span.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/client/operators/with_langtrace_chat_complete_span.ts index 9e32ba4b57bfe..ae9adc8fdc1fb 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/client/operators/with_langtrace_chat_complete_span.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/client/operators/with_langtrace_chat_complete_span.ts @@ -85,9 +85,9 @@ export function withLangtraceChatCompleteSpan({ span.setAttributes({ 'llm.token.counts': JSON.stringify({ - input_tokens: value.tokens.prompt, - output_tokens: value.tokens.completion, - total_tokens: value.tokens.total, + input_tokens: 'tokens' in value ? value.tokens.prompt : 0, + output_tokens: 'tokens' in value ? value.tokens.completion : 0, + total_tokens: 'tokens' in value ? value.tokens.total : 0, }), }); }) diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/index.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/index.ts index eb7eab19340ce..37323f0954df2 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/index.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/index.ts @@ -5,10 +5,8 @@ * 2.0. */ -import type { PluginStartContract as ActionsPluginStart } from '@kbn/actions-plugin/server/plugin'; import { createConcreteWriteIndex, getDataStreamAdapter } from '@kbn/alerting-plugin/server'; import type { CoreSetup, CoreStart, KibanaRequest, Logger } from '@kbn/core/server'; -import type { SecurityPluginStart } from '@kbn/security-plugin/server'; import { getSpaceIdFromPath } from '@kbn/spaces-plugin/common'; import type { TaskManagerSetupContract } from '@kbn/task-manager-plugin/server'; import { once } from 'lodash'; @@ -232,7 +230,7 @@ export class ObservabilityAIAssistantService { const [_, [coreStart, plugins]] = await Promise.all([ this.init(), this.core.getStartServices() as Promise< - [CoreStart, { security: SecurityPluginStart; actions: ActionsPluginStart }, unknown] + [CoreStart, ObservabilityAIAssistantPluginStartDependencies, unknown] >, ]); // user will not be found when executed from system connector context @@ -243,6 +241,7 @@ export class ObservabilityAIAssistantService { const basePath = coreStart.http.basePath.get(request); const { spaceId } = getSpaceIdFromPath(basePath, coreStart.http.basePath.serverBasePath); + const inferenceClient = plugins.inference.getClient({ request }); return new ObservabilityAIAssistantClient({ actionsClient: await plugins.actions.getActionsClientWithRequest(request), @@ -252,6 +251,7 @@ export class ObservabilityAIAssistantService { asInternalUser: coreStart.elasticsearch.client.asInternalUser, asCurrentUser: coreStart.elasticsearch.client.asScoped(request).asCurrentUser, }, + inferenceClient, logger: this.logger, user: user ? { diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/types.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/types.ts index 2e24cf25902e0..7b08d3e1c8f35 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/types.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/types.ts @@ -8,6 +8,7 @@ import type { FromSchema } from 'json-schema-to-ts'; import { Observable } from 'rxjs'; import type { AssistantScope } from '@kbn/ai-assistant-common'; +import type { ChatCompletionChunkEvent as InferenceChatCompletionChunkEvent } from '@kbn/inference-common'; import { ChatCompletionChunkEvent, ChatEvent } from '../../common/conversation_complete'; import type { CompatibleJSONSchema, @@ -47,7 +48,7 @@ export type FunctionCallChatFunction = ( Parameters[1], 'connectorId' | 'simulateFunctionCalling' | 'tracer' > -) => Observable; +) => Observable; type RespondFunction = ( options: { diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/types.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/types.ts index a8ab57b8ee53c..d77ac100b1514 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/server/types.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/types.ts @@ -26,6 +26,7 @@ import type { PluginSetupContract as AlertingPluginSetup, PluginStartContract as AlertingPluginStart, } from '@kbn/alerting-plugin/server'; +import type { InferenceServerStart } from '@kbn/inference-plugin/server'; import type { ObservabilityAIAssistantService } from './service'; export interface ObservabilityAIAssistantServerSetup { @@ -65,4 +66,5 @@ export interface ObservabilityAIAssistantPluginStartDependencies { cloud?: CloudStart; serverless?: ServerlessPluginStart; alerting: AlertingPluginStart; + inference: InferenceServerStart; } 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 750bf69477653..adf4d0b6bf36d 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/tsconfig.json +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/tsconfig.json @@ -46,6 +46,7 @@ "@kbn/management-settings-ids", "@kbn/ai-assistant-common", "@kbn/inference-common", + "@kbn/inference-plugin" ], "exclude": ["target/**/*"] } 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 210dee20339af..cefec5ae66758 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 @@ -15,12 +15,12 @@ import { StreamingChatResponseEventType, } from '@kbn/observability-ai-assistant-plugin/common'; import { createFunctionResponseMessage } from '@kbn/observability-ai-assistant-plugin/common/utils/create_function_response_message'; +import { convertMessagesForInference } from '@kbn/observability-ai-assistant-plugin/common/convert_messages_for_inference'; import { map } from 'rxjs'; import { v4 } from 'uuid'; import { RegisterInstructionCallback } from '@kbn/observability-ai-assistant-plugin/server/service/types'; import type { FunctionRegistrationParameters } from '..'; import { runAndValidateEsqlQuery } from './validate_esql_query'; -import { convertMessagesForInference } from '../../../common/convert_messages_for_inference'; export const QUERY_FUNCTION_NAME = 'query'; export const EXECUTE_QUERY_NAME = 'execute_query';