Skip to content

Commit e8263aa

Browse files
alvarezmelissa87elasticmachine
authored and
ymao1
committed
[ML][inference] AI Connector: allows user to manually set context length window (elastic#230005)
## Summary Part of elastic#228716 Because we're mutating the GenAI connector schemas by introducing a new field, the change needed to be done in two different serverless releases. elastic#228371 introduces the model changes for the connector schemas This PR adds the ability to set the `contextWindowLength` config property in the UI for the AI/Inference connector. This PR must be merged after elastic#228371 is deployed in production. The input field will only show up for providers with the 'chat_completion' supported task type. <img width="999" height="1273" alt="image" src="https://github.com/user-attachments/assets/1c7aa343-523e-48d5-9f4b-63c194f4b262" /> Enforces chat_completion task type when setting contextWindowLength: <img width="645" height="938" alt="image" src="https://github.com/user-attachments/assets/2df17065-95ae-4217-842f-a0e360a78b73" /> <img width="932" height="1190" alt="image" src="https://github.com/user-attachments/assets/eb694a47-5f94-4689-8cff-8dc1c4880433" /> ### Checklist Check the PR satisfies following conditions. Reviewers should verify this PR satisfies this list as well. - [ ] Any text added follows [EUI's writing guidelines](https://elastic.github.io/eui/#/guidelines/writing), uses sentence case text and includes [i18n support](https://github.com/elastic/kibana/blob/main/src/platform/packages/shared/kbn-i18n/README.md) - [ ] [Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html) was added for features that require explanation or tutorials - [ ] [Unit or functional tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html) were updated or added to match the most common scenarios - [ ] If a plugin configuration key changed, check if it needs to be allowlisted in the cloud and added to the [docker list](https://github.com/elastic/kibana/blob/main/src/dev/build/tasks/os_packages/docker_generator/resources/base/bin/kibana-docker) - [ ] This was checked for breaking HTTP API changes, and any breaking changes have been approved by the breaking-change committee. The `release_note:breaking` label should be applied in these situations. - [ ] [Flaky Test Runner](https://ci-stats.kibana.dev/trigger_flaky_test_runner/1) was used on any tests changed - [ ] The PR description includes the appropriate Release Notes section, and the correct `release_note:*` label is applied per the [guidelines](https://www.elastic.co/guide/en/kibana/master/contributing.html#kibana-release-notes-process) - [ ] Review the [backport guidelines](https://docs.google.com/document/d/1VyN5k91e5OVumlc0Gb9RPa3h1ewuPE705nRtioPiTvY/edit?usp=sharing) and apply applicable `backport:*` labels. Describe the risk, its severity, and mitigation for each identified risk. Invite stakeholders and evaluate how to proceed before merging. - [ ] [See some risk examples](https://github.com/elastic/kibana/blob/main/RISK_MATRIX.mdx) - [ ] ... --------- Co-authored-by: Elastic Machine <[email protected]>
1 parent 8d7140b commit e8263aa

File tree

11 files changed

+415
-62
lines changed

11 files changed

+415
-62
lines changed

‎x-pack/platform/packages/shared/kbn-inference-endpoint-ui-common/src/components/additional_options_fields.test.tsx‎

Lines changed: 225 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,225 @@
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 React from 'react';
9+
import { AdditionalOptionsFields } from './additional_options_fields';
10+
import { render, screen } from '@testing-library/react';
11+
import { I18nProvider } from '@kbn/i18n-react';
12+
import { Form, useForm } from '@kbn/es-ui-shared-plugin/static/forms/hook_form_lib';
13+
import type { Config } from '../types/types';
14+
import { CHAT_COMPLETION_TASK_TYPE } from '../constants';
15+
16+
const MockFormProvider = ({ children }: { children: React.ReactElement }) => {
17+
const { form } = useForm();
18+
19+
return (
20+
<I18nProvider>
21+
<Form form={form}>{children}</Form>
22+
</I18nProvider>
23+
);
24+
};
25+
26+
const textEmbeddingTask = 'text_embedding';
27+
28+
const mockConfig = {
29+
provider: 'openai',
30+
contextWindowLength: '',
31+
inferenceId: 'testId',
32+
providerConfig: {},
33+
taskType: CHAT_COMPLETION_TASK_TYPE,
34+
} as unknown as Config;
35+
const taskTypeOptions = [
36+
{ id: textEmbeddingTask, label: textEmbeddingTask, value: textEmbeddingTask },
37+
{
38+
id: CHAT_COMPLETION_TASK_TYPE,
39+
label: CHAT_COMPLETION_TASK_TYPE,
40+
value: CHAT_COMPLETION_TASK_TYPE,
41+
},
42+
];
43+
44+
describe('AdditionalOptionsFields', () => {
45+
describe('contextWindowLength', () => {
46+
it('should render contextWindowLength field when task type options include chat_completion', () => {
47+
render(
48+
<MockFormProvider>
49+
<AdditionalOptionsFields
50+
config={mockConfig}
51+
taskTypeOptions={taskTypeOptions}
52+
onTaskTypeOptionsSelect={jest.fn()}
53+
allowContextWindowLength={true}
54+
selectedTaskType={CHAT_COMPLETION_TASK_TYPE}
55+
/>
56+
</MockFormProvider>
57+
);
58+
59+
const inputField = screen.getByTestId('configuration-formrow-contextWindowLength');
60+
const inputLabel = screen.getByTestId('context-window-length-details-label');
61+
const isInputInvalid = inputField.getAttribute('aria-invalid');
62+
expect(inputField).toBeInTheDocument();
63+
expect(inputLabel).toBeInTheDocument();
64+
expect(isInputInvalid).toBeNull();
65+
});
66+
67+
it('should render contextWindowLength field as valid when selectedTaskType is chat_completion', () => {
68+
render(
69+
<MockFormProvider>
70+
<AdditionalOptionsFields
71+
config={{
72+
...mockConfig,
73+
taskType: CHAT_COMPLETION_TASK_TYPE,
74+
contextWindowLength: 500,
75+
}}
76+
taskTypeOptions={taskTypeOptions}
77+
onTaskTypeOptionsSelect={jest.fn()}
78+
allowContextWindowLength={true}
79+
selectedTaskType={CHAT_COMPLETION_TASK_TYPE}
80+
/>
81+
</MockFormProvider>
82+
);
83+
84+
const inputLabel = screen.getByTestId('context-window-length-details-label');
85+
const inputField = screen.getByTestId('configuration-formrow-contextWindowLength');
86+
const isInputInvalid = inputField.getAttribute('aria-invalid');
87+
const inputNumberField = screen.getByTestId('contextWindowLengthNumber');
88+
const inputValue = inputNumberField.getAttribute('value');
89+
expect(inputLabel).toBeInTheDocument();
90+
expect(isInputInvalid).toBeNull();
91+
expect(inputNumberField).toBeInTheDocument();
92+
expect(inputValue).toBe('500');
93+
});
94+
95+
it('should render contextWindowLength field with invalid indicator when selectedTaskType is not chat_completion', () => {
96+
render(
97+
<MockFormProvider>
98+
<AdditionalOptionsFields
99+
config={{ ...mockConfig, taskType: textEmbeddingTask, contextWindowLength: 500 }}
100+
taskTypeOptions={taskTypeOptions}
101+
onTaskTypeOptionsSelect={jest.fn()}
102+
allowContextWindowLength={true}
103+
selectedTaskType={textEmbeddingTask}
104+
/>
105+
</MockFormProvider>
106+
);
107+
108+
const inputField = screen.getByTestId('contextWindowLengthNumber');
109+
const isInvalid = inputField.getAttribute('aria-invalid');
110+
expect(isInvalid).toBe('true');
111+
});
112+
113+
it('should hide contextWindowLength field when chat_completion is not a task option', () => {
114+
render(
115+
<MockFormProvider>
116+
<AdditionalOptionsFields
117+
config={mockConfig}
118+
taskTypeOptions={[]}
119+
onTaskTypeOptionsSelect={jest.fn()}
120+
allowContextWindowLength={true}
121+
selectedTaskType={textEmbeddingTask}
122+
/>
123+
</MockFormProvider>
124+
);
125+
126+
const inputField = screen.queryByTestId('configuration-formrow-contextWindowLength');
127+
const inputLabel = screen.queryByTestId('context-window-length-details-label');
128+
expect(inputField).not.toBeInTheDocument();
129+
expect(inputLabel).not.toBeInTheDocument();
130+
});
131+
132+
it('should hide contextWindowLength field when allowContextWindowLength is false', () => {
133+
render(
134+
<MockFormProvider>
135+
<AdditionalOptionsFields
136+
config={mockConfig}
137+
taskTypeOptions={taskTypeOptions}
138+
onTaskTypeOptionsSelect={jest.fn()}
139+
allowContextWindowLength={false}
140+
selectedTaskType={textEmbeddingTask}
141+
/>
142+
</MockFormProvider>
143+
);
144+
145+
const inputField = screen.queryByTestId('configuration-formrow-contextWindowLength');
146+
const inputLabel = screen.queryByTestId('context-window-length-details-label');
147+
expect(inputField).not.toBeInTheDocument();
148+
expect(inputLabel).not.toBeInTheDocument();
149+
});
150+
});
151+
152+
describe('taskTypeOptions', () => {
153+
it('should render taskTypeOptions section', () => {
154+
render(
155+
<MockFormProvider>
156+
<AdditionalOptionsFields
157+
config={mockConfig}
158+
taskTypeOptions={taskTypeOptions}
159+
onTaskTypeOptionsSelect={jest.fn()}
160+
allowContextWindowLength={false}
161+
/>
162+
</MockFormProvider>
163+
);
164+
165+
const taskTypeTitle = screen.getByTestId('task-type-details-label');
166+
const taskTypeButtons = screen.getByTestId('taskTypeSelect');
167+
expect(taskTypeTitle).toBeInTheDocument();
168+
expect(taskTypeButtons).toBeInTheDocument();
169+
170+
taskTypeOptions.forEach(({ id }) => {
171+
const optionButton = screen.getByTestId(id);
172+
expect(optionButton).toBeInTheDocument();
173+
});
174+
});
175+
176+
it('should have single task type selection when only one task type option', () => {
177+
render(
178+
<MockFormProvider>
179+
<AdditionalOptionsFields
180+
config={mockConfig}
181+
taskTypeOptions={[taskTypeOptions[0]]}
182+
onTaskTypeOptionsSelect={jest.fn()}
183+
allowContextWindowLength={false}
184+
/>
185+
</MockFormProvider>
186+
);
187+
const taskTypeButton = screen.getByTestId('taskTypeSelectSingle');
188+
expect(taskTypeButton).toBeInTheDocument();
189+
});
190+
191+
it('should disable task type selection on edit', () => {
192+
render(
193+
<MockFormProvider>
194+
<AdditionalOptionsFields
195+
config={mockConfig}
196+
taskTypeOptions={taskTypeOptions}
197+
onTaskTypeOptionsSelect={jest.fn()}
198+
allowContextWindowLength={false}
199+
isEdit={true}
200+
/>
201+
</MockFormProvider>
202+
);
203+
const taskTypeButtonDisabled = screen.getByTestId('taskTypeSelectDisabled');
204+
expect(taskTypeButtonDisabled).toBeInTheDocument();
205+
});
206+
});
207+
208+
describe('inferenceId', () => {
209+
it('should render the inferenceId section', () => {
210+
render(
211+
<MockFormProvider>
212+
<AdditionalOptionsFields
213+
config={mockConfig}
214+
taskTypeOptions={taskTypeOptions}
215+
onTaskTypeOptionsSelect={jest.fn()}
216+
allowContextWindowLength={false}
217+
/>
218+
</MockFormProvider>
219+
);
220+
221+
const inputField = screen.getByTestId('inference-endpoint-input-field');
222+
expect(inputField).toBeInTheDocument();
223+
});
224+
});
225+
});

0 commit comments

Comments
 (0)