[8.x] [inference] Add support for inference connectors (#204541) (#20…

…5078) # Backport This will backport the following commits from `main` to `8.x`: - [[inference] Add support for inference connectors (#204541)](#204541)  ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sqren/backport)  --------- Co-authored-by: Pierre Gayvallet <[email protected]>
elastic · Dec 23, 2024 · a08a128 · a08a128
1 parent 9b9ce42
commit a08a128
Show file tree

Hide file tree

Showing 38 changed files with 984 additions and 297 deletions.
diff --git a/x-pack/packages/kbn-ai-assistant/src/chat/welcome_message.tsx b/x-pack/packages/kbn-ai-assistant/src/chat/welcome_message.tsx
@@ -10,7 +10,7 @@ import { css } from '@emotion/css';
 import { EuiFlexGroup, EuiFlexItem, EuiSpacer, useCurrentEuiBreakpoint } from '@elastic/eui';
 import type { ActionConnector } from '@kbn/triggers-actions-ui-plugin/public';
 import { GenerativeAIForObservabilityConnectorFeatureId } from '@kbn/actions-plugin/common';
-import { isSupportedConnectorType } from '@kbn/observability-ai-assistant-plugin/public';
+import { isSupportedConnectorType } from '@kbn/inference-common';
 import { AssistantBeacon } from '@kbn/ai-assistant-icon';
 import type { UseKnowledgeBaseResult } from '../hooks/use_knowledge_base';
 import type { UseGenAIConnectorsResult } from '../hooks/use_genai_connectors';

diff --git a/x-pack/packages/kbn-ai-assistant/tsconfig.json b/x-pack/packages/kbn-ai-assistant/tsconfig.json
@@ -38,6 +38,7 @@
     "@kbn/ml-plugin",
     "@kbn/share-plugin",
     "@kbn/ai-assistant-common",
+    "@kbn/inference-common",
     "@kbn/storybook",
     "@kbn/ai-assistant-icon",
   ]

diff --git a/x-pack/platform/packages/shared/ai-infra/inference-common/index.ts b/x-pack/platform/packages/shared/ai-infra/inference-common/index.ts
@@ -95,3 +95,9 @@ export {
 } from './src/errors';
 
 export { truncateList } from './src/truncate_list';
