Skip to content

Commit

Permalink
Revert "[Obs AI Assistant] Boost user prompt in recall (#184933)"
Browse files Browse the repository at this point in the history
This reverts commit baa22bb.
  • Loading branch information
delanni committed Jun 10, 2024
1 parent 8cd4858 commit 7090546
Show file tree
Hide file tree
Showing 11 changed files with 80 additions and 53 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,10 @@ export function createService({
return of(
createFunctionRequestMessage({
name: 'context',
args: {
queries: [],
categories: [],
},
}),
createFunctionResponseMessage({
name: 'context',
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -40,10 +40,34 @@ export function registerContextFunction({
description:
'This function provides context as to what the user is looking at on their screen, and recalled documents from the knowledge base that matches their query',
visibility: FunctionVisibility.Internal,
parameters: {
type: 'object',
properties: {
queries: {
type: 'array',
description: 'The query for the semantic search',
items: {
type: 'string',
},
},
categories: {
type: 'array',
description:
'Categories of internal documentation that you want to search for. By default internal documentation will be excluded. Use `apm` to get internal APM documentation, `lens` to get internal Lens documentation, or both.',
items: {
type: 'string',
enum: ['apm', 'lens'],
},
},
},
required: ['queries', 'categories'],
} as const,
},
async ({ messages, screenContexts, chat }, signal) => {
async ({ arguments: args, messages, screenContexts, chat }, signal) => {
const { analytics } = (await resources.context.core).coreStart;

const { queries, categories } = args;

async function getContext() {
const screenDescription = compact(
screenContexts.map((context) => context.screenDescription)
Expand All @@ -70,29 +94,38 @@ export function registerContextFunction({
messages.filter((message) => message.message.role === MessageRole.User)
);

const userPrompt = userMessage?.message.content;
const queries = [{ text: userPrompt, boost: 3 }, { text: screenDescription }].filter(
({ text }) => text
) as Array<{ text: string; boost?: number }>;
const nonEmptyQueries = compact(queries);

const queriesOrUserPrompt = nonEmptyQueries.length
? nonEmptyQueries
: compact([userMessage?.message.content]);

queriesOrUserPrompt.push(screenDescription);

const suggestions = await retrieveSuggestions({
client,
categories,
queries: queriesOrUserPrompt,
});

const suggestions = await retrieveSuggestions({ client, queries });
if (suggestions.length === 0) {
return { content };
return {
content,
};
}

try {
const { relevantDocuments, scores } = await scoreSuggestions({
suggestions,
screenDescription,
userPrompt,
queries: queriesOrUserPrompt,
messages,
chat,
signal,
logger: resources.logger,
});

analytics.reportEvent<RecallRanking>(RecallRankingEventType, {
prompt: queries.map((query) => query.text).join('|'),
prompt: queriesOrUserPrompt.join('|'),
scoredDocuments: suggestions.map((suggestion) => {
const llmScore = scores.find((score) => score.id === suggestion.id);
return {
Expand Down Expand Up @@ -145,12 +178,15 @@ export function registerContextFunction({
async function retrieveSuggestions({
queries,
client,
categories,
}: {
queries: Array<{ text: string; boost?: number }>;
queries: string[];
client: ObservabilityAIAssistantClient;
categories: Array<'apm' | 'lens'>;
}) {
const recallResponse = await client.recall({
queries,
categories,
});

return recallResponse.entries.map((entry) => omit(entry, 'labels', 'is_correction'));
Expand All @@ -172,16 +208,14 @@ const scoreFunctionArgumentsRt = t.type({
async function scoreSuggestions({
suggestions,
messages,
userPrompt,
screenDescription,
queries,
chat,
signal,
logger,
}: {
suggestions: Awaited<ReturnType<typeof retrieveSuggestions>>;
messages: Message[];
userPrompt: string | undefined;
screenDescription: string;
queries: string[];
chat: FunctionCallChatFunction;
signal: AbortSignal;
logger: Logger;
Expand All @@ -203,10 +237,7 @@ async function scoreSuggestions({
- The document contains new information not mentioned before in the conversation
Question:
${userPrompt}
Screen description:
${screenDescription}
${queries.join('\n')}
Documents:
${JSON.stringify(indexedSuggestions, null, 2)}`);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -65,16 +65,7 @@ const functionRecallRoute = createObservabilityAIAssistantServerRoute({
params: t.type({
body: t.intersection([
t.type({
queries: t.array(
t.intersection([
t.type({
text: t.string,
}),
t.partial({
boost: t.number,
}),
])
),
queries: t.array(nonEmptyStringRt),
}),
t.partial({
categories: t.array(t.string),
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -28,5 +28,9 @@ export function getContextFunctionRequestIfNeeded(

return createFunctionRequestMessage({
name: CONTEXT_FUNCTION_NAME,
args: {
queries: [],
categories: [],
},
});
}
Original file line number Diff line number Diff line change
Expand Up @@ -1232,6 +1232,7 @@ describe('Observability AI Assistant client', () => {
role: MessageRole.Assistant,
function_call: {
name: CONTEXT_FUNCTION_NAME,
arguments: JSON.stringify({ queries: [], categories: [] }),
trigger: MessageRole.Assistant,
},
},
Expand Down Expand Up @@ -1455,6 +1456,7 @@ describe('Observability AI Assistant client', () => {
role: MessageRole.Assistant,
function_call: {
name: CONTEXT_FUNCTION_NAME,
arguments: JSON.stringify({ queries: [], categories: [] }),
trigger: MessageRole.Assistant,
},
},
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -694,7 +694,7 @@ export class ObservabilityAIAssistantClient {
queries,
categories,
}: {
queries: Array<{ text: string; boost?: number }>;
queries: string[];
categories?: string[];
}): Promise<{ entries: RecalledEntry[] }> => {
return this.dependencies.knowledgeBaseService.recall({
Expand Down Expand Up @@ -757,9 +757,11 @@ export class ObservabilityAIAssistantClient {
};

fetchUserInstructions = async () => {
return this.dependencies.knowledgeBaseService.getUserInstructions(
const userInstructions = await this.dependencies.knowledgeBaseService.getUserInstructions(
this.dependencies.namespace,
this.dependencies.user
);

return userInstructions;
};
}
Original file line number Diff line number Diff line change
Expand Up @@ -303,20 +303,19 @@ export class KnowledgeBaseService {
user,
modelId,
}: {
queries: Array<{ text: string; boost?: number }>;
queries: string[];
categories?: string[];
namespace: string;
user?: { name: string };
modelId: string;
}): Promise<RecalledEntry[]> {
const query = {
bool: {
should: queries.map(({ text, boost = 1 }) => ({
should: queries.map((text) => ({
text_expansion: {
'ml.tokens': {
model_text: text,
model_id: modelId,
boost,
},
},
})),
Expand Down Expand Up @@ -386,7 +385,7 @@ export class KnowledgeBaseService {
uiSettingsClient,
modelId,
}: {
queries: Array<{ text: string; boost?: number }>;
queries: string[];
asCurrentUser: ElasticsearchClient;
uiSettingsClient: IUiSettingsClient;
modelId: string;
Expand Down Expand Up @@ -415,16 +414,15 @@ export class KnowledgeBaseService {
const vectorField = `${ML_INFERENCE_PREFIX}${field}_expanded.predicted_value`;
const modelField = `${ML_INFERENCE_PREFIX}${field}_expanded.model_id`;

return queries.map(({ text, boost = 1 }) => {
return queries.map((query) => {
return {
bool: {
should: [
{
text_expansion: {
[vectorField]: {
model_text: text,
model_text: query,
model_id: modelId,
boost,
},
},
},
Expand Down Expand Up @@ -472,7 +470,7 @@ export class KnowledgeBaseService {
asCurrentUser,
uiSettingsClient,
}: {
queries: Array<{ text: string; boost?: number }>;
queries: string[];
categories?: string[];
user?: { name: string };
namespace: string;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ describe('<ChatBody>', () => {
role: 'assistant',
function_call: {
name: CONTEXT_FUNCTION_NAME,
arguments: '{}',
arguments: '{"queries":[],"categories":[]}',
trigger: 'assistant',
},
content: '',
Expand Down Expand Up @@ -88,7 +88,7 @@ describe('<ChatBody>', () => {
role: 'assistant',
function_call: {
name: CONTEXT_FUNCTION_NAME,
arguments: '{}',
arguments: '{"queries":[],"categories":[]}',
trigger: 'assistant',
},
content: '',
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -193,6 +193,7 @@ export default function ApiTest({ getService }: FtrProviderContext) {
role: MessageRole.Assistant,
function_call: {
name: 'context',
arguments: JSON.stringify({ queries: [], categories: [] }),
trigger: MessageRole.Assistant,
},
},
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,6 @@ export default function ApiTest({ getService }: FtrProviderContext) {
format,
})
.set('kbn-xsrf', 'foo')
.set('elastic-api-version', '2023-10-31')
.send({
messages,
connectorId,
Expand All @@ -84,20 +83,13 @@ export default function ApiTest({ getService }: FtrProviderContext) {
if (err) {
return reject(err);
}
if (response.status !== 200) {
return reject(new Error(`${response.status}: ${JSON.stringify(response.body)}`));
}
return resolve(response);
});
});

const [conversationSimulator, titleSimulator] = await Promise.race([
Promise.all([
conversationInterceptor.waitForIntercept(),
titleInterceptor.waitForIntercept(),
]),
// make sure any request failures (like 400s) are properly propagated
responsePromise.then(() => []),
const [conversationSimulator, titleSimulator] = await Promise.all([
conversationInterceptor.waitForIntercept(),
titleInterceptor.waitForIntercept(),
]);

await titleSimulator.status(200);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,7 @@ export default function ApiTest({ getService, getPageObjects }: FtrProviderConte
content: '',
function_call: {
name: 'context',
arguments: '{}',
arguments: '{"queries":[],"categories":[]}',
trigger: MessageRole.Assistant,
},
},
Expand Down Expand Up @@ -290,6 +290,7 @@ export default function ApiTest({ getService, getPageObjects }: FtrProviderConte

expect(pick(contextRequest.function_call, 'name', 'arguments')).to.eql({
name: 'context',
arguments: JSON.stringify({ queries: [], categories: [] }),
});

expect(contextResponse.name).to.eql('context');
Expand Down Expand Up @@ -353,6 +354,7 @@ export default function ApiTest({ getService, getPageObjects }: FtrProviderConte

expect(pick(contextRequest.function_call, 'name', 'arguments')).to.eql({
name: 'context',
arguments: JSON.stringify({ queries: [], categories: [] }),
});

expect(contextResponse.name).to.eql('context');
Expand Down

0 comments on commit 7090546

Please sign in to comment.