Xiaohan-Tian
333 lines
11 KiB
TypeScript
333 lines
11 KiB
TypeScript
import { useState, useCallback } from 'react';
|
|
import { AgentCore } from '../agent/core/AgentCore';
|
|
import { KGCore } from '../core/KGCore';
|
|
import { KGMidiRegion } from '../core/region/KGMidiRegion';
|
|
import { createStreamingMessage, createMessage } from '../utils/chatMessageUtils';
|
|
import { useProjectStore } from '../stores/projectStore';
|
|
import type { ChatMessage } from '../types/projectTypes';
|
|
|
|
const TODO_TOOL_NAME = 'update_todo_list';
|
|
const ADD_NOTES_TOOL_NAME = 'add_notes';
|
|
|
|
interface PendingToolCall {
|
|
name: string;
|
|
arguments: Record<string, unknown> | null;
|
|
}
|
|
|
|
interface AddNotesToolArguments {
|
|
notes: Array<{
|
|
pitch: string;
|
|
start: number;
|
|
length: number;
|
|
velocity?: number;
|
|
}>;
|
|
region_id?: string;
|
|
}
|
|
|
|
const getBarNumberFromStartBeat = (beat: number, beatsPerBar: number): number => (
|
|
Math.floor(beat / beatsPerBar) + 1
|
|
);
|
|
|
|
const getBarNumberFromEndBeat = (beat: number, beatsPerBar: number): number => (
|
|
Math.max(1, Math.ceil(beat / beatsPerBar))
|
|
);
|
|
|
|
interface AddNotesSummaryData {
|
|
noteCount: number;
|
|
regionName: string;
|
|
trackName: string;
|
|
earliestNoteStartBar: number;
|
|
latestNoteEndBar: number;
|
|
}
|
|
|
|
const resolveTargetMidiRegion = (
|
|
regionId: string | undefined,
|
|
storeState: ReturnType<typeof useProjectStore.getState>,
|
|
): { region: KGMidiRegion; trackName: string } | undefined => {
|
|
const project = KGCore.instance().getCurrentProject();
|
|
const tracks = project.getTracks();
|
|
|
|
const findById = (candidateRegionId: string): { region: KGMidiRegion; trackName: string } | undefined => {
|
|
for (const track of tracks) {
|
|
const region = track.getRegions().find(candidate => candidate.getId() === candidateRegionId);
|
|
if (region instanceof KGMidiRegion) {
|
|
return {
|
|
region,
|
|
trackName: track.getName(),
|
|
};
|
|
}
|
|
}
|
|
|
|
return undefined;
|
|
};
|
|
|
|
if (regionId) {
|
|
return findById(regionId);
|
|
}
|
|
|
|
if (storeState.activeRegionId) {
|
|
const activeRegion = findById(storeState.activeRegionId);
|
|
if (activeRegion) {
|
|
return activeRegion;
|
|
}
|
|
}
|
|
|
|
const selectedRegionId = storeState.selectedRegionIds.at(-1);
|
|
if (selectedRegionId) {
|
|
const selectedRegion = findById(selectedRegionId);
|
|
if (selectedRegion) {
|
|
return selectedRegion;
|
|
}
|
|
}
|
|
|
|
const selectedItems = KGCore.instance().getSelectedItems();
|
|
for (const item of selectedItems) {
|
|
if (item instanceof KGMidiRegion) {
|
|
const track = tracks.find(candidate => candidate.getId() === item.getTrackId());
|
|
return {
|
|
region: item,
|
|
trackName: track?.getName() ?? `Track ${item.getTrackIndex() + 1}`,
|
|
};
|
|
}
|
|
}
|
|
|
|
return undefined;
|
|
};
|
|
|
|
const buildAddNotesSummary = (args: Record<string, unknown> | null): AddNotesSummaryData | undefined => {
|
|
if (!args) {
|
|
return undefined;
|
|
}
|
|
|
|
const typedArgs = args as AddNotesToolArguments;
|
|
if (!Array.isArray(typedArgs.notes) || typedArgs.notes.length === 0) {
|
|
return undefined;
|
|
}
|
|
|
|
const storeState = useProjectStore.getState();
|
|
const project = KGCore.instance().getCurrentProject();
|
|
const beatsPerBar = project.getTimeSignature().numerator ?? storeState.timeSignature.numerator;
|
|
const targetRegion = resolveTargetMidiRegion(typedArgs.region_id, storeState);
|
|
if (!targetRegion) {
|
|
return undefined;
|
|
}
|
|
|
|
const earliestNoteStartBeat = Math.min(...typedArgs.notes.map(note => note.start));
|
|
const latestNoteEndBeat = Math.max(...typedArgs.notes.map(note => note.start + note.length));
|
|
|
|
return {
|
|
noteCount: typedArgs.notes.length,
|
|
regionName: targetRegion.region.getName(),
|
|
trackName: targetRegion.trackName,
|
|
earliestNoteStartBar: getBarNumberFromStartBeat(earliestNoteStartBeat, beatsPerBar),
|
|
latestNoteEndBar: getBarNumberFromEndBeat(latestNoteEndBeat, beatsPerBar),
|
|
};
|
|
};
|
|
|
|
const buildToolResultDisplayContent = (
|
|
toolName: string,
|
|
success: boolean,
|
|
rawResult: string,
|
|
toolArgs: Record<string, unknown> | null,
|
|
): string => {
|
|
if (!success) {
|
|
return rawResult;
|
|
}
|
|
|
|
if (toolName === ADD_NOTES_TOOL_NAME) {
|
|
const summary = buildAddNotesSummary(toolArgs);
|
|
if (summary) {
|
|
return `Successfully created ${summary.noteCount} ${summary.noteCount === 1 ? 'note' : 'notes'} in region **${summary.regionName}** on track **${summary.trackName}**, spanning bars ${summary.earliestNoteStartBar} to ${summary.latestNoteEndBar}.`;
|
|
}
|
|
}
|
|
|
|
return rawResult;
|
|
};
|
|
|
|
interface StreamProcessorOptions {
|
|
onMessageUpdate: (messageId: string, updater: (msg: ChatMessage) => ChatMessage) => void;
|
|
onMessageAdd: (message: ChatMessage) => void;
|
|
onMessageRemove: (messageId: string) => void;
|
|
onProcessingChange: (isProcessing: boolean) => void;
|
|
}
|
|
|
|
interface StreamProcessorResult {
|
|
processStream: (input: string, logPrefix?: string) => Promise<string>;
|
|
abortController: AbortController | null;
|
|
isProcessing: boolean;
|
|
}
|
|
|
|
const PROCESSING_WAVE = '<span class="processing-wave">Processing...</span>';
|
|
|
|
export const useStreamProcessor = (options: StreamProcessorOptions): StreamProcessorResult => {
|
|
const { onMessageUpdate, onMessageAdd, onMessageRemove, onProcessingChange } = options;
|
|
const [abortController, setAbortController] = useState<AbortController | null>(null);
|
|
const [isProcessing, setIsProcessing] = useState(false);
|
|
|
|
const processStream = useCallback(async (input: string, logPrefix: string = 'USER'): Promise<string> => {
|
|
setIsProcessing(true);
|
|
onProcessingChange(true);
|
|
|
|
const controller = new AbortController();
|
|
setAbortController(controller);
|
|
|
|
// Track the current streaming message ID (mutable)
|
|
const initialStreamingMsg = createStreamingMessage();
|
|
let currentStreamingId = initialStreamingMsg.id;
|
|
const pendingToolCalls: PendingToolCall[] = [];
|
|
onMessageAdd(initialStreamingMsg);
|
|
|
|
try {
|
|
const agentCore = AgentCore.instance();
|
|
let assistantResponse = '';
|
|
let tokenCount = 0;
|
|
let hasTextContent = false;
|
|
let performanceInfo: ChatMessage['performanceInfo'];
|
|
|
|
console.log(`------------ ${logPrefix} ------------`);
|
|
console.log(input);
|
|
console.log('------------------------------');
|
|
|
|
for await (const chunk of agentCore.processUserInput(input)) {
|
|
if (controller.signal.aborted) {
|
|
return '';
|
|
}
|
|
|
|
if (chunk.type === 'text') {
|
|
assistantResponse += chunk.content;
|
|
tokenCount++;
|
|
hasTextContent = true;
|
|
|
|
onMessageUpdate(currentStreamingId, (msg) => ({
|
|
...msg,
|
|
content: `${PROCESSING_WAVE}${tokenCount > 0 ? ` ${tokenCount} tokens received.` : ''} click here to abort.`,
|
|
tokenCount
|
|
}));
|
|
} else if (chunk.type === 'tool_call' && chunk.toolCall) {
|
|
// Finalize or remove the current streaming message
|
|
if (hasTextContent) {
|
|
onMessageUpdate(currentStreamingId, (msg) => ({
|
|
...msg,
|
|
content: assistantResponse,
|
|
isStreaming: false,
|
|
tokenCount: undefined,
|
|
performanceInfo
|
|
}));
|
|
|
|
console.log('------------ ASSISTANT ------------');
|
|
console.log(assistantResponse);
|
|
console.log('-----------------------------------');
|
|
} else {
|
|
// No text before this tool call — remove the empty streaming placeholder
|
|
onMessageRemove(currentStreamingId);
|
|
}
|
|
|
|
// Show tool call in UI
|
|
const toolName = chunk.toolCall.function.name;
|
|
if (toolName === TODO_TOOL_NAME) {
|
|
continue;
|
|
}
|
|
let argsDisplay = '';
|
|
try {
|
|
const args = JSON.parse(chunk.toolCall.function.arguments);
|
|
argsDisplay = JSON.stringify(args, null, 2);
|
|
pendingToolCalls.push({ name: toolName, arguments: args });
|
|
} catch {
|
|
argsDisplay = chunk.toolCall.function.arguments;
|
|
pendingToolCalls.push({ name: toolName, arguments: null });
|
|
}
|
|
const toolCallMsg = {
|
|
...createMessage('assistant', `🔧 **Calling tool: ${toolName}**\n\n\`\`\`json\n${argsDisplay}\n\`\`\``),
|
|
isToolCallMessage: true,
|
|
};
|
|
onMessageAdd(toolCallMsg);
|
|
} else if (chunk.type === 'tool_result' && chunk.toolResult) {
|
|
// Show tool result in UI
|
|
const { name, success, result } = chunk.toolResult;
|
|
const pendingToolCallIndex = pendingToolCalls.findIndex(toolCall => toolCall.name === name);
|
|
const pendingToolCall = pendingToolCallIndex >= 0
|
|
? pendingToolCalls.splice(pendingToolCallIndex, 1)[0]
|
|
: undefined;
|
|
const toolResultDisplayContent = buildToolResultDisplayContent(
|
|
name,
|
|
success,
|
|
result,
|
|
pendingToolCall?.arguments ?? null,
|
|
);
|
|
const toolResultMsg = name === TODO_TOOL_NAME
|
|
? {
|
|
...createMessage('assistant', result),
|
|
toolName: name,
|
|
toolSuccess: success,
|
|
todoSnapshot: AgentCore.instance().getAgentState().getTodos().map(todo => ({ ...todo })),
|
|
}
|
|
: {
|
|
...createMessage('assistant', `${success ? '✅' : '❌'} **${name}**\n\n └── ${result}`),
|
|
toolName: name,
|
|
toolSuccess: success,
|
|
toolRawResult: result,
|
|
toolResultDisplayContent,
|
|
};
|
|
onMessageAdd(toolResultMsg);
|
|
|
|
// Reset for the next LLM turn in the agentic loop
|
|
assistantResponse = '';
|
|
tokenCount = 0;
|
|
hasTextContent = false;
|
|
performanceInfo = undefined;
|
|
|
|
// Create a fresh streaming placeholder for the next LLM response
|
|
const nextMsg = createStreamingMessage();
|
|
currentStreamingId = nextMsg.id;
|
|
onMessageAdd(nextMsg);
|
|
} else if (chunk.type === 'done') {
|
|
performanceInfo = chunk.performanceInfo;
|
|
// Finalize the streaming message
|
|
if (hasTextContent) {
|
|
onMessageUpdate(currentStreamingId, (msg) => ({
|
|
...msg,
|
|
content: assistantResponse,
|
|
isStreaming: false,
|
|
tokenCount: undefined,
|
|
performanceInfo
|
|
}));
|
|
} else {
|
|
// No text in final response — remove empty placeholder
|
|
onMessageRemove(currentStreamingId);
|
|
}
|
|
|
|
console.log('------------ ASSISTANT ------------');
|
|
console.log(assistantResponse);
|
|
console.log('-----------------------------------');
|
|
break;
|
|
}
|
|
}
|
|
|
|
return assistantResponse;
|
|
} catch (error) {
|
|
if (error instanceof Error && error.name === 'AbortError') {
|
|
return '';
|
|
}
|
|
|
|
console.error('Error processing stream:', error);
|
|
onMessageUpdate(currentStreamingId, (msg) => ({
|
|
...msg,
|
|
content: `Error: ${error instanceof Error ? error.message : 'Failed to process message'}`,
|
|
isStreaming: false,
|
|
tokenCount: undefined
|
|
}));
|
|
throw error;
|
|
} finally {
|
|
setAbortController(null);
|
|
setIsProcessing(false);
|
|
onProcessingChange(false);
|
|
}
|
|
}, [onMessageUpdate, onMessageAdd, onMessageRemove, onProcessingChange]);
|
|
|
|
return {
|
|
processStream,
|
|
abortController,
|
|
isProcessing
|
|
};
|
|
};
|