Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[inference] Add support for inference connectors #204541

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down
1 change: 1 addition & 0 deletions x-pack/packages/kbn-ai-assistant/tsconfig.json
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@
"@kbn/ml-plugin",
"@kbn/share-plugin",
"@kbn/ai-assistant-common",
"@kbn/inference-common",
"@kbn/storybook",
"@kbn/ai-assistant-icon",
]
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -95,3 +95,9 @@ export {
} from './src/errors';

export { truncateList } from './src/truncate_list';
export {
InferenceConnectorType,
isSupportedConnectorType,
isSupportedConnector,
type InferenceConnector,
} from './src/connectors';
Original file line number Diff line number Diff line change
@@ -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);
});
});
Original file line number Diff line number Diff line change
@@ -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;
}
Comment on lines +46 to +57
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Checking if a connector is compatible is no longer just based on its type, so I had to create that new check logic.

For inference connectors, we might eventually want to filter based on the provider, but for now I feel like filtering on completion tasks should be sufficient


/**
* 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>;
}
24 changes: 0 additions & 24 deletions x-pack/platform/plugins/shared/inference/common/connectors.ts

This file was deleted.

Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
3 changes: 1 addition & 2 deletions x-pack/platform/plugins/shared/inference/public/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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*/

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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';

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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', () => {
Expand All @@ -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);
});
});
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -23,6 +24,9 @@ export const getInferenceAdapter = (

case InferenceConnectorType.Bedrock:
return bedrockClaudeAdapter;

case InferenceConnectorType.Inference:
return inferenceAdapter;
}

return undefined;
Expand Down
Original file line number Diff line number Diff line change
@@ -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';
Loading
Loading