Skip to content

Commit 3dfa9e8

Browse files
fix(elastic-console): expose model context window sizes
Surface connector context window metadata from the inference plugin so OpenAI-compatible clients can size prompts per model. Co-authored-by: Cursor <cursoragent@cursor.com>
1 parent efcc42a commit 3dfa9e8

3 files changed

Lines changed: 134 additions & 0 deletions

File tree

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
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+
module.exports = {
9+
preset: '@kbn/test',
10+
rootDir: '../../../../..',
11+
roots: [
12+
'<rootDir>/x-pack/platform/plugins/shared/elastic_console/common',
13+
'<rootDir>/x-pack/platform/plugins/shared/elastic_console/public',
14+
'<rootDir>/x-pack/platform/plugins/shared/elastic_console/server',
15+
],
16+
};
Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
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 type { CoreSetup } from '@kbn/core/server';
9+
import { httpServerMock, httpServiceMock, loggingSystemMock } from '@kbn/core/server/mocks';
10+
import { type InferenceConnector, InferenceConnectorType } from '@kbn/inference-common';
11+
import { inferenceMock } from '@kbn/inference-plugin/server/mocks';
12+
import type { ElasticConsolePluginStart, ElasticConsoleStartDependencies } from '../types';
13+
import { isElasticConsoleEnabled } from './is_enabled';
14+
import { registerModelsRoute } from './models';
15+
16+
jest.mock('./is_enabled', () => ({
17+
isElasticConsoleEnabled: jest.fn(),
18+
}));
19+
20+
const isElasticConsoleEnabledMock = isElasticConsoleEnabled as jest.MockedFunction<
21+
typeof isElasticConsoleEnabled
22+
>;
23+
24+
const createConnector = (parts: Partial<InferenceConnector> = {}): InferenceConnector => ({
25+
connectorId: 'connector-1',
26+
name: 'Connector 1',
27+
type: InferenceConnectorType.OpenAI,
28+
config: {},
29+
capabilities: {},
30+
isInferenceEndpoint: false,
31+
isPreconfigured: false,
32+
...parts,
33+
});
34+
35+
describe('registerModelsRoute', () => {
36+
let coreSetup: jest.Mocked<
37+
Pick<CoreSetup<ElasticConsoleStartDependencies, ElasticConsolePluginStart>, 'getStartServices'>
38+
>;
39+
let inference: ReturnType<typeof inferenceMock.createStartContract>;
40+
let router: ReturnType<typeof httpServiceMock.createRouter>;
41+
const logger = loggingSystemMock.create().get();
42+
43+
beforeEach(() => {
44+
jest.spyOn(Date, 'now').mockReturnValue(1700000000000);
45+
46+
inference = inferenceMock.createStartContract();
47+
router = httpServiceMock.createRouter();
48+
coreSetup = {
49+
getStartServices: jest.fn().mockResolvedValue([{}, { inference }]),
50+
};
51+
isElasticConsoleEnabledMock.mockResolvedValue(true);
52+
});
53+
54+
afterEach(() => {
55+
jest.restoreAllMocks();
56+
});
57+
58+
it('registers the OpenAI-compatible models route', () => {
59+
registerModelsRoute({ router, coreSetup: coreSetup as CoreSetup, logger });
60+
61+
const [config] = router.get.mock.calls[0];
62+
expect(config.path).toBe('/internal/elastic_ramen/v1/models');
63+
expect(config.options?.access).toBe('internal');
64+
});
65+
66+
it('returns connector context window sizes when available', async () => {
67+
const request = httpServerMock.createKibanaRequest();
68+
const response = httpServerMock.createResponseFactory();
69+
inference.getConnectorList.mockResolvedValue([
70+
createConnector({
71+
connectorId: 'sonnet-4',
72+
type: InferenceConnectorType.Bedrock,
73+
capabilities: { contextWindowSize: 1000000 },
74+
}),
75+
createConnector({ connectorId: 'unknown-model' }),
76+
]);
77+
78+
registerModelsRoute({ router, coreSetup: coreSetup as CoreSetup, logger });
79+
80+
const [, handler] = router.get.mock.calls[0];
81+
await handler({}, request, response);
82+
83+
expect(inference.getConnectorList).toHaveBeenCalledWith(request);
84+
expect(response.ok).toHaveBeenCalledWith({
85+
body: {
86+
object: 'list',
87+
data: [
88+
{
89+
id: 'sonnet-4',
90+
object: 'model',
91+
created: 1700000000,
92+
owned_by: InferenceConnectorType.Bedrock,
93+
permission: [],
94+
root: 'sonnet-4',
95+
parent: null,
96+
context_window_size: 1000000,
97+
},
98+
{
99+
id: 'unknown-model',
100+
object: 'model',
101+
created: 1700000000,
102+
owned_by: InferenceConnectorType.OpenAI,
103+
permission: [],
104+
root: 'unknown-model',
105+
parent: null,
106+
context_window_size: undefined,
107+
},
108+
],
109+
},
110+
});
111+
112+
const [{ body }] = response.ok.mock.calls[0];
113+
const serializedBody = JSON.parse(JSON.stringify(body));
114+
expect(serializedBody.data[0].context_window_size).toBe(1000000);
115+
expect(serializedBody.data[1]).not.toHaveProperty('context_window_size');
116+
});
117+
});

x-pack/platform/plugins/shared/elastic_console/server/routes/models.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ export const registerModelsRoute = ({
4747
permission: [],
4848
root: connector.connectorId,
4949
parent: null,
50+
context_window_size: connector.capabilities.contextWindowSize,
5051
}));
5152

5253
return response.ok({

0 commit comments

Comments
 (0)