Skip to content

Commit

Permalink
[Obs AI Assistant] Update serverless tests to use cookie auth instead…
Browse files Browse the repository at this point in the history
… of API keys
  • Loading branch information
viduni94 committed Dec 6, 2024
1 parent e4b0b80 commit 1b295fe
Show file tree
Hide file tree
Showing 5 changed files with 68 additions and 55 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ export default function ApiTest({ getService }: FtrProviderContext) {
describe('List connectors', () => {
let roleAuthc: RoleCredentials;
let internalReqHeader: InternalRequestHeader;

before(async () => {
roleAuthc = await svlUserManager.createM2mApiKeyWithRoleScope('editor');
internalReqHeader = svlCommonApi.getInternalRequestHeader();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,7 @@ export default function ApiTest({ getService }: FtrProviderContext) {
let createResponse: Awaited<
SupertestReturnType<'POST /internal/observability_ai_assistant/conversation'>
>;

before(async () => {
createResponse = await observabilityAIAssistantAPIClient
.slsEditor({
Expand Down Expand Up @@ -126,6 +127,7 @@ export default function ApiTest({ getService }: FtrProviderContext) {
})
.expect(404);
});

it('returns the conversation', function () {
// delete user from response to avoid comparing it as it will be different in MKI
delete createResponse.body.user;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ export default function ApiTest({ getService }: FtrProviderContext) {
.expect(200);
expect(res.body).to.eql({});
});

describe('when managing a single entry', () => {
const knowledgeBaseEntry = {
id: 'my-doc-id-1',
Expand Down Expand Up @@ -129,13 +130,16 @@ export default function ApiTest({ getService }: FtrProviderContext) {
.expect(500);
});
});

describe('when managing multiple entries', () => {
before(async () => {
await clearKnowledgeBase(es);
});

afterEach(async () => {
await clearKnowledgeBase(es);
});

const knowledgeBaseEntries = [
{
id: 'my_doc_a',
Expand All @@ -153,6 +157,7 @@ export default function ApiTest({ getService }: FtrProviderContext) {
text: 'My content c',
},
];

it('returns 200 on create', async () => {
await observabilityAIAssistantAPIClient
.slsEditor({
Expand Down Expand Up @@ -221,6 +226,7 @@ export default function ApiTest({ getService }: FtrProviderContext) {
expect(entriesAsc[1].id).to.eql('my_doc_b');
expect(entriesAsc[2].id).to.eql('my_doc_c');
});

it('allows searching', async () => {
await observabilityAIAssistantAPIClient
.slsEditor({
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -71,26 +71,40 @@ export default function ApiTest({ getService }: FtrProviderContext) {
await clearKnowledgeBase(es);

const promises = [
{ roleAuthc: editorRoleAuthc, username: 'editor', isPublic: true },
{ roleAuthc: editorRoleAuthc, username: 'editor', isPublic: false },
{ roleAuthc: johnRoleAuthc, username: 'john', isPublic: true },
{ roleAuthc: johnRoleAuthc, username: 'john', isPublic: false },
].map(async ({ roleAuthc, username, isPublic }) => {
{ username: 'editor', isPublic: true },
{ username: 'editor', isPublic: false },
{ username: 'john', isPublic: true },
{ username: 'john', isPublic: false },
].map(async ({ username, isPublic }) => {
const visibility = isPublic ? 'Public' : 'Private';
await observabilityAIAssistantAPIClient
.slsUser({
endpoint: 'PUT /internal/observability_ai_assistant/kb/user_instructions',
params: {
body: {
id: `${visibility.toLowerCase()}-doc-from-${username}`,
text: `${visibility} user instruction from "${username}"`,
public: isPublic,

if (username === 'editor') {
await observabilityAIAssistantAPIClient
.slsEditor({
endpoint: 'PUT /internal/observability_ai_assistant/kb/user_instructions',
params: {
body: {
id: `${visibility.toLowerCase()}-doc-from-${username}`,
text: `${visibility} user instruction from "${username}"`,
public: isPublic,
},
},
},
roleAuthc,
internalReqHeader,
})
.expect(200);
})
.expect(200);
} else {
await observabilityAIAssistantAPIClient
.slsAdmin({
endpoint: 'PUT /internal/observability_ai_assistant/kb/user_instructions',
params: {
body: {
id: `${visibility.toLowerCase()}-doc-from-${username}`,
text: `${visibility} user instruction from "${username}"`,
public: isPublic,
},
},
})
.expect(200);
}
});

await Promise.all(promises);
Expand Down Expand Up @@ -126,11 +140,10 @@ export default function ApiTest({ getService }: FtrProviderContext) {
});

it('"john" can retrieve their own private instructions and the public instruction', async () => {
const res = await observabilityAIAssistantAPIClient.slsUser({
const res = await observabilityAIAssistantAPIClient.slsAdmin({
endpoint: 'GET /internal/observability_ai_assistant/kb/user_instructions',
roleAuthc: johnRoleAuthc,
internalReqHeader,
});

const instructions = res.body.userInstructions;

const sortByDocId = (data: any) => sortBy(data, 'doc_id');
Expand Down Expand Up @@ -161,7 +174,7 @@ export default function ApiTest({ getService }: FtrProviderContext) {
await clearKnowledgeBase(es);

await observabilityAIAssistantAPIClient
.slsUser({
.slsEditor({
endpoint: 'PUT /internal/observability_ai_assistant/kb/user_instructions',
params: {
body: {
Expand All @@ -170,13 +183,11 @@ export default function ApiTest({ getService }: FtrProviderContext) {
public: true,
},
},
roleAuthc: editorRoleAuthc,
internalReqHeader,
})
.expect(200);

await observabilityAIAssistantAPIClient
.slsUser({
.slsEditor({
endpoint: 'PUT /internal/observability_ai_assistant/kb/user_instructions',
params: {
body: {
Expand All @@ -185,18 +196,15 @@ export default function ApiTest({ getService }: FtrProviderContext) {
public: false,
},
},
roleAuthc: editorRoleAuthc,
internalReqHeader,
})
.expect(200);
});

it('updates the user instruction', async () => {
const res = await observabilityAIAssistantAPIClient.slsUser({
const res = await observabilityAIAssistantAPIClient.slsEditor({
endpoint: 'GET /internal/observability_ai_assistant/kb/user_instructions',
roleAuthc: editorRoleAuthc,
internalReqHeader,
});

const instructions = res.body.userInstructions;

expect(instructions).to.eql([
Expand All @@ -216,10 +224,12 @@ export default function ApiTest({ getService }: FtrProviderContext) {
const userInstructionText =
'Be polite and use language that is easy to understand. Never disagree with the user.';

async function getConversationForUser(roleAuthc: RoleCredentials) {
async function getConversationForUser(username: string) {
const user = username === 'editor' ? 'slsEditor' : 'slsAdmin';

// the user instruction is always created by "editor" user
await observabilityAIAssistantAPIClient
.slsUser({
.slsEditor({
endpoint: 'PUT /internal/observability_ai_assistant/kb/user_instructions',
params: {
body: {
Expand All @@ -228,8 +238,6 @@ export default function ApiTest({ getService }: FtrProviderContext) {
public: false,
},
},
roleAuthc: editorRoleAuthc,
internalReqHeader,
})
.expect(200);

Expand Down Expand Up @@ -257,36 +265,30 @@ export default function ApiTest({ getService }: FtrProviderContext) {
},
];

const createResponse = await observabilityAIAssistantAPIClient
.slsUser({
endpoint: 'POST /internal/observability_ai_assistant/chat/complete',
params: {
body: {
messages,
connectorId,
persist: true,
screenContexts: [],
scopes: ['observability'],
},
const createResponse = await observabilityAIAssistantAPIClient[user]({
endpoint: 'POST /internal/observability_ai_assistant/chat/complete',
params: {
body: {
messages,
connectorId,
persist: true,
screenContexts: [],
scopes: ['observability'],
},
roleAuthc,
internalReqHeader,
})
.expect(200);
},
}).expect(200);

await proxy.waitForAllInterceptorsSettled();
const conversationCreatedEvent = getConversationCreatedEvent(createResponse.body);
const conversationId = conversationCreatedEvent.conversation.id;

const res = await observabilityAIAssistantAPIClient.slsUser({
const res = await observabilityAIAssistantAPIClient[user]({
endpoint: 'GET /internal/observability_ai_assistant/conversation/{conversationId}',
params: {
path: {
conversationId,
},
},
roleAuthc,
internalReqHeader,
});

// wait for all interceptors to be settled
Expand Down Expand Up @@ -319,15 +321,15 @@ export default function ApiTest({ getService }: FtrProviderContext) {
});

it('adds the instruction to the system prompt', async () => {
const conversation = await getConversationForUser(editorRoleAuthc);
const conversation = await getConversationForUser('editor');
const systemMessage = conversation.messages.find(
(message) => message.message.role === MessageRole.System
)!;
expect(systemMessage.message.content).to.contain(userInstructionText);
});

it('does not add the instruction to the context', async () => {
const conversation = await getConversationForUser(editorRoleAuthc);
const conversation = await getConversationForUser('editor');
const contextMessage = conversation.messages.find(
(message) => message.message.name === CONTEXT_FUNCTION_NAME
);
Expand All @@ -341,7 +343,7 @@ export default function ApiTest({ getService }: FtrProviderContext) {
});

it('does not add the instruction conversation for other users', async () => {
const conversation = await getConversationForUser(johnRoleAuthc);
const conversation = await getConversationForUser('john');
const systemMessage = conversation.messages.find(
(message) => message.message.role === MessageRole.System
)!;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ export default function ApiTest({ getService }: FtrProviderContext) {
},
},
];

describe('/api/observability_ai_assistant/chat/complete', function () {
// TODO: https://github.com/elastic/kibana/issues/192751
this.tags(['skipMKI']);
Expand Down Expand Up @@ -108,6 +109,7 @@ export default function ApiTest({ getService }: FtrProviderContext) {
await titleSimulator.complete();

await conversationSimulator.status(200);

if (conversationSimulatorCallback) {
await conversationSimulatorCallback(conversationSimulator);
}
Expand Down

0 comments on commit 1b295fe

Please sign in to comment.