90 lines
2.4 KiB
JavaScript
90 lines
2.4 KiB
JavaScript
import { useState, useCallback, useRef } from 'react';
|
|
import { streamMessage } from '../api/orchestration';
|
|
|
|
export function useChat({ activeSession, appendMessage, updateLastMessage, refreshSessions }) {
|
|
const [streaming, setStreaming] = useState(false);
|
|
const [error, setError] = useState(null);
|
|
const [lastTokenCount, setLastTokenCount] = useState(0);
|
|
const [lastModel, setLastModel] = useState(null);
|
|
const cancelRef = useRef(null);
|
|
|
|
const sendMessage = useCallback(async (text, model) => {
|
|
if (!activeSession || !text.trim() || streaming) return;
|
|
|
|
setError(null);
|
|
|
|
// 1. Append user bubble immediately
|
|
appendMessage({
|
|
id: `user-${Date.now()}`,
|
|
role: 'user',
|
|
text,
|
|
});
|
|
|
|
// 2. Append empty assistant bubble — will be filled by stream
|
|
appendMessage({
|
|
id: `assistant-${Date.now()}`,
|
|
role: 'assistant',
|
|
text: '',
|
|
streaming: true,
|
|
});
|
|
|
|
setStreaming(true);
|
|
|
|
// 3. Open stream
|
|
cancelRef.current = streamMessage(
|
|
activeSession.external_id,
|
|
text,
|
|
model,
|
|
{
|
|
onChunk: (token) => {
|
|
updateLastMessage(msg => ({
|
|
...msg,
|
|
text: msg.text + token,
|
|
}));
|
|
},
|
|
|
|
onDone: ({ model: resolvedModel, tokenCount }) => {
|
|
// Mark bubble as complete
|
|
updateLastMessage(msg => ({ ...msg, streaming: false }));
|
|
setLastTokenCount(tokenCount);
|
|
setLastModel(resolvedModel);
|
|
setStreaming(false);
|
|
cancelRef.current = null;
|
|
|
|
// Refresh session list so new sessions appear in sidebar
|
|
refreshSessions();
|
|
},
|
|
|
|
onError: (err) => {
|
|
updateLastMessage(msg => ({
|
|
...msg,
|
|
text: msg.text || 'Something went wrong.',
|
|
streaming: false,
|
|
error: true,
|
|
}));
|
|
setError(err.message);
|
|
setStreaming(false);
|
|
cancelRef.current = null;
|
|
},
|
|
}
|
|
);
|
|
}, [activeSession, streaming, appendMessage, updateLastMessage, refreshSessions]);
|
|
|
|
const cancelStream = useCallback(() => {
|
|
if (cancelRef.current) {
|
|
cancelRef.current();
|
|
cancelRef.current = null;
|
|
updateLastMessage(msg => ({ ...msg, streaming: false }));
|
|
setStreaming(false);
|
|
}
|
|
}, [updateLastMessage]);
|
|
|
|
return {
|
|
sendMessage,
|
|
cancelStream,
|
|
streaming,
|
|
error,
|
|
lastTokenCount,
|
|
lastModel,
|
|
};
|
|
} |