fix: switch to processing status when receiving response tokens from LLM
This commit is contained in:
@@ -0,0 +1,25 @@
|
|||||||
|
import { fireEvent, render, screen } from '@testing-library/react';
|
||||||
|
import { describe, expect, it, vi } from 'vitest';
|
||||||
|
import AssistantMessage from './AssistantMessage';
|
||||||
|
|
||||||
|
describe('AssistantMessage', () => {
|
||||||
|
it.each([
|
||||||
|
'<span class="processing-wave">Thinking...</span> click here to abort.',
|
||||||
|
'<span class="processing-wave">Processing...</span> 3 tokens received. click here to abort.'
|
||||||
|
])('renders the abort action for streaming status content: %s', (content) => {
|
||||||
|
const onAbort = vi.fn();
|
||||||
|
|
||||||
|
render(
|
||||||
|
<AssistantMessage
|
||||||
|
content={content}
|
||||||
|
isStreaming
|
||||||
|
onAbort={onAbort}
|
||||||
|
/>
|
||||||
|
);
|
||||||
|
|
||||||
|
const abortButton = screen.getByRole('button', { name: 'click here to abort' });
|
||||||
|
expect(abortButton).toBeInTheDocument();
|
||||||
|
fireEvent.click(abortButton);
|
||||||
|
expect(onAbort).toHaveBeenCalledTimes(1);
|
||||||
|
});
|
||||||
|
});
|
||||||
@@ -44,22 +44,26 @@ const AssistantMessage: React.FC<AssistantMessageProps> = ({ content, isStreamin
|
|||||||
const prefillTps = formatTps(performanceInfo?.prefillTps);
|
const prefillTps = formatTps(performanceInfo?.prefillTps);
|
||||||
const generationTps = formatTps(performanceInfo?.generationTps);
|
const generationTps = formatTps(performanceInfo?.generationTps);
|
||||||
const hasPerformanceInfo = Boolean(prefillTps || generationTps);
|
const hasPerformanceInfo = Boolean(prefillTps || generationTps);
|
||||||
|
const processingWaveLabels = ['Thinking...', 'Processing...'];
|
||||||
|
|
||||||
const renderContent = () => {
|
const renderContent = () => {
|
||||||
// Handle special abort link for streaming messages
|
// Handle special abort link for streaming messages
|
||||||
if (isStreaming && onAbort && content.includes('click here to abort')) {
|
if (isStreaming && onAbort && content.includes('click here to abort')) {
|
||||||
const hasProcessingWave = content.includes('<span class="processing-wave">Thinking...</span>');
|
const processingWaveMarkup = processingWaveLabels
|
||||||
|
.map(label => `<span class="processing-wave">${label}</span>`)
|
||||||
|
.find(markup => content.includes(markup));
|
||||||
|
|
||||||
if (hasProcessingWave) {
|
if (processingWaveMarkup) {
|
||||||
const parts = content.split('click here to abort');
|
const parts = content.split('click here to abort');
|
||||||
const beforeAbort = parts[0].replace(
|
const beforeAbort = parts[0].replace(
|
||||||
'<span class="processing-wave">Thinking...</span>',
|
processingWaveMarkup,
|
||||||
''
|
''
|
||||||
);
|
);
|
||||||
|
const waveLabel = processingWaveLabels.find(label => processingWaveMarkup.includes(label)) ?? 'Thinking...';
|
||||||
|
|
||||||
return (
|
return (
|
||||||
<span>
|
<span>
|
||||||
<span className="processing-wave">Thinking...</span>
|
<span className="processing-wave">{waveLabel}</span>
|
||||||
{beforeAbort}
|
{beforeAbort}
|
||||||
<button onClick={onAbort} className="abort-link">
|
<button onClick={onAbort} className="abort-link">
|
||||||
click here to abort
|
click here to abort
|
||||||
|
|||||||
@@ -0,0 +1,92 @@
|
|||||||
|
import { act, renderHook } from '@testing-library/react';
|
||||||
|
import { describe, expect, it, vi } from 'vitest';
|
||||||
|
|
||||||
|
vi.mock('../agent/core/AgentCore', () => ({
|
||||||
|
AgentCore: {
|
||||||
|
instance: vi.fn()
|
||||||
|
}
|
||||||
|
}));
|
||||||
|
|
||||||
|
vi.mock('../utils/chatMessageUtils', () => ({
|
||||||
|
createStreamingMessage: () => ({
|
||||||
|
id: 'streaming-message',
|
||||||
|
role: 'assistant',
|
||||||
|
content: '<span class="processing-wave">Thinking...</span> click here to abort.',
|
||||||
|
isStreaming: true,
|
||||||
|
tokenCount: 0
|
||||||
|
}),
|
||||||
|
createMessage: (role: 'user' | 'assistant', content: string) => ({
|
||||||
|
id: `${role}-message`,
|
||||||
|
role,
|
||||||
|
content
|
||||||
|
})
|
||||||
|
}));
|
||||||
|
|
||||||
|
import { AgentCore } from '../agent/core/AgentCore';
|
||||||
|
import { useStreamProcessor } from './useStreamProcessor';
|
||||||
|
import type { ChatMessage } from '../types/projectTypes';
|
||||||
|
|
||||||
|
const flushMicrotasks = async (): Promise<void> => {
|
||||||
|
await Promise.resolve();
|
||||||
|
await Promise.resolve();
|
||||||
|
};
|
||||||
|
|
||||||
|
describe('useStreamProcessor', () => {
|
||||||
|
it('switches from Thinking to Processing after the first text token arrives', async () => {
|
||||||
|
let releaseDone!: () => void;
|
||||||
|
const doneGate = new Promise<void>((resolve) => {
|
||||||
|
releaseDone = resolve;
|
||||||
|
});
|
||||||
|
|
||||||
|
vi.spyOn(AgentCore, 'instance').mockReturnValue({
|
||||||
|
processUserInput: async function* () {
|
||||||
|
yield { type: 'text', content: 'Hello' };
|
||||||
|
await doneGate;
|
||||||
|
yield { type: 'done', content: '' };
|
||||||
|
}
|
||||||
|
} as unknown as AgentCore);
|
||||||
|
|
||||||
|
const messages = new Map<string, ChatMessage>();
|
||||||
|
const processingChanges: boolean[] = [];
|
||||||
|
|
||||||
|
const { result } = renderHook(() => useStreamProcessor({
|
||||||
|
onMessageAdd: (message) => {
|
||||||
|
messages.set(message.id, message);
|
||||||
|
},
|
||||||
|
onMessageUpdate: (messageId, updater) => {
|
||||||
|
const current = messages.get(messageId);
|
||||||
|
if (!current) {
|
||||||
|
throw new Error(`Missing message ${messageId}`);
|
||||||
|
}
|
||||||
|
messages.set(messageId, updater(current));
|
||||||
|
},
|
||||||
|
onMessageRemove: (messageId) => {
|
||||||
|
messages.delete(messageId);
|
||||||
|
},
|
||||||
|
onProcessingChange: (isProcessing) => {
|
||||||
|
processingChanges.push(isProcessing);
|
||||||
|
}
|
||||||
|
}));
|
||||||
|
|
||||||
|
let responsePromise!: Promise<string>;
|
||||||
|
await act(async () => {
|
||||||
|
responsePromise = result.current.processStream('test prompt');
|
||||||
|
await flushMicrotasks();
|
||||||
|
});
|
||||||
|
|
||||||
|
const streamingMessage = [...messages.values()][0];
|
||||||
|
expect(streamingMessage).toBeDefined();
|
||||||
|
expect(streamingMessage.content).toContain('<span class="processing-wave">Processing...</span>');
|
||||||
|
expect(streamingMessage.content).toContain('1 tokens received.');
|
||||||
|
expect(streamingMessage.content).toContain('click here to abort.');
|
||||||
|
expect(streamingMessage.tokenCount).toBe(1);
|
||||||
|
expect(processingChanges).toContain(true);
|
||||||
|
|
||||||
|
await act(async () => {
|
||||||
|
releaseDone();
|
||||||
|
await responsePromise;
|
||||||
|
});
|
||||||
|
|
||||||
|
expect(processingChanges.at(-1)).toBe(false);
|
||||||
|
});
|
||||||
|
});
|
||||||
@@ -16,6 +16,8 @@ interface StreamProcessorResult {
|
|||||||
isProcessing: boolean;
|
isProcessing: boolean;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
const PROCESSING_WAVE = '<span class="processing-wave">Processing...</span>';
|
||||||
|
|
||||||
export const useStreamProcessor = (options: StreamProcessorOptions): StreamProcessorResult => {
|
export const useStreamProcessor = (options: StreamProcessorOptions): StreamProcessorResult => {
|
||||||
const { onMessageUpdate, onMessageAdd, onMessageRemove, onProcessingChange } = options;
|
const { onMessageUpdate, onMessageAdd, onMessageRemove, onProcessingChange } = options;
|
||||||
const [abortController, setAbortController] = useState<AbortController | null>(null);
|
const [abortController, setAbortController] = useState<AbortController | null>(null);
|
||||||
@@ -56,7 +58,7 @@ export const useStreamProcessor = (options: StreamProcessorOptions): StreamProce
|
|||||||
|
|
||||||
onMessageUpdate(currentStreamingId, (msg) => ({
|
onMessageUpdate(currentStreamingId, (msg) => ({
|
||||||
...msg,
|
...msg,
|
||||||
content: `<span class="processing-wave">Thinking...</span>${tokenCount > 0 ? ` ${tokenCount} tokens received.` : ''} click here to abort.`,
|
content: `${PROCESSING_WAVE}${tokenCount > 0 ? ` ${tokenCount} tokens received.` : ''} click here to abort.`,
|
||||||
tokenCount
|
tokenCount
|
||||||
}));
|
}));
|
||||||
} else if (chunk.type === 'tool_call' && chunk.toolCall) {
|
} else if (chunk.type === 'tool_call' && chunk.toolCall) {
|
||||||
|
|||||||
+11
-3
@@ -17,11 +17,19 @@ vi.mock('../core/KGCore', () => ({
|
|||||||
getName: vi.fn().mockReturnValue('Test Project'),
|
getName: vi.fn().mockReturnValue('Test Project'),
|
||||||
getBpm: vi.fn().mockReturnValue(120),
|
getBpm: vi.fn().mockReturnValue(120),
|
||||||
getTimeSignature: vi.fn().mockReturnValue({ numerator: 4, denominator: 4 }),
|
getTimeSignature: vi.fn().mockReturnValue({ numerator: 4, denominator: 4 }),
|
||||||
getTracks: vi.fn().mockReturnValue([])
|
getTracks: vi.fn().mockReturnValue([]),
|
||||||
|
getMaxBars: vi.fn().mockReturnValue(64),
|
||||||
|
getBarWidthMultiplier: vi.fn().mockReturnValue(1),
|
||||||
|
getIsLooping: vi.fn().mockReturnValue(false),
|
||||||
|
getLoopingRange: vi.fn().mockReturnValue([0, 0])
|
||||||
}),
|
}),
|
||||||
getSelectedItems: vi.fn().mockReturnValue([]),
|
getSelectedItems: vi.fn().mockReturnValue([]),
|
||||||
setSelectedItems: vi.fn(),
|
setSelectedItems: vi.fn(),
|
||||||
executeCommand: vi.fn()
|
executeCommand: vi.fn(),
|
||||||
|
setPlayheadUpdateCallback: vi.fn(),
|
||||||
|
setPlaybackStateChangeCallback: vi.fn(),
|
||||||
|
setSelectionChangeCallback: vi.fn(),
|
||||||
|
onSelectionChanged: vi.fn()
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
}));
|
}));
|
||||||
@@ -66,4 +74,4 @@ beforeAll(() => {
|
|||||||
// Mock URL.createObjectURL (might be needed for file operations)
|
// Mock URL.createObjectURL (might be needed for file operations)
|
||||||
global.URL.createObjectURL = vi.fn(() => 'mocked-url');
|
global.URL.createObjectURL = vi.fn(() => 'mocked-url');
|
||||||
global.URL.revokeObjectURL = vi.fn();
|
global.URL.revokeObjectURL = vi.fn();
|
||||||
});
|
});
|
||||||
|
|||||||
Reference in New Issue
Block a user