Skip to content

Commit 3dcae51

Browse files
[inference] Add support for inference connectors (#204541)
## Summary ~Depends on~ #200249 merged! Fix #199082 - Add support for the `inference` stack connectors to the `inference` plugin (everything is inference) - Adapt the o11y assistant to use the `inference-common` utilities for connector filtering / compat checking ## How to test **1. Starts ES with the unified completion feature flag** ```sh yarn es snapshot --license trial ES_JAVA_OPTS="-Des.inference_unified_feature_flag_enabled=true" ``` **2. Enable the inference connector for Kibana** In the Kibana config file: ```yaml xpack.stack_connectors.enableExperimental: ['inferenceConnectorOn'] ``` **3. Start Dev Kibana** ```sh node scripts/kibana --dev --no-base-path ``` **4. Create an inference connector** Go to `http://localhost:5601/app/management/insightsAndAlerting/triggersActionsConnectors/connectors`, create an inference connector - Type: `AI connector` then - Service: `OpenAI` - API Key: Gwzk... Kidding, please ping someone - Model ID: `gpt-4o` - Task type: `completion` -> save **5. test the o11y assistant** Use the assistant as you would do for any other connector (just make sure the inference connector is selected as the one being used) and do your testing. --------- Co-authored-by: kibanamachine <[email protected]>
1 parent 4fa7b78 commit 3dcae51

File tree

37 files changed

+894
-273
lines changed

37 files changed

+894
-273
lines changed

x-pack/packages/kbn-ai-assistant/src/chat/welcome_message.tsx

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ import { css } from '@emotion/css';
1010
import { EuiFlexGroup, EuiFlexItem, EuiSpacer, useCurrentEuiBreakpoint } from '@elastic/eui';
1111
import type { ActionConnector } from '@kbn/triggers-actions-ui-plugin/public';
1212
import { GenerativeAIForObservabilityConnectorFeatureId } from '@kbn/actions-plugin/common';
13-
import { isSupportedConnectorType } from '@kbn/observability-ai-assistant-plugin/public';
13+
import { isSupportedConnectorType } from '@kbn/inference-common';
1414
import { AssistantBeacon } from '@kbn/ai-assistant-icon';
1515
import type { UseKnowledgeBaseResult } from '../hooks/use_knowledge_base';
1616
import type { UseGenAIConnectorsResult } from '../hooks/use_genai_connectors';

x-pack/packages/kbn-ai-assistant/tsconfig.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
"@kbn/ml-plugin",
3838
"@kbn/share-plugin",
3939
"@kbn/ai-assistant-common",
40+
"@kbn/inference-common",
4041
"@kbn/storybook",
4142
"@kbn/ai-assistant-icon",
4243
]

x-pack/platform/packages/shared/ai-infra/inference-common/index.ts

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,3 +95,9 @@ export {
9595
} from './src/errors';
9696

9797
export { truncateList } from './src/truncate_list';
98+
export {
99+
InferenceConnectorType,
100+
isSupportedConnectorType,
101+
isSupportedConnector,
102+
type InferenceConnector,
103+
} from './src/connectors';
Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
import {
9+
InferenceConnectorType,
10+
isSupportedConnectorType,
11+
isSupportedConnector,
12+
RawConnector,
13+
COMPLETION_TASK_TYPE,
14+
} from './connectors';
15+
16+
const createRawConnector = (parts: Partial<RawConnector>): RawConnector => {
17+
return {
18+
id: 'id',
19+
actionTypeId: 'connector-type',
20+
name: 'some connector',
21+
config: {},
22+
...parts,
23+
};
24+
};
25+
26+
describe('isSupportedConnectorType', () => {
27+
it('returns true for supported connector types', () => {
28+
expect(isSupportedConnectorType(InferenceConnectorType.OpenAI)).toBe(true);
29+
expect(isSupportedConnectorType(InferenceConnectorType.Bedrock)).toBe(true);
30+
expect(isSupportedConnectorType(InferenceConnectorType.Gemini)).toBe(true);
31+
expect(isSupportedConnectorType(InferenceConnectorType.Inference)).toBe(true);
32+
});
33+
it('returns false for unsupported connector types', () => {
34+
expect(isSupportedConnectorType('anything-else')).toBe(false);
35+
});
36+
});
37+
38+
describe('isSupportedConnector', () => {
39+
// TODO
40+
41+
it('returns true for OpenAI connectors', () => {
42+
expect(
43+
isSupportedConnector(createRawConnector({ actionTypeId: InferenceConnectorType.OpenAI }))
44+
).toBe(true);
45+
});
46+
47+
it('returns true for Bedrock connectors', () => {
48+
expect(
49+
isSupportedConnector(createRawConnector({ actionTypeId: InferenceConnectorType.Bedrock }))
50+
).toBe(true);
51+
});
52+
53+
it('returns true for Gemini connectors', () => {
54+
expect(
55+
isSupportedConnector(createRawConnector({ actionTypeId: InferenceConnectorType.Gemini }))
56+
).toBe(true);
57+
});
58+
59+
it('returns true for OpenAI connectors with the right taskType', () => {
60+
expect(
61+
isSupportedConnector(
62+
createRawConnector({
63+
actionTypeId: InferenceConnectorType.Inference,
64+
config: { taskType: COMPLETION_TASK_TYPE },
65+
})
66+
)
67+
).toBe(true);
68+
});
69+
70+
it('returns false for OpenAI connectors with a bad taskType', () => {
71+
expect(
72+
isSupportedConnector(
73+
createRawConnector({
74+
actionTypeId: InferenceConnectorType.Inference,
75+
config: { taskType: 'embeddings' },
76+
})
77+
)
78+
).toBe(false);
79+
});
80+
81+
it('returns false for OpenAI connectors without taskType', () => {
82+
expect(
83+
isSupportedConnector(
84+
createRawConnector({
85+
actionTypeId: InferenceConnectorType.Inference,
86+
config: {},
87+
})
88+
)
89+
).toBe(false);
90+
});
91+
});
Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
/**
9+
* The list of connector types that can be used with the inference APIs
10+
*/
11+
export enum InferenceConnectorType {
12+
OpenAI = '.gen-ai',
13+
Bedrock = '.bedrock',
14+
Gemini = '.gemini',
15+
Inference = '.inference',
16+
}
17+
18+
export const COMPLETION_TASK_TYPE = 'completion';
19+
20+
const allSupportedConnectorTypes = Object.values(InferenceConnectorType);
21+
22+
export interface InferenceConnector {
23+
type: InferenceConnectorType;
24+
name: string;
25+
connectorId: string;
26+
}
27+
28+
/**
29+
* Checks if a given connector type is compatible for inference.
30+
*
31+
* Note: this check is not sufficient to assert if a given connector can be
32+
* used for inference, as `.inference` connectors need additional check logic.
33+
* Please use `isSupportedConnector` instead when possible.
34+
*/
35+
export function isSupportedConnectorType(id: string): id is InferenceConnectorType {
36+
return allSupportedConnectorTypes.includes(id as InferenceConnectorType);
37+
}
38+
39+
/**
40+
* Checks if a given connector is compatible for inference.
41+
*
42+
* A connector is compatible if:
43+
* 1. its type is in the list of allowed types
44+
* 2. for inference connectors, if its taskType is "completion"
45+
*/
46+
export function isSupportedConnector(connector: RawConnector): connector is RawInferenceConnector {
47+
if (!isSupportedConnectorType(connector.actionTypeId)) {
48+
return false;
49+
}
50+
if (connector.actionTypeId === InferenceConnectorType.Inference) {
51+
const config = connector.config ?? {};
52+
if (config.taskType !== COMPLETION_TASK_TYPE) {
53+
return false;
54+
}
55+
}
56+
return true;
57+
}
58+
59+
/**
60+
* Connector types are living in the actions plugin and we can't afford
61+
* having dependencies from this package to some mid-level plugin,
62+
* so we're just using our own connector mixin type.
63+
*/
64+
export interface RawConnector {
65+
id: string;
66+
actionTypeId: string;
67+
name: string;
68+
config?: Record<string, any>;
69+
}
70+
71+
interface RawInferenceConnector {
72+
id: string;
73+
actionTypeId: InferenceConnectorType;
74+
name: string;
75+
config?: Record<string, any>;
76+
}