+export {
+  InferenceConnectorType,
+  isSupportedConnectorType,
+  isSupportedConnector,
+  type InferenceConnector,
+} from './src/connectors';
diff --git a/x-pack/platform/packages/shared/ai-infra/inference-common/src/connectors.test.ts b/x-pack/platform/packages/shared/ai-infra/inference-common/src/connectors.test.ts
@@ -0,0 +1,91 @@
+/*
+ * 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 {
+  InferenceConnectorType,
+  isSupportedConnectorType,
+  isSupportedConnector,
+  RawConnector,
+  COMPLETION_TASK_TYPE,
+} from './connectors';
+
+const createRawConnector = (parts: Partial<RawConnector>): RawConnector => {
+  return {
+    id: 'id',
+    actionTypeId: 'connector-type',
+    name: 'some connector',
+    config: {},
+    ...parts,
+  };
+};
+
+describe('isSupportedConnectorType', () => {
+  it('returns true for supported connector types', () => {
+    expect(isSupportedConnectorType(InferenceConnectorType.OpenAI)).toBe(true);
+    expect(isSupportedConnectorType(InferenceConnectorType.Bedrock)).toBe(true);
+    expect(isSupportedConnectorType(InferenceConnectorType.Gemini)).toBe(true);
+    expect(isSupportedConnectorType(InferenceConnectorType.Inference)).toBe(true);
+  });
+  it('returns false for unsupported connector types', () => {
+    expect(isSupportedConnectorType('anything-else')).toBe(false);
+  });
+});
+
+describe('isSupportedConnector', () => {
+  // TODO
+
+  it('returns true for OpenAI connectors', () => {
+    expect(
+      isSupportedConnector(createRawConnector({ actionTypeId: InferenceConnectorType.OpenAI }))
+    ).toBe(true);
+  });
+
+  it('returns true for Bedrock connectors', () => {
+    expect(
+      isSupportedConnector(createRawConnector({ actionTypeId: InferenceConnectorType.Bedrock }))
+    ).toBe(true);
+  });
+
+  it('returns true for Gemini connectors', () => {
+    expect(
+      isSupportedConnector(createRawConnector({ actionTypeId: InferenceConnectorType.Gemini }))
+    ).toBe(true);
+  });
+
+  it('returns true for OpenAI connectors with the right taskType', () => {
+    expect(
+      isSupportedConnector(
+        createRawConnector({
+          actionTypeId: InferenceConnectorType.Inference,
+          config: { taskType: COMPLETION_TASK_TYPE },
+        })
+      )
+    ).toBe(true);
+  });
+
+  it('returns false for OpenAI connectors with a bad taskType', () => {
+    expect(
+      isSupportedConnector(
+        createRawConnector({
+          actionTypeId: InferenceConnectorType.Inference,
+          config: { taskType: 'embeddings' },
+        })
+      )
+    ).toBe(false);
+  });
+
+  it('returns false for OpenAI connectors without taskType', () => {
+    expect(
+      isSupportedConnector(
+        createRawConnector({
+          actionTypeId: InferenceConnectorType.Inference,
+          config: {},
+        })
+      )
+    ).toBe(false);
+  });
+});
diff --git a/x-pack/platform/packages/shared/ai-infra/inference-common/src/connectors.ts b/x-pack/platform/packages/shared/ai-infra/inference-common/src/connectors.ts
@@ -0,0 +1,76 @@
+/*
+ * 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.
+ */
+
+/**
+ * The list of connector types that can be used with the inference APIs
+ */
+export enum InferenceConnectorType {
+  OpenAI = '.gen-ai',
+  Bedrock = '.bedrock',
+  Gemini = '.gemini',
+  Inference = '.inference',
+}
+
+export const COMPLETION_TASK_TYPE = 'completion';
+
+const allSupportedConnectorTypes = Object.values(InferenceConnectorType);
+
+export interface InferenceConnector {
+  type: InferenceConnectorType;
+  name: string;
+  connectorId: string;
+}
+
+/**
+ * Checks if a given connector type is compatible for inference.
+ *
+ * Note: this check is not sufficient to assert if a given connector can be
+ * used for inference, as `.inference` connectors need additional check logic.
+ * Please use `isSupportedConnector` instead when possible.
+ */
+export function isSupportedConnectorType(id: string): id is InferenceConnectorType {
+  return allSupportedConnectorTypes.includes(id as InferenceConnectorType);
+}
+
+/**
+ * Checks if a given connector is compatible for inference.
+ *
+ * A connector is compatible if:
+ * 1. its type is in the list of allowed types
+ * 2. for inference connectors, if its taskType is "completion"
+ */
+export function isSupportedConnector(connector: RawConnector): connector is RawInferenceConnector {
+  if (!isSupportedConnectorType(connector.actionTypeId)) {
+    return false;
+  }
+  if (connector.actionTypeId === InferenceConnectorType.Inference) {
+    const config = connector.config ?? {};
+    if (config.taskType !== COMPLETION_TASK_TYPE) {
+      return false;
+    }
+  }
+  return true;
+}
+
+/**
+ * Connector types are living in the actions plugin and we can't afford
+ * having dependencies from this package to some mid-level plugin,
+ * so we're just using our own connector mixin type.
+ */
+export interface RawConnector {
+  id: string;
+  actionTypeId: string;
+  name: string;
+  config?: Record<string, any>;
+}
+
+interface RawInferenceConnector {
+  id: string;
+  actionTypeId: InferenceConnectorType;
+  name: string;
+  config?: Record<string, any>;
+}
diff --git a/x-pack/platform/plugins/shared/inference/common/connectors.ts b/x-pack/platform/plugins/shared/inference/common/connectors.ts
diff --git a/x-pack/platform/plugins/shared/inference/common/http_apis.ts b/x-pack/platform/plugins/shared/inference/common/http_apis.ts
@@ -5,8 +5,12 @@
  * 2.0.
  */
 
