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

[8.x][Obs AI Assistant] Serverless API integration tests (#192219) #194582

Merged
merged 1 commit into from
Oct 1, 2024
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
1 change: 1 addition & 0 deletions .buildkite/ftr_oblt_serverless_configs.yml
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ disabled:
- x-pack/test_serverless/api_integration/test_suites/observability/config.feature_flags.ts
- x-pack/test_serverless/api_integration/test_suites/observability/common_configs/config.group1.ts
- x-pack/test_serverless/api_integration/test_suites/observability/fleet/config.ts
- x-pack/test_serverless/api_integration/test_suites/observability/ai_assistant/config.ts
- x-pack/test_serverless/functional/test_suites/observability/config.ts
- x-pack/test_serverless/functional/test_suites/observability/config.examples.ts
- x-pack/test_serverless/functional/test_suites/observability/config.feature_flags.ts
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,8 @@ export type StreamingChatResponseEvent =
| ConversationUpdateEvent
| MessageAddEvent
| ChatCompletionErrorEvent
| TokenCountEvent;
| TokenCountEvent
| BufferFlushEvent;

export type StreamingChatResponseEventWithoutError = Exclude<
StreamingChatResponseEvent,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -47,8 +47,8 @@ export interface RecalledEntry {
function isModelMissingOrUnavailableError(error: Error) {
return (
error instanceof errors.ResponseError &&
(error.body.error.type === 'resource_not_found_exception' ||
error.body.error.type === 'status_exception')
(error.body?.error?.type === 'resource_not_found_exception' ||
error.body?.error?.type === 'status_exception')
);
}
function isCreateModelValidationError(error: Error) {
Expand Down Expand Up @@ -127,7 +127,7 @@ export class KnowledgeBaseService {
};

const installModel = async () => {
this.dependencies.logger.info('Installing ELSER model');
this.dependencies.logger.info(`Installing ${elserModelId} model`);
try {
await this.dependencies.esClient.asInternalUser.ml.putTrainedModel(
{
Expand All @@ -146,12 +146,12 @@ export class KnowledgeBaseService {
throw error;
}
}
this.dependencies.logger.info('Finished installing ELSER model');
this.dependencies.logger.info(`Finished installing ${elserModelId} model`);
};

const pollForModelInstallCompleted = async () => {
await pRetry(async () => {
this.dependencies.logger.info('Polling installation of ELSER model');
this.dependencies.logger.info(`Polling installation of ${elserModelId} model`);
const modelInstalledAndReady = await isModelInstalledAndReady();
if (!modelInstalledAndReady) {
throwKnowledgeBaseNotReady({
Expand All @@ -169,7 +169,7 @@ export class KnowledgeBaseService {
wait_for: 'fully_allocated',
});
} catch (error) {
this.dependencies.logger.debug('Error starting model deployment');
this.dependencies.logger.debug(`Error starting ${elserModelId} model deployment`);
this.dependencies.logger.debug(error);
if (!isModelMissingOrUnavailableError(error)) {
throw error;
Expand All @@ -191,13 +191,13 @@ export class KnowledgeBaseService {
return Promise.resolve();
}

this.dependencies.logger.debug('Model is not allocated yet');
this.dependencies.logger.debug(`${elserModelId} model is not allocated yet`);
this.dependencies.logger.debug(() => JSON.stringify(response));

throw gatewayTimeout();
}, retryOptions);

this.dependencies.logger.info('Model is ready');
this.dependencies.logger.info(`${elserModelId} model is ready`);
this.ensureTaskScheduled();
};

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,8 +21,9 @@ export async function createKnowledgeBaseModel(ml: ReturnType<typeof MachineLear
field_names: ['text_field'],
},
};
await ml.api.importTrainedModel(TINY_ELSER.name, TINY_ELSER.id, config);
// necessary for MKI, check indices before importing model. compatible with stateful
await ml.api.assureMlStatsIndexExists();
await ml.api.importTrainedModel(TINY_ELSER.name, TINY_ELSER.id, config);
}

export async function deleteKnowledgeBaseModel(ml: ReturnType<typeof MachineLearningProvider>) {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,77 @@
/*
* 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 { ToolingLog } from '@kbn/tooling-log';
import type {
InternalRequestHeader,
RoleCredentials,
SupertestWithoutAuthProviderType,
} from '../../../../../shared/services';

export async function deleteActionConnector({
supertest,
connectorId,
log,
roleAuthc,
internalReqHeader,
}: {
supertest: SupertestWithoutAuthProviderType;
connectorId: string;
log: ToolingLog;
roleAuthc: RoleCredentials;
internalReqHeader: InternalRequestHeader;
}) {
try {
await supertest
.delete(`/api/actions/connector/${connectorId}`)
.set(roleAuthc.apiKeyHeader)
.set(internalReqHeader)
.expect(204);
} catch (e) {
log.error(`Failed to delete action connector with id ${connectorId} due to: ${e}`);
throw e;
}
}

export async function createProxyActionConnector({
log,
supertest,
port,
roleAuthc,
internalReqHeader,
}: {
log: ToolingLog;
supertest: SupertestWithoutAuthProviderType;
port: number;
roleAuthc: RoleCredentials;
internalReqHeader: InternalRequestHeader;
}) {
try {
const res = await supertest
.post('/api/actions/connector')
.set(roleAuthc.apiKeyHeader)
.set(internalReqHeader)
.send({
name: 'OpenAI Proxy',
connector_type_id: '.gen-ai',
config: {
apiProvider: 'OpenAI',
apiUrl: `http://localhost:${port}`,
},
secrets: {
apiKey: 'my-api-key',
},
})
.expect(200);

const connectorId = res.body.id as string;
return connectorId;
} catch (e) {
log.error(`Failed to create action connector due to: ${e}`);
throw e;
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
/*
* 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 { GenericFtrProviderContext } from '@kbn/test';
import { InheritedServices, InheritedFtrProviderContext } from '../../../../services';
import { ObservabilityAIAssistantApiClient } from './observability_ai_assistant_api_client';

export type ObservabilityAIAssistantServices = InheritedServices & {
observabilityAIAssistantAPIClient: (
context: InheritedFtrProviderContext
) => Promise<ObservabilityAIAssistantApiClient>;
};

export type FtrProviderContext = GenericFtrProviderContext<ObservabilityAIAssistantServices, {}>;
Original file line number Diff line number Diff line change
@@ -0,0 +1,191 @@
/*
* 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 type {
APIReturnType,
ObservabilityAIAssistantAPIClientRequestParamsOf,
ObservabilityAIAssistantAPIEndpoint,
} from '@kbn/observability-ai-assistant-plugin/public';
import { formatRequest } from '@kbn/server-route-repository';
import supertest from 'supertest';
import { Subtract } from 'utility-types';
import { format } from 'url';
import { Config } from '@kbn/test';
import { InheritedFtrProviderContext } from '../../../../services';
import type { InternalRequestHeader, RoleCredentials } from '../../../../../shared/services';

export function getObservabilityAIAssistantApiClient({
svlSharedConfig,
}: {
svlSharedConfig: Config;
}) {
const kibanaServer = svlSharedConfig.get('servers.kibana');
const cAuthorities = svlSharedConfig.get('servers.kibana.certificateAuthorities');

const url = format({
...kibanaServer,
auth: false, // don't use auth in serverless
});

return createObservabilityAIAssistantApiClient(supertest.agent(url, { ca: cAuthorities }));
}

type ObservabilityAIAssistantApiClientKey = 'slsUser';
export type ObservabilityAIAssistantApiClient = Record<
ObservabilityAIAssistantApiClientKey,
Awaited<ReturnType<typeof getObservabilityAIAssistantApiClient>>
>;
export function createObservabilityAIAssistantApiClient(st: supertest.Agent) {
return <TEndpoint extends ObservabilityAIAssistantAPIEndpoint>(
options: {
type?: 'form-data';
endpoint: TEndpoint;
roleAuthc: RoleCredentials;
internalReqHeader: InternalRequestHeader;
} & ObservabilityAIAssistantAPIClientRequestParamsOf<TEndpoint> & {
params?: { query?: { _inspect?: boolean } };
}
): SupertestReturnType<TEndpoint> => {
const { endpoint, type, roleAuthc, internalReqHeader } = options;

const params = 'params' in options ? (options.params as Record<string, any>) : {};

const { method, pathname, version } = formatRequest(endpoint, params.path);
const url = format({ pathname, query: params?.query });

const headers: Record<string, string> = { ...internalReqHeader, ...roleAuthc.apiKeyHeader };

if (version) {
headers['Elastic-Api-Version'] = version;
}

let res: supertest.Test;
if (type === 'form-data') {
const fields: Array<[string, any]> = Object.entries(params.body);
const formDataRequest = st[method](url)
.set(headers)
.set('Content-type', 'multipart/form-data');
for (const field of fields) {
void formDataRequest.field(field[0], field[1]);
}

res = formDataRequest;
} else if (params.body) {
res = st[method](url).send(params.body).set(headers);
} else {
res = st[method](url).set(headers);
}

return res as unknown as SupertestReturnType<TEndpoint>;
};
}

export type ObservabilityAIAssistantAPIClient = ReturnType<
typeof createObservabilityAIAssistantApiClient
>;

type WithoutPromise<T extends Promise<any>> = Subtract<T, Promise<any>>;

// this is a little intense, but without it, method overrides are lost
// e.g., {
// end(one:string)
// end(one:string, two:string)
// }
// would lose the first signature. This keeps up to eight signatures.
type OverloadedParameters<T> = T extends {
(...args: infer A1): any;
(...args: infer A2): any;
(...args: infer A3): any;
(...args: infer A4): any;
(...args: infer A5): any;
(...args: infer A6): any;
(...args: infer A7): any;
(...args: infer A8): any;
}
? A1 | A2 | A3 | A4 | A5 | A6 | A7 | A8
: T extends {
(...args: infer A1): any;
(...args: infer A2): any;
(...args: infer A3): any;
(...args: infer A4): any;
(...args: infer A5): any;
(...args: infer A6): any;
(...args: infer A7): any;
}
? A1 | A2 | A3 | A4 | A5 | A6 | A7
: T extends {
(...args: infer A1): any;
(...args: infer A2): any;
(...args: infer A3): any;
(...args: infer A4): any;
(...args: infer A5): any;
(...args: infer A6): any;
}
? A1 | A2 | A3 | A4 | A5 | A6
: T extends {
(...args: infer A1): any;
(...args: infer A2): any;
(...args: infer A3): any;
(...args: infer A4): any;
(...args: infer A5): any;
}
? A1 | A2 | A3 | A4 | A5
: T extends {
(...args: infer A1): any;
(...args: infer A2): any;
(...args: infer A3): any;
(...args: infer A4): any;
}
? A1 | A2 | A3 | A4
: T extends {
(...args: infer A1): any;
(...args: infer A2): any;
(...args: infer A3): any;
}
? A1 | A2 | A3
: T extends {
(...args: infer A1): any;
(...args: infer A2): any;
}
? A1 | A2
: T extends (...args: infer A) => any
? A
: any;

type OverrideReturnType<T extends (...args: any[]) => any, TNextReturnType> = (
...args: OverloadedParameters<T>
) => WithoutPromise<ReturnType<T>> & TNextReturnType;

type OverwriteThisMethods<T extends Record<string, any>, TNextReturnType> = TNextReturnType & {
[key in keyof T]: T[key] extends (...args: infer TArgs) => infer TReturnType
? TReturnType extends Promise<any>
? OverrideReturnType<T[key], TNextReturnType>
: (...args: TArgs) => TReturnType
: T[key];
};

export type SupertestReturnType<TEndpoint extends ObservabilityAIAssistantAPIEndpoint> =
OverwriteThisMethods<
WithoutPromise<supertest.Test>,
Promise<{
text: string;
status: number;
body: APIReturnType<TEndpoint>;
}>
>;

export async function getObservabilityAIAssistantApiClientService({
getService,
}: InheritedFtrProviderContext): Promise<ObservabilityAIAssistantApiClient> {
const svlSharedConfig = getService('config');
// defaults to elastic_admin user when used without auth
return {
slsUser: await getObservabilityAIAssistantApiClient({
svlSharedConfig,
}),
};
}
Loading