-
Notifications
You must be signed in to change notification settings - Fork 8.3k
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
pgayvallet
merged 16 commits into
elastic:main
from
pgayvallet:kbn-199082-inference-connector-support
Dec 23, 2024
Merged
Changes from all commits
Commits
Show all changes
16 commits
Select commit
Hold shift + click to select a range
978759f
[inference] Add support for inference connectors
pgayvallet ff00096
start factorizing with openai adapter
pgayvallet 777b914
extract processOpenAIStream
pgayvallet 986f554
[CI] Auto-commit changed files from 'node scripts/notice'
kibanamachine b1cd411
fix export
pgayvallet fbc5a9a
fix export
pgayvallet f6fb20d
revert local change
pgayvallet 2388085
Merge remote-tracking branch 'upstream/main' into kbn-199082-inferenc…
pgayvallet ec42ad3
use real action from PR
pgayvallet 3ee0da1
make the thing work with function calling
pgayvallet 408ac4d
self-review first pass
pgayvallet 87ea9bc
add some unit tests
pgayvallet 9581aa2
fix tests
pgayvallet 45622e9
lint
pgayvallet 8789d10
Merge remote-tracking branch 'upstream/main' into kbn-199082-inferenc…
pgayvallet b31e754
Merge remote-tracking branch 'upstream/main' into kbn-199082-inferenc…
pgayvallet File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
91 changes: 91 additions & 0 deletions
91
x-pack/platform/packages/shared/ai-infra/inference-common/src/connectors.test.ts
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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); | ||
}); | ||
}); |
76 changes: 76 additions & 0 deletions
76
x-pack/platform/packages/shared/ai-infra/inference-common/src/connectors.ts
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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; | ||
} | ||
|
||
/** | ||
* 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
24
x-pack/platform/plugins/shared/inference/common/connectors.ts
This file was deleted.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
8 changes: 8 additions & 0 deletions
8
x-pack/platform/plugins/shared/inference/server/chat_complete/adapters/inference/index.ts
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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'; |
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
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