-import type { FunctionCallingMode, Message, ToolOptions } from '@kbn/inference-common';
-import { InferenceConnector } from './connectors';
+import type {
+  FunctionCallingMode,
+  Message,
+  ToolOptions,
+  InferenceConnector,
+} from '@kbn/inference-common';
 
 export type ChatCompleteRequestBody = {
   connectorId: string;

diff --git a/x-pack/platform/plugins/shared/inference/public/types.ts b/x-pack/platform/plugins/shared/inference/public/types.ts
@@ -5,8 +5,7 @@
  * 2.0.
  */
 
-import type { ChatCompleteAPI, OutputAPI } from '@kbn/inference-common';
-import type { InferenceConnector } from '../common/connectors';
+import type { ChatCompleteAPI, OutputAPI, InferenceConnector } from '@kbn/inference-common';
 
 /* eslint-disable @typescript-eslint/no-empty-interface*/
 

diff --git a/x-pack/platform/plugins/shared/inference/scripts/util/kibana_client.ts b/x-pack/platform/plugins/shared/inference/scripts/util/kibana_client.ts
@@ -25,9 +25,9 @@ import {
   withoutOutputUpdateEvents,
   type ToolOptions,
   ChatCompleteOptions,
+  type InferenceConnector,
 } from '@kbn/inference-common';
 import type { ChatCompleteRequestBody } from '../../common/http_apis';
-import type { InferenceConnector } from '../../common/connectors';
 import { createOutputApi } from '../../common/output/create_output_api';
 import { eventSourceStreamIntoObservable } from '../../server/util/event_source_stream_into_observable';
 

diff --git a/...form/plugins/shared/inference/server/chat_complete/adapters/get_inference_adapter.test.ts b/...form/plugins/shared/inference/server/chat_complete/adapters/get_inference_adapter.test.ts
@@ -5,11 +5,12 @@
  * 2.0.
  */
 
-import { InferenceConnectorType } from '../../../common/connectors';
+import { InferenceConnectorType } from '@kbn/inference-common';
 import { getInferenceAdapter } from './get_inference_adapter';
 import { openAIAdapter } from './openai';
 import { geminiAdapter } from './gemini';
 import { bedrockClaudeAdapter } from './bedrock';
+import { inferenceAdapter } from './inference';
 
 describe('getInferenceAdapter', () => {
   it('returns the openAI adapter for OpenAI type', () => {
@@ -23,4 +24,8 @@ describe('getInferenceAdapter', () => {
   it('returns the bedrock adapter for Bedrock type', () => {
     expect(getInferenceAdapter(InferenceConnectorType.Bedrock)).toBe(bedrockClaudeAdapter);
   });
+
+  it('returns the inference adapter for Inference type', () => {
+    expect(getInferenceAdapter(InferenceConnectorType.Inference)).toBe(inferenceAdapter);
+  });
 });
diff --git a/.../platform/plugins/shared/inference/server/chat_complete/adapters/get_inference_adapter.ts b/.../platform/plugins/shared/inference/server/chat_complete/adapters/get_inference_adapter.ts
@@ -5,11 +5,12 @@
  * 2.0.
  */
 
-import { InferenceConnectorType } from '../../../common/connectors';
+import { InferenceConnectorType } from '@kbn/inference-common';
 import type { InferenceConnectorAdapter } from '../types';
 import { openAIAdapter } from './openai';
 import { geminiAdapter } from './gemini';
 import { bedrockClaudeAdapter } from './bedrock';
+import { inferenceAdapter } from './inference';
 
 export const getInferenceAdapter = (
   connectorType: InferenceConnectorType
@@ -23,6 +24,9 @@ export const getInferenceAdapter = (
 
     case InferenceConnectorType.Bedrock:
       return bedrockClaudeAdapter;
+
+    case InferenceConnectorType.Inference:
+      return inferenceAdapter;
   }
 
   return undefined;

diff --git a/x-pack/platform/plugins/shared/inference/server/chat_complete/adapters/inference/index.ts b/x-pack/platform/plugins/shared/inference/server/chat_complete/adapters/inference/index.ts
@@ -0,0 +1,8 @@
+/*
+ * 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 { inferenceAdapter } from './inference_adapter';