x-pack/platform/plugins/shared/inference/common/connectors.ts

Lines changed: 0 additions & 24 deletions
This file was deleted.

x-pack/platform/plugins/shared/inference/common/http_apis.ts

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,12 @@
55
* 2.0.
66
*/
77

8-
import type { FunctionCallingMode, Message, ToolOptions } from '@kbn/inference-common';
9-
import { InferenceConnector } from './connectors';
8+
import type {
9+
FunctionCallingMode,
10+
Message,
11+
ToolOptions,
12+
InferenceConnector,
13+
} from '@kbn/inference-common';
1014

1115
export type ChatCompleteRequestBody = {
1216
connectorId: string;

x-pack/platform/plugins/shared/inference/public/types.ts

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,7 @@
55
* 2.0.
66
*/
77

8-
import type { ChatCompleteAPI, OutputAPI } from '@kbn/inference-common';
9-
import type { InferenceConnector } from '../common/connectors';
8+
import type { ChatCompleteAPI, OutputAPI, InferenceConnector } from '@kbn/inference-common';
109

1110
/* eslint-disable @typescript-eslint/no-empty-interface*/
1211

x-pack/platform/plugins/shared/inference/scripts/util/kibana_client.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,9 +25,9 @@ import {
2525
withoutOutputUpdateEvents,
2626
type ToolOptions,
2727
ChatCompleteOptions,
28+
type InferenceConnector,
2829
} from '@kbn/inference-common';
2930
import type { ChatCompleteRequestBody } from '../../common/http_apis';
30-
import type { InferenceConnector } from '../../common/connectors';
3131
import { createOutputApi } from '../../common/output/create_output_api';
3232
import { eventSourceStreamIntoObservable } from '../../server/util/event_source_stream_into_observable';
3333

x-pack/platform/plugins/shared/inference/server/chat_complete/adapters/get_inference_adapter.test.ts

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,12 @@
55
* 2.0.
66
*/
77

8-
import { InferenceConnectorType } from '../../../common/connectors';
8+
import { InferenceConnectorType } from '@kbn/inference-common';
99
import { getInferenceAdapter } from './get_inference_adapter';
1010
import { openAIAdapter } from './openai';
1111
import { geminiAdapter } from './gemini';
1212
import { bedrockClaudeAdapter } from './bedrock';
13+
import { inferenceAdapter } from './inference';
1314

1415
describe('getInferenceAdapter', () => {
1516
it('returns the openAI adapter for OpenAI type', () => {
@@ -23,4 +24,8 @@ describe('getInferenceAdapter', () => {
2324
it('returns the bedrock adapter for Bedrock type', () => {
2425
expect(getInferenceAdapter(InferenceConnectorType.Bedrock)).toBe(bedrockClaudeAdapter);
2526
});
27+
28+
it('returns the inference adapter for Inference type', () => {
29+
expect(getInferenceAdapter(InferenceConnectorType.Inference)).toBe(inferenceAdapter);
30+
});
2631
});

0 commit comments

Comments
 (0)