[Obs AI Assistant] Use const for function names (#184830)

Some of the AI Assistant functions have very generic names. One example
is the `context` function. When stumbling upon code like

```
const contextRequest = functionClient.hasFunction('context')
```

... it is quite difficult to navigate to the context function
implementation without knowing it's exact location. Searching for
`context` is futile since it's a term used for many different things.

Using a constant to refer to a function makes it much easier to navigate
to where the function is registered. It also avoids typos but that's a
side-effect, not the main motivation for this.

---------

Co-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com>
This commit is contained in:
Søren Louv-Jansen 2024-06-06 18:55:40 +02:00 committed by GitHub
parent 966736e4b4
commit 5743aa09cf
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 75 additions and 50 deletions

View file

@ -26,6 +26,8 @@ import { parseSuggestionScores } from './parse_suggestion_scores';
const MAX_TOKEN_COUNT_FOR_DATA_ON_SCREEN = 1000;
export const CONTEXT_FUNCTION_NAME = 'context';
export function registerContextFunction({
client,
functions,
@ -34,7 +36,7 @@ export function registerContextFunction({
}: FunctionRegistrationParameters & { isKnowledgeBaseAvailable: boolean }) {
functions.registerFunction(
{
name: 'context',
name: CONTEXT_FUNCTION_NAME,
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,
@ -157,7 +159,7 @@ export function registerContextFunction({
.then(({ content, data }) => {
subscriber.next(
createFunctionResponseMessage({
name: 'context',
name: CONTEXT_FUNCTION_NAME,
content,
data,
})

View file

@ -9,13 +9,15 @@ import { FunctionRegistrationParameters } from '..';
import { FunctionVisibility } from '../../../common/functions/types';
import { getRelevantFieldNames } from './get_relevant_field_names';
export const GET_DATASET_INFO_FUNCTION_NAME = 'get_dataset_info';
export function registerGetDatasetInfoFunction({
resources,
functions,
}: FunctionRegistrationParameters) {
functions.registerFunction(
{
name: 'get_dataset_info',
name: GET_DATASET_INFO_FUNCTION_NAME,
visibility: FunctionVisibility.AssistantOnly,
description: `Use this function to get information about indices/datasets available and the fields available on them.

View file

@ -6,13 +6,17 @@
*/
import dedent from 'dedent';
import { registerContextFunction } from './context';
import { registerSummarizationFunction } from './summarize';
import { CONTEXT_FUNCTION_NAME, registerContextFunction } from './context';
import { registerSummarizationFunction, SUMMARIZE_FUNCTION_NAME } from './summarize';
import type { RegistrationCallback } from '../service/types';
import { registerElasticsearchFunction } from './elasticsearch';
import { registerGetDatasetInfoFunction } from './get_dataset_info';
import { GET_DATASET_INFO_FUNCTION_NAME, registerGetDatasetInfoFunction } from './get_dataset_info';
import { registerKibanaFunction } from './kibana';
import { registerExecuteConnectorFunction } from './execute_connector';
import { GET_DATA_ON_SCREEN_FUNCTION_NAME } from '../service/chat_function_client';
// cannot be imported from x-pack/plugins/observability_solution/observability_ai_assistant_app/server/functions/query/index.ts due to circular dependency
export const QUERY_FUNCTION_NAME = 'query';
export type FunctionRegistrationParameters = Omit<
Parameters<RegistrationCallback>[0],
@ -59,30 +63,30 @@ export const registerFunctions: RegistrationCallback = async ({
functions.registerInstruction(({ availableFunctionNames }) => {
const instructions: string[] = [];
if (availableFunctionNames.includes('get_dataset_info')) {
instructions.push(`You MUST use the get_dataset_info function ${
functions.hasFunction('get_apm_dataset_info') ? 'or get_apm_dataset_info' : ''
} function before calling the "query" or "changes" function.
if (availableFunctionNames.includes(GET_DATASET_INFO_FUNCTION_NAME)) {
instructions.push(`You MUST use the "${GET_DATASET_INFO_FUNCTION_NAME}" ${
functions.hasFunction('get_apm_dataset_info') ? 'or the get_apm_dataset_info' : ''
} function before calling the "${QUERY_FUNCTION_NAME}" or the "changes" functions.
If a function requires an index, you MUST use the results from the dataset info functions.`);
If a function requires an index, you MUST use the results from the dataset info functions.`);
}
if (availableFunctionNames.includes('get_data_on_screen')) {
instructions.push(`You have access to data on the screen by calling the "get_data_on_screen" function.
Use it to help the user understand what they are looking at. A short summary of what they are looking at is available in the return of the "context" function.
Data that is compact enough automatically gets included in the response for the "context" function.`);
if (availableFunctionNames.includes(GET_DATA_ON_SCREEN_FUNCTION_NAME)) {
instructions.push(`You have access to data on the screen by calling the "${GET_DATA_ON_SCREEN_FUNCTION_NAME}" function.
Use it to help the user understand what they are looking at. A short summary of what they are looking at is available in the return of the "${CONTEXT_FUNCTION_NAME}" function.
Data that is compact enough automatically gets included in the response for the "${CONTEXT_FUNCTION_NAME}" function.`);
}
if (isReady) {
if (availableFunctionNames.includes('summarize')) {
instructions.push(`You can use the "summarize" functions to store new information you have learned in a knowledge database.
if (availableFunctionNames.includes(SUMMARIZE_FUNCTION_NAME)) {
instructions.push(`You can use the "${SUMMARIZE_FUNCTION_NAME}" function to store new information you have learned in a knowledge database.
Only use this function when the user asks for it.
All summaries MUST be created in English, even if the conversation was carried out in a different language.`);
}
if (availableFunctionNames.includes('context')) {
if (availableFunctionNames.includes(CONTEXT_FUNCTION_NAME)) {
instructions.push(
`Additionally, you can use the "context" function to retrieve relevant information from the knowledge database.`
`Additionally, you can use the "${CONTEXT_FUNCTION_NAME}" function to retrieve relevant information from the knowledge database.`
);
}
} else {

View file

@ -8,13 +8,15 @@
import type { FunctionRegistrationParameters } from '.';
import { KnowledgeBaseEntryRole } from '../../common';
export const SUMMARIZE_FUNCTION_NAME = 'summarize';
export function registerSummarizationFunction({
client,
functions,
}: FunctionRegistrationParameters) {
functions.registerFunction(
{
name: 'summarize',
name: SUMMARIZE_FUNCTION_NAME,
description: `Use this function to store facts in the knowledge database if the user requests it.
You can score the learnings with a confidence metric, whether it is a correction on a previous learning.
An embedding will be created that you can recall later with a semantic search.

View file

@ -5,7 +5,7 @@
* 2.0.
*/
import dedent from 'dedent';
import { ChatFunctionClient } from '.';
import { ChatFunctionClient, GET_DATA_ON_SCREEN_FUNCTION_NAME } from '.';
import { FunctionVisibility } from '../../../common/functions/types';
describe('chatFunctionClient', () => {
@ -88,7 +88,7 @@ describe('chatFunctionClient', () => {
expect(functions[0]).toEqual({
definition: {
description: expect.any(String),
name: 'get_data_on_screen',
name: GET_DATA_ON_SCREEN_FUNCTION_NAME,
parameters: expect.any(Object),
visibility: FunctionVisibility.AssistantOnly,
},
@ -103,7 +103,7 @@ describe('chatFunctionClient', () => {
const result = await client.executeFunction({
chat: jest.fn(),
name: 'get_data_on_screen',
name: GET_DATA_ON_SCREEN_FUNCTION_NAME,
args: JSON.stringify({ data: ['my_dummy_data'] }),
messages: [],
signal: new AbortController().signal,

View file

@ -31,6 +31,8 @@ const ajv = new Ajv({
strict: false,
});
export const GET_DATA_ON_SCREEN_FUNCTION_NAME = 'get_data_on_screen';
export class ChatFunctionClient {
private readonly instructions: RegisteredInstruction[] = [];
private readonly functionRegistry: FunctionHandlerRegistry = new Map();
@ -46,7 +48,7 @@ export class ChatFunctionClient {
if (allData.length) {
this.registerFunction(
{
name: 'get_data_on_screen',
name: GET_DATA_ON_SCREEN_FUNCTION_NAME,
description: dedent(`Get data that is on the screen:
${allData.map((data) => `${data.name}: ${data.description}`).join('\n')}
`),

View file

@ -5,6 +5,7 @@
* 2.0.
*/
import { CONTEXT_FUNCTION_NAME } from '../../../../functions/context';
import { FunctionDefinition } from '../../../../../common';
import { TOOL_USE_END, TOOL_USE_START } from './constants';
@ -17,7 +18,7 @@ export function getSystemMessageInstructions({
return `In this environment, you have access to a set of tools you can use to answer the user's question.
${
functions?.find((fn) => fn.name === 'context')
functions?.find((fn) => fn.name === CONTEXT_FUNCTION_NAME)
? `The "context" tool is ALWAYS used after a user question. Even if it was used before, your job is to answer the last user question,
even if the "context" tool was executed after that. Consider the tools you need to answer the user's question.`
: ''

View file

@ -8,6 +8,7 @@
import { findLastIndex } from 'lodash';
import { Message, MessageAddEvent, MessageRole } from '../../../common';
import { createFunctionRequestMessage } from '../../../common/utils/create_function_request_message';
import { CONTEXT_FUNCTION_NAME } from '../../functions/context';
export function getContextFunctionRequestIfNeeded(
messages: Message[]
@ -19,14 +20,14 @@ export function getContextFunctionRequestIfNeeded(
const hasContextSinceLastUserMessage = messages
.slice(indexOfLastUserMessage)
.some((message) => message.message.name === 'context');
.some((message) => message.message.name === CONTEXT_FUNCTION_NAME);
if (hasContextSinceLastUserMessage) {
return undefined;
}
return createFunctionRequestMessage({
name: 'context',
name: CONTEXT_FUNCTION_NAME,
args: {
queries: [],
categories: [],

View file

@ -24,6 +24,7 @@ import {
StreamingChatResponseEventType,
} from '../../../common/conversation_complete';
import { createFunctionResponseMessage } from '../../../common/utils/create_function_response_message';
import { CONTEXT_FUNCTION_NAME } from '../../functions/context';
import { ChatFunctionClient } from '../chat_function_client';
import type { KnowledgeBaseService } from '../knowledge_base_service';
import { observableIntoStream } from '../util/observable_into_stream';
@ -145,7 +146,7 @@ describe('Observability AI Assistant client', () => {
functionClientMock.getFunctions.mockReturnValue([]);
functionClientMock.hasFunction.mockImplementation((name) => {
return name !== 'context';
return name !== CONTEXT_FUNCTION_NAME;
});
functionClientMock.hasAction.mockReturnValue(false);
@ -1230,7 +1231,7 @@ describe('Observability AI Assistant client', () => {
content: '',
role: MessageRole.Assistant,
function_call: {
name: 'context',
name: CONTEXT_FUNCTION_NAME,
arguments: JSON.stringify({ queries: [], categories: [] }),
trigger: MessageRole.Assistant,
},
@ -1248,7 +1249,7 @@ describe('Observability AI Assistant client', () => {
message: {
content: JSON.stringify([{ id: 'my_document', text: 'My document' }]),
role: MessageRole.User,
name: 'context',
name: CONTEXT_FUNCTION_NAME,
},
},
});
@ -1440,7 +1441,7 @@ describe('Observability AI Assistant client', () => {
it('executes the context function', async () => {
expect(functionClientMock.executeFunction).toHaveBeenCalledWith(
expect.objectContaining({ name: 'context' })
expect.objectContaining({ name: CONTEXT_FUNCTION_NAME })
);
});
@ -1454,7 +1455,7 @@ describe('Observability AI Assistant client', () => {
content: '',
role: MessageRole.Assistant,
function_call: {
name: 'context',
name: CONTEXT_FUNCTION_NAME,
arguments: JSON.stringify({ queries: [], categories: [] }),
trigger: MessageRole.Assistant,
},

View file

@ -53,6 +53,7 @@ import {
type Message,
} from '../../../common/types';
import { withoutTokenCountEvents } from '../../../common/utils/without_token_count_events';
import { CONTEXT_FUNCTION_NAME } from '../../functions/context';
import type { ChatFunctionClient } from '../chat_function_client';
import {
KnowledgeBaseEntryOperationType,
@ -271,7 +272,7 @@ export class ObservabilityAIAssistantClient {
]).pipe(
switchMap(([messagesWithUpdatedSystemMessage, knowledgeBaseInstructions]) => {
// if needed, inject a context function request here
const contextRequest = functionClient.hasFunction('context')
const contextRequest = functionClient.hasFunction(CONTEXT_FUNCTION_NAME)
? getContextFunctionRequestIfNeeded(messagesWithUpdatedSystemMessage)
: undefined;

View file

@ -21,6 +21,7 @@ import {
switchMap,
throwError,
} from 'rxjs';
import { CONTEXT_FUNCTION_NAME } from '../../../functions/context';
import { createFunctionNotFoundError, Message, MessageRole } from '../../../../common';
import {
createFunctionLimitExceededError,
@ -209,7 +210,7 @@ export function continueConversation({
function executeNextStep() {
if (isUserMessage) {
const operationName =
lastMessage.name && lastMessage.name !== 'context'
lastMessage.name && lastMessage.name !== CONTEXT_FUNCTION_NAME
? `function_response ${lastMessage.name}`
: 'user_message';

View file

@ -6,6 +6,7 @@
*/
import { Message } from '@kbn/observability-ai-assistant-plugin/common';
import { CONTEXT_FUNCTION_NAME } from '@kbn/observability-ai-assistant-plugin/server/functions/context';
import { reverseToLastUserMessage } from './chat_body';
describe('<ChatBody>', () => {
@ -38,7 +39,7 @@ describe('<ChatBody>', () => {
message: {
role: 'assistant',
function_call: {
name: 'context',
name: CONTEXT_FUNCTION_NAME,
arguments: '{"queries":[],"categories":[]}',
trigger: 'assistant',
},
@ -48,7 +49,7 @@ describe('<ChatBody>', () => {
{
message: {
role: 'user',
name: 'context',
name: CONTEXT_FUNCTION_NAME,
content: '[]',
},
},
@ -86,7 +87,7 @@ describe('<ChatBody>', () => {
message: {
role: 'assistant',
function_call: {
name: 'context',
name: CONTEXT_FUNCTION_NAME,
arguments: '{"queries":[],"categories":[]}',
trigger: 'assistant',
},
@ -96,7 +97,7 @@ describe('<ChatBody>', () => {
{
message: {
role: 'user',
name: 'context',
name: CONTEXT_FUNCTION_NAME,
content: '[]',
},
},

View file

@ -12,6 +12,7 @@ import { __IntlProvider as IntlProvider } from '@kbn/i18n-react';
import { ChatState, Message, MessageRole } from '@kbn/observability-ai-assistant-plugin/public';
import { createMockChatService } from './create_mock_chat_service';
import { KibanaContextProvider } from '@kbn/triggers-actions-ui-plugin/public/common/lib/kibana';
import { CONTEXT_FUNCTION_NAME } from '@kbn/observability-ai-assistant-plugin/server/functions/context';
const mockChatService = createMockChatService();
@ -135,7 +136,7 @@ describe('getTimelineItemsFromConversation', () => {
message: {
role: MessageRole.Assistant,
function_call: {
name: 'context',
name: CONTEXT_FUNCTION_NAME,
arguments: JSON.stringify({ queries: [], contexts: [] }),
trigger: MessageRole.Assistant,
},
@ -145,7 +146,7 @@ describe('getTimelineItemsFromConversation', () => {
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
name: 'context',
name: CONTEXT_FUNCTION_NAME,
content: JSON.stringify([]),
},
},
@ -430,7 +431,7 @@ describe('getTimelineItemsFromConversation', () => {
message: {
role: MessageRole.Assistant,
function_call: {
name: 'context',
name: CONTEXT_FUNCTION_NAME,
arguments: JSON.stringify({ queries: [], contexts: [] }),
trigger: MessageRole.User,
},
@ -440,7 +441,7 @@ describe('getTimelineItemsFromConversation', () => {
'@timestamp': new Date().toISOString(),
message: {
role: MessageRole.User,
name: 'context',
name: CONTEXT_FUNCTION_NAME,
content: JSON.stringify([]),
},
},

View file

@ -17,6 +17,8 @@ import {
import { getMetricChanges } from './get_metric_changes';
import { getLogChanges } from './get_log_changes';
export const CHANGES_FUNCTION_NAME = 'changes';
export function registerChangesFunction({
functions,
resources: {
@ -26,7 +28,7 @@ export function registerChangesFunction({
}: FunctionRegistrationParameters) {
functions.registerFunction(
{
name: 'changes',
name: CHANGES_FUNCTION_NAME,
description: 'Returns change points like spikes and dips for logs and metrics.',
parameters: changesFunctionParameters,
},

View file

@ -26,6 +26,8 @@ import type { FunctionRegistrationParameters } from '..';
import { correctCommonEsqlMistakes } from './correct_common_esql_mistakes';
import { runAndValidateEsqlQuery } from './validate_esql_query';
export const QUERY_FUNCTION_NAME = 'query';
const readFile = promisify(Fs.readFile);
const readdir = promisify(Fs.readdir);
@ -69,8 +71,8 @@ const loadEsqlDocs = once(async () => {
export function registerQueryFunction({ functions, resources }: FunctionRegistrationParameters) {
functions.registerInstruction(({ availableFunctionNames }) =>
availableFunctionNames.includes('query')
? `You MUST use the "query" function when the user wants to:
availableFunctionNames.includes(QUERY_FUNCTION_NAME)
? `You MUST use the "${QUERY_FUNCTION_NAME}" function when the user wants to:
- visualize data
- run any arbitrary query
- breakdown or filter ES|QL queries that are displayed on the current page
@ -78,11 +80,11 @@ export function registerQueryFunction({ functions, resources }: FunctionRegistra
- asks general questions about ES|QL
DO NOT UNDER ANY CIRCUMSTANCES generate ES|QL queries or explain anything about the ES|QL query language yourself.
DO NOT UNDER ANY CIRCUMSTANCES try to correct an ES|QL query yourself - always use the "query" function for this.
DO NOT UNDER ANY CIRCUMSTANCES try to correct an ES|QL query yourself - always use the "${QUERY_FUNCTION_NAME}" function for this.
If the user asks for a query, and one of the dataset info functions was called and returned no results, you should still call the query function to generate an example query.
Even if the "context" function was used before that, follow it up with the "query" function. If a query fails, do not attempt to correct it yourself. Again you should call the "query" function,
Even if the "${QUERY_FUNCTION_NAME}" function was used before that, follow it up with the "${QUERY_FUNCTION_NAME}" function. If a query fails, do not attempt to correct it yourself. Again you should call the "${QUERY_FUNCTION_NAME}" function,
even if it has been called before.
When the "visualize_query" function has been called, a visualization has been displayed to the user. DO NOT UNDER ANY CIRCUMSTANCES follow up a "visualize_query" function call with your own visualization attempt.
@ -132,7 +134,7 @@ export function registerQueryFunction({ functions, resources }: FunctionRegistra
);
functions.registerFunction(
{
name: 'query',
name: QUERY_FUNCTION_NAME,
description: `This function generates, executes and/or visualizes a query based on the user's request. It also explains how ES|QL works and how to convert queries from one language to another. Make sure you call one of the get_dataset functions first if you need index or field names. This function takes no input.`,
visibility: FunctionVisibility.AssistantOnly,
},
@ -159,7 +161,7 @@ export function registerQueryFunction({ functions, resources }: FunctionRegistra
await chat('classify_esql', {
messages: withEsqlSystemMessage().concat(
createFunctionResponseMessage({
name: 'query',
name: QUERY_FUNCTION_NAME,
content: {},
}).message,
{
@ -323,7 +325,7 @@ export function registerQueryFunction({ functions, resources }: FunctionRegistra
}
const queryFunctionResponseMessage = createFunctionResponseMessage({
name: 'query',
name: QUERY_FUNCTION_NAME,
content: {},
data: {
// add the included docs for debugging

View file

@ -8,6 +8,8 @@
import type { PluginConfigDescriptor, PluginInitializerContext } from '@kbn/core/server';
import type { ObservabilityAIAssistantAppConfig } from './config';
export { CHANGES_FUNCTION_NAME } from './functions/changes';
export { QUERY_FUNCTION_NAME } from './functions/query';
import { config as configSchema } from './config';
export type {
ObservabilityAIAssistantAppServerStart,