mirror of https://github.com/grafana/grafana
Alexander Akhmetov
4 months ago
parent
e216c2f29d
commit
e816303323
@ -0,0 +1,92 @@ |
||||
import { render, screen } from '@testing-library/react'; |
||||
import { FormProvider, useForm } from 'react-hook-form'; |
||||
import { Provider } from 'react-redux'; |
||||
import { configureStore } from 'app/store/configureStore'; |
||||
|
||||
import { GenAIAlertDescriptionButton } from './GenAIAlertDescriptionButton'; |
||||
import { GenAIButton } from './GenAIButton'; |
||||
|
||||
// Mock the GenAIButton component
|
||||
jest.mock('./GenAIButton', () => ({ |
||||
GenAIButton: jest.fn(() => <div data-testid="gen-ai-button">GenAIButton</div>), |
||||
Role: { |
||||
system: 'system', |
||||
user: 'user', |
||||
}, |
||||
})); |
||||
|
||||
// Mock react-hook-form
|
||||
const mockWatch = jest.fn(); |
||||
const FormProviderWrapper: React.FC<React.PropsWithChildren<{}>> = ({ children }) => { |
||||
const methods = useForm({ |
||||
defaultValues: { |
||||
name: 'Test Alert Rule', |
||||
type: 'test', |
||||
annotations: [{ key: 'description', value: 'Test description' }], |
||||
labels: [{ key: 'test', value: 'label' }], |
||||
queries: [{ model: { test: 'query' } }], |
||||
}, |
||||
}); |
||||
|
||||
// Track when watch method is used
|
||||
mockWatch.mockImplementation(() => {}); |
||||
// No need to replace methods.watch
|
||||
|
||||
return <FormProvider {...methods}>{children}</FormProvider>; |
||||
}; |
||||
|
||||
const mockStore = configureStore(); |
||||
|
||||
describe('GenAIAlertDescriptionButton', () => { |
||||
beforeEach(() => { |
||||
jest.clearAllMocks(); |
||||
}); |
||||
|
||||
it('should render GenAIButton', () => { |
||||
render( |
||||
<Provider store={mockStore}> |
||||
<FormProviderWrapper> |
||||
<GenAIAlertDescriptionButton onGenerate={jest.fn()} /> |
||||
</FormProviderWrapper> |
||||
</Provider> |
||||
); |
||||
|
||||
expect(screen.getByTestId('gen-ai-button')).toBeInTheDocument(); |
||||
}); |
||||
|
||||
it('should call GenAIButton with proper props', () => { |
||||
const onGenerate = jest.fn(); |
||||
render( |
||||
<Provider store={mockStore}> |
||||
<FormProviderWrapper> |
||||
<GenAIAlertDescriptionButton onGenerate={onGenerate} /> |
||||
</FormProviderWrapper> |
||||
</Provider> |
||||
); |
||||
|
||||
// Check that form values are used
|
||||
expect(GenAIButton).toHaveBeenCalledWith( |
||||
expect.objectContaining({ |
||||
messages: expect.any(Array), |
||||
}), |
||||
expect.anything() |
||||
); |
||||
|
||||
// Check that the messages array exists
|
||||
const props = jest.mocked(GenAIButton).mock.calls[0][0]; |
||||
expect(props.messages).toBeDefined(); |
||||
|
||||
// Since messages can be a function or array, we'd need runtime type checking
|
||||
// to safely test its contents, which would be better in an integration test
|
||||
|
||||
// Verify GenAIButton was called with expected props
|
||||
expect(GenAIButton).toHaveBeenCalledWith( |
||||
expect.objectContaining({ |
||||
onGenerate, |
||||
toggleTipTitle: 'Improve your alert rule description', |
||||
tooltip: 'Generate a description for this alert rule', |
||||
}), |
||||
expect.anything() |
||||
); |
||||
}); |
||||
}); |
||||
@ -0,0 +1,85 @@ |
||||
import { useMemo } from 'react'; |
||||
import { useFormContext } from 'react-hook-form'; |
||||
|
||||
import { t } from 'app/core/internationalization'; |
||||
|
||||
import { RuleFormValues } from '../../types/rule-form'; |
||||
import { Annotation } from '../../utils/constants'; |
||||
|
||||
import { GenAIButton, Message, Role } from './GenAIButton'; |
||||
|
||||
interface GenAIAlertDescriptionButtonProps { |
||||
onGenerate: (description: string) => void; |
||||
} |
||||
|
||||
const DESCRIPTION_GENERATION_STANDARD_PROMPT = |
||||
'You are an expert in creating Grafana Alert Rules.\n' + |
||||
'Your goal is to write a descriptive and concise alert rule description.\n' + |
||||
'The description should explain the purpose of the alert rule, what it monitors, and why someone should care about it.\n' + |
||||
'The description should be clear and understandable by someone who is not familiar with the system.\n' + |
||||
'IMPORTANT: Only use official Grafana alerting template variables: $labels, $values, $value, $expr, $humanizeValue, $humanizeDuration.\n' + |
||||
'DO NOT create variables that do not exist like $valueInt, $valueString, etc.\n' + |
||||
'For labels, always use the format {{ $labels.labelname }} with double curly braces.\n' + |
||||
'For values, use {{ $value }} or {{ $value | formatter }} with double curly braces,' + |
||||
'or use {{ .Value.<refID>.Value }} where redID is the refID from the queries.\n' + |
||||
'Use only existing labels that you can see in the provided label information.\n' + |
||||
'Keep the description under 300 characters.\n' + |
||||
'Respond with only the description of the alert rule.'; |
||||
|
||||
export const GenAIAlertDescriptionButton = ({ onGenerate }: GenAIAlertDescriptionButtonProps) => { |
||||
const { watch } = useFormContext<RuleFormValues>(); |
||||
const name = watch('name'); |
||||
const type = watch('type'); |
||||
const annotations = watch('annotations'); |
||||
const queries = watch('queries'); |
||||
|
||||
// Get the summary annotation if it exists
|
||||
const summary = annotations.find((annotation) => annotation.key === Annotation.summary)?.value || ''; |
||||
|
||||
const messages = useMemo(() => { |
||||
const queryModels = queries.map((q) => q.model); |
||||
// Filter out expression queries which have refId property
|
||||
const dataQueries = queryModels.filter((q) => !q.refId); |
||||
|
||||
// Get other important context
|
||||
const labels = watch('labels') || []; |
||||
|
||||
// Format labels and annotations for the prompt
|
||||
const labelsText = labels.length > 0 ? `Labels:\n${labels.map((l) => `- ${l.key}: ${l.value}`).join('\n')}\n` : ''; |
||||
|
||||
const annotationsText = |
||||
annotations.length > 0 |
||||
? `Annotations:\n${annotations |
||||
.filter((a) => a.key !== Annotation.description) |
||||
.map((a) => `- ${a.key}: ${a.value}`) |
||||
.join('\n')}\n` |
||||
: ''; |
||||
|
||||
const messages: Message[] = [ |
||||
{ |
||||
content: DESCRIPTION_GENERATION_STANDARD_PROMPT, |
||||
role: Role.system, |
||||
}, |
||||
{ |
||||
content: |
||||
`Alert Rule Name: ${name}\n` + |
||||
`Alert Rule Type: ${type}\n` + |
||||
`Alert Rule Summary: ${summary}\n` + |
||||
`${labelsText}` + |
||||
`${annotationsText}` + |
||||
`Queries: ${JSON.stringify(dataQueries, null, 2)}\n`, |
||||
role: Role.user, |
||||
}, |
||||
]; |
||||
return messages; |
||||
}, [name, type, summary, annotations, queries, watch]); |
||||
|
||||
return ( |
||||
<GenAIButton |
||||
messages={messages} |
||||
onGenerate={onGenerate} |
||||
toggleTipTitle={t('alerting.rule-editor.gen-ai.description-tip-title', 'Improve your alert rule description')} |
||||
tooltip={t('alerting.rule-editor.gen-ai.description-tooltip', 'Generate a description for this alert rule')} |
||||
/> |
||||
); |
||||
}; |
||||
@ -0,0 +1,92 @@ |
||||
import { render, screen } from '@testing-library/react'; |
||||
import { FormProvider, useForm } from 'react-hook-form'; |
||||
import { Provider } from 'react-redux'; |
||||
import { configureStore } from 'app/store/configureStore'; |
||||
|
||||
import { GenAIAlertSummaryButton } from './GenAIAlertSummaryButton'; |
||||
import { GenAIButton } from './GenAIButton'; |
||||
|
||||
// Mock the GenAIButton component
|
||||
jest.mock('./GenAIButton', () => ({ |
||||
GenAIButton: jest.fn(() => <div data-testid="gen-ai-button">GenAIButton</div>), |
||||
Role: { |
||||
system: 'system', |
||||
user: 'user', |
||||
}, |
||||
})); |
||||
|
||||
// Mock react-hook-form
|
||||
const mockWatch = jest.fn(); |
||||
const FormProviderWrapper: React.FC<React.PropsWithChildren<{}>> = ({ children }) => { |
||||
const methods = useForm({ |
||||
defaultValues: { |
||||
name: 'Test Alert Rule', |
||||
type: 'test', |
||||
annotations: [{ key: 'description', value: 'Test description' }], |
||||
labels: [{ key: 'test', value: 'label' }], |
||||
queries: [{ model: { test: 'query' } }], |
||||
}, |
||||
}); |
||||
|
||||
// Track when watch method is used
|
||||
mockWatch.mockImplementation(() => {}); |
||||
// No need to replace methods.watch
|
||||
|
||||
return <FormProvider {...methods}>{children}</FormProvider>; |
||||
}; |
||||
|
||||
const mockStore = configureStore(); |
||||
|
||||
describe('GenAIAlertSummaryButton', () => { |
||||
beforeEach(() => { |
||||
jest.clearAllMocks(); |
||||
}); |
||||
|
||||
it('should render GenAIButton', () => { |
||||
render( |
||||
<Provider store={mockStore}> |
||||
<FormProviderWrapper> |
||||
<GenAIAlertSummaryButton onGenerate={jest.fn()} /> |
||||
</FormProviderWrapper> |
||||
</Provider> |
||||
); |
||||
|
||||
expect(screen.getByTestId('gen-ai-button')).toBeInTheDocument(); |
||||
}); |
||||
|
||||
it('should call GenAIButton with proper props', () => { |
||||
const onGenerate = jest.fn(); |
||||
render( |
||||
<Provider store={mockStore}> |
||||
<FormProviderWrapper> |
||||
<GenAIAlertSummaryButton onGenerate={onGenerate} /> |
||||
</FormProviderWrapper> |
||||
</Provider> |
||||
); |
||||
|
||||
// Check that form values are used
|
||||
expect(GenAIButton).toHaveBeenCalledWith( |
||||
expect.objectContaining({ |
||||
messages: expect.any(Array), |
||||
}), |
||||
expect.anything() |
||||
); |
||||
|
||||
// Check that the messages array exists
|
||||
const props = jest.mocked(GenAIButton).mock.calls[0][0]; |
||||
expect(props.messages).toBeDefined(); |
||||
|
||||
// Since messages can be a function or array, we'd need runtime type checking
|
||||
// to safely test its contents, which would be better in an integration test
|
||||
|
||||
// Verify GenAIButton was called with expected props
|
||||
expect(GenAIButton).toHaveBeenCalledWith( |
||||
expect.objectContaining({ |
||||
onGenerate, |
||||
toggleTipTitle: 'Improve your alert rule summary', |
||||
tooltip: 'Generate a summary for this alert rule', |
||||
}), |
||||
expect.anything() |
||||
); |
||||
}); |
||||
}); |
||||
@ -0,0 +1,81 @@ |
||||
import { useMemo } from 'react'; |
||||
import { useFormContext } from 'react-hook-form'; |
||||
|
||||
import { t } from 'app/core/internationalization'; |
||||
|
||||
import { RuleFormValues } from '../../types/rule-form'; |
||||
import { GenAIButton, Message, Role } from './GenAIButton'; |
||||
|
||||
interface GenAIAlertSummaryButtonProps { |
||||
onGenerate: (summary: string) => void; |
||||
} |
||||
|
||||
const SUMMARY_GENERATION_STANDARD_PROMPT = |
||||
'You are an expert in creating Grafana Alert Rules.\n' + |
||||
'Your goal is to write a short, concise summary for an alert rule.\n' + |
||||
'The summary should briefly explain what happened and why someone should care about it.\n' + |
||||
'IMPORTANT: Only use official Grafana alerting template variables: $labels, $values, $value, $expr, $humanizeValue, $humanizeDuration.\n' + |
||||
'DO NOT create variables that do not exist like $valueInt, $valueString, etc.\n' + |
||||
'For labels, always use the format {{ $labels.labelname }} with double curly braces.\n' + |
||||
'For values, use {{ $value }} or {{ $value | formatter }} with double curly braces,' + |
||||
'or use {{ .Value.<refID>.Value }} where redID is the refID from the queries.\n' + |
||||
'Use only existing labels that you can see in the provided label information.\n' + |
||||
'Keep the summary under 100 characters and make it actionable.\n' + |
||||
'Respond with only the summary text.'; |
||||
|
||||
export const GenAIAlertSummaryButton = ({ onGenerate }: GenAIAlertSummaryButtonProps) => { |
||||
const { watch } = useFormContext<RuleFormValues>(); |
||||
const name = watch('name'); |
||||
const type = watch('type'); |
||||
const queries = watch('queries'); |
||||
|
||||
const messages = useMemo(() => { |
||||
const queryModels = queries.map((q) => q.model); |
||||
// Filter out expression queries which have refId property
|
||||
const dataQueries = queryModels.filter((q) => !q.refId); |
||||
|
||||
// Get other important context
|
||||
const annotations = watch('annotations') || []; |
||||
const labels = watch('labels') || []; |
||||
const description = annotations.find((a) => a.key === 'description')?.value || ''; |
||||
|
||||
// Format labels and annotations for the prompt
|
||||
const labelsText = labels.length > 0 ? `Labels:\n${labels.map((l) => `- ${l.key}: ${l.value}`).join('\n')}\n` : ''; |
||||
|
||||
const annotationsText = |
||||
annotations.length > 0 ? `Annotations:\n${annotations.map((a) => `- ${a.key}: ${a.value}`).join('\n')}\n` : ''; |
||||
|
||||
const messages: Message[] = [ |
||||
{ |
||||
content: SUMMARY_GENERATION_STANDARD_PROMPT, |
||||
role: Role.system, |
||||
}, |
||||
{ |
||||
content: |
||||
`Alert Rule Name: ${name}\n` + |
||||
`Alert Rule Type: ${type}\n` + |
||||
`${description}\n` + |
||||
`${labelsText}` + |
||||
`${annotationsText}` + |
||||
`Queries: ${JSON.stringify(dataQueries, null, 2)}\n` + |
||||
`Examples of VALID Grafana template usage:\n` + |
||||
`1. CPU usage on {{ $labels.instance }} is high ({{ $value }}%)\n` + |
||||
`2. Memory usage exceeds threshold: {{ $value | humanizePercentage }}\n` + |
||||
`3. Disk I/O on {{ $labels.device }} is {{ $value | humanize1024 }}/s\n` + |
||||
`4. Service {{ $labels.service }} has high error rate: {{ $value }}\n` + |
||||
`5. Network usage above threshold: {{ $value | humanizeBytes }}\n`, |
||||
role: Role.user, |
||||
}, |
||||
]; |
||||
return messages; |
||||
}, [name, type, queries, watch]); |
||||
|
||||
return ( |
||||
<GenAIButton |
||||
messages={messages} |
||||
onGenerate={onGenerate} |
||||
toggleTipTitle={t('alerting.rule-editor.gen-ai.summary-tip-title', 'Improve your alert rule summary')} |
||||
tooltip={t('alerting.rule-editor.gen-ai.summary-tooltip', 'Generate a summary for this alert rule')} |
||||
/> |
||||
); |
||||
}; |
||||
@ -0,0 +1,359 @@ |
||||
import { fireEvent, screen, waitFor } from '@testing-library/react'; |
||||
import userEvent from '@testing-library/user-event'; |
||||
import { Observable } from 'rxjs'; |
||||
import { render } from 'test/test-utils'; |
||||
|
||||
import { selectors } from '@grafana/e2e-selectors'; |
||||
|
||||
import { GenAIButton, GenAIButtonProps, Role } from './GenAIButton'; |
||||
import { StreamStatus, useLLMStream } from '../../../../dashboard/components/GenAI/hooks'; |
||||
|
||||
const mockedUseLLMStreamState = { |
||||
messages: [], |
||||
setMessages: jest.fn(), |
||||
reply: 'I am a robot', |
||||
streamStatus: StreamStatus.IDLE, |
||||
error: null, |
||||
value: null, |
||||
}; |
||||
|
||||
jest.mock('../../../../dashboard/components/GenAI/hooks', () => ({ |
||||
useLLMStream: jest.fn(() => mockedUseLLMStreamState), |
||||
StreamStatus: { |
||||
IDLE: 'idle', |
||||
GENERATING: 'generating', |
||||
}, |
||||
})); |
||||
|
||||
describe('GenAIButton', () => { |
||||
const onGenerate = jest.fn(); |
||||
|
||||
function setup(props: GenAIButtonProps = { onGenerate, messages: [] }) { |
||||
return render(<GenAIButton text="Auto-generate" {...props} />); |
||||
} |
||||
|
||||
describe('when LLM plugin is not configured', () => { |
||||
beforeAll(() => { |
||||
jest.mocked(useLLMStream).mockReturnValue({ |
||||
messages: [], |
||||
error: undefined, |
||||
streamStatus: StreamStatus.IDLE, |
||||
reply: 'Some completed genereated text', |
||||
setMessages: jest.fn(), |
||||
stopGeneration: jest.fn(), |
||||
value: { |
||||
enabled: false, |
||||
stream: new Observable().subscribe(), |
||||
}, |
||||
}); |
||||
}); |
||||
|
||||
it('should not render anything', async () => { |
||||
setup(); |
||||
|
||||
waitFor(async () => expect(await screen.findByText('Auto-generate')).not.toBeInTheDocument()); |
||||
}); |
||||
}); |
||||
|
||||
describe('when LLM plugin is properly configured, so it is enabled', () => { |
||||
const setMessagesMock = jest.fn(); |
||||
const setShouldStopMock = jest.fn(); |
||||
beforeEach(() => { |
||||
setMessagesMock.mockClear(); |
||||
setShouldStopMock.mockClear(); |
||||
|
||||
jest.mocked(useLLMStream).mockReturnValue({ |
||||
messages: [], |
||||
error: undefined, |
||||
streamStatus: StreamStatus.IDLE, |
||||
reply: 'Some completed generated text', |
||||
setMessages: setMessagesMock, |
||||
stopGeneration: setShouldStopMock, |
||||
value: { |
||||
enabled: true, |
||||
stream: new Observable().subscribe(), |
||||
}, |
||||
}); |
||||
}); |
||||
|
||||
it('should render text', async () => { |
||||
setup(); |
||||
|
||||
waitFor(async () => expect(await screen.findByText('Auto-generate')).toBeInTheDocument()); |
||||
}); |
||||
|
||||
it('should enable the button', async () => { |
||||
setup(); |
||||
waitFor(async () => expect(await screen.findByRole('button')).toBeEnabled()); |
||||
}); |
||||
|
||||
it('should send the configured messages', async () => { |
||||
setup({ onGenerate, messages: [{ content: 'Generate X', role: 'system' as Role }] }); |
||||
const generateButton = await screen.findByRole('button'); |
||||
|
||||
// Click the button
|
||||
await fireEvent.click(generateButton); |
||||
await waitFor(() => expect(generateButton).toBeEnabled()); |
||||
|
||||
// Wait for the loading state to be resolved
|
||||
expect(setMessagesMock).toHaveBeenCalledTimes(1); |
||||
expect(setMessagesMock).toHaveBeenCalledWith([{ content: 'Generate X', role: 'system' as Role }]); |
||||
}); |
||||
|
||||
it('should call the messages when they are provided as callback', async () => { |
||||
const onGenerate = jest.fn(); |
||||
const messages = jest.fn().mockReturnValue([{ content: 'Generate X', role: 'system' as Role }]); |
||||
const onClick = jest.fn(); |
||||
setup({ onGenerate, messages, temperature: 3, onClick }); |
||||
|
||||
const generateButton = await screen.findByRole('button'); |
||||
await fireEvent.click(generateButton); |
||||
|
||||
expect(messages).toHaveBeenCalledTimes(1); |
||||
expect(setMessagesMock).toHaveBeenCalledTimes(1); |
||||
expect(setMessagesMock).toHaveBeenCalledWith([{ content: 'Generate X', role: 'system' as Role }]); |
||||
}); |
||||
|
||||
it('should call the onClick callback', async () => { |
||||
const onGenerate = jest.fn(); |
||||
const onClick = jest.fn(); |
||||
const messages = [{ content: 'Generate X', role: 'system' as Role }]; |
||||
setup({ onGenerate, messages, temperature: 3, onClick }); |
||||
|
||||
const generateButton = await screen.findByRole('button'); |
||||
await fireEvent.click(generateButton); |
||||
|
||||
await waitFor(() => expect(onClick).toHaveBeenCalledTimes(1)); |
||||
}); |
||||
|
||||
it('should display the tooltip if provided', async () => { |
||||
const { getByRole, getByTestId } = setup({ |
||||
tooltip: 'This is a tooltip', |
||||
onGenerate, |
||||
messages: [], |
||||
}); |
||||
|
||||
// Wait for the check to be completed
|
||||
const button = getByRole('button'); |
||||
await userEvent.hover(button); |
||||
|
||||
const tooltip = await waitFor(() => getByTestId(selectors.components.Tooltip.container)); |
||||
expect(tooltip).toBeVisible(); |
||||
expect(tooltip).toHaveTextContent('This is a tooltip'); |
||||
}); |
||||
}); |
||||
|
||||
describe('when it is generating data', () => { |
||||
const setShouldStopMock = jest.fn(); |
||||
|
||||
beforeEach(() => { |
||||
jest.mocked(useLLMStream).mockReturnValue({ |
||||
messages: [], |
||||
error: undefined, |
||||
streamStatus: StreamStatus.GENERATING, |
||||
reply: 'Some incomplete generated text', |
||||
setMessages: jest.fn(), |
||||
stopGeneration: setShouldStopMock, |
||||
value: { |
||||
enabled: true, |
||||
stream: new Observable().subscribe(), |
||||
}, |
||||
}); |
||||
}); |
||||
|
||||
it('should render loading text', async () => { |
||||
setup(); |
||||
|
||||
waitFor(async () => expect(await screen.findByText('Auto-generate')).toBeInTheDocument()); |
||||
}); |
||||
|
||||
it('should enable the button', async () => { |
||||
setup(); |
||||
waitFor(async () => expect(await screen.findByRole('button')).toBeEnabled()); |
||||
}); |
||||
|
||||
it('shows the stop button while generating', async () => { |
||||
const { getByText, getByRole } = setup(); |
||||
const generateButton = getByText('Stop generating'); |
||||
|
||||
expect(generateButton).toBeVisible(); |
||||
await waitFor(() => expect(getByRole('button')).toBeEnabled()); |
||||
}); |
||||
|
||||
it('should not call onGenerate when the text is generating', async () => { |
||||
const onGenerate = jest.fn(); |
||||
setup({ onGenerate, messages: [] }); |
||||
|
||||
await waitFor(() => expect(onGenerate).not.toHaveBeenCalledTimes(1)); |
||||
}); |
||||
|
||||
it('should stop generating when clicking the button', async () => { |
||||
const onGenerate = jest.fn(); |
||||
const { getByText } = setup({ onGenerate, messages: [] }); |
||||
const generateButton = getByText('Stop generating'); |
||||
|
||||
await fireEvent.click(generateButton); |
||||
|
||||
expect(setShouldStopMock).toHaveBeenCalledTimes(1); |
||||
expect(onGenerate).not.toHaveBeenCalled(); |
||||
}); |
||||
}); |
||||
|
||||
describe('when it is completed from generating data', () => { |
||||
const setShouldStopMock = jest.fn(); |
||||
|
||||
beforeEach(() => { |
||||
const reply = 'Some completed generated text'; |
||||
const returnValue = { |
||||
messages: [], |
||||
error: undefined, |
||||
streamStatus: StreamStatus.COMPLETED, |
||||
reply, |
||||
setMessages: jest.fn(), |
||||
stopGeneration: setShouldStopMock, |
||||
value: { |
||||
enabled: true, |
||||
stream: new Observable().subscribe(), |
||||
}, |
||||
}; |
||||
|
||||
jest |
||||
.mocked(useLLMStream) |
||||
.mockImplementationOnce((options) => { |
||||
options?.onResponse?.(reply); |
||||
return returnValue; |
||||
}) |
||||
.mockImplementation(() => returnValue); |
||||
}); |
||||
|
||||
it('should render improve text ', async () => { |
||||
setup(); |
||||
|
||||
waitFor(async () => expect(await screen.findByText('Improve')).toBeInTheDocument()); |
||||
}); |
||||
|
||||
it('should enable the button', async () => { |
||||
setup(); |
||||
waitFor(async () => expect(await screen.findByRole('button')).toBeEnabled()); |
||||
}); |
||||
|
||||
it('should call onGenerate when the text is completed', async () => { |
||||
const onGenerate = jest.fn(); |
||||
setup({ onGenerate, messages: [] }); |
||||
|
||||
await waitFor(() => expect(onGenerate).toHaveBeenCalledTimes(1)); |
||||
expect(onGenerate).toHaveBeenCalledWith('Some completed generated text'); |
||||
}); |
||||
}); |
||||
|
||||
describe('when there is an error generating data', () => { |
||||
const setMessagesMock = jest.fn(); |
||||
const setShouldStopMock = jest.fn(); |
||||
beforeEach(() => { |
||||
setMessagesMock.mockClear(); |
||||
setShouldStopMock.mockClear(); |
||||
|
||||
jest.mocked(useLLMStream).mockReturnValue({ |
||||
messages: [], |
||||
error: new Error('Something went wrong'), |
||||
streamStatus: StreamStatus.IDLE, |
||||
reply: '', |
||||
setMessages: setMessagesMock, |
||||
stopGeneration: setShouldStopMock, |
||||
value: { |
||||
enabled: true, |
||||
stream: new Observable().subscribe(), |
||||
}, |
||||
}); |
||||
}); |
||||
|
||||
it('should render error state text', async () => { |
||||
setup(); |
||||
|
||||
waitFor(async () => expect(await screen.findByText('Retry')).toBeInTheDocument()); |
||||
}); |
||||
|
||||
it('should enable the button', async () => { |
||||
setup(); |
||||
waitFor(async () => expect(await screen.findByRole('button')).toBeEnabled()); |
||||
}); |
||||
|
||||
it('should retry when clicking', async () => { |
||||
const onGenerate = jest.fn(); |
||||
const messages = [{ content: 'Generate X', role: 'system' as Role }]; |
||||
const { getByText } = setup({ onGenerate, messages, temperature: 3 }); |
||||
const generateButton = getByText('Retry'); |
||||
|
||||
await fireEvent.click(generateButton); |
||||
|
||||
expect(setMessagesMock).toHaveBeenCalledTimes(1); |
||||
expect(setMessagesMock).toHaveBeenCalledWith(messages); |
||||
}); |
||||
|
||||
it('should display the error message as tooltip', async () => { |
||||
const { getByRole, getByTestId } = setup(); |
||||
|
||||
// Wait for the check to be completed
|
||||
const button = getByRole('button'); |
||||
await userEvent.hover(button); |
||||
|
||||
const tooltip = await waitFor(() => getByTestId(selectors.components.Tooltip.container)); |
||||
expect(tooltip).toBeVisible(); |
||||
|
||||
// The tooltip keeps interactive to be able to click the link
|
||||
await userEvent.hover(tooltip); |
||||
expect(tooltip).toBeVisible(); |
||||
expect(tooltip).toHaveTextContent( |
||||
'Failed to generate content using LLM. Please try again or if the problem persists, contact your organization admin.' |
||||
); |
||||
}); |
||||
|
||||
it('error message should overwrite the tooltip content passed in tooltip prop', async () => { |
||||
const { getByRole, getByTestId } = setup({ |
||||
tooltip: 'This is a tooltip', |
||||
onGenerate, |
||||
messages: [], |
||||
}); |
||||
|
||||
// Wait for the check to be completed
|
||||
const button = getByRole('button'); |
||||
await userEvent.hover(button); |
||||
|
||||
const tooltip = await waitFor(() => getByTestId(selectors.components.Tooltip.container)); |
||||
expect(tooltip).toBeVisible(); |
||||
|
||||
// The tooltip keeps interactive to be able to click the link
|
||||
await userEvent.hover(tooltip); |
||||
expect(tooltip).toBeVisible(); |
||||
expect(tooltip).toHaveTextContent( |
||||
'Failed to generate content using LLM. Please try again or if the problem persists, contact your organization admin.' |
||||
); |
||||
}); |
||||
|
||||
it('should call the onClick callback', async () => { |
||||
const onGenerate = jest.fn(); |
||||
const onClick = jest.fn(); |
||||
const messages = [{ content: 'Generate X', role: 'system' as Role }]; |
||||
setup({ onGenerate, messages, temperature: 3, onClick }); |
||||
|
||||
const generateButton = await screen.findByRole('button'); |
||||
await fireEvent.click(generateButton); |
||||
|
||||
await waitFor(() => expect(onClick).toHaveBeenCalledTimes(1)); |
||||
}); |
||||
|
||||
it('should call the messages when they are provided as callback', async () => { |
||||
const onGenerate = jest.fn(); |
||||
const messages = jest.fn().mockReturnValue([{ content: 'Generate X', role: 'system' as Role }]); |
||||
const onClick = jest.fn(); |
||||
setup({ onGenerate, messages, temperature: 3, onClick }); |
||||
|
||||
const generateButton = await screen.findByRole('button'); |
||||
await fireEvent.click(generateButton); |
||||
|
||||
expect(messages).toHaveBeenCalledTimes(1); |
||||
expect(setMessagesMock).toHaveBeenCalledTimes(1); |
||||
expect(setMessagesMock).toHaveBeenCalledWith([{ content: 'Generate X', role: 'system' as Role }]); |
||||
}); |
||||
}); |
||||
}); |
||||
@ -0,0 +1,226 @@ |
||||
import { css } from '@emotion/css'; |
||||
import { useCallback, useState } from 'react'; |
||||
import * as React from 'react'; |
||||
|
||||
import { GrafanaTheme2 } from '@grafana/data'; |
||||
import { llm } from '@grafana/llm'; |
||||
import { Button, Spinner, Text, Toggletip, Tooltip, useStyles2 } from '@grafana/ui'; |
||||
import { Trans } from 'app/core/internationalization'; |
||||
|
||||
import { StreamStatus, useLLMStream } from '../../../../dashboard/components/GenAI/hooks'; |
||||
import { DEFAULT_LLM_MODEL, sanitizeReply } from '../../../../dashboard/components/GenAI/utils'; |
||||
|
||||
/** |
||||
* Message roles for LLM interactions |
||||
*/ |
||||
export enum Role { |
||||
// System content cannot be overwritten by user prompts.
|
||||
'system' = 'system', |
||||
// User content is the content that the user has entered.
|
||||
// This content can be overwritten by following prompt.
|
||||
'user' = 'user', |
||||
} |
||||
|
||||
const GenAIHistory = ({ |
||||
history, |
||||
onApplySuggestion, |
||||
}: { |
||||
history: string[]; |
||||
onApplySuggestion: (suggestion: string) => void; |
||||
}) => { |
||||
return ( |
||||
<div> |
||||
{history.map((suggestion, i) => ( |
||||
<div key={i} style={{ marginBottom: '8px' }}> |
||||
<Text>{suggestion}</Text> |
||||
<Button size="sm" variant="secondary" onClick={() => onApplySuggestion(suggestion)}> |
||||
<Trans>Apply</Trans> |
||||
</Button> |
||||
</div> |
||||
))} |
||||
</div> |
||||
); |
||||
}; |
||||
|
||||
export type Message = llm.Message; |
||||
|
||||
export interface GenAIButtonProps { |
||||
text?: string; |
||||
toggleTipTitle?: string; |
||||
onClick?: (e: React.MouseEvent<HTMLButtonElement>) => void; |
||||
// Messages to send to the LLM plugin
|
||||
messages: Message[] | (() => Message[]); |
||||
// Callback function that the LLM plugin streams responses to
|
||||
onGenerate: (response: string) => void; |
||||
// Temperature for the LLM plugin. Default is 1.
|
||||
// Closer to 0 means more conservative, closer to 1 means more creative.
|
||||
temperature?: number; |
||||
model?: llm.Model; |
||||
disabled?: boolean; |
||||
tooltip?: string; |
||||
} |
||||
|
||||
export const STOP_GENERATION_TEXT = <Trans>Stop generating</Trans>; |
||||
|
||||
export const GenAIButton = ({ |
||||
text = <Trans>Auto-generate</Trans>, |
||||
toggleTipTitle = '', |
||||
onClick: onClickProp, |
||||
model = DEFAULT_LLM_MODEL, |
||||
messages, |
||||
onGenerate, |
||||
temperature = 1, |
||||
disabled, |
||||
tooltip, |
||||
}: GenAIButtonProps) => { |
||||
const styles = useStyles2(getStyles); |
||||
|
||||
const [history, setHistory] = useState<string[]>([]); |
||||
const unshiftHistoryEntry = useCallback((historyEntry: string) => { |
||||
setHistory((h) => [historyEntry, ...h]); |
||||
}, []); |
||||
|
||||
const onResponse = useCallback( |
||||
(reply: string) => { |
||||
const sanitizedReply = sanitizeReply(reply); |
||||
onGenerate(sanitizedReply); |
||||
unshiftHistoryEntry(sanitizedReply); |
||||
}, |
||||
[onGenerate, unshiftHistoryEntry] |
||||
); |
||||
|
||||
const { setMessages, stopGeneration, value, error, streamStatus } = useLLMStream({ |
||||
model, |
||||
temperature, |
||||
onResponse, |
||||
}); |
||||
|
||||
const [showHistory, setShowHistory] = useState(false); |
||||
const hasHistory = history.length > 0; |
||||
const isFirstHistoryEntry = !hasHistory; |
||||
|
||||
const isGenerating = streamStatus === StreamStatus.GENERATING; |
||||
const isButtonDisabled = disabled || (value && !value.enabled && !error); |
||||
|
||||
const showTooltip = error || tooltip ? undefined : false; |
||||
const tooltipContent = error ? ( |
||||
<Trans> |
||||
Failed to generate content using LLM. Please try again or if the problem persists, contact your organization |
||||
admin. |
||||
</Trans> |
||||
) : ( |
||||
tooltip || '' |
||||
); |
||||
|
||||
const onClick = (e: React.MouseEvent<HTMLButtonElement>) => { |
||||
if (streamStatus === StreamStatus.GENERATING) { |
||||
stopGeneration(); |
||||
} else { |
||||
if (!hasHistory) { |
||||
onClickProp?.(e); |
||||
setMessages(getMessages()); |
||||
} else { |
||||
setShowHistory(true); |
||||
} |
||||
} |
||||
}; |
||||
|
||||
// The button is disabled if the plugin is not installed or enabled
|
||||
if (!value?.enabled) { |
||||
return null; |
||||
} |
||||
|
||||
const onApplySuggestion = (suggestion: string) => { |
||||
onGenerate(suggestion); |
||||
setShowHistory(false); |
||||
}; |
||||
|
||||
const getIcon = () => { |
||||
if (isGenerating) { |
||||
return undefined; |
||||
} |
||||
|
||||
if (error || (value && !value.enabled)) { |
||||
return 'exclamation-circle'; |
||||
} |
||||
|
||||
return 'ai'; |
||||
}; |
||||
|
||||
const getText = () => { |
||||
let buttonText = text; |
||||
|
||||
if (error) { |
||||
buttonText = <Trans>Retry</Trans>; |
||||
} |
||||
|
||||
if (isGenerating) { |
||||
buttonText = STOP_GENERATION_TEXT; |
||||
} |
||||
|
||||
if (hasHistory) { |
||||
buttonText = <Trans>Improve</Trans>; |
||||
} |
||||
|
||||
return buttonText; |
||||
}; |
||||
|
||||
const button = ( |
||||
<Button |
||||
icon={getIcon()} |
||||
onClick={onClick} |
||||
fill="text" |
||||
size="sm" |
||||
disabled={isButtonDisabled} |
||||
variant={error ? 'destructive' : 'primary'} |
||||
> |
||||
{getText()} |
||||
</Button> |
||||
); |
||||
|
||||
const getMessages = () => (typeof messages === 'function' ? messages() : messages); |
||||
|
||||
const renderButtonWithToggletip = () => { |
||||
if (hasHistory) { |
||||
const title = <Text element="p">{toggleTipTitle}</Text>; |
||||
|
||||
return ( |
||||
<Toggletip |
||||
title={title} |
||||
content={<GenAIHistory history={history} onApplySuggestion={onApplySuggestion} />} |
||||
placement="left-start" |
||||
fitContent={true} |
||||
show={showHistory} |
||||
onClose={() => setShowHistory(false)} |
||||
onOpen={() => setShowHistory(true)} |
||||
> |
||||
{button} |
||||
</Toggletip> |
||||
); |
||||
} |
||||
|
||||
return button; |
||||
}; |
||||
|
||||
return ( |
||||
<div className={styles.wrapper}> |
||||
{isGenerating && <Spinner size="sm" className={styles.spinner} />} |
||||
{isFirstHistoryEntry ? ( |
||||
<Tooltip show={showTooltip} interactive content={tooltipContent}> |
||||
{button} |
||||
</Tooltip> |
||||
) : ( |
||||
renderButtonWithToggletip() |
||||
)} |
||||
</div> |
||||
); |
||||
}; |
||||
|
||||
const getStyles = (theme: GrafanaTheme2) => ({ |
||||
wrapper: css({ |
||||
display: 'flex', |
||||
}), |
||||
spinner: css({ |
||||
color: theme.colors.text.link, |
||||
}), |
||||
}); |
||||
Loading…
Reference in new issue